anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Accelerate Breeze-TTS-2 non-streaming inference (#7752)

Browse Source
This commit is contained in:
Yueqian Lin authored and GitHub committed 2026-09-22 17:10:15 -07:00
1 parent d2aa07159c
commit 3f473a358b
84 files changed
+8162 -5593
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/cuda/test-merge.yml
+20
View File
@@ -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:
.buildkite/cuda/test-ready.yml
+19
View File
@@ -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:
benchmarks/tts/evidence/breeze_tts_2_20260918.json
+757
View File
@@ -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"
]
}
}
benchmarks/tts/evidence/breeze_tts_2_sampling_graphs_20260918.json
+2735
View File
File diff suppressed because it is too large. Load diff
benchmarks/tts/model_configs.yaml
+11
View File
@@ -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
docs/design/feature/preprocess_phase_contract.md
+20 -2
View File
@@ -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
docs/models/supported_models.md
+1 -1
View File
@@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
docs/source/performance/breeze_tts_2_gpu_bubbles.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 234 KiB

examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py
+18 -35
View File
@@ -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,37 +104,18 @@ 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,
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,
)
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,
)
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:
"""Extract the final complete audio tensor from multimodal output.
@@ -180,6 +158,11 @@ def main():
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
t_start = time.perf_counter()
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
examples/online_serving/text_to_speech/README.md
+1 -1
View File
@@ -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
recipes/BreezeBlue/Breeze-TTS-2.md
+270 -173
View File
@@ -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" \
export VLLM_WSL2_ENABLE_PIN_MEMORY=1
```
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 --fail http://127.0.0.1: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",
"input": "你好，欢迎使用语音合成系统。",
"instructions": "年轻女性，声音温暖自然，使用标准普通话。",
"extra_params": {"guidance_scale": 4.0},
"response_format": "wav",
"sample_rate": 24000
}' --output breeze_plain.wav
"seed": 42,
"max_new_tokens": 250
}' --output breeze.wav
```
Voice design (instruction only, no reference audio):
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,<base64-encoded-reference>",
"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": "(laugh) Welcome to our story time.",
"instructions": "A warm young woman, clear voice, lively delivery.",
"response_format": "wav"
}' --output breeze_design.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 cloning (`ref_audio` and `ref_text` must be provided together):
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" \
curl --fail http://127.0.0.1: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
"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
```
Voice direction (reference + instruction):
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 -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
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}}'
```
Streaming raw PCM (`stream_format="audio"`):
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.
## Pipeline and optimization
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.
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.
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
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
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
```
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.
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.
## Notes
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.
- **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.
The shared serving benchmark supports default speech, Voice Design and
Voice Clone. For Voice Design:
## Supported features
```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
```
| 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 |
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.
requirements/common.txt
+2
View File
@@ -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
tests/config/test_environment_variables.py
+2 -2
View File
@@ -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,
}
tests/core/sched/test_omni_cfg_ar_scheduler.py
+102
View File
@@ -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()
tests/e2e/online_serving/test_breeze_tts_2.py
+81 -73
View File
@@ -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"],
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,
),
id="breeze_tts_2",
)
]
@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 = {
def test_chinese_voice_design(omni_server, online_client, run_level) -> None:
online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": get_prompt(),
"stream": False,
"timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
"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",
"voice": "S0",
"sample_rate": 24000,
"min_audio_bytes": _MIN_AUDIO_BYTES,
"min_audio_bytes": 24000,
"seed": 42,
"max_new_tokens": 8 if run_level == "core_model" else 250,
}
online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
)
@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 = {
def test_english_streaming(omni_server, online_client, run_level) -> None:
online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": get_prompt(),
"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",
"timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
"response_format": "wav",
"voice": "S0",
"sample_rate": 24000,
"min_audio_bytes": _MIN_AUDIO_BYTES,
# 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,
}
online_client.send_audio_speech_request(request_config)
)
@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,
}
)
tests/entrypoints/openai_api/test_breeze_tts_2_adapter.py
+185
View File
@@ -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}]<ins_bos>{DEFAULT_INSTRUCTION}<ins_eos>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]<ins_bos><ins_eos>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
tests/entrypoints/openai_api/test_serving_speech.py
+19 -5
View File
@@ -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={
try:
yield OmniRequestOutput(
request_id="req-close",
final_output_type="audio",
_multimodal_output={
"audio": torch.zeros(16, dtype=torch.float32),
"sr": 24000,
}
},
)
await asyncio.sleep(0)
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)
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
tests/entrypoints/openai_api/test_tts_detection.py
+1
View File
@@ -140,6 +140,7 @@ _PIPELINE_STAGES = [
"audio_vae",
"aura",
"breeze_tts_2",
"breeze_tts_2_codec",
"code2wav",
"cosyvoice3_code2wav",
"cosyvoice3_talker",
tests/entrypoints/test_async_omni.py
+82
View File
@@ -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."""
tests/entrypoints/test_omni_entrypoints.py
+2
View File
@@ -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():
tests/helpers/assertions.py
+5 -1
View File
@@ -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):
tests/helpers/client.py
+2
View File
@@ -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",
tests/helpers/tests/test_assertions.py
+21
View File
@@ -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"),
[
tests/model_executor/models/breeze_tts_2/__init__.py
View File
Whitespace-only changes.
tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py
-128
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_batching.py
-95
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_codec_loading.py
-68
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_configuration.py
-113
View File
@@ -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)
tests/model_executor/models/breeze_tts_2/test_instruction.py
-157
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_local_modules.py
-152
View File
@@ -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))
tests/model_executor/models/breeze_tts_2/test_model_registration.py
-28
View File
@@ -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()
tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py
-36
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_prompt_builder.py
-189
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_reference_voice.py
-187
View File
@@ -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)]
tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py
-122
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py
-133
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_streaming.py
-243
View File
@@ -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 == {}
tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py
-243
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_weight_loading.py
-73
View File
@@ -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
tests/model_executor/models/breeze_tts_2_helpers.py
+54
View File
@@ -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),
},
}
tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py
+11 -3
View File
@@ -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,
tests/model_executor/models/test_breeze_tts_2.py
+459
View File
@@ -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]<ins_bos>Speak happily.<ins_eos>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]]
tests/model_executor/models/test_breeze_tts_2_first_code_integration.py
+184
View File
@@ -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
tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py
+307
View File
@@ -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)
tests/model_executor/models/test_breeze_tts_2_graphs.py
+405
View File
@@ -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]
tests/model_executor/models/test_breeze_tts_2_preprocess.py
+119
View File
@@ -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])
tests/model_executor/models/test_breeze_tts_2_registration.py
+50
View File
@@ -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
tests/model_executor/models/breeze_tts_2/test_sampling_eos.py → tests/model_executor/models/test_breeze_tts_2_sampling_eos.py
+1 -1
View File
@@ -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(
tests/worker/test_omni_gpu_model_runner.py
+152
View File
@@ -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)
vllm_omni/config/environment_variable_inventory.py
-1
View File
@@ -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",
vllm_omni/core/sched/omni_cfg_ar_scheduler.py
+155
View File
@@ -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)
vllm_omni/deploy/breeze_tts.yaml
+62
View File
@@ -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
vllm_omni/deploy/breeze_tts_2.yaml
+40 -34
View File
@@ -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
vllm_omni/deploy/breeze_tts_2_throughput.yaml
+62
View File
@@ -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
vllm_omni/deploy/breeze_tts_throughput.yaml
+63
View File
@@ -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
vllm_omni/engine/arg_utils.py
+2 -4
View File
@@ -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)
vllm_omni/entrypoints/async_omni.py
+20 -7
View File
@@ -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,6 +314,9 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
except (asyncio.CancelledError, GeneratorExit):
self._record_request_failure_once(request_id, reason="client_disconnect")
# 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
@@ -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."""
vllm_omni/entrypoints/async_omni_base.py
+6 -1
View File
@@ -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)
vllm_omni/entrypoints/client_request_state.py
+1
View File
@@ -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
vllm_omni/entrypoints/openai/serving_speech.py
+2
View File
@@ -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):
vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py
+149 -223
View File
@@ -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
vllm_omni/model_executor/models/breeze_tts_2/__init__.py
+7
View File
@@ -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"]
vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py
-149
View File
@@ -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"]
vllm_omni/model_executor/models/breeze_tts_2/code2wav.py
+13
View File
@@ -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",)
vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze.py
+47
View File
@@ -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
vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py
-366
View File
@@ -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",
]
vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py
+409
View File
@@ -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)
vllm_omni/model_executor/models/breeze_tts_2/first_code_sampler.py
+102
View File
@@ -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))
vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze.py
+407
View File
@@ -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
vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py
-425
View File
@@ -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"]
vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py
-523
View File
@@ -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",
]
vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py
-641
View File
@@ -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"]
vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py
-350
View File
@@ -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"]
vllm_omni/model_executor/models/breeze_tts_2/pipeline.py
+10 -29
View File
@@ -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"]
vllm_omni/model_executor/models/breeze_tts_2/prompt.py
+95
View File
@@ -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}<ins_bos>{instructions}<ins_eos>{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
vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py
-336
View File
@@ -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 = "<ins_bos>"
INSTRUCTION_EOS = "<ins_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"]
vllm_omni/model_executor/models/breeze_tts_2/reference_encoder.py
+165
View File
@@ -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)
vllm_omni/model_executor/models/breeze_tts_2/text_encoder_graph.py
+122
View File
@@ -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)]
vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py
+5 -2
View File
@@ -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(
vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py
+9 -2
View File
@@ -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,11 +412,13 @@ class CUDAGraphDecoderWrapper:
icl_capture_shapes = self._get_icl_capture_shapes()
for batch_size in self.capture_batch_sizes:
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)
@@ -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)
vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
+2
View File
@@ -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,
vllm_omni/model_executor/models/registry.py
+2 -10
View File
@@ -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",
),
}
vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py
+69 -225
View File
@@ -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"]
vllm_omni/worker/gpu_model_runner.py
+38 -1
View File
@@ -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