diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 021f5d678..baf64a2cc 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -80,6 +80,26 @@ steps: - group: ":card_index_dividers: E2E Test" depends_on: upload-merge-pipeline steps: + - label: "TTS · Breeze-TTS-2 Voice Design" + timeout_in_minutes: 20 + source_file_dependencies: + - tests/e2e/online_serving/test_breeze_tts_2.py + - tests/helpers/client.py + - tests/helpers/media.py + - vllm_omni/model_executor/models/breeze_tts_2/ + - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py + - vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py + - vllm_omni/deploy/breeze_tts_2.yaml + - vllm_omni/deploy/breeze_tts_2_throughput.yaml + - vllm_omni/deploy/breeze_tts.yaml + - vllm_omni/deploy/breeze_tts_throughput.yaml + - vllm_omni/config/pipeline_registry.py + - vllm_omni/engine/arg_utils.py + commands: + - pytest -s -v tests/e2e/online_serving/test_breeze_tts_2.py -m 'advanced_model and tts' --run-level=advanced_model + mirror_hardwares: l4_1 + - label: "TTS · Qwen3-TTS CustomVoice Test" source_file_dependencies: tts_qwen3_tts_function commands: diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 5fe28722f..7a83acc0b 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -198,6 +198,25 @@ steps: " mirror_hardwares: l4_1 + - label: "TTS · Breeze-TTS-2 Voice Design" + timeout_in_minutes: 20 + source_file_dependencies: + - tests/e2e/online_serving/test_breeze_tts_2.py + - tests/helpers/client.py + - vllm_omni/model_executor/models/breeze_tts_2/ + - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py + - vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py + - vllm_omni/deploy/breeze_tts_2.yaml + - vllm_omni/deploy/breeze_tts_2_throughput.yaml + - vllm_omni/deploy/breeze_tts.yaml + - vllm_omni/deploy/breeze_tts_throughput.yaml + - vllm_omni/config/pipeline_registry.py + - vllm_omni/engine/arg_utils.py + commands: + - pytest -s -v tests/e2e/online_serving/test_breeze_tts_2.py -m 'core_model and tts' --run-level=core_model + mirror_hardwares: l4_1 + - label: "Diffusion · Wan22 Test" source_file_dependencies: diffusion_wan22_function commands: diff --git a/benchmarks/tts/evidence/breeze_tts_2_20260918.json b/benchmarks/tts/evidence/breeze_tts_2_20260918.json new file mode 100644 index 000000000..03d1caa43 --- /dev/null +++ b/benchmarks/tts/evidence/breeze_tts_2_20260918.json @@ -0,0 +1,757 @@ +{ + "measurement_scope": "Historical measurements for PR7752, not a main-branch comparison or a measurement of subsequent review fixes.", + "environment": { + "gpu": "NVIDIA H100 80 GB", + "gpu_count": 1, + "python": "3.12", + "pytorch": "2.13.0+cu132", + "vllm": "0.29.0", + "transformers": "5.14.1", + "nsight_systems": "2026.5.1" + }, + "model": { + "id": "BreezeBlue/Breeze-TTS-2", + "revision": "3e28c5151381a722f1d8661b4118c298caa77aa4" + }, + "workload": { + "texts": [ + "Hello, this is Breeze TTS 2 running on vLLM Omni.", + "你好，今天天气很好，我们一起出去走走吧。" + ], + "alternating_texts": true, + "seed": 42, + "guidance_scale": 1.0, + "repetition_penalty": 1.1, + "top_p": 1.0, + "sampled": { + "temperature": 0.9, + "top_k": 50 + }, + "greedy": { + "temperature": 0.0, + "top_k": 0 + }, + "stream": false, + "response_format": "pcm", + "sample_rate": 24000, + "channels": 1, + "pcm_bits": 16 + }, + "cumulative_diagnostic": { + "baseline_sha": "5bda85c371e8c0c7728c2f7f5865d3ddd79b9d99", + "combined_sha": "37a04897cf8566c63e82eb13c27c064b9bcd4222", + "intermediate_variants": "Ordered patch subsets of the combined commit, not independent commits.", + "patch_sha256": { + "batched_eos": "2f107ea8e890e03cc138ef2f69e693e3474d3aaf0eb4051c00886cbaa9785b3d", + "greedy_graph": "4e783313260ddc2e36c8a74ba9d95dc5d78338443612ebf89854e31b8be6933b", + "batched_embedding": "90557e4ec73fdedcff51fd717d5e8c6477e71087fd97ac52534542234f5d1d3a" + }, + "sequence_slots": 8, + "kv_gib": 2, + "codec_chunk_frames": 16, + "max_new_tokens": 48, + "temperature": 0.0, + "cohorts_per_variant": 1, + "nsight_injection_present": true, + "collection_enabled": false, + "warmup": "Two eight-request full-response warmups, then four full-response 64-request cohorts, each preceded by one concurrency-sized warmup.", + "variants": [ + { + "variant": "baseline", + "makespan_s": 0.9672094583511353, + "requests": 8, + "concurrency": 8, + "audio_s": 30.72, + "failures": 0 + }, + { + "variant": "plus_batched_eos", + "makespan_s": 0.9280132218264043, + "requests": 8, + "concurrency": 8, + "audio_s": 30.72, + "failures": 0 + }, + { + "variant": "plus_greedy_graph", + "makespan_s": 0.8857334200292826, + "requests": 8, + "concurrency": 8, + "audio_s": 30.72, + "failures": 0 + }, + { + "variant": "plus_batched_embedding", + "makespan_s": 0.8435833756811917, + "requests": 8, + "concurrency": 8, + "audio_s": 30.72, + "failures": 0 + } + ], + "limits": [ + "Single cumulative diagnostic samples; no repeat spread.", + "Nsight injection remains when collection is disabled.", + "The final variant used graph tracing configuration; prior variants used node tracing configuration.", + "Do not interpret these timings as isolated uninstrumented per-optimization gains." + ] + }, + "mechanical_attribution": { + "batched_eos": { + "stage0_stream_synchronization_calls": [ + 773, + 438 + ], + "device_to_host_copies": [ + 841, + 506 + ], + "capture_mode": "node" + }, + "greedy_graph": { + "radix_sort_kernels": [ + 735, + 0 + ], + "capture_mode": "node" + }, + "combined_graph_capture": { + "captured_wall_ms": [ + 955.4114495404065, + 855.6996420957148 + ], + "gpu_envelope_ms": [ + 926.633897, + 826.396835 + ], + "activity_upper_ms": [ + 651.130806, + 628.864685 + ], + "idle_lower_ms": [ + 275.50309100000004, + 197.53215 + ], + "stage0_explicit_kernel_launch_count": [ + 7087, + 6198 + ], + "stage0_graph_launch_count": [ + 156, + 156 + ] + }, + "limits": [ + "Host API durations overlap GPU activity and are not additional idle time.", + "Graph envelopes hide internal gaps: activity is an upper bound and idle is a lower bound.", + "Stage 1 has eight versus nine graph launches across the graph-mode traces.", + "No separate final-capacity stage breakdown was measured." + ] + }, + "combined_uninstrumented_ab": { + "baseline_sha": "5bda85c371e8c0c7728c2f7f5865d3ddd79b9d99", + "head_sha": "37a04897cf8566c63e82eb13c27c064b9bcd4222", + "sequence_slots": 8, + "kv_gib": 2, + "codec_chunk_frames": 16, + "requests_per_group": 64, + "max_new_tokens": 384, + "baseline_rounds": 2, + "head_rounds": 1, + "profiler_configured": false, + "nsight_injection_present": false, + "warmup": "One excluded 64-request full-response wave for each temperature/concurrency family, plus an excluded 2*C-request full-response cohort before every group.", + "rows": [ + { + "build": "plain-baseline", + "phase": "round-0", + "temperature": 0.0, + "concurrency": 8, + "requests": 64, + "makespan_s": 10.77213971177116, + "req_per_s": 5.941252314994074, + "aggregate_rtf": 0.030157166046391827, + "audio_s": 357.2, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 305.88, + "endpoint_trimmed_aggregate_rtf": 0.03521688149526337, + "quiet_tail_at_least_3s": 7, + "max_quiet_tail_s": 3.24 + }, + { + "build": "plain-baseline", + "phase": "round-0", + "temperature": 0.0, + "concurrency": 16, + "requests": 64, + "makespan_s": 11.076066234149039, + "req_per_s": 5.778224745774739, + "aggregate_rtf": 0.029092420240988227, + "audio_s": 380.72, + "caps": 2, + "failures": 0, + "endpoint_trimmed_audio_s": 315.48, + "endpoint_trimmed_aggregate_rtf": 0.035108616185333584, + "quiet_tail_at_least_3s": 6, + "max_quiet_tail_s": 26.04 + }, + { + "build": "plain-baseline", + "phase": "round-0", + "temperature": 0.9, + "concurrency": 8, + "requests": 64, + "makespan_s": 8.603743681218475, + "req_per_s": 7.438622345259851, + "aggregate_rtf": 0.03389435739528236, + "audio_s": 253.84, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 239.02, + "endpoint_trimmed_aggregate_rtf": 0.035995915325991445, + "quiet_tail_at_least_3s": 0, + "max_quiet_tail_s": 1.3 + }, + { + "build": "plain-baseline", + "phase": "round-0", + "temperature": 0.9, + "concurrency": 16, + "requests": 64, + "makespan_s": 8.66502623911947, + "req_per_s": 7.386013409984042, + "aggregate_rtf": 0.033143460216950235, + "audio_s": 261.44, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 242.14, + "endpoint_trimmed_aggregate_rtf": 0.03578519137325296, + "quiet_tail_at_least_3s": 0, + "max_quiet_tail_s": 1.3 + }, + { + "build": "plain-baseline", + "phase": "round-1", + "temperature": 0.0, + "concurrency": 8, + "requests": 64, + "makespan_s": 13.904579282738268, + "req_per_s": 4.602800178172402, + "aggregate_rtf": 0.0392784725501081, + "audio_s": 354.0, + "caps": 1, + "failures": 0, + "endpoint_trimmed_audio_s": 285.36, + "endpoint_trimmed_aggregate_rtf": 0.048726448285457906, + "quiet_tail_at_least_3s": 8, + "max_quiet_tail_s": 25.38 + }, + { + "build": "plain-baseline", + "phase": "round-1", + "temperature": 0.0, + "concurrency": 16, + "requests": 64, + "makespan_s": 10.060411099344492, + "req_per_s": 6.361569061941222, + "aggregate_rtf": 0.02935460754944121, + "audio_s": 342.72, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 289.04, + "endpoint_trimmed_aggregate_rtf": 0.03480629359031446, + "quiet_tail_at_least_3s": 10, + "max_quiet_tail_s": 3.16 + }, + { + "build": "plain-baseline", + "phase": "round-1", + "temperature": 0.9, + "concurrency": 8, + "requests": 64, + "makespan_s": 8.535738762002438, + "req_per_s": 7.4978864494894575, + "aggregate_rtf": 0.033363581777683074, + "audio_s": 255.84, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 237.9, + "endpoint_trimmed_aggregate_rtf": 0.03587952401009852, + "quiet_tail_at_least_3s": 0, + "max_quiet_tail_s": 2.12 + }, + { + "build": "plain-baseline", + "phase": "round-1", + "temperature": 0.9, + "concurrency": 16, + "requests": 64, + "makespan_s": 8.584469985216856, + "req_per_s": 7.4553234049642105, + "aggregate_rtf": 0.033355882752629996, + "audio_s": 257.36, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 242.57999999999998, + "endpoint_trimmed_aggregate_rtf": 0.035388201769382704, + "quiet_tail_at_least_3s": 0, + "max_quiet_tail_s": 1.94 + }, + { + "build": "plain-batched", + "phase": "round-0", + "temperature": 0.0, + "concurrency": 8, + "requests": 64, + "makespan_s": 12.983129982370883, + "req_per_s": 4.929473870083891, + "aggregate_rtf": 0.03320153943936907, + "audio_s": 391.04, + "caps": 2, + "failures": 0, + "endpoint_trimmed_audio_s": 312.72, + "endpoint_trimmed_aggregate_rtf": 0.041516788124747, + "quiet_tail_at_least_3s": 11, + "max_quiet_tail_s": 28.4 + }, + { + "build": "plain-batched", + "phase": "round-0", + "temperature": 0.0, + "concurrency": 16, + "requests": 64, + "makespan_s": 9.772029620129615, + "req_per_s": 6.549304749155183, + "aggregate_rtf": 0.027698496655696188, + "audio_s": 352.8, + "caps": 1, + "failures": 0, + "endpoint_trimmed_audio_s": 286.26, + "endpoint_trimmed_aggregate_rtf": 0.034136902187275955, + "quiet_tail_at_least_3s": 6, + "max_quiet_tail_s": 26.04 + }, + { + "build": "plain-batched", + "phase": "round-0", + "temperature": 0.9, + "concurrency": 8, + "requests": 64, + "makespan_s": 8.171920725144446, + "req_per_s": 7.831696140061215, + "aggregate_rtf": 0.03169376638669115, + "audio_s": 257.84, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 241.72, + "endpoint_trimmed_aggregate_rtf": 0.03380738344011437, + "quiet_tail_at_least_3s": 0, + "max_quiet_tail_s": 1.22 + }, + { + "build": "plain-batched", + "phase": "round-0", + "temperature": 0.9, + "concurrency": 16, + "requests": 64, + "makespan_s": 8.02551719499752, + "req_per_s": 7.974563937124526, + "aggregate_rtf": 0.030886380830501534, + "audio_s": 259.84000000000003, + "caps": 0, + "failures": 0, + "endpoint_trimmed_audio_s": 243.0, + "endpoint_trimmed_aggregate_rtf": 0.03302681973250008, + "quiet_tail_at_least_3s": 0, + "max_quiet_tail_s": 1.78 + } + ], + "limits": [ + "Ranges compare the head sample with two baseline samples; they are not confidence intervals.", + "Full-response greedy output lengths and frame-limit incidence vary, so no general greedy throughput gain is claimed.", + "No transcription was run on these historical A/B outputs." + ] + }, + "capacity": { + "runtime_sha": "7bc0b3e3a2e41ef44e8efb30fd927e93bf6a48a7", + "main_base_sha": "6bda549ac23a667efbe44a54cf3f6d57e6e24da2", + "sequence_slots": [ + 8, + 16, + 32 + ], + "kv_gib": [ + 2, + 4, + 8 + ], + "codec_graph_buckets": [ + [ + 1, + 2, + 4, + 8 + ], + [ + 1, + 2, + 4, + 8, + 16 + ], + [ + 1, + 2, + 4, + 8, + 16, + 32 + ] + ], + "backbone_capture_sizes": [ + 1, + 2, + 4, + 8, + 32, + 64, + 128, + 256 + ], + "codec_chunk_frames": 16, + "requests_per_group": 128, + "cohorts_per_point": 1, + "excluded_warmups_per_configuration": { + "cohorts": 4, + "requests_per_cohort": 64, + "covers": "Both client concurrency levels and both sampling settings" + }, + "profiler_configured": false, + "nsight_injection_present": false, + "rows": [ + { + "slots": 8, + "concurrency": 8, + "observed_peak_client_requests": 8, + "temperature": 0.0, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 6.7812081269297915, + "makespan_s": 18.87569259107113, + "aggregate_rtf": 0.02841355459879445, + "audio_s": 664.32, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 8, + "concurrency": 16, + "observed_peak_client_requests": 16, + "temperature": 0.0, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 5.8798058352164615, + "makespan_s": 21.76942633604631, + "aggregate_rtf": 0.029975526459636356, + "audio_s": 726.24, + "near_cap_count": 2, + "exact_cap_count": 2, + "fixed48_verified": null + }, + { + "slots": 8, + "concurrency": 8, + "observed_peak_client_requests": 8, + "temperature": 0.9, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 7.72991106982811, + "makespan_s": 16.559052082709968, + "aggregate_rtf": 0.0320365502296664, + "audio_s": 516.88, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 8, + "concurrency": 16, + "observed_peak_client_requests": 16, + "temperature": 0.9, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 7.945057322591292, + "makespan_s": 16.110645248088986, + "aggregate_rtf": 0.030958196095482296, + "audio_s": 520.4, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 8, + "concurrency": 8, + "observed_peak_client_requests": 8, + "temperature": 0.0, + "max_new_tokens": 48, + "requests": 128, + "failures": 0, + "req_per_s": 9.850579526557848, + "makespan_s": 12.994159344118088, + "aggregate_rtf": 0.026436684863521503, + "audio_s": 491.52, + "near_cap_count": 128, + "exact_cap_count": 128, + "fixed48_verified": true + }, + { + "slots": 8, + "concurrency": 16, + "observed_peak_client_requests": 16, + "temperature": 0.0, + "max_new_tokens": 48, + "requests": 128, + "failures": 0, + "req_per_s": 10.665146583658053, + "makespan_s": 12.00171033712104, + "aggregate_rtf": 0.02441754218978076, + "audio_s": 491.52, + "near_cap_count": 128, + "exact_cap_count": 128, + "fixed48_verified": true + }, + { + "slots": 16, + "concurrency": 16, + "observed_peak_client_requests": 16, + "temperature": 0.0, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 7.754874623816668, + "makespan_s": 16.505747186020017, + "aggregate_rtf": 0.023979758231665527, + "audio_s": 688.32, + "near_cap_count": 1, + "exact_cap_count": 1, + "fixed48_verified": null + }, + { + "slots": 16, + "concurrency": 32, + "observed_peak_client_requests": 32, + "temperature": 0.0, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 10.104109318018903, + "makespan_s": 12.668113138061017, + "aggregate_rtf": 0.018577125085143444, + "audio_s": 681.92, + "near_cap_count": 1, + "exact_cap_count": 1, + "fixed48_verified": null + }, + { + "slots": 16, + "concurrency": 16, + "observed_peak_client_requests": 16, + "temperature": 0.9, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 10.920674745703012, + "makespan_s": 11.72088748915121, + "aggregate_rtf": 0.021831484669108944, + "audio_s": 536.88, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 16, + "concurrency": 32, + "observed_peak_client_requests": 32, + "temperature": 0.9, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 11.210975235811437, + "makespan_s": 11.417383172083646, + "aggregate_rtf": 0.02120613516360261, + "audio_s": 538.4, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 16, + "concurrency": 16, + "observed_peak_client_requests": 16, + "temperature": 0.0, + "max_new_tokens": 48, + "requests": 128, + "failures": 0, + "req_per_s": 15.14487085057202, + "makespan_s": 8.451706274878234, + "aggregate_rtf": 0.017195040435543284, + "audio_s": 491.52, + "near_cap_count": 128, + "exact_cap_count": 128, + "fixed48_verified": true + }, + { + "slots": 16, + "concurrency": 32, + "observed_peak_client_requests": 32, + "temperature": 0.0, + "max_new_tokens": 48, + "requests": 128, + "failures": 0, + "req_per_s": 16.563937910502617, + "makespan_s": 7.727630995213985, + "aggregate_rtf": 0.015721905507840955, + "audio_s": 491.52, + "near_cap_count": 128, + "exact_cap_count": 128, + "fixed48_verified": true + }, + { + "slots": 32, + "concurrency": 32, + "observed_peak_client_requests": 32, + "temperature": 0.0, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 15.246181432777554, + "makespan_s": 8.395544849336147, + "aggregate_rtf": 0.012390119317202106, + "audio_s": 677.6, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 32, + "concurrency": 64, + "observed_peak_client_requests": 64, + "temperature": 0.0, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 11.166074835908939, + "makespan_s": 11.463294119108468, + "aggregate_rtf": 0.015826284127331107, + "audio_s": 724.32, + "near_cap_count": 1, + "exact_cap_count": 1, + "fixed48_verified": null + }, + { + "slots": 32, + "concurrency": 32, + "observed_peak_client_requests": 32, + "temperature": 0.9, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 17.17219888865503, + "makespan_s": 7.453908543102443, + "aggregate_rtf": 0.014336645143680651, + "audio_s": 519.92, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 32, + "concurrency": 64, + "observed_peak_client_requests": 64, + "temperature": 0.9, + "max_new_tokens": 384, + "requests": 128, + "failures": 0, + "req_per_s": 17.514343665613783, + "makespan_s": 7.3082955572754145, + "aggregate_rtf": 0.013677750331777613, + "audio_s": 534.32, + "near_cap_count": 0, + "exact_cap_count": 0, + "fixed48_verified": null + }, + { + "slots": 32, + "concurrency": 32, + "observed_peak_client_requests": 32, + "temperature": 0.0, + "max_new_tokens": 48, + "requests": 128, + "failures": 0, + "req_per_s": 25.82908453652114, + "makespan_s": 4.955653763841838, + "aggregate_rtf": 0.010082303393232905, + "audio_s": 491.52, + "near_cap_count": 128, + "exact_cap_count": 128, + "fixed48_verified": true + }, + { + "slots": 32, + "concurrency": 64, + "observed_peak_client_requests": 64, + "temperature": 0.0, + "max_new_tokens": 48, + "requests": 128, + "failures": 0, + "req_per_s": 28.81314277607874, + "makespan_s": 4.442417163401842, + "aggregate_rtf": 0.00903812085653044, + "audio_s": 491.52, + "near_cap_count": 128, + "exact_cap_count": 128, + "fixed48_verified": true + } + ], + "limits": [ + "This scales slots, KV reservation and codec graph buckets together.", + "One cohort per point; no repeated capacity samples or confidence intervals.", + "The 48-frame cohorts deliberately truncate speech and are excluded from full-utterance quality claims.", + "Peak device memory was not measured.", + "No capacity claim extends to cloning, CFG guidance above one, long prompts or other GPUs." + ] + }, + "quality": { + "full_response_pcm_checked": 1536, + "sampled_full_responses": 768, + "sampled_frame_limit_hits": 0, + "sampled_quiet_tails_at_least_three_seconds": 0, + "greedy_frame_limit_hits": 5, + "asr_selected_unique_clips": 47, + "asr_selected_sampled_clips": 24, + "sampled_asr": "All selected sampled transcripts matched after accounting for TTS 2 versus TTS2 spacing.", + "endpoint_method": "Trim leading/trailing quiet using 20 ms RMS windows at or below -40 dBFS; keep internal pauses and original makespans.", + "limits": [ + "ASR selects median-duration and longest English/Chinese clips per group, deduplicated by WAV hash.", + "Selected ASR and endpoint measurements do not establish perceptual quality parity.", + "Capped greedy clips have extra ASR words with unreliable tail alignment.", + "No listening or MOS assessment was performed." + ] + }, + "metrics": { + "requests_per_second": "Completed requests divided by cohort wall time", + "aggregate_rtf": "Cohort wall time divided by total returned audio duration", + "timing_exclusions": [ + "PCM validation", + "hashing", + "WAV file writes", + "post-timing audio analysis" + ] + } +} diff --git a/benchmarks/tts/evidence/breeze_tts_2_sampling_graphs_20260918.json b/benchmarks/tts/evidence/breeze_tts_2_sampling_graphs_20260918.json new file mode 100644 index 000000000..81aa68b0a --- /dev/null +++ b/benchmarks/tts/evidence/breeze_tts_2_sampling_graphs_20260918.json @@ -0,0 +1,2735 @@ +{ + "schema_version": 1, + "performance_audit_status": "complete", + "quality_audit_status": "targeted_checks_complete_with_severe_ASR_mismatches_in_baseline_and_final", + "code": { + "baseline_commit": "ec6f20e7ce82eb5aade6c0ea4c9a616b5e80bf16", + "candidate_commit": "ca22df8b3a6459154271873006caaebc829ec5a9", + "all_seven_changed_runtime_and_test_files_match_tested_candidate": true, + "candidate_gpu_tests": { + "passed": 50, + "skipped": 0 + } + }, + "hardware": { + "gpu_count": 1, + "gpu_model": "NVIDIA H100 80GB HBM3", + "driver_version": "580.159.04" + }, + "runtime": { + "python": "3.12.3", + "packages": { + "torch": "2.13.0+cu132", + "vllm": "0.29.0", + "transformers": "5.14.1", + "triton": "3.7.1", + "numpy": "2.3.5", + "flashinfer-python": "0.6.18" + }, + "vllm_omni_distribution_metadata_version": "0.29.0rc1", + "vllm_omni_runtime_source": "Code commits above; distribution metadata does not replace the source revision." + }, + "model": { + "id": "BreezeBlue/Breeze-TTS-2", + "revision": "3e28c5151381a722f1d8661b4118c298caa77aa4", + "config_sha256": "57849eb756ca1602efe89afcc5c70379dd127f8deeb9f655204d4340190ea929" + }, + "deployment": { + "repository_path": "vllm_omni/deploy/breeze_tts_2_throughput.yaml", + "rendered_sha256": "add2dd15fda6a8049b5fc940e94cb61e6dd87b8e8c7c6a84cc899d641b8fcef3", + "rendered_yaml": "async_chunk: true\nconnectors:\n shm:\n name: SharedMemoryConnector\n extra:\n codec_streaming: true\n codec_chunk_frames: 16\n initial_codec_chunk_frames: 16\n codec_chunk_ramp:\n - 16\n codec_left_context_frames: 0\n decode_cudagraph_batch_sizes:\n - 1\n - 2\n - 4\n - 8\n - 16\n - 32\n connector_get_sleep_s: 0.001\n connector_get_max_wait_first_chunk: 1000\n connector_get_max_wait: 300\nstages:\n- stage_id: 0\n devices: '0'\n max_num_seqs: 32\n max_model_len: 2048\n max_num_batched_tokens: 4096\n gpu_memory_utilization: 0.65\n kv_cache_memory_bytes: 8589934592\n enforce_eager: false\n compilation_config:\n mode: 3\n cudagraph_mode: FULL_AND_PIECEWISE\n cudagraph_capture_sizes:\n - 1\n - 2\n - 4\n - 8\n - 32\n - 64\n - 128\n - 256\n enable_prefix_caching: false\n enable_chunked_prefill: false\n async_scheduling: false\n engine_extras:\n additional_config:\n cfg_prefill_delay_tokens: 0\n output_connectors:\n to_stage_1: shm\n default_sampling_params:\n temperature: 0.9\n top_p: 1.0\n top_k: 50\n repetition_penalty: 1.1\n max_tokens: 1500\n seed: 42\n- stage_id: 1\n devices: '0'\n max_num_seqs: 32\n max_model_len: 32768\n max_num_batched_tokens: 32768\n gpu_memory_utilization: 0.15\n enforce_eager: false\n enable_prefix_caching: false\n async_scheduling: true\n input_connectors:\n from_stage_0: shm\n default_sampling_params:\n temperature: 0.0\n top_p: 1.0\n top_k: -1\n repetition_penalty: 1.0\n max_tokens: 32768\n" + }, + "dataset": { + "repository_path": "benchmarks/build_dataset/seed_tts_smoke/en/meta.lst", + "sha256": "17e5a91c6889ffef16957f7a5367897494a1b9f989cce195d2681c8e9c537bb3", + "texts": 30, + "language": "English", + "kind": "Checked-in synthetic smoke texts; not original Seed test-en recordings." + }, + "protocol": { + "api": "POST /v1/audio/speech", + "stream": false, + "response_format": "pcm", + "audio_sample_rate_hz": 24000, + "channels": 1, + "sample_width_bytes": 2, + "server_slots_per_stage": 32, + "client_concurrencies": [ + 32, + 64 + ], + "closed_loop": true, + "rounds_per_arm": 3, + "requests_per_cohort": 300, + "measured_cohorts": 24, + "measured_requests_per_arm": 3600, + "request_order": "30 texts cyclically repeated ten times; complete request bodies match across arms, repeated rounds and client concurrency.", + "seed": "42 + zero-based request index; resets for each cohort.", + "natural_sampled": { + "temperature": 0.9, + "top_k": 50, + "top_p": 1.0, + "max_new_tokens": 384 + }, + "capped_greedy": { + "temperature": 0.0, + "top_k": 0, + "top_p": 1.0, + "max_new_tokens": 48, + "minimum_tokens": null, + "fixed_work_valid": false + }, + "shared_sampling": { + "repetition_penalty": 1.1, + "guidance_scale": 1.0 + }, + "warmup": "Fresh server per arm. Before timing: 150 requests each in sampled C32, capped-greedy C32, sampled C64, capped-greedy C64 order; 600 excluded warmup requests per arm.", + "cohort_order": "Round 0: sampled C32, greedy C32, sampled C64, greedy C64. Rotate left one position in round 1 and two in round 2.", + "arms_order": "All baseline rounds and C1 control precede candidate validation, final rounds and final C1 control; arms are sequential, not interleaved.", + "timing": "Request start through complete PCM receipt. Cohort wall is latest response end minus earliest request start. PCM joining, analysis, hashing and writes follow completion of all cohort requests.", + "profiling": "Both timing arms uninjected; separate profiles excluded.", + "asr": "No ASR during timing arms; quality work follows both servers stopping.", + "memory_polling": "Identical 1 Hz management polling in both arms; whole-run scope stated separately.", + "source_guards": "Six identical execution helpers and full source/config/model/runtime/GPU checks; final timing bound to the exact successful 50-test candidate.", + "workload_scope": { + "reference_audio": false, + "guidance_scales": [ + 1.0 + ], + "texts": "30 synthetic English smoke texts; not original Seed-TTS or a published blog benchmark.", + "untested_performance_scopes": [ + "voice cloning", + "guidance scale greater than one", + "long prompts", + "multiple GPUs" + ], + "startup_and_warmup": "Server startup and all four warmup waves are excluded from measured makespans." + } + }, + "metric_definitions": { + "rps": "Successful responses / cohort makespan.", + "rtf": "Cohort makespan / total full returned PCM duration, including quiet tails.", + "p95_response_s": "Per-cohort statistics.quantiles(request wall seconds,n=20,method=\"inclusive\")[18]. Summary is median/minimum/maximum of three cohort p95s.", + "summary": "Three per-cohort observations per arm/family: median and observed full min/max ranges, no confidence intervals.", + "total_audio_percent_change": "Change in SUMMED audio duration across all three rounds: 900 responses per arm/family.", + "near_cap": "PCM-derived frames >= max_new_tokens-1; exact-cap counts reported separately.", + "frames": "PCM byte length /3840 (1920 mono int16 samples/frame). No RVQ token sequence recorded.", + "interior": "Last initial request start through last request start; counts completions and their full delivered clips. This is not GPU work isolated to the window." + }, + "individual_measured_cohorts": [ + { + "arm": "baseline", + "round": 0, + "mode": "greedy48", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 12.409082953352481, + "audio_s": 1138.1599999999999, + "rps": 24.175839675481495, + "rtf": 0.010902757919231464, + "p95_response_s": 1.3660403012065216, + "mean_duration_s": 3.7938666666666663, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 276, + "at_max_frames": 274, + "exact48": 274, + "pcm_derived_frames": 14227.0, + "peak_inflight": 32, + "interior_rps": 23.55702600182998, + "interior_rtf": 0.011197488491907423, + "interior_mean_inflight": 31.999404196554725, + "fixed48_valid": false, + "max_response_s": 1.3778274049982429 + }, + { + "arm": "baseline", + "round": 0, + "mode": "greedy48", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 11.4890843112953, + "audio_s": 1137.6, + "rps": 26.11174153409772, + "rtf": 0.010099406040168162, + "p95_response_s": 2.568560073385015, + "mean_duration_s": 3.792, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 275, + "at_max_frames": 269, + "exact48": 269, + "pcm_derived_frames": 14220.0, + "peak_inflight": 64, + "interior_rps": 25.370435925769396, + "interior_rtf": 0.010412561279692619, + "interior_mean_inflight": 63.99934963725613, + "fixed48_valid": false, + "max_response_s": 2.6073009888641536 + }, + { + "arm": "baseline", + "round": 0, + "mode": "sampled", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 19.73926813295111, + "audio_s": 1372.48, + "rps": 15.19813186483873, + "rtf": 0.014382190001275873, + "p95_response_s": 2.6196906621102243, + "mean_duration_s": 4.574933333333333, + "median_duration_s": 4.56, + "min_duration_s": 2.64, + "max_duration_s": 7.76, + "caps": 0, + "at_max_frames": 0, + "exact48": 11, + "pcm_derived_frames": 17156.0, + "peak_inflight": 32, + "interior_rps": 14.769683037456982, + "interior_rtf": 0.014883914366667608, + "interior_mean_inflight": 31.99964170182187, + "fixed48_valid": false, + "max_response_s": 3.423801572062075 + }, + { + "arm": "baseline", + "round": 0, + "mode": "sampled", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 19.836500237695873, + "audio_s": 1368.8, + "rps": 15.123635540804791, + "rtf": 0.01449189088084152, + "p95_response_s": 4.799736037221737, + "mean_duration_s": 4.562666666666667, + "median_duration_s": 4.48, + "min_duration_s": 2.88, + "max_duration_s": 8.72, + "caps": 0, + "at_max_frames": 0, + "exact48": 5, + "pcm_derived_frames": 17110.0, + "peak_inflight": 64, + "interior_rps": 14.826302139784737, + "interior_rtf": 0.014957955135430911, + "interior_mean_inflight": 63.99961896669195, + "fixed48_valid": false, + "max_response_s": 5.861642682924867 + }, + { + "arm": "baseline", + "round": 1, + "mode": "greedy48", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 12.368925544898957, + "audio_s": 1140.8799999999999, + "rps": 24.25432984546684, + "rtf": 0.010841565760552344, + "p95_response_s": 1.3458482225192712, + "mean_duration_s": 3.802933333333333, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 278, + "at_max_frames": 276, + "exact48": 276, + "pcm_derived_frames": 14261.0, + "peak_inflight": 32, + "interior_rps": 23.594010932160216, + "interior_rtf": 0.011157971292675595, + "interior_mean_inflight": 31.99941429464469, + "fixed48_valid": false, + "max_response_s": 1.3915124712511897 + }, + { + "arm": "baseline", + "round": 1, + "mode": "greedy48", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 11.264604894910008, + "audio_s": 1140.8, + "rps": 26.632092541084788, + "rtf": 0.009874303028497552, + "p95_response_s": 2.4218830559635536, + "mean_duration_s": 3.8026666666666666, + "median_duration_s": 3.84, + "min_duration_s": 2.8, + "max_duration_s": 3.84, + "caps": 282, + "at_max_frames": 282, + "exact48": 282, + "pcm_derived_frames": 14260.0, + "peak_inflight": 64, + "interior_rps": 25.794812158986286, + "interior_rtf": 0.010210167083200961, + "interior_mean_inflight": 63.999343316676544, + "fixed48_valid": false, + "max_response_s": 2.4386735460720956 + }, + { + "arm": "baseline", + "round": 1, + "mode": "sampled", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 20.184422803111374, + "audio_s": 1385.04, + "rps": 14.862946685488366, + "rtf": 0.01457316958579635, + "p95_response_s": 2.753425784688443, + "mean_duration_s": 4.6168, + "median_duration_s": 4.52, + "min_duration_s": 2.4, + "max_duration_s": 8.0, + "caps": 0, + "at_max_frames": 0, + "exact48": 11, + "pcm_derived_frames": 17313.0, + "peak_inflight": 32, + "interior_rps": 14.506207588669962, + "interior_rtf": 0.01503976665662243, + "interior_mean_inflight": 31.99963270833469, + "fixed48_valid": false, + "max_response_s": 3.552142067812383 + }, + { + "arm": "baseline", + "round": 1, + "mode": "sampled", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 19.947364503052086, + "audio_s": 1378.88, + "rps": 15.039580790439654, + "rtf": 0.014466352766775995, + "p95_response_s": 4.897944782185368, + "mean_duration_s": 4.596266666666667, + "median_duration_s": 4.48, + "min_duration_s": 2.88, + "max_duration_s": 10.24, + "caps": 0, + "at_max_frames": 0, + "exact48": 9, + "pcm_derived_frames": 17236.0, + "peak_inflight": 64, + "interior_rps": 14.772988235869104, + "interior_rtf": 0.014789570445164352, + "interior_mean_inflight": 63.99964299078361, + "fixed48_valid": false, + "max_response_s": 6.380287105217576 + }, + { + "arm": "baseline", + "round": 2, + "mode": "greedy48", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 12.094073274172843, + "audio_s": 1140.24, + "rps": 24.805538481452444, + "rtf": 0.010606603236312392, + "p95_response_s": 1.3324332737131044, + "mean_duration_s": 3.8008, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 276, + "at_max_frames": 273, + "exact48": 273, + "pcm_derived_frames": 14253.0, + "peak_inflight": 32, + "interior_rps": 24.0738475426106, + "interior_rtf": 0.010936431591609813, + "interior_mean_inflight": 31.999390705533155, + "fixed48_valid": false, + "max_response_s": 1.3492237930186093 + }, + { + "arm": "baseline", + "round": 2, + "mode": "greedy48", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 11.363486644811928, + "audio_s": 1142.32, + "rps": 26.400347831355695, + "rtf": 0.009947726245545845, + "p95_response_s": 2.411577110388316, + "mean_duration_s": 3.807733333333333, + "median_duration_s": 3.84, + "min_duration_s": 2.8, + "max_duration_s": 3.84, + "caps": 283, + "at_max_frames": 283, + "exact48": 283, + "pcm_derived_frames": 14279.0, + "peak_inflight": 64, + "interior_rps": 25.24913933766535, + "interior_rtf": 0.010414089097638338, + "interior_mean_inflight": 63.99934789676571, + "fixed48_valid": false, + "max_response_s": 2.4508438971824944 + }, + { + "arm": "baseline", + "round": 2, + "mode": "sampled", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 19.982479984872043, + "audio_s": 1371.52, + "rps": 15.013151532098034, + "rtf": 0.014569587016501432, + "p95_response_s": 2.6867772510275243, + "mean_duration_s": 4.571733333333333, + "median_duration_s": 4.48, + "min_duration_s": 2.72, + "max_duration_s": 10.0, + "caps": 0, + "at_max_frames": 0, + "exact48": 7, + "pcm_derived_frames": 17144.0, + "peak_inflight": 32, + "interior_rps": 14.688290200283893, + "interior_rtf": 0.014979169569545266, + "interior_mean_inflight": 31.99963069200315, + "fixed48_valid": false, + "max_response_s": 4.393639067653567 + }, + { + "arm": "baseline", + "round": 2, + "mode": "sampled", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 19.736425609793514, + "audio_s": 1363.84, + "rps": 15.200320763813254, + "rtf": 0.014471217745331942, + "p95_response_s": 4.845615285425447, + "mean_duration_s": 4.546133333333333, + "median_duration_s": 4.4, + "min_duration_s": 2.72, + "max_duration_s": 11.68, + "caps": 0, + "at_max_frames": 0, + "exact48": 11, + "pcm_derived_frames": 17048.0, + "peak_inflight": 64, + "interior_rps": 14.76624455355767, + "interior_rtf": 0.015027830388430046, + "interior_mean_inflight": 63.9996463755415, + "fixed48_valid": false, + "max_response_s": 7.120478071272373 + }, + { + "arm": "final", + "round": 0, + "mode": "greedy48", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 12.865571803878993, + "audio_s": 1141.92, + "rps": 23.31804637781816, + "rtf": 0.011266613951834622, + "p95_response_s": 1.4281870554899796, + "mean_duration_s": 3.8064000000000004, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 280, + "at_max_frames": 276, + "exact48": 276, + "pcm_derived_frames": 14274.0, + "peak_inflight": 32, + "interior_rps": 22.500397275173285, + "interior_rtf": 0.011688354533888935, + "interior_mean_inflight": 31.999437066682333, + "fixed48_valid": false, + "max_response_s": 1.464728930965066 + }, + { + "arm": "final", + "round": 0, + "mode": "greedy48", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 11.269619753584266, + "audio_s": 1141.6, + "rps": 26.620241548485783, + "rtf": 0.009871776238248307, + "p95_response_s": 2.4245060281828046, + "mean_duration_s": 3.805333333333333, + "median_duration_s": 3.84, + "min_duration_s": 2.88, + "max_duration_s": 3.84, + "caps": 283, + "at_max_frames": 279, + "exact48": 279, + "pcm_derived_frames": 14270.0, + "peak_inflight": 64, + "interior_rps": 25.810591381249292, + "interior_rtf": 0.010193914144174995, + "interior_mean_inflight": 63.99934948513199, + "fixed48_valid": false, + "max_response_s": 2.4315964658744633 + }, + { + "arm": "final", + "round": 0, + "mode": "sampled", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 18.745703546330333, + "audio_s": 1377.92, + "rps": 16.003667147437, + "rtf": 0.013604348254129654, + "p95_response_s": 2.5659151571802794, + "mean_duration_s": 4.593066666666667, + "median_duration_s": 4.48, + "min_duration_s": 2.64, + "max_duration_s": 9.76, + "caps": 0, + "at_max_frames": 0, + "exact48": 9, + "pcm_derived_frames": 17224.0, + "peak_inflight": 32, + "interior_rps": 15.772488640297134, + "interior_rtf": 0.013955919853733439, + "interior_mean_inflight": 31.999625586638548, + "fixed48_valid": false, + "max_response_s": 4.05552648473531 + }, + { + "arm": "final", + "round": 0, + "mode": "sampled", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 18.54867773130536, + "audio_s": 1379.92, + "rps": 16.173659618533225, + "rtf": 0.013441850057470985, + "p95_response_s": 4.537079988559708, + "mean_duration_s": 4.599733333333334, + "median_duration_s": 4.56, + "min_duration_s": 2.8, + "max_duration_s": 8.48, + "caps": 0, + "at_max_frames": 0, + "exact48": 6, + "pcm_derived_frames": 17249.0, + "peak_inflight": 64, + "interior_rps": 15.8409283531812, + "interior_rtf": 0.013839659582945052, + "interior_mean_inflight": 63.99963034685142, + "fixed48_valid": false, + "max_response_s": 5.520832907408476 + }, + { + "arm": "final", + "round": 1, + "mode": "greedy48", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 12.254824021831155, + "audio_s": 1139.68, + "rps": 24.48015568934894, + "rtf": 0.010752863980969354, + "p95_response_s": 1.3326363097177818, + "mean_duration_s": 3.7989333333333337, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 277, + "at_max_frames": 274, + "exact48": 274, + "pcm_derived_frames": 14246.0, + "peak_inflight": 32, + "interior_rps": 23.823985180226373, + "interior_rtf": 0.011062433628333746, + "interior_mean_inflight": 31.999409786066966, + "fixed48_valid": false, + "max_response_s": 1.361410547979176 + }, + { + "arm": "final", + "round": 1, + "mode": "greedy48", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 11.139470681082457, + "audio_s": 1141.12, + "rps": 26.931261689971794, + "rtf": 0.00976187489578875, + "p95_response_s": 2.3560594292357564, + "mean_duration_s": 3.803733333333333, + "median_duration_s": 3.84, + "min_duration_s": 2.8, + "max_duration_s": 3.84, + "caps": 282, + "at_max_frames": 282, + "exact48": 282, + "pcm_derived_frames": 14264.0, + "peak_inflight": 64, + "interior_rps": 26.111513201255992, + "interior_rtf": 0.010082729464473302, + "interior_mean_inflight": 63.99933949378279, + "fixed48_valid": false, + "max_response_s": 2.371493713930249 + }, + { + "arm": "final", + "round": 1, + "mode": "sampled", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 18.653851190116256, + "audio_s": 1373.1200000000001, + "rps": 16.08246988476862, + "rtf": 0.013585011645097482, + "p95_response_s": 2.5459276892710476, + "mean_duration_s": 4.577066666666667, + "median_duration_s": 4.48, + "min_duration_s": 2.8, + "max_duration_s": 8.88, + "caps": 0, + "at_max_frames": 0, + "exact48": 4, + "pcm_derived_frames": 17164.0, + "peak_inflight": 32, + "interior_rps": 15.806562888397778, + "interior_rtf": 0.01389570439926151, + "interior_mean_inflight": 31.999616261200515, + "fixed48_valid": false, + "max_response_s": 3.6991110420785844 + }, + { + "arm": "final", + "round": 1, + "mode": "sampled", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 18.535222947131842, + "audio_s": 1369.76, + "rps": 16.18540013549836, + "rtf": 0.013531730337527628, + "p95_response_s": 4.556147839757614, + "mean_duration_s": 4.5658666666666665, + "median_duration_s": 4.4, + "min_duration_s": 2.64, + "max_duration_s": 8.8, + "caps": 0, + "at_max_frames": 0, + "exact48": 11, + "pcm_derived_frames": 17122.0, + "peak_inflight": 64, + "interior_rps": 15.81525110648647, + "interior_rtf": 0.014033163705711225, + "interior_mean_inflight": 63.99962328469041, + "fixed48_valid": false, + "max_response_s": 5.544348424300551 + }, + { + "arm": "final", + "round": 2, + "mode": "greedy48", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 12.365936156827956, + "audio_s": 1142.1599999999999, + "rps": 24.260193178690514, + "rtf": 0.010826798484299887, + "p95_response_s": 1.423822610313073, + "mean_duration_s": 3.8071999999999995, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 282, + "at_max_frames": 280, + "exact48": 280, + "pcm_derived_frames": 14277.0, + "peak_inflight": 32, + "interior_rps": 23.562173703329602, + "interior_rtf": 0.01115026548503734, + "interior_mean_inflight": 31.99942383477762, + "fixed48_valid": false, + "max_response_s": 1.47030667681247 + }, + { + "arm": "final", + "round": 2, + "mode": "greedy48", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 11.290032208897173, + "audio_s": 1141.68, + "rps": 26.57211197002461, + "rtf": 0.009888963815515007, + "p95_response_s": 2.424221107456833, + "mean_duration_s": 3.8056, + "median_duration_s": 3.84, + "min_duration_s": 2.96, + "max_duration_s": 3.84, + "caps": 282, + "at_max_frames": 280, + "exact48": 280, + "pcm_derived_frames": 14271.0, + "peak_inflight": 64, + "interior_rps": 25.776552207561373, + "interior_rtf": 0.010202825738838942, + "interior_mean_inflight": 63.99933511364126, + "fixed48_valid": false, + "max_response_s": 2.4331844090484083 + }, + { + "arm": "final", + "round": 2, + "mode": "sampled", + "client_concurrency": 32, + "requests": 300, + "failures": 0, + "makespan_s": 18.93886736407876, + "audio_s": 1395.04, + "rps": 15.840440414562925, + "rtf": 0.013575859734544357, + "p95_response_s": 2.6509435122599827, + "mean_duration_s": 4.650133333333333, + "median_duration_s": 4.48, + "min_duration_s": 2.56, + "max_duration_s": 25.84, + "caps": 0, + "at_max_frames": 0, + "exact48": 11, + "pcm_derived_frames": 17438.0, + "peak_inflight": 32, + "interior_rps": 15.46414657722741, + "interior_rtf": 0.013988772091516137, + "interior_mean_inflight": 31.999631729721422, + "fixed48_valid": false, + "max_response_s": 10.408001323696226 + }, + { + "arm": "final", + "round": 2, + "mode": "sampled", + "client_concurrency": 64, + "requests": 300, + "failures": 0, + "makespan_s": 18.556426780298352, + "audio_s": 1373.84, + "rps": 16.166905598362, + "rtf": 0.013506978090824516, + "p95_response_s": 4.500672037573532, + "mean_duration_s": 4.579466666666667, + "median_duration_s": 4.44, + "min_duration_s": 2.72, + "max_duration_s": 14.0, + "caps": 0, + "at_max_frames": 0, + "exact48": 5, + "pcm_derived_frames": 17173.0, + "peak_inflight": 64, + "interior_rps": 15.927726434800704, + "interior_rtf": 0.013815576532996239, + "interior_mean_inflight": 63.999623723744996, + "fixed48_valid": false, + "max_response_s": 7.473742607049644 + } + ], + "summaries": { + "sampled-c32": { + "baseline": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 19.982479984872043, + "min": 19.73926813295111, + "max": 20.184422803111374 + }, + "audio_s": { + "median": 1372.48, + "min": 1371.52, + "max": 1385.04 + }, + "rps": { + "median": 15.013151532098034, + "min": 14.862946685488366, + "max": 15.19813186483873 + }, + "rtf": { + "median": 0.014569587016501432, + "min": 0.014382190001275873, + "max": 0.01457316958579635 + }, + "p95_response_s": { + "median": 2.6867772510275243, + "min": 2.6196906621102243, + "max": 2.753425784688443 + }, + "mean_duration_s": { + "median": 4.574933333333333, + "min": 4.571733333333333, + "max": 4.6168 + }, + "median_duration_s": { + "median": 4.52, + "min": 4.48, + "max": 4.56 + }, + "min_duration_s": { + "median": 2.64, + "min": 2.4, + "max": 2.72 + }, + "max_duration_s": { + "median": 8.0, + "min": 7.76, + "max": 10.0 + }, + "caps": { + "median": 0, + "min": 0, + "max": 0 + }, + "at_max_frames": { + "median": 0, + "min": 0, + "max": 0 + }, + "exact48": { + "median": 11, + "min": 7, + "max": 11 + }, + "pcm_derived_frames": { + "median": 17156.0, + "min": 17144.0, + "max": 17313.0 + }, + "peak_inflight": { + "median": 32, + "min": 32, + "max": 32 + }, + "interior_rps": { + "median": 14.688290200283893, + "min": 14.506207588669962, + "max": 14.769683037456982 + }, + "interior_rtf": { + "median": 0.014979169569545266, + "min": 0.014883914366667608, + "max": 0.01503976665662243 + }, + "interior_mean_inflight": { + "median": 31.99963270833469, + "min": 31.99963069200315, + "max": 31.99964170182187 + }, + "total_requests": 900, + "total_audio_s": 4129.04, + "total_caps": 0, + "total_at_max_frames": 0, + "total_exact48": 29, + "all_fixed48_valid": false + }, + "final": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 18.745703546330333, + "min": 18.653851190116256, + "max": 18.93886736407876 + }, + "audio_s": { + "median": 1377.92, + "min": 1373.1200000000001, + "max": 1395.04 + }, + "rps": { + "median": 16.003667147437, + "min": 15.840440414562925, + "max": 16.08246988476862 + }, + "rtf": { + "median": 0.013585011645097482, + "min": 0.013575859734544357, + "max": 0.013604348254129654 + }, + "p95_response_s": { + "median": 2.5659151571802794, + "min": 2.5459276892710476, + "max": 2.6509435122599827 + }, + "mean_duration_s": { + "median": 4.593066666666667, + "min": 4.577066666666667, + "max": 4.650133333333333 + }, + "median_duration_s": { + "median": 4.48, + "min": 4.48, + "max": 4.48 + }, + "min_duration_s": { + "median": 2.64, + "min": 2.56, + "max": 2.8 + }, + "max_duration_s": { + "median": 9.76, + "min": 8.88, + "max": 25.84 + }, + "caps": { + "median": 0, + "min": 0, + "max": 0 + }, + "at_max_frames": { + "median": 0, + "min": 0, + "max": 0 + }, + "exact48": { + "median": 9, + "min": 4, + "max": 11 + }, + "pcm_derived_frames": { + "median": 17224.0, + "min": 17164.0, + "max": 17438.0 + }, + "peak_inflight": { + "median": 32, + "min": 32, + "max": 32 + }, + "interior_rps": { + "median": 15.772488640297134, + "min": 15.46414657722741, + "max": 15.806562888397778 + }, + "interior_rtf": { + "median": 0.013955919853733439, + "min": 0.01389570439926151, + "max": 0.013988772091516137 + }, + "interior_mean_inflight": { + "median": 31.999625586638548, + "min": 31.999616261200515, + "max": 31.999631729721422 + }, + "total_requests": 900, + "total_audio_s": 4146.08, + "total_caps": 0, + "total_at_max_frames": 0, + "total_exact48": 24, + "all_fixed48_valid": false + }, + "median_percent_change": { + "rps": 6.597652819404698, + "rtf": -6.75774385566884, + "p95_response_s": -4.49840394476404, + "audio_s": 0.39636278852879236, + "mean_duration_s": 0.39636278852881457, + "interior_rps": 7.381379488214934 + }, + "total_audio_percent_change": 0.41268672621239677, + "matched_pcm_hashes": 14, + "matched_durations": 112 + }, + "greedy48-c32": { + "baseline": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 12.368925544898957, + "min": 12.094073274172843, + "max": 12.409082953352481 + }, + "audio_s": { + "median": 1140.24, + "min": 1138.1599999999999, + "max": 1140.8799999999999 + }, + "rps": { + "median": 24.25432984546684, + "min": 24.175839675481495, + "max": 24.805538481452444 + }, + "rtf": { + "median": 0.010841565760552344, + "min": 0.010606603236312392, + "max": 0.010902757919231464 + }, + "p95_response_s": { + "median": 1.3458482225192712, + "min": 1.3324332737131044, + "max": 1.3660403012065216 + }, + "mean_duration_s": { + "median": 3.8008, + "min": 3.7938666666666663, + "max": 3.802933333333333 + }, + "median_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "min_duration_s": { + "median": 2.96, + "min": 2.96, + "max": 2.96 + }, + "max_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "caps": { + "median": 276, + "min": 276, + "max": 278 + }, + "at_max_frames": { + "median": 274, + "min": 273, + "max": 276 + }, + "exact48": { + "median": 274, + "min": 273, + "max": 276 + }, + "pcm_derived_frames": { + "median": 14253.0, + "min": 14227.0, + "max": 14261.0 + }, + "peak_inflight": { + "median": 32, + "min": 32, + "max": 32 + }, + "interior_rps": { + "median": 23.594010932160216, + "min": 23.55702600182998, + "max": 24.0738475426106 + }, + "interior_rtf": { + "median": 0.011157971292675595, + "min": 0.010936431591609813, + "max": 0.011197488491907423 + }, + "interior_mean_inflight": { + "median": 31.999404196554725, + "min": 31.999390705533155, + "max": 31.99941429464469 + }, + "total_requests": 900, + "total_audio_s": 3419.2799999999997, + "total_caps": 830, + "total_at_max_frames": 823, + "total_exact48": 823, + "all_fixed48_valid": false + }, + "final": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 12.365936156827956, + "min": 12.254824021831155, + "max": 12.865571803878993 + }, + "audio_s": { + "median": 1141.92, + "min": 1139.68, + "max": 1142.1599999999999 + }, + "rps": { + "median": 24.260193178690514, + "min": 23.31804637781816, + "max": 24.48015568934894 + }, + "rtf": { + "median": 0.010826798484299887, + "min": 0.010752863980969354, + "max": 0.011266613951834622 + }, + "p95_response_s": { + "median": 1.423822610313073, + "min": 1.3326363097177818, + "max": 1.4281870554899796 + }, + "mean_duration_s": { + "median": 3.8064000000000004, + "min": 3.7989333333333337, + "max": 3.8071999999999995 + }, + "median_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "min_duration_s": { + "median": 2.96, + "min": 2.96, + "max": 2.96 + }, + "max_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "caps": { + "median": 280, + "min": 277, + "max": 282 + }, + "at_max_frames": { + "median": 276, + "min": 274, + "max": 280 + }, + "exact48": { + "median": 276, + "min": 274, + "max": 280 + }, + "pcm_derived_frames": { + "median": 14274.0, + "min": 14246.0, + "max": 14277.0 + }, + "peak_inflight": { + "median": 32, + "min": 32, + "max": 32 + }, + "interior_rps": { + "median": 23.562173703329602, + "min": 22.500397275173285, + "max": 23.823985180226373 + }, + "interior_rtf": { + "median": 0.01115026548503734, + "min": 0.011062433628333746, + "max": 0.011688354533888935 + }, + "interior_mean_inflight": { + "median": 31.99942383477762, + "min": 31.999409786066966, + "max": 31.999437066682333 + }, + "total_requests": 900, + "total_audio_s": 3423.76, + "total_caps": 839, + "total_at_max_frames": 830, + "total_exact48": 830, + "all_fixed48_valid": false + }, + "median_percent_change": { + "rps": 0.024174377362840005, + "rtf": -0.13620981118971365, + "p95_response_s": 5.793698463846297, + "audio_s": 0.14733740265207018, + "mean_duration_s": 0.14733740265207018, + "interior_rps": -0.13493775569637378 + }, + "total_audio_percent_change": 0.13102173557006935, + "matched_pcm_hashes": 1, + "matched_durations": 804 + }, + "sampled-c64": { + "baseline": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 19.836500237695873, + "min": 19.736425609793514, + "max": 19.947364503052086 + }, + "audio_s": { + "median": 1368.8, + "min": 1363.84, + "max": 1378.88 + }, + "rps": { + "median": 15.123635540804791, + "min": 15.039580790439654, + "max": 15.200320763813254 + }, + "rtf": { + "median": 0.014471217745331942, + "min": 0.014466352766775995, + "max": 0.01449189088084152 + }, + "p95_response_s": { + "median": 4.845615285425447, + "min": 4.799736037221737, + "max": 4.897944782185368 + }, + "mean_duration_s": { + "median": 4.562666666666667, + "min": 4.546133333333333, + "max": 4.596266666666667 + }, + "median_duration_s": { + "median": 4.48, + "min": 4.4, + "max": 4.48 + }, + "min_duration_s": { + "median": 2.88, + "min": 2.72, + "max": 2.88 + }, + "max_duration_s": { + "median": 10.24, + "min": 8.72, + "max": 11.68 + }, + "caps": { + "median": 0, + "min": 0, + "max": 0 + }, + "at_max_frames": { + "median": 0, + "min": 0, + "max": 0 + }, + "exact48": { + "median": 9, + "min": 5, + "max": 11 + }, + "pcm_derived_frames": { + "median": 17110.0, + "min": 17048.0, + "max": 17236.0 + }, + "peak_inflight": { + "median": 64, + "min": 64, + "max": 64 + }, + "interior_rps": { + "median": 14.772988235869104, + "min": 14.76624455355767, + "max": 14.826302139784737 + }, + "interior_rtf": { + "median": 0.014957955135430911, + "min": 0.014789570445164352, + "max": 0.015027830388430046 + }, + "interior_mean_inflight": { + "median": 63.99964299078361, + "min": 63.99961896669195, + "max": 63.9996463755415 + }, + "total_requests": 900, + "total_audio_s": 4111.52, + "total_caps": 0, + "total_at_max_frames": 0, + "total_exact48": 25, + "all_fixed48_valid": false + }, + "final": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 18.54867773130536, + "min": 18.535222947131842, + "max": 18.556426780298352 + }, + "audio_s": { + "median": 1373.84, + "min": 1369.76, + "max": 1379.92 + }, + "rps": { + "median": 16.173659618533225, + "min": 16.166905598362, + "max": 16.18540013549836 + }, + "rtf": { + "median": 0.013506978090824516, + "min": 0.013441850057470985, + "max": 0.013531730337527628 + }, + "p95_response_s": { + "median": 4.537079988559708, + "min": 4.500672037573532, + "max": 4.556147839757614 + }, + "mean_duration_s": { + "median": 4.579466666666667, + "min": 4.5658666666666665, + "max": 4.599733333333334 + }, + "median_duration_s": { + "median": 4.44, + "min": 4.4, + "max": 4.56 + }, + "min_duration_s": { + "median": 2.72, + "min": 2.64, + "max": 2.8 + }, + "max_duration_s": { + "median": 8.8, + "min": 8.48, + "max": 14.0 + }, + "caps": { + "median": 0, + "min": 0, + "max": 0 + }, + "at_max_frames": { + "median": 0, + "min": 0, + "max": 0 + }, + "exact48": { + "median": 6, + "min": 5, + "max": 11 + }, + "pcm_derived_frames": { + "median": 17173.0, + "min": 17122.0, + "max": 17249.0 + }, + "peak_inflight": { + "median": 64, + "min": 64, + "max": 64 + }, + "interior_rps": { + "median": 15.8409283531812, + "min": 15.81525110648647, + "max": 15.927726434800704 + }, + "interior_rtf": { + "median": 0.013839659582945052, + "min": 0.013815576532996239, + "max": 0.014033163705711225 + }, + "interior_mean_inflight": { + "median": 63.999623723744996, + "min": 63.99962328469041, + "max": 63.99963034685142 + }, + "total_requests": 900, + "total_audio_s": 4123.52, + "total_caps": 0, + "total_at_max_frames": 0, + "total_exact48": 22, + "all_fixed48_valid": false + }, + "median_percent_change": { + "rps": 6.942934289148828, + "rtf": -6.663154901517987, + "p95_response_s": -6.367308972995566, + "audio_s": 0.36820572764464377, + "mean_duration_s": 0.36820572764464377, + "interior_rps": 7.229005399998334 + }, + "total_audio_percent_change": 0.29186286336926326, + "matched_pcm_hashes": 7, + "matched_durations": 112 + }, + "greedy48-c64": { + "baseline": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 11.363486644811928, + "min": 11.264604894910008, + "max": 11.4890843112953 + }, + "audio_s": { + "median": 1140.8, + "min": 1137.6, + "max": 1142.32 + }, + "rps": { + "median": 26.400347831355695, + "min": 26.11174153409772, + "max": 26.632092541084788 + }, + "rtf": { + "median": 0.009947726245545845, + "min": 0.009874303028497552, + "max": 0.010099406040168162 + }, + "p95_response_s": { + "median": 2.4218830559635536, + "min": 2.411577110388316, + "max": 2.568560073385015 + }, + "mean_duration_s": { + "median": 3.8026666666666666, + "min": 3.792, + "max": 3.807733333333333 + }, + "median_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "min_duration_s": { + "median": 2.8, + "min": 2.8, + "max": 2.96 + }, + "max_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "caps": { + "median": 282, + "min": 275, + "max": 283 + }, + "at_max_frames": { + "median": 282, + "min": 269, + "max": 283 + }, + "exact48": { + "median": 282, + "min": 269, + "max": 283 + }, + "pcm_derived_frames": { + "median": 14260.0, + "min": 14220.0, + "max": 14279.0 + }, + "peak_inflight": { + "median": 64, + "min": 64, + "max": 64 + }, + "interior_rps": { + "median": 25.370435925769396, + "min": 25.24913933766535, + "max": 25.794812158986286 + }, + "interior_rtf": { + "median": 0.010412561279692619, + "min": 0.010210167083200961, + "max": 0.010414089097638338 + }, + "interior_mean_inflight": { + "median": 63.99934789676571, + "min": 63.999343316676544, + "max": 63.99934963725613 + }, + "total_requests": 900, + "total_audio_s": 3420.72, + "total_caps": 840, + "total_at_max_frames": 834, + "total_exact48": 834, + "all_fixed48_valid": false + }, + "final": { + "requests": { + "median": 300, + "min": 300, + "max": 300 + }, + "failures": { + "median": 0, + "min": 0, + "max": 0 + }, + "makespan_s": { + "median": 11.269619753584266, + "min": 11.139470681082457, + "max": 11.290032208897173 + }, + "audio_s": { + "median": 1141.6, + "min": 1141.12, + "max": 1141.68 + }, + "rps": { + "median": 26.620241548485783, + "min": 26.57211197002461, + "max": 26.931261689971794 + }, + "rtf": { + "median": 0.009871776238248307, + "min": 0.00976187489578875, + "max": 0.009888963815515007 + }, + "p95_response_s": { + "median": 2.424221107456833, + "min": 2.3560594292357564, + "max": 2.4245060281828046 + }, + "mean_duration_s": { + "median": 3.805333333333333, + "min": 3.803733333333333, + "max": 3.8056 + }, + "median_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "min_duration_s": { + "median": 2.88, + "min": 2.8, + "max": 2.96 + }, + "max_duration_s": { + "median": 3.84, + "min": 3.84, + "max": 3.84 + }, + "caps": { + "median": 282, + "min": 282, + "max": 283 + }, + "at_max_frames": { + "median": 280, + "min": 279, + "max": 282 + }, + "exact48": { + "median": 280, + "min": 279, + "max": 282 + }, + "pcm_derived_frames": { + "median": 14270.0, + "min": 14264.0, + "max": 14271.0 + }, + "peak_inflight": { + "median": 64, + "min": 64, + "max": 64 + }, + "interior_rps": { + "median": 25.810591381249292, + "min": 25.776552207561373, + "max": 26.111513201255992 + }, + "interior_rtf": { + "median": 0.010193914144174995, + "min": 0.010082729464473302, + "max": 0.010202825738838942 + }, + "interior_mean_inflight": { + "median": 63.99933949378279, + "min": 63.99933511364126, + "max": 63.99934948513199 + }, + "total_requests": 900, + "total_audio_s": 3424.3999999999996, + "total_caps": 847, + "total_at_max_frames": 841, + "total_exact48": 841, + "all_fixed48_valid": false + }, + "median_percent_change": { + "rps": 0.8329197726286042, + "rtf": -0.7634911277493606, + "p95_response_s": 0.09653857924816212, + "audio_s": 0.07012622720896644, + "mean_duration_s": 0.07012622720896644, + "interior_rps": 1.7349148306624818 + }, + "total_audio_percent_change": 0.10757969082531016, + "matched_pcm_hashes": 164, + "matched_durations": 844 + } + }, + "memory": { + "scope": "1 Hz whole-GPU resident memory across startup, four warmups, all measured rounds, post-timing C1 control, and shutdown; not peak allocated memory and not a measured-window-only peak.", + "arms": { + "baseline": { + "samples": 499, + "observed_duration_s": 498.07, + "max_sampled_resident_MiB": 30081, + "max_sampled_resident_GiB": 29.3759765625, + "max_poll_gap_s": 1.001, + "median_poll_gap_s": 1.0, + "last_sample_MiB": 4, + "csv_sha256": "17134f74aa2854bd5d792da1a451879dea086f145eea971bd5fe5e153c14ebb1" + }, + "final": { + "samples": 479, + "observed_duration_s": 478.068, + "max_sampled_resident_MiB": 29245, + "max_sampled_resident_GiB": 28.5595703125, + "max_poll_gap_s": 1.001, + "median_poll_gap_s": 1.0, + "last_sample_MiB": 4, + "csv_sha256": "51d8ce22b6632a07725285379e4d1fd23b8b972ea0329bd65e161751a6baf320" + } + }, + "max_sampled_resident_change_MiB": -836, + "max_sampled_resident_change_pct": -2.7791629267644002, + "interpretation": "These are observed whole-run sampled resident maxima, including startup and warmup. The difference is not attributed to memory savings from the optimization." + }, + "c1_control": { + "requests": 60, + "request_bodies_match": true, + "matching_counts": { + "pcm_bytes": 60, + "generated_frame_estimate": 60, + "pcm_sha256": 60 + }, + "wav_payloads_independently_verified": { + "baseline": 60, + "final": 60 + }, + "note": "Frame counts are derived from PCM length; no RVQ token sequence was recorded. Recorded hashes are distinguished from downloaded WAV verification. Excluded from throughput.", + "position": "60 sampled requests at concurrency one after measured rounds per arm; excluded from throughput.", + "waveform_verification_status": "All 120 saved C1 WAV payloads were independently decoded and checked locally against their recorded PCM hashes; all 60 cross-arm request bodies, PCM-derived frame counts and PCM hashes match.", + "saved_wav_payloads_validated_post_timing": 120, + "quality_caveat": "Byte-identical output is not necessarily correct speech: one matched control pair has a severe ASR mismatch in the targeted follow-up." + }, + "quality_observations": { + "status": "Endpoint, C1 and selected ASR checks completed; severe ASR mismatches confirmed in baseline and final outputs.", + "natural_c32_max_audio_s": { + "baseline": 10.0, + "final": 25.84 + }, + "natural_c64_max_audio_s": { + "baseline": 11.68, + "final": 14.0 + }, + "final_c32_longest_response_s": 10.408001323696226, + "final_c32_longest_pcm_derived_frames": 323, + "outlier_reached_384_frame_cap": false, + "caveat": "The final C64 14.0-second output has 27 normalized word edits against 13 reference words. Same-input/seed checks also find severe baseline mismatches and a mismatched byte-identical C1 pair. The final C32 25.84-second output transcribes expected words, but its excess duration remains unexplained. These selected checks establish neither quality equivalence nor a change in failure rates.", + "outlier_endpoint_metrics": [ + { + "arm": "final", + "client_concurrency": 32, + "duration_s": 25.84, + "leading_quiet_s": 0.02, + "trailing_quiet_s": 0.4, + "endpoint_trimmed_duration_s": 25.42, + "generated_frame_estimate": 323.0, + "wall_s": 10.408001323696226 + }, + { + "arm": "final", + "client_concurrency": 64, + "duration_s": 14.0, + "leading_quiet_s": 0.04, + "trailing_quiet_s": 0.04, + "endpoint_trimmed_duration_s": 13.92, + "generated_frame_estimate": 175.0, + "wall_s": 7.473742607049644 + }, + { + "arm": "baseline", + "client_concurrency": 32, + "duration_s": 5.44, + "leading_quiet_s": 0.02, + "trailing_quiet_s": 0.82, + "endpoint_trimmed_duration_s": 4.6, + "generated_frame_estimate": 68.0, + "wall_s": 2.19121035002172 + } + ], + "current_outlier_ASR": [ + { + "arm": "baseline", + "client_concurrency": 32, + "audio_duration_s": 5.44, + "endpoint_trimmed_audio_s": 4.6, + "word_edits": 0, + "reference_words": 12, + "normalized_word_error_rate": 0.0, + "exact_normalized_match": true, + "finding": "Paired baseline transcribed expected words." + }, + { + "arm": "final", + "client_concurrency": 32, + "audio_duration_s": 25.84, + "endpoint_trimmed_audio_s": 25.42, + "word_edits": 0, + "reference_words": 12, + "normalized_word_error_rate": 0.0, + "exact_normalized_match": true, + "finding": "Expected words transcribed; large duration anomaly remains unresolved." + }, + { + "arm": "final", + "client_concurrency": 64, + "audio_duration_s": 14.0, + "endpoint_trimmed_audio_s": 13.92, + "word_edits": 27, + "reference_words": 13, + "normalized_word_error_rate": 2.076923076923077, + "exact_normalized_match": false, + "finding": "ASR transcribed unrelated content; serious content mismatch." + } + ], + "controlled_same_input_seed_followup": "Complete 13-membership diagnostic: severe baseline mismatches and a mismatched byte-identical C1 pair confirm the issue predates the candidate relative to this baseline. Failure-rate changes and quality equivalence remain unproven." + }, + "limitations": [ + "Three observed rounds per arm; medians and full min/max ranges, no confidence intervals or statistical significance claim.", + "Arms ran sequentially in one allocation, each with a fresh server and identical four-wave warmup; temporal effects remain possible.", + "Greedy48 is a capped greedy workload; early EOS prevents a fixed-work claim whenever all_fixed48_valid is false.", + "Natural output duration and caps affect wall/audio RTF. ASR and quiet-tail validation are separate evidence; no quality equivalence inferred here.", + "Interior completion rates omit ramp-up/drain but include whole delivered clips crossing window boundaries; not a GPU-work measure.", + "Greedy median throughput changes only +0.02% at C32 and +0.83% at C64; ranges overlap. C32 median per-cohort p95 response increases. The 48-token cap permits early EOS and is not a fixed-work control.", + "Gains are incremental versus the stated baseline only; no multiplication by historical gains.", + "C1 recorded hash equality does not establish bitwise concurrent trajectory equality or perceptual quality.", + "Whole-run sampled resident memory maxima are not peak allocator usage or measured-window-only peaks.", + "Performance evidence covers guidance scale 1, no reference audio, and 30 synthetic English smoke texts only. It establishes no new voice-cloning, higher-guidance, long-prompt or multi-GPU performance claim.", + "Startup and warmup are excluded from throughput timing; no cold-start improvement is claimed.", + "Severe ASR mismatches occur in both baseline and final outputs, including one byte-identical C1 pair. This confirms a pre-existing baseline problem but establishes neither failure-rate changes, quality equivalence nor causal attribution. The final C32 duration anomaly remains unexplained." + ], + "endpoint_quality": { + "waveforms_validated_by_post_timing_scan": 7200, + "group_count": 24, + "method": "20ms endpoint RMS windows <= -40dBFS; internal pauses retained; WAVs unchanged; trimmed RTF is a diagnostic sensitivity check, not speech detection or a replacement timing metric.", + "families": { + "sampled-c32": { + "baseline": { + "raw_audio_s": 4129.04, + "endpoint_trimmed_audio_s": 3846.56, + "endpoint_quiet_fraction": 0.06841299672563114, + "endpoint_trimmed_rtf": { + "median": 0.015680875390857187, + "min": 0.015360102819197814, + "max": 0.01568163481932419 + }, + "max_trailing_quiet_s": 1.54, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0 + }, + "final": { + "raw_audio_s": 4146.08, + "endpoint_trimmed_audio_s": 3880.1400000000003, + "endpoint_quiet_fraction": 0.06414251533979075, + "endpoint_trimmed_rtf": { + "median": 0.014528399685595633, + "min": 0.014491443388230744, + "max": 0.014539698190213456 + }, + "max_trailing_quiet_s": 1.34, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0 + } + }, + "sampled-c64": { + "baseline": { + "raw_audio_s": 4111.52, + "endpoint_trimmed_audio_s": 3836.66, + "endpoint_quiet_fraction": 0.06685118885473026, + "endpoint_trimmed_rtf": { + "median": 0.015514306549576187, + "min": 0.015510448744788451, + "max": 0.015515933418093543 + }, + "max_trailing_quiet_s": 2.04, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0 + }, + "final": { + "raw_audio_s": 4123.52, + "endpoint_trimmed_audio_s": 3841.3, + "endpoint_quiet_fraction": 0.06844152568679192, + "endpoint_trimmed_rtf": { + "median": 0.01450060700187415, + "min": 0.014416817761002147, + "max": 0.014537429762456347 + }, + "max_trailing_quiet_s": 1.46, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0 + } + }, + "greedy48-c32": { + "baseline": { + "raw_audio_s": 3419.2799999999997, + "endpoint_trimmed_audio_s": 3314.56, + "endpoint_quiet_fraction": 0.030626330689501827, + "endpoint_trimmed_rtf": { + "median": 0.01117902963097769, + "min": 0.010945258900026102, + "max": 0.011248670141550168 + }, + "max_trailing_quiet_s": 1.78, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0 + }, + "final": { + "raw_audio_s": 3423.76, + "endpoint_trimmed_audio_s": 3315.88, + "endpoint_quiet_fraction": 0.03150921793583666, + "endpoint_trimmed_rtf": { + "median": 0.011153931029244703, + "min": 0.011152539823979037, + "max": 0.011607545971489015 + }, + "max_trailing_quiet_s": 3.08, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 1 + } + }, + "greedy48-c64": { + "baseline": { + "raw_audio_s": 3420.72, + "endpoint_trimmed_audio_s": 3317.9799999999996, + "endpoint_quiet_fraction": 0.030034612596178654, + "endpoint_trimmed_rtf": { + "median": 0.010294318704194309, + "min": 0.01015213404613458, + "max": 0.01040169148359978 + }, + "max_trailing_quiet_s": 1.5, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0 + }, + "final": { + "raw_audio_s": 3424.3999999999996, + "endpoint_trimmed_audio_s": 3326.52, + "endpoint_quiet_fraction": 0.02858310944983053, + "endpoint_trimmed_rtf": { + "median": 0.010161415751703484, + "min": 0.010035378354518348, + "max": 0.01019471231750449 + }, + "max_trailing_quiet_s": 1.18, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0 + } + } + }, + "individual_groups": [ + { + "arm": "baseline", + "round": 0, + "mode": "sampled", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1372.48, + "endpoint_trimmed_audio_s": 1285.1, + "raw_aggregate_rtf": 0.014382190001275873, + "endpoint_trimmed_aggregate_rtf": 0.015360102819197814, + "endpoint_quiet_fraction": 0.0636657729074377, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.28 + }, + { + "arm": "baseline", + "round": 0, + "mode": "sampled", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1368.8, + "endpoint_trimmed_audio_s": 1278.46, + "raw_aggregate_rtf": 0.01449189088084152, + "endpoint_trimmed_aggregate_rtf": 0.015515933418093543, + "endpoint_quiet_fraction": 0.065999415546464, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.26 + }, + { + "arm": "baseline", + "round": 0, + "mode": "greedy48", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1138.1599999999999, + "endpoint_trimmed_audio_s": 1103.1599999999999, + "raw_aggregate_rtf": 0.010902757919231464, + "endpoint_trimmed_aggregate_rtf": 0.011248670141550168, + "endpoint_quiet_fraction": 0.03075138820552471, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.78 + }, + { + "arm": "baseline", + "round": 0, + "mode": "greedy48", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1137.6, + "endpoint_trimmed_audio_s": 1104.54, + "raw_aggregate_rtf": 0.010099406040168162, + "endpoint_trimmed_aggregate_rtf": 0.01040169148359978, + "endpoint_quiet_fraction": 0.029061181434599112, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.5 + }, + { + "arm": "baseline", + "round": 1, + "mode": "sampled", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1385.04, + "endpoint_trimmed_audio_s": 1287.2, + "raw_aggregate_rtf": 0.01457316958579635, + "endpoint_trimmed_aggregate_rtf": 0.015680875390857187, + "endpoint_quiet_fraction": 0.07064055911742616, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.32 + }, + { + "arm": "baseline", + "round": 1, + "mode": "sampled", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1378.88, + "endpoint_trimmed_audio_s": 1285.74, + "raw_aggregate_rtf": 0.014466352766775995, + "endpoint_trimmed_aggregate_rtf": 0.015514306549576187, + "endpoint_quiet_fraction": 0.0675475748433512, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.8 + }, + { + "arm": "baseline", + "round": 1, + "mode": "greedy48", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1140.8799999999999, + "endpoint_trimmed_audio_s": 1106.44, + "raw_aggregate_rtf": 0.010841565760552344, + "endpoint_trimmed_aggregate_rtf": 0.01117902963097769, + "endpoint_quiet_fraction": 0.030187223897342255, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.4 + }, + { + "arm": "baseline", + "round": 1, + "mode": "greedy48", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1140.8, + "endpoint_trimmed_audio_s": 1109.58, + "raw_aggregate_rtf": 0.009874303028497552, + "endpoint_trimmed_aggregate_rtf": 0.01015213404613458, + "endpoint_quiet_fraction": 0.02736676016830297, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.16 + }, + { + "arm": "baseline", + "round": 2, + "mode": "sampled", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1371.52, + "endpoint_trimmed_audio_s": 1274.26, + "raw_aggregate_rtf": 0.014569587016501432, + "endpoint_trimmed_aggregate_rtf": 0.01568163481932419, + "endpoint_quiet_fraction": 0.07091402239850676, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.54 + }, + { + "arm": "baseline", + "round": 2, + "mode": "sampled", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1363.84, + "endpoint_trimmed_audio_s": 1272.46, + "raw_aggregate_rtf": 0.014471217745331942, + "endpoint_trimmed_aggregate_rtf": 0.015510448744788451, + "endpoint_quiet_fraction": 0.06700199436884084, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 2.04 + }, + { + "arm": "baseline", + "round": 2, + "mode": "greedy48", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1140.24, + "endpoint_trimmed_audio_s": 1104.96, + "raw_aggregate_rtf": 0.010606603236312392, + "endpoint_trimmed_aggregate_rtf": 0.010945258900026102, + "endpoint_quiet_fraction": 0.03094085455693536, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.68 + }, + { + "arm": "baseline", + "round": 2, + "mode": "greedy48", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1142.32, + "endpoint_trimmed_audio_s": 1103.86, + "raw_aggregate_rtf": 0.009947726245545845, + "endpoint_trimmed_aggregate_rtf": 0.010294318704194309, + "endpoint_quiet_fraction": 0.03366832411233283, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.5 + }, + { + "arm": "final", + "round": 0, + "mode": "sampled", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1377.92, + "endpoint_trimmed_audio_s": 1290.28, + "raw_aggregate_rtf": 0.013604348254129654, + "endpoint_trimmed_aggregate_rtf": 0.014528399685595633, + "endpoint_quiet_fraction": 0.0636031119368324, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.34 + }, + { + "arm": "final", + "round": 0, + "mode": "sampled", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1379.92, + "endpoint_trimmed_audio_s": 1286.6, + "raw_aggregate_rtf": 0.013441850057470985, + "endpoint_trimmed_aggregate_rtf": 0.014416817761002147, + "endpoint_quiet_fraction": 0.0676271088179026, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.28 + }, + { + "arm": "final", + "round": 0, + "mode": "greedy48", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1141.92, + "endpoint_trimmed_audio_s": 1108.3799999999999, + "raw_aggregate_rtf": 0.011266613951834622, + "endpoint_trimmed_aggregate_rtf": 0.011607545971489015, + "endpoint_quiet_fraction": 0.029371584699453716, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.32 + }, + { + "arm": "final", + "round": 0, + "mode": "greedy48", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1141.6, + "endpoint_trimmed_audio_s": 1109.06, + "raw_aggregate_rtf": 0.009871776238248307, + "endpoint_trimmed_aggregate_rtf": 0.010161415751703484, + "endpoint_quiet_fraction": 0.0285038542396636, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.16 + }, + { + "arm": "final", + "round": 1, + "mode": "sampled", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1373.1200000000001, + "endpoint_trimmed_audio_s": 1282.96, + "raw_aggregate_rtf": 0.013585011645097482, + "endpoint_trimmed_aggregate_rtf": 0.014539698190213456, + "endpoint_quiet_fraction": 0.06566068515497558, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.3 + }, + { + "arm": "final", + "round": 1, + "mode": "sampled", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1369.76, + "endpoint_trimmed_audio_s": 1275.0, + "raw_aggregate_rtf": 0.013531730337527628, + "endpoint_trimmed_aggregate_rtf": 0.014537429762456347, + "endpoint_quiet_fraction": 0.06918000233617567, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.46 + }, + { + "arm": "final", + "round": 1, + "mode": "greedy48", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1139.68, + "endpoint_trimmed_audio_s": 1098.7, + "raw_aggregate_rtf": 0.010752863980969354, + "endpoint_trimmed_aggregate_rtf": 0.011153931029244703, + "endpoint_quiet_fraction": 0.035957461743647355, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 1, + "max_trailing_quiet_s": 3.08 + }, + { + "arm": "final", + "round": 1, + "mode": "greedy48", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1141.12, + "endpoint_trimmed_audio_s": 1110.02, + "raw_aggregate_rtf": 0.00976187489578875, + "endpoint_trimmed_aggregate_rtf": 0.010035378354518348, + "endpoint_quiet_fraction": 0.027253925967470477, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.16 + }, + { + "arm": "final", + "round": 2, + "mode": "sampled", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1395.04, + "endpoint_trimmed_audio_s": 1306.9, + "raw_aggregate_rtf": 0.013575859734544357, + "endpoint_trimmed_aggregate_rtf": 0.014491443388230744, + "endpoint_quiet_fraction": 0.0631809840578047, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.32 + }, + { + "arm": "final", + "round": 2, + "mode": "sampled", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1373.84, + "endpoint_trimmed_audio_s": 1279.7, + "raw_aggregate_rtf": 0.013506978090824516, + "endpoint_trimmed_aggregate_rtf": 0.01450060700187415, + "endpoint_quiet_fraction": 0.06852326326209739, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.28 + }, + { + "arm": "final", + "round": 2, + "mode": "greedy48", + "client_concurrency": 32, + "count": 300, + "raw_audio_s": 1142.1599999999999, + "endpoint_trimmed_audio_s": 1108.8, + "raw_aggregate_rtf": 0.010826798484299887, + "endpoint_trimmed_aggregate_rtf": 0.011152539823979037, + "endpoint_quiet_fraction": 0.029207816768228534, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.5 + }, + { + "arm": "final", + "round": 2, + "mode": "greedy48", + "client_concurrency": 64, + "count": 300, + "raw_audio_s": 1141.68, + "endpoint_trimmed_audio_s": 1107.44, + "raw_aggregate_rtf": 0.009888963815515007, + "endpoint_trimmed_aggregate_rtf": 0.01019471231750449, + "endpoint_quiet_fraction": 0.029990890617335862, + "all_quiet_count": 0, + "trailing_quiet_at_least_3s": 0, + "max_trailing_quiet_s": 1.18 + } + ], + "current_outliers": [ + { + "arm": "final", + "client_concurrency": 32, + "duration_s": 25.84, + "leading_quiet_s": 0.02, + "trailing_quiet_s": 0.4, + "endpoint_trimmed_duration_s": 25.42, + "generated_frame_estimate": 323.0, + "wall_s": 10.408001323696226 + }, + { + "arm": "final", + "client_concurrency": 64, + "duration_s": 14.0, + "leading_quiet_s": 0.04, + "trailing_quiet_s": 0.04, + "endpoint_trimmed_duration_s": 13.92, + "generated_frame_estimate": 175.0, + "wall_s": 7.473742607049644 + }, + { + "arm": "baseline", + "client_concurrency": 32, + "duration_s": 5.44, + "leading_quiet_s": 0.02, + "trailing_quiet_s": 0.82, + "endpoint_trimmed_duration_s": 4.6, + "generated_frame_estimate": 68.0, + "wall_s": 2.19121035002172 + } + ], + "independent_local_outlier_waveform_checks": [ + { + "arm": "final", + "client_concurrency": 32, + "duration_s": 25.84, + "leading_quiet_s": 0.02, + "trailing_quiet_s": 0.4, + "endpoint_trimmed_duration_s": 25.42, + "pcm_sha256_verified": true + }, + { + "arm": "final", + "client_concurrency": 64, + "duration_s": 14.0, + "leading_quiet_s": 0.04, + "trailing_quiet_s": 0.04, + "endpoint_trimmed_duration_s": 13.920000000000002, + "pcm_sha256_verified": true + }, + { + "arm": "baseline", + "client_concurrency": 32, + "duration_s": 5.44, + "leading_quiet_s": 0.02, + "trailing_quiet_s": 0.8200000000000001, + "endpoint_trimmed_duration_s": 4.6000000000000005, + "pcm_sha256_verified": true + } + ] + }, + "asr_validation": { + "status": "selected_checks_complete_with_baseline_and_final_quality_limitations", + "model_id": "deepdml/faster-whisper-large-v3-turbo-ct2", + "model_revision": "4df90f75321148c3a29a9e2351b7ddf8f5b115a8", + "device": "cpu", + "compute_type": "int8", + "cpu_threads": 8, + "versions": { + "faster-whisper": "1.2.1", + "ctranslate2": "4.8.2" + }, + "primary_selection": "One upper-duration-median clip per each of 30 prompts in round 0 sampled C32, each arm.", + "normalization": "NFKC, case folding, non-letter/non-number separators; no numeric word expansion.", + "primary": { + "baseline": { + "clips": 30, + "word_edits": 25, + "reference_words": 384, + "exact_normalized_matches": 18, + "normalized_word_error_rate": 0.06510416666666667 + }, + "final": { + "clips": 30, + "word_edits": 20, + "reference_words": 384, + "exact_normalized_matches": 21, + "normalized_word_error_rate": 0.052083333333333336 + } + }, + "primary_audio_memberships": 60, + "supplemental_current_outliers_required": 3, + "evidence_sha256": "abae4040710aef2e489996375e0e30b03159a4ddf5fbacccbdb0e40c18b9c4e5", + "limitations": [ + "A targeted transcription sanity check of selected clips, not a full-corpus ASR score, perceptual-quality evaluation or MOS.", + "Numeric spelling and recognition differences contribute to normalized word edits.", + "No quality-improvement or concurrent-output equivalence claim." + ], + "supplemental_current_outliers": [ + { + "arm": "baseline", + "client_concurrency": 32, + "audio_duration_s": 5.44, + "endpoint_trimmed_audio_s": 4.6, + "word_edits": 0, + "reference_words": 12, + "normalized_word_error_rate": 0.0, + "exact_normalized_match": true, + "finding": "Paired baseline transcribed expected words." + }, + { + "arm": "final", + "client_concurrency": 32, + "audio_duration_s": 25.84, + "endpoint_trimmed_audio_s": 25.42, + "word_edits": 0, + "reference_words": 12, + "normalized_word_error_rate": 0.0, + "exact_normalized_match": true, + "finding": "Expected words transcribed; large duration anomaly remains unresolved." + }, + { + "arm": "final", + "client_concurrency": 64, + "audio_duration_s": 14.0, + "endpoint_trimmed_audio_s": 13.92, + "word_edits": 27, + "reference_words": 13, + "normalized_word_error_rate": 2.076923076923077, + "exact_normalized_match": false, + "finding": "ASR transcribed unrelated content; serious content mismatch." + } + ], + "supplemental_memberships": 3, + "same_input_seed_followup_status": "Complete; severe baseline and final ASR mismatches confirmed. Failure rates and quality equivalence remain unproven.", + "retry_kept_serving_versions_unchanged": true, + "same_input_seed_followup": { + "status": "complete_targeted_diagnostic", + "scope": "Only the 12 sampled cohorts and C1 controls of this matched comparison; final means candidate commit ca22df8b3a6459154271873006caaebc829ec5a9.", + "input_text": "Our quarterly engineering review begins promptly at two in the afternoon on Thursday.", + "sampling_seed": 94, + "selection": "Post-hoc follow-up of a problematic input/seed across both arms, all rounds and both client concurrency levels, plus one representative of its byte-identical C1 pair.", + "asr_memberships": 13, + "new_transcriptions": 12, + "reused_current_outlier_transcriptions": 1, + "local_waveforms_independently_hashed": 13, + "c1_counterpart_payload_bytes_also_compared": true, + "model_id": "deepdml/faster-whisper-large-v3-turbo-ct2", + "model_revision": "4df90f75321148c3a29a9e2351b7ddf8f5b115a8", + "versions": { + "faster-whisper": "1.2.1", + "ctranslate2": "4.8.2" + }, + "rows": [ + { + "arm": "baseline", + "round": 0, + "client_concurrency": 32, + "audio_duration_s": 7.76, + "word_edits": 1, + "reference_words": 13, + "normalized_word_error_rate": 0.07692307692307693, + "transcript": "Our quarterly engineering review begins promptly at 2 in the afternoon on Thursday." + }, + { + "arm": "baseline", + "round": 0, + "client_concurrency": 64, + "audio_duration_s": 7.36, + "word_edits": 1, + "reference_words": 13, + "normalized_word_error_rate": 0.07692307692307693, + "transcript": "Our quarterly engineering review begins promptly at 2 in the afternoon on Thursday." + }, + { + "arm": "baseline", + "round": 1, + "client_concurrency": 32, + "audio_duration_s": 5.76, + "word_edits": 1, + "reference_words": 13, + "normalized_word_error_rate": 0.07692307692307693, + "transcript": "Uh, quarterly engineering review begins promptly at two in the afternoon on Thursday." + }, + { + "arm": "baseline", + "round": 1, + "client_concurrency": 64, + "audio_duration_s": 10.24, + "word_edits": 16, + "reference_words": 13, + "normalized_word_error_rate": 1.2307692307692308, + "transcript": "I come to the engineering review, it begins from the S2 and I have to do it on those days." + }, + { + "arm": "baseline", + "round": 2, + "client_concurrency": 32, + "audio_duration_s": 10.0, + "word_edits": 149, + "reference_words": 13, + "normalized_word_error_rate": 11.461538461538462, + "transcript": "Thanks, I'll take you with 4.4th and 5.5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th and 5th" + }, + { + "arm": "baseline", + "round": 2, + "client_concurrency": 64, + "audio_duration_s": 5.04, + "word_edits": 1, + "reference_words": 13, + "normalized_word_error_rate": 0.07692307692307693, + "transcript": "Our quarterly engineering review begins promptly at 2 in the afternoon on Thursday." + }, + { + "arm": "final", + "round": 0, + "client_concurrency": 32, + "audio_duration_s": 6.72, + "word_edits": 1, + "reference_words": 13, + "normalized_word_error_rate": 0.07692307692307693, + "transcript": "Our quarterly engineering review begins promptly at 2 in the afternoon on Thursday." + }, + { + "arm": "final", + "round": 0, + "client_concurrency": 64, + "audio_duration_s": 5.76, + "word_edits": 13, + "reference_words": 13, + "normalized_word_error_rate": 1.0, + "transcript": "They come here while they can't get me, walk back up here." + }, + { + "arm": "final", + "round": 1, + "client_concurrency": 32, + "audio_duration_s": 6.96, + "word_edits": 1, + "reference_words": 13, + "normalized_word_error_rate": 0.07692307692307693, + "transcript": "Our quarterly engineering review begins promptly at 2 in the afternoon on Thursday." + }, + { + "arm": "final", + "round": 1, + "client_concurrency": 64, + "audio_duration_s": 7.04, + "word_edits": 5, + "reference_words": 13, + "normalized_word_error_rate": 0.38461538461538464, + "transcript": "Thank you, Jeff Boyden. Our quarterly engineering review begins promptly at 2 in the afternoon on Thursday." + }, + { + "arm": "final", + "round": 2, + "client_concurrency": 32, + "audio_duration_s": 8.0, + "word_edits": 2, + "reference_words": 13, + "normalized_word_error_rate": 0.15384615384615385, + "transcript": "A quarterly engineering review begins promptly at 2 in the afternoon on Thursday." + }, + { + "arm": "final", + "round": 2, + "client_concurrency": 64, + "audio_duration_s": 14.0, + "word_edits": 27, + "reference_words": 13, + "normalized_word_error_rate": 2.076923076923077, + "transcript": "I can't get more than we're working them. Don't care, I'll say. Don't look like. Don't know. Just, just. Don't." + }, + { + "arm": "baseline_and_final_identical_C1_pair", + "round": null, + "client_concurrency": 1, + "audio_duration_s": 10.16, + "word_edits": 26, + "reference_words": 13, + "normalized_word_error_rate": 2.0, + "transcript": "I can't just walk in there so much, and I'm always going to be fighting a deal, being young, roughly, men, death, with men." + } + ], + "findings": [ + "Severe ASR mismatches occur in baseline outputs, including 16 and 149 word edits over 13 reference words. The issue predates this candidate relative to the stated baseline.", + "The byte-identical baseline/final C1 representative has 26 word edits over 13 reference words: bitwise agreement does not establish correct speech content.", + "Final concurrent outputs also show serious ASR mismatches; the previously observed C32 duration anomaly remains unexplained." + ], + "limitations": [ + "This selected diagnostic cannot estimate failure rates, establish quality equivalence, or attribute changes causally to batching or the optimization.", + "ASR can introduce recognition errors and repetitions; these are transcription sanity findings, not perceptual-quality scores or MOS.", + "The frozen 60-clip primary sample and this targeted diagnostic remain separate." + ] + } + }, + "profiling": { + "method": { + "tool": "Nsight Systems 2026.5.1", + "graph_trace": "graph:host-only", + "trace": "cuda,nvtx", + "sampling": "sampled generation, temperature 0.9/top-k 50/top-p 1", + "cpu_context_switches": "disabled", + "client_concurrency": 32, + "capture_target_seconds": 15, + "capture_begins_after_completions": 128, + "request_cap": 2400, + "configuration": "Same model, deployment and workload as the uninjected comparison; four matched warmups and a 300-request collection-disabled control precede each capture.", + "timings_excluded_from_throughput": true, + "cpu_sampling": "disabled" + }, + "arms": { + "baseline": { + "capture_duration_s": 15.051548019, + "both_stages_verified": true, + "sustained_window_valid": true, + "gpu_activity_envelope_s": 15.01934332, + "recorded_activity_upper_bound_s": 10.267107808, + "recorded_idle_lower_bound_s": 4.752235512, + "recorded_activity_upper_fraction": 0.6835923241949036, + "fully_contained_32_request_scopes": [ + { + "scope": "breeze.first_code rows=32", + "fully_contained_ranges": 183, + "cpu_span_total_ms": 2206.124574, + "cpu_span_mean_ms": 12.055325540983606, + "external_kernel_records": 157530, + "external_kernels_per_range_distribution": { + "852": 10, + "864": 106, + "858": 63, + "840": 2, + "846": 2 + }, + "graph_envelope_records": 0, + "graph_envelopes_per_range_distribution": { + "0": 183 + } + }, + { + "scope": "breeze.depth rows=32 sampled=1", + "fully_contained_ranges": 106, + "cpu_span_total_ms": 847.075044, + "cpu_span_mean_ms": 7.991274000000001, + "external_kernel_records": 50880, + "external_kernels_per_range_distribution": { + "480": 106 + }, + "graph_envelope_records": 106, + "graph_envelopes_per_range_distribution": { + "1": 106 + } + }, + { + "scope": "breeze.depth.rng rows=32 draws=480", + "fully_contained_ranges": 106, + "cpu_span_total_ms": 689.3856850000001, + "cpu_span_mean_ms": 6.50363853773585, + "external_kernel_records": 50880, + "external_kernels_per_range_distribution": { + "480": 106 + }, + "graph_envelope_records": 0, + "graph_envelopes_per_range_distribution": { + "0": 106 + } + }, + { + "scope": "breeze.first_code.eos_readback rows=32", + "fully_contained_ranges": 183, + "cpu_span_total_ms": 84.379238, + "cpu_span_mean_ms": 0.46108873224043717, + "external_kernel_records": 183, + "external_kernels_per_range_distribution": { + "1": 183 + }, + "graph_envelope_records": 0, + "graph_envelopes_per_range_distribution": { + "0": 183 + } + } + ] + }, + "final": { + "capture_duration_s": 15.089370444, + "both_stages_verified": true, + "sustained_window_valid": true, + "gpu_activity_envelope_s": 15.028040847, + "recorded_activity_upper_bound_s": 11.987511476, + "recorded_idle_lower_bound_s": 3.040529371, + "recorded_activity_upper_fraction": 0.7976762638619677, + "fully_contained_32_request_scopes": [ + { + "scope": "breeze.first_code rows=32", + "fully_contained_ranges": 191, + "cpu_span_total_ms": 978.859137, + "cpu_span_mean_ms": 5.124916947643979, + "external_kernel_records": 66937, + "external_kernels_per_range_distribution": { + "347": 51, + "353": 126, + "335": 2, + "341": 12 + }, + "graph_envelope_records": 191, + "graph_envelopes_per_range_distribution": { + "1": 191 + } + }, + { + "scope": "breeze.first_code.eos_readback rows=32", + "fully_contained_ranges": 191, + "cpu_span_total_ms": 662.0009520000001, + "cpu_span_mean_ms": 3.465973570680629, + "external_kernel_records": 191, + "external_kernels_per_range_distribution": { + "1": 191 + }, + "graph_envelope_records": 0, + "graph_envelopes_per_range_distribution": { + "0": 191 + } + }, + { + "scope": "breeze.depth rows=32 sampled=1", + "fully_contained_ranges": 126, + "cpu_span_total_ms": 288.402938, + "cpu_span_mean_ms": 2.2889122063492064, + "external_kernel_records": 8064, + "external_kernels_per_range_distribution": { + "64": 126 + }, + "graph_envelope_records": 126, + "graph_envelopes_per_range_distribution": { + "1": 126 + } + }, + { + "scope": "breeze.depth.rng rows=32 draws=480", + "fully_contained_ranges": 126, + "cpu_span_total_ms": 22.377353, + "cpu_span_mean_ms": 0.17759803968253968, + "external_kernel_records": 0, + "external_kernels_per_range_distribution": { + "0": 126 + }, + "graph_envelope_records": 0, + "graph_envelopes_per_range_distribution": { + "0": 126 + } + } + ] + } + }, + "limits": [ + "One separately instrumented capture per arm; profiler overhead changes host timing.", + "CPU annotation spans include waiting and overlap for nested scopes; do not sum them or equate their differences with serving speedup.", + "External kernel records exclude graph-internal nodes. Captured sampling kernels still execute on the GPU. Replay seed/offset prologue kernels remain external.", + "Both stages covered on one GPU before drain. Different iteration counts are normalized by fully contained 32-request ranges.", + "Graph envelopes hide internal gaps. Recorded activity is an upper bound and recorded idle is a lower bound, not SM utilization. Capturing more work changes those bounds.", + "Faster submission moves some waiting into EOS readback; it does not remove the dependency on completed sampling." + ] + } +} diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 53c84e66a..b8921c1d3 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -7,6 +7,17 @@ # is tracked here. models: + BreezeBlue/Breeze-TTS-2: + supported_tasks: [default_voice, voice_design, voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + task_extra_body: + default_voice: {} + voice_design: + task_type: VoiceDesign + voice_clone: + task_type: Base + IndexTeam/IndexTTS-2.5: supported_tasks: [voice_clone] backend: openai-audio-speech diff --git a/docs/design/feature/preprocess_phase_contract.md b/docs/design/feature/preprocess_phase_contract.md index d795c9c23..31e320489 100644 --- a/docs/design/feature/preprocess_phase_contract.md +++ b/docs/design/feature/preprocess_phase_contract.md @@ -42,7 +42,8 @@ Scheduler / current input batch prefill / other spans single-token decode preprocess preprocess or preprocess_decode_batch | | - prompt embedding runner-managed talker_mtp + prompt embedding decode embeddings + | talker_mtp (if present) \ / model forward ``` @@ -96,9 +97,24 @@ computed per-row fields are not promised at that earlier boundary. ### Routing guarantee With `has_preprocess=True`, the batched decode hook is eligible only when it is -callable, the runner has `talker_mtp`, the scheduled span is one token, and +callable, the scheduled span is one token, and `_omni_is_prefill` is false. Otherwise the normal `preprocess` hook runs. +The hook's return format depends on whether the runner manages MTP: + +| Runner capability | `preprocess_decode_batch` return tuple | +| --- | --- | +| Without `talker_mtp` | `(input_ids, inputs_embeds, updates)` | +| With `talker_mtp` | `(input_ids, inputs_embeds, hidden_states, text_step, updates)` | + +Each tensor's batch dimension and each entry in the `updates` list correspond +to one physical request row, including separate CFG branches. Non-MTP batches +contain contiguous single-token decode spans; the runner flushes a batch before +processing any other span. It copies embeddings into its existing input buffer +and applies each update dictionary to the corresponding request. Returning the +original `input_ids` object lets it skip an unchanged token-ID copy. The existing +five-result MTP format remains supported. + Runner-managed MTP is likewise restricted to single-token decode rows. A prefill row, including a one-token tail, is never added to that MTP batch: the runner does not consume its `mtp_inputs`, generate codes for it, or overwrite its @@ -130,6 +146,8 @@ formalizes that existing behavior. Prefer upgrading the matching vLLM and vLLM-Omni release pair and consuming the explicit phase keys. Do not silently fall back to `input_ids.numel() > 1` when the model requires chunking correctness. For a custom or backported runtime, verify both metadata and MTP guards. +Older runners can restrict `preprocess_decode_batch` to MTP models; a non-MTP +model must retain its scalar `preprocess` implementation for those runtimes. ## Correctness and testing plan diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 1f543a6d2..1d9d52f77 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -72,7 +72,7 @@ th { | `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | -| `BreezeForConditionalGeneration` | Breeze-TTS-2 (TTS, voice clone, voice direction) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) | +| `BreezeForConditionalGeneration` | Breeze-TTS-2 (voice design, voice clone, voice direction, CFG) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) | | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — | | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/docs/source/performance/breeze_tts_2_gpu_bubbles.png b/docs/source/performance/breeze_tts_2_gpu_bubbles.png new file mode 100644 index 000000000..707f2742b Binary files /dev/null and b/docs/source/performance/breeze_tts_2_gpu_bubbles.png differ diff --git a/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py b/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py index 942ead792..05063d106 100644 --- a/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py +++ b/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py @@ -14,12 +14,10 @@ from pathlib import Path import soundfile as sf import torch -from transformers import AutoConfig +from transformers import AutoTokenizer from vllm_omni import Omni -from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( - BreezeTTS2PromptBuilder, -) +from vllm_omni.model_executor.models.breeze_tts_2.prompt import DEFAULT_INSTRUCTION, build_breeze_prompt from vllm_omni.utils.tracking_parser import TrackingArgumentParser REPO_ROOT = Path(__file__).resolve().parents[4] @@ -95,8 +93,7 @@ def resolve_template(args) -> str: def reference_payload(ref_audio_path: str | None) -> dict: """Load the reference clip into the payload keys the prompt builder reads. - ``BreezeTTS2PromptBuilder`` consumes ``ref_audio`` (a waveform) plus - ``ref_audio_sample_rate``; a bare path field is not read. + The prompt helper consumes a waveform and its sample rate. """ if ref_audio_path is None: return {} @@ -107,36 +104,17 @@ def reference_payload(ref_audio_path: str | None) -> dict: def build_prompt(args) -> dict: - config = AutoConfig.from_pretrained(args.model, trust_remote_code=True) - reference_encoder = None - if args.ref_audio is not None: - from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( - BreezeReferenceAudioTokenizer, - ) - - reference_encoder = BreezeReferenceAudioTokenizer.from_pretrained( - args.model, - num_codebooks=int(getattr(config, "num_codebooks", 16)), - device_map="cpu", - ) - builder = BreezeTTS2PromptBuilder.from_pretrained( - args.model, - config, - reference_audio_encoder=reference_encoder, + tokenizer = AutoTokenizer.from_pretrained(args.model) + payload = reference_payload(args.ref_audio) + reference = (payload["ref_audio"], payload["ref_audio_sample_rate"]) if payload else None + return build_breeze_prompt( + tokenizer, + args.text, + DEFAULT_INSTRUCTION if args.instruction is None else args.instruction, + speaker=args.voice, + ref_audio=reference, + ref_text=args.ref_text, ) - payload = { - "text": args.text, - "speaker": args.voice, - **reference_payload(args.ref_audio), - } - if args.instruction: - payload["instruction"] = args.instruction - if args.ref_text: - payload["ref_text"] = args.ref_text - prompt = builder.build(payload, template=resolve_template(args)) - if args.max_new_tokens is not None: - prompt["additional_information"]["breeze_max_new_frames"] = int(args.max_new_tokens) - return prompt def extract_audio(multimodal_output: dict) -> torch.Tensor: @@ -180,7 +158,12 @@ def main(): print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens") t_start = time.perf_counter() - outputs = engine.generate([prompt]) + if args.max_new_tokens is not None: + sampling = [params.clone() for params in engine.resolve_sampling_params_list(None)] + sampling[0].max_tokens = args.max_new_tokens + outputs = engine.generate([prompt], sampling_params_list=sampling) + else: + outputs = engine.generate([prompt]) elapsed = time.perf_counter() - t_start request_output = outputs[0] diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 711fdec6b..cadc100f4 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -1061,7 +1061,7 @@ vllm serve BreezeBlue/Breeze-TTS-2 --omni --port 8091 ./breeze_tts_2/run_server.sh ``` -The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming, 8-frame inter-stage chunks). +The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming with 1, 2, 4, then 5-frame inter-stage chunks). ### Sending requests diff --git a/recipes/BreezeBlue/Breeze-TTS-2.md b/recipes/BreezeBlue/Breeze-TTS-2.md index fb26ce91c..7fafed649 100644 --- a/recipes/BreezeBlue/Breeze-TTS-2.md +++ b/recipes/BreezeBlue/Breeze-TTS-2.md @@ -1,210 +1,307 @@ # Breeze-TTS-2 -> Text-to-speech serving (plain / voice design / voice clone / voice direction) +[Breeze-TTS-2](https://huggingface.co/BreezeBlue/Breeze-TTS-2) generates 24 kHz +English and Chinese speech. The integration supports Voice Design from text +instructions, Voice Clone from a reference recording and transcript, and +Voice Direction combining reference conditioning with instructions. -## Summary +## Installation and serving -- Vendor: BreezeBlue -- Model: `BreezeBlue/Breeze-TTS-2` (`BreezeForConditionalGeneration`) -- Task: Text-to-speech with speaker tags, natural-language style instructions, single-reference voice cloning, and reference+instruction "voice direction" -- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API (streaming and non-streaming); offline `Omni` example -- Hardware: 1x NVIDIA L20 48GB (CUDA) -- Maintainer: Community - -## When to use this recipe - -Use this recipe when you want a known-good starting point for serving -Breeze-TTS-2 with vLLM-Omni. Breeze-TTS-2 is a two-stage AR TTS pipeline: - -| Stage | Components | Output | -| ----- | ---------- | ------ | -| 0 (talker, `LLM_AR`) | T5Gemma2 text encoder + Qwen3 backbone + 2052-way codebook-0 head + 12-layer depth decoder | 16-codebook codec frames | -| 1 (codec, `LLM_GENERATION`) | Bundled Qwen3-TTS audio tokenizer decoder (Mimi fallback) | 24 kHz mono waveform | - -Four prompt modes are selected automatically from the request fields: - -| Mode | Trigger | Fields | -| ---- | ------- | ------ | -| `tts_plain` | only `input` (default `voice=S0`) | `input`, `voice` | -| `tts_instruction` (voice design) | `input` + `instructions` | style/emotion/pace directives, not spoken | -| `ref_clone_tata` (clone) | `ref_audio` + `ref_text` + `input` | reference clip and its exact transcript | -| `ref_edit_tata` (voice direction) | reference trio + `instructions` | keep the reference timbre, direct the delivery | - -Current scope: greedy sampling with `cfg_scale=1.0` and a default -`repetition_penalty=1.1` on generated codebook-0 tokens; CFG ≠ 1.0 and -`negative_prompt` are rejected with explicit client errors until companion -request support lands. - -The [official inference example](https://github.com/breezeblue-ai/breeze-tts/blob/main/infer.py) -also uses repetition penalty `1.1`, but its -[runtime defaults](https://github.com/breezeblue-ai/breeze-tts/blob/main/breeze_infer/runtime.py) -sample both the backbone and depth decoder with temperature `0.9` and top-k -`50`. The [official voice-design examples](https://huggingface.co/BreezeBlue/Breeze-TTS-2#-voice-design) -use CFG `4`. Those sampling and guidance settings are outside the current -port's scope, so the examples below do not establish parity with the official -voice-design outputs. - -## Supported model contract - -| Item | Value | -| ---- | ----- | -| Tasks | Text-to-speech through `/v1/audio/speech` (online) and the offline `Omni` example | -| Modes | plain, voice design (`instructions`), voice clone (`ref_audio` + `ref_text`), voice direction (reference + `instructions`) | -| Languages | English and Chinese (model card) | -| Reference audio | exactly one clip per request with its exact transcript; encoded by the bundled Qwen3-TTS tokenizer to 16 codebooks at 12.5 Hz | -| Output | 24 kHz mono; `wav` or `pcm`; raw audio streaming (`stream_format="audio"`) or SSE `speech.audio.delta` (`stream_format="sse"`) | -| Length | prompt bounded by stage-0 `max_model_len=4096`; synthesis bounded by `max_new_tokens` codec frames (default 2048, about 164 s) | -| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0` and `repetition_penalty=1.1`; other guidance values and `negative_prompt` are rejected | -| Deployment profile | `vllm_omni/deploy/breeze_tts_2.yaml`: two stages on one GPU, async-chunk streaming with 8-frame codec chunks, stage 0 `gpu_memory_utilization=0.80`, stage 1 `0.15` | - -## References - -- Issue: [#6656 [New Model]: Breeze TTS 2](https://github.com/vllm-project/vllm-omni/issues/6656) -- Related examples under `examples/`: - [`examples/online_serving/text_to_speech/breeze_tts_2/`](../../examples/online_serving/text_to_speech/breeze_tts_2/), - [`examples/offline_inference/text_to_speech/breeze_tts_2/`](../../examples/offline_inference/text_to_speech/breeze_tts_2/) - -## Hardware - -- Accelerator model and per-device memory: 1x NVIDIA L20 48GB -- Number of devices: 1 (both stages are placed on device 0) -- Device interconnect: not applicable (single device) -- Host memory: no special requirement; per-request codec state is small -- Qualification scope: functional serving of all four prompt modes, streaming PCM, and 4 concurrent mixed-mode requests (greedy, `cfg_scale=1.0`). No other accelerator is qualified by this recipe. - -## Software environment - -- OS: Linux -- Python: 3.12.12 -- Driver / runtime: NVIDIA driver 580.82.07, CUDA 13.0 (PyTorch 2.13.0+cu130) -- vLLM version: 0.28.0 -- vLLM-Omni version or commit: the commit that introduced this recipe (branch based on `main` @ `67ef0b64`) -- Transformers: 5.14.1 - -## Command - -Start the server from the repository root. The deploy config -(`vllm_omni/deploy/breeze_tts_2.yaml`, async chunk streaming by default) -auto-loads from the checkpoint's `model_type`; pass it explicitly if you -customized it: +Follow the project's source installation guide for vLLM and vLLM-Omni. +The integration loads the released checkpoint directly. The original Breeze +repository and the `qwen-tts` package are not required. The common requirements +include `soxr` to reproduce reference-audio resampling. ```bash -vllm serve BreezeBlue/Breeze-TTS-2 \ - --deploy-config vllm_omni/deploy/breeze_tts_2.yaml \ - --omni --port 8091 +vllm-omni serve BreezeBlue/Breeze-TTS-2 --omni --host 127.0.0.1 --port 8091 ``` -## Verification +The default configuration uses one CUDA GPU, BF16 generation and codec +computation, and FP32 reference encoding. Decoded waveform tensors are converted +to FP32 before response formatting. It targets GPUs with at least 16 GiB +available for the deployment; other workloads consume additional memory. -Plain synthesis (speaker tag `S0`..`S9`, default `S0`): +On WSL2, set vLLM's supported pinned-memory option before starting the server: ```bash -curl -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "model": "BreezeBlue/Breeze-TTS-2", - "input": "Hello, this is Breeze TTS 2 running on vLLM-Omni.", - "voice": "S0", - "response_format": "wav", - "sample_rate": 24000 - }' --output breeze_plain.wav +export VLLM_WSL2_ENABLE_PIN_MEMORY=1 ``` -Voice design (instruction only, no reference audio): +Native Linux enables pinned memory by default. Startup warmup exercises plain +generation, CFG, reference conditioning and two-request batches. New text +lengths or batch sizes can require additional graph +capture or compilation; warm the workload before measuring steady performance. + +## Voice Design ```bash -curl -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "input": "(laugh) Welcome to our story time.", - "instructions": "A warm young woman, clear voice, lively delivery.", - "response_format": "wav" - }' --output breeze_design.wav +curl --fail http://127.0.0.1:8091/v1/audio/speech \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "BreezeBlue/Breeze-TTS-2", + "input": "你好，欢迎使用语音合成系统。", + "instructions": "年轻女性，声音温暖自然，使用标准普通话。", + "extra_params": {"guidance_scale": 4.0}, + "response_format": "wav", + "seed": 42, + "max_new_tokens": 250 + }' --output breeze.wav ``` -Voice cloning (`ref_audio` and `ref_text` must be provided together): +The `voice` field selects a speaker tag, such as `S0` or `[S0]`; omitting it +uses `[S0]`. These are checkpoint prompt tags, not a catalog of named voices. +Omitting `instructions` uses `Speak clearly and naturally.`. An explicitly +empty instruction remains empty. Vocal-event tokens, including `(laughs)` +and `[笑]`, are passed to the checkpoint unchanged. + +## Voice Clone and Voice Direction + +Add `ref_audio` and its exact `ref_text` transcript. Reference audio accepts the +shared speech API's audio URLs and base64 data URLs, either as a string or +a one-item list. The shared resolver's +format, duration and local-file access restrictions apply. + +```json +{ + "model": "BreezeBlue/Breeze-TTS-2", + "input": "Please read this new sentence clearly and naturally.", + "ref_audio": "data:audio/wav;base64,", + "ref_text": "The words actually spoken in the reference recording.", + "instructions": "Keep the reference voice. Speak gently and calmly.", + "extra_params": {"guidance_scale": 4.0}, + "response_format": "wav", + "seed": 42 +} +``` + +Use `guidance_scale: 1.0` with the default instruction for ordinary cloning. +Reference audio plus a delivery instruction and CFG provides Voice Direction. +When explicitly specified, `task_type` is `Base` for reference conditioning +and `VoiceDesign` without reference audio. The legacy `CustomVoice` value +also remains accepted. + +Stereo references are averaged to mono and resampled to 24 kHz with soxr HQ. +Reference text is encoded independently from the target text. The reference +codec preserves full causal attention during offline encoding, as used by the +released Breeze runtime; newer Transformers sliding-window defaults would +change long references. Reference convolution uses full FP32 precision. + +## Sampling and output + +Supported `extra_params`: + +| Parameter | Default | Meaning | +| --- | --- | --- | +| `guidance_scale` | 1.0 | Any finite positive value; values other than 1 use two CFG branches | +| `temperature` | 0.9 | Zero selects greedy generation | +| `top_k` | 50 | Zero or -1 disables top-k filtering | +| `top_p` | 1.0 | Nucleus sampling probability in (0, 1] | +| `repetition_penalty` | 1.1 | Positive penalty applied to the first codebook's generated history | + +`cfg_scale` remains an alias for `guidance_scale`; `guidance_scale` takes +precedence when both are supplied. + +CFG combines conditional and unconditional logits in the Qwen3 output head +and all 15 remaining depth codebooks. Both branches retain reference audio +and reference text; the unconditional branch removes only target instructions. + +`max_new_tokens` limits generated audio frames, including EOS. Each frame is +80 ms; the default is 1500 tokens, subject to the remaining 2048-token context. +Both CFG branches stop at the same generation limit despite unequal prompts. +Each request owns its seed, generator, audio-code history and codec state. +Different GPU kernels and batch shapes can still change sampled outputs. + +For streamed PCM, use `"stream": true`, `"stream_format": "audio"`, and +`"response_format": "pcm"`. Output is mono signed 16-bit little-endian PCM at +24 kHz. Complete WAV responses accumulate the same internal streaming output. +With the default deployment, the first chunks contain 1, 2, 4 and 5 codec +frames, then remain at 5 frames. + +## Non-streaming throughput profile + +For concurrent complete responses, select the throughput deployment from the +repository root. ```bash -curl -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "input": "This is new text spoken in the cloned voice.", - "ref_audio": "file:///path/to/reference.wav", - "ref_text": "The exact transcript of the reference audio.", - "response_format": "wav" - }' --output breeze_clone.wav +vllm-omni serve BreezeBlue/Breeze-TTS-2 --omni \ + --deploy-config vllm_omni/deploy/breeze_tts_2_throughput.yaml \ + --host 127.0.0.1 --port 8091 ``` -Voice direction (reference + instruction): +Request a complete WAV with `"stream": false` and omit `stream_format`. +Setting `"stream_format": "audio"` enables streaming even when `stream` is false. ```bash -curl -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "input": "We need to discuss what happened last night.", - "ref_audio": "file:///path/to/reference.wav", - "ref_text": "The exact transcript of the reference audio.", - "instructions": "Speak slowly with a restrained, serious tone.", - "response_format": "wav" - }' --output breeze_direction.wav +curl --fail http://127.0.0.1:8091/v1/audio/speech \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "BreezeBlue/Breeze-TTS-2", + "input": "Hello, welcome to the speech synthesis demonstration.", + "stream": false, + "response_format": "wav", + "seed": 42, + "max_new_tokens": 250 + }' --output breeze-complete.wav ``` -Streaming raw PCM (`stream_format="audio"`): +The profile was tested on one H100 80 GB. It sets `max_num_seqs: 32` on both +stages and reserves 8 GiB of stage-0 KV cache for thirty-two 2048-token contexts. +These are physical sequence slots: stage 0 can run up to thirty-two plain +requests (`guidance_scale: 1.0`) or sixteen +CFG requests, since each CFG request occupies two slots. Mixed workloads share +the same limit; additional requests queue. The default deployment's memory +target does not apply to this larger profile. + +The internal codec uses fixed 16-frame chunks from the first chunk onward, +with a shorter final chunk when needed. `cfg_prefill_delay_tokens: 0` removes +the default eight-frame admission delay, so waiting requests can enter as soon +as scheduling and KV capacity permit. The stateful codec still requires +`async_chunk: true` for these complete HTTP responses. Warm the expected text +lengths and concurrency before measuring the profile. + +The shared benchmark defaults to streaming. Override both `stream` and +`stream_format` to measure complete responses, and set the generation limit +explicitly in the request body: ```bash -curl -N -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Streaming output from Breeze TTS 2.", - "stream": true, - "stream_format": "audio", - "response_format": "pcm", - "sample_rate": 24000 - }' --output breeze_stream.pcm +python benchmarks/tts/bench_tts.py \ + --model BreezeBlue/Breeze-TTS-2 --task default_voice --locale en \ + --dataset-path benchmarks/build_dataset/seed_tts_smoke \ + --host 127.0.0.1 --port 8091 --concurrency 1 2 4 8 16 32 64 \ + --num-prompts 128 --num-warmups 64 \ + --output-dir results/breeze-nonstream -- \ + --extra-body '{"stream":false,"stream_format":null,"response_format":"pcm","max_new_tokens":384,"seed":42,"extra_params":{"temperature":0.9,"top_k":50,"top_p":1,"repetition_penalty":1.1,"guidance_scale":1}}' ``` -Expected results: finite, non-silent 24 kHz mono WAV files; the streaming example -returns raw PCM bytes. Use `stream_format="sse"` for `speech.audio.delta` -events containing base64-encoded audio. On 1x L20 (greedy, single request after -warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time. +Use the same dataset and sampling settings when comparing deployment profiles. +Verify completed and failed request counts in each result. For complete +responses, compare throughput and RTF; first-response timing represents +completion latency. Run ASR separately: the benchmark's WER mode forces streaming. -## Notes +## Pipeline and optimization -- **Sample rate**: only `24000` is accepted; other values fail validation - before inference. -- **Speed / seeds / multi-reference**: `speed` must be `1.0`; - `task_type=VoiceDesign`, `ref_audio_2`, and speaker embeddings are not - supported yet. -- **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and - `negative_prompt` return a client error. -- **Repetition penalty**: the default `1.1` matches upstream's penalty on - generated codebook-0 tokens and discourages repeated silence. Override it - with `extra_params: {"repetition_penalty": 1.0}` to disable it, or another - finite positive value. Text and reference audio are excluded from its history. - The legacy `tools/breeze_tts_2_upstream_golden.py` uses penalty `1.0`; set - the same serving override when comparing against those golden frames. -- **Length budget**: prompt length is bounded by stage 0 `max_model_len=4096`; - synthesis length by `max_new_tokens` (default 2048 frames ≈ 164 s at the - 12.5 Hz codec frame rate) — set `max_new_tokens` for shorter caps. -- **Checkpoint layout**: the checkpoint must contain the `audio_tokenizer/` - subdirectory (bundled Qwen3-TTS codec) for the default streaming path; the - Mimi fallback is used only when that directory is absent and does not - support async-chunk streaming. -- **License**: the reference inference code is Apache 2.0; the checkpoints are - distributed under the separate BreezeBlue Research and Non-Commercial - License. Commercial use requires written authorization from BreezeBlue. +Stage 0 batches T5Gemma2 text encoding, optional reference-audio encoding, +vLLM's native Qwen3 backbone with paged KV caching, and Breeze's Llama depth +decoder. Text segments up to 128 tokens use padded CUDA Graphs with separate +attention masks. Longer segments use dynamically compiled encoder layers +without retaining a separate graph workspace for each input length. +The depth decoder uses fused QKV/MLP projections, compiled layer kernels, +fixed frame-local KV buffers and a CUDA Graph covering all 15 depth steps. +Graphs pad requests to power-of-two buckets, keeping CFG branches separate. +Compiled layers reuse kernels across batch sizes, avoiding a new compilation +for every shrinking batch. Only real requests advance their random state, +and returned tensors exclude padded rows. +Positive temperature, top-k and top-p are mutable GPU inputs to the captured +sampler. Changing those settings, or changing a positive CFG scale within the +paired path, reuses the graph. Greedy decoding uses a separate graph that skips +sorting and random sampling. Each request owns its random-number generator. -## Supported features +Reference encoding uses dynamic-shape compilation with CUDA Graphs disabled +for that path, so large reference-convolution workspaces remain reusable by +the allocator. Every convolution retains each recording's actual boundary; +the final downsampler repeats its last valid state for partial frames, matching +independent encoding. Static convolution metadata removes GPU-to-CPU padding +calculations. Startup compiles reference batches of one and two. +Waveform storage is aligned to 1920-sample frames, with that frame width +marked static for compilation. Actual sample lengths still govern every +convolution boundary. This avoids nested dynamic ceil/mod expressions that +otherwise make compiler tiling spend minutes simplifying reference shapes. -| Feature | Status | Notes | -| ------- | ------ | ----- | -| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks; raw audio or SSE `speech.audio.delta` | -| Voice cloning | ✓ | one reference clip plus transcript per request | -| Voice design / voice direction | ✓ | `instructions` field, with or without a reference | -| Classifier-free guidance | ✗ | `cfg_scale` must be `1.0`; `negative_prompt` rejected (follow-up) | -| Non-greedy sampling | ✗ | greedy decoding only | -| Prefix caching | ✗ | disabled in the deploy config | -| Tensor / pipeline parallelism | untested | single-GPU profile only | -| Quantization | ✗ | not supported in this release | -| CUDA graphs | not tuned | depth decoder and stage-1 codec run eagerly | -| Platforms | CUDA | qualified on 1x NVIDIA L20 | +Stage 1 reuses Qwen3-TTS's stateful codec and its CUDA Graph decoder with the +checkpoint's `audio_tokenizer/` weights. Shared-memory chunks connect stages. +Only the generated-code decode path is captured: reference conditioning was +already consumed in stage 0, so the Qwen ICL reference-prefix graphs and their +large convolution workspaces are unnecessary for Breeze. + +The Qwen3 backbone uses compiled piecewise prefill and full decode graphs. +Stage 0 uses synchronous scheduling to keep CFG branches at the same +generation step; stage 1 uses asynchronous scheduling. CFG admission reserves +enough KV pages for both branches' complete generation and preserves their +actual position IDs. Cancellation finishes both branches. + +The default admission policy lets active requests produce eight audio frames +before adding another prefill. This gives playback a buffer before text and +reference encoding share the device with an existing stream. It increases a +queued request's first-audio latency; benchmark it separately from an idle server. + +The default `max_num_seqs: 2` counts physical branches: it permits two plain +requests or one CFG request at a time. Additional CFG requests queue. +Prefix caching and chunked prefill are disabled because prompt token IDs +reserve positions for externally encoded embeddings. The deployment requires +`async_chunk: true`, including for complete HTTP responses. + +The accelerated implementation retains the `breeze_tts_2` model package, +`breeze_tts_2` and `breeze_tts_2_codec` stage names, and default +`breeze_tts_2.yaml` deployment. The shorter `breeze_tts.yaml` and +`breeze_tts_throughput.yaml` deployment filenames remain compatibility aliases. +For a copied custom configuration, use the current `breeze_tts_2.yaml` as +the starting point: codec chunk controls use `codec_chunk_frames`, +`initial_codec_chunk_frames` and `codec_chunk_ramp`. +Defaults now use a 2048-token context and temperature 0.9; set temperature +to zero explicitly when comparing greedy generation with the earlier runtime. + +The supported execution scope is one CUDA GPU, TP=1, PP=1 and unquantized +released weights. Named voices, speaker embeddings, dual CFG and multiple +reference recordings are not supported. The original public fast runtime +also does not expose dual CFG. + +## Offline inference and validation + +`vllm_omni.model_executor.models.breeze_tts_2.prompt.build_breeze_prompt` +accepts a checkpoint tokenizer, text, instructions, sampling parameters, +and optional `ref_audio=(waveform, sample_rate)` plus `ref_text`. Pass its +result to `Omni.generate()` using the default Breeze deployment. Waveform +arrays use soundfile's `(samples, channels)` convention. + +CPU tests compare cached depth attention and CFG against Hugging Face Llama +and cover request-owned sampling, payload transport, chunk delivery and CFG +admission. GPU tests compare padded batched text graphs with independent +eager encoders and exercise depth graphs across CFG scales and batch shapes. + +```bash +pytest -v tests/model_executor/models/test_breeze_tts_2.py \ + tests/model_executor/models/test_breeze_tts_2_registration.py \ + tests/core/sched/test_omni_cfg_ar_scheduler.py -m 'core_model and cpu' +pytest -v tests/model_executor/models/test_breeze_tts_2_graphs.py -m 'core_model and cuda' +pytest -v tests/e2e/online_serving/test_breeze_tts_2.py \ + -m 'core_model and tts' --run-level=core_model +pytest -v tests/e2e/online_serving/test_breeze_tts_2.py \ + -m 'advanced_model and tts' --run-level=advanced_model +``` + +Serving tests use the shared client for Chinese CFG design, English PCM, +concurrent CFG requests, reference cloning and direction. Core tests use dummy +weights and eight-frame requests; advanced tests use real weights and the +shared content checks. These checks do not establish perceptual quality parity. +The Chinese content check uses Whisper large-v3 as its primary transcriber: +small misrecognizes homophones on the fixed test clip. The expected text and +shared similarity threshold are unchanged; there is no failed-check retry. +Ensure enough host memory and disk space for the large-v3 test model when +Whisper runs on CPU alongside the serving process. + +The English PCM test uses a 0 dB harmonic-to-noise floor because the original +eager runtime scores 0.56 dB on that same prompt and seed, below the shared +1 dB floor. Raw PCM is not transcribed by the shared helper. Separate matched +evaluations must retain audio and compare intelligibility, speaker similarity, +instruction effects, first audio, RTF and playback underruns. + +The shared serving benchmark supports default speech, Voice Design and +Voice Clone. For Voice Design: + +```bash +python benchmarks/tts/bench_tts.py \ + --model BreezeBlue/Breeze-TTS-2 --task voice_design --locale en \ + --host 127.0.0.1 --port 8091 --concurrency 1 2 \ + --num-prompts 4 --num-warmups 1 --request-seed 42 --output-len 250 \ + --output-dir results/breeze +``` + +For a local model path, add `--served-model-name /path/to/Breeze-TTS-2`. +For cloning, select `--task voice_clone --dataset-path /path/to/seed-tts-eval` +with the desired locale. Reference conditioning with instructions and other +CFG scales also require identical requests and sampling settings in the +reference runtime and this deployment. + +The reference source and released weights have separate licenses. Consult +the [reference repository](https://github.com/breezeblue-ai/breeze-tts) and +checkpoint model card. diff --git a/requirements/common.txt b/requirements/common.txt index f143e5e22..5896bc6ff 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -15,6 +15,8 @@ kernels==0.16.1 safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 +# Breeze reference conditioning must use the released codec's soxr resampling. +soxr>=0.5.0.post1 cache-dit==1.5.0 tqdm>=4.66.0 torchsde>=0.2.6 diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 201c9d0e7..081c4d203 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 30, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 59, + EnvironmentVariableCategory.MODEL_SPECIFIC: 58, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 2, } @@ -181,7 +181,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): ModelEnvironmentVariableDisposition.PROMOTE: 31, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 16, + ModelEnvironmentVariableDisposition.INTERNALIZE: 15, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7, } diff --git a/tests/core/sched/test_omni_cfg_ar_scheduler.py b/tests/core/sched/test_omni_cfg_ar_scheduler.py new file mode 100644 index 000000000..e9ff8ea53 --- /dev/null +++ b/tests/core/sched/test_omni_cfg_ar_scheduler.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +from vllm import SamplingParams +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import Request, RequestStatus + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_cfg_ar_scheduler import OmniCFGARScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _request(name: str, prompt_length: int) -> Request: + return Request(name, [0] * prompt_length, SamplingParams(max_tokens=100), None) + + +def _scheduler(monkeypatch, waiting, running=()): + scheduler = OmniCFGARScheduler.__new__(OmniCFGARScheduler) + scheduler.policy = SchedulingPolicy.FCFS + scheduler.waiting = create_request_queue(scheduler.policy) + for request in waiting: + scheduler.waiting.add_request(request) + scheduler.skipped_waiting = create_request_queue(scheduler.policy) + scheduler.running = list(running) + scheduler.requests = {request.request_id: request for request in [*waiting, *running]} + scheduler.max_model_len = 128 + scheduler.max_num_running_reqs = 2 + scheduler.max_num_scheduled_tokens = 256 + scheduler.block_size = 16 + scheduler._capacity = 16 + scheduler._pairs = {"pair": {"cond": "cond", "uncond": "uncond"}} + scheduler._request_pairs = {"cond": "pair", "uncond": "pair"} + scheduler._prefill_min_output_tokens = 0 + + def advance(self, throttle_prefills=False): + self.running.extend(self.waiting) + self.waiting = create_request_queue(self.policy) + scheduled = {} + for request in self.running: + count = request.num_tokens - request.num_computed_tokens + request.num_computed_tokens += count + request.append_output_token_ids(0) + scheduled[request.request_id] = count + output = SchedulerOutput.make_empty() + output.num_scheduled_tokens = scheduled + output.total_num_scheduled_tokens = sum(scheduled.values()) + return output + + monkeypatch.setattr(OmniARScheduler, "schedule", advance) + return scheduler + + +def test_unequal_prompts_keep_real_positions_and_share_context_limit(monkeypatch): + cond, uncond = _request("cond", 80), _request("uncond", 30) + scheduler = _scheduler(monkeypatch, [cond, uncond]) + output = scheduler.schedule() + assert output.num_scheduled_tokens == {"cond": 80, "uncond": 30} + assert (cond.max_tokens, uncond.max_tokens) == (48, 48) + for _ in range(4): + assert scheduler.schedule().num_scheduled_tokens == {"cond": 1, "uncond": 1} + assert cond.num_computed_tokens - uncond.num_computed_tokens == 50 + + +def test_waiting_pair_cannot_be_split_or_starved_by_later_single(monkeypatch): + cond, uncond = _request("cond", 40), _request("uncond", 20) + running, later = _request("running", 10), _request("later", 10) + scheduler = _scheduler(monkeypatch, [cond, uncond, later], [running]) + assert scheduler.schedule().num_scheduled_tokens == {"running": 10} + assert list(scheduler.waiting) == [cond, uncond, later] + scheduler.running.clear() + del scheduler.requests["running"] + assert scheduler.schedule().num_scheduled_tokens == {"cond": 40, "uncond": 20} + assert list(scheduler.waiting) == [later] + + +def test_cancel_either_branch_finishes_both_and_releases_pair(monkeypatch): + cond, uncond = _request("cond", 40), _request("uncond", 20) + scheduler = _scheduler(monkeypatch, [cond, uncond]) + + def finish(self, ids, status): + assert set(ids) == {"cond", "uncond"} + assert status == RequestStatus.FINISHED_ABORTED + finished = [self.requests.pop(rid) for rid in ids] + self.waiting.remove_requests(finished) + return finished + + monkeypatch.setattr(OmniARScheduler, "finish_requests", finish) + assert len(scheduler.finish_requests("uncond", RequestStatus.FINISHED_ABORTED)) == 2 + assert not scheduler._pairs and not scheduler._request_pairs and not scheduler.waiting + + +def test_new_prefill_waits_for_initial_streaming_audio(monkeypatch): + request = _request("plain", 20) + scheduler = _scheduler(monkeypatch, [], [request]) + scheduler._prefill_min_output_tokens = 8 + request.append_output_token_ids([0] * 7) + assert scheduler._should_defer_waiting_admission() + request.append_output_token_ids(0) + assert not scheduler._should_defer_waiting_admission() diff --git a/tests/e2e/online_serving/test_breeze_tts_2.py b/tests/e2e/online_serving/test_breeze_tts_2.py index 29cd43b68..fb469768c 100644 --- a/tests/e2e/online_serving/test_breeze_tts_2.py +++ b/tests/e2e/online_serving/test_breeze_tts_2.py @@ -1,100 +1,108 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -""" -E2E Online tests for Breeze-TTS-2 with text input and audio output. - -These tests verify the /v1/audio/speech endpoint works correctly with -actual model inference, not mocks. Breeze-TTS-2 plain mode needs only -``input`` and a speaker tag (``voice`` defaults to ``S0``); reference -cloning requires ``ref_audio`` + ``ref_text`` and is not exercised here. -""" - -import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +"""Breeze design, reference conditioning and CFG through the shared speech client.""" import pytest from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path MODEL = "BreezeBlue/Breeze-TTS-2" -DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0 -MAX_CONCURRENT = 4 -# ~0.5 s of 24 kHz mono PCM_16 in WAV (~24k payload + header). -_MIN_AUDIO_BYTES = 40_000 - - -def get_prompt(prompt_type="text"): - """English prompt for zero-shot TTS.""" - prompts = { - "text": "The weather is nice today, perfect for a walk in the park.", - } - return prompts.get(prompt_type, prompts["text"]) - - -tts_server_params = [ - pytest.param( - OmniServerParams( - model=MODEL, - stage_config_path=get_deploy_config_path("breeze_tts_2.yaml"), - server_args=["--trust-remote-code", "--disable-log-stats"], - ), - id="breeze_tts_2", - ) +pytestmark = [ + pytest.mark.core_model, + pytest.mark.advanced_model, + pytest.mark.tts, + pytest.mark.parametrize( + "omni_server", + [OmniServerParams(model=MODEL, stage_config_path=get_deploy_config_path("breeze_tts_2.yaml"))], + indirect=True, + ), ] -@pytest.mark.slow -@pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_server, online_client) -> None: - """ - Test plain text-to-audio via OpenAI API. - Deploy Setting: breeze_tts_2.yaml - Input Modal: text - Output Modal: audio - Input Setting: stream=False - Datasets: few requests (max_num_seqs=8) - """ - request_config = { - "model": omni_server.model, - "input": get_prompt(), - "stream": False, - "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, - "response_format": "wav", - "voice": "S0", - "sample_rate": 24000, - "min_audio_bytes": _MIN_AUDIO_BYTES, - } - online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) +def test_chinese_voice_design(omni_server, online_client, run_level) -> None: + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "你好，欢迎使用语音合成系统。今天的天气很好。", + "instructions": "年轻女性，使用标准普通话，语速适中。", + "extra_params": {"guidance_scale": 4.0}, + "transcript_language": "zh", + # Whisper-small transcribes "合成" as "和程" on this fixed clip. + # Use the stronger primary ASR without changing the expected text + # or similarity threshold, or retrying after a failed assertion. + "transcript_model": "large-v3", + "response_format": "wav", + "sample_rate": 24000, + "min_audio_bytes": 24000, + "seed": 42, + "max_new_tokens": 8 if run_level == "core_model" else 250, + } + ) -@pytest.mark.slow -@pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_002(omni_server, online_client) -> None: - """ - Test streaming text-to-audio via OpenAI API (async-chunk SSE deltas). - Deploy Setting: breeze_tts_2.yaml - Input Modal: text - Output Modal: audio - Input Setting: stream=True - Datasets: single request - """ - request_config = { - "model": omni_server.model, - "input": get_prompt(), - "stream": True, - "stream_format": "audio", - "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, - "response_format": "wav", - "voice": "S0", - "sample_rate": 24000, - "min_audio_bytes": _MIN_AUDIO_BYTES, - } - online_client.send_audio_speech_request(request_config) +def test_english_streaming(omni_server, online_client, run_level) -> None: + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello, this is a demonstration of natural speech synthesis.", + "instructions": "A warm, clear female voice speaking English at a moderate pace.", + "transcript_language": "en", + "response_format": "pcm", + "sample_rate": 24000, + "min_audio_bytes": 24000, + "stream": True, + "stream_format": "audio", + # The reference eager runtime scores 0.56 dB on this same prompt + # (seed 42, CFG=1); the shared 1 dB floor rejects it as well. + "min_hnr_db": 0.0, + "seed": 42, + "max_new_tokens": 8 if run_level == "core_model" else 250, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_two_concurrent_requests(omni_server, online_client, run_level) -> None: + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Each request keeps its own audio generation state.", + "instructions": "A clear English voice.", + "extra_params": {"guidance_scale": 4.0}, + "transcript_language": "en", + "response_format": "wav", + "sample_rate": 24000, + "min_audio_bytes": 24000, + "seed": 123, + "max_new_tokens": 8 if run_level == "core_model" else 250, + }, + request_num=2, + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("guidance_scale", [1.0, 4.0]) +def test_reference_cloning_and_direction(omni_server, online_client, run_level, guidance_scale) -> None: + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Please read this sentence clearly and naturally.", + "instructions": "Keep the reference voice. Speak gently and calmly.", + "ref_audio": get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True), + "ref_text": "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you.", + "task_type": "Base", + "extra_params": {"guidance_scale": guidance_scale}, + "transcript_language": "en", + "response_format": "wav", + "sample_rate": 24000, + "min_audio_bytes": 24000, + "seed": 42, + "max_new_tokens": 8 if run_level == "core_model" else 250, + } + ) diff --git a/tests/entrypoints/openai_api/test_breeze_tts_2_adapter.py b/tests/entrypoints/openai_api/test_breeze_tts_2_adapter.py new file mode 100644 index 000000000..159700475 --- /dev/null +++ b/tests/entrypoints/openai_api/test_breeze_tts_2_adapter.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU regression coverage for Breeze's public speech request compatibility.""" + +from concurrent.futures import Executor +from dataclasses import dataclass +from unittest.mock import AsyncMock, Mock + +import numpy as np +import pytest +from pydantic import ValidationError +from vllm import SamplingParams + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter +from vllm_omni.model_executor.models.breeze_tts_2.prompt import DEFAULT_INSTRUCTION + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +REFERENCE = "data:audio/wav;base64,cHVibGljLXRlc3QtZGF0YQ==" + + +@dataclass +class ModelConfigStub: + max_model_len: int + + +@dataclass +class EngineStub: + model_config: ModelConfigStub + + +@dataclass +class ServerStub: + _tts_executor: Executor | None + _validate_ref_audio_format: Mock + _resolve_ref_audio: AsyncMock + + +@pytest.fixture +def adapter(): + engine = EngineStub(model_config=ModelConfigStub(max_model_len=32)) + server = ServerStub( + _tts_executor=None, + _validate_ref_audio_format=Mock(return_value=None), + _resolve_ref_audio=AsyncMock(return_value=(np.ones(1920, dtype=np.float32), 24000, None)), + ) + result = BreezeTTS2Adapter(SpeechServingContext(server=server, engine_client=engine)) + result.tokenizer = Mock() + result.tokenizer.encode.side_effect = lambda text, **kwargs: list(text.encode("utf-8")) + return result + + +@pytest.mark.parametrize("voice,tag", [(None, "S0"), ("default", "S0"), ("", "S0"), ("S3", "S3"), ("[S2]", "S2")]) +def test_speaker_tags_reach_the_conditioning_prompt(adapter, voice, tag): + request = OpenAICreateSpeechRequest(input="Hello.", voice=voice) + assert adapter.validate(request) is None + prompt = adapter._build_prompt(request, SamplingParams(), None) + target = bytes(prompt["additional_information"]["breeze_prompt"]["target_ids"]).decode() + assert target == f"[{tag}]{DEFAULT_INSTRUCTION}Hello." + + +def test_explicit_empty_instruction_remains_empty(adapter): + request = OpenAICreateSpeechRequest(input="Hello.", instructions="") + prompt = adapter._build_prompt(request, SamplingParams(), None) + target = bytes(prompt["additional_information"]["breeze_prompt"]["target_ids"]).decode() + assert target == "[S0]Hello." + + +@pytest.mark.parametrize( + "extra,expected", + [({"cfg_scale": 4.0}, 4.0), ({"guidance_scale": 2.0, "cfg_scale": 4.0}, 2.0)], +) +def test_cfg_alias_and_explicit_guidance_precedence(adapter, extra, expected): + request = OpenAICreateSpeechRequest(input="Hello.", extra_params=extra) + assert adapter.validate(request) is None + prompt = adapter._build_prompt(request, SamplingParams(), None) + assert prompt["additional_information"]["breeze_prompt"]["guidance_scale"] == expected + assert prompt["additional_information"]["cfg_group"]["role"] == "cond" + + +@pytest.mark.parametrize("reference", [None, REFERENCE, [REFERENCE]]) +def test_custom_voice_and_false_x_vector_are_compatible(adapter, reference): + request = OpenAICreateSpeechRequest( + input="Hello.", + task_type="CustomVoice", + x_vector_only_mode=False, + ref_audio=reference, + ref_text="Reference words." if reference is not None else None, + ) + assert adapter.validate(request) is None + if reference is not None: + adapter.ctx.server._validate_ref_audio_format.assert_called_once_with(REFERENCE) + + +@pytest.mark.asyncio +async def test_single_reference_list_is_resolved_once_and_preserved_in_prompt(adapter): + request = OpenAICreateSpeechRequest( + input="New words.", voice="S1", ref_audio=[REFERENCE], ref_text="Reference words.", task_type="Base" + ) + assert adapter.validate(request) is None + prepared = await adapter.build(request, [SamplingParams()], has_inline_ref_audio=True) + adapter.ctx.server._resolve_ref_audio.assert_awaited_once_with(REFERENCE) + info = prepared.prompt["additional_information"] + assert bytes(info["breeze_prompt"]["reference_ids"]).decode() == "[S1]Reference words." + assert info["breeze_prompt"]["reference_frames"] == 1 + np.testing.assert_array_equal(info["reference_waveform"].numpy(), np.ones(1920, dtype=np.float32)) + + +@pytest.mark.parametrize( + "kwargs,fragment", + [ + ({"input": " \n"}, "empty"), + ({"speed": 1.2}, "speed"), + ({"language": "English"}, "language"), + ({"ref_audio_2": REFERENCE}, "one reference"), + ({"speaker_embedding": [0.1, 0.2]}, "one reference"), + ({"x_vector_only_mode": True}, "x_vector_only_mode"), + ({"ref_audio": [REFERENCE, REFERENCE], "ref_text": "words"}, "exactly one"), + ({"ref_audio": REFERENCE}, "transcript"), + ({"ref_audio": REFERENCE, "ref_text": " "}, "transcript"), + ({"ref_text": "words"}, "requires ref_audio"), + ({"task_type": "Base"}, "VoiceDesign"), + ({"ref_audio": REFERENCE, "ref_text": "words", "task_type": "VoiceDesign"}, "Base"), + ({"extra_params": {"unknown": 1}}, "Unsupported"), + ({"extra_params": {"temperature": True}}, "finite number"), + ({"extra_params": {"temperature": float("nan")}}, "finite number"), + ({"extra_params": {"cfg_scale": float("inf")}}, "finite number"), + ({"extra_params": {"cfg_scale": 0}}, "positive"), + ({"extra_params": {"guidance_scale": 0, "cfg_scale": 4}}, "positive"), + ({"extra_params": {"repetition_penalty": 0}}, "positive"), + ({"extra_params": {"temperature": -0.1}}, "temperature"), + ({"extra_params": {"top_p": 0}}, "top_p"), + ({"extra_params": {"top_p": 1.1}}, "top_p"), + ({"extra_params": {"top_k": 1.5}}, "top_k"), + ({"extra_params": {"top_k": -2}}, "top_k"), + ], +) +def test_unsupported_or_invalid_inputs_fail_before_build(adapter, kwargs, fragment): + request = OpenAICreateSpeechRequest(**{"input": "Hello.", **kwargs}) + error = adapter.validate(request) + assert error is not None and fragment in error + adapter.ctx.server._resolve_ref_audio.assert_not_awaited() + + +def test_empty_reference_list_is_rejected_by_request_schema(): + with pytest.raises(ValidationError, match="list cannot be empty"): + OpenAICreateSpeechRequest(input="Hello.", ref_audio=[]) + + +def test_reference_format_validation_error_is_preserved(adapter): + adapter.ctx.server._validate_ref_audio_format.return_value = "Unsupported reference format" + request = OpenAICreateSpeechRequest(input="Hello.", ref_audio=[REFERENCE], ref_text="words") + assert adapter.validate(request) == "Unsupported reference format" + + +@pytest.mark.parametrize("requested,expected", [(None, 9), (4, 4), (200, 9)]) +def test_sampling_override_clones_defaults_and_caps_remaining_context(adapter, requested, expected): + defaults = [SamplingParams(max_tokens=20, top_k=50), SamplingParams(max_tokens=100)] + request = OpenAICreateSpeechRequest( + input="Hello.", max_new_tokens=requested, extra_params={"repetition_penalty": 1.7} + ) + result = adapter.apply_sampling_overrides(defaults, request, {"prompt_token_ids": [0] * 23}) + assert result is not defaults and result[0] is not defaults[0] + assert result[0].max_tokens == expected + assert result[0].repetition_penalty == 1.7 + assert (defaults[0].max_tokens, defaults[0].repetition_penalty) == (20, 1.0) + assert result[1].max_tokens == defaults[1].max_tokens == 100 + + +def test_sampling_override_normalizes_disabled_top_k_without_mutating_default(adapter): + default = SamplingParams(max_tokens=20) + default.top_k = 0 + result = adapter.apply_sampling_overrides([default], OpenAICreateSpeechRequest(input="Hello.")) + assert result[0].top_k == -1 + assert default.top_k == 0 + + +@pytest.mark.parametrize("prompt_length", [32, 33]) +def test_filled_context_is_rejected_without_mutating_sampling_defaults(adapter, prompt_length): + default = SamplingParams(max_tokens=20) + request = OpenAICreateSpeechRequest(input="Hello.", max_new_tokens=4) + with pytest.raises(ValueError, match="fills the model context"): + adapter.apply_sampling_overrides([default], request, {"prompt_token_ids": [0] * prompt_length}) + assert default.max_tokens == 20 diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 7295418fb..573816a0e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -16,6 +16,7 @@ from inspect import Signature, signature from pathlib import Path from types import SimpleNamespace +import anyio import numpy as np import pytest import torch @@ -5779,21 +5780,34 @@ class TestTTSAsyncOffloading: @pytest.mark.asyncio async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server): + closed = asyncio.Event() + async def pcm_generator(): - yield SimpleNamespace( - multimodal_output={ - "audio": torch.zeros(16, dtype=torch.float32), - "sr": 24000, - } - ) - await asyncio.sleep(0) + try: + yield OmniRequestOutput( + request_id="req-close", + final_output_type="audio", + _multimodal_output={ + "audio": torch.zeros(16, dtype=torch.float32), + "sr": 24000, + }, + ) + finally: + # Engine abort waits for stage acknowledgments. Retain an + # actual cancellation checkpoint to cover ASGI cancel scopes. + await anyio.sleep(0) + closed.set() qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = ("artifact-close", False) - stream = qwen3_tts_server._generate_audio_chunks(pcm_generator(), "req-close") + engine_stream = pcm_generator() + stream = qwen3_tts_server._generate_audio_chunks(engine_stream, "req-close") assert await anext(stream) - await stream.aclose() + with anyio.CancelScope() as scope: + scope.cancel() + await stream.aclose() + assert closed.is_set() assert "req-close" not in qwen3_tts_server._request_ref_audio_artifact_keys assert ("artifact-close", False) not in qwen3_tts_server._ref_audio_model_artifact_ready diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 8ecb137a2..cc88969d1 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -140,6 +140,7 @@ _PIPELINE_STAGES = [ "audio_vae", "aura", "breeze_tts_2", + "breeze_tts_2_codec", "code2wav", "cosyvoice3_code2wav", "cosyvoice3_talker", diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index fdac862ca..6bbf26d1a 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -5,6 +5,7 @@ import asyncio import re from types import SimpleNamespace +import anyio import pytest from vllm.lora.request import LoRARequest from vllm.sampling_params import RequestOutputKind, SamplingParams @@ -545,6 +546,87 @@ def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker): asyncio.run(run_test()) +@pytest.mark.cpu +def test_generate_waits_for_abort_inside_cancelled_asgi_scope(): + async def run_test(): + submitted = anyio.Event() + acknowledged = [] + + async def add_request(**kwargs): + submitted.set() + + async def abort(request_ids, timeout=None): + assert timeout == async_omni_mod.ABORT_TIMEOUT_S + await anyio.sleep(0) + acknowledged.extend(request_ids) + + omni = get_async_omni_instance(fake_add_request=add_request, fake_abort_request=abort) + + async def collect(): + async for _ in omni.generate( + prompt={"prompt": "prompt"}, + request_id="cancel-asgi", + sampling_params_list=[SamplingParams()], + output_modalities=["image"], + ): + pass + + async with anyio.create_task_group() as group: + group.start_soon(collect) + await submitted.wait() + group.cancel_scope.cancel() + assert len(acknowledged) == 1 and acknowledged[0].startswith("cancel-asgi-") + assert omni.request_states == {} + + asyncio.run(run_test()) + + +@pytest.mark.cpu +def test_generate_bounds_abort_inside_cancelled_asgi_scope(monkeypatch): + monkeypatch.setattr(async_omni_mod, "ABORT_TIMEOUT_S", 0.05) + + async def run_test(): + submitted = anyio.Event() + abort_started = anyio.Event() + abort_finished = anyio.Event() + + async def add_request(**kwargs): + submitted.set() + + async def abort(request_ids, timeout=None): + assert len(request_ids) == 1 + assert timeout == 0.05 + abort_started.set() + try: + await asyncio.Future() + finally: + abort_finished.set() + + omni = get_async_omni_instance(fake_add_request=add_request, fake_abort_request=abort) + + async def collect(): + async for _ in omni.generate( + prompt={"prompt": "prompt"}, + request_id="cancel-asgi-hanging-abort", + sampling_params_list=[SamplingParams()], + output_modalities=["image"], + ): + pass + + async def cancel_request(): + async with anyio.create_task_group() as group: + group.start_soon(collect) + await submitted.wait() + group.cancel_scope.cancel() + + await asyncio.wait_for(cancel_request(), timeout=1.0) + assert abort_started.is_set() + assert abort_finished.is_set() + assert omni.request_states == {} + + asyncio.run(run_test()) + + @pytest.mark.cpu def test_generate_yields_streaming_diffusion_chunks_before_final(): """AsyncOmni.generate yields every intermediate diffusion chunk before the final one.""" diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 1b33850da..cf22c79f1 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -1158,6 +1158,7 @@ def test_async_omni_errored_property_alive(): ) assert omni.errored is False + assert omni.is_stopped is False def test_async_omni_errored_property_dead_engine(): @@ -1168,6 +1169,7 @@ def test_async_omni_errored_property_dead_engine(): ) assert omni.errored is True + assert omni.is_stopped is True def test_async_omni_errored_false_when_stage_dead(): diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 5df9dfa7e..ba1053928 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -780,7 +780,11 @@ def _resolve_audio_transcript( audio_bytes = getattr(response, "audio_bytes", None) if not audio_bytes: return None - return convert_audio_bytes_to_text(audio_bytes, language=request_config.get("transcript_language")) + return convert_audio_bytes_to_text( + audio_bytes, + model_size=request_config.get("transcript_model", "small"), + language=request_config.get("transcript_language"), + ) def assert_omni_response(response: Any, request_config: dict[str, Any], run_level): diff --git a/tests/helpers/client.py b/tests/helpers/client.py index e1042321a..e9b523a43 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -1444,6 +1444,7 @@ class OnlineOmniClient: - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body) - min_audio_bytes: optional minimum ``len(audio_bytes)`` checked in ``assert_audio_speech_response`` - transcript_expected_text: local expected spoken text; defaults to ``input`` + - transcript_model: primary Whisper model for content checks; defaults to ``small`` - timeout: request timeout in seconds (float, optional, default 120.0) - stream: whether to use streaming API (bool, optional, default False) @@ -1467,6 +1468,7 @@ class OnlineOmniClient: "task_type", "ref_text", "ref_audio", + "extra_params", "language", "max_new_tokens", "seed", diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py index 55956a6f2..660a7752f 100644 --- a/tests/helpers/tests/test_assertions.py +++ b/tests/helpers/tests/test_assertions.py @@ -11,6 +11,7 @@ from tests.helpers.assertions import ( _resolve_audio_transcript, assert_audio_speech_response, ) +from tests.helpers.client import OmniResponse pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -75,6 +76,26 @@ def test_resolve_transcript_leaves_language_unset_by_default(monkeypatch): assert captured["model_size"] == "small" +def test_resolve_transcript_uses_configured_primary_model(monkeypatch): + captured = _capture_transcribe(monkeypatch) + response = OmniResponse(audio_content=None, audio_bytes=b"fake-wav") + + _resolve_audio_transcript( + response, + { + "input": "你好", + "response_format": "wav", + "transcript_language": "zh", + "transcript_model": "large-v3", + }, + "advanced_model", + speech_api=True, + ) + + assert captured["model_size"] == "large-v3" + assert captured["language"] == "zh" + + @pytest.mark.parametrize( ("request_config", "expected_text"), [ diff --git a/tests/model_executor/models/breeze_tts_2/__init__.py b/tests/model_executor/models/breeze_tts_2/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py b/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py deleted file mode 100644 index 75dcbb1fb..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py +++ /dev/null @@ -1,128 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from types import SimpleNamespace - -import numpy as np -import pytest -import torch - -from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( - BreezeReferenceAudioTokenizer, - resolve_audio_tokenizer_path, -) - - -def _patch_hf_api(monkeypatch, snapshot_download): - """Point the resolver's tagged HF helper at a fake API object.""" - from vllm.transformers_utils import repo_utils - - monkeypatch.setattr(repo_utils, "hf_api", lambda: SimpleNamespace(snapshot_download=snapshot_download)) - - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _Tokenizer: - def encode(self, audio, sr=None): - assert isinstance(audio, np.ndarray) - assert sr == 16000 - return {"audio_codes": [torch.arange(32).reshape(16, 2)]} - - -def test_reference_audio_tokenizer_normalizes_codebook_major_output(): - adapter = BreezeReferenceAudioTokenizer(_Tokenizer()) - - codes = adapter.encode(np.zeros(32, dtype=np.float32), 16000) - - assert tuple(codes.shape) == (2, 16) - assert codes.dtype == torch.int16 - assert codes.device.type == "cpu" - - -def test_reference_audio_tokenizer_rejects_waveform_without_sample_rate(): - adapter = BreezeReferenceAudioTokenizer(_Tokenizer()) - - try: - adapter.encode(np.zeros(32, dtype=np.float32)) - except ValueError as exc: - assert "sample_rate is required" in str(exc) - else: - raise AssertionError("missing sample rate should be rejected") - - -def test_resolve_prefers_local_directory_without_hub_access(tmp_path, monkeypatch): - bundled = tmp_path / "audio_tokenizer" - bundled.mkdir() - - def _no_hub(*_args, **_kwargs): - raise AssertionError("local directories must not touch the hub") - - _patch_hf_api(monkeypatch, _no_hub) - - assert resolve_audio_tokenizer_path(str(tmp_path)) == bundled - - -def test_resolve_local_directory_without_bundled_tokenizer_returns_none(tmp_path, monkeypatch): - def _no_hub(*_args, **_kwargs): - raise AssertionError("local directories must not touch the hub") - - _patch_hf_api(monkeypatch, _no_hub) - - assert resolve_audio_tokenizer_path(str(tmp_path)) is None - - -def test_resolve_repo_id_uses_hf_snapshot(tmp_path, monkeypatch): - snapshot_root = tmp_path / "snapshot" - (snapshot_root / "audio_tokenizer").mkdir(parents=True) - calls = [] - - def _fake_snapshot_download(repo_id, **kwargs): - calls.append((repo_id, kwargs)) - if kwargs.get("local_files_only"): - raise FileNotFoundError("cold cache") - return str(snapshot_root) - - _patch_hf_api(monkeypatch, _fake_snapshot_download) - - resolved = resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") - - assert resolved == snapshot_root / "audio_tokenizer" - # The cold-cache probe runs first; the online fetch carries the pattern. - assert calls[0] == ( - "BreezeBlue/Breeze-TTS-2", - {"allow_patterns": ["audio_tokenizer/*"], "local_files_only": True}, - ) - assert calls[1] == ("BreezeBlue/Breeze-TTS-2", {"allow_patterns": ["audio_tokenizer/*"]}) - - -def test_resolve_repo_id_prefers_complete_cached_snapshot(tmp_path, monkeypatch): - snapshot_root = tmp_path / "snapshot" - (snapshot_root / "audio_tokenizer").mkdir(parents=True) - calls = [] - - def _fake_snapshot_download(repo_id, **kwargs): - calls.append((repo_id, kwargs)) - assert kwargs.get("local_files_only"), "cached snapshot must not hit the network" - return str(snapshot_root) - - _patch_hf_api(monkeypatch, _fake_snapshot_download) - - resolved = resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") - - assert resolved == snapshot_root / "audio_tokenizer" - assert len(calls) == 1 - - -def test_resolve_repo_id_without_bundled_tokenizer_returns_none(tmp_path, monkeypatch): - snapshot_root = tmp_path / "snapshot" - snapshot_root.mkdir() - - def _fake_snapshot_download(repo_id, **kwargs): - if kwargs.get("local_files_only"): - raise FileNotFoundError("cold cache") - return str(snapshot_root) - - _patch_hf_api(monkeypatch, _fake_snapshot_download) - - assert resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") is None diff --git a/tests/model_executor/models/breeze_tts_2/test_batching.py b/tests/model_executor/models/breeze_tts_2/test_batching.py deleted file mode 100644 index 41fc0b9ba..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_batching.py +++ /dev/null @@ -1,95 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from typing import Any - -import pytest -import torch - -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( - BreezeTTS2TalkerForGeneration, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _RequestHead: - def __call__(self, hidden: torch.Tensor) -> torch.Tensor: - # Hidden row 0 selects codec id 2; row 1 selects codec id 3. - code0 = 2 + int(hidden.reshape(-1)[0].item()) - logits = torch.zeros((hidden.shape[0], 9), dtype=hidden.dtype) - logits[:, code0] = 1.0 - return logits - - -def _model() -> BreezeTTS2TalkerForGeneration: - model = object.__new__(BreezeTTS2TalkerForGeneration) - model.num_codebooks = 4 - model.codebook_vocab_size = 8 - model.codebook_size = 8 - model.codebook_eos_token_id = 0 - model.audio_token_id = 100 - model.audio_eos_token_id = 101 - model.lm_head = _RequestHead() - model.logits_processor = lambda head, hidden: head(hidden) - model._codec_disallowed_mask = None - model._batch_state = None - model._batch_state_spans = None - model._async_chunk = True - model._golden_dump_dir = None - model._generate_depth_codes = lambda _hidden, code0: _frame(code0).unsqueeze(0) - return model - - -def _frame(code0: int) -> torch.Tensor: - return torch.tensor([code0, code0 + 1, code0 + 2, code0 + 3], dtype=torch.long) - - -def _step(model: BreezeTTS2TalkerForGeneration, infos: list[dict[str, Any]], hidden: torch.Tensor): - return model.make_omni_output( - hidden, - model_intermediate_buffer=infos, - request_token_spans=[(0, 1), (1, 2)], - ) - - -def test_two_requests_accumulate_and_emit_independent_frames(): - model = _model() - infos = [ - { - "breeze_generated_frames": 1, - "breeze_max_new_frames": 2, - "breeze_audio_codes": _frame(2).reshape(1, 4), - }, - {"breeze_generated_frames": 0, "breeze_max_new_frames": 3}, - ] - hidden = torch.tensor([[0.0], [1.0]], dtype=torch.float32) - - first = _step(model, infos, hidden) - assert [tuple(item.shape) for item in first.multimodal_outputs["codes"]["audio"]] == [(1, 4), (1, 4)] - assert torch.equal(first.multimodal_outputs["codes"]["audio"][0][0], _frame(2)) - assert torch.equal(first.multimodal_outputs["codes"]["audio"][1][0], _frame(3)) - - second = _step(model, infos, hidden) - # Request 0 reaches its two-frame budget and emits only a finish marker; - # request 1 still has budget and emits its next tail frame. - assert second.multimodal_outputs["codes"]["audio"][0].numel() == 0 - assert torch.equal(second.multimodal_outputs["codes"]["audio"][1][0], _frame(3)) - # Async streaming only maintains the per-frame tail; the cumulative - # buffer is deliberately not grown (request 0 keeps its pre-seeded - # entry untouched, request 1 never gains one). - assert infos[0]["breeze_audio_codes"].shape == (1, 4) - assert "breeze_audio_codes" not in infos[1] - assert infos[0]["breeze_force_eos"] is True - assert infos[1]["breeze_force_eos"] is False - - -def test_finished_requests_only_clear_model_owned_batch_mapping(): - model = _model() - model._batch_state = [{"id": "a"}, {"id": "b"}] - model._batch_state_spans = [(0, 1), (1, 2)] - - model.on_requests_finished(["a", "b"]) - - assert model._batch_state is None - assert model._batch_state_spans is None diff --git a/tests/model_executor/models/breeze_tts_2/test_codec_loading.py b/tests/model_executor/models/breeze_tts_2/test_codec_loading.py deleted file mode 100644 index f558d067a..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_codec_loading.py +++ /dev/null @@ -1,68 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from typing import Any - -import pytest -import torch -from torch import nn - -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_codec import ( - BreezeTTS2MimiCodec, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _BundledTokenizer: - def __init__(self) -> None: - self.model = nn.Linear(2, 2) - - def named_parameters(self): - for name, parameter in self.model.named_parameters(): - yield f"model.{name}", parameter - - def decode(self, _encoded: Any) -> tuple[list[torch.Tensor], int]: - return [torch.zeros(4)], 24_000 - - -def test_bundled_qwen_tokenizer_consumes_unused_root_fallback_weights(tmp_path): - bundled = tmp_path / "audio_tokenizer" - bundled.mkdir() - codec = object.__new__(BreezeTTS2MimiCodec) - codec.model_path = str(tmp_path) - codec.vllm_config = None - # load_weights reuses the path resolved once in __init__. - codec._tokenizer_path = bundled - codec._codec = None - codec._audio_tokenizer = _BundledTokenizer() - codec._loaded_local_weights = False - - weights = [ - ("backbone_model.layers.0.norm.weight", torch.ones(2)), - ("codec_model.decoder.layers.0.conv.weight", torch.ones(2)), - ("lm_head.weight", torch.ones(2, 2)), - ] - - loaded = codec.load_weights(iter(weights)) - - assert "_audio_tokenizer.model.weight" in loaded - assert "_audio_tokenizer.model.bias" in loaded - # The bundled tokenizer was preloaded, so load_weights must not try to copy - # the root Mimi fallback into the nonexistent ``self._codec`` module. - assert codec._codec is None - - -def test_codec_warmup_with_single_token_returns_empty_audio(): - codec = object.__new__(BreezeTTS2MimiCodec) - codec._num_codebooks = 4 - codec._codebook_size = 8 - codec._sample_rate = 24_000 - codec._audio_tokenizer = None - codec._codec = None - codec._async_chunk = False - - output = codec.forward(torch.tensor([1], dtype=torch.long), seq_token_counts=[1]) - - assert output.multimodal_outputs["model_outputs"][0].numel() == 0 - assert output.multimodal_outputs["sr"][0].item() == 24_000 diff --git a/tests/model_executor/models/breeze_tts_2/test_configuration.py b/tests/model_executor/models/breeze_tts_2/test_configuration.py deleted file mode 100644 index 950f420cf..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_configuration.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import pytest -from transformers.models.qwen3 import Qwen3Config - -from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( - BreezeTTS2Config, - BreezeTTS2DepthDecoderConfig, - BreezeTTS2TextEncoderConfig, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _config() -> BreezeTTS2Config: - return BreezeTTS2Config( - backbone_config={ - "model_type": "qwen3", - "hidden_size": 32, - "intermediate_size": 64, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 8, - "vocab_size": 128, - "max_position_embeddings": 512, - }, - text_encoder_config={ - "model_type": "t5gemma2_text", - "hidden_size": 24, - "intermediate_size": 48, - "num_hidden_layers": 6, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "head_dim": 8, - "vocab_size": 256, - }, - depth_decoder_config={ - "model_type": "breeze_depth_decoder_model", - "hidden_size": 32, - "backbone_hidden_size": 32, - "audio_embed_size": 32, - "intermediate_size": 64, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 8, - "num_codebooks": 16, - "vocab_size": 2051, - "rope_scaling": { - "rope_type": "llama3", - "factor": 32.0, - "low_freq_factor": 0.001953125, - "high_freq_factor": 0.0078125, - "original_max_position_embeddings": 16, - }, - }, - codec_config={"model_type": "mimi", "codebook_size": 2048, "num_quantizers": 32}, - num_codebooks=16, - vocab_size=2051, - text_vocab_size=262158, - ) - - -def test_nested_configs_and_text_config_use_qwen_backbone(): - config = _config() - - assert isinstance(config.backbone_config, Qwen3Config) - assert config.get_text_config() is config.backbone_config - assert config.get_text_config().vocab_size == 128 - assert isinstance(config.text_encoder_config, BreezeTTS2TextEncoderConfig) - assert isinstance(config.depth_decoder_config, BreezeTTS2DepthDecoderConfig) - assert config.num_codebooks == config.depth_decoder_config.num_codebooks == 16 - assert config.codec_config["codebook_size"] == 2048 - - -def test_transformers_five_roundtrip_preserves_nested_configs_and_rope(): - config = _config() - - roundtrip = BreezeTTS2Config(**dict(config.to_dict())) - - assert isinstance(roundtrip.backbone_config, Qwen3Config) - assert isinstance(roundtrip.text_encoder_config, BreezeTTS2TextEncoderConfig) - assert isinstance(roundtrip.depth_decoder_config, BreezeTTS2DepthDecoderConfig) - assert roundtrip.depth_decoder_config.rope_parameters["rope_theta"] == 500_000.0 - - -def test_nested_autoconfig_registration_is_idempotent(): - import importlib - - module = importlib.import_module("vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2") - importlib.reload(module) - - assert module.BreezeTTS2Config.model_type == "breeze" - - -def test_text_encoder_defaults_to_upstream_layer_pattern(): - config = BreezeTTS2TextEncoderConfig(num_hidden_layers=6) - - assert config.layer_types == [ - "sliding_attention", - "sliding_attention", - "sliding_attention", - "sliding_attention", - "sliding_attention", - "full_attention", - ] - - -def test_invalid_codebook_layout_is_rejected(): - with pytest.raises(ValueError, match="positive"): - BreezeTTS2DepthDecoderConfig(num_codebooks=0) diff --git a/tests/model_executor/models/breeze_tts_2/test_instruction.py b/tests/model_executor/models/breeze_tts_2/test_instruction.py deleted file mode 100644 index c7f2fef66..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_instruction.py +++ /dev/null @@ -1,157 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from types import SimpleNamespace - -import pytest -from vllm import SamplingParams - -from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter -from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( - BreezeTTS2PromptBuilder, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _Tokenizer: - unk_token_id = -1 - - def convert_tokens_to_ids(self, token): - return {"<|AUDIO|>": 262144, "<|audio_eos|>": 262145}.get(token, self.unk_token_id) - - def __call__(self, texts, *, add_special_tokens=True, padding=False): - if isinstance(texts, str): - texts = [texts] - ids = [[1000 + ord(char) % 1000 for char in text] for text in texts] - if add_special_tokens: - ids = [[2, *row, 1] for row in ids] - return {"input_ids": ids} - - def decode(self, ids, *, skip_special_tokens=False): - del skip_special_tokens - return "".join(chr(int(token_id) - 1000) for token_id in ids if token_id not in (1, 2)) - - -def test_instruction_prompt_remains_one_independent_text_segment(): - builder = BreezeTTS2PromptBuilder( - _Tokenizer(), - SimpleNamespace( - audio_token_id=262144, - audio_eos_token_id=262145, - pad_token_id=0, - backbone_config=SimpleNamespace(vocab_size=100000), - num_codebooks=16, - codec_config={"codebook_size": 2048}, - ), - ) - - prompt = builder.build( - { - "template": "tts_instruction", - "speaker": "S2", - "instruction": "Use a bright and friendly delivery.", - "text": "Hello from Breeze.", - } - ) - - info = prompt["additional_information"] - assert info["template"] == "tts_instruction" - assert "input_values" not in info - assert bool(info["text_ids_mask"].all()) - assert info["text_ids_len"].tolist() == [len(info["prompt_ids"])] - - -def _request(extra_params=None, **overrides): - base = dict( - input="Hello.", - ref_text=None, - ref_audio=None, - task_type=None, - speed=1.0, - max_new_tokens=None, - language=None, - speaker_embedding=None, - x_vector_only_mode=None, - extra_params=extra_params or {}, - ) - base.update(overrides) - return SimpleNamespace(**base) - - -def _adapter() -> BreezeTTS2Adapter: - adapter = object.__new__(BreezeTTS2Adapter) - adapter.ctx = SimpleNamespace(server=SimpleNamespace()) - return adapter - - -def test_cfg_scale_one_is_explicitly_supported_without_companion_request(): - adapter = _adapter() - - assert adapter.validate(_request({"guidance_scale": 1.0})) is None - assert adapter.validate(_request({"cfg_scale": 1.0})) is None - - -def test_non_unit_cfg_is_rejected_until_negative_branch_exists(): - adapter = _adapter() - - assert "guidance_scale=1.0" in adapter.validate(_request({"guidance_scale": 2.0})) - assert "guidance_scale=1.0" in adapter.validate(_request({"cfg_scale": 0.5})) - assert "negative_prompt" in adapter.validate(_request({"negative_prompt": "noisy"})) - - -def test_unsupported_conditioning_fields_are_rejected_not_ignored(): - adapter = _adapter() - - assert "language" in adapter.validate(_request(language="en")) - assert "speaker_embedding" in adapter.validate(_request(speaker_embedding=[0.1, 0.2])) - assert "x_vector_only_mode" in adapter.validate(_request(x_vector_only_mode=True)) - # Absent fields must keep the request valid. - assert adapter.validate(_request()) is None - assert adapter.validate(_request(x_vector_only_mode=False)) is None - - -def test_multiple_reference_clips_are_rejected_not_silently_truncated(): - adapter = _adapter() - adapter.ctx.server._validate_ref_audio_format = lambda _uri: None - - error = adapter.validate(_request(ref_audio=["a.wav", "b.wav"], ref_text="hello")) - assert "exactly one reference clip" in error - assert "2" in error - # A single-element list stays valid and reaches format validation. - assert adapter.validate(_request(ref_audio=["a.wav"], ref_text="hello")) is None - - -def test_non_greedy_sampling_overrides_are_rejected(): - adapter = _adapter() - - # Greedy values, including the deploy-config defaults, stay valid. - assert adapter.validate(_request({"temperature": 0.0, "top_p": 1.0, "top_k": -1})) is None - assert adapter.validate(_request({"temperature": 0})) is None - # The talker ignores the sampled code0, so any non-greedy override must be - # refused instead of silently producing greedy audio. - assert "greedy" in adapter.validate(_request({"temperature": 0.7})) - assert "greedy" in adapter.validate(_request({"top_p": 0.9})) - assert "greedy" in adapter.validate(_request({"top_k": 50})) - assert "temperature must be a number" in adapter.validate(_request({"temperature": "warm"})) - - -@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf"), "invalid"]) -def test_invalid_repetition_penalties_are_rejected(value): - assert "repetition_penalty" in _adapter().validate(_request({"repetition_penalty": value})) - - -@pytest.mark.parametrize("value", [0.9, 1.0, 1.1, "1.2"]) -@pytest.mark.parametrize("max_new_tokens", [None, 32]) -def test_repetition_penalty_override_reaches_standard_sampling_metadata(value, max_new_tokens): - adapter = _adapter() - request = _request({"repetition_penalty": value}, max_new_tokens=max_new_tokens) - original = [SamplingParams(max_tokens=128, repetition_penalty=1.1)] - - assert adapter.validate(request) is None - updated = adapter.apply_sampling_overrides(original, request) - - assert updated[0].repetition_penalty == float(value) - assert updated[0].max_tokens == (128 if max_new_tokens is None else max_new_tokens) - assert original[0].repetition_penalty == 1.1 - assert original[0].max_tokens == 128 diff --git a/tests/model_executor/models/breeze_tts_2/test_local_modules.py b/tests/model_executor/models/breeze_tts_2/test_local_modules.py deleted file mode 100644 index 2ce54b09a..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_local_modules.py +++ /dev/null @@ -1,152 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import pytest -import torch - -from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( - BreezeTTS2DepthDecoderConfig, - BreezeTTS2TextEncoderConfig, -) -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_depth import ( - BreezeDepthDecoderForCausalLM, -) -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text import ( - BreezeTTS2TextEncoder, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _depth_config() -> BreezeTTS2DepthDecoderConfig: - return BreezeTTS2DepthDecoderConfig( - hidden_size=16, - backbone_hidden_size=16, - audio_embed_size=16, - intermediate_size=32, - num_hidden_layers=1, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=4, - max_position_embeddings=8, - num_codebooks=4, - vocab_size=8, - rope_scaling={ - "rope_type": "llama3", - "rope_theta": 100.0, - "factor": 1.0, - "low_freq_factor": 0.5, - "high_freq_factor": 1.0, - "original_max_position_embeddings": 4, - }, - ) - - -def test_depth_decoder_position_specific_logits_and_offsets(): - config = _depth_config() - decoder = BreezeDepthDecoderForCausalLM(config) - - output = decoder( - input_ids=torch.tensor([[0, 2]]), - backbone_last_hidden_state=torch.randn(1, 16), - use_cache=False, - ) - - assert output.logits.shape == (1, 1, 8) - with torch.no_grad(): - embedding = decoder.model.embed_tokens.weight - expected_first = embedding[2] @ decoder.model.inputs_embeds_projector.weight.T - # The first hidden row is replaced by the backbone state, so compare the - # second input row (codebook 0, id 2 + 0 * vocab_size). - actual_input = decoder.model.embed_tokens(torch.tensor([[2]]))[0, 0] - actual_projected = actual_input @ decoder.model.inputs_embeds_projector.weight.T - assert torch.allclose(expected_first, actual_projected) - - -def _text_config() -> BreezeTTS2TextEncoderConfig: - return BreezeTTS2TextEncoderConfig( - hidden_size=16, - intermediate_size=32, - num_hidden_layers=2, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=8, - vocab_size=32, - layer_types=["sliding_attention", "full_attention"], - rope_parameters={ - "sliding_attention": {"rope_type": "default", "rope_theta": 100.0}, - "full_attention": {"rope_type": "default", "rope_theta": 100.0}, - }, - ) - - -def test_text_encoder_masks_right_padding_independently(monkeypatch): - import vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text as text_module - - class _ParallelEmbedding(torch.nn.Embedding): - def __init__(self, num_embeddings: int, embedding_dim: int): - super().__init__(num_embeddings, embedding_dim) - - class _QKVLinear(torch.nn.Module): - def __init__( - self, - hidden_size: int, - head_size: int, - total_num_heads: int, - total_num_kv_heads: int, - **_: object, - ): - super().__init__() - output_size = (total_num_heads + 2 * total_num_kv_heads) * head_size - self.linear = torch.nn.Linear(hidden_size, output_size, bias=False) - - def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: - return self.linear(hidden_states), None - - class _RowLinear(torch.nn.Module): - def __init__(self, input_size: int, output_size: int, **_: object): - super().__init__() - self.linear = torch.nn.Linear(input_size, output_size, bias=False) - - def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: - return self.linear(hidden_states), None - - class _MergedLinear(torch.nn.Module): - def __init__(self, input_size: int, output_sizes: list[int], **_: object): - super().__init__() - self.linear = torch.nn.Linear(input_size, sum(output_sizes), bias=False) - - def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: - return self.linear(hidden_states), None - - monkeypatch.setattr(text_module, "get_tensor_model_parallel_world_size", lambda: 1) - monkeypatch.setattr(text_module, "VocabParallelEmbedding", _ParallelEmbedding) - monkeypatch.setattr(text_module, "QKVParallelLinear", _QKVLinear) - monkeypatch.setattr(text_module, "RowParallelLinear", _RowLinear) - monkeypatch.setattr(text_module, "MergedColumnParallelLinear", _MergedLinear) - - config = _text_config() - encoder = BreezeTTS2TextEncoder(config).eval() - ids = torch.tensor([[3, 4]]) - padded_ids = torch.tensor([[3, 4, 0, 0]]) - mask = torch.tensor([[True, True, False, False]]) - - with torch.inference_mode(): - short = encoder(ids) - padded = encoder(padded_ids, attention_mask=mask) - - assert short.shape == (1, 2, 16) - assert padded.shape == (1, 4, 16) - assert torch.allclose(short[0, :2], padded[0, :2], atol=1e-5) - - -def test_text_encoder_weight_loader_reports_fused_target_names(): - encoder = object.__new__(BreezeTTS2TextEncoder) - parameter = torch.nn.Parameter(torch.zeros(2, 2)) - parameter.weight_loader = lambda param, tensor, shard_id: param.data.copy_(tensor) - encoder.__dict__["named_parameters"] = lambda: [("layers.0.qkv_proj.weight", parameter)] - - loaded = encoder.load_weights([("layers.0.q_proj.weight", torch.ones(2, 2))]) - - assert loaded == {"layers.0.qkv_proj.weight"} - assert torch.equal(parameter.detach(), torch.ones(2, 2)) diff --git a/tests/model_executor/models/breeze_tts_2/test_model_registration.py b/tests/model_executor/models/breeze_tts_2/test_model_registration.py deleted file mode 100644 index 12a2b31f3..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_model_registration.py +++ /dev/null @@ -1,28 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import pytest -from vllm.model_executor.models.registry import ModelRegistry - -from vllm_omni.model_executor.models.registry import OmniModelRegistry - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_breeze_architectures_resolve_to_omni_classes(): - assert "BreezeForConditionalGeneration" in OmniModelRegistry.get_supported_archs() - assert "BreezeTTS2MimiCodec" in OmniModelRegistry.get_supported_archs() - - talker = OmniModelRegistry._try_load_model_cls("BreezeForConditionalGeneration") - codec = OmniModelRegistry._try_load_model_cls("BreezeTTS2MimiCodec") - assert talker.__name__ == "BreezeTTS2TalkerForGeneration" - assert codec.__name__ == "BreezeTTS2MimiCodec" - - -def test_plugin_registration_publishes_architectures_to_upstream_registry(): - from vllm_omni.engine.arg_utils import register_omni_models_to_vllm - - register_omni_models_to_vllm() - - assert "BreezeForConditionalGeneration" in ModelRegistry.get_supported_archs() - assert "BreezeTTS2MimiCodec" in ModelRegistry.get_supported_archs() diff --git a/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py b/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py deleted file mode 100644 index 069296942..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py +++ /dev/null @@ -1,36 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from pathlib import Path - -import pytest -import yaml - -from vllm_omni.config.pipeline_registry import OMNI_PIPELINES -from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_breeze_pipeline_is_registered_with_talker_and_codec_stages(): - assert OMNI_PIPELINES["breeze"] is BREEZE_TTS_2_PIPELINE - assert [stage.model_stage for stage in BREEZE_TTS_2_PIPELINE.stages] == [ - "breeze_tts_2", - "breeze_tts_2_codec", - ] - assert BREEZE_TTS_2_PIPELINE.stages[1].requires_full_payload_input is True - assert BREEZE_TTS_2_PIPELINE.stages[0].async_chunk_process_next_stage_input_func.endswith( - "talker2codec_async_chunk" - ) - assert BREEZE_TTS_2_PIPELINE.stages[1].retains_state_across_chunks is True - - -def test_breeze_default_deploy_is_streaming(): - deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml" - config = yaml.safe_load(deploy_path.read_text(encoding="utf-8")) - assert config["async_chunk"] is True - assert len(config["stages"]) == 2 - # Stage 0 must prefill each prompt in one step: the talker builds a codec - # frame from the last hidden row of every scheduled request. - assert config["stages"][0]["enable_chunked_prefill"] is False - assert config["connectors"]["connector_of_shared_memory"]["extra"]["breeze_codec_chunk_frames"] == 8 diff --git a/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py b/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py deleted file mode 100644 index 3bd26fce6..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py +++ /dev/null @@ -1,189 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from types import SimpleNamespace - -import pytest -import torch - -from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( - BreezeTTS2PromptBuilder, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _Tokenizer: - pad_token_id = 0 - eos_token_id = 1 - unk_token_id = -1 - _special = { - "<|AUDIO|>": 262144, - "<|audio_eos|>": 262145, - } - - def convert_tokens_to_ids(self, token): - return self._special.get(token, self.unk_token_id) - - def __call__(self, texts, *, add_special_tokens=True, padding=False): - if isinstance(texts, str): - texts = [texts] - ids = [] - for text in texts: - values = [1000 + ord(char) % 1000 for char in text] - if add_special_tokens: - values = [2, *values, 1] - ids.append(values) - return {"input_ids": ids} - - def decode(self, ids, *, skip_special_tokens=False): - del skip_special_tokens - rendered = [] - for token_id in ids: - if token_id == 262144: - rendered.append("<|AUDIO|>") - elif token_id == 262145: - rendered.append("<|audio_eos|>") - elif token_id in (2, 1): - continue - else: - rendered.append(chr(int(token_id) - 1000)) - return "".join(rendered) - - -class _RenderingTokenizer(_Tokenizer): - """Tiny tokenizer that preserves Breeze placeholders after decode.""" - - def __call__(self, texts, *, add_special_tokens=True, padding=False): - if isinstance(texts, str): - texts = [texts] - ids = [] - for text in texts: - row = [] - cursor = 0 - while cursor < len(text): - if text.startswith("<|AUDIO|>", cursor): - row.append(262144) - cursor += len("<|AUDIO|>") - elif text.startswith("<|audio_eos|>", cursor): - row.append(262145) - cursor += len("<|audio_eos|>") - else: - row.append(1000 + ord(text[cursor]) % 1000) - cursor += 1 - if add_special_tokens: - row = [2, *row, 1] - ids.append(row) - return {"input_ids": ids} - - -class _ReferenceEncoder: - def __init__(self): - self.calls = [] - - def encode(self, audio, sample_rate=None): - self.calls.append((audio, sample_rate)) - return torch.arange(32, dtype=torch.int16).reshape(2, 16) - - -def _config(): - return SimpleNamespace( - audio_token_id=262144, - audio_eos_token_id=262145, - pad_token_id=0, - backbone_config=SimpleNamespace(vocab_size=100000), - num_codebooks=16, - codec_config={"codebook_size": 2048}, - ) - - -def test_builder_accepts_missing_codec_config(): - config = _config() - config.codec_config = None - - builder = BreezeTTS2PromptBuilder(_Tokenizer(), config) - - assert builder.codebook_size == 2048 - - -def test_build_instruction_prompt(): - builder = BreezeTTS2PromptBuilder(_Tokenizer(), _config()) - - prompt = builder.build({"text": "hello", "instruction": "happy", "speaker": "S0"}) - - assert prompt["prompt_token_ids"] - info = prompt["additional_information"] - assert "input_values" not in info - assert set(prompt["prompt_token_ids"]) == {0} - assert info["prompt_ids"].numel() == len(prompt["prompt_token_ids"]) - assert info["text_ids_mask"].dtype == torch.bool - assert info["text_ids_len"].tolist() == [len(prompt["prompt_token_ids"])] - assert bool(info["text_ids_mask"].all()) - - -def test_build_reference_prompt_and_reuse_encoder(): - encoder = _ReferenceEncoder() - builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config(), encoder) - - prompt = builder.build( - { - "template": "ref_edit_tata", - "text": "target", - "instruction": "calm", - "ref_text": "reference", - "ref_audio": "ref.wav", - "ref_audio_sample_rate": 16000, - } - ) - - info = prompt["additional_information"] - codes = info["input_values"] - assert tuple(codes.shape) == (2, 16) - assert codes.dtype == torch.int16 - assert set(prompt["prompt_token_ids"]) == {0} - assert info["prompt_ids"].tolist().count(262144) == 2 - assert info["prompt_ids"].tolist().count(262145) == 1 - assert int((~info["text_ids_mask"]).sum()) == 3 - assert info["text_ids_len"].numel() == 2 - assert encoder.calls == [("ref.wav", 16000)] - - -def test_build_reference_prompt_preserves_special_tokens_after_rendering(): - encoder = _ReferenceEncoder() - builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config(), encoder) - - prompt = builder.build( - { - "template": "ref_clone_tata", - "text": "target", - "ref_text": "reference", - "ref_audio": "ref.wav", - "ref_audio_sample_rate": 16000, - } - ) - - info = prompt["additional_information"] - prompt_ids = info["prompt_ids"].tolist() - assert prompt_ids.count(262144) == 2 - assert prompt_ids.count(262145) == 1 - assert int((~info["text_ids_mask"]).sum()) == 3 - assert len(info["text_ids_len"]) == 2 - assert sum(info["text_ids_len"].tolist()) == int(info["text_ids_mask"].sum()) - - -def test_build_reference_prompt_accepts_preencoded_codes_mapping(): - builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config()) - codes = torch.arange(32, dtype=torch.int16).reshape(2, 16) - - prompt = builder.build( - { - "template": "ref_clone_tata", - "text": "target", - "ref_text": "reference", - "codes": {"ref": codes}, - } - ) - - info = prompt["additional_information"] - assert torch.equal(info["input_values"], codes) - assert info["ref_code_len"] == 2 diff --git a/tests/model_executor/models/breeze_tts_2/test_reference_voice.py b/tests/model_executor/models/breeze_tts_2/test_reference_voice.py deleted file mode 100644 index d29b863c4..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_reference_voice.py +++ /dev/null @@ -1,187 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import importlib.util -from pathlib import Path -from types import SimpleNamespace - -import numpy as np -import pytest -import soundfile as sf -import torch - -from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( - BreezeTTS2PromptBuilder, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - -_EXAMPLE_PATH = ( - Path(__file__).resolve().parents[4] - / "examples" - / "offline_inference" - / "text_to_speech" - / "breeze_tts_2" - / "end2end.py" -) - - -class _Tokenizer: - unk_token_id = -1 - - def convert_tokens_to_ids(self, token): - return { - "<|AUDIO|>": 262144, - "<|audio_eos|>": 262145, - }.get(token, self.unk_token_id) - - def __call__(self, texts, *, add_special_tokens=True, padding=False): - if isinstance(texts, str): - texts = [texts] - ids = [[1000 + ord(char) % 1000 for char in text] for text in texts] - for row_index, text in enumerate(texts): - if "<|AUDIO|>" in text or "<|audio_eos|>" in text: - cursor = 0 - row = [] - while cursor < len(text): - if text.startswith("<|AUDIO|>", cursor): - row.append(262144) - cursor += len("<|AUDIO|>") - elif text.startswith("<|audio_eos|>", cursor): - row.append(262145) - cursor += len("<|audio_eos|>") - else: - row.append(1000 + ord(text[cursor]) % 1000) - cursor += 1 - ids[row_index] = row - if add_special_tokens: - ids = [[2, *row, 1] for row in ids] - return {"input_ids": ids} - - def decode(self, ids, *, skip_special_tokens=False): - del skip_special_tokens - rendered = [] - for token_id in ids: - if token_id == 262144: - rendered.append("<|AUDIO|>") - elif token_id == 262145: - rendered.append("<|audio_eos|>") - elif token_id in (1, 2): - continue - else: - rendered.append(chr(int(token_id) - 1000)) - return "".join(rendered) - - -class _ReferenceEncoder: - def __init__(self) -> None: - self.calls: list[tuple[object, int | None]] = [] - - def encode(self, audio, sample_rate=None): - self.calls.append((audio, sample_rate)) - return torch.arange(32, dtype=torch.int16).reshape(2, 16) - - -def _builder() -> tuple[BreezeTTS2PromptBuilder, _ReferenceEncoder]: - encoder = _ReferenceEncoder() - builder = BreezeTTS2PromptBuilder( - _Tokenizer(), - SimpleNamespace( - audio_token_id=262144, - audio_eos_token_id=262145, - pad_token_id=0, - backbone_config=SimpleNamespace(vocab_size=100000), - num_codebooks=16, - codec_config={"codebook_size": 2048}, - ), - encoder, - ) - return builder, encoder - - -def test_reference_voice_uses_exact_transcript_and_audio_codes(): - builder, encoder = _builder() - - prompt = builder.build( - { - "template": "ref_clone_tata", - "speaker": "S1", - "ref_text": "exact reference transcript", - "ref_audio": "reference.wav", - "ref_audio_sample_rate": 24_000, - "text": "new target text", - } - ) - - info = prompt["additional_information"] - assert tuple(info["input_values"].shape) == (2, 16) - assert info["ref_code_len"] == 2 - assert info["prompt_ids"].tolist().count(262144) == 2 - assert info["text_ids_len"].numel() == 2 - assert int(info["text_ids_mask"].sum()) == int(info["text_ids_len"].sum()) - assert encoder.calls == [("reference.wav", 24_000)] - - -def test_reference_plus_instruction_composes_voice_direction_prompt(): - builder, _ = _builder() - - prompt = builder.build( - { - "template": "ref_edit_tata", - "speaker": "S0", - "ref_text": "exact reference transcript", - "ref_audio": "reference.wav", - "instruction": "speak calmly and warmly", - "text": "new target text", - } - ) - - info = prompt["additional_information"] - assert info["template"] == "ref_edit_tata" - assert tuple(info["input_values"].shape) == (2, 16) - assert info["prompt_ids"].tolist().count(262144) == 2 - assert info["text_ids_len"].numel() == 2 - - -def test_offline_example_reference_payload_feeds_prompt_builder(tmp_path): - """Smoke: the offline example's payload keys flow into the real builder. - - Loads ``end2end.py`` (voxtral-style importlib), reads a real stereo wav - through its ``reference_payload`` helper, and feeds the result into the - real ``BreezeTTS2PromptBuilder`` — guarding the example↔builder field - contract that previously broke via an unread ``ref_audio_path`` key. - """ - spec = importlib.util.spec_from_file_location("breeze_tts_2_end2end", _EXAMPLE_PATH) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - sample_rate = 16_000 - t = np.arange(sample_rate, dtype=np.float32) / sample_rate - tone = np.sin(2 * np.pi * 220.0 * t) - ref_path = tmp_path / "reference.wav" - sf.write(str(ref_path), np.stack([tone, tone], axis=1), sample_rate) - - payload = module.reference_payload(str(ref_path)) - - # The helper emits exactly the keys the prompt builder consumes. - assert set(payload) == {"ref_audio", "ref_audio_sample_rate"} - assert payload["ref_audio_sample_rate"] == sample_rate - assert payload["ref_audio"].ndim == 1 # stereo down-mixed to one clip - - builder, encoder = _builder() - prompt = builder.build( - { - "template": "ref_clone_tata", - "speaker": "S0", - "ref_text": "exact reference transcript", - "text": "new target text", - **payload, - } - ) - - info = prompt["additional_information"] - assert tuple(info["input_values"].shape) == (2, 16) - assert info["ref_code_len"] == 2 - assert info["prompt_ids"].tolist().count(262144) == 2 - # The builder received the example's waveform and sample rate verbatim. - assert encoder.calls == [(payload["ref_audio"], sample_rate)] diff --git a/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py b/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py deleted file mode 100644 index 0edaabc6f..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py +++ /dev/null @@ -1,122 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from typing import Any - -import pytest -import torch -from vllm.v1.sample.logits_processor import LogitsProcessors -from vllm.v1.sample.metadata import SamplingMetadata - -from tests.model_executor.models.breeze_tts_2.test_talker_state_machine import _depth_codes, _talker - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _model(first_score=11.0, second_score=10.5): - model = _talker() - model._async_chunk = True - logits = torch.full((1, 9), -100.0) - logits[0, 2:4] = torch.tensor([first_score, second_score]) - model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone() - model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) - return model - - -def _step(model, infos, penalties, *, no_penalties=False): - return model.make_omni_output( - torch.zeros(len(infos), 16), - model_intermediate_buffer=infos, - request_token_spans=[(i, i + 1) for i in range(len(infos))], - sampling_metadata=SamplingMetadata( - temperature=None, - all_greedy=True, - all_random=False, - top_p=None, - top_k=None, - generators={}, - max_num_logprobs=None, - no_penalties=no_penalties, - frequency_penalties=torch.zeros(len(infos)), - presence_penalties=torch.zeros(len(infos)), - repetition_penalties=torch.tensor(penalties), - # Text/reference conditioning must never be counted as generated - # codec history, even if the scheduler ids overlap the codec ids. - prompt_token_ids=torch.tensor([[2, 2]] * len(infos)), - output_token_ids=[list(info.get("breeze_generated_code0_ids", [])) for info in infos], - allowed_token_ids_mask=None, - bad_words_token_ids={}, - logitsprocs=LogitsProcessors(), - ), - ) - - -@pytest.mark.parametrize("scores", [(11.0, 10.5), (-1.0, -1.05)]) -def test_penalty_applies_once_per_generated_token_and_respects_logit_sign(scores): - model = _model(*scores) - info = {"prompt_ids": torch.tensor([2]), "input_values": torch.tensor([[2, 3, 4, 5]])} - chosen = [] - - for _ in range(4): - output = _step(model, [info], [1.1]) - chosen.append(output.multimodal_outputs["codes"]["audio"][0][0, 0].item()) - - assert chosen == [2, 3, 2, 2] - - -def test_request_history_and_penalty_are_independent(): - model = _model() - first: dict[str, Any] = {} - second: dict[str, Any] = {} - _step(model, [first], [1.1]) - output = _step(model, [first, second], [1.1, 1.1]) - assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2] - - # Reverse the batch order; a request's history follows its state, while - # each penalty follows the current sampling-metadata row. - output = _step(model, [second, first], [1.1, 1.0]) - assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2] - - -def test_disabled_penalties_do_not_read_uninitialized_sampling_buffer(): - model = _model() - info: dict[str, Any] = {} - for _ in range(2): - output = _step(model, [info], [float("nan")], no_penalties=True) - assert output.multimodal_outputs["codes"]["audio"][0][0, 0].item() == 2 - - -def test_scheduler_logits_preserve_selected_frame_under_further_sampling_transforms(): - model = _model() - info: dict[str, Any] = {} - _step(model, [info], [1.1]) - output = _step(model, [info], [1.1]) - selected = output.multimodal_outputs["codes"]["audio"][0][0, 0].item() - logits = model.compute_logits(torch.zeros(1, 16)) - - assert selected == 3 - assert logits.argmax(dim=-1).item() == selected - assert torch.isfinite(logits).sum().item() == 1 - assert logits[0, selected].item() == 0.0 - # The scheduler may apply repetition again and rescale temperature. A - # finite zero for the selected token preserves a valid point mass. - for temperature in (0.1, 0.9, 2.0): - probabilities = torch.softmax(logits / 1.1 / temperature, dim=-1) - assert torch.isfinite(probabilities).all() - assert probabilities[0, selected].item() == 1.0 - - -def test_repetition_penalty_can_select_eos_without_emitting_a_frame(): - model = _model() - info: dict[str, Any] = {} - _step(model, [info], [1.1]) - logits = torch.full((1, 9), -100.0) - logits[0, 2] = 11.0 - logits[0, 8] = 10.5 - model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone() - - output = _step(model, [info], [1.1]) - - assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0 - assert info["breeze_generated_frames"] == 1 - assert model.compute_logits(torch.zeros(1, 16)).argmax(dim=-1).item() == 8 diff --git a/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py b/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py deleted file mode 100644 index b9795f798..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from types import SimpleNamespace - -import pytest -import torch - -from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import ( - talker2codec, - talker2codec_full_payload, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_talker2codec_flattens_codebook_major(): - codes = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int16) - output = SimpleNamespace(multimodal_output={"codes": {"audio": codes}}) - - prompts = talker2codec([output]) - - assert prompts[0]["prompt_token_ids"] == [1, 3, 5, 2, 4, 6] - assert torch.equal(prompts[0]["multi_modal_data"]["codes"]["audio"], codes.transpose(0, 1)) - - -def test_talker2codec_keeps_empty_request_terminal(): - output = SimpleNamespace(multimodal_output={"codes": {}}) - - prompts = talker2codec([output]) - - assert prompts == [{"prompt_token_ids": []}] - - -def test_talker2codec_accepts_engine_core_output_wrapper(): - codes = torch.tensor([[7, 8]], dtype=torch.int16) - output = SimpleNamespace(outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": codes}})]) - - prompts = talker2codec([output]) - - assert prompts[0]["prompt_token_ids"] == [7, 8] - - -def test_talker2codec_unwraps_per_request_audio_container(): - codes = torch.tensor([[7, 8], [9, 10]], dtype=torch.int16) - output = SimpleNamespace(multimodal_output={"codes": {"audio": [codes]}}) - - prompts = talker2codec([output]) - - assert prompts[0]["prompt_token_ids"] == [7, 9, 8, 10] - - -def test_talker2codec_accepts_flat_full_payload(): - output = {"codes.audio": torch.tensor([1, 3, 2, 4], dtype=torch.int16)} - - prompts = talker2codec([output]) - - assert prompts[0]["prompt_token_ids"] == [1, 3, 2, 4] - assert torch.equal( - prompts[0]["multi_modal_data"]["codes"]["audio"], - torch.tensor([1, 3, 2, 4], dtype=torch.int16), - ) - - -def test_talker2codec_full_payload_transposes_accumulated_frames(): - payload = talker2codec_full_payload( - transfer_manager=None, - pooling_output={"codes.audio": torch.tensor([[1, 2], [3, 4], [5, 6]])}, - request=None, - ) - - assert payload is not None - assert torch.equal( - payload.codes.audio, - torch.tensor([1, 3, 5, 2, 4, 6]), - ) - - -def test_talker2codec_full_payload_accepts_runner_pooling_output(): - payload = talker2codec_full_payload( - pooling_output={"codes.audio": torch.tensor([[9, 8]], dtype=torch.int16)}, - request=None, - ) - - assert payload is not None - assert torch.equal(payload.codes.audio, torch.tensor([9, 8])) - - -def test_full_payload_zero_frame_finish_sends_explicit_terminal_marker(): - payload = talker2codec_full_payload( - pooling_output={"codes.audio": torch.empty(0, 16, dtype=torch.long)}, - request=None, - is_finished=True, - ) - - assert payload is not None - assert payload.codes.audio.numel() == 0 - assert payload.meta.code_flat_numel == 0 - assert bool(payload.meta.finished.item()) is True - - -def test_full_payload_missing_codes_finish_sends_explicit_terminal_marker(): - payload = talker2codec_full_payload( - pooling_output={}, - request=None, - is_finished=True, - ) - - assert payload is not None - assert payload.codes.audio.numel() == 0 - assert bool(payload.meta.finished.item()) is True - - -def test_full_payload_empty_without_finish_stays_silent(): - payload = talker2codec_full_payload( - pooling_output={"codes.audio": torch.empty(0, 16, dtype=torch.long)}, - request=None, - is_finished=False, - ) - - assert payload is None - - -def test_full_payload_carries_finished_flag_on_nonempty_frames(): - payload = talker2codec_full_payload( - pooling_output={"codes.audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)}, - request=None, - is_finished=True, - ) - - assert payload is not None - assert payload.codes.audio.numel() == 4 - assert bool(payload.meta.finished.item()) is True diff --git a/tests/model_executor/models/breeze_tts_2/test_streaming.py b/tests/model_executor/models/breeze_tts_2/test_streaming.py deleted file mode 100644 index a630b90bc..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_streaming.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import copy -from collections import defaultdict -from types import SimpleNamespace - -import pytest -import torch - -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_codec import ( - BreezeTTS2MimiCodec, -) -from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import ( - talker2codec_async_chunk, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _TransferManager: - def __init__(self, chunk_frames: int = 4): - self.connector = SimpleNamespace(config={"breeze_codec_chunk_frames": chunk_frames}) - self.code_prompt_token_ids: defaultdict[str, list[torch.Tensor]] = defaultdict(list) - - -class _Request: - def __init__(self, request_id: str, finished: bool = False): - self.external_req_id = request_id - self._finished = finished - - def is_finished(self): - return self._finished - - -def test_async_processor_sends_only_unemitted_tail_and_flushes_finish(): - manager = _TransferManager(chunk_frames=4) - request = _Request("req-1") - - for code0 in range(3): - frame = torch.tensor([code0, code0 + 4, code0 + 8, code0 + 12], dtype=torch.long) - assert talker2codec_async_chunk(manager, {"codes": {"audio": frame}}, request) is None - - payload = talker2codec_async_chunk( - manager, - {"codes": {"audio": torch.tensor([3, 7, 11, 15], dtype=torch.long)}}, - request, - ) - assert payload is not None - assert payload.codes.audio.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] - assert bool(payload.meta.finished.item()) is False - assert manager.code_prompt_token_ids["req-1"] == [] - - assert talker2codec_async_chunk(manager, {"codes": {"audio": torch.arange(4)}}, request) is None - payload = talker2codec_async_chunk(manager, None, _Request("req-1", finished=True)) - assert payload.codes.audio.tolist() == [0, 1, 2, 3] - assert bool(payload.meta.finished.item()) is True - assert bool(payload.meta.stream_finished.item()) is True - assert payload.meta.codec_streaming is True - - -class _Decoder: - def __init__(self): - self.calls = [] - - def batched_chunked_decode(self, codes, lengths, caches=None, **kwargs): - self.calls.append((codes.shape, tuple(lengths), tuple(caches), kwargs)) - return [torch.ones(length * 1920, dtype=torch.float32) for length in lengths] - - -def test_stateful_codec_batches_request_chunks_and_releases_finished_state(): - decoder = _Decoder() - codec = object.__new__(BreezeTTS2MimiCodec) - codec._async_chunk = True - codec._num_codebooks = 4 - codec._codebook_size = 8 - codec._sample_rate = 24_000 - codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) - codec._decoder_state_cache = {} - - input_ids = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype=torch.long) - output = codec.forward( - input_ids, - runtime_additional_information=[ - # ``stream_finished`` is the model-visible flag: the connector - # receiver strips ``meta.finished`` while merging runtime info. - {"meta": {"request_id": "live", "stream_finished": False}}, - {"meta": {"request_id": "done", "stream_finished": True}}, - ], - seq_token_counts=[4, 4], - request_ids=["scheduler-live", "scheduler-done"], - ) - - assert decoder.calls[0][0] == (2, 4, 1) - assert [len(item) for item in output.multimodal_outputs["model_outputs"]] == [1920, 1920] - assert set(codec._decoder_state_cache) == {"scheduler-live"} - - codec.on_requests_finished(["scheduler-live"]) - assert codec._decoder_state_cache == {} - - -def test_stateful_codec_rejects_partial_frame_chunk(): - decoder = _Decoder() - codec = object.__new__(BreezeTTS2MimiCodec) - codec._async_chunk = True - codec._num_codebooks = 4 - codec._codebook_size = 8 - codec._sample_rate = 24_000 - codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) - codec._decoder_state_cache = {} - - # 5 codes cannot form a whole frame; a silent state reset would corrupt - # the rest of the stream, so the codec must fail loudly. - with pytest.raises(ValueError, match="not divisible"): - codec.forward( - torch.tensor([0, 1, 2, 3, 4], dtype=torch.long), - runtime_additional_information=[{"meta": {"request_id": "bad", "finished": False}}], - seq_token_counts=[5], - request_ids=["scheduler-bad"], - ) - assert decoder.calls == [] - - -@pytest.mark.parametrize("token_counts", [(4, 0), (0, 4), (0, 4, 0, 8, 0)]) -def test_stateful_codec_terminal_marker_pops_state_without_decoding(token_counts): - decoder = _Decoder() - codec = object.__new__(BreezeTTS2MimiCodec) - codec._async_chunk = True - codec._num_codebooks = 4 - codec._codebook_size = 8 - codec._sample_rate = 24_000 - codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) - request_ids = [f"scheduler-{index}" for index in range(len(token_counts))] - codec._decoder_state_cache = {request_id: {} for request_id in request_ids} - - output = codec.forward( - torch.arange(sum(token_counts), dtype=torch.long) % codec._codebook_size, - runtime_additional_information=[ - # Plain ``finished`` covers the fallback for direct in-process - # callers that do not go through the connector receiver. - {"meta": {"finished": count == 0}} - for count in token_counts - ], - seq_token_counts=token_counts, - request_ids=request_ids, - ) - - # Empty terminal markers release their state without consuming an output - # slot belonging to another request, regardless of their batch position. - assert len(decoder.calls) == 1 - assert decoder.calls[0][0][0] == sum(count > 0 for count in token_counts) - assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [ - count // codec._num_codebooks * 1920 for count in token_counts - ] - live_request_ids = [request_id for request_id, count in zip(request_ids, token_counts) if count] - assert list(codec._decoder_state_cache) == live_request_ids - assert all(codec._decoder_state_cache[request_id]["prefix_frames"] == 0 for request_id in live_request_ids) - - -@pytest.mark.parametrize("configured_chunk_frames", [None, 4]) -def test_stateful_codec_advances_decoder_cache_after_sliding_window(monkeypatch, tmp_path, configured_chunk_frames): - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer - from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( - Qwen3TTSTokenizerV2DecoderConfig, - ) - from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( - Qwen3TTSTokenizerV2Decoder, - ) - - torch.manual_seed(0) - config = Qwen3TTSTokenizerV2DecoderConfig( - codebook_size=32, - hidden_size=16, - latent_dim=16, - codebook_dim=16, - num_attention_heads=2, - num_key_value_heads=2, - intermediate_size=32, - num_hidden_layers=1, - num_quantizers=2, - decoder_dim=32, - upsample_rates=(2,), - upsampling_ratios=(2,), - sliding_window=72, - ) - decoder = Qwen3TTSTokenizerV2Decoder(config).eval() - # Codebooks initialize to zeros before checkpoint loading; populate them - # so different codec IDs produce distinguishable cached frame histories. - with torch.no_grad(): - for name, parameter in decoder.quantizer.named_parameters(): - if name.endswith("embedding_sum"): - parameter.normal_() - tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder, get_output_sample_rate=lambda: 24_000)) - monkeypatch.setattr(Qwen3TTSTokenizer, "from_pretrained", lambda *args, **kwargs: tokenizer) - connector_extra = {} if configured_chunk_frames is None else {"breeze_codec_chunk_frames": configured_chunk_frames} - chunk_frames = configured_chunk_frames or 8 - codec = object.__new__(BreezeTTS2MimiCodec) - codec.vllm_config = SimpleNamespace( - model_config=SimpleNamespace(stage_connector_config={"extra": connector_extra}), - device_config=SimpleNamespace(device="cpu"), - ) - codec._tokenizer_path = tmp_path / "audio_tokenizer" - codec._audio_tokenizer = None - codec._async_chunk = True - codec._num_codebooks = config.num_quantizers - codec._codebook_size = config.codebook_size - codec._decoder_state_cache = {} - codec.load_weights(iter(())) - - # At least two chunks after the decoder's rolling window expose a stale - # cache even when the first padded chunk still produces plausible audio. - total_frames = config.sliding_window + 2 * chunk_frames - tail_frames = 3 - codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, total_frames + tail_frames)) - for end in range(chunk_frames, total_frames + 1, chunk_frames): - flat = codes[0, :, end - chunk_frames : end].reshape(-1) - codec.forward( - flat, - runtime_additional_information=[{"meta": {"stream_finished": False}}], - seq_token_counts=[flat.numel()], - request_ids=["scheduler-live"], - ) - state = codec._decoder_state_cache["scheduler-live"] - assert state["suffix_frames"] == min(end - chunk_frames, config.sliding_window) + chunk_frames - with torch.inference_mode(): - expected_quantized = decoder.quantizer.decode(codes[..., max(0, end - config.sliding_window) : end]) - torch.testing.assert_close(state["suffix_quantized"], expected_quantized) - - # A short terminal chunk must emit only its real frames, then release - # state. Its waveform should agree with the unpadded single-request path. - with torch.inference_mode(): - expected_tail = decoder.chunked_decode(codes[..., -tail_frames:], caches=copy.deepcopy(state)) - flat_tail = codes[0, :, -tail_frames:].reshape(-1) - output = codec.forward( - flat_tail, - runtime_additional_information=[{"meta": {"stream_finished": True}}], - seq_token_counts=[flat_tail.numel()], - request_ids=["scheduler-live"], - ) - waveform = output.multimodal_outputs["model_outputs"][0] - assert waveform.numel() == tail_frames * decoder.total_upsample - torch.testing.assert_close(waveform, expected_tail.reshape(-1), atol=1e-5, rtol=1e-4) - assert codec._decoder_state_cache == {} diff --git a/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py b/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py deleted file mode 100644 index 2f4a13df1..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from typing import Any - -import pytest -import torch - -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( - BreezeTTS2TalkerForGeneration, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _Head: - def __init__(self) -> None: - self.code0 = 2 - - def __call__(self, hidden: torch.Tensor) -> torch.Tensor: - logits = torch.zeros((hidden.shape[0], 9), dtype=hidden.dtype) - logits[:, self.code0] = 1.0 - return logits - - -class _Embedding: - def __init__(self) -> None: - self.weight = torch.randn(4 * 8, 16) - - def __call__(self, ids: torch.Tensor) -> torch.Tensor: - return self.weight[ids.reshape(-1)].reshape(*ids.shape, -1) - - -def _talker() -> BreezeTTS2TalkerForGeneration: - model = object.__new__(BreezeTTS2TalkerForGeneration) - model.num_codebooks = 4 - model.codebook_vocab_size = 8 - model.codebook_size = 8 - model.hidden_size = 16 - model.codebook_eos_token_id = 0 - model.audio_token_id = 100 - model.audio_eos_token_id = 101 - # Bypass nn.Module registration; this fake is created with object.__new__. - model.__dict__["embed_audio_tokens"] = _Embedding() - model.audio_token_offsets = torch.arange(4) * 8 - model.lm_head = _Head() - model.logits_processor = lambda head, hidden: head(hidden) - model._codec_disallowed_mask = None - model._batch_state = None - model._batch_state_spans = None - model._async_chunk = False - model._golden_dump_dir = None - return model - - -def _depth_codes(_model: Any, _hidden: torch.Tensor, code0: int) -> torch.Tensor: - if code0 >= 8: - return torch.empty(0, dtype=torch.long) - return torch.tensor([code0, code0 + 1, code0 + 2, code0 + 3], dtype=torch.long).unsqueeze(0) - - -def test_prefill_hidden_generates_and_stores_first_complete_frame(): - model = _talker() - model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) - info = {"_omni_is_prefill": True, "breeze_generated_frames": 0, "breeze_max_new_frames": 3} - hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) - - output = model.make_omni_output( - hidden, - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - - codes = output.multimodal_outputs["codes"]["audio"][0] - assert torch.equal(codes, torch.tensor([[2, 3, 4, 5]])) - assert torch.equal(info["breeze_current_frame"], codes[0]) - assert info["breeze_generated_frames"] == 1 - assert info["breeze_force_eos"] is False - - -def test_first_decode_consumes_full_frame_instead_of_scalar_eos_fallback(): - model = _talker() - model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) - info = {"_omni_is_prefill": True, "breeze_generated_frames": 0, "breeze_max_new_frames": 3} - hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) - output = model.make_omni_output( - hidden, - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - frame = output.multimodal_outputs["codes"]["audio"][0][0] - decode_info = dict(info) - decode_info.pop("_omni_is_prefill", None) - - safe_id, decode_embeds, update = model.preprocess( - torch.tensor([2], dtype=torch.long), - None, - **decode_info, - ) - - assert safe_id.item() == 2 - assert update == {} - expected = model.embed_input_ids(frame.reshape(1, 1, -1)).reshape(1, -1) - assert torch.equal(decode_embeds, expected) - - -def test_real_decode_rejects_missing_current_frame(): - model = _talker() - info = { - "breeze_generated_frames": 1, - "breeze_audio_codes": torch.tensor([[2, 3, 4, 5]]), - } - - try: - model.preprocess(torch.tensor([2], dtype=torch.long), None, **info) - except RuntimeError as exc: - assert "breeze_current_frame" in str(exc) - assert "codebooks 1..15" in str(exc) - else: - raise AssertionError("missing breeze_current_frame must fail explicitly") - - -def test_embed_input_ids_rejects_scalar_codec_id(): - model = _talker() - - try: - model.embed_input_ids(torch.tensor([[2]])) - except ValueError as exc: - assert "complete codec frames" in str(exc) - else: - raise AssertionError("scalar codec ids must not be embedded") - - -def test_decode_accumulates_frames_then_natural_eos_stops_without_fake_frame(): - model = _talker() - model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) - info = {"breeze_generated_frames": 1, "breeze_max_new_frames": 3} - info["breeze_audio_codes"] = torch.tensor([[2, 3, 4, 5]]) - model.lm_head.code0 = 3 - hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) - - output = model.make_omni_output( - hidden, - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) - - # Main-head class 8 is the Breeze backbone EOS (one past vocab_size=8). - model.lm_head.code0 = 8 - output = model.make_omni_output( - hidden, - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) - assert info["breeze_force_eos"] is True - assert info["breeze_force_eos"] is True - - logits = model.compute_logits(hidden) - assert torch.equal(logits[0, :8], torch.full((8,), float("-inf"))) - assert logits[0, 8].item() == 0.0 - - -def test_frame_budget_forces_eos_on_followup_step(): - model = _talker() - model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) - info = {"breeze_generated_frames": 2, "breeze_max_new_frames": 2} - info["breeze_audio_codes"] = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 0]]) - hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) - - output = model.make_omni_output( - hidden, - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - - assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) - assert info["breeze_generated_frames"] == 2 - assert info["breeze_force_eos"] is True - - -def test_async_chunk_mode_skips_full_accumulation(): - model = _talker() - model._async_chunk = True - model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) - info = {"breeze_generated_frames": 0, "breeze_max_new_frames": 8} - hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) - - for _ in range(3): - output = model.make_omni_output( - hidden, - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 4) - - # Streaming keeps only the per-frame tail; the cumulative buffer (and its - # per-step torch.cat) must not exist without the opt-in golden dump. - assert "breeze_audio_codes" not in info - assert info["breeze_generated_frames"] == 3 - assert torch.equal(info["breeze_current_frame"], torch.tensor([2, 3, 4, 5])) - - -def test_async_chunk_mode_keeps_accumulation_for_golden_dump(tmp_path): - model = _talker() - model._async_chunk = True - model._golden_dump_dir = str(tmp_path) - model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) - info = {"breeze_generated_frames": 0, "breeze_max_new_frames": 8} - hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) - - for _ in range(3): - output = model.make_omni_output( - hidden, - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 4) - - assert info["breeze_audio_codes"].shape == (3, 4) - assert torch.equal(info["breeze_audio_codes"][-1], torch.tensor([2, 3, 4, 5])) - - -def test_opt_in_golden_dump_writes_terminal_frames(tmp_path): - model = _talker() - model._golden_dump_dir = str(tmp_path) - info = { - "template": "tts_plain", - "prompt_ids": torch.tensor([1, 2, 3]), - "breeze_audio_codes": torch.tensor([[1, 2, 3, 4]]), - "breeze_generated_frames": 1, - "breeze_force_eos": True, - } - - model._dump_golden_frames(info) - - dump_path = next(tmp_path.glob("breeze_*.pt")) - payload = torch.load(dump_path, weights_only=False) - assert payload["template"] == "tts_plain" - assert torch.equal(payload["prompt_ids"], torch.tensor([1, 2, 3])) - assert torch.equal(payload["codes"], torch.tensor([[1, 2, 3, 4]])) - assert payload["finished"] is True diff --git a/tests/model_executor/models/breeze_tts_2/test_weight_loading.py b/tests/model_executor/models/breeze_tts_2/test_weight_loading.py deleted file mode 100644 index f4c0da429..000000000 --- a/tests/model_executor/models/breeze_tts_2/test_weight_loading.py +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from types import SimpleNamespace - -import pytest -import torch -from torch import nn - -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( - BreezeTTS2TalkerForGeneration, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _PaddedHead(nn.Module): - def __init__(self) -> None: - super().__init__() - self.weight = nn.Parameter(torch.zeros(2_112, 32), requires_grad=False) - - -class _DepthDecoder: - def __init__(self) -> None: - self.loaded: list[tuple[str, torch.Tensor]] = [] - - def load_weights(self, weights): - self.loaded.extend(weights) - return {"model.embed_tokens.weight"} - - -class _TalkerHost(nn.Module): - def __init__(self) -> None: - super().__init__() - self.config = SimpleNamespace(tie_codebooks_embeddings=True) - self.codebook_vocab_size = 2_051 - self.hidden_size = 32 - self.lm_head = _PaddedHead() - self.embed_audio_tokens = nn.Embedding(4 * 2_051, 32) - self.depth_decoder = _DepthDecoder() - - -def _padded_loader(parameter: nn.Parameter, weight: torch.Tensor) -> None: - assert parameter.shape != weight.shape - parameter[: weight.shape[0]].copy_(weight) - - -def test_lm_head_loader_accepts_vocab_padding(): - host = _TalkerHost() - host.lm_head.weight.weight_loader = _padded_loader - checkpoint_weight = torch.full((2_052, 32), 1.5) - - loaded = BreezeTTS2TalkerForGeneration.load_weights( - host, - iter([("lm_head.weight", checkpoint_weight)]), - ) - - assert "lm_head.weight" in loaded - assert torch.equal(host.lm_head.weight[:2_052], checkpoint_weight) - assert torch.count_nonzero(host.lm_head.weight[2_052:]).item() == 0 - - -def test_tied_audio_embedding_reports_both_strict_loading_aliases(): - host = _TalkerHost() - checkpoint_weight = torch.full((4 * 2_051, 32), 2.0) - - loaded = BreezeTTS2TalkerForGeneration.load_weights( - host, - iter([("depth_decoder.model.embed_tokens.weight", checkpoint_weight)]), - ) - - assert "embed_audio_tokens.weight" in loaded - assert "depth_decoder.model.embed_tokens.weight" in loaded diff --git a/tests/model_executor/models/breeze_tts_2_helpers.py b/tests/model_executor/models/breeze_tts_2_helpers.py new file mode 100644 index 000000000..e584c003d --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2_helpers.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Weight-free Breeze talker and request helpers shared by model tests.""" + +from dataclasses import dataclass +from unittest.mock import Mock + +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.first_code_sampler import BreezeFirstCodeSampler +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration + + +@dataclass +class TalkerConfig: + vocab_size: int + eos_token_id: int + + +@dataclass +class DepthStub: + generate_frames: Mock + + +def _small_talker(): + model = BreezeForConditionalGeneration.__new__(BreezeForConditionalGeneration) + torch.nn.Module.__init__(model) + model._first_code_sampler = BreezeFirstCodeSampler() + model.config = TalkerConfig(vocab_size=8, eos_token_id=7) + model.num_codebooks, model.codebook_size, model.hidden_size = 3, 4, 2 + model.lm_head = torch.nn.Linear(2, 8, bias=False) + with torch.no_grad(): + model.lm_head.weight.zero_() + model.lm_head.weight[1] = torch.tensor([10.0, 0.0]) + model.lm_head.weight[2] = torch.tensor([9.5, 0.0]) + model.lm_head.weight[3] = torch.tensor([0.0, 10.0]) + model.lm_head.weight[7] = torch.tensor([-10.0, -10.0]) + model.depth_decoder = DepthStub( + generate_frames=Mock(side_effect=lambda hidden, first, **kwargs: first[:, None].repeat(1, 3)) + ) + return model + + +def _request_info(request_id): + return { + "global_request_id": [request_id], + "breeze_prompt": {"role": "cond", "guidance_scale": 1.0}, + "breeze_sampling": {"temperature": 0.0, "top_k": 0, "top_p": 1.0, "repetition_penalty": 1.1}, + "breeze_state": { + "generator": torch.Generator().manual_seed(42), + "history": torch.empty(1, 0, dtype=torch.long), + "current": torch.zeros(1, 3, dtype=torch.long), + }, + } diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index 8ad2b519f..581eb977e 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -23,7 +23,7 @@ class _FakeDecoder(nn.Module): def __init__(self, total_upsample: int = _TOTAL_UPSAMPLE): super().__init__() self.total_upsample = total_upsample - self.decode_calls: list[dict[str, int | tuple[int, ...]]] = [] + self.decode_calls: list[dict[str, object]] = [] self.batched_decode_calls: list[dict[str, object]] = [] self.decode_codes: list[torch.Tensor] = [] self.cudagraph_calls: list[dict[str, int | torch.device]] = [] @@ -91,8 +91,8 @@ class _FakeDecoder(nn.Module): wav_len = frames * self.total_upsample + 6 wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1) offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000 - batched_wav = wav.expand(batch, 1, wav_len) + offsets - return [batched_wav[row, :, : lengths[row] * self.total_upsample] for row in range(batch)] + batched_outputs = wav.expand(batch, 1, wav_len) + offsets + return [batched_outputs[row, :, : lengths[row] * self.total_upsample] for row in range(batch)] def enable_cudagraph(self, **kwargs): self.cudagraph_calls.append(kwargs) @@ -577,6 +577,7 @@ def test_decode_chunking_override_is_passed_to_cudagraph(): _load_weights_noop(model) assert model.decoder.cudagraph_calls[-1] == { + "capture_modes": ("icl", "xvec"), "capture_batch_sizes": None, "stateless_capture_sizes": None, "device": torch.device("cuda"), @@ -607,6 +608,13 @@ def test_cudagraph_batch_config_is_preserved(): assert call["capture_batch_sizes"] == [1, 2, 4, 8] +def test_cudagraph_captures_only_modes_used_by_the_model(): + model = _make_model(async_chunk=True, device=torch.device("cuda")) + model.decoder_cudagraph_modes = ("xvec",) + _load_weights_noop(model) + assert model.decoder.cudagraph_calls[-1]["capture_modes"] == ("xvec",) + + def test_stateless_cudagraph_capture_sizes_are_passed_from_config(): model = _make_model( async_chunk=False, diff --git a/tests/model_executor/models/test_breeze_tts_2.py b/tests/model_executor/models/test_breeze_tts_2.py new file mode 100644 index 000000000..dde3e0e65 --- /dev/null +++ b/tests/model_executor/models/test_breeze_tts_2.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Check depth attention against HF Llama and request-local sampling state.""" + +from collections import defaultdict + +import numpy as np +import pytest +import torch +import torch.nn.functional as F +from transformers import LlamaConfig, MimiConfig +from transformers.models.llama.modeling_llama import LlamaModel +from transformers.models.mimi.modeling_mimi import MimiConv1d + +from tests.model_executor.models.breeze_tts_2_helpers import _request_info, _small_talker +from vllm_omni.engine.serialization import deserialize_additional_information, serialize_additional_information +from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import ( + BreezeDepthDecoder, + sample_graph_logits, + sample_logits, +) +from vllm_omni.model_executor.models.breeze_tts_2.prompt import CFG_UNCOND_SUFFIX, build_breeze_prompt +from vllm_omni.model_executor.models.breeze_tts_2.reference_encoder import BreezeReferenceConv +from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import expand_cfg_prompts, talker2code2wav_async_chunk + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("parameter", ["temperature", "top_p", "repetition_penalty"]) +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_offline_prompt_rejects_nonfinite_sampling(mocker, parameter, value): + with pytest.raises(ValueError, match="sampling"): + build_breeze_prompt(mocker.Mock(), "Hello.", **{parameter: value}) + + +@pytest.mark.parametrize("stride", [2, 3]) +def test_reference_convolution_batch_preserves_partial_final_frames(stride): + torch.manual_seed(7) + config = MimiConfig(use_causal_conv=True, pad_mode="constant") + first = MimiConv1d(config, 1, 4, kernel_size=2 * stride, stride=stride) + last = MimiConv1d(config, 4, 4, kernel_size=4, stride=2, pad_mode="replicate", bias=False) + batched_first = BreezeReferenceConv(first, stride) + batched_last = BreezeReferenceConv(last, 2 * stride) + lengths = torch.tensor([37, 62]) + batched_first.sample_lengths = batched_last.sample_lengths = lengths + waves = [torch.randn(1, 1, int(length)) for length in lengths] + padded = torch.cat([F.pad(wave, (0, 96 - wave.shape[-1])) for wave in waves]) + actual = batched_last(batched_first(padded)) + for row, waveform in enumerate(waves): + expected = last(first(waveform)) + torch.testing.assert_close(actual[row : row + 1, :, : expected.shape[-1]], expected) + assert not actual[row, :, expected.shape[-1] :].count_nonzero() + + +def test_reference_conditioning_survives_transport_and_cfg_keeps_true_prompt_lengths(mocker) -> None: + tokenizer = mocker.Mock() + tokenizer.encode.side_effect = lambda text, **kwargs: list(text.encode("utf-8")) + waveform = np.zeros(1921, dtype=np.float32) + prompt = build_breeze_prompt( + tokenizer, "Hello.", "Speak happily.", ref_audio=(waveform, 24000), ref_text="Reference.", guidance_scale=4 + ) + companion = expand_cfg_prompts(prompt, None)[0].prompt + assert len(prompt["prompt_token_ids"]) > len(companion["prompt_token_ids"]) + positive = deserialize_additional_information(serialize_additional_information(prompt["additional_information"])) + negative = deserialize_additional_information(serialize_additional_information(companion["additional_information"])) + assert positive["breeze_prompt"]["target_ids"] == list(b"[S0]Speak happily.Hello.") + assert negative["breeze_prompt"]["target_ids"] == list(b"[S0]Hello.") + assert positive["breeze_prompt"]["reference_ids"] == negative["breeze_prompt"]["reference_ids"] + assert positive["breeze_prompt"]["reference_frames"] == 2 + assert positive["cfg_group"]["role"] == "cond" + assert negative["cfg_group"]["role"] == "uncond" + torch.testing.assert_close(positive["reference_waveform"], torch.from_numpy(waveform), atol=0, rtol=0) + torch.testing.assert_close(negative["reference_waveform"], positive["reference_waveform"], atol=0, rtol=0) + + +@pytest.fixture +def depth() -> BreezeDepthDecoder: + torch.manual_seed(42) + config = LlamaConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + vocab_size=35, + max_position_embeddings=16, + rope_theta=500000.0, + rope_scaling={ + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 0.001953125, + "high_freq_factor": 0.0078125, + "original_max_position_embeddings": 16, + }, + ) + config.num_codebooks = 4 + config.audio_embed_size = 48 + config._attn_implementation = "eager" + model = BreezeDepthDecoder(config).eval() + torch.nn.init.normal_(model.codebooks_head, std=0.1) + return model + + +@pytest.fixture +def hf_depth(depth: BreezeDepthDecoder) -> LlamaModel: + reference = LlamaModel(depth.config).eval() + weights = {} + for name, parameter in depth.named_parameters(): + if name.startswith("layers."): + if name.endswith("qkv.weight"): + layer = depth.layers[int(name.split(".")[1])] + for part, tensor in zip( + ("q", "k", "v"), parameter.split([layer.q_size, layer.kv_size, layer.kv_size]), strict=True + ): + weights[name.replace("qkv.weight", f"self_attn.{part}_proj.weight")] = tensor + elif name.endswith("gate_up.weight"): + for part, tensor in zip(("gate", "up"), parameter.chunk(2), strict=True): + weights[name.replace("gate_up.weight", f"mlp.{part}_proj.weight")] = tensor + else: + weights[name.replace(".o_proj.", ".self_attn.o_proj.").replace(".down_proj.", ".mlp.down_proj.")] = ( + parameter + ) + elif name == "norm.weight": + weights[name] = parameter + result = reference.load_state_dict(weights, strict=False) + assert result.missing_keys == ["embed_tokens.weight"] + assert result.unexpected_keys == [] + return reference + + +@pytest.mark.parametrize("batch", [1, 2]) +@torch.inference_mode() +def test_depth_cached_attention_matches_hf_full_sequence( + depth: BreezeDepthDecoder, hf_depth: LlamaModel, batch: int +) -> None: + """Wrong cache positions or fused-QKV ordering change the reference hidden states.""" + embeddings = torch.randn(batch, 4, depth.config.hidden_size) + expected = hf_depth(inputs_embeds=embeddings, use_cache=False).last_hidden_state + shape = (batch, 2, 4, 8) + caches = [(torch.zeros(shape), torch.zeros(shape)) for _ in depth.layers] + actual = [] + for start, end in [(0, 2), (2, 3), (3, 4)]: + x = embeddings[:, start:end] + for layer, cache in zip(depth.layers, caches, strict=True): + x = layer( + x, + depth.positions[start:end], + cache, + depth.causal_mask[start:end, :4][None, None], + depth.rope_cos[start:end][None, None], + depth.rope_sin[start:end][None, None], + ) + actual.append(depth.norm(x)) + torch.testing.assert_close(torch.cat(actual, 1), expected, atol=2e-6, rtol=2e-5) + + +@pytest.mark.parametrize("guidance_scale", [0.5, 4.0]) +@torch.inference_mode() +def test_cfg_depth_matches_hf_full_prefix_at_every_codebook( + depth: BreezeDepthDecoder, hf_depth: LlamaModel, guidance_scale: float +) -> None: + hidden = torch.randn(4, depth.config.audio_embed_size) + first = torch.tensor([7, 11]) + prefix = torch.cat((hidden[:, None], depth.embed_tokens(first.repeat(2)[:, None])), dim=1) + expected = [first] + for codebook in range(depth.num_codebooks - 1): + states = hf_depth(inputs_embeds=depth.inputs_embeds_projector(prefix), use_cache=False).last_hidden_state + logits = F.linear(states[:, -1].float(), depth.codebooks_head[codebook].T) + cond, uncond = logits.chunk(2) + token = (uncond + guidance_scale * (cond - uncond))[:, : depth.vocab_size - 3].argmax(-1) + expected.append(token) + embed = depth.embed_tokens((token.repeat(2) + (codebook + 1) * depth.vocab_size)[:, None]) + prefix = torch.cat((prefix, embed), dim=1) + actual = depth._generate_frame( + hidden, first, depth._allocate_cache(hidden), 0.0, 0, 1.0, torch.Generator(), guidance_scale=guidance_scale + ) + torch.testing.assert_close(actual, torch.stack(expected, dim=-1), atol=0, rtol=0) + + +@torch.inference_mode() +def test_frame_generation_does_not_reuse_another_requests_cache(depth: BreezeDepthDecoder) -> None: + hidden = torch.randn(1, 48) + first = torch.tensor([7]) + + def generate(value: torch.Tensor) -> torch.Tensor: + return depth.generate_frame( + value, first, temperature=0.9, top_k=10, top_p=0.8, generator=torch.Generator().manual_seed(99) + ) + + expected = generate(hidden) + generate(hidden * -3) + torch.testing.assert_close(generate(hidden), expected, atol=0, rtol=0) + + +def test_top_k_keeps_tied_candidates_and_nucleus_keeps_crossing_token() -> None: + generator = torch.Generator().manual_seed(42) + logits = torch.tensor([[0.0, 0.0, -100.0]]).expand(128, -1) + sampled = sample_logits(logits, 1.0, 1, 1.0, generator) + assert set(sampled.tolist()) == {0, 1} + sampled = sample_logits(torch.tensor([[10.0, 0.0]]), 1.0, 0, 0.1, generator) + assert sampled.item() == 0 + + +@pytest.mark.parametrize("parameters", [(0.0, 0, 1.0), (0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95)]) +def test_graph_sampler_matches_scalar_filtering_with_identical_noise(parameters) -> None: + generator = torch.Generator().manual_seed(42) + logits = torch.randn(32, 2051, generator=generator) + logits[:, -3:] = -torch.inf + logits[:, 10:14] = 4.0 + noise = torch.empty_like(logits).exponential_(generator=generator) + expected = sample_logits(logits, *parameters, generator, noise) + actual = sample_graph_logits(logits, torch.tensor(parameters), noise) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("top_p", [0.25, 0.5, 0.75]) +@pytest.mark.parametrize("top_k", [0, 2]) +def test_graph_sampler_excludes_token_after_exact_nucleus_boundary(top_p, top_k) -> None: + # Equal probabilities make the CDF boundary exact in float32. Bias the + # shared noise toward the first token outside that nucleus so retaining + # one extra candidate deterministically changes the sampled token. + logits = torch.tensor([[0.0, 0.0, 0.0, 0.0, -torch.inf, -torch.inf, -torch.inf]]) + noise = torch.ones_like(logits) + first_excluded = int(top_p * 4) + noise[:, first_excluded] = 0.001 + parameters = (1.0, top_k, top_p) + expected = sample_logits(logits, *parameters, torch.Generator().manual_seed(42), noise) + actual = sample_graph_logits(logits, torch.tensor(parameters), noise) + + assert expected.item() != first_excluded + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_codec_chunks_preserve_zero_frames_and_flush_each_request_once(mocker) -> None: + manager = mocker.Mock() + manager.connector.config = {"extra": {"codec_chunk_frames": 5, "initial_codec_chunk_frames": 5}} + manager.code_prompt_token_ids = defaultdict(list) + manager.request_payload = {} + first = mocker.Mock(external_req_id="first") + first.is_finished.return_value = False + second = mocker.Mock(external_req_id="second") + second.is_finished.return_value = False + zeros = {"codes": {"audio": torch.zeros(3, 16, dtype=torch.long)}} + assert talker2code2wav_async_chunk(manager, zeros, first) is None + assert talker2code2wav_async_chunk(manager, zeros, second) is None + codes = torch.arange(64).reshape(4, 16) + output = talker2code2wav_async_chunk(manager, {"codes": {"audio": codes}}, first) + expected = torch.cat((torch.zeros(3, 16, dtype=torch.long), codes)) + torch.testing.assert_close(output.codes.audio.reshape(16, -1).T, expected) + final = talker2code2wav_async_chunk(manager, None, first, is_finished=True) + assert final.codes.audio.numel() == 0 + assert final.meta.finished.item() + other = talker2code2wav_async_chunk(manager, None, second, is_finished=True) + torch.testing.assert_close(other.codes.audio, torch.zeros(48, dtype=torch.long)) + + +def test_codec_chunk_ramp_keeps_frame_order_and_flushes_partial_tail(mocker) -> None: + manager = mocker.Mock() + manager.connector.config = {"extra": {"codec_chunk_frames": 5, "codec_chunk_ramp": [1, 2, 4, 5]}} + manager.code_prompt_token_ids = defaultdict(list) + manager.request_payload = {} + request = mocker.Mock(external_req_id="ramp") + request.is_finished.return_value = False + outputs = [] + for index in range(15): + result = talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.full((1, 16), index)}}, request) + if result is not None: + outputs.append(result.codes.audio.reshape(16, -1).T) + tail = talker2code2wav_async_chunk(manager, None, request, is_finished=True) + outputs.append(tail.codes.audio.reshape(16, -1).T) + assert [len(part) for part in outputs] == [1, 2, 4, 5, 3] + torch.testing.assert_close(torch.cat(outputs), torch.arange(15)[:, None].expand(-1, 16)) + assert tail.meta.finished.item() + + +@torch.inference_mode() +def test_talker_history_and_generators_follow_requests_across_reordering_and_new_admission(): + model = _small_talker() + first, second = _request_info("first"), _request_info("second") + # Unequal spans must select the last hidden row of each request. + model.make_omni_output( + torch.tensor([[0.0, 1.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1.0]]), + model_intermediate_buffer=[first, second], + request_token_spans=[(0, 3), (3, 4)], + ) + reordered = model.make_omni_output( + torch.tensor([[0.0, 1.0], [1.0, 0.0]]), + model_intermediate_buffer=[second, first], + request_token_spans=[(0, 1), (1, 2)], + ) + assert [frame[0, 0].item() for frame in reordered.multimodal_outputs["codes"]["audio"]] == [3, 2] + assert first["breeze_state"]["history"].tolist() == [[1, 2]] + assert second["breeze_state"]["history"].tolist() == [[3, 3]] + generators = model.depth_decoder.generate_frames.call_args.kwargs["generators"] + assert generators[0] is second["breeze_state"]["generator"] + assert generators[1] is first["breeze_state"]["generator"] + + newcomer = _request_info("newcomer") + admitted = model.make_omni_output( + torch.tensor([[1.0, 0.0], [1.0, 0.0]]), + model_intermediate_buffer=[newcomer, first], + request_token_spans=[(0, 1), (1, 2)], + ) + assert [frame[0, 0].item() for frame in admitted.multimodal_outputs["codes"]["audio"]] == [1, 1] + assert newcomer["breeze_state"]["history"].tolist() == [[1]] + assert first["breeze_state"]["history"].tolist() == [[1, 2, 1]] + assert second["breeze_state"]["history"].tolist() == [[3, 3]] + + +@torch.inference_mode() +def test_talker_eos_does_not_shift_live_state_or_poison_the_next_batch(): + model = _small_talker() + finished, live = _request_info("finished"), _request_info("live") + output = model.make_omni_output( + torch.tensor([[-1.0, -1.0], [0.0, 1.0]]), + model_intermediate_buffer=[finished, live], + request_token_spans=[(0, 1), (1, 2)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert finished["breeze_state"]["history"].numel() == 0 + assert live["breeze_state"]["current"].tolist() == [[3, 3, 3]] + assert model.compute_logits(torch.zeros(2, 2)).argmax(-1).tolist() == [7, 0] + + model.make_omni_output(torch.tensor([[0.0, 1.0]]), model_intermediate_buffer=[live], request_token_spans=[(0, 1)]) + assert live["breeze_state"]["history"].tolist() == [[3, 3]] + assert model.compute_logits(torch.zeros(1, 2)).argmax(-1).tolist() == [0] + + +@pytest.mark.parametrize("temperature", [0.0, 0.9]) +@torch.inference_mode() +def test_talker_mixed_eos_and_cfg_matches_independent_request_state(temperature): + """Batch EOS decisions must not mix CFG state or change a request's RNG stream.""" + finished_companion = "finished_cfg" + CFG_UNCOND_SUFFIX + live_companion = "live_cfg" + CFG_UNCOND_SUFFIX + order = ["finished_cfg", "live_plain", finished_companion, "finished_plain", live_companion, "live_cfg"] + hidden_by_id = { + "finished_cfg": [-1.0, -1.0], + finished_companion: [-1.0, -1.0], + "finished_plain": [-1.0, -1.0], + "live_plain": [0.0, 1.0], + "live_cfg": [1.0, 0.0], + live_companion: [1.0, 0.0], + } + + def make_infos(): + infos = {request_id: _request_info(request_id) for request_id in order} + for index, (request_id, info) in enumerate(infos.items()): + info["breeze_sampling"].update(temperature=temperature, top_k=1) + if "cfg" not in request_id: + # Interleave distinct sampler groups while each request still + # matches its independently decoded state and RNG trajectory. + info["breeze_sampling"].update(temperature=temperature * 0.5, top_p=0.8) + info["breeze_state"]["generator"].manual_seed(42 + index) + info["breeze_state"]["history"] = torch.tensor([[3, 3]]) + if "cfg" in request_id: + info["breeze_prompt"]["guidance_scale"] = 4.0 + if request_id.endswith(CFG_UNCOND_SUFFIX): + info["breeze_prompt"]["role"] = "uncond" + return infos + + def sample_depth(hidden, first, *, temperature, generators, **kwargs): + # Consume request-local RNG state in the downstream phase as well as + # codebook-0 sampling, without requiring depth CUDA graphs on CPU. + tail = ( + torch.stack([torch.randint(0, 4, (2,), generator=generator) for generator in generators]) + if temperature > 0 + else first[:, None].expand(-1, 2) + ) + return torch.cat((first[:, None], tail), dim=1) + + batched, independent = _small_talker(), _small_talker() + batched.depth_decoder.generate_frames.side_effect = sample_depth + independent.depth_decoder.generate_frames.side_effect = sample_depth + batch_infos, independent_infos = make_infos(), make_infos() + initial_rng = {rid: info["breeze_state"]["generator"].get_state() for rid, info in batch_infos.items()} + saved_frames = [] + # The second step drops both EOS requests and moves the live CFG companion + # ahead of its parent, exercising shrinking and reordered batches. + for step_ids in (order, [live_companion, "live_cfg", "live_plain"]): + expected_audio, expected_tokens = {}, {} + for request_id in step_ids: + if request_id.endswith(CFG_UNCOND_SUFFIX): + continue + pair = [request_id] + if independent_infos[request_id]["breeze_prompt"]["guidance_scale"] != 1.0: + pair.append(request_id + CFG_UNCOND_SUFFIX) + result = independent.make_omni_output( + torch.tensor([hidden_by_id[rid] for rid in pair]), + model_intermediate_buffer=[independent_infos[rid] for rid in pair], + request_token_spans=[(i, i + 1) for i in range(len(pair))], + ) + tokens = independent.compute_logits(torch.zeros(len(pair), 2)).argmax(-1) + for index, rid in enumerate(pair): + expected_audio[rid] = result.multimodal_outputs["codes"]["audio"][index].clone() + expected_tokens[rid] = tokens[index] + + result = batched.make_omni_output( + torch.tensor([hidden_by_id[rid] for rid in step_ids]), + model_intermediate_buffer=[batch_infos[rid] for rid in step_ids], + request_token_spans=[(i, i + 1) for i in range(len(step_ids))], + ) + actual_tokens = batched.compute_logits(torch.zeros(len(step_ids), 2)).argmax(-1) + for index, rid in enumerate(step_ids): + actual = result.multimodal_outputs["codes"]["audio"][index] + torch.testing.assert_close(actual, expected_audio[rid], atol=0, rtol=0) + torch.testing.assert_close(actual_tokens[index], expected_tokens[rid], atol=0, rtol=0) + saved_frames.append((actual, actual.clone())) + for rid in order: + actual, expected = batch_infos[rid]["breeze_state"], independent_infos[rid]["breeze_state"] + for key in ("history", "current"): + torch.testing.assert_close(actual[key], expected[key], atol=0, rtol=0) + torch.testing.assert_close(actual["generator"].get_state(), expected["generator"].get_state()) + + # Unconditional companions never own sampling; an EOS request consumes + # only codebook-0 sampling, with no downstream depth draws. + for rid in (finished_companion, live_companion): + torch.testing.assert_close(batch_infos[rid]["breeze_state"]["generator"].get_state(), initial_rng[rid]) + for rid in ("finished_cfg", "finished_plain"): + reference_generator = torch.Generator().set_state(initial_rng[rid]) + sampling = batch_infos[rid]["breeze_sampling"] + sample_logits( + torch.zeros(1, batched.config.vocab_size), + sampling["temperature"], + sampling["top_k"], + sampling["top_p"], + reference_generator, + ) + torch.testing.assert_close( + batch_infos[rid]["breeze_state"]["generator"].get_state(), reference_generator.get_state() + ) + + for rid in ("finished_cfg", finished_companion, "finished_plain", live_companion): + assert batch_infos[rid]["breeze_state"]["history"].tolist() == [[3, 3]] + assert batch_infos["live_cfg"]["breeze_state"]["history"].tolist() == [[3, 3, 1, 2]] + assert batch_infos["live_plain"]["breeze_state"]["history"].tolist() == [[3, 3, 3, 3]] + assert batch_infos["live_cfg"]["breeze_state"]["current"] is batch_infos[live_companion]["breeze_state"]["current"] + for original, snapshot in saved_frames: + torch.testing.assert_close(original, snapshot, atol=0, rtol=0) + + +@torch.inference_mode() +def test_talker_all_eos_preserves_history_and_skips_depth(): + model = _small_talker() + infos = [_request_info("first"), _request_info("second")] + for info in infos: + info["breeze_state"]["history"] = torch.tensor([[1, 2]]) + output = model.make_omni_output( + torch.tensor([[-1.0, -1.0], [-2.0, -2.0]]), + model_intermediate_buffer=infos, + request_token_spans=[(0, 1), (1, 2)], + ) + + model.depth_decoder.generate_frames.assert_not_called() + assert model.compute_logits(torch.zeros(2, 2)).argmax(-1).tolist() == [7, 7] + assert all(frame.numel() == 0 for frame in output.multimodal_outputs["codes"]["audio"]) + for info in infos: + assert info["breeze_state"]["history"].tolist() == [[1, 2]] + assert info["breeze_state"]["current"].tolist() == [[0, 0, 0]] diff --git a/tests/model_executor/models/test_breeze_tts_2_first_code_integration.py b/tests/model_executor/models/test_breeze_tts_2_first_code_integration.py new file mode 100644 index 000000000..ce138303d --- /dev/null +++ b/tests/model_executor/models/test_breeze_tts_2_first_code_integration.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Check only the caller's grouping seam, with the existing weight-free fixture.""" + +import pytest +import torch +import torch.nn.functional as F + +from tests.model_executor.models.breeze_tts_2_helpers import _request_info, _small_talker +from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits +from vllm_omni.model_executor.models.breeze_tts_2.prompt import CFG_UNCOND_SUFFIX + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@torch.inference_mode() +def test_first_code_groups_preserve_logits_request_order_cfg_and_eos(monkeypatch): + """Grouping changes dispatch only; head/history arithmetic and routing stay exact.""" + model = _small_talker() + cfg_a, cfg_b = "cfg_a" + CFG_UNCOND_SUFFIX, "cfg_b" + CFG_UNCOND_SUFFIX + order = ["plain_a", "plain_b", cfg_a, "finished_a", "cfg_a", "cfg_b", cfg_b, "plain_c"] + tuples = {"a": (0.9, 3, 1.0), "b": (0.7, 2, 0.8), "c": (0.0, 0, 1.0)} + chosen = {"plain_a": 1, "plain_b": 2, "finished_a": 7, "cfg_a": 3, "cfg_b": 0, "plain_c": 2} + hidden = { + "plain_a": [1.0, -0.5], + "plain_b": [-0.5, 1.0], + cfg_a: [0.25, 0.5], + "finished_a": [-1.0, -1.0], + "cfg_a": [0.75, -0.5], + "cfg_b": [0.5, 0.75], + cfg_b: [-0.25, 0.5], + "plain_c": [0.5, -0.75], + } + infos = {rid: _request_info(rid) for rid in order} + for index, (rid, info) in enumerate(infos.items()): + parent = rid.removesuffix(CFG_UNCOND_SUFFIX) + parameters = tuples[parent[-1]] + info["breeze_sampling"].update( + temperature=parameters[0], + top_k=parameters[1], + top_p=parameters[2], + repetition_penalty=2.0 if parent.startswith("cfg") else 1.25, + ) + info["breeze_state"]["generator"].manual_seed(100 + index) + info["breeze_state"]["history"] = torch.tensor([[1, 3, 1]]) + if parent.startswith("cfg"): + info["breeze_prompt"]["guidance_scale"] = 4.0 if parent == "cfg_a" else 2.0 + if rid.endswith(CFG_UNCOND_SUFFIX): + info["breeze_prompt"]["role"] = "uncond" + original_states = { + rid: {key: value.clone() for key, value in info["breeze_state"].items() if key in ("history", "current")} + for rid, info in infos.items() + } + request_for_generator = {id(info["breeze_state"]["generator"]): rid for rid, info in infos.items()} + + class RecordingSampler: + def __init__(self): + self.calls = [] + + def sample(self, logits, parameters, generators): + ids = [request_for_generator[id(generator)] for generator in generators] + self.calls.append((parameters, ids, [row.clone() for row in logits])) + return [torch.tensor([chosen[rid]], dtype=torch.long) for rid in ids] + + recorder = RecordingSampler() + model._first_code_sampler = recorder + # Build the scalar pre-grouping reference before spying on actual head + # calls. Both positive and negative repeated logits exercise the penalty. + expected_logits = {} + for rid in chosen: + row = torch.tensor([hidden[rid]], dtype=torch.bfloat16).float() + scores = F.linear(row, model.lm_head.weight) + scale = infos[rid]["breeze_prompt"]["guidance_scale"] + if scale != 1.0: + uncond = F.linear( + torch.tensor([hidden[rid + CFG_UNCOND_SUFFIX]], dtype=torch.bfloat16).float(), model.lm_head.weight + ) + scores = uncond + scale * (scores - uncond) + scores[:, model.codebook_size : model.config.eos_token_id] = -torch.inf + history = infos[rid]["breeze_state"]["history"] + selected = scores.gather(1, history) + penalty = infos[rid]["breeze_sampling"]["repetition_penalty"] + scores.scatter_(1, history, torch.where(selected < 0, selected * penalty, selected / penalty)) + expected_logits[rid] = scores + + actual_linear, head_inputs = F.linear, [] + + def scalar_head(inputs, weight, bias=None): + assert inputs.dtype == torch.float32 and inputs.shape == (1, model.hidden_size) + assert weight is model.lm_head.weight + head_inputs.append(inputs.clone()) + return actual_linear(inputs, weight, bias) + + # Prefill-like unequal spans must still select each physical row's last + # hidden state, including a CFG companion preceding its parent. + flat_hidden: list[list[float]] = [] + spans: list[tuple[int, int]] = [] + for index, rid in enumerate(order): + start = len(flat_hidden) + flat_hidden.extend([[-99.0, -99.0]] * (index % 3)) + flat_hidden.append(hidden[rid]) + spans.append((start, len(flat_hidden))) + values = torch.tensor(flat_hidden, dtype=torch.bfloat16) + with monkeypatch.context() as head_guard: + head_guard.setattr(F, "linear", scalar_head) + output = model.make_omni_output( + values, model_intermediate_buffer=[infos[rid] for rid in order], request_token_spans=spans + ) + assert output.text_hidden_states is values + assert len(head_inputs) == len(chosen) + 2 + assert [(parameters, ids) for parameters, ids, _ in recorder.calls] == [ + (tuples["a"], ["plain_a", "finished_a", "cfg_a"]), + (tuples["b"], ["plain_b", "cfg_b"]), + (tuples["c"], ["plain_c"]), + ] + for _, ids, rows in recorder.calls: + for rid, row in zip(ids, rows, strict=True): + assert row.dtype == torch.float32 and row.shape == (1, model.config.vocab_size) + torch.testing.assert_close(row, expected_logits[rid], atol=0, rtol=0) + + # Depth group creation must follow original conditioned order, not the + # first-code sampler's grouped traversal (A,A,A,B,B,C). + live_order = ["plain_a", "plain_b", "cfg_a", "cfg_b", "plain_c"] + calls = model.depth_decoder.generate_frames.call_args_list + assert len(calls) == len(live_order) + for call, rid in zip(calls, live_order, strict=True): + depth_hidden, first = call.args + scale = infos[rid]["breeze_prompt"]["guidance_scale"] + expected_hidden = [hidden[rid]] + ([hidden[rid + CFG_UNCOND_SUFFIX]] if scale != 1 else []) + torch.testing.assert_close(depth_hidden, torch.tensor(expected_hidden, dtype=torch.bfloat16), atol=0, rtol=0) + assert first.tolist() == [chosen[rid]] + assert call.kwargs["generators"] == [infos[rid]["breeze_state"]["generator"]] + assert call.kwargs["guidance_scale"] == scale + assert tuple(call.kwargs[key] for key in ("temperature", "top_k", "top_p")) == tuples[rid[-1]] + assert model.compute_logits(torch.zeros(len(order), 2)).argmax(-1).tolist() == [0, 0, 0, 7, 0, 0, 0, 0] + for index, rid in enumerate(order): + state = infos[rid]["breeze_state"] + audio = output.multimodal_outputs["codes"]["audio"][index] + if rid in live_order: + assert state["history"].tolist() == [[1, 3, 1, chosen[rid]]] + assert audio.tolist() == [[chosen[rid]] * 3] + else: + assert audio.numel() == 0 + torch.testing.assert_close(state["history"], original_states[rid]["history"], atol=0, rtol=0) + torch.testing.assert_close( + infos["finished_a"]["breeze_state"]["current"], original_states["finished_a"]["current"], atol=0, rtol=0 + ) + for parent, companion in (("cfg_a", cfg_a), ("cfg_b", cfg_b)): + assert infos[parent]["breeze_state"]["current"] is infos[companion]["breeze_state"]["current"] + + +@torch.inference_mode() +def test_first_code_shared_generator_preserves_order_across_parameter_groups(monkeypatch): + """A shared RNG stream must retain A,B,A request order across tuple groups.""" + model = _small_talker() + infos = [_request_info(rid) for rid in ("first_a", "middle_b", "last_a")] + parameters = [(0.9, 3, 1.0), (0.7, 2, 0.8), (0.9, 3, 1.0)] + actual_generator = torch.Generator().manual_seed(2**63 + 19) + expected_generator = torch.Generator().manual_seed(2**63 + 19) + hidden = torch.tensor([[0.2, 0.0], [0.0, 0.2], [0.1, 0.2]], dtype=torch.bfloat16) + expected = [] + for index, (info, values) in enumerate(zip(infos, parameters, strict=True)): + info["breeze_sampling"].update(temperature=values[0], top_k=values[1], top_p=values[2]) + info["breeze_state"]["generator"] = actual_generator + logits = F.linear(hidden[index : index + 1].float(), model.lm_head.weight) + logits[:, model.codebook_size : model.config.eos_token_id] = -torch.inf + expected.append(sample_logits(logits, *values, expected_generator)) + assert all(token.item() != model.config.eos_token_id for token in expected) + sample, calls = model._first_code_sampler.sample, [] + + def record_sample(logits, settings, generators): + calls.append((settings, len(logits), [generator is actual_generator for generator in generators])) + return sample(logits, settings, generators) + + monkeypatch.setattr(model._first_code_sampler, "sample", record_sample) + output = model.make_omni_output( + hidden, model_intermediate_buffer=infos, request_token_spans=[(0, 1), (1, 2), (2, 3)] + ) + assert calls == [(values, 1, [True]) for values in parameters] + torch.testing.assert_close(actual_generator.get_state(), expected_generator.get_state(), atol=0, rtol=0) + for info, token, audio in zip(infos, expected, output.multimodal_outputs["codes"]["audio"], strict=True): + torch.testing.assert_close(audio, token[:, None].repeat(1, 3), atol=0, rtol=0) + torch.testing.assert_close(info["breeze_state"]["history"], token[:, None], atol=0, rtol=0) + assert info["breeze_state"]["generator"] is actual_generator diff --git a/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py b/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py new file mode 100644 index 000000000..660737765 --- /dev/null +++ b/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""First-code sampler token, RNG, graph-cache and fallback contracts.""" + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits +from vllm_omni.model_executor.models.breeze_tts_2.first_code_sampler import BreezeFirstCodeSampler + +pytestmark = [pytest.mark.core_model] + +PARAMETERS = (0.9, 50, 1.0) +VOCAB = 2052 +EOS = 2051 + + +def _logits(count: int, *, device: str, seed: int = 100, vocab: int = VOCAB) -> list[torch.Tensor]: + source = torch.Generator(device="cpu").manual_seed(seed) + values = torch.randn((count, vocab), generator=source, dtype=torch.float32) * 3 + if vocab == VOCAB: + values[:, 2048:EOS] = -torch.inf + return list(values.to(device).split(1)) + + +def _generators(seeds: list[int], device: str) -> list[torch.Generator]: + return [torch.Generator(device=device).manual_seed(seed) for seed in seeds] + + +def _reference(rows, parameters, generators) -> list[torch.Tensor]: + return [sample_logits(row, *parameters, generator) for row, generator in zip(rows, generators, strict=True)] + + +def _assert_outputs(actual, expected, rows) -> None: + assert isinstance(actual, list) and len(actual) == len(expected) + for result, reference, row in zip(actual, expected, rows, strict=True): + assert result.shape == (1,) and result.dtype == torch.long and result.device == row.device + torch.testing.assert_close(result, reference, atol=0, rtol=0) + + +def _assert_states(actual, expected) -> None: + for left, right in zip(actual, expected, strict=True): + torch.testing.assert_close(left.get_state(), right.get_state(), atol=0, rtol=0) + + +@pytest.mark.cpu +def test_empty_and_cpu_sampling_preserve_eager_contract() -> None: + sampler = BreezeFirstCodeSampler() + assert sampler.sample([], PARAMETERS, []) == [] + assert not sampler._graphs + rows = _logits(3, device="cpu") + before = [row.clone() for row in rows] + actual_generators = _generators([42, 2**63 + 19, 2**64 - 23], "cpu") + expected_generators = _generators([42, 2**63 + 19, 2**64 - 23], "cpu") + for parameters in (PARAMETERS, (0.7, 20, 0.8), (1.2, 0, 0.95)): + expected = _reference(rows, parameters, expected_generators) + _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows) + _assert_states(actual_generators, expected_generators) + for row, original in zip(rows, before, strict=True): + torch.testing.assert_close(row, original, atol=0, rtol=0) + assert not sampler._graphs + + +@pytest.mark.cpu +def test_greedy_never_samples_or_advances_rng(monkeypatch) -> None: + sampler = BreezeFirstCodeSampler() + rows = _logits(3, device="cpu") + generators = _generators([42, 43, 44], "cpu") + original_states = [generator.get_state().clone() for generator in generators] + + def unexpected_multinomial(*_args, **_kwargs): + raise AssertionError("Greedy sampling must not call multinomial") + + monkeypatch.setattr(torch, "multinomial", unexpected_multinomial) + for parameters in ((0.0, 0, 1.0), (0.0, 1, 0.2)): + _assert_outputs(sampler.sample(rows, parameters, generators), [row.argmax(-1) for row in rows], rows) + for generator, state in zip(generators, original_states, strict=True): + torch.testing.assert_close(generator.get_state(), state, atol=0, rtol=0) + assert not sampler._graphs + + +@pytest.mark.cpu +@pytest.mark.parametrize( + "invalid", ["extra_generator", "missing_generator", "empty_rows", "multirow", "mixed_vocab", "mixed_dtype"] +) +def test_invalid_batch_fails_before_advancing_any_generator(invalid: str) -> None: + sampler = BreezeFirstCodeSampler() + rows = _logits(2, device="cpu") + generators = _generators([42, 43], "cpu") + if invalid == "extra_generator": + rows = rows[:1] + elif invalid == "missing_generator": + generators = generators[:1] + elif invalid == "empty_rows": + rows = [] + elif invalid == "multirow": + rows[1] = rows[1].expand(2, -1) + elif invalid == "mixed_vocab": + rows[1] = rows[1][:, :-1] + else: + rows[1] = rows[1].double() + states = [generator.get_state().clone() for generator in generators] + with pytest.raises(ValueError): + sampler.sample(rows, PARAMETERS, generators) + for generator, state in zip(generators, states, strict=True): + torch.testing.assert_close(generator.get_state(), state, atol=0, rtol=0) + assert not sampler._graphs + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None: + sampler = BreezeFirstCodeSampler() + seeds = list(range(42, 77)) + [2**63 + 19, 2**64 - 23, 101] + actual_generators = _generators(seeds, "cuda") + expected_generators = _generators(seeds, "cuda") + # Torch also accepts an unspecified CUDA index. Exercise it alongside the + # indexed generators created from serving request tensors. + indexed_device = torch.device("cuda", torch.accelerator.current_device_index()) + for pool in (actual_generators, expected_generators): + pool[1] = torch.Generator(device=indexed_device).manual_seed(seeds[1]) + orders = ( + [0], + [2, 0, 1], + list(range(16)), + [36, *range(16)], + [35, *range(30)], + list(reversed(range(32))), + [35, 1, 2], + ) + retained: list[tuple[torch.Tensor, torch.Tensor]] = [] + global_state = torch.cuda.get_rng_state().clone() + for sweep in range(2): + if sweep: + for pool in (actual_generators, expected_generators): + pool[2] = torch.Generator(device="cuda").manual_seed(2**64 - 101) + for case, order in enumerate(orders): + rows = _logits(len(order), device="cuda", seed=1000 + sweep * 10 + case) + if sweep and case == 3: + for index, row in enumerate(rows): + row.fill_(-torch.inf) + row[0, EOS if index % 2 else 2] = 0 + original = [row.clone() for row in rows] + live_actual = [actual_generators[index] for index in order] + live_expected = [expected_generators[index] for index in order] + expected = _reference(rows, PARAMETERS, live_expected) + actual = sampler.sample(rows, PARAMETERS, live_actual) + _assert_outputs(actual, expected, rows) + _assert_states(actual_generators, expected_generators) + retained.extend((value, reference.clone()) for value, reference in zip(actual, expected, strict=True)) + for row, before in zip(rows, original, strict=True): + torch.testing.assert_close(row, before, atol=0, rtol=0) + # Preserve the subsequent 15 independent depth draws, including + # skipping them for EOS. All inactive/companion generators remain + # in the full-pool comparison above and below. + for index, token in enumerate(torch.cat(expected).cpu().tolist()): + if token == EOS: + continue + left = torch.empty((15, 2051), device="cuda") + right = torch.empty_like(left) + for codebook in range(15): + left[codebook].exponential_(generator=live_actual[index]) + right[codebook].exponential_(generator=live_expected[index]) + torch.testing.assert_close(left, right, atol=0, rtol=0) + _assert_states(actual_generators, expected_generators) + # B1/3/16/17/31/32 must use exactly buckets 1/4/16/32. + assert len(sampler._graphs) == 4 + rows = _logits(3, device="cuda", seed=2026) + expected = _reference(rows, PARAMETERS, expected_generators[:3]) + + def unexpected_multinomial(*_args, **_kwargs): + raise AssertionError("Warmed graph replay must not call Python multinomial") + + with monkeypatch.context() as replay_guard: + replay_guard.setattr(torch, "multinomial", unexpected_multinomial) + _assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators[:3]), expected, rows) + _assert_outputs( + sampler.sample(rows, (0.0, 1, 0.2), actual_generators[:3]), [row.argmax(-1) for row in rows], rows + ) + _assert_states(actual_generators, expected_generators) + assert len(sampler._graphs) == 4 + for value, expected in retained: + torch.testing.assert_close(value, expected, atol=0, rtol=0) + torch.testing.assert_close(torch.cuda.get_rng_state(), global_state, atol=0, rtol=0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None: + sampler = BreezeFirstCodeSampler() + actual_generators = _generators([42, 43, 44], "cuda") + expected_generators = _generators([42, 43, 44], "cuda") + cases = ((0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95), (1.0, 1, 1.0), (1.0, 0, 0.5)) + for index, parameters in enumerate(cases): + rows = _logits(3, device="cuda", seed=200 + index) + if index >= 3: + for row in rows: + row.fill_(-torch.inf) + row[0, :4] = 0 # top-k ties and exact top-p=0.5 boundary. + expected = _reference(rows, parameters, expected_generators) + _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows) + _assert_states(actual_generators, expected_generators) + assert len(sampler._graphs) == len(cases) + rows = _logits(3, device="cuda", vocab=17) + expected = _reference(rows, PARAMETERS, expected_generators) + _assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators), expected, rows) + _assert_states(actual_generators, expected_generators) + assert len(sampler._graphs) == len(cases) + 1 + + class UnavailableGraph: + def __init__(self, *_args, **_kwargs): + raise AssertionError("Unsupported graph API must select the eager fallback") + + fallback = BreezeFirstCodeSampler() + expected = _reference(rows, PARAMETERS, expected_generators) + with monkeypatch.context() as api_guard: + api_guard.setattr(torch.cuda, "CUDAGraph", UnavailableGraph) + _assert_outputs(fallback.sample(rows, PARAMETERS, actual_generators), expected, rows) + _assert_states(actual_generators, expected_generators) + assert not fallback._graphs + + # Validate every row before the first eager/captured draw, including an + # invalid device that appears only after an otherwise valid CUDA row. + invalid = [rows[0], rows[1].cpu(), rows[2]] + with pytest.raises(ValueError): + sampler.sample(invalid, PARAMETERS, actual_generators) + _assert_states(actual_generators, expected_generators) + # A late invalid generator must likewise fail before row zero samples. + cpu_generator = torch.Generator(device="cpu").manual_seed(77) + cpu_reference = torch.Generator(device="cpu").manual_seed(77) + mixed_generators = [actual_generators[0], cpu_generator, actual_generators[2]] + with pytest.raises(ValueError): + sampler.sample(rows, PARAMETERS, mixed_generators) + _assert_states(mixed_generators, [expected_generators[0], cpu_reference, expected_generators[2]]) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> None: + sampler = BreezeFirstCodeSampler() + rows = _logits(3, device="cuda") + # The duplicate-generator guard must precede lookup of an existing graph. + sampler.sample(rows, PARAMETERS, _generators([100, 101, 102], "cuda")) + cached = dict(sampler._graphs) + assert len(cached) == 1 + actual = torch.Generator(device="cuda").manual_seed(2**64 - 23) + reference = torch.Generator(device="cuda").manual_seed(2**64 - 23) + multinomial = torch.multinomial + calls = [] + + def eager_spy(*args, **kwargs): + calls.append(kwargs.get("generator")) + return multinomial(*args, **kwargs) + + for order in ([0, 1, 2], [2, 0, 1]): + ordered = [rows[index] for index in order] + expected = _reference(ordered, PARAMETERS, [reference] * 3) + calls.clear() + with monkeypatch.context() as fallback_guard: + fallback_guard.setattr(torch, "multinomial", eager_spy) + _assert_outputs(sampler.sample(ordered, PARAMETERS, [actual] * 3), expected, ordered) + assert len(calls) == 3 and all(generator is actual for generator in calls) + _assert_states([actual], [reference]) + assert set(sampler._graphs) == set(cached) + assert all(sampler._graphs[key] is value for key, value in cached.items()) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) -> None: + sampler = BreezeFirstCodeSampler() + rows = _logits(1, device="cuda") + actual_generators = _generators([42], "cuda") + expected_generators = _generators([42], "cuda") + for top_k in range(1, 33): + parameters = (0.9, top_k, 1.0) + expected = _reference(rows, parameters, expected_generators) + _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows) + _assert_states(actual_generators, expected_generators) + cached = dict(sampler._graphs) + assert len(cached) == 32 + parameters = (0.9, 33, 1.0) + expected = _reference(rows, parameters, expected_generators) + multinomial = torch.multinomial + calls = [] + + def eager_spy(*args, **kwargs): + calls.append(kwargs.get("generator")) + return multinomial(*args, **kwargs) + + with monkeypatch.context() as fallback_guard: + fallback_guard.setattr(torch, "multinomial", eager_spy) + _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows) + assert len(calls) == 1 and calls[0] is actual_generators[0] + _assert_states(actual_generators, expected_generators) + assert set(sampler._graphs) == set(cached) + assert all(sampler._graphs[key] is value for key, value in cached.items()) + expected = _reference(rows, (0.9, 1, 1.0), expected_generators) + + def unexpected_multinomial(*_args, **_kwargs): + raise AssertionError("Cache exhaustion must retain existing captured entries") + + with monkeypatch.context() as replay_guard: + replay_guard.setattr(torch, "multinomial", unexpected_multinomial) + _assert_outputs(sampler.sample(rows, (0.9, 1, 1.0), actual_generators), expected, rows) + _assert_states(actual_generators, expected_generators) diff --git a/tests/model_executor/models/test_breeze_tts_2_graphs.py b/tests/model_executor/models/test_breeze_tts_2_graphs.py new file mode 100644 index 000000000..72c2cfeea --- /dev/null +++ b/tests/model_executor/models/test_breeze_tts_2_graphs.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Check captured cache ownership and padded text masking against eager models.""" + +import pytest +import torch +from transformers import LlamaConfig, T5Gemma2TextConfig +from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2TextEncoder +from vllm.platforms import current_platform + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import BreezeDepthDecoder, sample_logits +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration +from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import ( + BreezeTextEncoderCompiled, + BreezeTextEncoderGraph, +) + +pytestmark = [pytest.mark.core_model] + + +@pytest.fixture +def full_precision_matmul(): + # Different padded shapes can select different TF32 kernels. Compare + # masking and replay semantics independently of that reduced precision. + original_precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + try: + yield + finally: + torch.set_float32_matmul_precision(original_precision) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_text_graph_masks_padding_and_local_attention_across_replays(full_precision_matmul) -> None: + torch.manual_seed(42) + config = T5Gemma2TextConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + sliding_window=16, + layer_types=["sliding_attention", "full_attention"], + dropout_rate=0.0, + ) + config._attn_implementation = "sdpa" + encoder = T5Gemma2TextEncoder(config).to("cuda").eval() + encoder.embed_tokens.eoi_token_index = 17 + encoder.embed_tokens.eoi_embedding.fill_(0.4) + projection = torch.nn.Linear(32, 48, bias=False).to("cuda").eval() + graph = BreezeTextEncoderGraph(encoder, projection, 64) + for length in (31, 51, 17, 31): + prompt = torch.randint(0, 64, (1, length), device="cuda") + prompt[:, 3] = 17 + expected = projection(encoder(input_ids=prompt).last_hidden_state)[0] + actual = graph.run(prompt) + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5) + + batched = BreezeTextEncoderGraph(encoder, projection, 64, batch_size=3) + prompts = [torch.randint(0, 64, (length,), device="cuda") for length in (31, 51, 17)] + outputs = batched.run_batch(prompts) + for prompt, actual in zip(prompts, outputs, strict=True): + expected = projection(encoder(input_ids=prompt[None]).last_hidden_state)[0] + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5) + + compiled = BreezeTextEncoderCompiled(encoder, projection) + long_prompts = [torch.randint(0, 64, (length,)) for length in (211, 173, 129)] + for batch in ([long_prompts[0]], long_prompts, [long_prompts[1]]): + outputs = compiled.run_batch(batch) + for prompt, actual in zip(batch, outputs, strict=True): + expected = projection(encoder(input_ids=prompt[None].to("cuda")).last_hidden_state)[0] + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5) + + +def _make_depth_decoder(vocab_size: int = 35) -> BreezeDepthDecoder: + config = LlamaConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + vocab_size=vocab_size, + max_position_embeddings=16, + rope_theta=500000.0, + rope_scaling={ + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 0.001953125, + "high_freq_factor": 0.0078125, + "original_max_position_embeddings": 16, + }, + ) + config.num_codebooks = 16 + config.audio_embed_size = 48 + depth = BreezeDepthDecoder(config).to("cuda").eval() + torch.nn.init.normal_(depth.codebooks_head, std=0.1) + return depth + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_depth_graph_survives_interleaved_batches_and_workspace_allocations(full_precision_matmul, monkeypatch) -> None: + torch.manual_seed(42) + depth = _make_depth_decoder() + hidden = torch.randn(2, 48, device="cuda") + first = torch.tensor([7, 11], device="cuda") + generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in (42, 99)] + expected = torch.cat( + [ + depth.generate_frame(hidden[i : i + 1], first[i : i + 1], temperature=0, top_k=0, top_p=1, generator=g) + for i, g in enumerate(generators) + ] + ) + for batch in (1, 2, 1, 2): + # Unowned allocations made before graph capture would be recycled here. + pressure = [torch.full((batch, 2, depth.num_codebooks, 8), float("nan"), device="cuda") for _ in range(32)] + actual = depth.generate_frames( + hidden[:batch], first[:batch], temperature=0, top_k=0, top_p=1, generators=generators[:batch] + ) + torch.testing.assert_close(actual, expected[:batch], atol=0, rtol=0) + del pressure + + def sample(batch: int) -> torch.Tensor: + return depth.generate_frames( + hidden[:batch], + first[:batch], + temperature=0.9, + top_k=10, + top_p=0.8, + generators=[torch.Generator(device="cuda").manual_seed(seed) for seed in (42, 99)[:batch]], + ) + + sampled = sample(1) + sampled_graph = depth._graphs[(1, 1, False)] + greedy_graph = depth._graphs[(1, 1, True)] + assert sampled_graph is not greedy_graph + sample(2) + torch.testing.assert_close(sample(1), sampled, atol=0, rtol=0) + for temperature, top_k, top_p in ((0.0, 0, 1.0), (0.7, 20, 0.95), (0.0, 1, 0.2), (1.2, 0, 0.8)): + rng_before = generators[0].get_state() + actual = depth.generate_frames( + hidden[:1], first[:1], temperature=temperature, top_k=top_k, top_p=top_p, generators=generators[:1] + ) + if temperature == 0: + assert depth._graphs[(1, 1, True)] is greedy_graph + torch.testing.assert_close(actual, expected[:1], atol=0, rtol=0) + torch.testing.assert_close(generators[0].get_state(), rng_before, atol=0, rtol=0) + else: + assert depth._graphs[(1, 1, False)] is sampled_graph + torch.testing.assert_close(sample(1), sampled, atol=0, rtol=0) + + graph = None + for scale in (4.0, 0.5, 4.0): + expected_cfg = depth._generate_frame( + hidden, first[:1], depth._allocate_cache(hidden), 0, 0, 1, generators[0], guidance_scale=scale + ) + actual_cfg = depth.generate_frames( + hidden, first[:1], temperature=0, top_k=0, top_p=1, generators=generators[:1], guidance_scale=scale + ) + torch.testing.assert_close(actual_cfg, expected_cfg, atol=0, rtol=0) + if graph is not None: + assert depth._graphs[(2, 2, True)] is graph + graph = depth._graphs[(2, 2, True)] + sample(2) + + # Use a still-cold sampled CFG entry to exercise platforms without + # captured generator support, while reusing the existing compiled model. + fallback_key = (2, 2, False) + assert fallback_key not in depth._graphs + reference_generator = torch.Generator(device="cuda").set_state(generators[0].get_state()) + inactive_rng = generators[1].get_state() + exponential = torch.Tensor.exponential_ + observed_generators = [] + + def record_eager_noise(tensor, *args, **kwargs): + observed_generators.append(kwargs.get("generator")) + return exponential(tensor, *args, **kwargs) + + fallback_graph = None + retained = [] + with monkeypatch.context() as platform_guard: + platform_guard.setattr(current_platform, "is_cuda", lambda: False) + for scale in (4.0, 0.5, 4.0): + expected = depth._generate_frame( + hidden, + first[:1], + depth._allocate_cache(hidden), + 0.9, + 10, + 0.8, + reference_generator, + guidance_scale=scale, + ) + observed_generators.clear() + with monkeypatch.context() as noise_spy: + noise_spy.setattr(torch.Tensor, "exponential_", record_eager_noise) + actual = depth.generate_frames( + hidden, + first[:1], + temperature=0.9, + top_k=10, + top_p=0.8, + generators=generators[:1], + guidance_scale=scale, + ) + entry = depth._graphs[fallback_key] + assert entry.noise_generators is None + assert len(observed_generators) == depth.num_codebooks - 1 + assert all(generator is generators[0] for generator in observed_generators) + if fallback_graph is not None: + assert entry is fallback_graph + fallback_graph = entry + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(generators[0].get_state(), reference_generator.get_state(), atol=0, rtol=0) + torch.testing.assert_close(generators[1].get_state(), inactive_rng, atol=0, rtol=0) + retained.append((actual, expected.clone())) + for actual, expected in retained: + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_matmul, monkeypatch) -> None: + def reject_eager_noise(*args, **kwargs): + raise AssertionError("Sampled depth replay must not issue eager RNG draws") + + torch.manual_seed(17) + # Match the checkpoint's multinomial draw size, including reserved IDs. + depth = _make_depth_decoder(vocab_size=2051) + hidden = torch.randn(32, 48, device="cuda") + uncond = torch.randn_like(hidden) + first = torch.arange(32, device="cuda") + first_logits = torch.randn(32, depth.vocab_size, device="cuda") + first_logits[:, -3:] = -torch.inf + seeds = [42 + index for index in range(32)] + seeds[0], seeds[-1] = 2**63 + 17, 2**64 - 1 + generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds] + reference_generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds] + + expected_greedy = torch.cat( + [ + depth.generate_frame( + hidden[index : index + 1], + first[index : index + 1], + temperature=0, + top_k=0, + top_p=1, + generator=reference_generators[index], + ) + for index in range(32) + ] + ) + retained: list[tuple[torch.Tensor, torch.Tensor]] = [] + for batch in (*range(1, 33), 17, 11, 3, 1): + rows = list(reversed(range(batch))) + actual = depth.generate_frames( + hidden[rows], first[rows], temperature=0, top_k=0, top_p=1, generators=[generators[row] for row in rows] + ) + expected = expected_greedy[rows] + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + retained.append((actual, expected.clone())) + assert set(depth._graphs) == {(batch, 1, True) for batch in (1, 2, 4, 8, 16, 32)} + for generator, reference in zip(generators, reference_generators, strict=True): + torch.testing.assert_close(generator.get_state(), reference.get_state(), atol=0, rtol=0) + + retired_generators: list[tuple[torch.Generator, torch.Tensor]] = [] + for temperature, guidance_scale in ((0.9, 1.0), (0.0, 4.0), (0.9, 4.0)): + batches: tuple[int, ...] = (3, 5, 11, 7, 3, 1) + if guidance_scale == 1: + batches = (16, 17, 31, 32, 17, 31, 16, *batches) + for step, batch in enumerate(batches): + if guidance_scale == 1 and step == 4: + # New requests reuse rows in the same 32-request graph. Their + # seeds must replace the previous owners, not just offsets. + for row, seed in ((0, 2**64 - 1), (16, 2**63 + 321)): + retired_generators.append((generators[row], generators[row].get_state())) + generators[row] = torch.Generator(device="cuda").manual_seed(seed) + reference_generators[row] = torch.Generator(device="cuda").manual_seed(seed) + rows = list(reversed(range(batch))) + batch_hidden = hidden[rows] + if guidance_scale != 1: + batch_hidden = torch.cat((batch_hidden, uncond[rows])) + expected, first_codes = [], [] + for row in rows: + # The talker and depth decoder share each request's generator. + # Exercise a codebook-0 sample immediately before graph replay. + row_first = sample_logits(first_logits[row : row + 1], temperature, 10, 0.8, generators[row]) + reference_first = sample_logits( + first_logits[row : row + 1], temperature, 10, 0.8, reference_generators[row] + ) + torch.testing.assert_close(row_first, reference_first, atol=0, rtol=0) + first_codes.append(row_first) + row_hidden = hidden[row : row + 1] + if guidance_scale != 1: + row_hidden = torch.cat((row_hidden, uncond[row : row + 1])) + expected.append( + depth._generate_frame( + row_hidden, + reference_first, + depth._allocate_cache(row_hidden), + temperature, + 10, + 0.8, + reference_generators[row], + guidance_scale=guidance_scale if guidance_scale != 1 else None, + ) + ) + request_bucket = 1 << (batch - 1).bit_length() + branches = 2 if guidance_scale != 1 else 1 + graph_key = (request_bucket * branches, branches, temperature == 0) + with monkeypatch.context() as rng_guard: + if temperature > 0 and graph_key in depth._graphs: + # A warmed replay must generate noise within the CUDA + # graph, without dispatching new RNG kernels from Python. + rng_guard.setattr(torch.Tensor, "exponential_", reject_eager_noise) + actual = depth.generate_frames( + batch_hidden, + torch.cat(first_codes), + temperature=temperature, + top_k=10, + top_p=0.8, + generators=[generators[row] for row in rows], + guidance_scale=guidance_scale, + ) + expected_frames = torch.cat(expected) + torch.testing.assert_close(actual, expected_frames, atol=0, rtol=0) + retained.append((actual, expected_frames.clone())) + # Check the next codebook-0 sample as well as the depth result: a + # replay must return the advanced RNG stream to its live request. + for row in rows: + actual_next = sample_logits(first_logits[row : row + 1], temperature, 10, 0.8, generators[row]) + expected_next = sample_logits( + first_logits[row : row + 1], temperature, 10, 0.8, reference_generators[row] + ) + torch.testing.assert_close(actual_next, expected_next, atol=0, rtol=0) + # Padding must not draw RNG, and inactive request streams must not advance. + for generator, reference in zip(generators, reference_generators, strict=True): + torch.testing.assert_close(generator.get_state(), reference.get_state(), atol=0, rtol=0) + for generator, state in retired_generators: + torch.testing.assert_close(generator.get_state(), state, atol=0, rtol=0) + + expected_keys = {(batch, 1, True) for batch in (1, 2, 4, 8, 16, 32)} + expected_keys.update((batch, 1, False) for batch in (1, 4, 8, 16, 32)) + expected_keys.update((batch, 2, greedy) for batch in (2, 8, 16, 32) for greedy in (False, True)) + assert set(depth._graphs) == expected_keys + for (batch_bucket, branches, greedy), entry in depth._graphs.items(): + if greedy: + assert entry.noise_generators is None + else: + assert entry.noise_generators is not None + assert len(entry.noise_generators) == batch_bucket // branches + assert {id(generator) for generator in entry.noise_generators}.isdisjoint( + id(generator) for generator in generators + ) + for actual, expected in retained: + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_batched_decode_embedding_matches_scalar_bfloat16() -> None: + # Use the checkpoint's RVQ table dimensions without loading its backbone. + model = BreezeForConditionalGeneration.__new__(BreezeForConditionalGeneration) + torch.nn.Module.__init__(model) + model.num_codebooks, model.hidden_size = 16, 2048 + model.register_buffer("offsets", torch.arange(16, device="cuda") * 2051) + model.depth_decoder = torch.nn.Module() + torch.manual_seed(17) + model.depth_decoder.embed_tokens = torch.nn.Embedding(16 * 2051, 2048, device="cuda", dtype=torch.bfloat16) + for seed in (17, 42, 99): + generator = torch.Generator(device="cuda").manual_seed(seed) + frames = torch.randint(0, 2051, (32, 16), device="cuda", generator=generator) + frames[1] = frames[0] + for batch in (1, 2, 3, 4, 5, 6, 7, 8, 16, 32, 7, 1): + # Reverse physical rows and revisit smaller batches after expansion. + current_frames = [row[None] for row in frames[:batch].flip(0)] + if batch > 1: + # CFG branches share their generated audio frame. + current_frames[-2] = current_frames[-1] + infos = [ + { + "_omni_is_prefill": False, + "breeze_state": {"current": current}, + "global_request_id": [str(index)], + "breeze_prompt": {}, + "breeze_sampling": {}, + } + for index, current in enumerate(current_frames) + ] + ids = torch.zeros(batch, device="cuda", dtype=torch.long) + expected = torch.cat( + [ + model.preprocess(input_ids=ids[index : index + 1], input_embeds=None, **info)[1] + for index, info in enumerate(infos) + ] + ) + output_ids, actual, updates = model.preprocess_decode_batch(input_ids=ids, req_infos=infos) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert output_ids is ids + assert updates == [{} for _ in infos] diff --git a/tests/model_executor/models/test_breeze_tts_2_preprocess.py b/tests/model_executor/models/test_breeze_tts_2_preprocess.py new file mode 100644 index 000000000..7b2316e39 --- /dev/null +++ b/tests/model_executor/models/test_breeze_tts_2_preprocess.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU contracts for Breeze's decode-only embedding preprocessing.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _embedding_talker(dtype): + model = BreezeForConditionalGeneration.__new__(BreezeForConditionalGeneration) + torch.nn.Module.__init__(model) + model.num_codebooks, model.hidden_size = 16, 8 + model.register_buffer("offsets", torch.arange(model.num_codebooks) * 7) + model.depth_decoder = torch.nn.Module() + model.depth_decoder.embed_tokens = torch.nn.Embedding(16 * 7, 8, dtype=dtype) + with torch.no_grad(): + model.depth_decoder.embed_tokens.weight.copy_( + torch.randn(16 * 7, 8, generator=torch.Generator().manual_seed(17)) + ) + return model + + +def _decode_info(request_id, current, seed): + return { + "global_request_id": [request_id], + "breeze_prompt": {"role": "cond", "guidance_scale": 1.0}, + "breeze_sampling": {"temperature": 0.9, "top_k": 4, "top_p": 1.0, "repetition_penalty": 1.1}, + "breeze_state": { + "current": current, + "history": torch.tensor([[1, 2, 3]], dtype=torch.long), + "generator": torch.Generator().manual_seed(seed), + }, + "_omni_is_prefill": False, + } + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@torch.inference_mode() +def test_decode_batch_matches_scalar_after_reordering_shrinking_and_cfg(dtype): + model = _embedding_talker(dtype) + shared_frame = torch.arange(16, dtype=torch.long).reshape(1, 16) % 7 + infos = [ + _decode_info("unconditional", shared_frame, 1), + _decode_info("other", (shared_frame + 1) % 7, 2), + _decode_info("conditioned", shared_frame, 3), + _decode_info("last", (shared_frame + 3) % 7, 4), + ] + infos[0]["breeze_prompt"] = {"role": "uncond", "guidance_scale": 4.0} + infos[2]["breeze_prompt"]["guidance_scale"] = 4.0 + saved = [] + for indices in ([0, 1, 2, 3], [2, 0, 1], [1]): + selected = [infos[index] for index in indices] + # Placeholder IDs deliberately differ from the RVQ frame IDs. + input_ids = torch.arange(len(selected), dtype=torch.long) + 100 + expected = torch.cat( + [ + model.preprocess(input_ids=input_ids[row : row + 1], input_embeds=None, **info)[1] + for row, info in enumerate(selected) + ] + ) + state_before = [ + ( + info["breeze_state"]["current"].clone(), + info["breeze_state"]["history"].clone(), + info["breeze_state"]["generator"].get_state(), + ) + for info in infos + ] + output_ids, embeds, updates = model.preprocess_decode_batch(input_ids=input_ids, req_infos=selected) + assert output_ids is input_ids + torch.testing.assert_close(embeds, expected, rtol=0, atol=0) + assert updates == [{} for _ in selected] + for info, (current, history, rng) in zip(infos, state_before, strict=True): + state = info["breeze_state"] + assert torch.equal(state["current"], current) + assert torch.equal(state["history"], history) + assert torch.equal(state["generator"].get_state(), rng) + saved.append((embeds, embeds.clone())) + # Simulate the next generated frame. CFG branches share the new frame. + next_frame = (infos[2]["breeze_state"]["current"] + 2) % 7 + infos[0]["breeze_state"]["current"] = next_frame + infos[2]["breeze_state"]["current"] = next_frame + infos[1]["breeze_state"]["current"] = (infos[1]["breeze_state"]["current"] + 1) % 7 + assert infos[0]["breeze_state"]["current"] is infos[2]["breeze_state"]["current"] + for output, snapshot in saved: + assert torch.equal(output, snapshot) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_decode_batch_empty_preserves_embedding_shape_and_dtype(dtype): + model = _embedding_talker(dtype) + ids = torch.empty(0, dtype=torch.long) + output_ids, embeds, updates = model.preprocess_decode_batch(input_ids=ids, req_infos=[]) + assert output_ids is ids + assert embeds.shape == (0, 8) + assert embeds.dtype == dtype + assert embeds.device == model.depth_decoder.embed_tokens.weight.device + assert updates == [] + + +@pytest.mark.parametrize("invalid", ["count", "state", "frame_shape", "prefill"]) +def test_decode_batch_rejects_invalid_request_state(invalid): + model = _embedding_talker(torch.float32) + info = _decode_info("request", torch.zeros(1, 16, dtype=torch.long), 42) + ids = torch.tensor([0]) + if invalid == "count": + ids = torch.tensor([0, 1]) + elif invalid == "state": + del info["breeze_state"] + elif invalid == "frame_shape": + info["breeze_state"]["current"] = torch.zeros(2, 16, dtype=torch.long) + else: + info["_omni_is_prefill"] = True + with pytest.raises(ValueError, match="Breeze decode"): + model.preprocess_decode_batch(input_ids=ids, req_infos=[info]) diff --git a/tests/model_executor/models/test_breeze_tts_2_registration.py b/tests/model_executor/models/test_breeze_tts_2_registration.py new file mode 100644 index 000000000..9d81230d0 --- /dev/null +++ b/tests/model_executor/models/test_breeze_tts_2_registration.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The accelerated runtime preserves Breeze-TTS-2's model identifiers.""" + +from pathlib import Path + +import pytest +import yaml + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.entrypoints.openai.tts_adapters import detect_tts_model_type, resolve_adapter +from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter +from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_breeze_tts_2_pipeline_keeps_checkpoint_and_stage_identifiers(): + assert OMNI_PIPELINES["breeze"] is BREEZE_TTS_2_PIPELINE + assert BREEZE_TTS_2_PIPELINE.model_arch == "BreezeForConditionalGeneration" + assert BREEZE_TTS_2_PIPELINE.default_deploy_config_name == "breeze_tts_2.yaml" + assert [stage.model_stage for stage in BREEZE_TTS_2_PIPELINE.stages] == [ + "breeze_tts_2", + "breeze_tts_2_codec", + ] + + +def test_breeze_tts_2_adapter_remains_registered(): + assert resolve_adapter("breeze_tts_2") is BreezeTTS2Adapter + + +@pytest.mark.parametrize( + "model_stage,model_arch", + [ + ("breeze_tts_2", None), + ("breeze_tts_2_codec", None), + ("breeze_tts_2", "BreezeForConditionalGeneration"), + (None, "BreezeForConditionalGeneration"), + ], +) +def test_breeze_tts_2_detection_keeps_existing_identifiers(model_stage, model_arch): + assert detect_tts_model_type(model_stage, model_arch) == "breeze_tts_2" + + +@pytest.mark.parametrize("suffix", ["", "_throughput"]) +def test_breeze_deploy_aliases_match_canonical_configuration(suffix): + deploy_dir = Path(__file__).parents[3] / "vllm_omni" / "deploy" + canonical = yaml.safe_load((deploy_dir / f"breeze_tts_2{suffix}.yaml").read_text()) + alias = yaml.safe_load((deploy_dir / f"breeze_tts{suffix}.yaml").read_text()) + assert canonical == alias diff --git a/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py b/tests/model_executor/models/test_breeze_tts_2_sampling_eos.py similarity index 98% rename from tests/model_executor/models/breeze_tts_2/test_sampling_eos.py rename to tests/model_executor/models/test_breeze_tts_2_sampling_eos.py index 8cd389d17..2ca0b268f 100644 --- a/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py +++ b/tests/model_executor/models/test_breeze_tts_2_sampling_eos.py @@ -19,7 +19,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @pytest.fixture def talker_sampling_params(): - deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml" + deploy_path = Path(__file__).parents[3] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml" deploy = yaml.safe_load(deploy_path.read_text(encoding="utf-8")) params = SamplingParams( **merge_sampling_constraints( diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index ee51b0f58..b05c6eab2 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -3,6 +3,7 @@ from contextlib import contextmanager from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -1030,3 +1031,154 @@ def test_preprocess_one_token_chunked_prefill_tail_then_decode(monkeypatch, batc assert route == ("batch" if batched_decode and not is_prefill else "normal") torch.testing.assert_close(embeds, torch.full((1, 4), 42.0 if is_prefill else 43.0)) assert ("codes" in runner.model_intermediate_buffer["r"]) == (not is_prefill) + + +class DecodeOnlyPreprocessModel: + has_preprocess = True + + def __init__(self, *, rewrite_ids=False, invalid_output=None): + self.rewrite_ids = rewrite_ids + self.invalid_output = invalid_output + self.calls = [] + + @staticmethod + def embed_input_ids(input_ids, **kwargs): + return input_ids.float().view(-1, 1).expand(-1, 4).clone() + + def preprocess(self, input_ids, input_embeds, **info): + self.calls.append(("scalar", [info["request_id"]], [info["_omni_is_prefill"]])) + return input_ids, self.embed_input_ids(input_ids), {} + + def preprocess_decode_batch(self, *, input_ids, req_infos): + self.calls.append( + ("batch", [info["request_id"] for info in req_infos], [info["_omni_is_prefill"] for info in req_infos]) + ) + ids = input_ids + 100 if self.rewrite_ids else input_ids + embeds = self.embed_input_ids(ids) + 10 + updates = [{"decode_marker": info["request_id"]} for info in req_infos] + if self.invalid_output == "embeddings": + embeds = embeds[:-1] + elif self.invalid_output == "ids": + ids = ids[:-1] + elif self.invalid_output == "updates": + updates = updates[:-1] + return ids, embeds, updates + + +def _without_mtp_buffers(runner): + runner.has_talker_mtp = False + for name in ("talker_mtp", "talker_mtp_input_ids", "talker_mtp_inputs_embeds", "last_talker_hidden", "text_step"): + delattr(runner, name) + + +@pytest.mark.parametrize("rewrite_ids", [False, True], ids=["original-ids", "replaced-ids"]) +@pytest.mark.parametrize("order", ["interleaved", "reversed", "adjacent"]) +def test_decode_batch_without_mtp_preserves_mixed_rows_and_buffers(monkeypatch, rewrite_ids, order): + rows = [ + ("cfg-uncond", 5, 5, 1), + ("prefill-tail", 5, 4, 1), + ("cfg-cond", 5, 6, 1), + ("single-prompt", 1, 0, 1), + ("multi-token-decode", 5, 5, 2), + ] + if order == "reversed": + rows.reverse() + elif order == "adjacent": + rows = [rows[1], rows[0], rows[2], *rows[3:]] + runner, scheduled, _, _ = _make_phase_runner(monkeypatch, rows, batched_decode=False) + _without_mtp_buffers(runner) + model = DecodeOnlyPreprocessModel(rewrite_ids=rewrite_ids) + runner.model = model + total = scheduled.total_num_scheduled_tokens + runner.inputs_embeds = DummyBuffer(torch.full((total + 2, 4), -99.0)) + ids_buffer = runner.input_ids.gpu + embeds_buffer = runner.inputs_embeds.gpu + output_ids, output_embeds, *_ = runner._preprocess(scheduled, total) + assert output_ids.data_ptr() == ids_buffer.data_ptr() + assert output_embeds.data_ptr() == embeds_buffer.data_ptr() + assert torch.equal(embeds_buffer[total:], torch.full((2, 4), -99.0)) + seen = [ + (route, rid, prefill) for route, ids, phases in model.calls for rid, prefill in zip(ids, phases, strict=True) + ] + offset = 0 + for (route, rid, prefill), (expected_id, prompt, computed, span) in zip(seen, rows, strict=True): + eligible = span == 1 and computed >= prompt + assert rid == expected_id + assert prefill == (computed < prompt) + assert route == ("batch" if eligible else "scalar") + expected_ids = torch.arange(offset, offset + span) + (100 if eligible and rewrite_ids else 0) + assert torch.equal(output_ids[offset : offset + span], expected_ids) + expected_embeds = expected_ids.float().view(-1, 1).expand(-1, 4) + (10 if eligible else 0) + torch.testing.assert_close(output_embeds[offset : offset + span], expected_embeds) + info = runner.model_intermediate_buffer[rid] + assert info.get("decode_marker") == (rid if eligible else None) + offset += span + + +def test_decode_batch_without_mtp_accepts_wrapper_input_ids_none(monkeypatch): + rows = [("first", 5, 5, 1), ("second", 5, 6, 1)] + runner, scheduled, _, _ = _make_phase_runner(monkeypatch, rows, batched_decode=False) + _without_mtp_buffers(runner) + model = DecodeOnlyPreprocessModel() + runner.model = model + runner.supports_mm_inputs = True + runner.maybe_get_ec_connector_output = _noop_forward_context + runner.encoder_cache = {} + runner._execute_mm_encoder = lambda output: None + runner._gather_mm_embeddings = lambda output: ([], None) + runner._extract_mm_kwargs = lambda output: {} + runner._prepare_mm_inputs = lambda size: (None, runner.inputs_embeds.gpu[:size]) + ids, embeds, *_ = runner._preprocess(scheduled, 2) + assert ids.data_ptr() == runner.input_ids.gpu.data_ptr() + assert model.calls == [("batch", ["first", "second"], [False, False])] + torch.testing.assert_close(embeds, torch.tensor([[10.0] * 4, [11.0] * 4])) + + +@pytest.mark.parametrize( + "offsets", + [ + pytest.param([0, 2, 3], id="gap"), + pytest.param([1, 0, 2], id="reversed"), + pytest.param([0, 0, 1], id="duplicate"), + pytest.param([-1, 0, 1], id="negative-contiguous-start"), + pytest.param([1, 2, 3], id="contiguous-out-of-bounds"), + ], +) +def test_decode_batch_without_mtp_rejects_invalid_offsets_before_model_hook(monkeypatch, offsets): + rows = [("first", 5, 5, 1), ("second", 5, 6, 1)] + runner, scheduled, _, _ = _make_phase_runner(monkeypatch, rows, batched_decode=False) + _without_mtp_buffers(runner) + model = DecodeOnlyPreprocessModel(rewrite_ids=True) + runner.model = model + scalar_hook = Mock(wraps=model.preprocess) + batch_hook = Mock(wraps=model.preprocess_decode_batch) + monkeypatch.setattr(model, "preprocess", scalar_hook) + monkeypatch.setattr(model, "preprocess_decode_batch", batch_hook) + runner.query_start_loc.cpu = torch.tensor(offsets, dtype=torch.int32) + runner.inputs_embeds = DummyBuffer(torch.full((4, 4), -99.0)) + ids_before = runner.input_ids.gpu.clone() + embeds_before = runner.inputs_embeds.gpu.clone() + + with pytest.raises( + RuntimeError, + match="Non-MTP batched decode preprocessing requires contiguous in-bounds token offsets", + ): + runner._preprocess(scheduled, scheduled.total_num_scheduled_tokens) + + # These checks remain active when the runner is imported under python -O. + scalar_hook.assert_not_called() + batch_hook.assert_not_called() + torch.testing.assert_close(runner.input_ids.gpu, ids_before, rtol=0, atol=0) + torch.testing.assert_close(runner.inputs_embeds.gpu, embeds_before, rtol=0, atol=0) + + +@pytest.mark.parametrize("invalid", ["embeddings", "ids", "updates"]) +def test_decode_batch_without_mtp_rejects_incomplete_outputs(monkeypatch, invalid): + rows = [("first", 5, 5, 1), ("second", 5, 6, 1)] + runner, scheduled, _, _ = _make_phase_runner(monkeypatch, rows, batched_decode=False) + _without_mtp_buffers(runner) + runner.model = DecodeOnlyPreprocessModel(invalid_output=invalid) + before = runner.inputs_embeds.gpu.clone() + with pytest.raises(ValueError, match="decode preprocessing"): + runner._preprocess(scheduled, 2) + assert torch.equal(runner.inputs_embeds.gpu, before) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index c3871bace..048e340f3 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -194,7 +194,6 @@ _MODEL_INTERNALIZE = ( "MING_TTS_STAGE1_FINAL_LOG", "MOSS_TTS_DEBUG_STOP", "NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", - "VLLM_OMNI_BREEZE_GOLDEN_DIR", "VLLM_OMNI_QWEN_VISUAL_DEBUG_DIR", "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS", "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS_LOG_EVERY", diff --git a/vllm_omni/core/sched/omni_cfg_ar_scheduler.py b/vllm_omni/core/sched/omni_cfg_ar_scheduler.py new file mode 100644 index 000000000..f1a72ac81 --- /dev/null +++ b/vllm_omni/core/sched/omni_cfg_ar_scheduler.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Atomic CFG admission for full-attention AR models with unequal prompts. + +Each branch retains its real prompt length and position IDs. Admission reserves +capacity for the complete generation of both branches, so KV pressure cannot +preempt one branch while its partner continues. Reservations account for pages; +the ordinary paged attention manager still allocates pages as tokens arrive. +""" + +from collections.abc import Iterable +from typing import Any + +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import Request, RequestStatus + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.engine.serialization import deserialize_additional_information + + +class OmniCFGARScheduler(OmniARScheduler): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + if self.scheduler_config.async_scheduling or self.scheduler_config.enable_chunked_prefill: + raise ValueError("Atomic CFG scheduling requires synchronous scheduling and complete prefills") + if self.policy != SchedulingPolicy.FCFS or self.cache_config.enable_prefix_caching: + raise ValueError("Atomic CFG scheduling requires FCFS and disabled prefix caching") + if self.vllm_config.speculative_config is not None or self.connector is not None: + raise ValueError("Atomic CFG scheduling does not support speculative decoding or KV connectors") + if len(self.kv_cache_manager.kv_cache_config.kv_cache_groups) != 1: + raise ValueError("Atomic CFG reservations require one full-attention KV group") + self._capacity = self.kv_cache_manager.block_pool.num_gpu_blocks - 1 - self.kv_cache_manager.watermark_blocks + pair_pages = 2 * -(-self.max_model_len // self.block_size) + if self._capacity < pair_pages or self.max_num_running_reqs < 2: + raise ValueError("Atomic CFG requires KV capacity and max_num_seqs for two maximum-length branches") + if self.max_num_scheduled_tokens < 2 * self.max_model_len: + raise ValueError("Atomic CFG requires max_num_batched_tokens >= 2 * max_model_len") + self._pairs: dict[str, dict[str, str]] = {} + self._request_pairs: dict[str, str] = {} + self._prefill_min_output_tokens = int(self.vllm_config.additional_config.get("cfg_prefill_delay_tokens", 0)) + if self._prefill_min_output_tokens < 0: + raise ValueError("cfg_prefill_delay_tokens cannot be negative") + + def _should_defer_waiting_admission(self) -> bool: + # Let fresh streaming requests accumulate playback audio before + # another request's text/reference encoding occupies the same device. + return any(request.num_output_tokens < self._prefill_min_output_tokens for request in self.running) + + def add_request(self, request: Request) -> None: + payload = getattr(request, "additional_information", None) + info = deserialize_additional_information(payload) if payload is not None else {} + group = info.get("cfg_group") + if group is not None: + role = group["role"] + suffix = group["uncond_suffix"] + if role not in ("cond", "uncond") or not suffix: + raise ValueError("CFG groups require cond/uncond roles and a companion suffix") + external_id = request.external_req_id + if role == "uncond": + if not external_id.endswith(suffix): + raise ValueError("CFG companion request ID does not match its suffix") + pair_id = external_id[: -len(suffix)] + else: + pair_id = external_id + pair = self._pairs.setdefault(pair_id, {}) + if role in pair: + raise ValueError("Duplicate CFG branch") + pair[role] = request.request_id + self._request_pairs[request.request_id] = pair_id + super().add_request(request) + + def _reserved_pages(self, request: Request) -> int: + length = min(self.max_model_len, request.num_prompt_tokens + request.max_tokens) + return -(-length // self.block_size) + + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + self._drop_aborted_queued_requests() + self._prune_pairs() + original_waiting = self.waiting + admitted = create_request_queue(self.policy) + pending = {request.request_id: request for request in original_waiting} + pages = self._capacity - sum(self._reserved_pages(request) for request in self.running) + slots = self.max_num_running_reqs - len(self.running) + tokens = self.max_num_scheduled_tokens - len(self.running) + chosen: set[str] = set() + for request in original_waiting: + if request.request_id in chosen: + continue + pair_id = self._request_pairs.get(request.request_id) + if pair_id is None: + group = [request] + else: + pair = self._pairs[pair_id] + if len(pair) != 2 or any(rid not in pending for rid in pair.values()): + continue + group = [pending[pair[role]] for role in ("cond", "uncond")] + # Offline callers also need both unequal-length prompts to + # stop together when the longer branch reaches the context. + limit = min(min(member.max_tokens, self.max_model_len - member.num_prompt_tokens) for member in group) + for member in group: + member.max_tokens = limit + needed_pages = sum(self._reserved_pages(member) for member in group) + needed_tokens = sum(member.num_tokens - member.num_computed_tokens for member in group) + if len(group) > slots or needed_pages > pages or needed_tokens > tokens: + # Do not let later single-branch requests repeatedly occupy + # the spare slot and starve an earlier complete CFG pair. + break + for member in group: + admitted.add_request(member) + chosen.add(member.request_id) + pages -= needed_pages + tokens -= needed_tokens + slots -= len(group) + original_waiting.remove_requests([pending[rid] for rid in chosen]) + self.waiting = admitted + try: + output = super().schedule(throttle_prefills) + finally: + original_waiting.prepend_requests(self.waiting) + self.waiting = original_waiting + for pair in self._pairs.values(): + if len(pair) != 2: + continue + ids = list(pair.values()) + scheduled = [rid in output.num_scheduled_tokens for rid in ids] + if any(scheduled) != all(scheduled): + raise RuntimeError("CFG atomic admission invariant violated: only one branch was scheduled") + if all(scheduled): + progress = [ + self.requests[rid].num_computed_tokens - self.requests[rid].num_prompt_tokens for rid in ids + ] + if progress[0] != progress[1]: + raise RuntimeError("CFG branches reached different generation steps") + return output + + def finish_requests(self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus) -> list[Request]: + if request_ids is None: + ids = None + else: + ids = {request_ids} if isinstance(request_ids, str) else set(request_ids) + for rid in tuple(ids): + pair_id = self._request_pairs.get(rid) + if pair_id is not None: + ids.update(self._pairs[pair_id].values()) + finished = super().finish_requests(ids, finished_status) + self._prune_pairs() + return finished + + def _prune_pairs(self) -> None: + for pair_id, pair in list(self._pairs.items()): + if all(rid not in self.requests for rid in pair.values()): + del self._pairs[pair_id] + for rid in pair.values(): + self._request_pairs.pop(rid, None) diff --git a/vllm_omni/deploy/breeze_tts.yaml b/vllm_omni/deploy/breeze_tts.yaml new file mode 100644 index 000000000..35e43430c --- /dev/null +++ b/vllm_omni/deploy/breeze_tts.yaml @@ -0,0 +1,62 @@ +# Compatibility alias for breeze_tts_2.yaml; keep both profiles synchronized. +# Breeze-TTS-2, one GPU, voice design, cloning, and instruction guidance. +# Both streamed and complete HTTP responses use the stateful codec pipeline. +async_chunk: true +connectors: + shm: + name: SharedMemoryConnector + extra: + codec_streaming: true + codec_chunk_frames: 5 + initial_codec_chunk_frames: 1 + codec_chunk_ramp: [1, 2, 4, 5] + codec_left_context_frames: 0 + decode_cudagraph_batch_sizes: [1, 2] + connector_get_sleep_s: 0.001 + connector_get_max_wait_first_chunk: 1000 + connector_get_max_wait: 300 +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 2 + max_model_len: 2048 + max_num_batched_tokens: 4096 + gpu_memory_utilization: 0.65 + kv_cache_memory_bytes: 536870912 + enforce_eager: false + compilation_config: + mode: 3 + cudagraph_mode: FULL_AND_PIECEWISE + cudagraph_capture_sizes: [1, 2, 32, 64, 128, 256] + enable_prefix_caching: false + enable_chunked_prefill: false + async_scheduling: false + engine_extras: + additional_config: + cfg_prefill_delay_tokens: 8 + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 0.9 + top_p: 1.0 + top_k: 50 + repetition_penalty: 1.1 + max_tokens: 1500 + seed: 42 + - stage_id: 1 + devices: "0" + max_num_seqs: 2 + max_model_len: 32768 + max_num_batched_tokens: 32768 + gpu_memory_utilization: 0.15 + enforce_eager: false + enable_prefix_caching: false + async_scheduling: true + input_connectors: + from_stage_0: shm + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + max_tokens: 32768 diff --git a/vllm_omni/deploy/breeze_tts_2.yaml b/vllm_omni/deploy/breeze_tts_2.yaml index fd75dd69e..0f6561265 100644 --- a/vllm_omni/deploy/breeze_tts_2.yaml +++ b/vllm_omni/deploy/breeze_tts_2.yaml @@ -1,55 +1,61 @@ -# Breeze-TTS-2: talker -> stateful Qwen3-TTS codec decoder. -# Stage 0 batches codec frames and sends incremental codebook-major chunks. +# Breeze-TTS-2, one GPU, voice design, cloning, and instruction guidance. +# Both streamed and complete HTTP responses use the stateful codec pipeline. async_chunk: true -trust_remote_code: true - connectors: - connector_of_shared_memory: + shm: name: SharedMemoryConnector extra: codec_streaming: true - connector_get_sleep_s: 0.01 - connector_get_max_wait_first_chunk: 3000 + codec_chunk_frames: 5 + initial_codec_chunk_frames: 1 + codec_chunk_ramp: [1, 2, 4, 5] + codec_left_context_frames: 0 + decode_cudagraph_batch_sizes: [1, 2] + connector_get_sleep_s: 0.001 + connector_get_max_wait_first_chunk: 1000 connector_get_max_wait: 300 - breeze_codec_chunk_frames: 8 - stages: - stage_id: 0 - max_num_seqs: 8 - gpu_memory_utilization: 0.80 - enable_prefix_caching: false - # The talker derives one codec frame from the last hidden row of every - # scheduled request. A partially scheduled prompt would yield a frame from - # a mid-prompt hidden state, so keep each prefill in a single step. - enable_chunked_prefill: false - async_scheduling: true + devices: "0" + max_num_seqs: 2 + max_model_len: 2048 max_num_batched_tokens: 4096 - max_model_len: 4096 - devices: "0" + gpu_memory_utilization: 0.65 + kv_cache_memory_bytes: 536870912 + enforce_eager: false + compilation_config: + mode: 3 + cudagraph_mode: FULL_AND_PIECEWISE + cudagraph_capture_sizes: [1, 2, 32, 64, 128, 256] + enable_prefix_caching: false + enable_chunked_prefill: false + async_scheduling: false + engine_extras: + additional_config: + cfg_prefill_delay_tokens: 8 output_connectors: - to_stage_1: connector_of_shared_memory + to_stage_1: shm default_sampling_params: - temperature: 0.0 + temperature: 0.9 top_p: 1.0 - top_k: -1 - # Match upstream's penalty on generated codebook-0 history. + top_k: 50 repetition_penalty: 1.1 - max_tokens: 2048 - # min_tokens also masks the tokenizer's EOS 1 despite ignore_eos. - # Keep valid codec token 1 and an immediate backbone EOS available. - min_tokens: 0 - + max_tokens: 1500 + seed: 42 - stage_id: 1 - max_num_seqs: 8 + devices: "0" + max_num_seqs: 2 + max_model_len: 32768 + max_num_batched_tokens: 32768 gpu_memory_utilization: 0.15 - enforce_eager: true + enforce_eager: false enable_prefix_caching: false async_scheduling: true - max_num_batched_tokens: 65536 - max_model_len: 65536 - devices: "0" + input_connectors: + from_stage_0: shm default_sampling_params: temperature: 0.0 top_p: 1.0 top_k: -1 - max_tokens: 65536 + repetition_penalty: 1.0 + max_tokens: 32768 diff --git a/vllm_omni/deploy/breeze_tts_2_throughput.yaml b/vllm_omni/deploy/breeze_tts_2_throughput.yaml new file mode 100644 index 000000000..df432550d --- /dev/null +++ b/vllm_omni/deploy/breeze_tts_2_throughput.yaml @@ -0,0 +1,62 @@ +# Breeze-TTS-2 non-streaming throughput profile, tested on one H100 80 GB. +# Both streamed and complete HTTP responses use the stateful codec pipeline. +async_chunk: true +connectors: + shm: + name: SharedMemoryConnector + extra: + codec_streaming: true + codec_chunk_frames: 16 + initial_codec_chunk_frames: 16 + codec_chunk_ramp: [16] + codec_left_context_frames: 0 + decode_cudagraph_batch_sizes: [1, 2, 4, 8, 16, 32] + connector_get_sleep_s: 0.001 + connector_get_max_wait_first_chunk: 1000 + connector_get_max_wait: 300 +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 32 + max_model_len: 2048 + max_num_batched_tokens: 4096 + gpu_memory_utilization: 0.65 + # Reserve enough KV for thirty-two full-length requests. + kv_cache_memory_bytes: 8589934592 + enforce_eager: false + compilation_config: + mode: 3 + cudagraph_mode: FULL_AND_PIECEWISE + cudagraph_capture_sizes: [1, 2, 4, 8, 32, 64, 128, 256] + enable_prefix_caching: false + enable_chunked_prefill: false + async_scheduling: false + engine_extras: + additional_config: + cfg_prefill_delay_tokens: 0 + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 0.9 + top_p: 1.0 + top_k: 50 + repetition_penalty: 1.1 + max_tokens: 1500 + seed: 42 + - stage_id: 1 + devices: "0" + max_num_seqs: 32 + max_model_len: 32768 + max_num_batched_tokens: 32768 + gpu_memory_utilization: 0.15 + enforce_eager: false + enable_prefix_caching: false + async_scheduling: true + input_connectors: + from_stage_0: shm + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + max_tokens: 32768 diff --git a/vllm_omni/deploy/breeze_tts_throughput.yaml b/vllm_omni/deploy/breeze_tts_throughput.yaml new file mode 100644 index 000000000..3e0883d11 --- /dev/null +++ b/vllm_omni/deploy/breeze_tts_throughput.yaml @@ -0,0 +1,63 @@ +# Compatibility alias for breeze_tts_2_throughput.yaml; keep both profiles synchronized. +# Breeze-TTS-2 non-streaming throughput profile, tested on one H100 80 GB. +# Both streamed and complete HTTP responses use the stateful codec pipeline. +async_chunk: true +connectors: + shm: + name: SharedMemoryConnector + extra: + codec_streaming: true + codec_chunk_frames: 16 + initial_codec_chunk_frames: 16 + codec_chunk_ramp: [16] + codec_left_context_frames: 0 + decode_cudagraph_batch_sizes: [1, 2, 4, 8, 16, 32] + connector_get_sleep_s: 0.001 + connector_get_max_wait_first_chunk: 1000 + connector_get_max_wait: 300 +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 32 + max_model_len: 2048 + max_num_batched_tokens: 4096 + gpu_memory_utilization: 0.65 + # Reserve enough KV for thirty-two full-length requests. + kv_cache_memory_bytes: 8589934592 + enforce_eager: false + compilation_config: + mode: 3 + cudagraph_mode: FULL_AND_PIECEWISE + cudagraph_capture_sizes: [1, 2, 4, 8, 32, 64, 128, 256] + enable_prefix_caching: false + enable_chunked_prefill: false + async_scheduling: false + engine_extras: + additional_config: + cfg_prefill_delay_tokens: 0 + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 0.9 + top_p: 1.0 + top_k: 50 + repetition_penalty: 1.1 + max_tokens: 1500 + seed: 42 + - stage_id: 1 + devices: "0" + max_num_seqs: 32 + max_model_len: 32768 + max_num_batched_tokens: 32768 + gpu_memory_utilization: 0.15 + enforce_eager: false + enable_prefix_caching: false + async_scheduling: true + input_connectors: + from_stage_0: shm + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + max_tokens: 32768 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 4d44354d0..b9ed35de0 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -54,9 +54,7 @@ def _register_omni_hf_configs() -> None: try: from transformers import AutoConfig - from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( - BreezeTTS2Config, - ) + from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze import BreezeConfig from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import ( IndexTTS2Config, IndexTTS25Config, @@ -95,6 +93,7 @@ def _register_omni_hf_configs() -> None: _CONFIG_REGISTRY = None for model_type, config_cls in [ + ("breeze", BreezeConfig), ("dense", MingDenseConfig), ("bailingmm", MingMoeConfig), ("indextts2", IndexTTS2Config), @@ -108,7 +107,6 @@ def _register_omni_hf_configs() -> None: ("glm_tts", GLMTTSConfig), ("omnivoice", OmniVoiceConfig), ("voxcpm2", VoxCPM2Config), - ("breeze", BreezeTTS2Config), ]: try: AutoConfig.register(model_type, config_cls) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 7a3f4337e..438cf08c8 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -14,8 +14,9 @@ import asyncio import time import uuid from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast +import anyio from vllm import TokensPrompt from vllm.engine.protocol import EngineClient, StreamingInput from vllm.logger import init_logger @@ -75,6 +76,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient): ... print(output) """ + engine: AsyncOmniEngine + def _create_engine(self, **engine_kwargs: Any) -> AsyncOmniEngine: return AsyncOmniEngine(**engine_kwargs) @@ -265,7 +268,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): first_chunk_submitted = asyncio.get_running_loop().create_future() input_stream_task = await self._add_streaming_input_request( request_id=request_id, - input_stream=prompt, + input_stream=cast(AsyncGenerator, prompt), sampling_params_list=req_sp_list, final_stage_id=final_stage_id_for_e2e, final_output_stage_ids=final_output_stage_ids, @@ -285,7 +288,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient): lora_request=lora_request, ) submit_ts = time.time() - req_state.metrics.stage_first_ts[0] = submit_ts + stage_first_ts = cast(list[float | None], req_state.metrics.stage_first_ts) + stage_first_ts[0] = submit_ts req_start_ts[request_id] = submit_ts if admitting: await self._release_generate_admission() @@ -310,7 +314,10 @@ class AsyncOmni(AsyncOmniBase, EngineClient): except (asyncio.CancelledError, GeneratorExit): self._record_request_failure_once(request_id, reason="client_disconnect") - await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) + # ASGI cancellation also affects subsequent awaits. Shield the + # bounded engine abort before removing the local request state. + with anyio.CancelScope(shield=True): + await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: @@ -359,6 +366,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): # only check thinker's sampling params now stage0_params = sampling_params_list[0] self._validate_streaming_input_sampling_params(stage0_params) + stage0_params = cast(SamplingParams, stage0_params) req_state = self.request_states[request_id] has_submitted_first_chunk = False @@ -801,7 +809,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient): async with self._pause_cond: return self._paused - async def start_profile( + # EngineClient exposes async operations; OmniBase implements the sync API. + async def start_profile( # type: ignore[override] self, profile_prefix: str | None = None, stages: list[int] | None = None, @@ -816,7 +825,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): """ return await self.collective_rpc(method="profile", args=(True, profile_prefix), stage_ids=stages) - async def stop_profile(self, stages: list[int] | None = None) -> list[Any]: + async def stop_profile(self, stages: list[int] | None = None) -> list[Any]: # type: ignore[override] """Stop profiling specified stages. Uses vLLM-compatible profile(is_start=False) interface. @@ -1103,10 +1112,14 @@ class AsyncOmni(AsyncOmniBase, EngineClient): """Get tokenizer for the comprehension stage.""" stage_index = self._get_comprehension_stage_index() if stage_index is not None: - tokenizer = self.engine.output_processors[stage_index].tokenizer + processor = self.engine.output_processors[stage_index] + assert processor is not None + tokenizer = processor.tokenizer if tokenizer is not None: return tokenizer - return self.input_processor.tokenizer # type: ignore[return-value] + processor = self.input_processor + assert processor is not None + return processor.tokenizer async def is_tracing_enabled(self) -> bool: """Check if tracing is enabled.""" diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index 35572e303..c4c1bc491 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -71,6 +71,9 @@ class AsyncEventResolver: if tid is None and isinstance(ack, dict): tid = ack.get("task_id") + if tid is None: + logger.warning("Received ACK without a task_id") + return async with self._lock: task_info = self._pending_tasks.get(tid) @@ -382,6 +385,7 @@ class AsyncOmniBase(OmniBase): if should_continue: continue + assert req_state is not None req_state.stage_id = stage_id # Route to the per-request queue @@ -563,7 +567,8 @@ class AsyncOmniBase(OmniBase): # ==================== EngineClient Interface ==================== - async def check_health(self) -> None: + # The async entrypoints expose the synchronous base health check as an awaitable. + async def check_health(self) -> None: # type: ignore[override] """Check engine health by verifying the Orchestrator process is alive.""" OmniBase.check_health(self) diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index ca9e96bab..9a8391052 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -22,6 +22,7 @@ class ClientRequestState: self.final_stage_id: int | None = final_stage_id self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None + self.input_stream_task: asyncio.Task | None = None # Request-scoped idempotency guard for Prometheus failure counters. self.failure_recorded = False # Wall-clock time at which the user's request arrived in the engine diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index a32eeeac6..c044a67b6 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1807,6 +1807,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): finally: if not artifact_ready: self._discard_ref_audio_artifact_warmup(request_id) + # Disconnects can arrive while suspended at yield. Closing the + # engine iterator must survive the cancelled ASGI scope. close = getattr(generator, "aclose", None) if close is not None: with anyio.CancelScope(shield=True): diff --git a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py index 03f0d3c4b..673d43e5b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py @@ -1,262 +1,188 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Breeze-TTS-2 serving adapter for synchronous full-payload generation.""" - import asyncio -import copy import math -from typing import TYPE_CHECKING, Any import numpy as np +import pybase64 as base64 +from transformers import AutoTokenizer, PreTrainedTokenizerBase +from vllm.inputs import TokensPrompt +from vllm.logger import init_logger +from vllm.sampling_params import SamplingParams +from vllm.utils.async_utils import make_async +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, PreparedRequest, + SpeechServingContext, apply_max_new_tokens, - conditioning_cache_salt, -) -from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( - BreezeReferenceAudioTokenizer, -) -from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( - BreezeTTS2PromptBuilder, ) +from vllm_omni.model_executor.models.breeze_tts_2.prompt import DEFAULT_INSTRUCTION, build_breeze_prompt -if TYPE_CHECKING: - from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +logger = init_logger(__name__) @register_tts_adapter class BreezeTTS2Adapter(ARTTSAdapter): - """Build Breeze prompts before scheduler submission. - - The adapter owns only request-facing concerns. Model workers still own the - stage-0 embeddings and stage-1 waveform decode. - """ - - stage_keys = frozenset({"breeze_tts_2"}) - model_archs = frozenset({"BreezeForConditionalGeneration"}) name = "breeze_tts_2" - # Breeze's architecture and stage key are unique, so this detector cannot - # collide with another adapter. Lower priority runs first; 9 keeps it - # ahead of the default-priority (100) detectors. + stage_keys = frozenset({"breeze_tts_2", "breeze_tts_2_codec"}) + model_archs = frozenset({"BreezeForConditionalGeneration"}) detect_priority = 9 supported_output_sample_rates = frozenset({24000}) - def _load_supported_speakers(self) -> set[str]: - # Breeze accepts arbitrary speaker tags such as S0/S1 in its prompt; - # there is no finite checkpoint speaker inventory to advertise. - return set() + def __init__(self, ctx: SpeechServingContext) -> None: + super().__init__(ctx) + self.tokenizer: PreTrainedTokenizerBase | None = None + self._build_async = make_async(self._build_prompt, executor=ctx.server._tts_executor) - def _load_codec_frame_rate(self) -> float | None: - assert self.ctx.engine_client is not None - config = self.ctx.engine_client.model_config.hf_config - codec = getattr(config, "codec_config", None) - if isinstance(codec, dict): - frame_rate = codec.get("_frame_rate") - if frame_rate is not None: - return float(frame_rate) - return 12.5 + def _load_codec_frame_rate(self) -> float: + # Breeze-TTS-2 uses 1,920 waveform samples per frame at 24 kHz. + return 24000 / 1920 - def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - if not request.input or not request.input.strip(): - return "Breeze-TTS-2 input text cannot be empty" - if request.ref_text is not None and not request.ref_text.strip(): - return "Breeze-TTS-2 ref_text cannot be empty" - if request.ref_text is not None and request.ref_audio is None: - return "Breeze-TTS-2 ref_text requires ref_audio" - if request.ref_audio is not None and (request.ref_text is None or not request.ref_text.strip()): - return "Breeze-TTS-2 ref_audio requires ref_text" - if isinstance(request.ref_audio, list) and len(request.ref_audio) != 1: - # The prompt path conditions on exactly one reference waveform; - # extra clips would be silently discarded otherwise. - return f"Breeze-TTS-2 supports exactly one reference clip, got {len(request.ref_audio)}" - if request.task_type == "Base" and request.ref_audio is None: - return "Breeze-TTS-2 Base task requires ref_audio and ref_text" - if request.task_type == "VoiceDesign": - return "Breeze-TTS-2 does not support task_type=VoiceDesign" + async def warmup(self) -> None: + """Prime text, depth and reference paths through the ordinary runner.""" + server = self.ctx.server + logger.info("Warming Breeze text graphs, CFG branches, reference encoder and streaming codec") + + async def generate(text: str, request_id: str, **kwargs: object) -> bytes: + request = OpenAICreateSpeechRequest( + input=text, model=server.model_name, response_format="wav", seed=42, max_new_tokens=64, **kwargs + ) + audio, _ = await server._generate_audio_bytes(request, request_id=request_id) + return audio + + reference_text = "Welcome to this demonstration of clear and natural speech synthesis." + reference = await generate(reference_text, "breeze-warmup-single") + await generate( + "Hello.", + "breeze-warmup-cfg", + extra_params={"guidance_scale": 4.0, "temperature": 0.7, "top_k": 100, "top_p": 0.8}, + ) + await generate( + "Hello again.", + "breeze-warmup-reference", + ref_audio="data:audio/wav;base64," + base64.b64encode(reference).decode("ascii"), + ref_text=reference_text, + extra_params={"guidance_scale": 4.0}, + ) + await asyncio.gather( + generate(reference_text, "breeze-warmup-batch-a"), + generate("Thank you for listening to this clear and natural voice demonstration.", "breeze-warmup-batch-b"), + ) + logger.info("Breeze speech warmup complete") + + def validate(self, request: OpenAICreateSpeechRequest) -> str | None: + if not request.input.strip(): + return "Input text cannot be empty" if request.speed is not None and request.speed != 1.0: - return "Breeze-TTS-2 does not support speed adjustment" + return "Breeze does not support speed adjustment" if request.language is not None: - return "Breeze-TTS-2 does not support 'language'; the prompt language follows the input text" - if request.speaker_embedding is not None: - return "Breeze-TTS-2 does not support 'speaker_embedding' (voice-cloning helper field)" + return "Breeze infers the language from the input text; the language field is not supported" + if request.ref_audio_2 is not None or request.speaker_embedding is not None: + return "Breeze accepts one reference recording and its transcript" + has_reference = request.ref_audio is not None + if has_reference: + reference = request.ref_audio + if isinstance(reference, list): + if len(reference) != 1: + return "Breeze supports exactly one reference recording" + reference = reference[0] + if not isinstance(reference, str) or not isinstance(request.ref_text, str) or not request.ref_text.strip(): + return "Breeze voice cloning requires one ref_audio URL and a non-empty ref_text transcript" + error = self.ctx.server._validate_ref_audio_format(reference) + if error: + return error + elif request.ref_text is not None: + return "Breeze ref_text requires ref_audio" + expected_task = "Base" if has_reference else "VoiceDesign" + if request.task_type not in (None, "CustomVoice", expected_task): + return f"Breeze requires task_type='{expected_task}' for this conditioning" if request.x_vector_only_mode: - return "Breeze-TTS-2 does not support 'x_vector_only_mode' (voice-cloning helper field)" - extra_params = request.extra_params or {} - guidance_scale = extra_params.get("guidance_scale", extra_params.get("cfg_scale", 1.0)) - try: - unsupported_guidance = guidance_scale is not None and float(guidance_scale) != 1.0 - except (TypeError, ValueError): - return "Breeze-TTS-2 guidance_scale must be a number" - if unsupported_guidance: - return "Breeze-TTS-2 currently supports only guidance_scale=1.0" - if extra_params.get("negative_prompt"): - return "Breeze-TTS-2 does not support negative_prompt/CFG yet" - # The talker selects codebook 0 by argmax inside make_omni_output and - # ignores the scheduler's sampled token, so non-greedy overrides would - # not change the audio while still letting the sampler draw an early - # EOS. Reject them until the talker consumes the sampled code. - greedy_error = _validate_greedy_sampling(extra_params) - if greedy_error is not None: - return greedy_error - repetition_penalty = extra_params.get("repetition_penalty") - if repetition_penalty is not None: - try: - repetition_penalty = float(repetition_penalty) - except (TypeError, ValueError): - return "Breeze-TTS-2 repetition_penalty must be finite and positive" - if not math.isfinite(repetition_penalty) or repetition_penalty <= 0: - return "Breeze-TTS-2 repetition_penalty must be finite and positive" - if request.max_new_tokens is not None: - if request.max_new_tokens < self.max_new_tokens_min: - return f"max_new_tokens must be at least {self.max_new_tokens_min}" - if request.max_new_tokens > self.max_new_tokens_max: - return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" - if request.ref_audio is not None: - ref_audio = request.ref_audio[0] if isinstance(request.ref_audio, list) else request.ref_audio - return self.ctx.server._validate_ref_audio_format(ref_audio) + return "Breeze does not support x_vector_only_mode" + extra = request.extra_params or {} + supported = {"guidance_scale", "cfg_scale", "temperature", "top_k", "top_p", "repetition_penalty"} + if unknown := set(extra) - supported: + return f"Unsupported Breeze parameters: {sorted(unknown)}" + for key, value in extra.items(): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return f"Breeze {key} must be a finite number" + if extra.get("guidance_scale", extra.get("cfg_scale", 1.0)) <= 0 or extra.get("repetition_penalty", 1.1) <= 0: + return "Breeze guidance_scale and repetition_penalty must be positive" + if extra.get("temperature", 0.9) < 0 or not 0 < extra.get("top_p", 1.0) <= 1: + return "Breeze requires temperature >= 0 and 0 < top_p <= 1" + top_k = extra.get("top_k", 50) + if not isinstance(top_k, int) or top_k < -1: + return "Breeze top_k must be -1, 0, or a positive integer" return None + def _build_prompt( + self, + request: OpenAICreateSpeechRequest, + sampling: SamplingParams, + reference: tuple[np.ndarray, int] | None, + ) -> TokensPrompt: + if self.tokenizer is None: + engine_client = self.ctx.engine_client + if engine_client is None: + raise RuntimeError("Breeze speech serving requires an engine client") + model = engine_client.model_config.model + self.tokenizer = AutoTokenizer.from_pretrained(model, config=engine_client.model_config.hf_config) + extra = request.extra_params or {} + return build_breeze_prompt( + self.tokenizer, + request.input, + DEFAULT_INSTRUCTION if request.instructions is None else request.instructions, + speaker=request.voice if request.voice and request.voice != "default" else "S0", + ref_audio=reference, + ref_text=request.ref_text, + guidance_scale=float(extra.get("guidance_scale", extra.get("cfg_scale", 1.0))), + temperature=float(extra.get("temperature", sampling.temperature)), + top_k=max(int(extra.get("top_k", sampling.top_k)), 0), + top_p=float(extra.get("top_p", sampling.top_p)), + repetition_penalty=float(extra.get("repetition_penalty", sampling.repetition_penalty)), + ) + async def build( self, - request: "OpenAICreateSpeechRequest", - sampling_params_list: list, + request: OpenAICreateSpeechRequest, + sampling_params_list: list[SamplingParams], has_inline_ref_audio: bool, ) -> PreparedRequest: - del has_inline_ref_audio - server = self.ctx.server - builder = await self._get_builder() - template = self._resolve_template(request) - payload: dict[str, Any] = { - "text": request.input, - "template": template, - "speaker": request.voice or "S0", - } - if request.instructions: - payload["instruction"] = request.instructions - if request.ref_text: - payload["ref_text"] = request.ref_text - + reference = None if request.ref_audio is not None: - ref_audio = request.ref_audio[0] if isinstance(request.ref_audio, list) else request.ref_audio - wav_list, sr, cache_key = await server._resolve_ref_audio(ref_audio) - payload["ref_audio"] = np.asarray(wav_list, dtype=np.float32) - payload["ref_audio_sample_rate"] = int(sr) - else: - cache_key = None - - prompt = await asyncio.to_thread(builder.build, payload, template=template) - # The AR scheduler's completion budget is measured in codebook-0 - # frames for Breeze. Preserve it in the mutable request metadata so - # the talker can emit its own EOS before a scheduler length cutoff. - max_new_frames = request.max_new_tokens - if max_new_frames is None and sampling_params_list: - max_new_frames = getattr(sampling_params_list[0], "max_tokens", None) - if max_new_frames is not None: - prompt["additional_information"]["breeze_max_new_frames"] = int(max_new_frames) - tts_params = { - "template": [template], - "text": [request.input], - } - if cache_key: - tts_params["ref_audio_cache_key"] = cache_key - prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) - return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) - - async def _get_builder(self) -> BreezeTTS2PromptBuilder: - server = self.ctx.server - cached = getattr(server, "_breeze_tts_2_prompt_builder", None) - if cached is not None: - return cached - # Multiple speech requests can arrive while the first tokenizer load - # is still in progress. Serialize the one-time CPU initialization so - # model files are not opened and decoded repeatedly. - lock = getattr(server, "_breeze_tts_2_prompt_builder_lock", None) - if lock is None: - lock = asyncio.Lock() - server._breeze_tts_2_prompt_builder_lock = lock - async with lock: - cached = getattr(server, "_breeze_tts_2_prompt_builder", None) - if cached is not None: - return cached - model_path = server.engine_client.model_config.model - config = server.engine_client.model_config.hf_config - audio_tokenizer = await asyncio.to_thread( - BreezeReferenceAudioTokenizer.from_pretrained, - model_path, - num_codebooks=int(getattr(config, "num_codebooks", 16)), - codebook_size=self._codebook_size(config), - device_map="cpu", - ) - builder = await asyncio.to_thread( - BreezeTTS2PromptBuilder.from_pretrained, - model_path, - config, - reference_audio_encoder=audio_tokenizer, - ) - server._breeze_tts_2_prompt_builder = builder - return builder - - @staticmethod - def _codebook_size(config: Any) -> int: - codec = getattr(config, "codec_config", None) - if isinstance(codec, dict): - return int(codec.get("codebook_size", 2048)) - return int(getattr(codec, "codebook_size", 2048)) - - @staticmethod - def _resolve_template(request: "OpenAICreateSpeechRequest") -> str: - if request.ref_audio is not None: - return "ref_edit_tata" if request.instructions else "ref_clone_tata" - return "tts_instruction" if request.instructions else "tts_plain" + audio = request.ref_audio[0] if isinstance(request.ref_audio, list) else request.ref_audio + waveform, sample_rate, _ = await self.ctx.server._resolve_ref_audio(audio) + reference = (np.asarray(waveform, dtype=np.float32), sample_rate) + prompt = await self._build_async(request, sampling_params_list[0], reference) + return PreparedRequest(prompt=prompt, model_type=self.name) def apply_sampling_overrides( self, - sampling_params_list: list, - request: "OpenAICreateSpeechRequest", - prompt: dict[str, Any] | None = None, + sampling_params_list: list[SamplingParams], + request: OpenAICreateSpeechRequest, + prompt: dict | None = None, request_id: str | None = None, - ) -> list: - del prompt, request_id - sampling_params_list = apply_max_new_tokens(sampling_params_list, request) - repetition_penalty = (request.extra_params or {}).get("repetition_penalty") - if repetition_penalty is not None: - # The generic speech handler copies extra_params to extra_args; - # promote this value so the worker receives sampling metadata. - # Only this scalar changes; preserve the caller's params without - # copying nested state or the unchanged codec-stage parameters. - sampling_params_list = [copy.copy(sampling_params_list[0]), *sampling_params_list[1:]] - sampling_params_list[0].repetition_penalty = float(repetition_penalty) - return sampling_params_list - - -def _validate_greedy_sampling(extra_params: dict[str, Any]) -> str | None: - """Return an error unless the sampling overrides keep decoding greedy.""" - temperature = extra_params.get("temperature") - if temperature is not None: - try: - if float(temperature) != 0.0: - return "Breeze-TTS-2 currently supports only greedy decoding (temperature=0)" - except (TypeError, ValueError): - return "Breeze-TTS-2 temperature must be a number" - top_p = extra_params.get("top_p") - if top_p is not None: - try: - if float(top_p) != 1.0: - return "Breeze-TTS-2 currently supports only greedy decoding (top_p=1.0)" - except (TypeError, ValueError): - return "Breeze-TTS-2 top_p must be a number" - top_k = extra_params.get("top_k") - if top_k is not None: - try: - if int(top_k) not in (-1, 0, 1): - return "Breeze-TTS-2 currently supports only greedy decoding (top_k=-1)" - except (TypeError, ValueError): - return "Breeze-TTS-2 top_k must be an integer" - return None - - -__all__ = ["BreezeTTS2Adapter"] + ) -> list[SamplingParams]: + params = apply_max_new_tokens(sampling_params_list, request) + params = list(params) + params[0] = params[0].clone() + if params[0].top_k == 0: + params[0].top_k = -1 + extra = request.extra_params or {} + if "repetition_penalty" in extra: + params[0].repetition_penalty = extra["repetition_penalty"] + if prompt is not None: + engine_client = self.ctx.engine_client + if engine_client is None: + raise RuntimeError("Breeze speech serving requires an engine client") + model_config = engine_client.model_config + available = model_config.max_model_len - len(prompt["prompt_token_ids"]) + if available <= 0: + raise ValueError("Breeze conditioning fills the model context; shorten the text or reference audio") + # Both CFG branches must stop after the same number of frames, + # including when the conditioned prompt reaches the context limit. + params[0].max_tokens = min(params[0].max_tokens, available) + return params diff --git a/vllm_omni/model_executor/models/breeze_tts_2/__init__.py b/vllm_omni/model_executor/models/breeze_tts_2/__init__.py index e69de29bb..f50b655cc 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/__init__.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from .code2wav import BreezeCode2Wav +from .modeling_breeze import BreezeForConditionalGeneration + +__all__ = ["BreezeCode2Wav", "BreezeForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py b/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py deleted file mode 100644 index ea2796a53..000000000 --- a/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py +++ /dev/null @@ -1,149 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Reference-audio encoding helpers for Breeze-TTS-2. - -The Breeze prompt uses a Qwen3-TTS tokenizer only for reference audio. This -module deliberately does not contain text tokenization or waveform decoding; -the latter belongs to the pipeline's codec stage. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from pathlib import Path -from typing import Any - -import numpy as np -import torch - - -def resolve_audio_tokenizer_path(model_path: str) -> Path | None: - """Resolve the bundled ``audio_tokenizer`` dir for a local path or HF repo id. - - Local directories (dev checkouts) resolve directly: present returns the - subdirectory, absent returns ``None`` so callers can fall back. Repo ids - resolve through the vLLM-tagged HF helpers: an already-complete cached - snapshot is preferred (offline deployments keep working), and a partial - cache falls through to a targeted online fetch restricted to the bundled - tokenizer subtree. Returns ``None`` when the checkpoint carries no - ``audio_tokenizer`` subdirectory. - """ - root = Path(model_path) - if root.is_dir(): - bundled = root / "audio_tokenizer" - return bundled if bundled.is_dir() else None - - from vllm.transformers_utils.repo_utils import hf_api - - kwargs: dict[str, object] = {"allow_patterns": ["audio_tokenizer/*"]} - try: - cached = Path(hf_api().snapshot_download(model_path, local_files_only=True, **kwargs)) - except Exception: - cached = None - if cached is not None and (cached / "audio_tokenizer").is_dir(): - return cached / "audio_tokenizer" - online = Path(hf_api().snapshot_download(model_path, **kwargs)) - bundled = online / "audio_tokenizer" - return bundled if bundled.is_dir() else None - - -class BreezeReferenceAudioTokenizer: - """Small adapter that normalizes Qwen3-TTS encoder output to ``(T, Q)``. - - The underlying tokenizer is created once per serving worker and is reused - for all requests. Keeping this wrapper independent from the - prompt builder makes it straightforward to inject a fake encoder in tests - and to replace the upstream Qwen3 implementation later. - """ - - def __init__(self, tokenizer: Any, *, num_codebooks: int = 16, codebook_size: int = 2048) -> None: - self.tokenizer = tokenizer - self.num_codebooks = int(num_codebooks) - self.codebook_size = int(codebook_size) - if self.num_codebooks <= 0 or self.codebook_size <= 0: - raise ValueError("num_codebooks and codebook_size must be positive") - - @classmethod - def from_pretrained( - cls, - model_path: str, - *, - num_codebooks: int = 16, - codebook_size: int = 2048, - **kwargs: Any, - ) -> BreezeReferenceAudioTokenizer: - """Load the bundled ``audio_tokenizer`` with the existing Qwen wrapper.""" - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import ( - Qwen3TTSTokenizer, - ) - - resolved = resolve_audio_tokenizer_path(model_path) - if resolved is None: - raise FileNotFoundError( - "Breeze reference audio tokenizer was not found for " - f"{model_path} (expected an audio_tokenizer/ subdirectory)" - ) - tokenizer = Qwen3TTSTokenizer.from_pretrained(str(resolved), **kwargs) - return cls(tokenizer, num_codebooks=num_codebooks, codebook_size=codebook_size) - - @torch.inference_mode() - def encode(self, audio: Any, sample_rate: int | None = None) -> torch.Tensor: - """Encode one waveform or audio source into ``(T, Q)`` codes.""" - if isinstance(audio, torch.Tensor): - audio = audio.detach().cpu().numpy() - elif isinstance(audio, list) and audio and all(isinstance(x, (int, float)) for x in audio): - audio = np.asarray(audio, dtype=np.float32) - if isinstance(audio, np.ndarray) and audio.ndim > 1: - # Accept both common layouts: (samples, channels) and - # (channels, samples). Qwen3TTSTokenizer expects one 1-D clip. - channel_first = audio.shape[0] <= 8 and audio.shape[-1] > audio.shape[0] - audio = audio.mean(axis=0 if channel_first else -1) - - if isinstance(audio, np.ndarray) and sample_rate is None: - raise ValueError("sample_rate is required when ref_audio is a numpy waveform") - - encoded = self.tokenizer.encode(audio, sr=sample_rate) - return self._normalize_codes(encoded) - - def _normalize_codes(self, encoded: Any) -> torch.Tensor: - codes = getattr(encoded, "audio_codes", None) - if codes is None and isinstance(encoded, Mapping): - codes = encoded.get("audio_codes") - if codes is None and isinstance(encoded, (tuple, list)): - codes = encoded[0] if encoded else None - if codes is None: - raise ValueError("reference audio tokenizer returned no audio_codes") - - if isinstance(codes, (tuple, list)) and codes: - # Qwen3TTSTokenizer returns one tensor per input waveform. - if len(codes) != 1: - raise ValueError(f"only one reference waveform is supported, got {len(codes)}") - codes = codes[0] - codes = torch.as_tensor(codes) - if codes.ndim == 3: - if codes.shape[0] != 1: - raise ValueError(f"only one reference waveform is supported, got {tuple(codes.shape)}") - codes = codes[0] - if codes.ndim != 2: - raise ValueError(f"reference audio codes must be 2D, got {tuple(codes.shape)}") - if codes.shape[0] == 0: - raise ValueError("reference audio tokenizer returned no frames") - if codes.shape[-1] != self.num_codebooks: - if codes.shape[0] == self.num_codebooks: - codes = codes.transpose(0, 1) - else: - raise ValueError( - f"reference audio codes must have {self.num_codebooks} codebooks, got {tuple(codes.shape)}" - ) - codes = codes.to(device="cpu", dtype=torch.long).contiguous() - if codes.numel() and (int(codes.min()) < 0 or int(codes.max()) >= self.codebook_size): - raise ValueError( - f"reference audio code outside [0, {self.codebook_size}): " - f"min={int(codes.min())}, max={int(codes.max())}" - ) - # int16 is enough for Mimi/Qwen3-TTS codes and halves IPC/cache memory. - return codes.to(dtype=torch.int16) - - -__all__ = ["BreezeReferenceAudioTokenizer", "resolve_audio_tokenizer_path"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/code2wav.py b/vllm_omni/model_executor/models/breeze_tts_2/code2wav.py new file mode 100644 index 000000000..74b862364 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/code2wav.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import Qwen3TTSCode2Wav + + +class BreezeCode2Wav(Qwen3TTSCode2Wav): + """Breeze bundles the same Qwen3 codec under audio_tokenizer/.""" + + tokenizer_subfolder = "audio_tokenizer" + # Reference conditioning is consumed by the AR stage. The codec receives + # generated codes only, so it never decodes an ICL reference prefix. + decoder_cudagraph_modes = ("xvec",) diff --git a/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze.py b/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze.py new file mode 100644 index 000000000..d889262c3 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Configuration for the released Breeze-TTS-2 checkpoint.""" + +from transformers import LlamaConfig, PretrainedConfig, Qwen3Config +from transformers.models.t5gemma2.configuration_t5gemma2 import T5Gemma2TextConfig + + +class BreezeConfig(PretrainedConfig): + model_type = "breeze" + + def __init__( + self, + backbone_config: dict | None = None, + depth_decoder_config: dict | None = None, + text_encoder_config: dict | None = None, + num_codebooks: int = 16, + audio_vocab_size: int = 2051, + audio_embed_size: int = 2048, + max_position_embeddings: int = 2048, + rope_scaling: dict | None = None, + rope_theta: float = 500000.0, + **kwargs: object, + ) -> None: + self.max_position_embeddings = max_position_embeddings + self.rope_parameters = {"rope_theta": rope_theta, **(rope_scaling or {"rope_type": "default"})} + self.num_codebooks = num_codebooks + self.audio_vocab_size = audio_vocab_size + self.audio_embed_size = audio_embed_size + self.backbone_config = Qwen3Config(**(backbone_config or {})) + self.depth_decoder_config = LlamaConfig(**(depth_decoder_config or {})) + self.text_encoder_config = T5Gemma2TextConfig(**(text_encoder_config or {})) + # The scheduler samples continuation/EOS tokens. Text encoder IDs are + # carried in the prompt payload and are outside this vocabulary. + self.vocab_size = audio_vocab_size + 1 + self.eos_token_id = audio_vocab_size + self.backbone_config.vocab_size = self.vocab_size + self.backbone_config.bos_token_id = 0 + self.backbone_config.pad_token_id = 0 + self.backbone_config.eos_token_id = self.eos_token_id + self.backbone_config.tie_word_embeddings = False + super().__init__(**kwargs) + self.vocab_size = audio_vocab_size + 1 + self.eos_token_id = audio_vocab_size + + def get_text_config(self, decoder: bool = False) -> Qwen3Config: + return self.backbone_config diff --git a/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py b/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py deleted file mode 100644 index 2c04768a9..000000000 --- a/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py +++ /dev/null @@ -1,366 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Configuration classes for Breeze-TTS-2.""" - -from __future__ import annotations - -from typing import Any - -from transformers import AutoConfig -from transformers.configuration_utils import PretrainedConfig -from transformers.models.qwen3 import Qwen3Config - - -class BreezeTTS2Config(PretrainedConfig): - """Configuration for ``BreezeForConditionalGeneration``.""" - - model_type = "breeze" - - def __init__( - self, - backbone_config: dict | Qwen3Config | None = None, - text_encoder_config: dict | BreezeTTS2TextEncoderConfig | None = None, - depth_decoder_config: dict | BreezeTTS2DepthDecoderConfig | None = None, - codec_config: dict | None = None, - architectures: list[str] | None = None, - attention_bias: bool = False, - attention_dropout: float = 0.0, - audio_embed_size: int = 2048, - audio_eos_token_id: int = 262145, - audio_num_codebooks: int = 16, - audio_token_id: int = 262144, - audio_vocab_size: int = 2051, - backbone_flavor: str = "llama-1B", - backbone_model_name_or_path: str | None = None, - backbone_model_type: str = "qwen3", - bos_token_id: int = 2, - codebook_eos_token_id: int = 0, - codebook_pad_token_id: int = 2050, - decoder_flavor: str = "llama-100M", - depth_header_loss_weight: float = 1.0, - dtype: str = "bfloat16", - eos_token_id: int = 1, - head_dim: int = 128, - hidden_act: str = "silu", - hidden_size: int = 2048, - initializer_range: float = 0.02, - intermediate_size: int = 6144, - max_position_embeddings: int = 2048, - mlp_bias: bool = False, - num_attention_heads: int = 16, - num_codebooks: int = 16, - num_hidden_layers: int = 28, - num_key_value_heads: int = 8, - pad_token_id: int = 0, - rms_norm_eps: float = 1e-5, - rope_scaling: dict | None = None, - rope_theta: float = 500_000.0, - text_encoder_bucket_max_length_ratio: float = 4.0, - text_encoder_lora_config: dict | None = None, - text_encoder_proj_type: str = "linear", - text_encoder_special_tokens_config: dict | None = None, - text_vocab_size: int = 262158, - tie_codebooks_embeddings: bool = True, - tie_word_embeddings: bool = False, - use_cache: bool = True, - vocab_size: int = 2051, - sampling_rate: int = 24000, - codec_model_name_or_path: str = "kyutai/mimi", - **kwargs: Any, - ) -> None: - serialized_rope_parameters = kwargs.pop("rope_parameters", None) - if rope_scaling is None and serialized_rope_parameters is not None: - rope_scaling = serialized_rope_parameters - # Breeze stores the Qwen3 backbone in backbone_config. This follows the - # MOSS-TTS pattern and deliberately does not declare sub_configs. - if backbone_config is None: - backbone_config = {} - if isinstance(backbone_config, dict): - backbone_config = dict(backbone_config) - backbone_config.pop("model_type", None) - self.backbone_config = Qwen3Config(**backbone_config) - else: - self.backbone_config = backbone_config - - if text_encoder_config is None: - text_encoder_config = {} - if isinstance(text_encoder_config, dict): - text_encoder_config = dict(text_encoder_config) - text_encoder_config.pop("model_type", None) - self.text_encoder_config = BreezeTTS2TextEncoderConfig(**text_encoder_config) - else: - self.text_encoder_config = text_encoder_config - - if depth_decoder_config is None: - depth_decoder_config = {} - if isinstance(depth_decoder_config, dict): - depth_decoder_config = dict(depth_decoder_config) - depth_decoder_config.pop("model_type", None) - self.depth_decoder_config = BreezeTTS2DepthDecoderConfig(**depth_decoder_config) - else: - self.depth_decoder_config = depth_decoder_config - - # Nested configs must exist before the parent constructor: Transformers - # 5.x may call get_text_config() during config validation. - super().__init__( - architectures=architectures, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - pad_token_id=pad_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.audio_embed_size = audio_embed_size - self.audio_eos_token_id = audio_eos_token_id - self.audio_num_codebooks = audio_num_codebooks - self.audio_token_id = audio_token_id - self.audio_vocab_size = audio_vocab_size - # Keep Mimi's nested config as a plain checkpoint dictionary. Stage 1 - # owns codec construction; stage 0 only needs codebook_size here. - self.codec_config = codec_config - self.backbone_flavor = backbone_flavor - self.backbone_model_name_or_path = backbone_model_name_or_path - self.backbone_model_type = backbone_model_type - self.codebook_eos_token_id = codebook_eos_token_id - self.codebook_pad_token_id = codebook_pad_token_id - self.decoder_flavor = decoder_flavor - self.depth_header_loss_weight = depth_header_loss_weight - self.dtype = dtype - self.head_dim = head_dim - self.hidden_act = hidden_act - self.hidden_size = hidden_size - self.initializer_range = initializer_range - self.intermediate_size = intermediate_size - self.max_position_embeddings = max_position_embeddings - self.mlp_bias = mlp_bias - self.num_attention_heads = num_attention_heads - self.num_codebooks = num_codebooks - self.num_hidden_layers = num_hidden_layers - self.num_key_value_heads = num_key_value_heads - self.rms_norm_eps = rms_norm_eps - if rope_scaling is not None: - # Preserve the upstream legacy field while also satisfying the - # normalized ``rope_parameters`` validation added in Transformers 5. - rope_scaling = dict(rope_scaling) - rope_scaling.setdefault("rope_type", rope_scaling.get("type", "default")) - rope_scaling.setdefault("rope_theta", rope_theta) - self.rope_scaling = rope_scaling - self.rope_theta = rope_theta - self.text_encoder_bucket_max_length_ratio = text_encoder_bucket_max_length_ratio - self.text_encoder_lora_config = text_encoder_lora_config or {} - self.text_encoder_proj_type = text_encoder_proj_type - self.text_encoder_special_tokens_config = text_encoder_special_tokens_config or {} - self.text_vocab_size = text_vocab_size - self.tie_codebooks_embeddings = tie_codebooks_embeddings - self.use_cache = use_cache - self.vocab_size = vocab_size - self.sampling_rate = sampling_rate - self.codec_model_name_or_path = codec_model_name_or_path - - if not getattr(self, "architectures", None): - self.architectures = ["BreezeForConditionalGeneration"] - - def get_text_config(self, **_: Any) -> Qwen3Config: - return self.backbone_config - - -class BreezeTTS2TextEncoderConfig(PretrainedConfig): - """Configuration for Breeze's T5Gemma2 text encoder.""" - - model_type = "t5gemma2_text" - - def __init__( - self, - hidden_size: int = 1152, - intermediate_size: int = 6912, - num_hidden_layers: int = 26, - num_attention_heads: int = 4, - num_key_value_heads: int = 1, - head_dim: int = 256, - max_position_embeddings: int = 32768, - hidden_activation: str = "gelu_pytorch_tanh", - rms_norm_eps: float = 1e-6, - rope_theta: float = 10_000.0, - sliding_window: int = 512, - query_pre_attn_scalar: int = 256, - eoi_token_index: int = 256000, - vocab_size: int = 262158, - layer_types: list[str] | None = None, - rope_parameters: dict | None = None, - attention_bias: bool = False, - attention_dropout: float = 0.0, - dropout_rate: float = 0.0, - initializer_range: float = 0.02, - tie_word_embeddings: bool = True, - use_bidirectional_attention: bool = False, - use_cache: bool = True, - requires_grad: bool = False, - preferred_attn_implementation: str = "flash_attention_2", - dtype: str = "bfloat16", - is_encoder_decoder: bool = False, - bos_token_id: int = 2, - eos_token_id: int = 1, - pad_token_id: int = 0, - **kwargs: Any, - ) -> None: - kwargs.pop("model_type", None) - - super().__init__( - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - pad_token_id=pad_token_id, - is_encoder_decoder=is_encoder_decoder, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.max_position_embeddings = max_position_embeddings - self.hidden_activation = hidden_activation - self.hidden_act = hidden_activation - self.rms_norm_eps = rms_norm_eps - self.rope_theta = rope_theta - self.sliding_window = sliding_window - self.query_pre_attn_scalar = query_pre_attn_scalar - self.eoi_token_index = eoi_token_index - self.vocab_size = vocab_size - if layer_types is None: - # T5Gemma2 uses sliding attention except every sixth layer. - layer_types = [ - "full_attention" if (index + 1) % 6 == 0 else "sliding_attention" for index in range(num_hidden_layers) - ] - self.layer_types = list(layer_types) - if rope_parameters is None: - rope_parameters = { - "full_attention": { - "factor": 8.0, - "rope_theta": 1_000_000.0, - "rope_type": "linear", - }, - "sliding_attention": { - "rope_theta": 10_000.0, - "rope_type": "default", - }, - } - self.rope_parameters = dict(rope_parameters) - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.dropout_rate = dropout_rate - self.initializer_range = initializer_range - self.use_bidirectional_attention = use_bidirectional_attention - self.use_cache = use_cache - self.requires_grad = requires_grad - self.preferred_attn_implementation = preferred_attn_implementation - self.dtype = dtype - - -class BreezeTTS2DepthDecoderConfig(PretrainedConfig): - """Configuration for Breeze's per-frame depth decoder.""" - - model_type = "breeze_depth_decoder_model" - - def __init__( - self, - hidden_size: int = 1024, - backbone_hidden_size: int = 2048, - audio_embed_size: int = 2048, - intermediate_size: int = 8192, - num_hidden_layers: int = 12, - num_attention_heads: int = 8, - num_key_value_heads: int = 2, - head_dim: int = 128, - max_position_embeddings: int = 33, - hidden_act: str = "silu", - rms_norm_eps: float = 1e-5, - rope_theta: float = 500_000.0, - rope_scaling: dict | None = None, - attention_bias: bool = False, - attention_dropout: float = 0.0, - mlp_bias: bool = False, - initializer_range: float = 0.02, - use_cache: bool = True, - num_codebooks: int = 16, - vocab_size: int = 2051, - codebook_loss_weights: list[int] | None = None, - dtype: str = "bfloat16", - tie_word_embeddings: bool = False, - pad_token_id: int | None = 0, - **kwargs: Any, - ) -> None: - kwargs.pop("model_type", None) - # Transformers 5 serializes the normalized field as - # ``rope_parameters`` while the upstream Breeze checkpoint calls it - # ``rope_scaling``. Normalize before ``PretrainedConfig`` validates - # so ``to_dict() -> BreezeTTS2DepthDecoderConfig(**dict)`` round-trips. - serialized_rope_parameters = kwargs.pop("rope_parameters", None) - if rope_scaling is None: - rope_scaling = { - "factor": 32.0, - "high_freq_factor": 0.0078125, - "low_freq_factor": 0.001953125, - "original_max_position_embeddings": 16, - "rope_type": "llama3", - } - if codebook_loss_weights is None: - codebook_loss_weights = [3] - if num_codebooks <= 0 or vocab_size <= 0: - raise ValueError("num_codebooks and vocab_size must be positive") - rope_parameters = dict(serialized_rope_parameters or rope_scaling) - rope_parameters.setdefault("rope_type", rope_parameters.get("type", "default")) - rope_parameters.setdefault("rope_theta", rope_theta) - - super().__init__( - vocab_size=vocab_size, - tie_word_embeddings=tie_word_embeddings, - pad_token_id=pad_token_id, - **kwargs, - ) - # The vLLM depth decoder is a plain ``nn.Module``, so it does not get - # Transformers' PreTrainedModel default attention initialization. - self._attn_implementation = getattr(self, "_attn_implementation", None) or "eager" - self.hidden_size = hidden_size - self.backbone_hidden_size = backbone_hidden_size - self.audio_embed_size = audio_embed_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.max_position_embeddings = max_position_embeddings - self.hidden_act = hidden_act - self.rms_norm_eps = rms_norm_eps - self.rope_theta = rope_theta - self.rope_parameters = dict(getattr(self, "rope_parameters", None) or rope_parameters) - self.rope_parameters.setdefault("rope_theta", rope_theta) - # Transformers 4 exposes this as a plain attribute; Transformers 5 - # implements ``rope_scaling`` as a property backed by rope_parameters. - if not hasattr(type(self), "rope_scaling"): - self.rope_scaling = self.rope_parameters - self.rope_type = rope_scaling.get("rope_type", "default") - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.mlp_bias = mlp_bias - self.initializer_range = initializer_range - self.use_cache = use_cache - self.num_codebooks = num_codebooks - self.vocab_size = vocab_size - self.codebook_loss_weights = list(codebook_loss_weights) - self.dtype = dtype - - -AutoConfig.register("breeze", BreezeTTS2Config, exist_ok=True) - - -__all__ = [ - "BreezeTTS2Config", - "BreezeTTS2DepthDecoderConfig", - "BreezeTTS2TextEncoderConfig", -] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py b/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py new file mode 100644 index 000000000..70b7262f1 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Breeze's within-frame transformer, with frame-local KV state. + +Attention follows the reference BF16 matmul / FP32 softmax ordering. Each +frame overwrites its causal cache before reading it. Captured workspaces +belong to their graphs; generated frames and RNG state belong to requests. +""" + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig +from vllm.platforms import current_platform + +from vllm_omni.platforms import current_omni_platform + + +def sample_logits( + logits: torch.Tensor, + temperature: float, + top_k: int, + top_p: float, + generator: torch.Generator, + noise: torch.Tensor | None = None, +) -> torch.Tensor: + """Sample using scalar parameters, optionally with pre-drawn noise. + + Keep the filtering contract aligned with ``sample_graph_logits``: scale + in FP32, retain top-k ties, and remove nucleus candidates only when the + cumulative mass before them is >= top_p. Temperature zero selects argmax + without drawing RNG. The parity tests cover ties and exact boundaries. + """ + if temperature == 0: + return logits.argmax(-1) + scores = logits.float() / temperature + if top_k > 0: + threshold = scores.topk(min(top_k, scores.shape[-1]), dim=-1).values[..., -1:] + scores = scores.masked_fill(scores < threshold, -torch.inf) + if top_p < 1: + values, indices = scores.sort(dim=-1, descending=True) + probabilities = values.softmax(-1) + cumulative = probabilities.cumsum(-1) + # Keep the first candidate crossing the nucleus boundary as in HF. + remove = cumulative - probabilities >= top_p + scores = scores.scatter(-1, indices, values.masked_fill(remove, -torch.inf)) + probabilities = scores.softmax(-1) + if noise is not None: + return (probabilities / noise).argmax(-1) + return torch.multinomial(probabilities, 1, generator=generator).squeeze(-1) + + +def sample_graph_logits(logits: torch.Tensor, parameters: torch.Tensor, noise: torch.Tensor) -> torch.Tensor: + """Sample with mutable GPU parameters and a request-owned RNG draw. + + Sorting once supplies both the top-k threshold and the nucleus order. + Filtering keeps ties at the top-k threshold, and keeps the first token + crossing the nucleus boundary. The three reserved codec IDs remain in + the RNG workspace, with zero probability, as in the reference sampler. + + Replay can change parameters, so selection uses tensor operations rather + than ``sample_logits``' Python branches on scalar parameters. Keep their + FP32 scaling and top-k/top-p boundaries aligned without sharing that + control flow. This path consumes supplied noise; its caller owns RNG + advancement and must avoid advancing request RNG for greedy decoding. + """ + temperature, top_k, top_p = parameters.unbind() + scaled = logits.float() / torch.where(temperature > 0, temperature, 1.0) + values, indices = scaled.sort(dim=-1, descending=True) + cutoff = torch.where(top_k > 0, top_k, logits.shape[-1]).long().clamp(1, logits.shape[-1]) - 1 + threshold = values.gather(-1, cutoff.expand(values.shape[0], 1)) + values = values.masked_fill(values < threshold, -torch.inf) + probabilities = values.softmax(-1) + cumulative = probabilities.cumsum(-1) + # Use the same boundary and arithmetic as the eager sampler, including + # an exactly representable cumulative probability equal to top_p. + remove = (cumulative - probabilities >= top_p) & (top_p < 1) + values = values.masked_fill(remove, -torch.inf) + # Argmax of p / Exp(1) is the single-sample multinomial algorithm. + selected = (values.softmax(-1) / noise.gather(-1, indices)).argmax(-1, keepdim=True) + sampled = indices.gather(-1, selected).squeeze(-1) + return torch.where(temperature > 0, sampled, logits.argmax(-1)) + + +class BreezeRMSNorm(nn.Module): + def __init__(self, size: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(size)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x.float() + y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps) + return self.weight * y.to(x.dtype) + + +class BreezeDepthLayer(nn.Module): + def __init__(self, config: PretrainedConfig) -> None: + super().__init__() + self.heads = config.num_attention_heads + self.kv_heads = config.num_key_value_heads + self.dim = config.head_dim + self.q_size = self.heads * self.dim + self.kv_size = self.kv_heads * self.dim + self.qkv = nn.Linear(config.hidden_size, self.q_size + 2 * self.kv_size, bias=False) + self.o_proj = nn.Linear(self.q_size, config.hidden_size, bias=False) + self.gate_up = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + self.input_layernorm = BreezeRMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = BreezeRMSNorm(config.hidden_size, config.rms_norm_eps) + + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + cache: tuple[torch.Tensor, torch.Tensor], + mask: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + batch, length, _ = x.shape + residual = x + q, k, v = self.qkv(self.input_layernorm(x)).split([self.q_size, self.kv_size, self.kv_size], -1) + q = q.view(batch, length, self.heads, self.dim).transpose(1, 2) + k = k.view(batch, length, self.kv_heads, self.dim).transpose(1, 2) + v = v.view(batch, length, self.kv_heads, self.dim).transpose(1, 2) + q1, q2 = q.chunk(2, -1) + k1, k2 = k.chunk(2, -1) + q = q * cos + torch.cat((-q2, q1), -1) * sin + k = k * cos + torch.cat((-k2, k1), -1) * sin + cache[0].index_copy_(2, positions, k) + cache[1].index_copy_(2, positions, v) + # Fold query heads sharing one KV head into its query dimension. + # This avoids materializing repeated keys/values for grouped-query + # attention while preserving BF16 matmul and FP32 softmax ordering. + grouped_q = q.reshape(batch, self.kv_heads, self.heads // self.kv_heads * length, self.dim) + scores = ( + torch.matmul(grouped_q, cache[0].transpose(-1, -2)).reshape(batch, self.heads, length, cache[0].shape[2]) + * self.dim**-0.5 + ) + scores = scores.masked_fill(~mask, torch.finfo(q.dtype).min) + probabilities = scores.softmax(-1, dtype=torch.float32).to(q.dtype) + grouped_probabilities = probabilities.reshape(batch, self.kv_heads, -1, cache[0].shape[2]) + y = torch.matmul(grouped_probabilities, cache[1]).reshape(batch, self.heads, length, self.dim) + y = y.transpose(1, 2).reshape(batch, length, self.q_size) + x = residual + self.o_proj(y) + gate, up = self.gate_up(self.post_attention_layernorm(x)).chunk(2, -1) + return x + self.down_proj(F.silu(gate) * up) + + +class BreezeDepthDecoder(nn.Module): + def __init__(self, config: PretrainedConfig) -> None: + super().__init__() + self.config = config + self.num_codebooks = config.num_codebooks + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(self.num_codebooks * self.vocab_size, config.audio_embed_size) + self.inputs_embeds_projector = nn.Linear(config.audio_embed_size, config.hidden_size, bias=False) + self.layers = nn.ModuleList([BreezeDepthLayer(config) for _ in range(config.num_hidden_layers)]) + self.norm = BreezeRMSNorm(config.hidden_size, config.rms_norm_eps) + self.codebooks_head = nn.Parameter( + torch.empty(self.num_codebooks - 1, config.hidden_size, self.vocab_size, dtype=torch.float32) + ) + rope = config.rope_parameters + inv = 1.0 / (rope["rope_theta"] ** (torch.arange(0, config.head_dim, 2).float() / config.head_dim)) + wavelength = 2 * torch.pi / inv + low = rope["original_max_position_embeddings"] / rope["low_freq_factor"] + high = rope["original_max_position_embeddings"] / rope["high_freq_factor"] + smooth = (rope["original_max_position_embeddings"] / wavelength - rope["low_freq_factor"]) / ( + rope["high_freq_factor"] - rope["low_freq_factor"] + ) + scaled = (1 - smooth) * inv / rope["factor"] + smooth * inv + inv = torch.where(wavelength > low, inv / rope["factor"], torch.where(wavelength < high, inv, scaled)) + angles = torch.outer(torch.arange(self.num_codebooks).float(), inv) + angles = torch.cat((angles, angles), -1) + self.register_buffer("rope_cos", angles.cos(), persistent=False) + self.register_buffer("rope_sin", angles.sin(), persistent=False) + self.register_buffer("positions", torch.arange(self.num_codebooks), persistent=False) + self.register_buffer( + "causal_mask", + torch.arange(self.num_codebooks)[None, :] <= torch.arange(self.num_codebooks)[:, None], + persistent=False, + ) + self._graphs: dict[tuple[int, int, bool], BreezeDepthGraph] = {} + self._compiled_layer = None + self._compiled_sampler = None + + def _allocate_cache(self, hidden: torch.Tensor) -> list[tuple[torch.Tensor, torch.Tensor]]: + shape = (hidden.shape[0], self.config.num_key_value_heads, self.num_codebooks, self.config.head_dim) + # Masked future positions still participate in the value matmul. + # Initialize them to finite values so zero probabilities cannot meet NaN. + return [(hidden.new_zeros(shape), hidden.new_zeros(shape)) for _ in self.layers] + + def generate_frame( + self, + hidden: torch.Tensor, + first: torch.Tensor, + *, + temperature: float, + top_k: int, + top_p: float, + generator: torch.Generator, + ) -> torch.Tensor: + return self._generate_frame(hidden, first, self._allocate_cache(hidden), temperature, top_k, top_p, generator) + + def _generate_frame( + self, + hidden: torch.Tensor, + first: torch.Tensor, + caches: list[tuple[torch.Tensor, torch.Tensor]], + temperature: float, + top_k: int, + top_p: float, + generator: torch.Generator, + noise: torch.Tensor | None = None, + guidance_scale: float | torch.Tensor | None = None, + sampling_parameters: torch.Tensor | None = None, + greedy: bool = False, + ) -> torch.Tensor: + batch = hidden.shape[0] + branches = 2 if guidance_scale is not None else 1 + branch_first = first.repeat(branches) + x = torch.cat((hidden.unsqueeze(1), self.embed_tokens(branch_first.reshape(batch, 1))), 1) + frames = [first] + start = 0 + for codebook in range(self.num_codebooks - 1): + x = self.inputs_embeds_projector(x) + end = start + x.shape[1] + cos = self.rope_cos[start:end].to(x.dtype)[None, None] + sin = self.rope_sin[start:end].to(x.dtype)[None, None] + for layer, cache in zip(self.layers, caches, strict=True): + args = (x, self.positions[start:end], cache, self.causal_mask[start:end][None, None], cos, sin) + if noise is None and not greedy: + x = layer(*args) + else: + x = self._compiled_layer(layer, *args) + logits = F.linear(self.norm(x[:, -1]).float(), self.codebooks_head[codebook].T.float()) + if guidance_scale is not None: + cond, uncond = logits.chunk(2, dim=0) + logits = uncond + guidance_scale * (cond - uncond) + logits[:, self.vocab_size - 3 :] = -torch.inf + if greedy: + token = logits.argmax(-1) + elif noise is None: + token = sample_logits(logits, temperature, top_k, top_p, generator) + else: + token = self._compiled_sampler(logits, sampling_parameters, noise[:, codebook]) + frames.append(token) + x = self.embed_tokens((token.repeat(branches) + (codebook + 1) * self.vocab_size).reshape(batch, 1)) + start = end + return torch.stack(frames, -1) + + def generate_frames( + self, + hidden: torch.Tensor, + first: torch.Tensor, + *, + temperature: float, + top_k: int, + top_p: float, + generators: list[torch.Generator], + guidance_scale: float = 1.0, + ) -> torch.Tensor: + batch = hidden.shape[0] + branches = 2 if guidance_scale != 1.0 else 1 + requests = first.shape[0] + if batch != requests * branches or len(generators) != requests: + raise ValueError("Depth generation requires one token and generator per request and aligned CFG branches") + if requests == 0: + return first.new_empty((0, self.num_codebooks)) + # Buckets bound captured workspace shapes as live requests arrive and + # finish. Compiled layers share code across those batch dimensions. + request_bucket = 1 << (requests - 1).bit_length() + batch_bucket = request_bucket * branches + # Greedy decoding needs no sorting, probability filtering, or RNG + # workspace. Capture it separately from the mutable sampled path. + key = (batch_bucket, branches, temperature == 0) + parameters = (temperature, top_k, top_p) + entry = self._graphs.get(key) + if entry is None: + if self._compiled_layer is None: + self._compiled_layer = torch.compile( + BreezeDepthLayer.forward, + fullgraph=True, + dynamic=True, + options={"epilogue_fusion": False, "max_autotune": True}, + ) + self._compiled_sampler = torch.compile(sample_graph_logits, fullgraph=True, dynamic=False) + entry = BreezeDepthGraph( + self, + hidden.new_empty((batch_bucket, hidden.shape[-1])), + first.new_empty((request_bucket,)), + parameters, + guidance_scale, + ) + self._graphs[key] = entry + if entry.parameter_values != parameters: + entry.parameters.copy_(entry.parameters.new_tensor(parameters)) + entry.parameter_values = parameters + if batch == batch_bucket: + entry.hidden.copy_(hidden) + else: + # CFG graphs keep all conditioned rows before all unconditional + # rows. Pad each branch separately so matching rows stay paired. + for branch in range(branches): + entry.hidden[branch * request_bucket : branch * request_bucket + requests].copy_( + hidden[branch * requests : (branch + 1) * requests] + ) + entry.first[:requests].copy_(first) + if entry.guidance_scale is not None and entry.guidance_value != guidance_scale: + entry.guidance_scale.fill_(guidance_scale) + entry.guidance_value = guidance_scale + if entry.noise_generators is not None: + # Graphs retain placeholder generator states. Copy each live + # request's seed and post-codebook-0 offset into its current row; + # replacing a generator state would invalidate the captured RNG. + for placeholder, generator in zip(entry.noise_generators, generators): + placeholder.manual_seed(generator.initial_seed()) + placeholder.set_offset(generator.get_offset()) + elif temperature > 0: + assert entry.noise is not None + for row, generator in enumerate(generators): + # Match the 15 independent exponential draws used by + # torch.multinomial(..., num_samples=1), preserving each + # request's RNG stream even when requests are batched. + for codebook in range(self.num_codebooks - 1): + entry.noise[row, codebook].exponential_(generator=generator) + entry.graph.replay() + if entry.noise_generators is not None: + # Replay advances host offsets without synchronizing the device. + # Padded rows have no request state to update. + for placeholder, generator in zip(entry.noise_generators, generators): + generator.set_offset(placeholder.get_offset()) + # The tiny RVQ result outlives this replay in request state and the + # transfer queue. The graph owns and overwrites its output workspace. + return entry.output[:requests].clone() + + +class BreezeDepthGraph: + def __init__( + self, + model: BreezeDepthDecoder, + hidden: torch.Tensor, + first: torch.Tensor, + parameters: tuple[float, int, float], + guidance_scale: float = 1.0, + ) -> None: + self.parameter_values = parameters + self.greedy = parameters[0] == 0 + self.parameters = hidden.new_tensor(parameters, dtype=torch.float32) + self.guidance_value = guidance_scale + self.guidance_scale = hidden.new_tensor(guidance_scale, dtype=torch.float32) if guidance_scale != 1.0 else None + self.hidden = torch.zeros_like(hidden) + self.first = torch.zeros_like(first) + self.noise = ( + None + if self.greedy + else hidden.new_ones((first.shape[0], model.num_codebooks - 1, model.vocab_size), dtype=torch.float32) + ) + self.noise_generators: list[torch.Generator] | None = None + if not self.greedy and current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state"): + self.noise_generators = [torch.Generator(device=hidden.device) for _ in range(first.shape[0])] + # Materialize the distribution kernel before graph capture. + self._fill_noise() + # Captured kernels retain addresses, not the Python tensors allocated + # before capture. Keep every KV allocation alive with its graph. + self.caches = model._allocate_cache(hidden) + generator = torch.Generator(device=hidden.device) + for _ in range(3): + model._generate_frame( + self.hidden, + self.first, + self.caches, + *parameters, + generator, + self.noise, + self.guidance_scale, + self.parameters, + greedy=self.greedy, + ) + current_omni_platform.synchronize() + self.graph = torch.cuda.CUDAGraph() + if self.noise_generators is not None: + for generator in self.noise_generators: + self.graph.register_generator_state(generator) + with torch.cuda.graph( + self.graph, pool=current_platform.get_global_graph_pool(), capture_error_mode="thread_local" + ): + if self.noise_generators is not None: + self._fill_noise() + self.output = model._generate_frame( + self.hidden, + self.first, + self.caches, + *parameters, + generator, + self.noise, + self.guidance_scale, + self.parameters, + greedy=self.greedy, + ) + + def _fill_noise(self) -> None: + assert self.noise is not None and self.noise_generators is not None + for row, generator in enumerate(self.noise_generators): + # Preserve the shape and ordering of independent multinomial + # draws; one larger draw would consume a different RNG stream. + for codebook in range(self.noise.shape[1]): + self.noise[row, codebook].exponential_(generator=generator) diff --git a/vllm_omni/model_executor/models/breeze_tts_2/first_code_sampler.py b/vllm_omni/model_executor/models/breeze_tts_2/first_code_sampler.py new file mode 100644 index 000000000..87d9e45eb --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/first_code_sampler.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Capture first-code sampling while preserving each request's RNG stream.""" + +import torch +from vllm.platforms import current_platform + +from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits +from vllm_omni.platforms import current_omni_platform + + +class BreezeFirstCodeSampler: + def __init__(self) -> None: + self._graphs: dict[tuple[torch.device, int, int, tuple[float, int, float]], _SamplerGraph] = {} + + def sample( + self, + logits: list[torch.Tensor], + parameters: tuple[float, int, float], + generators: list[torch.Generator], + ) -> list[torch.Tensor]: + if len(logits) != len(generators): + raise ValueError("First-code sampling requires one generator per logit row") + if not logits: + return [] + first = logits[0] + for row, generator in zip(logits, generators, strict=True): + if ( + row.ndim != 2 + or row.shape[0] != 1 + or row.shape[1] == 0 + or row.shape != first.shape + or row.dtype != torch.float32 + or row.device != first.device + or generator.device.type != row.device.type + or (generator.device.index is not None and generator.device.index != row.device.index) + ): + raise ValueError("First-code sampling requires aligned FP32 [1, vocab] rows and generators") + capture = ( + parameters[0] > 0 + and first.is_cuda + and current_platform.is_cuda() + and hasattr(torch.cuda.CUDAGraph, "register_generator_state") + and len({id(generator) for generator in generators}) == len(generators) + ) + if capture: + bucket = 1 << (len(logits) - 1).bit_length() + key = (first.device, first.shape[1], bucket, parameters) + entry = self._graphs.get(key) + # Sampling parameters are static graph branches. Bound their cache + # without evicting buffers retained by existing captured graphs. + with torch.accelerator.device_index(first.device.index): + if entry is None and len(self._graphs) < 32: + entry = _SamplerGraph(first, bucket, parameters) + self._graphs[key] = entry + if entry is not None: + return entry.replay(logits, generators) + # Shared generators must retain sequential draw ordering, including + # when a matching graph is already cached for independent requests. + return [sample_logits(row, *parameters, generator) for row, generator in zip(logits, generators, strict=True)] + + +class _SamplerGraph: + def __init__(self, row: torch.Tensor, bucket: int, parameters: tuple[float, int, float]) -> None: + self.logits = row.new_zeros((bucket, row.shape[1])) + self.generators = [torch.Generator(device=row.device) for _ in range(bucket)] + + def sample_rows() -> torch.Tensor: + # Keep the original one-row softmax, filtering, multinomial + # validation and RNG draw geometry; batched reductions can change + # seeded output trajectories even when probabilities are close. + return torch.cat( + [ + sample_logits(self.logits[index : index + 1], *parameters, generator) + for index, generator in enumerate(self.generators) + ] + ) + + for _ in range(3): + sample_rows() + current_omni_platform.synchronize() + self.graph = torch.cuda.CUDAGraph() + self.capture_stream = torch.cuda.Stream(device=row.device) + for generator in self.generators: + self.graph.register_generator_state(generator) + with torch.cuda.graph(self.graph, stream=self.capture_stream): + self.output = sample_rows() + + def replay(self, logits: list[torch.Tensor], generators: list[torch.Generator]) -> list[torch.Tensor]: + live = len(logits) + self.logits[:live].copy_(torch.cat(logits)) + # Captures retain the placeholder state objects. Update their values + # rather than replacing those objects with request-owned RNG states. + for placeholder, request in zip(self.generators[:live], generators, strict=True): + placeholder.manual_seed(request.initial_seed()) + placeholder.set_offset(request.get_offset()) + self.graph.replay() + for placeholder, request in zip(self.generators[:live], generators, strict=True): + request.set_offset(placeholder.get_offset()) + # Padded rows have valid logits but no request state to advance. Clone + # the live results so later replays cannot overwrite a returned token. + return list(self.output[:live].clone().split(1)) diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze.py new file mode 100644 index 000000000..c8fdc5148 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze.py @@ -0,0 +1,407 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Breeze-TTS-2: T5Gemma2 prompt encoding, paged Qwen3, and RVQ decoding.""" + +from collections.abc import Iterable +from typing import Any, TypedDict, cast + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2TextEncoder +from vllm.config import VllmConfig +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.utils import record_function_or_nullcontext + +from vllm_omni.config.model import OmniModelConfig +from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze import BreezeConfig +from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import BreezeDepthDecoder +from vllm_omni.model_executor.models.breeze_tts_2.first_code_sampler import BreezeFirstCodeSampler +from vllm_omni.model_executor.models.breeze_tts_2.prompt import CFG_UNCOND_SUFFIX +from vllm_omni.model_executor.models.breeze_tts_2.reference_encoder import BreezeReferenceEncoder +from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import ( + BreezeTextEncoderCompiled, + BreezeTextEncoderGraph, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + + +class BreezeSampling(TypedDict): + temperature: float + top_k: int + top_p: float + repetition_penalty: float + + +class BreezeState(TypedDict): + generator: torch.Generator + history: torch.Tensor + current: torch.Tensor + + +class BreezeForConditionalGeneration(nn.Module): + have_multimodal_outputs = True + has_preprocess = True + has_postprocess = False + requires_full_prefix_cached_hidden_states = False + omni_pooler_payload_include_hidden = False + use_async_omni_output = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config = cast(BreezeConfig, vllm_config.model_config.hf_config) + parallel = vllm_config.parallel_config + if parallel.tensor_parallel_size != 1 or parallel.pipeline_parallel_size != 1: + raise ValueError("Breeze-TTS-2 currently supports TP=1 and PP=1") + if vllm_config.cache_config.enable_prefix_caching: + raise ValueError("Breeze prompt placeholders require enable_prefix_caching=False") + if vllm_config.scheduler_config.enable_chunked_prefill: + raise ValueError("Breeze text encoding requires enable_chunked_prefill=False") + if not cast(OmniModelConfig, vllm_config.model_config).async_chunk: + raise ValueError("Breeze requires async_chunk=True for its stateful codec stage") + if vllm_config.model_config.quantization is not None: + raise ValueError("The initial Breeze integration does not support quantization") + if self.config.backbone_model_type != "qwen3" or self.config.text_encoder_proj_type != "linear": + raise ValueError("This integration supports the Breeze-TTS-2 Qwen3/linear checkpoint") + if self.config.num_codebooks != 16 or self.config.audio_vocab_size != 2051: + raise ValueError("Breeze-TTS-2 requires 16 codebooks with 2048 audio codes and 3 reserved IDs") + self.num_codebooks = self.config.num_codebooks + self.codebook_size = self.config.audio_vocab_size - 3 + self.hidden_size = self.config.backbone_config.hidden_size + backbone_config = vllm_config.with_hf_config(self.config.backbone_config, architectures=["Qwen3ForCausalLM"]) + backbone_config.model_config.hf_text_config = backbone_config.model_config.hf_config + self.model = Qwen3Model(vllm_config=backbone_config, prefix=f"{prefix}.model".strip(".")) + # Breeze replaces the backbone token embedding with summed RVQ embeddings. + self.model.embed_tokens = nn.Identity() + text_config = self.config.text_encoder_config + text_config._attn_implementation = "sdpa" + self.text_encoder = T5Gemma2TextEncoder(text_config) + self.text_encoder_proj = nn.Linear(text_config.hidden_size, self.hidden_size, bias=False) + self._long_text_encoder = BreezeTextEncoderCompiled(self.text_encoder, self.text_encoder_proj) + self.reference_encoder = BreezeReferenceEncoder(vllm_config) + self.depth_decoder = BreezeDepthDecoder(self.config.depth_decoder_config) + self._first_code_sampler = BreezeFirstCodeSampler() + self.lm_head = nn.Linear(self.hidden_size, self.config.vocab_size, bias=False, dtype=torch.float32) + self.register_buffer( + "offsets", torch.arange(self.num_codebooks) * self.config.audio_vocab_size, persistent=False + ) + self.gpu_resident_buffer_keys = { + ("breeze_state", "current"), + ("breeze_state", "history"), + ("breeze_prepared", "embeds"), + } + self._next_tokens: torch.Tensor | None = None + self._text_graphs: dict[tuple[int, int], BreezeTextEncoderGraph] = {} + + @torch.inference_mode() + def preprocess_batch( + self, *, req_ids: list[str], model_intermediate_buffer: dict[str, dict[str, Any]], device: torch.device + ) -> None: + pending = [ + model_intermediate_buffer[rid] for rid in req_ids if "breeze_state" not in model_intermediate_buffer[rid] + ] + if not pending: + return + segment_ids: dict[tuple[int, ...], None] = {} + references: dict[str, dict[str, Any]] = {} + reference_keys: dict[str, str] = {} + for info in pending: + conditioning = info["breeze_prompt"] + segment_ids[tuple(conditioning["target_ids"])] = None + if "reference_waveform" in info: + segment_ids[tuple(conditioning["reference_ids"])] = None + request_id = info["global_request_id"][0] + key = request_id.removesuffix(CFG_UNCOND_SUFFIX) if conditioning["role"] == "uncond" else request_id + reference_keys[request_id] = key + references.setdefault(key, info) + if not self._text_graphs: + # The serving warmup pays these captures before accepting speech + # requests. Cover the default two-row deployment's distinct text + # segments, including two independent reference recordings. + for batch in range(1, min(4, 2 * self.vllm_config.scheduler_config.max_num_seqs) + 1): + for size in (32, 64, 128): + self._text_graphs[(batch, size)] = BreezeTextEncoderGraph( + self.text_encoder, self.text_encoder_proj, size, batch_size=batch + ) + self.reference_encoder.warmup(min(2, self.vllm_config.scheduler_config.max_num_seqs)) + self._long_text_encoder.warmup(min(4, 2 * self.vllm_config.scheduler_config.max_num_seqs)) + length = max(len(ids) for ids in segment_ids) + bucket = max(32, 1 << (length - 1).bit_length()) + with record_function_or_nullcontext("breeze:text_encoder"): + prompts = [torch.tensor(ids, device="cpu", dtype=torch.long) for ids in segment_ids] + if bucket > 128: + encoded = self._long_text_encoder.run_batch(prompts) + # Large prefill activations must not raise the decoder graph + # pool's permanent high-water mark. Return free prefill + # workspace to the device before starting streamed decode. + torch.accelerator.empty_cache() + else: + graph_key = (len(segment_ids), bucket) + if graph_key not in self._text_graphs: + self._text_graphs[graph_key] = BreezeTextEncoderGraph( + self.text_encoder, self.text_encoder_proj, bucket, batch_size=len(segment_ids) + ) + encoded = self._text_graphs[graph_key].run_batch(prompts) + segments = dict(zip(segment_ids, encoded, strict=True)) + reference_embeds: dict[str, torch.Tensor] = {} + if references: + with record_function_or_nullcontext("breeze:reference_encoder"): + codes = self.reference_encoder.encode_batch( + [info["reference_waveform"] for info in references.values()], + [info["breeze_prompt"]["reference_frames"] for info in references.values()], + ) + for key, code in zip(references, codes, strict=True): + # Audio EOS is a zero RVQ frame, distinct from generation EOS. + eos = code.new_full((1, self.num_codebooks), self.config.codebook_eos_token_id) + audio = self.depth_decoder.embed_tokens(torch.cat((code, eos)) + self.offsets).sum(1) + reference_ids = tuple(references[key]["breeze_prompt"]["reference_ids"]) + reference_embeds[key] = torch.cat((segments[reference_ids], audio), dim=0) + for info in pending: + embeds = segments[tuple(info["breeze_prompt"]["target_ids"])] + request_id = info["global_request_id"][0] + if request_id in reference_keys: + embeds = torch.cat((reference_embeds[reference_keys[request_id]], embeds), dim=0) + info["breeze_prepared"] = {"embeds": embeds} + + def embed_input_ids(self, input_ids: torch.Tensor, **_: object) -> torch.Tensor: + # vLLM's memory profiling does not have a request payload. + return self.lm_head.weight.new_zeros((input_ids.numel(), self.hidden_size), dtype=self.model.norm.weight.dtype) + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + *, + breeze_prompt: dict[str, Any], + breeze_sampling: BreezeSampling, + global_request_id: list[str], + breeze_prepared: dict[str, torch.Tensor] | None = None, + breeze_state: BreezeState | None = None, + _omni_is_prefill: bool, + _omni_seed: int | None = None, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, object]]: + if _omni_is_prefill: + if breeze_prepared is None: + raise RuntimeError("Breeze prefill requires the runner's batch preprocessing hook") + embeds = breeze_prepared.pop("embeds") + if embeds.shape[0] != input_ids.numel(): + raise ValueError("Breeze prompt payload and scheduled token count differ") + generator = torch.Generator(device=input_ids.device) + if _omni_seed is None: + generator.seed() + else: + generator.manual_seed(_omni_seed) + state: BreezeState = { + "generator": generator, + "history": input_ids.new_empty((1, 0)), + "current": input_ids.new_empty((1, self.num_codebooks)), + } + return input_ids, embeds, {"breeze_state": state} + if breeze_state is None or input_ids.numel() != 1: + raise ValueError("Breeze decode requires initialized per-request state and one token") + embeds = self.depth_decoder.embed_tokens(breeze_state["current"] + self.offsets).sum(1) + return input_ids, embeds, {} + + def preprocess_decode_batch( + self, + *, + input_ids: torch.Tensor, + req_infos: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, list[dict[str, Any]]]: + """Embed one current RVQ frame per physical decode row, including CFG.""" + if input_ids.numel() != len(req_infos): + raise ValueError("Breeze decode requires one token per request") + frames: list[torch.Tensor] = [] + for info in req_infos: + state = info.get("breeze_state") + if info.get("_omni_is_prefill") or state is None: + raise ValueError("Breeze decode requires initialized per-request state") + current = state["current"] + if current.shape != (1, self.num_codebooks): + raise ValueError("Breeze decode requires one RVQ frame per request") + frames.append(current) + if not frames: + embeds = self.depth_decoder.embed_tokens.weight.new_empty((0, self.hidden_size)) + else: + current = torch.cat(frames, dim=0) + embeds = self.depth_decoder.embed_tokens(current + self.offsets).sum(1) + return input_ids, embeds, [{} for _ in req_infos] + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: object, + ) -> torch.Tensor | IntermediateTensors: + return self.model(None, positions, intermediate_tensors, inputs_embeds) + + def compute_logits( + self, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata | None = None + ) -> torch.Tensor: + logits = hidden_states.new_full((hidden_states.shape[0], self.config.vocab_size), -torch.inf) + if self._next_tokens is None: + logits[:, 0] = 0 + else: + if self._next_tokens.numel() != hidden_states.shape[0]: + raise RuntimeError("Breeze sampled rows do not match the current request batch") + logits.scatter_(1, self._next_tokens[:, None], 0) + return logits + + def make_omni_output( + self, + model_outputs: torch.Tensor, + *, + model_intermediate_buffer: list[dict[str, object]] | None = None, + request_token_spans: list[tuple[int, int]] | None = None, + **_: object, + ) -> OmniOutput: + if not model_intermediate_buffer: + self._next_tokens = None + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs={}) + if request_token_spans is None or len(request_token_spans) != len(model_intermediate_buffer): + raise RuntimeError("Breeze requires one hidden-state span per request") + self._next_tokens = torch.zeros(len(model_intermediate_buffer), device=model_outputs.device, dtype=torch.long) + output_codes = [ + torch.empty((0, self.num_codebooks), device=model_outputs.device, dtype=torch.long) + for _ in model_intermediate_buffer + ] + infos = model_intermediate_buffer + request_rows = {cast(list[str], info["global_request_id"])[0]: index for index, info in enumerate(infos)} + prepared_rows: list[tuple[list[int], torch.Tensor, torch.Tensor, tuple[float, int, float, float]]] = [] + for index, (info, (start, end)) in enumerate(zip(infos, request_token_spans, strict=True)): + if not 0 <= start < end <= model_outputs.shape[0]: + raise RuntimeError("Invalid Breeze request span") + conditioning = cast(dict[str, Any], info["breeze_prompt"]) + if conditioning["role"] == "uncond": + parent = cast(list[str], info["global_request_id"])[0].removesuffix(CFG_UNCOND_SUFFIX) + if parent not in request_rows: + raise RuntimeError("Breeze CFG companion was scheduled without its conditioned branch") + continue + hidden = model_outputs[end - 1 : end] + state = cast(BreezeState, info["breeze_state"]) + sampling = cast(BreezeSampling, info["breeze_sampling"]) + logits = F.linear(hidden.float(), self.lm_head.weight) + indices = [index] + guidance_scale = float(conditioning["guidance_scale"]) + if guidance_scale != 1.0: + request_id = cast(list[str], info["global_request_id"])[0] + companion = request_rows.get(request_id + CFG_UNCOND_SUFFIX) + if companion is None: + raise RuntimeError("Breeze CFG request was scheduled without its unconditional branch") + if infos[companion]["breeze_sampling"] != sampling: + raise RuntimeError("Breeze CFG branches have different sampling parameters") + _, companion_end = request_token_spans[companion] + uncond = model_outputs[companion_end - 1 : companion_end] + uncond_logits = F.linear(uncond.float(), self.lm_head.weight) + logits = uncond_logits + guidance_scale * (logits - uncond_logits) + hidden = torch.cat((hidden, uncond)) + indices.append(companion) + logits[:, self.codebook_size : self.config.eos_token_id] = -torch.inf + history = state["history"] + if history.numel(): + selected = logits.gather(1, history) + penalty = sampling["repetition_penalty"] + logits.scatter_(1, history, torch.where(selected < 0, selected * penalty, selected / penalty)) + parameters = (sampling["temperature"], sampling["top_k"], sampling["top_p"], guidance_scale) + prepared_rows.append((indices, hidden, logits, parameters)) + + generators = [cast(BreezeState, infos[row[0][0]]["breeze_state"])["generator"] for row in prepared_rows] + sampling_groups: dict[tuple[float, int, float], list[int]] = {} + for position, (_, _, _, (temperature, top_k, top_p, _)) in enumerate(prepared_rows): + sampling_groups.setdefault((temperature, top_k, top_p), []).append(position) + positions_by_group = list(sampling_groups.values()) + if len({id(generator) for generator in generators}) != len(generators): + # Normally each request owns a generator. Preserve sequential draw + # ordering even if callers supply shared state across parameter groups. + positions_by_group = [[position] for position in range(len(prepared_rows))] + first_codes: dict[int, torch.Tensor] = {} + for positions in positions_by_group: + temperature, top_k, top_p, _ = prepared_rows[positions[0]][3] + first = self._first_code_sampler.sample( + [prepared_rows[position][2] for position in positions], + (temperature, top_k, top_p), + [generators[position] for position in positions], + ) + first_codes.update(zip(positions, first, strict=True)) + sampled_rows = [ + (indices, hidden, first_codes[position], parameters) + for position, (indices, hidden, _, parameters) in enumerate(prepared_rows) + ] + # Read back EOS once after all first-code samples. Restore request order + # before grouping depth work so its RNG and CFG ordering stay unchanged. + first_ids = torch.cat([row[2] for row in sampled_rows]).to("cpu").tolist() if sampled_rows else [] + groups: dict[tuple[float, int, float, float], list[tuple[list[int], torch.Tensor, torch.Tensor]]] = {} + for (indices, hidden, first, parameters), first_id in zip(sampled_rows, first_ids, strict=True): + if first_id == self.config.eos_token_id: + self._next_tokens[indices] = self.config.eos_token_id + continue + state = cast(BreezeState, infos[indices[0]]["breeze_state"]) + state["history"] = torch.cat((state["history"], first[:, None]), dim=1) + groups.setdefault(parameters, []).append((indices, hidden, first)) + for (temperature, top_k, top_p, guidance_scale), rows in groups.items(): + # Depth batches place every conditioned row before every + # unconditional row. Both branches consume the same sampled codes. + hidden = torch.stack([row[1] for row in rows]).transpose(0, 1).reshape(-1, self.hidden_size) + with record_function_or_nullcontext("breeze:depth_decoder"): + frames = self.depth_decoder.generate_frames( + hidden, + torch.cat([row[2] for row in rows]), + temperature=temperature, + top_k=top_k, + top_p=top_p, + generators=[cast(BreezeState, infos[row[0][0]]["breeze_state"])["generator"] for row in rows], + guidance_scale=guidance_scale, + ) + for offset, (indices, _, _) in enumerate(rows): + frame = frames[offset : offset + 1] + for index in indices: + cast(BreezeState, infos[index]["breeze_state"])["current"] = frame + output_codes[indices[0]] = frame + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs={"codes": {"audio": output_codes}}) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + backbone: list[tuple[str, torch.Tensor]] = [] + depth: dict[str, torch.Tensor] = {} + loaded: set[str] = set() + params = dict(self.named_parameters()) + for name, value in weights: + if name.startswith(("backbone_model.layers.", "backbone_model.norm.")): + backbone.append((name.removeprefix("backbone_model."), value)) + elif name.startswith("depth_decoder."): + depth[name] = value + elif name.startswith(("text_encoder.", "text_encoder_proj.", "lm_head.")): + default_weight_loader(params[name], value) + loaded.add(name) + elif not name.startswith(("codec_model.", "backbone_model.embed_tokens.", "embed_text_tokens.")): + raise ValueError(f"Unexpected Breeze checkpoint parameter: {name}") + loaded.update("model." + name for name in self.model.load_weights(backbone)) + for target, parameter in self.depth_decoder.named_parameters(): + if target == "codebooks_head": + value = depth["depth_decoder.codebooks_head.weight"] + elif target.endswith("qkv.weight"): + layer = target.split(".")[1] + value = torch.cat( + [depth[f"depth_decoder.model.layers.{layer}.self_attn.{p}_proj.weight"] for p in ("q", "k", "v")] + ) + elif target.endswith("gate_up.weight"): + layer = target.split(".")[1] + value = torch.cat( + [depth[f"depth_decoder.model.layers.{layer}.mlp.{p}_proj.weight"] for p in ("gate", "up")] + ) + else: + source = target.replace(".o_proj.", ".self_attn.o_proj.").replace(".down_proj.", ".mlp.down_proj.") + value = depth["depth_decoder.model." + source] + default_weight_loader(parameter, value) + loaded.add("depth_decoder." + target) + loaded.update("reference_encoder." + name for name in self.reference_encoder.load_weights(self.vllm_config)) + missing = set(params) - loaded + if missing: + raise ValueError(f"Uninitialized Breeze parameters: {sorted(missing)}") + return loaded diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py deleted file mode 100644 index 860d64c44..000000000 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py +++ /dev/null @@ -1,425 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Breeze-TTS-2 stage-1 codec (codec codes -> waveform). - -The talker emits a frame-major ``(T, 16)`` matrix. The stage input -processor flattens it codebook-major; this module restores the ``(T, 16)`` -layout and delegates decoding to Breeze's bundled Qwen3-TTS tokenizer. A -Mimi fallback remains available for checkpoints that omit ``audio_tokenizer``. -""" - -from __future__ import annotations - -from collections.abc import Iterable, Mapping -from typing import Any - -import torch -from torch import nn -from vllm.config import VllmConfig - -from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( - resolve_audio_tokenizer_path, -) -from vllm_omni.model_executor.models.output_templates import OmniOutput - - -def _split_flat_ids(ids: torch.Tensor, counts: list[int] | None) -> list[torch.Tensor]: - if counts is None or len(counts) <= 1: - return [ids] - boundaries = [0] - for count in counts: - boundaries.append(boundaries[-1] + int(count)) - return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(counts))] - - -class BreezeTTS2MimiCodec(nn.Module): - """Synchronous Breeze codec stage used by ``GenerationModelRunner``.""" - - input_modalities = "audio" - requires_raw_input_tokens = True - # Stateful chunk decoding needs the scheduler-side request id, which is - # also used by on_requests_finished to release decoder state. - requires_request_ids = True - have_multimodal_outputs = True - has_preprocess = False - has_postprocess = False - enable_update_additional_information = True - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - del prefix - self.vllm_config = vllm_config - self.config = vllm_config.model_config.hf_config - self.model_path = vllm_config.model_config.model - self._async_chunk = bool(getattr(vllm_config.model_config, "async_chunk", False)) - # Resolve once: a local directory works directly, an HF repo id - # resolves to its cached snapshot. ``None`` means the checkpoint has - # no bundled tokenizer and the Mimi fallback applies. - self._tokenizer_path = resolve_audio_tokenizer_path(self.model_path) - if self._async_chunk and self._tokenizer_path is None: - raise RuntimeError( - "Breeze-TTS-2 async_chunk requires the bundled Qwen3-TTS audio_tokenizer; " - "the stateless Mimi fallback cannot preserve chunk decoder state" - ) - - self._audio_tokenizer = None - codec_config = getattr(self.config, "codec_config", None) - if isinstance(codec_config, Mapping): - codec_kwargs = dict(codec_config) - elif codec_config is not None: - codec_kwargs = codec_config.to_dict() - else: - codec_kwargs = {} - codec_kwargs.pop("model_type", None) - codec_kwargs.pop("architectures", None) - - # Breeze checkpoints normally carry a Qwen3-TTS decoder in this - # subdirectory. Avoid constructing a second, unused Mimi network in - # that case; this saves both startup time and worker memory. - self._codec = None - if self._tokenizer_path is None: - # Import lazily: stage 0 does not need the Mimi implementation. - from transformers import MimiConfig, MimiModel - - self._codec = MimiModel(MimiConfig(**codec_kwargs)).eval() - for parameter in self._codec.parameters(): - parameter.requires_grad_(False) - - self._num_codebooks = int(getattr(self.config, "num_codebooks", 16)) - self._codebook_size = int(codec_kwargs.get("codebook_size", 2048)) - self._sample_rate = int(codec_kwargs.get("sampling_rate", getattr(self.config, "sampling_rate", 24000))) - self._decoder_state_cache: dict[str, dict[str, Any]] = {} - - def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: - """Return stable dummy embeddings; this stage never samples.""" - if input_ids.numel() == 0: - return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) - return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) - - def compute_logits( - self, - hidden_states: torch.Tensor | OmniOutput, - sampling_metadata: Any = None, - ) -> None: - del hidden_states, sampling_metadata - return None - - @torch.inference_mode() - def forward( - self, - input_ids: torch.Tensor | None = None, - runtime_additional_information: list[dict[str, Any]] | None = None, - **kwargs: Any, - ) -> OmniOutput: - empty = torch.empty(0, dtype=torch.float32) - sr = torch.tensor(self._sample_rate, dtype=torch.int32) - if input_ids is None or input_ids.numel() == 0: - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": [empty], "sr": [sr]}, - ) - # The generation runner's profiling/dummy invocation has no request ids - # or runtime payload. Real requests always provide both. - if runtime_additional_information is None and kwargs.get("request_ids") is None: - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": [empty], "sr": [sr]}, - ) - - ids = input_ids.reshape(-1).to(dtype=torch.long) - counts = kwargs.get("seq_token_counts") - if counts is not None: - counts = [int(item) for item in counts] - requests = _split_flat_ids(ids, counts) - waveforms: list[torch.Tensor] = [] - sample_rates: list[torch.Tensor] = [] - if self._async_chunk: - return self._forward_stateful_chunks( - ids=ids, - requests=requests, - runtime_additional_information=runtime_additional_information, - empty=empty, - sr=sr, - **kwargs, - ) - for index, request_ids in enumerate(requests): - runtime_info = ( - runtime_additional_information[index] - if runtime_additional_information is not None and index < len(runtime_additional_information) - else None - ) - flat = self._payload_codes(request_ids, runtime_info) - if flat.numel() == 0: - waveforms.append(empty.to(device=ids.device)) - sample_rates.append(sr) - continue - if flat.numel() % self._num_codebooks != 0: - raise ValueError( - f"Breeze codec input length {int(flat.numel())} is not divisible by " - f"num_codebooks {self._num_codebooks}" - ) - if int(flat.min()) < 0 or int(flat.max()) >= self._codebook_size: - raise ValueError( - f"Breeze codec id outside [0, {self._codebook_size}): min={int(flat.min())}, max={int(flat.max())}" - ) - frames = flat.numel() // self._num_codebooks - codes = flat.reshape(self._num_codebooks, frames).unsqueeze(0) - waveform = self._decode_codes(codes) - if waveform.ndim == 3: - waveform = waveform[0, 0] - elif waveform.ndim == 2: - waveform = waveform[0] - if waveform.ndim != 1: - raise ValueError(f"Breeze Mimi waveform must be 1D, got {tuple(waveform.shape)}") - # Serving/output serialization consumes CPU tensors. Move only - # the final waveform off the worker device; codec inference itself - # remains on the accelerator. - waveforms.append(waveform.detach().to(device="cpu", dtype=torch.float32).contiguous()) - sample_rates.append(sr) - - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, - ) - - def _forward_stateful_chunks( - self, - *, - ids: torch.Tensor, - requests: list[torch.Tensor], - runtime_additional_information: list[dict[str, Any]] | None, - empty: torch.Tensor, - sr: torch.Tensor, - **kwargs: Any, - ) -> OmniOutput: - """Decode incremental codebook-major chunks with per-request state.""" - decoder = self._audio_tokenizer.model.decoder - scheduler_request_ids = kwargs.get("request_ids") - if scheduler_request_ids is None: - raise ValueError("Breeze async_chunk codec inputs require scheduler request_ids") - request_ids = [str(request_id) for request_id in scheduler_request_ids] - if len(request_ids) != len(requests): - raise ValueError( - f"Breeze codec request id count {len(request_ids)} does not match input request count {len(requests)}" - ) - valid: list[tuple[str, torch.Tensor, bool]] = [] - valid_indices: list[int] = [] - for index, request_ids_tensor in enumerate(requests): - runtime_info = ( - runtime_additional_information[index] - if runtime_additional_information is not None and index < len(runtime_additional_information) - else {} - ) - state_id = request_ids[index] - # The connector receiver strips ``meta.finished`` while merging - # runtime metadata, so the stream-terminated signal arrives on - # the model-visible ``stream_finished`` field (the same contract - # CosyVoice3 / Audex use). Fall back to ``finished`` for direct - # in-process callers that still set only the transport flag. - finished = self._meta_bool(runtime_info, "stream_finished") or self._meta_bool(runtime_info, "finished") - flat = self._payload_codes(request_ids_tensor, runtime_info) - if flat.numel() == 0: - # Terminal marker: the stage-0 processor sends an empty - # payload with ``finished=True`` once generation ends. - self._decoder_state_cache.pop(state_id, None) - continue - if flat.numel() % self._num_codebooks != 0: - # A partial frame means the transport payload is corrupt. - # Silently resetting decoder state would corrupt the rest of - # the stream; fail loudly instead. The scheduler aborts the - # request and ``on_requests_finished`` releases the state. - raise ValueError( - f"Breeze codec chunk length {int(flat.numel())} is not divisible by " - f"num_codebooks {self._num_codebooks}" - ) - if int(flat.min()) < 0 or int(flat.max()) >= self._codebook_size: - raise ValueError( - f"Breeze codec id outside [0, {self._codebook_size}): min={int(flat.min())}, max={int(flat.max())}" - ) - frames = int(flat.numel()) // self._num_codebooks - valid.append((state_id, flat.reshape(self._num_codebooks, frames), finished)) - valid_indices.append(index) - - waveforms = [empty.to(device=ids.device)] * len(requests) - sample_rates = [sr] * len(requests) - if not valid: - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, - ) - states: list[dict[str, Any]] = [] - max_frames = max(codes.shape[-1] for _, codes, _ in valid) - batched_codes = valid[0][1].new_zeros((len(valid), self._num_codebooks, max_frames)) - for row, (state_id, row_codes, _) in enumerate(valid): - batched_codes[row, :, : row_codes.shape[-1]].copy_(row_codes) - state = self._decoder_state_cache.get(state_id) - if state is None: - state = {} - self._decoder_state_cache[state_id] = state - state.setdefault("prefix_frames", 0) - states.append(state) - - decoded = decoder.batched_chunked_decode( - batched_codes, - [int(codes.shape[-1]) for _, codes, _ in valid], - caches=states, - chunk_size=300, - left_context_size=25, - ) - if len(decoded) != len(valid): - raise RuntimeError(f"Breeze codec returned {len(decoded)} chunks for {len(valid)} requests") - for request_index, waveform in zip(valid_indices, decoded, strict=True): - if waveform.ndim == 2 and waveform.shape[0] == 1: - waveform = waveform[0] - elif waveform.ndim != 1: - waveform = waveform.reshape(-1) - waveforms[request_index] = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous() - - for state_id, _, finished in valid: - if finished: - self._decoder_state_cache.pop(state_id, None) - - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, - ) - - @staticmethod - def _meta_bool(info: Mapping[str, Any] | None, key: str) -> bool: - if not isinstance(info, Mapping): - return False - meta = info.get("meta", {}) - if not isinstance(meta, Mapping): - return False - value = meta.get(key, False) - if isinstance(value, torch.Tensor): - return bool(value.reshape(-1)[0].item()) if value.numel() else False - return bool(value) - - def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: - for request_id in finished_req_ids: - self._decoder_state_cache.pop(str(request_id), None) - - def _decode_codes(self, codes: torch.Tensor) -> torch.Tensor: - """Decode one ``(1, Q, T)`` sequence using Breeze's bundled codec.""" - if self._audio_tokenizer is not None: - wavs, sample_rate = self._audio_tokenizer.decode({"audio_codes": [codes[0].transpose(0, 1).contiguous()]}) - self._sample_rate = int(sample_rate) - if not wavs: - raise RuntimeError("Breeze Qwen3-TTS tokenizer returned no waveform") - return torch.as_tensor(wavs[0], device=codes.device, dtype=torch.float32).reshape(-1) - - if self._codec is None: - raise RuntimeError("Breeze codec is not initialized; load_weights() was not completed") - decoded = self._codec.decode(codes) - waveform = getattr(decoded, "audio_values", decoded) - if isinstance(waveform, Mapping): - waveform = waveform.get("audio_values") - if waveform is None: - raise RuntimeError("MimiModel.decode returned no audio_values") - return torch.as_tensor(waveform, device=codes.device, dtype=torch.float32) - - @staticmethod - def _payload_codes(input_ids: torch.Tensor, runtime_info: Mapping[str, Any] | None) -> torch.Tensor: - if isinstance(runtime_info, Mapping): - codes = runtime_info.get("codes") - if isinstance(codes, Mapping): - audio = codes.get("audio") - if audio is not None: - return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1) - # Full-payload transport may flatten nested OmniPayload keys at - # the wire boundary and restore them only after this hook runs. - audio = runtime_info.get("codes.audio") - if audio is not None: - return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1) - return input_ids - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load Breeze's bundled Qwen3-TTS codec, or fallback Mimi weights.""" - if self._tokenizer_path is not None: - # The root checkpoint contains both the talker weights and a Mimi - # ``codec_model.*`` fallback. When the bundled Qwen3-TTS tokenizer - # is present, it owns decoding; consume the root iterator without - # copying the unused fallback into worker memory. - if self._audio_tokenizer is None: - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import ( - Qwen3TTSTokenizer, - ) - - device = self.vllm_config.device_config.device - self._audio_tokenizer = Qwen3TTSTokenizer.from_pretrained( - str(self._tokenizer_path), - device_map=str(device), - ) - self._sample_rate = int(self._audio_tokenizer.model.get_output_sample_rate()) - if self._async_chunk: - connector_config = getattr(self.vllm_config.model_config, "stage_connector_config", None) - extra = ( - connector_config.get("extra", connector_config) - if isinstance(connector_config, Mapping) - else getattr(connector_config, "extra", None) - ) - value = extra.get("breeze_codec_chunk_frames") if isinstance(extra, Mapping) else None - try: - chunk_frames = int(value) if value is not None else 8 - except (TypeError, ValueError) as exc: - raise ValueError(f"Invalid Breeze codec_chunk_frames={value!r}") from exc - if chunk_frames <= 0: - raise ValueError(f"Breeze codec_chunk_frames must be positive, got {chunk_frames}") - # Qwen treats chunks shorter than its configured window as - # terminal tails and does not advance their rolling cache. - # Match the producer's chunk size instead of its default 25. - decoder = self._audio_tokenizer.model.decoder - decoder._initial_codec_chunk_frames = chunk_frames - decoder._incremental_chunk_frames = chunk_frames - decoder._incremental_chunk_ramp = [] - for _ in weights: - pass - tokenizer_module = getattr(self._audio_tokenizer, "model", None) - named_parameters = getattr(tokenizer_module, "named_parameters", None) - if callable(named_parameters): - return {f"_audio_tokenizer.model.{name}" for name, _ in named_parameters()} - return set() - - params = dict(self._codec.named_parameters()) if self._codec is not None else {} - loaded: set[str] = set() - for name, tensor in weights: - if name.startswith("codec_model."): - target_name = name[len("codec_model.") :] - else: - continue - parameter = params.get(target_name) - if parameter is None: - raise ValueError(f"Unknown Breeze Mimi weight: {target_name}") - loader = getattr(parameter, "weight_loader", None) - if loader is None: - parameter.data.copy_(tensor.to(device=parameter.device, dtype=parameter.dtype)) - else: - loader(parameter, tensor) - loaded.add(target_name) - - if loaded: - return {f"_codec.{name}" for name in loaded} - - # Breeze inference uses the bundled Qwen3-TTS tokenizer for both - # reference encoding and generated-code decoding. Load it once per - # Stage-1 worker; unlike the request-side wrapper this instance is - # placed on the worker device and reused for the entire batch. - # A few Breeze distributions omit codec tensors and point to the - # standalone kyutai/mimi repository instead. Load it once, never per - # request, and fail loudly if it cannot be resolved. - codec_name = getattr(self.config, "codec_model_name_or_path", None) or "kyutai/mimi" - from transformers import AutoModel - - self._codec = AutoModel.from_pretrained(codec_name).eval() - for parameter in self._codec.parameters(): - parameter.requires_grad_(False) - self._sample_rate = int(getattr(self._codec.config, "sampling_rate", self._sample_rate)) - self._codebook_size = int( - getattr(getattr(self._codec.config, "quantizer", None), "cardinality", self._codebook_size) - ) - return {f"_codec.{name}" for name, _ in self._codec.named_parameters()} - - -__all__ = ["BreezeTTS2MimiCodec"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py deleted file mode 100644 index 559c88b07..000000000 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py +++ /dev/null @@ -1,523 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -from collections.abc import Callable -from inspect import signature -from typing import Any - -import torch -from torch import nn -from transformers.activations import ACT2FN -from transformers.cache_utils import Cache, DynamicCache -from transformers.integrations import use_kernel_forward_from_hub -from transformers.masking_utils import create_causal_mask -from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from transformers.processing_utils import Unpack -from transformers.utils import TransformersKwargs, logging -from transformers.utils.deprecation import deprecate_kwarg - -logger = logging.get_logger(__name__) - - -def _create_depth_causal_mask( - *, - config: Any, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None, - past_key_values: Cache | None, - position_ids: torch.LongTensor | None, - cache_position: torch.LongTensor, -) -> torch.Tensor | None: - """Bridge Transformers 4.57 and 5.x ``create_causal_mask`` contracts.""" - parameters = signature(create_causal_mask).parameters - kwargs: dict[str, Any] = { - "config": config, - "attention_mask": attention_mask, - "past_key_values": past_key_values, - "position_ids": position_ids, - } - if "inputs_embeds" in parameters: - kwargs["inputs_embeds"] = inputs_embeds - else: - kwargs["input_embeds"] = inputs_embeds - kwargs["cache_position"] = cache_position - return create_causal_mask(**kwargs) - - -class BreezeDepthDecoderModel(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - if hasattr(config, "audio_embed_size") and config.audio_embed_size: - self.audio_embed_size = config.audio_embed_size - else: - self.audio_embed_size = config.backbone_hidden_size - self.embed_tokens = nn.Embedding((config.num_codebooks * config.vocab_size), self.audio_embed_size) - self.layers = nn.ModuleList( - [BreezeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - self.norm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = BreezeRotaryEmbedding(config=config) - self.gradient_checkpointing = False - self.inputs_embeds_projector = nn.Linear(self.audio_embed_size, config.hidden_size, bias=False) - - if config.backbone_hidden_size != self.audio_embed_size: - self.backbone_hidden_state_projector = nn.Linear( - config.backbone_hidden_size, self.audio_embed_size, bias=False - ) - else: - self.backbone_hidden_state_projector = None - - # Weights are loaded from the Breeze checkpoint by the parent talker. - # Avoid PreTrainedModel.post_init here because this module is a plain - # nn.Module used for eager per-frame inference. - - def forward( - self, - input_ids: torch.LongTensor = None, - backbone_last_hidden_state: torch.FloatTensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - use_cache: bool | None = None, - cache_position: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> tuple | BaseModelOutputWithPast: - r""" - backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*): - The last hidden state of the backbone model. This is required when - the first codebook token (generated by the backbone model) is - provided in the `input_ids` argument. - """ - if position_ids is not None and not torch.compiler.is_compiling(): - logger.warning_once( - "Custom `position_ids` were provided but will be ignored. " - "Breeze depth decoder derives them from `cache_position`; " - "they must be identical across the batch." - ) - position_ids = None - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds.") - - if use_cache and past_key_values is None: - past_key_values = DynamicCache(config=self.config) - - if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - inputs_seq_length = inputs_embeds.shape[1] if inputs_embeds is not None else input_ids.shape[1] - device = inputs_embeds.device if inputs_embeds is not None else input_ids.device - cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_seq_length, device=device) - - if inputs_embeds is None: - codebook_idxs = torch.clamp(cache_position - 1, min=0) - offset = codebook_idxs * self.vocab_size - inputs_embeds = self.embed_tokens(input_ids + offset) - - input_ids_are_first_codebook = cache_position[0] == 0 - if backbone_last_hidden_state is not None: - if self.backbone_hidden_state_projector is not None: - inputs_embeds[:, 0] = self.backbone_hidden_state_projector(backbone_last_hidden_state) - else: - inputs_embeds[:, 0] = backbone_last_hidden_state - else: - if not torch.compiler.is_compiling() and input_ids_are_first_codebook: - logger.warning( - "When the first codebook token is provided, " - "`backbone_last_hidden_state` is required for correct inference." - ) - - inputs_embeds = self.inputs_embeds_projector(inputs_embeds) - - causal_mask = _create_depth_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - cache_position=cache_position, - ) - - hidden_states = inputs_embeds - - # create position embeddings to be shared across the decoder layers - position_ids = cache_position.unsqueeze(0) - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - for decoder_layer in self.layers[: self.config.num_hidden_layers]: - hidden_states = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - - hidden_states = self.norm(hidden_states) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=past_key_values if use_cache else None, - ) - - -@use_kernel_forward_from_hub("RMSNorm") -class BreezeRMSNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - """ - BreezeRMSNorm is equivalent to T5LayerNorm - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" - - -class BreezeRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - - def __init__(self, config, device=None): - super().__init__() - # BC: "rope_type" was originally "type" - if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): - self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) - else: - self.rope_type = "default" - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - - self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - - inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.original_inv_freq = self.inv_freq - - @torch.no_grad() - @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) - def forward(self, x, position_ids): - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) - position_ids_expanded = position_ids[:, None, :].float() - - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - with torch.autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() * self.attention_scaling - sin = emb.sin() * self.attention_scaling - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - -class BreezeMLP(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - position_ids (`torch.Tensor`, *optional*): - Deprecated and unused. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] - attn_weights = attn_weights + causal_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - -class BreezeAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config, layer_idx: int): - super().__init__() - self.config = config - self.layer_idx = layer_idx - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads - self.scaling = self.head_dim**-0.5 - self.attention_dropout = config.attention_dropout - self.is_causal = True - - self.q_proj = nn.Linear( - config.hidden_size, - config.num_attention_heads * self.head_dim, - bias=config.attention_bias, - ) - self.k_proj = nn.Linear( - config.hidden_size, - config.num_key_value_heads * self.head_dim, - bias=config.attention_bias, - ) - self.v_proj = nn.Linear( - config.hidden_size, - config.num_key_value_heads * self.head_dim, - bias=config.attention_bias, - ) - self.o_proj = nn.Linear( - config.num_attention_heads * self.head_dim, - config.hidden_size, - bias=config.attention_bias, - ) - - @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None, - past_key_values: Cache | None = None, - cache_position: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> tuple[torch.Tensor, torch.Tensor]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) - key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) - - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - - if past_key_values is not None: - # sin and cos are specific to RoPE models; cache_position needed for the static cache - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) - - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - -class BreezeDecoderLayer(nn.Module): - def __init__(self, config, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - - self.self_attn = BreezeAttention(config=config, layer_idx=layer_idx) - - self.mlp = BreezeMLP(config) - self.input_layernorm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - use_cache: bool | None = False, - cache_position: torch.LongTensor | None = None, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, # necessary, but kept here for BC - **kwargs: Unpack[TransformersKwargs], - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - # Self Attention - hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = residual + hidden_states - - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -class BreezeCodebooksHead(nn.Module): - """Position-specific heads for codebooks 1..15.""" - - def __init__(self, hidden_size: int, num_codebooks: int, vocab_size: int) -> None: - super().__init__() - self.num_codebooks = int(num_codebooks) - self.weight = nn.Parameter(torch.empty(self.num_codebooks - 1, hidden_size, vocab_size)) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # ``hidden_states`` contains the placeholder at position 0 followed - # by codebook-0..N tokens. Return logits for codebook-1..N. - depth_states = hidden_states[:, 1 : self.num_codebooks, :] - if depth_states.shape[1] > self.num_codebooks - 1: - raise ValueError(f"depth decoder received too many input positions: {hidden_states.shape[1]}") - return torch.einsum("bkh,khv->bkv", depth_states, self.weight[: depth_states.shape[1]]) - - -class BreezeDepthDecoderForCausalLM(nn.Module): - """Eager depth decoder used by the vLLM stage-0 talker.""" - - def __init__(self, config) -> None: - super().__init__() - self.config = config - self.model = BreezeDepthDecoderModel(config) - self.codebooks_head = BreezeCodebooksHead( - config.hidden_size, - config.num_codebooks, - config.vocab_size, - ) - - def forward( - self, - input_ids: torch.Tensor, - backbone_last_hidden_state: torch.Tensor, - *, - use_cache: bool = False, - **kwargs, - ) -> CausalLMOutputWithPast: - outputs = self.model( - input_ids=input_ids, - backbone_last_hidden_state=backbone_last_hidden_state, - use_cache=use_cache, - **kwargs, - ) - logits = self.codebooks_head(outputs.last_hidden_state) - return CausalLMOutputWithPast( - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=getattr(outputs, "hidden_states", None), - attentions=getattr(outputs, "attentions", None), - ) - - def load_weights(self, weights): - params_dict = dict(self.named_parameters()) - loaded: set[str] = set() - for name, tensor in weights: - param = params_dict.get(name) - if param is None: - continue - if name == "codebooks_head.weight": - expected = ( - int(self.config.num_codebooks) - 1, - int(self.config.hidden_size), - int(self.config.vocab_size), - ) - actual = tuple(int(dim) for dim in tensor.shape) - if actual != expected: - raise ValueError(f"Breeze depth weight {name!r} has shape {actual}, expected {expected}") - loader = getattr(param, "weight_loader", None) - if loader is None: - param.data.copy_(tensor.to(device=param.device, dtype=param.dtype)) - else: - loader(param, tensor) - loaded.add(name) - return loaded - - -__all__ = [ - "BreezeCodebooksHead", - "BreezeDepthDecoderForCausalLM", - "BreezeDepthDecoderModel", -] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py deleted file mode 100644 index 765df5fbf..000000000 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py +++ /dev/null @@ -1,641 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Breeze-TTS-2 stage-0 talker for vLLM-Omni. - -The scheduler samples Breeze codebook-0 tokens one at a time. The other 15 -codebooks are kept in the request-local Omni buffer: after each backbone step, -the eager depth decoder completes the frame and the next backbone step -consumes the summed 16-codebook embedding. This preserves the upstream Breeze -generation order without importing HuggingFace GenerationMixin. -""" - -from __future__ import annotations - -import hashlib -import math -import os -from collections.abc import Iterable, Mapping, Sequence -from pathlib import Path -from typing import Any - -import torch -from torch import nn -from torch.nn.utils.rnn import pad_sequence -from vllm.config import VllmConfig -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead -from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.qwen3 import Qwen3Model -from vllm.sequence import IntermediateTensors -from vllm.v1.sample.metadata import SamplingMetadata - -from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( - BreezeTTS2Config, -) -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_depth import ( - BreezeDepthDecoderForCausalLM, -) -from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text import ( - BreezeTTS2TextEncoder, -) -from vllm_omni.model_executor.models.output_templates import OmniOutput - - -def _prefixed(prefix: str, name: str) -> str: - return f"{prefix}.{name}" if prefix else name - - -def _normalize_request_infos(infos: Any) -> list[dict[str, Any]] | None: - """Normalize runner metadata to one mutable dict per request.""" - if not isinstance(infos, Sequence) or isinstance(infos, (str, bytes, bytearray)): - return None - if not all(isinstance(info, dict) for info in infos): - return None - return infos - - -def _iter_request_rows( - states: Sequence[dict[str, Any]], - spans: Sequence[tuple[int, int]] | None, - num_rows: int, -) -> Iterable[tuple[dict[str, Any], int, int]]: - """Map request state to logits rows for full or sampled hidden layouts.""" - if num_rows <= 0 or not states: - return - if spans is not None and len(spans) == len(states): - if all(0 <= int(start) <= int(end) <= num_rows for start, end in spans): - for state, (start, end) in zip(states, spans, strict=True): - if int(start) < int(end): - yield state, int(start), int(end) - return - # The runner passes full-sequence spans to make_omni_output but only - # one sampled row per request to compute_logits. - if num_rows == len(states): - for index, state in enumerate(states): - yield state, index, index + 1 - return - if num_rows == len(states): - for index, state in enumerate(states): - yield state, index, index + 1 - return - raise RuntimeError( - "Breeze talker cannot align request state with logits rows: " - f"states={len(states)}, rows={num_rows}, spans={spans}" - ) - - -def _check_weight_shape(name: str, tensor: torch.Tensor, expected: tuple[int, ...]) -> None: - """Fail early for Breeze-owned, non-sharded checkpoint tensors.""" - actual = tuple(int(dim) for dim in tensor.shape) - if actual != expected: - raise ValueError(f"Breeze weight {name!r} has shape {actual}, expected {expected}") - - -class BreezeTTS2TalkerForGeneration(nn.Module): - """Stage-0 Breeze talker: text/audio conditioning -> codec frames.""" - - have_multimodal_outputs = True - has_preprocess = True - has_postprocess = True - # Stage 0 emits a tiny per-frame codec payload. Let the AR runner snapshot - # it through the async-output path instead of the final-only client path. - use_async_omni_output = True - eager_omni_postprocess_before_async_output = True - omni_pooler_payload_include_hidden = False - input_modalities = "text" - gpu_resident_buffer_keys: set[tuple[str, str]] = { - ("hidden_states", "last"), - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - self.vllm_config = vllm_config - self.config: BreezeTTS2Config = vllm_config.model_config.hf_config - self._async_chunk = bool(vllm_config.model_config.async_chunk) - self._golden_dump_dir = os.environ.get("VLLM_OMNI_BREEZE_GOLDEN_DIR") - - backbone_config = self.config.backbone_config - backbone_vllm_config = vllm_config.with_hf_config( - backbone_config, - architectures=["Qwen3ForCausalLM"], - ) - backbone_vllm_config.model_config.hf_text_config = backbone_config - self.model = Qwen3Model( - vllm_config=backbone_vllm_config, - prefix=_prefixed(prefix, "model"), - ) - # Upstream's Qwen3 adapter keeps only decoder layers/norm/RoPE and - # replaces the standard text embedding with Breeze's summed audio/text - # tables. Breeze always supplies ``inputs_embeds`` through preprocess, - # so remove the unused Qwen embedding instead of allocating and then - # reporting an uninitialized 151936x2048 parameter at load time. - self.model.embed_tokens = nn.Identity() - - self.hidden_size = int(backbone_config.hidden_size) - self.num_codebooks = int(self.config.num_codebooks) - self.codebook_vocab_size = int(self.config.vocab_size) - codec_config = self.config.codec_config - self.codebook_size = ( - int(codec_config.get("codebook_size", self.codebook_vocab_size)) - if isinstance(codec_config, Mapping) - else self.codebook_vocab_size - ) - self.codebook_eos_token_id = int(self.config.codebook_eos_token_id) - self.audio_token_id = int(self.config.audio_token_id) - self.audio_eos_token_id = int(self.config.audio_eos_token_id) - self.text_pad_token_id = int(self.config.pad_token_id) - - self.lm_head = ParallelLMHead( - self.codebook_vocab_size + 1, - self.hidden_size, - bias=False, - prefix=_prefixed(prefix, "lm_head"), - ) - self.logits_processor = LogitsProcessor(self.codebook_vocab_size + 1) - - # Upstream Breeze uses embedding(code_i + i * vocab_size) and sums - # over codebooks before entering the Qwen3 backbone. - self.embed_audio_tokens = nn.Embedding( - self.num_codebooks * self.codebook_vocab_size, - self.hidden_size, - ) - # Keep Breeze's standalone text table even though the normal path - # uses T5Gemma2 outputs. It is part of the checkpoint contract and is - # also the correct fallback when a checkpoint omits text_encoder. - self.embed_text_tokens = nn.Embedding(int(self.config.text_vocab_size), self.hidden_size) - self.register_buffer( - "audio_token_offsets", - torch.arange(self.num_codebooks, dtype=torch.long) * self.codebook_vocab_size, - persistent=False, - ) - - self.text_encoder = BreezeTTS2TextEncoder(self.config.text_encoder_config) - self.text_encoder_proj = nn.Linear( - int(self.config.text_encoder_config.hidden_size), - self.hidden_size, - bias=False, - ) - self.depth_decoder = BreezeDepthDecoderForCausalLM(self.config.depth_decoder_config) - if bool(self.config.tie_codebooks_embeddings): - # The upstream HF model ties these tables. Sharing the Parameter - # also prevents a tied checkpoint from leaving the depth input - # embedding randomly initialized when only one key is serialized. - self.depth_decoder.model.embed_tokens.weight = self.embed_audio_tokens.weight - - self._codec_disallowed_mask: torch.Tensor | None = None - self._batch_state: list[dict[str, Any]] | None = None - self._batch_state_spans: Sequence[tuple[int, int]] | None = None - - # ------------------------------------------------------------------ - # Embedding and backbone hooks - # ------------------------------------------------------------------ - - def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: - """Embed complete generated codec frames; scalar ids are invalid.""" - if input_ids.ndim < 1 or input_ids.shape[-1] != self.num_codebooks: - raise ValueError( - "Breeze embed_input_ids requires complete codec frames with shape " - f"(..., {self.num_codebooks}), got {tuple(input_ids.shape)}" - ) - codes = input_ids.to(dtype=torch.long) - offsets = self.audio_token_offsets.to(device=codes.device) - ids = codes.clamp(0, self.codebook_vocab_size - 1) + offsets - return self.embed_audio_tokens(ids).sum(dim=-2) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **_: Any, - ) -> torch.Tensor | IntermediateTensors: - return self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - ) - - # ------------------------------------------------------------------ - # Prompt/text conditioning - # ------------------------------------------------------------------ - - def _encode_text_segments( - self, - input_ids: torch.Tensor, - text_mask: torch.Tensor, - text_lengths: torch.Tensor, - ) -> torch.Tensor: - text_ids = input_ids[text_mask.to(device=input_ids.device, dtype=torch.bool)].to(torch.long) - lengths = [int(x) for x in text_lengths.reshape(-1).tolist() if int(x) > 0] - if int(text_ids.numel()) != sum(lengths): - raise ValueError(f"Breeze text mask/length mismatch: {text_ids.numel()} != {sum(lengths)}") - if not lengths: - return input_ids.new_empty((0, self.hidden_size), dtype=self.text_encoder.dtype) - segments = list(torch.split(text_ids, lengths)) - padded = pad_sequence(segments, batch_first=True, padding_value=0) - attn = torch.zeros_like(padded, dtype=torch.bool) - for row, length in enumerate(lengths): - attn[row, :length] = True - hidden = self.text_encoder(padded, attention_mask=attn) - projected = self.text_encoder_proj(hidden) - return torch.cat([projected[row, :length] for row, length in enumerate(lengths)], dim=0) - - def _build_prompt_embeddings(self, prompt_ids: torch.Tensor, info: Mapping[str, Any]) -> torch.Tensor: - device = next(self.parameters()).device - prompt_ids = prompt_ids.reshape(-1).to(device=device, dtype=torch.long) - text_mask = torch.as_tensor(info.get("text_ids_mask"), device=device, dtype=torch.bool).reshape(-1) - text_lengths = torch.as_tensor(info.get("text_ids_len"), device=device, dtype=torch.long) - if text_mask.numel() != prompt_ids.numel(): - raise ValueError("Breeze text_ids_mask length does not match prompt length") - output = torch.zeros( - (prompt_ids.numel(), self.hidden_size), - device=device, - dtype=self.embed_text_tokens.weight.dtype, - ) - text_embeds = self._encode_text_segments(prompt_ids, text_mask, text_lengths) - output[text_mask] = text_embeds.to(dtype=output.dtype) - - audio_mask = prompt_ids == self.audio_token_id - input_values = info.get("input_values") - if bool(audio_mask.any()): - if input_values is None: - raise ValueError("Breeze audio placeholders require input_values") - codes = torch.as_tensor(input_values, device=device, dtype=torch.long) - if codes.ndim != 2 or codes.shape[-1] != self.num_codebooks: - raise ValueError(f"Breeze input_values must be (T, {self.num_codebooks})") - if int(codes.shape[0]) != int(audio_mask.sum()): - raise ValueError("Breeze reference code count does not match audio placeholders") - output[audio_mask] = self.embed_input_ids(codes).to(dtype=output.dtype) - - eos_mask = prompt_ids == self.audio_eos_token_id - if bool(eos_mask.any()): - eos_codes = torch.full( - (1, self.num_codebooks), - self.codebook_eos_token_id, - dtype=torch.long, - device=device, - ) - output[eos_mask] = self.embed_input_ids(eos_codes).expand(int(eos_mask.sum()), -1) - return output - - def preprocess( - self, - input_ids: torch.Tensor, - input_embeds: torch.Tensor | None, - **info_dict: Any, - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Build prefill embeddings and consume one completed frame per decode.""" - del input_embeds - info = dict(info_dict) - span_len = int(input_ids.shape[0]) - is_prefill = bool(info.get("_omni_is_prefill", span_len > 1)) - if is_prefill: - full_ids = info.get("prompt_ids") - if not isinstance(full_ids, torch.Tensor): - raise RuntimeError("Breeze prefill is missing prompt_ids") - full_embeds = info.get("breeze_prompt_embeds") - if not isinstance(full_embeds, torch.Tensor): - full_embeds = self._build_prompt_embeddings(full_ids, info).detach().cpu().contiguous() - offset = int(info.get("breeze_prefill_offset", info.get("_omni_num_computed_tokens", 0)) or 0) - take = full_embeds[offset : offset + span_len].to(device=input_ids.device) - if take.shape[0] != span_len: - raise RuntimeError("Breeze prompt embedding slice is shorter than scheduled span") - return ( - torch.full_like(input_ids, self.text_pad_token_id), - take, - { - "breeze_prompt_embeds": full_embeds, - "breeze_prefill_offset": offset + span_len, - "breeze_generated_frames": int(info.get("breeze_generated_frames", 0) or 0), - "breeze_max_new_frames": int(info.get("breeze_max_new_frames", -1)), - }, - ) - - frame = info.get("breeze_current_frame") - if isinstance(frame, torch.Tensor) and frame.numel() == self.num_codebooks: - frame = frame.to(device=input_ids.device, dtype=torch.long).reshape(1, 1, -1) - embeds = self.embed_input_ids(frame).reshape(1, -1) - safe_id = frame[..., 0].reshape_as(input_ids).to(dtype=input_ids.dtype) - return safe_id, embeds, {} - has_breeze_state = any( - key in info - for key in ( - "breeze_current_frame", - "breeze_audio_codes", - "breeze_generated_frames", - "breeze_prompt_embeds", - ) - ) - if has_breeze_state: - # A real decode must consume the full frame produced from the - # preceding backbone hidden state. A scalar fallback would silently - # fill codebooks 1..15 with EOS and corrupt AR state. - raise RuntimeError( - "Breeze decode is missing breeze_current_frame; refusing scalar fallback " - "because it would fill codebooks 1..15 with EOS" - ) - # Runner-only dummy execution has no Breeze codec state. It must not - # invent a semantically invalid audio frame; a stable zero embedding is - # sufficient for shape/profiling warmup. - dummy_embeds = torch.zeros( - (*input_ids.shape, self.hidden_size), - device=input_ids.device, - dtype=self.embed_audio_tokens.weight.dtype, - ) - return input_ids, dummy_embeds, {} - - def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: - if hidden_states.numel() == 0: - return {} - return {"hidden_states": {"last": hidden_states[-1].detach()}} - - # ------------------------------------------------------------------ - # Main head, depth decoder and Omni output - # ------------------------------------------------------------------ - - def compute_logits( - self, - hidden_states: torch.Tensor | OmniOutput, - sampling_metadata: SamplingMetadata | None = None, - ) -> torch.Tensor | None: - del sampling_metadata - if isinstance(hidden_states, OmniOutput): - hidden_states = hidden_states.text_hidden_states - if hidden_states is None: - return None - logits = self._main_head_logits(hidden_states) - - # ``make_omni_output`` already selected codebook 0 and completed its - # frame. Return that same decision to the scheduler, including EOS. - # A finite zero survives the sampler's repetition/temperature passes - # without applying the repetition penalty twice or changing the frame. - # These logits describe deterministic codec selection, not model - # confidence; sampled-token logprobs are consequently zero. - states = self._batch_state or [] - if states: - for state, row_start, row_end in _iter_request_rows( - states, - getattr(self, "_batch_state_spans", None), - int(logits.shape[0]), - ): - selected_code0 = state.get("breeze_selected_code0") - if selected_code0 is not None: - row = logits[row_start:row_end] - row.fill_(float("-inf")) - row[:, int(selected_code0)] = 0.0 - return logits - - def _main_head_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - logits = self.logits_processor(self.lm_head, hidden_states) - if logits is None: - raise RuntimeError("Breeze main head returned no logits") - if self._codec_disallowed_mask is None or self._codec_disallowed_mask.device != logits.device: - mask = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) - if self.codebook_size < self.codebook_vocab_size: - mask[self.codebook_size : self.codebook_vocab_size] = True - self._codec_disallowed_mask = mask - return logits.masked_fill(self._codec_disallowed_mask, float("-inf")) - - @torch.no_grad() - def _generate_depth_codes(self, backbone_hidden: torch.Tensor, code0: int) -> torch.Tensor: - sequence = torch.tensor([[0, code0]], dtype=torch.long, device=backbone_hidden.device) - generated = [int(code0)] - for _ in range(self.num_codebooks - 1): - depth_out = self.depth_decoder( - input_ids=sequence, - backbone_last_hidden_state=backbone_hidden.reshape(1, -1), - use_cache=False, - ) - logits = depth_out.logits[:, -1, :].clone() - if self.codebook_size < logits.shape[-1]: - logits[:, self.codebook_size :] = float("-inf") - next_code = int(torch.argmax(logits, dim=-1).item()) - generated.append(next_code) - sequence = torch.cat([sequence, torch.tensor([[next_code]], device=sequence.device)], dim=1) - return torch.tensor(generated, dtype=torch.long, device=backbone_hidden.device).unsqueeze(0) - - @staticmethod - def _accumulate_codes(info: dict[str, Any], frame: torch.Tensor) -> torch.Tensor: - """Append one ``(1, Q)`` frame to the request's cumulative codes.""" - accumulated = info.get("breeze_audio_codes") - accumulated = ( - torch.cat([accumulated.to(frame.device), frame], dim=0) - if isinstance(accumulated, torch.Tensor) and accumulated.numel() - else frame - ) - info["breeze_audio_codes"] = accumulated - return accumulated - - def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: - if isinstance(model_outputs, OmniOutput): - self._batch_state = None - self._batch_state_spans = None - return model_outputs - hidden = model_outputs - infos = kwargs.get("model_intermediate_buffer", []) - spans = kwargs.get("request_token_spans") - info_items = _normalize_request_infos(infos) - if info_items is None: - self._batch_state = None - self._batch_state_spans = None - return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) - # The runner invokes compute_logits after this method, so retain the - # same request ordering and spans for per-request EOS masking. - self._batch_state = info_items - self._batch_state_spans = spans - per_request_codes: list[torch.Tensor] = [] - # The runner passes the same per-request sampling metadata to this - # hook and compute_logits, including offline SamplingParams overrides. - # When no_penalties is true vLLM does not initialize this GPU buffer. - sampling_metadata = kwargs.get("sampling_metadata") - repetition_penalties = None - if sampling_metadata is not None and not sampling_metadata.no_penalties: - repetition_penalties = sampling_metadata.repetition_penalties.tolist() - - for index, info in enumerate(info_items): - if spans is not None and index < len(spans): - row_start = max(0, int(spans[index][0])) - row_end = min(int(spans[index][1]), int(hidden.shape[0])) - elif len(info_items) == 1: - row_start, row_end = 0, int(hidden.shape[0]) - else: - raise RuntimeError("Breeze make_omni_output requires request_token_spans for batched decode") - if row_end <= row_start: - per_request_codes.append(hidden.new_empty((0, self.num_codebooks), dtype=torch.long)) - continue - row_hidden = hidden[row_end - 1] - if self._golden_dump_dir and bool(info.get("_omni_is_prefill", False)): - self._dump_golden_prefill(info, row_hidden) - main_logits = self._main_head_logits(row_hidden.reshape(1, -1)).float() - penalty = float(repetition_penalties[index]) if repetition_penalties is not None else 1.0 - if not math.isfinite(penalty) or penalty <= 0: - raise ValueError("Breeze repetition_penalty must be finite and positive") - # Match upstream: penalize each previously generated codebook-0 - # id once, excluding the text prompt and reference audio codes. - history = info.setdefault("breeze_generated_code0_ids", []) - if penalty != 1.0 and history: - history_ids = torch.tensor(history, device=main_logits.device, dtype=torch.long) - scores = main_logits[0, history_ids] - main_logits[0, history_ids] = torch.where(scores < 0, scores * penalty, scores / penalty) - code0 = int(torch.argmax(main_logits, dim=-1).item()) - generated_frames = int(info.get("breeze_generated_frames", 0) or 0) - max_new_frames = int(info.get("breeze_max_new_frames", -1)) - terminal = code0 == self.codebook_vocab_size or (max_new_frames > 0 and generated_frames >= max_new_frames) - info["breeze_selected_code0"] = self.codebook_vocab_size if terminal else code0 - if terminal: - info["breeze_force_eos"] = True - existing = info.get("breeze_audio_codes") - per_request_codes.append( - existing - if isinstance(existing, torch.Tensor) and existing.numel() and not self._async_chunk - else hidden.new_empty((0, self.num_codebooks), dtype=torch.long) - ) - continue - if code0 < 0 or code0 >= self.codebook_size: - raise RuntimeError(f"Breeze main head produced invalid codec id {code0}") - - frame = self._generate_depth_codes(row_hidden, code0) - history.append(code0) - info["breeze_current_frame"] = frame[0].detach() - info["breeze_generated_frames"] = generated_frames + 1 - info["breeze_force_eos"] = False - if self._async_chunk: - # Streaming consumers only ever read the per-frame tail. The - # cumulative buffer exists for full-payload transport and the - # opt-in golden dump; maintaining it here would also cost an - # O(T^2) tensor-cat per request. - if self._golden_dump_dir: - self._accumulate_codes(info, frame) - per_request_codes.append(frame) - else: - per_request_codes.append(self._accumulate_codes(info, frame)) - if self._golden_dump_dir: - for info in info_items: - self._dump_golden_frames(info) - return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"codes": {"audio": per_request_codes}}) - - def _dump_golden_frames(self, info: dict[str, Any]) -> None: - """Write terminal codec frames for opt-in golden-parity diagnostics.""" - codes = info.get("breeze_audio_codes") - if not info.get("breeze_force_eos", False) or not isinstance(codes, torch.Tensor) or codes.numel() == 0: - return - prompt_ids = info.get("prompt_ids") - if not isinstance(prompt_ids, torch.Tensor): - raise RuntimeError("Breeze golden dump is missing prompt_ids") - digest = hashlib.sha256(prompt_ids.detach().cpu().numpy().tobytes()).hexdigest()[:16] - output_dir = Path(self._golden_dump_dir) # type: ignore[arg-type] - output_dir.mkdir(parents=True, exist_ok=True) - torch.save( - { - "template": info.get("template"), - "prompt_ids": prompt_ids.cpu(), - "codes": codes.detach().cpu().to(torch.long).contiguous(), - "generated_frames": int(info.get("breeze_generated_frames", 0) or 0), - "finished": True, - }, - output_dir / f"breeze_{digest}.pt", - ) - - def _dump_golden_prefill(self, info: dict[str, Any], hidden: torch.Tensor) -> None: - """Write the terminal-prompt hidden row for opt-in parity diagnostics.""" - prompt_ids = info.get("prompt_ids") - if not isinstance(prompt_ids, torch.Tensor): - raise RuntimeError("Breeze golden dump is missing prompt_ids") - digest = hashlib.sha256(prompt_ids.detach().cpu().numpy().tobytes()).hexdigest()[:16] - output_dir = Path(self._golden_dump_dir) # type: ignore[arg-type] - output_dir.mkdir(parents=True, exist_ok=True) - torch.save( - { - "template": info.get("template"), - "prompt_ids": prompt_ids.cpu(), - "hidden": hidden.detach().reshape(-1).cpu().to(torch.float32).contiguous(), - }, - output_dir / f"prefill_{digest}.pt", - ) - - def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: - """Drop cross-step references retained for the finished AR requests.""" - del finished_req_ids - # Codec/frame state lives in request-local additional_information and is - # freed with the Request. Only the short logits-to-state mapping is - # model-owned and must not leak into the next micro-batch. - self._batch_state = None - self._batch_state_spans = None - - # ------------------------------------------------------------------ - # Weight loading - # ------------------------------------------------------------------ - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loaded: set[str] = set() - params = dict(self.named_parameters()) - backbone_weights: list[tuple[str, torch.Tensor]] = [] - text_weights: list[tuple[str, torch.Tensor]] = [] - depth_weights: list[tuple[str, torch.Tensor]] = [] - for name, tensor in weights: - # This table belongs to Breeze's custom audio embedding module, - # although its checkpoint key is nested under backbone_model. - # Handle it before forwarding the remaining backbone weights to - # Qwen3Model. - if name == "backbone_model.embed_tokens.embed_audio_tokens.weight": - target = "embed_audio_tokens.weight" - param = params.get(target) - if param is None: - raise ValueError(f"Missing Breeze parameter: {target}") - _check_weight_shape(name, tensor, tuple(param.shape)) - default_weight_loader(param, tensor) - loaded.add(target) - continue - if name.startswith("backbone_model."): - # Qwen3Model.load_weights() is called on the nested module and - # therefore expects names relative to that module (the same - # convention used by MOSS-TTS), not the outer ``model.`` - # namespace used by ``named_parameters()``. - backbone_weights.append((name[len("backbone_model.") :], tensor)) - continue - if name.startswith("text_encoder."): - text_name = name[len("text_encoder.") :] - text_weights.append((text_name, tensor)) - continue - if name.startswith("depth_decoder."): - depth_weights.append((name[len("depth_decoder.") :], tensor)) - continue - if name in {"lm_head.weight", "embed_text_tokens.weight"} or name.startswith("text_encoder_proj."): - target = name - else: - continue - param = params.get(target) - if param is None: - continue - # ParallelLMHead stores a TP-local shard while checkpoint tensors - # are global; validate against the global Breeze head shape. - expected_shape = ( - (self.codebook_vocab_size + 1, self.hidden_size) if name == "lm_head.weight" else tuple(param.shape) - ) - _check_weight_shape(name, tensor, expected_shape) - loader = getattr(param, "weight_loader", default_weight_loader) - loader(param, tensor) - loaded.add(target) - - if backbone_weights: - loaded.update("model." + key for key in self.model.load_weights(backbone_weights)) - if text_weights: - loaded.update("text_encoder." + key for key in self.text_encoder.load_weights(text_weights)) - if depth_weights: - loaded.update("depth_decoder." + key for key in self.depth_decoder.load_weights(depth_weights)) - if bool(getattr(self.config, "tie_codebooks_embeddings", False)): - # The checkpoint serializes only one side of the tied table. - # ``named_parameters(remove_duplicate=True)`` exposes it as the - # talker-owned ``embed_audio_tokens.weight``, so report both - # aliases to the strict weight tracker. - if "depth_decoder.model.embed_tokens.weight" in loaded: - loaded.add("embed_audio_tokens.weight") - if "embed_audio_tokens.weight" in loaded: - loaded.add("depth_decoder.model.embed_tokens.weight") - return loaded - - -__all__ = ["BreezeTTS2TalkerForGeneration"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py deleted file mode 100644 index 0276e8278..000000000 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py +++ /dev/null @@ -1,350 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Breeze T5Gemma2 text encoder. - -The Breeze checkpoint contains a T5Gemma2 encoder, not the older T5Gemma -implementation used by some diffusion models. This module keeps the encoder -local to Breeze so checkpoint names and attention semantics stay explicit, -while using vLLM tensor-parallel layers for the expensive projections. -""" - -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -import torch -import torch.nn.functional as F -from torch import nn -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import get_act_fn -from vllm.model_executor.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import default_weight_loader - - -class _RMSNorm(nn.Module): - """T5Gemma2 RMSNorm, whose checkpoint stores a zero-centered weight.""" - - def __init__(self, dim: int, eps: float) -> None: - super().__init__() - self.weight = nn.Parameter(torch.zeros(dim)) - self.eps = float(eps) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - normalized = hidden_states.float() * torch.rsqrt( - hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + self.eps - ) - return (normalized * (1.0 + self.weight.float())).to(hidden_states.dtype) - - -class _RotaryEmbedding(nn.Module): - def __init__(self, config: Any) -> None: - super().__init__() - self.config = config - for layer_type in dict.fromkeys(config.layer_types): - params = dict(config.rope_parameters.get(layer_type, {})) - base = float(params.get("rope_theta", 10_000.0)) - factor = float(params.get("factor", 1.0)) - dim = int(config.head_dim) - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - if params.get("rope_type") == "linear": - inv_freq = inv_freq / factor - elif params.get("rope_type", "default") != "default": - raise ValueError(f"Unsupported T5Gemma2 rope type: {params.get('rope_type')!r}") - self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False) - - def forward( - self, - hidden_states: torch.Tensor, - position_ids: torch.Tensor, - layer_type: str, - ) -> tuple[torch.Tensor, torch.Tensor]: - inv_freq = getattr(self, f"{layer_type}_inv_freq") - expanded = inv_freq[None, :, None].to(device=hidden_states.device) - positions = position_ids[:, None, :].float() - freqs = (expanded.float() @ positions).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - return emb.cos().to(hidden_states.dtype), emb.sin().to(hidden_states.dtype) - - -def _rotate_half(value: torch.Tensor) -> torch.Tensor: - first, second = value[..., : value.shape[-1] // 2], value[..., value.shape[-1] // 2 :] - return torch.cat((-second, first), dim=-1) - - -def _apply_rotary( - query: torch.Tensor, - key: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - cos = cos.unsqueeze(1) - sin = sin.unsqueeze(1) - return (query * cos) + (_rotate_half(query) * sin), (key * cos) + (_rotate_half(key) * sin) - - -def _repeat_kv(hidden_states: torch.Tensor, repeats: int) -> torch.Tensor: - if repeats == 1: - return hidden_states - batch, heads, seq_len, head_dim = hidden_states.shape - return ( - hidden_states[:, :, None, :, :] - .expand(batch, heads, repeats, seq_len, head_dim) - .reshape(batch, heads * repeats, seq_len, head_dim) - ) - - -def _attention_mask( - padding_mask: torch.Tensor | None, - *, - batch_size: int, - seq_len: int, - layer_type: str, - sliding_window: int | None, - dtype: torch.dtype, - device: torch.device, -) -> torch.Tensor | None: - if layer_type == "full_attention" and padding_mask is None: - return None - valid = ( - torch.ones((batch_size, seq_len), device=device, dtype=torch.bool) - if padding_mask is None - else padding_mask.to(device=device, dtype=torch.bool) - ) - allowed = valid[:, None, None, :].expand(batch_size, 1, seq_len, seq_len) - if layer_type == "sliding_attention": - if sliding_window is None: - raise ValueError("T5Gemma2 sliding_attention requires sliding_window") - q_idx = torch.arange(seq_len, device=device)[:, None] - kv_idx = torch.arange(seq_len, device=device)[None, :] - left = (int(sliding_window) + 1) // 2 - right = int(sliding_window) // 2 + 1 - distance = q_idx - kv_idx - local = ((distance >= 0) & (distance < left)) | ((distance < 0) & (-distance < right)) - allowed = allowed & local[None, None, :, :] - mask = torch.zeros((batch_size, 1, seq_len, seq_len), device=device, dtype=dtype) - mask.masked_fill_(~allowed, torch.finfo(dtype).min) - return mask - - -class _SelfAttention(nn.Module): - def __init__(self, config: Any) -> None: - super().__init__() - tp_size = get_tensor_model_parallel_world_size() - if config.num_attention_heads % tp_size != 0: - raise ValueError("T5Gemma2 attention heads must be divisible by tensor parallel size") - self.layer_type = config.layer_type - self.head_dim = int(config.head_dim) - self.num_heads = int(config.num_attention_heads) // tp_size - self.num_kv_heads = max(1, int(config.num_key_value_heads) // tp_size) - self.num_key_value_groups = self.num_heads // self.num_kv_heads - self.scaling = float(config.query_pre_attn_scalar) ** -0.5 - self.qkv_proj = QKVParallelLinear( - hidden_size=config.hidden_size, - head_size=self.head_dim, - total_num_heads=config.num_attention_heads, - total_num_kv_heads=config.num_key_value_heads, - bias=config.attention_bias, - ) - self.o_proj = RowParallelLinear( - input_size=config.num_attention_heads * self.head_dim, - output_size=config.hidden_size, - bias=config.attention_bias, - input_is_parallel=True, - ) - self.q_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) - self.k_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) - self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q_size = self.num_heads * self.head_dim - kv_size = self.num_kv_heads * self.head_dim - query, key, value = qkv.split((q_size, kv_size, kv_size), dim=-1) - shape = (*hidden_states.shape[:-1], -1, self.head_dim) - query = self.q_norm(query.view(shape).transpose(1, 2)) - key = self.k_norm(key.view(shape).transpose(1, 2)) - value = value.view(shape).transpose(1, 2) - query, key = _apply_rotary(query, key, *position_embeddings) - key = _repeat_kv(key, self.num_key_value_groups) - value = _repeat_kv(value, self.num_key_value_groups) - # SDPA scales by head_dim**-0.5; adjust Q to T5Gemma2's scalar. - query = query * (self.scaling / (self.head_dim**-0.5)) - output = F.scaled_dot_product_attention( - query, - key, - value, - attn_mask=attention_mask, - dropout_p=0.0, - ) - output = output.transpose(1, 2).contiguous().reshape(*hidden_states.shape[:-1], -1) - output, _ = self.o_proj(output) - return output - - -class _MLP(nn.Module): - def __init__(self, config: Any) -> None: - super().__init__() - self.gate_up_proj = MergedColumnParallelLinear( - input_size=config.hidden_size, - output_sizes=[config.intermediate_size, config.intermediate_size], - bias=False, - ) - self.down_proj = RowParallelLinear( - input_size=config.intermediate_size, - output_size=config.hidden_size, - bias=False, - input_is_parallel=True, - ) - self.act_fn = get_act_fn(config.hidden_activation) - self.dropout = nn.Dropout(float(config.dropout_rate)) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - gate_up, _ = self.gate_up_proj(hidden_states) - gate, up = gate_up.chunk(2, dim=-1) - output, _ = self.down_proj(self.dropout(self.act_fn(gate) * up)) - return output - - -class _EncoderLayer(nn.Module): - def __init__(self, config: Any, layer_idx: int) -> None: - super().__init__() - layer_config = type("LayerConfig", (), dict(vars(config)))() - layer_config.layer_type = config.layer_types[layer_idx] - self.attention_type = layer_config.layer_type - self.pre_self_attn_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) - self.self_attn = _SelfAttention(layer_config) - self.post_self_attn_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) - self.pre_feedforward_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) - self.mlp = _MLP(config) - self.post_feedforward_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) - self.dropout = nn.Dropout(float(config.dropout_rate)) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None, - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.pre_self_attn_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, position_embeddings, attention_mask) - hidden_states = residual + self.dropout(self.post_self_attn_layernorm(hidden_states)) - residual = hidden_states - hidden_states = self.pre_feedforward_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + self.dropout(self.post_feedforward_layernorm(hidden_states)) - return hidden_states - - -class BreezeTTS2TextEncoder(nn.Module): - """Frozen T5Gemma2 encoder returning ``(B, S, hidden_size)`` states.""" - - def __init__(self, config: Any) -> None: - super().__init__() - self.config = config - self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) - self.embed_scale = float(config.hidden_size**0.5) - self.eoi_token_index = int(getattr(config, "eoi_token_index", 256000)) - # T5Gemma2 stores this learned replacement vector under the input - # embedding module (``embed_tokens.eoi_embedding``), so keep the same - # parameter path for checkpoint loading and tied-weight inspection. - self.embed_tokens.eoi_embedding = nn.Parameter(torch.zeros(config.hidden_size)) - self.layers = nn.ModuleList([_EncoderLayer(config, i) for i in range(config.num_hidden_layers)]) - self.norm = _RMSNorm(config.hidden_size, config.rms_norm_eps) - self.rotary_emb = _RotaryEmbedding(config) - self.dropout = nn.Dropout(float(config.dropout_rate)) - - @property - def dtype(self) -> torch.dtype: - return next(self.parameters()).dtype - - def forward( - self, - input_ids: torch.Tensor, - attention_mask: torch.Tensor | None = None, - ) -> torch.Tensor: - if input_ids.ndim != 2: - raise ValueError(f"Breeze text encoder expects 2D input_ids, got {tuple(input_ids.shape)}") - if attention_mask is not None and attention_mask.shape != input_ids.shape: - raise ValueError( - "Breeze text encoder attention_mask must match input_ids: " - f"{tuple(attention_mask.shape)} != {tuple(input_ids.shape)}" - ) - hidden_states = self.embed_tokens(input_ids) * self.embed_scale - eoi_mask = input_ids == self.eoi_token_index - if bool(eoi_mask.any()): - hidden_states = torch.where( - eoi_mask.unsqueeze(-1), - self.embed_tokens.eoi_embedding.to(dtype=hidden_states.dtype), - hidden_states, - ) - hidden_states = self.dropout(hidden_states) - positions = torch.arange(input_ids.shape[1], device=input_ids.device, dtype=torch.long).unsqueeze(0) - position_embeddings = { - layer_type: self.rotary_emb(hidden_states, positions, layer_type) - for layer_type in set(self.config.layer_types) - } - masks = { - layer_type: _attention_mask( - attention_mask, - batch_size=input_ids.shape[0], - seq_len=input_ids.shape[1], - layer_type=layer_type, - sliding_window=self.config.sliding_window, - dtype=hidden_states.dtype, - device=input_ids.device, - ) - for layer_type in set(self.config.layer_types) - } - for layer in self.layers: - hidden_states = layer( - hidden_states, - position_embeddings[layer.attention_type], - masks[layer.attention_type], - ) - return self.dropout(self.norm(hidden_states)) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = ( - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ) - params = dict(self.named_parameters()) - loaded: set[str] = set() - for name, tensor in weights: - target = name - shard_id: str | int | None = None - for packed_name, source_name, source_id in stacked_params_mapping: - marker = f".{source_name}." - if marker in name: - target = name.replace(marker, f".{packed_name}.") - shard_id = source_id - break - parameter = params.get(target) - if parameter is None: - continue - loader = getattr(parameter, "weight_loader", default_weight_loader) - if shard_id is None: - loader(parameter, tensor) - else: - loader(parameter, tensor, shard_id) - loaded.add(target) - return loaded - - -__all__ = ["BreezeTTS2TextEncoder"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py index 2b538df01..c52d34a6c 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py @@ -1,21 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Synchronous Breeze-TTS-2 pipeline: talker -> Mimi codec.""" - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -_PROC = "vllm_omni.model_executor.stage_input_processors.breeze_tts_2" +from vllm_omni.config.stage_config import PipelineConfig, StageExecutionType, StagePipelineConfig BREEZE_TTS_2_PIPELINE = PipelineConfig( model_type="breeze", - default_deploy_config_name="breeze_tts_2.yaml", model_arch="BreezeForConditionalGeneration", - hf_architectures=("BreezeForConditionalGeneration",), + default_deploy_config_name="breeze_tts_2.yaml", stages=( StagePipelineConfig( stage_id=0, @@ -24,34 +15,24 @@ BREEZE_TTS_2_PIPELINE = PipelineConfig( input_sources=(), owns_tokenizer=True, engine_output_type="latent", - async_chunk_process_next_stage_input_func=f"{_PROC}.talker2codec_async_chunk", - custom_process_next_stage_input_func=f"{_PROC}.talker2codec_full_payload", - sampling_constraints={ - "detokenize": False, - # Text EOS 1 is a valid codec token; only the backbone's - # explicit codec EOS below may terminate audio generation. - "ignore_eos": True, - # The talker has already selected the emitted frame's token; - # min_tokens must not mask that token in the sampler. - "min_tokens": 0, - "stop_token_ids": [2051], - }, + scheduler_cls="vllm_omni.core.sched.omni_cfg_ar_scheduler.OmniCFGARScheduler", + prompt_expand_func="vllm_omni.model_executor.stage_input_processors.breeze_tts_2.expand_cfg_prompts", + async_chunk_process_next_stage_input_func=( + "vllm_omni.model_executor.stage_input_processors.breeze_tts_2.talker2code2wav_async_chunk" + ), + sampling_constraints={"detokenize": False, "stop_token_ids": [2051], "ignore_eos": True, "min_tokens": 0}, ), StagePipelineConfig( stage_id=1, model_stage="breeze_tts_2_codec", + model_arch="BreezeCode2Wav", execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), final_output=True, final_output_type="audio", engine_output_type="audio", - model_arch="BreezeTTS2MimiCodec", - sync_process_input_func=f"{_PROC}.talker2codec", - sampling_constraints={"detokenize": True}, - requires_full_payload_input=True, retains_state_across_chunks=True, + sampling_constraints={"detokenize": True}, ), ), ) - -__all__ = ["BREEZE_TTS_2_PIPELINE"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/prompt.py b/vllm_omni/model_executor/models/breeze_tts_2/prompt.py new file mode 100644 index 000000000..2af9a4854 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/prompt.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Breeze's instruction and reference-audio templates for both public APIs.""" + +import math +from typing import Any + +import numpy as np +import torch +from transformers import PreTrainedTokenizerBase +from vllm.inputs import TokensPrompt, tokens_input +from vllm.multimodal.audio import AudioResampler + +DEFAULT_INSTRUCTION = "Speak clearly and naturally." +SAMPLE_RATE = 24000 +ENCODE_DOWNSAMPLE_RATE = 1920 +CFG_UNCOND_SUFFIX = "__cfg_uncond" + + +def build_breeze_prompt( + tokenizer: PreTrainedTokenizerBase, + text: str, + instructions: str = DEFAULT_INSTRUCTION, + *, + speaker: str = "S0", + ref_audio: tuple[np.ndarray, int] | None = None, + ref_text: str | None = None, + guidance_scale: float = 1.0, + temperature: float = 0.9, + top_k: int = 50, + top_p: float = 1.0, + repetition_penalty: float = 1.1, +) -> TokensPrompt: + if not text.strip(): + raise ValueError("Breeze input text cannot be empty") + if not math.isfinite(guidance_scale) or guidance_scale <= 0: + raise ValueError("Breeze guidance_scale must be finite and positive") + if top_k == -1: + top_k = 0 + if ( + not all(math.isfinite(value) for value in (temperature, top_p, repetition_penalty)) + or temperature < 0 + or top_k < 0 + or not 0 < top_p <= 1 + or repetition_penalty <= 0 + ): + raise ValueError("Invalid Breeze sampling parameters") + if (ref_audio is None) != (ref_text is None): + raise ValueError("Breeze voice cloning requires both ref_audio and ref_text") + speaker_tag = speaker if speaker.startswith("[") and speaker.endswith("]") else f"[{speaker}]" + target = tokenizer.encode(f"{speaker_tag}{instructions}{text}", add_special_tokens=True) + conditioning: dict[str, Any] = {"target_ids": target, "guidance_scale": guidance_scale, "role": "cond"} + prefix_length = 0 + if ref_audio is not None: + if not ref_text or not ref_text.strip(): + raise ValueError("Breeze reference transcript cannot be empty") + waveform, sample_rate = ref_audio + waveform = np.asarray(waveform, dtype=np.float32) + if waveform.ndim == 2: + # Public waveform input follows soundfile: (samples, channels). + waveform = waveform.mean(axis=1) + if waveform.ndim != 1 or not waveform.size or not np.isfinite(waveform).all() or sample_rate <= 0: + raise ValueError("Breeze reference audio must contain finite samples at a positive sample rate") + if sample_rate != SAMPLE_RATE: + # Match the released tokenizer's librosa/soxr_hq resampling, + # including ceil-length handling for non-integral rate ratios. + length = math.ceil(waveform.size * SAMPLE_RATE / sample_rate) + waveform = AudioResampler(target_sr=SAMPLE_RATE, method="soxr").resample(waveform, orig_sr=sample_rate) + waveform = np.pad(waveform, (0, max(0, length - waveform.size)))[:length] + reference_ids = tokenizer.encode(f"{speaker_tag}{ref_text}", add_special_tokens=True) + frames = math.ceil(waveform.size / ENCODE_DOWNSAMPLE_RATE) + conditioning.update( + reference_ids=reference_ids, + reference_frames=frames, + ) + prefix_length = len(reference_ids) + frames + 1 + if guidance_scale != 1.0: + conditioning["negative_ids"] = tokenizer.encode(f"{speaker_tag}{text}", add_special_tokens=True) + prompt = tokens_input(prompt_token_ids=[0] * (prefix_length + len(target))) + prompt["additional_information"] = { + "breeze_prompt": conditioning, + "breeze_sampling": { + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + "repetition_penalty": repetition_penalty, + }, + } + if ref_audio is not None: + # Tensor payloads belong at the transport root (or in a declared + # tensor namespace), rather than inside JSON-valued model metadata. + prompt["additional_information"]["reference_waveform"] = torch.from_numpy(np.ascontiguousarray(waveform)) + if guidance_scale != 1.0: + prompt["additional_information"]["cfg_group"] = {"role": "cond", "uncond_suffix": CFG_UNCOND_SUFFIX} + return prompt diff --git a/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py b/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py deleted file mode 100644 index 32c8145bb..000000000 --- a/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py +++ /dev/null @@ -1,336 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Breeze-TTS-2 request-to-prompt conversion. - -This module is intentionally limited to request preparation. It creates the -real prompt token ids needed by the scheduler and stores the text/audio masks -and reference codes in ``additional_information``. The stage-0 model later -turns those values into embeddings; this keeps tokenization independent from -the Qwen3/text-encoder implementation and avoids changing sequence length in -the model ``preprocess`` hook. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Any - -import numpy as np -import torch -from transformers import AutoTokenizer - -from vllm_omni.inputs.data import OmniTokensPrompt - -from .audio_tokenizer import BreezeReferenceAudioTokenizer - -AUDIO_TAG = "<|AUDIO|>" -AUDIO_EOS = "<|audio_eos|>" -INSTRUCTION_BOS = "" -INSTRUCTION_EOS = "" - - -@dataclass(frozen=True) -class _TextSegment: - text: str - - -@dataclass(frozen=True) -class _AudioSegment: - source: Any - sample_rate: int | None - - -class BreezeTTS2PromptBuilder: - """Build one Breeze prompt without depending on model modules. - - ``tokenizer`` is a normal HuggingFace tokenizer. ``reference_audio_encoder`` - is injected so the builder can use the serving worker's cached audio - encoder and never reload it per request. The returned ``OmniTokensPrompt`` - contains scheduler token ids and compact CPU metadata for stage-0 - ``preprocess``. - """ - - def __init__( - self, - tokenizer: Any, - config: Any, - reference_audio_encoder: BreezeReferenceAudioTokenizer | None = None, - ) -> None: - self.tokenizer = tokenizer - self.audio_encoder = reference_audio_encoder - self.audio_token_id = int(config.audio_token_id) - self.audio_eos_token_id = int(config.audio_eos_token_id) - # The Breeze tokenizer vocabulary is larger than the Qwen3 backbone - # vocabulary. ``prompt_token_ids`` are scheduler bookkeeping only, so - # keep them in the backbone vocabulary and retain the real Breeze ids - # in ``additional_information['prompt_ids']`` for the talker. - self.scheduler_token_id = int(config.pad_token_id) - scheduler_vocab_size = int(config.backbone_config.vocab_size) - if self.scheduler_token_id >= scheduler_vocab_size: - raise ValueError( - "prompt_pad_token_id must be smaller than the backbone vocabulary: " - f"{self.scheduler_token_id} >= {scheduler_vocab_size}" - ) - self.num_codebooks = int(config.num_codebooks) - codec_config = getattr(config, "codec_config", None) - if isinstance(codec_config, Mapping): - self.codebook_size = int(codec_config.get("codebook_size", 2048)) - elif codec_config is not None: - self.codebook_size = int(getattr(codec_config, "codebook_size", 2048)) - else: - # ``codec_config`` is optional in BreezeTTS2Config; a minimal - # checkpoint may omit it. Breeze's codec uses 2048-entry - # codebooks, the same fallback the talker and codec stages apply. - self.codebook_size = 2048 - if self.num_codebooks <= 0 or self.codebook_size <= 0: - raise ValueError("num_codebooks and codec codebook_size must be positive") - self._validate_audio_special_tokens() - - @classmethod - def from_pretrained( - cls, - model_path: str, - config: Any, - *, - reference_audio_encoder: BreezeReferenceAudioTokenizer | None = None, - tokenizer_kwargs: Mapping[str, Any] | None = None, - ) -> BreezeTTS2PromptBuilder: - """Load the text tokenizer once and construct a reusable builder.""" - kwargs = dict(tokenizer_kwargs or {}) - kwargs.setdefault("use_fast", True) - kwargs.setdefault("trust_remote_code", True) - tokenizer = AutoTokenizer.from_pretrained(model_path, **kwargs) - return cls(tokenizer, config, reference_audio_encoder) - - def build( - self, - request: Mapping[str, Any], - *, - template: str | None = None, - ) -> OmniTokensPrompt: - """Convert one request to a vLLM-Omni token prompt. - - Supported templates are ``tts_plain``, ``tts_instruction``, - ``ref_clone_tata`` and ``ref_edit_tata``. When omitted, the template - is inferred from the presence of instruction/reference fields. - """ - template_name = template or request.get("template") - if not isinstance(template_name, str) or not template_name: - template_name = _infer_template(request) - segments = self._build_segments(template_name, request) - - # Breeze's reference implementation first adds tokenizer special - # tokens to every text segment, decodes those segments back to text, - # inserts the audio placeholders, and tokenizes the complete string - # once. This matters at segment boundaries: tokenizing each raw text - # fragment independently can produce different boundary tokens. - rendered_segments: list[tuple[str, bool, int]] = [] - reference_codes: list[torch.Tensor] = [] - - for segment in segments: - if isinstance(segment, _TextSegment): - ids = self._encode_one_text(segment.text, add_special_tokens=True) - if not ids: - raise ValueError("Breeze text tokenizer returned an empty segment") - rendered = str(self.tokenizer.decode(ids, skip_special_tokens=False)) - rendered_segments.append((rendered, True, 0)) - continue - - codes = self._resolve_reference_codes(segment) - reference_codes.append(codes) - audio_text = AUDIO_TAG * int(codes.shape[0]) + AUDIO_EOS - rendered_segments.append((audio_text, False, int(codes.shape[0]) + 1)) - - prompt_ids: list[int] = [] - text_mask: list[bool] = [] - text_lengths: list[int] = [] - rendered_prompt = "".join(item[0] for item in rendered_segments) - prompt_ids = self._encode_one_text(rendered_prompt, add_special_tokens=False) - offset = 0 - for rendered, is_text, expected_len in rendered_segments: - segment_ids = self._encode_one_text(rendered, add_special_tokens=False) - actual_len = len(segment_ids) - if not is_text and actual_len != expected_len: - raise ValueError( - "Breeze audio placeholder token length changed after rendering: " - f"expected {expected_len}, got {actual_len}" - ) - if not is_text and segment_ids != [ - *([self.audio_token_id] * (expected_len - 1)), - self.audio_eos_token_id, - ]: - raise ValueError("Breeze tokenizer did not preserve audio placeholder ids") - text_mask.extend([is_text] * actual_len) - if is_text: - text_lengths.append(actual_len) - offset += actual_len - if offset != len(prompt_ids) or len(text_mask) != len(prompt_ids): - raise RuntimeError("Breeze rendered prompt length accounting error") - if not prompt_ids: - raise ValueError("Breeze prompt is empty") - - input_values = torch.cat(reference_codes, dim=0) if reference_codes else None - if input_values is not None and int(input_values.shape[0]) != prompt_ids.count(self.audio_token_id): - raise RuntimeError("reference code count does not match audio placeholder count") - - additional_information: dict[str, Any] = { - # The model runner may invoke preprocess on a chunk of a long - # prompt. Keep one compact CPU copy so the talker can reconstruct - # the full embedding buffer before slicing that chunk. - "prompt_ids": torch.tensor(prompt_ids, dtype=torch.long), - "attention_mask": torch.ones(len(prompt_ids), dtype=torch.bool), - "text_ids_mask": torch.tensor(text_mask, dtype=torch.bool), - "text_ids_len": torch.tensor(text_lengths, dtype=torch.long), - "template": template_name, - } - if input_values is not None: - # ``input_values`` is the name used by the upstream Breeze model. - # Keep the tensor unbatched here: Omni transports one request at a - # time and the talker adds the batch dimension after collation. - additional_information["input_values"] = input_values - additional_information["ref_code_len"] = int(input_values.shape[0]) - - if self.scheduler_token_id < 0: - raise ValueError("prompt_pad_token_id must be non-negative") - - return OmniTokensPrompt( - # Never expose Breeze's high text/audio ids to the Qwen3 scheduler - # or its token-range validation. The real ids remain available in - # additional_information and are consumed by ``preprocess``. - prompt_token_ids=[self.scheduler_token_id] * len(prompt_ids), - additional_information=additional_information, - ) - - def _build_segments(self, template: str, request: Mapping[str, Any]) -> list[_TextSegment | _AudioSegment]: - speaker = _speaker_prefix(request.get("speaker", "S0")) - text = _required_text(request, "text") - instruction = request.get("instruction") - ref_text = request.get("ref_text") - - if template == "tts_plain": - return [_TextSegment(f"{speaker}{text}")] - if template == "tts_instruction": - if not isinstance(instruction, str) or not instruction: - raise ValueError("tts_instruction requires a non-empty instruction") - return [_TextSegment(f"{speaker}{INSTRUCTION_BOS}{instruction}{INSTRUCTION_EOS}{text}")] - if template in ("ref_clone_tata", "ref_edit_tata"): - if not isinstance(ref_text, str) or not ref_text: - raise ValueError(f"{template} requires a non-empty ref_text") - audio = self._audio_segment(request) - if template == "ref_clone_tata": - return [_TextSegment(f"{speaker}{ref_text}"), audio, _TextSegment(f"{speaker}{text}")] - if not isinstance(instruction, str) or not instruction: - raise ValueError("ref_edit_tata requires a non-empty instruction") - return [ - _TextSegment(f"{speaker}{ref_text}"), - audio, - _TextSegment(f"{speaker}{INSTRUCTION_BOS}{instruction}{INSTRUCTION_EOS}{text}"), - ] - raise ValueError( - f"unknown Breeze template {template!r}; expected tts_plain, tts_instruction, " - "ref_clone_tata or ref_edit_tata" - ) - - def _audio_segment(self, request: Mapping[str, Any]) -> _AudioSegment: - source = request.get("ref_audio") - if source is None: - source = request.get("ref_audio_codes") - if source is None: - raw_codes = request.get("codes") - source = raw_codes.get("ref") if isinstance(raw_codes, Mapping) else None - if source is None: - raise ValueError("reference template requires ref_audio or ref_audio_codes") - sample_rate = request.get("ref_audio_sample_rate") - return _AudioSegment(source, int(sample_rate) if sample_rate is not None else None) - - def _resolve_reference_codes(self, segment: _AudioSegment) -> torch.Tensor: - source = segment.source - if _looks_like_codes(source, self.num_codebooks): - codes = torch.as_tensor(source) - return _normalize_codes(codes, self.num_codebooks, self.codebook_size) - if self.audio_encoder is None: - raise RuntimeError("reference audio was provided but no reference_audio_encoder is configured") - codes = self.audio_encoder.encode(source, segment.sample_rate) - return _normalize_codes(codes, self.num_codebooks, self.codebook_size) - - def _encode_one_text(self, text: str, *, add_special_tokens: bool) -> list[int]: - encoded = self.tokenizer(text, add_special_tokens=add_special_tokens, padding=False) - ids = encoded["input_ids"] if isinstance(encoded, Mapping) else encoded.input_ids - if isinstance(ids, torch.Tensor): - ids = ids.tolist() - return _as_int_list(ids) - - def _validate_audio_special_tokens(self) -> None: - convert = self.tokenizer.convert_tokens_to_ids - for token, expected in ((AUDIO_TAG, self.audio_token_id), (AUDIO_EOS, self.audio_eos_token_id)): - actual = convert(token) - unk = getattr(self.tokenizer, "unk_token_id", None) - if actual is None or (unk is not None and actual == unk) or int(actual) != expected: - raise ValueError(f"tokenizer id for {token!r} is {actual}, expected config value {expected}") - - -def _infer_template(request: Mapping[str, Any]) -> str: - has_ref = ( - request.get("ref_audio") is not None - or request.get("ref_audio_codes") is not None - or bool(request.get("ref_text")) - ) - has_instruction = bool(request.get("instruction")) - if has_ref and has_instruction: - return "ref_edit_tata" - if has_ref: - return "ref_clone_tata" - return "tts_instruction" if has_instruction else "tts_plain" - - -def _required_text(request: Mapping[str, Any], key: str) -> str: - value = request.get(key) - if not isinstance(value, str) or not value.strip(): - raise ValueError(f"Breeze request requires non-empty {key}") - return value - - -def _speaker_prefix(value: Any) -> str: - if value in (None, ""): - return "" - value = str(value) - return value if value.startswith("[") and value.endswith("]") else f"[{value}]" - - -def _as_int_list(value: Any) -> list[int]: - if isinstance(value, torch.Tensor): - value = value.tolist() - elif isinstance(value, np.ndarray): - value = value.tolist() - if not isinstance(value, list): - raise ValueError("tokenizer input_ids must be a list") - if value and isinstance(value[0], (list, tuple, np.ndarray, torch.Tensor)): - if len(value) != 1: - raise ValueError("tokenizer input_ids contains multiple sequences") - return _as_int_list(value[0]) - return [int(item) for item in value] - - -def _looks_like_codes(value: Any, num_codebooks: int) -> bool: - if isinstance(value, (torch.Tensor, np.ndarray)): - return value.ndim == 2 and value.shape[-1] == num_codebooks - return False - - -def _normalize_codes(codes: Any, num_codebooks: int, codebook_size: int) -> torch.Tensor: - codes = torch.as_tensor(codes) - if codes.ndim != 2: - raise ValueError(f"reference audio codes must be 2D, got {tuple(codes.shape)}") - if codes.shape[0] == 0: - raise ValueError("reference audio codes must contain at least one frame") - if codes.shape[-1] != num_codebooks: - raise ValueError(f"expected {num_codebooks} codebooks, got {tuple(codes.shape)}") - codes = codes.to(device="cpu", dtype=torch.long).contiguous() - if int(codes.min()) < 0 or int(codes.max()) >= codebook_size: - raise ValueError(f"reference code outside [0, {codebook_size})") - return codes.to(dtype=torch.int16) - - -__all__ = ["BreezeTTS2PromptBuilder"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/reference_encoder.py b/vllm_omni/model_executor/models/breeze_tts_2/reference_encoder.py new file mode 100644 index 000000000..d695b1cfa --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/reference_encoder.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The bundled Qwen3 codec encoder used by Breeze voice cloning.""" + +import torch +from torch import nn +from transformers.models.mimi.modeling_mimi import MimiConv1d +from vllm.config import VllmConfig +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Config, +) +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Encoder, +) + + +class BreezeReferenceEncoder(nn.Module): + def __init__(self, vllm_config: VllmConfig) -> None: + super().__init__() + config = Qwen3TTSTokenizerV2Config.from_pretrained( + vllm_config.model_config.model, + subfolder="audio_tokenizer", + revision=vllm_config.model_config.revision, + ) + if config.input_sample_rate != 24000 or config.encode_downsample_rate != 1920: + raise ValueError("Breeze requires the bundled 24 kHz, 12.5 Hz Qwen3 codec") + config.encoder_config._attn_implementation = "sdpa" + # The released Breeze runtime encodes references in FP32. Quantizer + # boundaries affect speaker conditioning, so preserve that precision. + self.encoder = Qwen3TTSTokenizerV2Encoder._from_config(config.encoder_config, dtype=torch.float32).eval() + self._convolutions: list[BreezeReferenceConv] = [] + hop = 1 + for name, module in list(self.encoder.named_modules()): + if isinstance(module, MimiConv1d): + hop *= module.conv.stride[0] + convolution = BreezeReferenceConv(module, hop) + self._convolutions.append(convolution) + parent, _, attribute = name.rpartition(".") + self.encoder.get_submodule(parent).add_module(attribute, convolution) + if hop != config.encode_downsample_rate: + raise ValueError("Breeze reference convolution strides do not match the codec frame rate") + self.num_quantizers = config.encoder_valid_num_quantizers + # Reference convolutions have large temporary workspaces. Compile + # without a graph's permanently reserved workspace for every length. + self._compiled_encode = torch.compile( + self._encode_impl, + fullgraph=True, + dynamic=True, + options={"triton.cudagraphs": False, "epilogue_fusion": False}, + ) + + def forward(self, waveform: torch.Tensor, frames: int) -> torch.Tensor: + return self.encode_batch([waveform], [frames])[0] + + def encode_batch(self, waveforms: list[torch.Tensor], frame_counts: list[int]) -> list[torch.Tensor]: + device = next(self.encoder.parameters()).device + waveform = nn.utils.rnn.pad_sequence(waveforms, batch_first=True) + waveform = nn.functional.pad(waveform, (0, -waveform.shape[1] % 1920)) + waveform_frames = waveform.reshape(len(waveforms), -1, 1920).to(device=device, dtype=torch.float32) + lengths = torch.tensor([waveform.numel() for waveform in waveforms], device=device, dtype=torch.long) + with torch.backends.cudnn.flags(allow_tf32=False): + encoded = self._compiled_encode(waveform_frames, lengths) + return [ + encoded[row, :, :frames].transpose(0, 1).clone(memory_format=torch.contiguous_format) + for row, frames in enumerate(frame_counts) + ] + + def warmup(self, max_batch_size: int) -> None: + device = next(self.encoder.parameters()).device + with torch.backends.cudnn.flags(allow_tf32=False): + for batch in range(1, max_batch_size + 1): + waveform_frames = torch.zeros((batch, 32, 1920), device=device, dtype=torch.float32) + lengths = torch.full((batch,), 32 * 1920, device=device, dtype=torch.long) + # Initialize the codec's normalized codebooks before tracing. + self._encode_impl(waveform_frames, lengths) + self._compiled_encode(waveform_frames, lengths) + torch.accelerator.empty_cache() + + def _encode_impl(self, waveform_frames: torch.Tensor, sample_lengths: torch.Tensor) -> torch.Tensor: + # Expose the complete 1920-sample stride product to the compiler. + # Nested ceil/mod expressions for arbitrary waveform lengths otherwise + # cause expensive symbolic simplification during kernel tiling. Every + # convolution below masks the true per-recording boundary, so aligning + # the storage to full frames preserves partial final-frame semantics. + torch._dynamo.mark_static(waveform_frames, 2) + waveform = waveform_frames.flatten(1) + # The released Breeze reference encoder uses full causal attention + # during offline encoding. Recent Transformers Mimi defaults to a + # sliding mask, which changes references longer than ten seconds. + # Supply the actual reference mask rather than changing the checkpoint + # configuration or relying on a Transformers-version default. + for convolution in self._convolutions: + convolution.sample_lengths = sample_lengths + hidden = self.encoder.encoder(waveform.unsqueeze(1)).transpose(1, 2) + length = hidden.shape[1] + causal_mask = torch.ones((length, length), device=hidden.device, dtype=torch.bool).tril()[None, None] + hidden = self.encoder.encoder_transformer( + hidden, attention_mask=causal_mask, use_cache=False, return_dict=True + ).last_hidden_state + hidden = self.encoder.downsample(hidden.transpose(1, 2)) + return self.encoder.quantizer.encode(hidden, self.num_quantizers).transpose(0, 1) + + def load_weights(self, vllm_config: VllmConfig) -> set[str]: + loader = DefaultModelLoader(vllm_config.load_config) + source = DefaultModelLoader.Source( + model_or_path=vllm_config.model_config.model, + revision=vllm_config.model_config.revision, + subfolder="audio_tokenizer", + ) + parameters = dict(self.named_parameters()) + buffers = dict(self.named_buffers()) + loaded: set[str] = set() + for name, value in loader._get_weights_iterator(source): + if name.startswith("decoder."): + continue + if name in parameters: + default_weight_loader(parameters[name], value) + elif name in buffers: + buffers[name].copy_(value) + else: + raise ValueError(f"Unexpected reference codec weight: {name}") + loaded.add(name) + required = set(parameters) | {name for name in buffers if ".quantizer." in name} + missing = required - loaded + if missing: + raise ValueError(f"Uninitialized reference codec weights or codebooks: {sorted(missing)}") + return loaded + + +class BreezeReferenceConv(nn.Module): + """Offline Mimi convolution with padding derived from static shape metadata.""" + + def __init__(self, source: MimiConv1d, output_hop: int) -> None: + super().__init__() + if not source.causal or source.pad_mode not in ("constant", "replicate"): + raise ValueError("Breeze reference encoding requires causal zero/replicate-padded convolutions") + self.conv = source.conv + self.pad_mode = source.pad_mode + self.stride = self.conv.stride[0] + self.padding = (self.conv.kernel_size[0] - 1) * self.conv.dilation[0] + 1 - self.stride + self.output_hop = output_hop + self.sample_lengths: torch.Tensor | None = None + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + # HF stores these integers in GPU tensors, introducing a device sync + # when F.pad converts them back to integers. Shapes determine padding. + if self.sample_lengths is None: + raise RuntimeError("Reference convolution requires the batch's input lengths") + if self.pad_mode == "replicate": + # Mimi's final downsampler repeats the actual last encoder state + # for a partial frame. Each row has its own boundary in a bucket. + input_hop = self.output_hop // self.stride + lengths = (self.sample_lengths + input_hop - 1) // input_hop + indices = torch.arange(hidden.shape[-1], device=hidden.device)[None, :] + indices = torch.minimum(indices, lengths[:, None] - 1) + hidden = hidden.gather(2, indices[:, None, :].expand(-1, hidden.shape[1], -1)) + right = -hidden.shape[-1] % self.stride + hidden = nn.functional.pad(hidden, (self.padding, right), mode=self.pad_mode) + hidden = self.conv(hidden) + lengths = (self.sample_lengths + self.output_hop - 1) // self.output_hop + invalid = torch.arange(hidden.shape[-1], device=hidden.device)[None, :] >= lengths[:, None] + return hidden.masked_fill(invalid[:, None, :], 0) diff --git a/vllm_omni/model_executor/models/breeze_tts_2/text_encoder_graph.py b/vllm_omni/model_executor/models/breeze_tts_2/text_encoder_graph.py new file mode 100644 index 000000000..a687f51bc --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/text_encoder_graph.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Padded T5Gemma2 prefill graphs with explicit bidirectional masks.""" + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.models.t5gemma2.modeling_t5gemma2 import ( + T5Gemma2EncoderLayer, + T5Gemma2TextEncoder, + sliding_window_mask_function, +) +from vllm.platforms import current_platform + +from vllm_omni.platforms import current_omni_platform + + +class BreezeTextEncoderCompiled: + """Compiled prefill layers without retaining large activation graph pools.""" + + def __init__(self, encoder: T5Gemma2TextEncoder, projection: nn.Linear) -> None: + self.encoder = encoder + self.projection = projection + self._layer = torch.compile( + T5Gemma2EncoderLayer.forward, + fullgraph=True, + dynamic=True, + options={"triton.cudagraphs": False, "epilogue_fusion": False}, + ) + + def run_batch(self, prompts: list[torch.Tensor]) -> list[torch.Tensor]: + encoder = self.encoder + device = self.projection.weight.device + dtype = self.projection.weight.dtype + lengths = [prompt.numel() for prompt in prompts] + ids = nn.utils.rnn.pad_sequence(prompts, batch_first=True).to(device) + positions = torch.arange(ids.shape[1], device=device)[None] + invalid = positions >= torch.tensor(lengths, device=device)[:, None] + mask = torch.zeros((len(prompts), 1, ids.shape[1], ids.shape[1]), device=device, dtype=dtype) + mask.masked_fill_(invalid[:, None, None, :], torch.finfo(dtype).min) + mask_function = sliding_window_mask_function(encoder.config.sliding_window, is_causal=False) + allowed = mask_function(0, 0, positions[:, :, None], positions[:, None, :]) + local_mask = mask.masked_fill(~allowed[:, None], torch.finfo(dtype).min) + masks = {"full_attention": mask, "sliding_attention": local_mask} + embedding = encoder.embed_tokens + hidden = F.embedding(ids, embedding.weight) * embedding.embed_scale.to(dtype) + hidden = torch.where( + (ids == embedding.eoi_token_index).unsqueeze(-1), embedding.eoi_embedding.to(dtype), hidden + ) + position_embeddings = { + kind: encoder.rotary_emb(hidden, positions, kind) for kind in set(encoder.config.layer_types) + } + for layer in encoder.layers: + kind = layer.attention_type + hidden = self._layer(layer, hidden, position_embeddings[kind], masks[kind], positions) + output = self.projection(encoder.norm(hidden)) + return [output[row, :length].clone() for row, length in enumerate(lengths)] + + def warmup(self, max_batch_size: int) -> None: + # Dynamo specializes size-one dimensions; a second batch exercises + # the general batch dimension used by paired and reference prompts. + for batch in sorted({1, max_batch_size}): + self.run_batch([torch.zeros(256, dtype=torch.long) for _ in range(batch)]) + torch.accelerator.empty_cache() + + +class BreezeTextEncoderGraph: + def __init__(self, encoder: T5Gemma2TextEncoder, projection: nn.Linear, size: int, batch_size: int = 1) -> None: + device = projection.weight.device + dtype = projection.weight.dtype + self.ids = torch.zeros((batch_size, size), device=device, dtype=torch.long) + self.positions = torch.arange(size, device=device).unsqueeze(0) + self.full_mask = torch.zeros((batch_size, 1, size, size), device=device, dtype=dtype) + mask_function = sliding_window_mask_function(encoder.config.sliding_window, is_causal=False) + allowed = mask_function(0, 0, self.positions[:, :, None], self.positions[:, None, :]) + self.local_mask_base = torch.zeros_like(self.full_mask).masked_fill_(~allowed[:, None], torch.finfo(dtype).min) + self.local_mask = self.local_mask_base.clone() + self.masks = {"full_attention": self.full_mask, "sliding_attention": self.local_mask} + for _ in range(3): + self._encode(encoder, projection) + current_omni_platform.synchronize() + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph( + self.graph, pool=current_platform.get_global_graph_pool(), capture_error_mode="thread_local" + ): + self.output = self._encode(encoder, projection) + + def _encode(self, encoder: T5Gemma2TextEncoder, projection: nn.Linear) -> torch.Tensor: + embedding = encoder.embed_tokens + hidden = F.embedding(self.ids, embedding.weight) * embedding.embed_scale.to(embedding.weight.dtype) + # HF's boolean-index assignment materializes a data-dependent index + # list. Express the same EOI substitution with a fixed-shape select. + hidden = torch.where( + (self.ids == embedding.eoi_token_index).unsqueeze(-1), + embedding.eoi_embedding.to(hidden.dtype), + hidden, + ) + encoded = encoder(inputs_embeds=hidden, attention_mask=self.masks, position_ids=self.positions) + return projection(encoded.last_hidden_state) + + def run(self, prompt: torch.Tensor) -> torch.Tensor: + return self.run_batch([prompt[0]])[0] + + def run_batch(self, prompts: list[torch.Tensor]) -> list[torch.Tensor]: + if len(prompts) != self.ids.shape[0]: + raise ValueError("Breeze text batch differs from the captured graph size") + lengths = [prompt.numel() for prompt in prompts] + if not all(0 < length <= self.ids.shape[1] for length in lengths): + raise ValueError("Breeze text segment does not fit its graph bucket") + self.ids.zero_() + for row, (prompt, length) in enumerate(zip(prompts, lengths, strict=True)): + self.ids[row, :length].copy_(prompt) + invalid = self.positions >= torch.tensor(lengths, device=self.ids.device)[:, None] + invalid = invalid[:, None, None, :] + self.full_mask.zero_() + self.full_mask.masked_fill_(invalid, torch.finfo(self.full_mask.dtype).min) + self.local_mask.copy_(self.local_mask_base) + self.local_mask.masked_fill_(invalid, torch.finfo(self.local_mask.dtype).min) + self.graph.replay() + # Subsequent prefills and backbone/depth graphs reuse graph storage. + # The runner must receive independently owned prompt embeddings. + return [self.output[row, :length].clone() for row, length in enumerate(lengths)] diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index ce9f0b04a..8bb9e30d5 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -58,6 +58,8 @@ class Qwen3TTSCode2Wav(nn.Module): via the SpeechTokenizer decoder directly (bypassing HF wrapper overhead).""" input_modalities = "audio" + tokenizer_subfolder = "speech_tokenizer" + decoder_cudagraph_modes: tuple[str, ...] = ("icl", "xvec") # Ask the model runner for the scheduler-side request IDs. Stateful # decoder caches must use the same IDs delivered by on_requests_finished; @@ -105,7 +107,7 @@ class Qwen3TTSCode2Wav(nn.Module): # load_weights(). tok_config = Qwen3TTSTokenizerV2Config.from_pretrained( self.model_path, - subfolder="speech_tokenizer", + subfolder=self.tokenizer_subfolder, ) dec_config = tok_config.decoder_config self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(dec_config) @@ -194,6 +196,7 @@ class Qwen3TTSCode2Wav(nn.Module): ) self.decoder.enable_cudagraph( + capture_modes=self.decoder_cudagraph_modes, capture_batch_sizes=decode_cudagraph_batch_sizes, stateless_capture_sizes=decode_cudagraph_capture_sizes, device=device, @@ -566,7 +569,7 @@ class Qwen3TTSCode2Wav(nn.Module): source = DefaultModelLoader.Source( model_or_path=self.model_path, revision=self.vllm_config.model_config.revision, - subfolder="speech_tokenizer", + subfolder=self.tokenizer_subfolder, ) subfolder_weights = model_loader._get_weights_iterator(source) loaded = AutoWeightsLoader(self).load_weights( diff --git a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py index 3c73f5d13..a87af5f85 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py @@ -1,5 +1,6 @@ # Copyright 2026 The Alibaba Qwen team. # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ CUDA Graph wrapper for Qwen3TTSTokenizerV2Decoder. @@ -45,6 +46,7 @@ class CUDAGraphDecoderWrapper: def __init__( self, decoder: torch.nn.Module, + capture_modes: tuple[str, ...] = ("icl", "xvec"), capture_batch_sizes: list[int] | None = None, stateless_capture_sizes: list[int] | None = None, num_quantizers: int = 8, @@ -57,6 +59,9 @@ class CUDAGraphDecoderWrapper: decode_left_context: int = 25, ): self.decoder = decoder + if not capture_modes or set(capture_modes) - {"icl", "xvec"}: + raise ValueError("capture_modes must contain icl and/or xvec") + self.capture_modes = capture_modes self.capture_batch_sizes = sorted(set(capture_batch_sizes or [1])) self.num_quantizers = num_quantizers self.enabled = enabled @@ -407,16 +412,18 @@ class CUDAGraphDecoderWrapper: icl_capture_shapes = self._get_icl_capture_shapes() for batch_size in self.capture_batch_sizes: - try: - self._capture_icl_prefix(batch_size, device, dtype) - logger.info("Captured ICL prefix CUDA Graph for batch=%d", batch_size) - except Exception: - logger.warning("Failed to capture ICL prefix CUDA Graph for batch=%d", batch_size, exc_info=True) - try: - self._capture_xvec_prefix(batch_size, device, dtype) - logger.info("Captured xvec prefix CUDA Graph for batch=%d", batch_size) - except Exception: - logger.warning("Failed to capture xvec prefix CUDA Graph for batch=%d", batch_size, exc_info=True) + if "icl" in self.capture_modes: + try: + self._capture_icl_prefix(batch_size, device, dtype) + logger.info("Captured ICL prefix CUDA Graph for batch=%d", batch_size) + except Exception: + logger.warning("Failed to capture ICL prefix CUDA Graph for batch=%d", batch_size, exc_info=True) + if "xvec" in self.capture_modes: + try: + self._capture_xvec_prefix(batch_size, device, dtype) + logger.info("Captured xvec prefix CUDA Graph for batch=%d", batch_size) + except Exception: + logger.warning("Failed to capture xvec prefix CUDA Graph for batch=%d", batch_size, exc_info=True) logger.info( "Starting CUDA Graph warmup for %d shapes: batch_sizes=%s seq_lens=%s", @@ -425,7 +432,7 @@ class CUDAGraphDecoderWrapper: self.icl_capture_sizes, ) - for batch_size, size in icl_capture_shapes: + for batch_size, size in icl_capture_shapes if "icl" in self.capture_modes else (): try: caches = self._make_dummy_icl_cache(batch_size, device, next(self.decoder.parameters()).dtype) self._capture_combined_suffix("icl", batch_size, size, caches, device, dtype) @@ -439,7 +446,7 @@ class CUDAGraphDecoderWrapper: ) model_dtype = next(self.decoder.parameters()).dtype - for batch_size in self.capture_batch_sizes: + for batch_size in self.capture_batch_sizes if "xvec" in self.capture_modes else (): for size, previous_frames in self._xvec_previous_frames_by_target.items(): try: caches = self._make_dummy_xvec_cache(batch_size, device, model_dtype) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index d01eae654..7d6d721bd 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -999,6 +999,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): def enable_cudagraph( self, + capture_modes: tuple[str, ...] = ("icl", "xvec"), capture_batch_sizes: list[int] | None = None, stateless_capture_sizes: list[int] | None = None, device: torch.device | None = None, @@ -1021,6 +1022,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): self._cudagraph_wrapper = CUDAGraphDecoderWrapper( decoder=self, + capture_modes=capture_modes, capture_batch_sizes=capture_batch_sizes, stateless_capture_sizes=stateless_capture_sizes, num_quantizers=self.config.num_quantizers, diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 468ed5302..71b2792e3 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -9,6 +9,8 @@ from vllm.model_executor.models.registry import ( ) _OMNI_MODELS = { + "BreezeForConditionalGeneration": ("breeze_tts_2", "modeling_breeze", "BreezeForConditionalGeneration"), + "BreezeCode2Wav": ("breeze_tts_2", "code2wav", "BreezeCode2Wav"), "Qwen2_5OmniForConditionalGeneration": ( "qwen2_5_omni", "qwen2_5_omni", @@ -476,16 +478,6 @@ _OMNI_MODELS = { "acoustic", "MiniMaxMusic3AcousticForConditionalGeneration", ), - "BreezeForConditionalGeneration": ( - "breeze_tts_2", - "modeling_breeze_tts_2_talker", - "BreezeTTS2TalkerForGeneration", - ), - "BreezeTTS2MimiCodec": ( - "breeze_tts_2", - "modeling_breeze_tts_2_codec", - "BreezeTTS2MimiCodec", - ), } diff --git a/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py b/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py index ae393948a..696dbca48 100644 --- a/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py +++ b/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py @@ -1,244 +1,88 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Transfer raw Breeze RVQ frames to the stateful Qwen3 codec.""" -"""Stage-0 to stage-1 input conversion for Breeze-TTS-2. - -The prompt builder and the reference-audio tokenizer live under the Breeze -model directory. This module intentionally contains only the inter-stage -contract: flatten the generated ``(T, Q)`` codec matrix into the codebook-major -token sequence consumed by the stage-1 codec model. -""" - -from __future__ import annotations - -from collections.abc import Mapping, Sequence -from typing import Any +from typing import TYPE_CHECKING, Any import torch -from vllm.logger import init_logger +from vllm.v1.request import Request -from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct -from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct +from vllm_omni.model_executor.models.breeze_tts_2.prompt import CFG_UNCOND_SUFFIX +from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt -logger = init_logger(__name__) - -# ``BreezeTTS2TalkerForGeneration`` emits the complete sequence accumulated -# so far on every decode step. The generic full-payload accumulator normally -# concatenates rank-2 tensors, which would duplicate all earlier frames. -# Replace the latest complete sequence instead. -_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes.audio"}) -_DEFAULT_CODEC_CHUNK_FRAMES = 8 +if TYPE_CHECKING: + from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import OmniChunkTransferAdapter -def _extract_audio_codes(output: Any) -> torch.Tensor | None: - """Extract ``codes.audio`` from the supported Omni output shapes.""" - # Depending on the engine path, source_outputs contains either the - # per-request output itself or an EngineCoreOutput with ``outputs[0]``. - if isinstance(output, Mapping): - multimodal = output.get("multimodal_output", output.get("multimodal_outputs")) - if multimodal is None and "codes" in output: - multimodal = output - if multimodal is None and "codes.audio" in output: - audio = output["codes.audio"] - if isinstance(audio, (list, tuple)): - if len(audio) != 1: - raise ValueError(f"expected one Breeze audio-code tensor, got {len(audio)}") - audio = audio[0] - audio = torch.as_tensor(audio) - if audio.ndim not in (1, 2): - raise ValueError(f"Breeze audio codes must be rank 1 or 2, got {tuple(audio.shape)}") - return audio.to(device="cpu", dtype=torch.long).contiguous() - else: - multimodal = getattr(output, "multimodal_output", None) - if multimodal is None: - multimodal = getattr(output, "multimodal_outputs", None) - if multimodal is None: - candidates = getattr(output, "outputs", None) - if isinstance(candidates, Sequence) and candidates: - return _extract_audio_codes(candidates[0]) - if not isinstance(multimodal, Mapping): - return None - codes = multimodal.get("codes") - audio = codes.get("audio") if isinstance(codes, Mapping) else multimodal.get("codes.audio") - if audio is None: - return None - # ``OmniOutput`` stores one tensor per request under ``codes.audio``. - # ``process_engine_inputs`` invokes this function once per source output, - # so unwrap that single-item container before tensor conversion. - if isinstance(audio, (list, tuple)): - if len(audio) != 1: - raise ValueError(f"expected one Breeze audio-code tensor, got {len(audio)}") - audio = audio[0] - audio = torch.as_tensor(audio) - if audio.ndim == 3 and audio.shape[0] == 1: - audio = audio[0] - if audio.ndim not in (1, 2): - raise ValueError( - f"Breeze audio codes must have shape (T, Q) or codebook-major (Q*T,), got {tuple(audio.shape)}" - ) - return audio.to(device="cpu", dtype=torch.long).contiguous() +def expand_cfg_prompts(prompt: dict[str, Any] | str, sampling_params: Any) -> list[ExpandedPrompt]: + if not isinstance(prompt, dict): + raise ValueError("Breeze requires build_breeze_prompt to encode its conditioning") + info = prompt["additional_information"] + conditioning = info["breeze_prompt"] + if conditioning["guidance_scale"] == 1.0: + return [] + negative = conditioning["negative_ids"] + companion_conditioning = {**conditioning, "role": "uncond", "target_ids": negative} + companion_conditioning.pop("negative_ids") + length = len(prompt["prompt_token_ids"]) - len(conditioning["target_ids"]) + len(negative) + companion = { + "prompt_token_ids": [0] * length, + "additional_information": { + **info, + "breeze_prompt": companion_conditioning, + "cfg_group": {"role": "uncond", "uncond_suffix": CFG_UNCOND_SUFFIX}, + }, + } + return [ExpandedPrompt(prompt=companion, role="uncond", request_id_suffix=CFG_UNCOND_SUFFIX)] -def talker2codec( - source_outputs: Sequence[Any], - prompt: Any = None, - _requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """Build synchronous stage-1 prompts from completed stage-0 outputs. - - Breeze's Mimi decoder receives one flat codebook-major sequence. The - generated matrix is therefore transposed from ``(T, Q)`` to ``(Q, T)`` - before flattening. No audio decoding, text tokenization, or streaming - state is handled here. - """ - del prompt, _requires_multimodal_data - results: list[OmniTokensPrompt] = [] - for source_index, output in enumerate(source_outputs): - codes = _extract_audio_codes(output) - if codes is None or codes.numel() == 0: - logger.warning("Breeze stage output %d contains no audio codes", source_index) - results.append(OmniTokensPrompt(prompt_token_ids=[])) - continue - # Connector full-payload builders send a codebook-major flat tensor; - # direct in-process stage outputs carry frame-major ``(T, Q)``. - if codes.ndim == 1: - codebook_major_data = codes.contiguous() - codebook_major = codebook_major_data - else: - codebook_major_data = codes.transpose(0, 1).contiguous() - codebook_major = codebook_major_data.reshape(-1) - results.append( - OmniTokensPrompt( - prompt_token_ids=codebook_major.tolist(), - multi_modal_data={"codes": {"audio": codebook_major_data}}, - ) - ) - return results - - -def talker2codec_full_payload( - transfer_manager: Any = None, - pooling_output: Mapping[str, Any] | None = None, - request: Any = None, +def talker2code2wav_async_chunk( + transfer_manager: "OmniChunkTransferAdapter", + multimodal_output: OmniPayload | None, + request: Request, is_finished: bool = False, ) -> OmniPayloadStruct | None: - """Build the final Stage-1 payload from accumulated talker codes. - - ``pooling_output`` is the AR runner's flattened per-request payload. Mimi - consumes codebook-major IDs, so the frame-major talker matrix is transposed - before flattening. A request that finishes without any codec frame still - gets an explicit empty payload with ``finished=True``: returning ``None`` - would leave Stage 1 without a terminal marker, and the connector receiver - requeues marker-less empty chunks forever. - """ - del transfer_manager, request - audio = None - if isinstance(pooling_output, Mapping): - audio = pooling_output.get("codes.audio") - if audio is None: - nested = pooling_output.get("codes") - if isinstance(nested, Mapping): - audio = nested.get("audio") - - codes = torch.empty(0, dtype=torch.long) if audio is None else torch.as_tensor(audio, dtype=torch.long) - if codes.numel() == 0: - if not is_finished: - return None - return OmniPayloadStruct( - codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), - meta=MetaStruct( - code_flat_numel=0, - finished=torch.tensor(True, dtype=torch.bool), - ), - ) - if codes.ndim == 2: - flat = codes.transpose(0, 1).contiguous().reshape(-1) - elif codes.ndim == 1: - flat = codes.contiguous() - else: - raise ValueError(f"Breeze audio codes must be rank 1 or 2, got {tuple(codes.shape)}") - return OmniPayloadStruct( - codes=CodesStruct(audio=flat), - meta=MetaStruct( - code_flat_numel=int(flat.numel()), - finished=torch.tensor(bool(is_finished), dtype=torch.bool), - ), - ) - - -def _extract_current_audio_frame( - multimodal_output: Mapping[str, Any] | None, -) -> torch.Tensor | None: - """Extract the single ``(Q,)`` frame emitted by async Stage 0.""" - if not isinstance(multimodal_output, Mapping): - return None - codes = multimodal_output.get("codes") - audio = codes.get("audio") if isinstance(codes, Mapping) else None - if audio is None: - return None - frame = torch.as_tensor(audio, dtype=torch.long).reshape(-1) - if frame.numel() == 0: - return None - return frame.cpu().contiguous() - - -def _codec_chunk_frames(transfer_manager: Any) -> int: - connector = getattr(transfer_manager, "connector", None) - # SharedMemoryConnector stores its ConnectorSpec.extra directly in - # ``connector.config``; there is no second ``extra`` wrapper at runtime. - config = getattr(connector, "config", {}) or {} - value = config.get("breeze_codec_chunk_frames") - try: - chunk_frames = int(value) if value is not None else _DEFAULT_CODEC_CHUNK_FRAMES - except (TypeError, ValueError) as exc: - raise ValueError(f"Invalid Breeze codec_chunk_frames={value!r}") from exc - if chunk_frames <= 0: - raise ValueError(f"Breeze codec_chunk_frames must be positive, got {chunk_frames}") - return chunk_frames - - -def talker2codec_async_chunk( - transfer_manager: Any, - multimodal_output: Mapping[str, Any] | None, - request: Any, - is_finished: bool = False, -) -> OmniPayloadStruct | None: - """Batch Stage-0 frames and send only the un-emitted codebook-major tail.""" request_id = request.external_req_id - finished = bool(is_finished or request.is_finished()) - frame = _extract_current_audio_frame(multimodal_output) - if frame is not None: - transfer_manager.code_prompt_token_ids[request_id].append(frame) - - pending = transfer_manager.code_prompt_token_ids[request_id] - chunk_frames = _codec_chunk_frames(transfer_manager) - if not pending and not finished: + if request_id.endswith(CFG_UNCOND_SUFFIX): return None - if len(pending) < chunk_frames and not finished: - return None - - if pending: - # Stage 0 supplies frame-major ``(T, Q)`` rows. The Qwen3-TTS decoder - # consumes ``(Q, T)``, so transpose before flattening the emitted tail. - tail = torch.stack(pending, dim=0) - flat = tail.transpose(0, 1).contiguous().reshape(-1) + finished = is_finished or request.is_finished() + frames = transfer_manager.code_prompt_token_ids[request_id] + if multimodal_output is not None: + codes = multimodal_output.get("codes", {}).get("audio") + if codes is not None and codes.numel(): + if codes.ndim != 2 or codes.shape[1] != 16: + raise ValueError("Breeze codec payload must have shape (frames, 16)") + rows = codes.to(device="cpu", dtype=torch.long) + if not bool(((rows >= 0) & (rows < 2048)).all()): + raise ValueError("Breeze codec payload contains a reserved or invalid code") + frames.extend(rows.unbind(0)) + connector = transfer_manager.connector + if connector is None: + raise RuntimeError("Breeze streaming requires a stage connector") + config = connector.config + extra = config.get("extra", config) + state = transfer_manager.request_payload.setdefault(request_id, {"frames": 0, "chunks": 0}) + emitted = state["frames"] + ramp = extra.get("codec_chunk_ramp", []) + if state["chunks"] < len(ramp): + chunk_size = int(ramp[state["chunks"]]) else: - flat = torch.empty(0, dtype=torch.long) - pending.clear() - + chunk_size = int( + extra.get("initial_codec_chunk_frames", 1) if emitted == 0 else extra.get("codec_chunk_frames", 5) + ) + if chunk_size <= 0: + raise ValueError("codec_chunk_frames must be positive") + pending = len(frames) - emitted + if not finished and pending < chunk_size: + return None + if pending < 0: + raise RuntimeError("Breeze codec transfer emitted more frames than generated") + # The shared codec retains its causal state; each frame is delivered once. + flat = torch.stack(frames[emitted:]).T.contiguous().reshape(-1) if pending else torch.empty(0, dtype=torch.long) + state["frames"] = len(frames) + state["chunks"] += 1 return OmniPayloadStruct( codes=CodesStruct(audio=flat), - meta=MetaStruct( - request_id=request_id, - code_flat_numel=int(flat.numel()), - codec_streaming=True, - finished=torch.tensor(finished, dtype=torch.bool), - # The chunk adapter strips ``meta.finished`` before it reaches the - # stage-1 model; ``stream_finished`` survives and is what the - # codec uses to release per-request decoder state. - stream_finished=torch.tensor(finished, dtype=torch.bool), - ), + meta=MetaStruct(request_id=request_id, left_context_size=0, finished=torch.tensor(finished)), ) - - -__all__ = ["talker2codec", "talker2codec_async_chunk", "talker2codec_full_payload"] diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 742382938..80eaf3e67 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1770,6 +1770,43 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): req_ids_b = [item[0] for item in decode_batch_items] start_offsets_b = [item[1] for item in decode_batch_items] req_infos_b = [item[2] for item in decode_batch_items] + if not self.has_talker_mtp: + # Non-MTP hooks return (ids, embeds, per-request updates), + # matching scalar preprocess without the MTP-only tensors. + # Every non-decode row flushes the group, so these one-token + # spans are contiguous in the runner's capture buffers. + start = start_offsets_b[0] + end = start + len(req_ids_b) + if start < 0 or end > preprocess_input_ids.shape[0] or start_offsets_b != list(range(start, end)): + raise RuntimeError( + "Non-MTP batched decode preprocessing requires contiguous in-bounds token offsets" + ) + ids_b = preprocess_input_ids[start:end] + req_input_ids, req_embeds, updates = batch_decode_preprocess( + input_ids=ids_b, + req_infos=req_infos_b, + ) + if ( + req_input_ids.numel() != len(req_ids_b) + or req_embeds.ndim != 2 + or req_embeds.shape[0] != len(req_ids_b) + or len(updates) != len(req_ids_b) + ): + raise ValueError("Batched decode preprocessing must return one output and update per request") + if inputs_embeds is None: + inputs_embeds = torch.empty( + (preprocess_input_ids.shape[0], req_embeds.shape[-1]), + device=req_embeds.device, + dtype=req_embeds.dtype, + ) + inputs_embeds[start:end].copy_(req_embeds) + if req_input_ids is not ids_b: + preprocess_input_ids[start:end].copy_(req_input_ids.reshape(-1)) + for req_id_b, update_dict_b in zip(req_ids_b, updates, strict=True): + self._update_intermediate_buffer(req_id_b, update_dict_b) + decode_batch_items.clear() + return + ids_b = torch.stack( [preprocess_input_ids[offset : offset + 1].reshape(-1)[0] for offset in start_offsets_b] ) @@ -1839,7 +1876,7 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): # Seed, so a model that samples inside forward() can be # reproducible: vLLM's own sampler seeding does not reach it. req_infos["_omni_seed"] = getattr(sampling_params, "seed", None) - if callable(batch_decode_preprocess) and self.has_talker_mtp and span_len == 1 and not is_prefill: + if callable(batch_decode_preprocess) and span_len == 1 and not is_prefill: decode_batch_items.append((req_id, s, req_infos)) continue