From bdb1d996bee26e34f59f0f7fe24f928b4bb21996 Mon Sep 17 00:00:00 2001 From: zzh <67720593+BLANKETusers@users.noreply.github.com> Date: Wed, 12 Aug 2026 10:40:05 +0800 Subject: [PATCH 01/41] [Bugfix] Fix HunyuanImage3 accuracy test (#5981) Signed-off-by: zzh <943967662@qq.com> --- vllm_omni/diffusion/layers/norm.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 83897eb15..8d91df8f0 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -65,9 +65,10 @@ class LayerNorm(nn.LayerNorm, CustomOp): class RMSNorm(CustomOp): - def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None: + def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype | None = None) -> None: super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) + weight_dtype = dtype or torch.get_default_dtype() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype)) self.hidden_size = hidden_size self.variance_epsilon = eps From b792578d890e73d0371ac7d0d5eb120b1ab90d63 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Wed, 12 Aug 2026 11:32:36 +0800 Subject: [PATCH 02/41] [CI][MiniCPM-o] Add MiniCPM-o 4.5 perf coverage to the ready gate (#6079) Signed-off-by: amy-why-3459 --- .buildkite/cuda/test-ready.yml | 12 +++ .buildkite/npu/test-npu-ready.yml | 9 +++ .../perf/tests/test_minicpmo_4_5_ready.json | 80 +++++++++++++++++++ 3 files changed, 101 insertions(+) create mode 100644 tests/dfx/perf/tests/test_minicpmo_4_5_ready.json diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index fa83fe0f7..f416631ff 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -227,6 +227,18 @@ steps: " mirror_hardwares: h100_1 + - label: "Omni · MiniCPM-o 4.5 Perf Test" + timeout_in_minutes: 60 + source_file_dependencies: + - tests/dfx/perf/tests/test_minicpmo_4_5_ready.json + - vllm_omni/deploy/minicpmo_4_5.yaml + - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda' + mirror_hardwares: h100_1 + - label: "TTS · VoxCPM2 Test" source_file_dependencies: - tests/e2e/online_serving/test_voxcpm2_tts.py diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index fd7f6f546..f23ba7120 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -17,6 +17,15 @@ steps: VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" commands: - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and npu and A3 and omni' --run-level 'core_model' + - label: "Omni · MiniCPM-o 4.5 Perf Test" + timeout_in_minutes: 60 + mirror_hardwares: a3_npu_2 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and npu and A3 and omni' - group: ":card_index_dividers: TTS Model Test" key: ready-tts-group diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json new file mode 100644 index 000000000..e327be510 --- /dev/null +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json @@ -0,0 +1,80 @@ +[ + { + "test_name": "test_minicpmo_4_5_ready", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100", + "npu": "A3" + }, + "num_cards": 1 + } + }, + "core_model", + "omni" + ], + "server_params": { + "model": "openbmb/MiniCPM-o-4_5", + "extra_cli_args": [ + "--deploy-config", + "vllm_omni/deploy/minicpmo_4_5.yaml", + "--trust-remote-code" + ] + }, + "benchmark_params": [ + { + "dataset_name": "seed-tts", + "dataset_path": "zhaochenyang20/seed-tts-eval", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 10, + 40 + ], + "max_concurrency": [ + 1, + 4 + ], + "no_oversample": true, + "trust_remote_code": true, + "seed_tts_locale": "en", + "extra_body": { + "modalities": [ + "text", + "audio" + ], + "chat_template_kwargs": { + "enable_thinking": false, + "use_tts_template": true + } + }, + "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration", + "baseline": { + "H100": { + "request_throughput": [ + 0.7803, + 1.1008 + ], + "mean_ttft_ms": [ + 166.9764, + 307.7133 + ], + "mean_e2el_ms": [ + 1361.1682, + 3800.6332 + ], + "mean_audio_ttfp_ms": [ + 635.8099, + 1615.6591 + ], + "mean_audio_rtf": [ + 0.3151, + 0.8568 + ] + } + } + } + ] + } +] From 854bff7fe57383b0aae49ff4a2a59ca0003bcc03 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Wed, 12 Aug 2026 05:36:24 +0200 Subject: [PATCH 03/41] [Bugfix] Fix Wan spatial reshard boundary (#6062) Signed-off-by: Rahul Steiger Co-authored-by: Rahul Steiger --- .../distributed/test_wan_spatial_shard.py | 122 ++++++++++++++++++ .../autoencoders/wan_spatial_shard.py | 25 +++- 2 files changed, 145 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py index 72a2a45f5..beb5ac2db 100644 --- a/tests/diffusion/distributed/test_wan_spatial_shard.py +++ b/tests/diffusion/distributed/test_wan_spatial_shard.py @@ -231,6 +231,128 @@ def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.Mon assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0])) +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("split_dim", ["height", "width"]) +@pytest.mark.parametrize( + ("full_extent", "world_size"), + [ + pytest.param(5, 4, id="four-ranks-empty-tail"), + pytest.param(20, 8, id="production-eight-ranks-empty-tail"), + pytest.param(39, 2, id="two-ranks-partial-tail"), + ], +) +def test_reshard_from_trimmed_extent_clamps_and_reconstructs( + monkeypatch: pytest.MonkeyPatch, + split_dim: str, + full_extent: int, + world_size: int, +): + local_extent = (full_extent + world_size - 1) // world_size + shape = (1, 1, 1, full_extent, 1) if split_dim == "height" else (1, 1, 1, 1, full_extent) + x = torch.arange(full_extent, dtype=torch.float32).reshape(shape) + outputs = [] + + for rank in range(world_size): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group, rank=rank: (rank, world_size)) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=full_extent, + local_input_extent=local_extent, + split_dim=split_dim, + rank=rank, + world_size=world_size, + ) + ) + try: + out = wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=local_extent, + split_dim=split_dim, + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + start = rank * local_extent + valid_extent = min(local_extent, max(0, full_extent - start)) + expected = torch.zeros(local_extent, dtype=torch.float32) + if valid_extent: + expected[:valid_extent] = torch.arange(start, start + valid_extent, dtype=torch.float32) + assert out.shape[_spatial_dim_for_test(split_dim)] == local_extent + assert torch.equal(out.flatten(), expected) + outputs.append(out) + + reconstructed = torch.cat(outputs, dim=_spatial_dim_for_test(split_dim)) + reconstructed = reconstructed.narrow(_spatial_dim_for_test(split_dim), 0, full_extent) + assert torch.equal(reconstructed, x) + + +def _spatial_dim_for_test(split_dim: str) -> int: + return -2 if split_dim == "height" else -1 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_reshard_from_trimmed_extent_rejects_rank_mismatch(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2)) + x = torch.zeros((1, 1, 1, 1, 4), dtype=torch.float32) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=4, + local_input_extent=2, + split_dim="width", + rank=0, + world_size=2, + ) + ) + try: + with pytest.raises(RuntimeError, match="group_rank=1, context_rank=0"): + wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=2, + split_dim="width", + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.MonkeyPatch): + class ShrinkingWanAttentionBlock(torch.nn.Module): + def forward(self, x): + return x[..., :-1] + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 2)) + monkeypatch.setattr( + wan_spatial_shard, + "all_gather_along_dim", + lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim), + ) + module = ShrinkingWanAttentionBlock() + wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width") + x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=4, + local_input_extent=2, + split_dim="width", + rank=0, + world_size=2, + ) + ) + try: + with pytest.raises( + RuntimeError, + match=r"local_extent=2, gathered_extent=4, trimmed_extent=4, output_extent=3", + ): + module(x) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + @pytest.mark.core_model @pytest.mark.cpu def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch): diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index bda728540..07edad694 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -203,8 +203,17 @@ def reshard_from_trimmed_extent( return x dim = _spatial_dim(split_dim) + ctx = _SPATIAL_SHARD_CONTEXT.get() + if ctx is not None and ctx.rank != rank: + raise RuntimeError( + "Wan VAE spatial-shard rank mismatch while resharding: " + f"group_rank={rank}, context_rank={ctx.rank}, split_dim={split_dim!r}." + ) + valid_extent = _local_valid_extent(local_extent) - start = rank * local_extent + actual_extent = x.shape[dim] + start = min(rank * local_extent, actual_extent) + valid_extent = min(valid_extent, actual_extent - start) local = _narrow_along_dim(x, dim, start, valid_extent).contiguous() if valid_extent < local_extent: local = _pad_along_dim(local, local_extent - valid_extent, dim=dim) @@ -594,10 +603,22 @@ def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_di dim = _spatial_dim(split_dim) local_extent = x.shape[dim] gathered = all_gather_along_dim(x, group=group, dim=dim).contiguous() + gathered_extent = gathered.shape[dim] full_extent = _current_full_extent(local_extent) if full_extent is not None: - gathered = _narrow_along_dim(gathered, dim, 0, full_extent).contiguous() + trim_extent = min(full_extent, gathered_extent) + gathered = _narrow_along_dim(gathered, dim, 0, trim_extent).contiguous() + else: + trim_extent = gathered_extent out = orig_forward(gathered, *args, **kwargs) + if out.shape[dim] != trim_extent: + raise RuntimeError( + "Wan VAE attention changed the spatial extent during global gather: " + f"rank={_rank_world(group)[0]}, split_dim={split_dim!r}, " + f"local_extent={local_extent}, gathered_extent={gathered_extent}, " + f"trimmed_extent={trim_extent}, output_extent={out.shape[dim]}, " + f"context={_SPATIAL_SHARD_CONTEXT.get()!r}." + ) return reshard_from_trimmed_extent(out, local_extent=local_extent, split_dim=split_dim, group=group) module.forward = MethodType(_forward, module) From 3c3989ed04300d396b3f6b44e9c0b70a0f9cabe0 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Wed, 12 Aug 2026 11:37:01 +0800 Subject: [PATCH 04/41] docs: align quantization overview with navigation (#6074) Signed-off-by: Hongsheng Liu --- docs/user_guide/quantization/overview.md | 30 ++++++++++++------------ 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index ee36f5b60..9ceadda5c 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -12,20 +12,20 @@ For the internal architecture and backend extension points, see the | Mode | Guide | Description | Methods | |------|-------|-------------|---------| -| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 | +| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, BitsAndBytes W4, MXFP8 W8A8, MXFP4 W4A4 | | Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA | -| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale | +| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale | ## Hardware Support -| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim | -|--------|----------|-----------|----------|------------|------------|------|-----------|-------------| -| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | -| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ | -| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ | -| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | +| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim | +|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------| +| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | +| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ | +| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | +| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. FP8 on Ampere may use a weight-only path where available. @@ -43,10 +43,10 @@ otherwise. |--------|-------|------|----------------|--------| | FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | +| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Online W4 weight-only | Z-Image-Turbo; Qwen-Image/Wan2.2 not validated | Validated for Z-Image-Turbo | | ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints | | MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 | | MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy | -| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists | | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven | | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path | @@ -61,9 +61,9 @@ in BF16 unless the model guide explicitly adds support. |--------|-------|-------|----------------|--------| | ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path | | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | +| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | -| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints | | msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | @@ -76,10 +76,10 @@ attached to the intended stage rather than applied globally. |--------|-------|-------|----------------|--------| | FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | +| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated | | ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation | | MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | | MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | -| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed | | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed | | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation | @@ -106,7 +106,7 @@ config = build_quant_config({ | Component | Default quantized? | Notes | |-----------|--------------------|-------| -| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim | +| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim | | Text encoder | No | Keep BF16 unless a method-specific guide documents support | | VAE | No | Keep BF16; storage-only paths are method-specific | | Scheduler/tokenizer | No | Loaded from the base model repository | @@ -139,8 +139,8 @@ from vllm_omni.quantization import build_quant_config build_quant_config("fp8") build_quant_config({"method": "fp8", "activation_scheme": "static"}) +build_quant_config("bitsandbytes") build_quant_config("auto-round", bits=4, group_size=128) -build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"}) build_quant_config({"transformer": {"method": "fp8"}, "vae": None}) build_quant_config(None) ``` From a15b5b4bcda33a9fbd24017381d1f9423d1e840d Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Wed, 12 Aug 2026 11:37:34 +0800 Subject: [PATCH 05/41] [CI/Build] Add examples policy to PR skills (#6046) Signed-off-by: Hongsheng Liu --- .claude/skills/precheck-pr/SKILL.md | 5 +- .../precheck-pr/references/checklists.md | 13 ++++ .../precheck-pr/references/examples-policy.md | 68 +++++++++++++++++++ .claude/skills/review-pr/SKILL.md | 8 +++ 4 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 .claude/skills/precheck-pr/references/examples-policy.md diff --git a/.claude/skills/precheck-pr/SKILL.md b/.claude/skills/precheck-pr/SKILL.md index cd86ec4aa..4c8086029 100644 --- a/.claude/skills/precheck-pr/SKILL.md +++ b/.claude/skills/precheck-pr/SKILL.md @@ -1,6 +1,6 @@ --- name: precheck-pr -description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub. +description: Self-check your branch before creating a PR — catch dead code, prevent new model-specific Python examples, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub. --- # PR Pre-Check @@ -65,6 +65,8 @@ Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type **Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope. +**Also run the Examples-Policy check on every PR** using [references/examples-policy.md](references/examples-policy.md). Inspect only Python paths introduced by the diff. Treat a new model-, checkpoint-, vendor-, or family-specific Python example as blocking; do not report pre-existing example debt. + ### Step 5: Print Report ``` @@ -77,6 +79,7 @@ Pre-check report for ───────────────── ────── PR title format ✓ Code quality ⚠ 1 broad except, 2 Any hints + Examples policy ✗ new model-specific Python example PR desc integrity ✓ Registry/config ✓ Dead code ⚠ 2 warnings diff --git a/.claude/skills/precheck-pr/references/checklists.md b/.claude/skills/precheck-pr/references/checklists.md index a232403a7..b65a66713 100644 --- a/.claude/skills/precheck-pr/references/checklists.md +++ b/.claude/skills/precheck-pr/references/checklists.md @@ -47,6 +47,19 @@ Diff-scoped sweep for five fragility patterns that are pervasive in the existing Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row. +--- + +## All PRs: Examples Policy (run on every PR, quick & full) + +Inspect Python files introduced under `examples/` using +[examples-policy.md](examples-policy.md). This is a diff-scoped ratchet: existing +model-specific examples may be edited or removed without triggering the check. + +- [ ] **No new model-specific Python examples:** the diff adds, copies, or renames no Python file whose path or behavior is specific to one model, checkpoint, vendor, or model family. ✗ for any violation. Route inference through a shared task/protocol entrypoint, model request behavior through `vllm_omni/model_extras`, and commands through task docs or `recipes/`. +- [ ] **New generic example code is justified:** any new task-neutral Python file under `examples/` is an actual user-facing entrypoint rather than reusable production logic, a benchmark, an app, a setup tool, or a reproducer. ⚠ when the placement is questionable. + +Roll the result into the report as a separate **Examples policy** dimension. + ### Conventions (eyeball — no grep) Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check: diff --git a/.claude/skills/precheck-pr/references/examples-policy.md b/.claude/skills/precheck-pr/references/examples-policy.md new file mode 100644 index 000000000..bbd183961 --- /dev/null +++ b/.claude/skills/precheck-pr/references/examples-policy.md @@ -0,0 +1,68 @@ +# Examples Python Policy + +Prevent new model-specific Python examples without forcing cleanup of the +existing backlog. Run this check on every PR in quick and full modes. + +## Scope + +Set the merge base as in the main workflow, then list Python paths introduced +by the PR, including additions, copies, and renames: + +```bash +BASE="$(git merge-base HEAD origin/main)" +git diff --name-status --diff-filter=ACR --find-renames --find-copies \ + "${BASE}"...HEAD -- 'examples/**/*.py' +``` + +Inspect destination paths for copies and renames. Existing model-specific files +that are only modified or deleted are grandfathered and must not be reported. + +## Blocking classification + +Mark ✗ when a newly introduced Python file does any of the following: + +- Names or scopes a path to one model, checkpoint, vendor, or model family. +- Hard-codes a model identifier, model-specific prompt, request shape, output + conversion, or launch configuration behind an otherwise generic filename. +- Duplicates an existing task/protocol runner to demonstrate one model. +- Adds a per-model helper below a generic task hub such as `text_to_speech/`. + +Do not infer safety from a generic filename alone. For example, a file named +`text_to_audio.py` is still model-specific if it only implements one model's +contract. + +## Acceptable additions + +A new Python example may pass when it is a genuinely model-neutral, +user-facing task or protocol entrypoint. It must accept model selection through +configuration and keep model-specific behavior out of the script. + +Use ⚠ when the file is model-neutral but reusable logic appears misplaced under +`examples/`. Route code by purpose: + +| Code purpose | Preferred location | +|---|---| +| Model prompt, defaults, request/output adaptation | `vllm_omni/model_extras/` or production model modules | +| Runnable model command and validation evidence | Task documentation or `recipes/` | +| UI or interactive application | `apps/` | +| Benchmark or evaluator | `benchmarks/` | +| Regression reproducer | `tests/` | +| Download, conversion, or setup utility | `tools/` | + +If a model cannot use an existing shared runner, report the missing generic +capability. Do not solve the gap by adding a model-specific Python example. + +## Report format + +Add one row to the pre-check report: + +```text +Examples policy ✓ no new Python example paths +Examples policy ⚠ generic helper belongs in tools/ +Examples policy ✗ examples/offline_inference/new_model/end2end.py is model-specific +Examples policy ✗ examples/offline_inference/x_to_y_model_name.py is model-specific +``` + +For a blocker, name the path, the model-specific behavior, and the appropriate +shared runner or destination. Do not propose deleting unrelated existing +examples as part of the PR. diff --git a/.claude/skills/review-pr/SKILL.md b/.claude/skills/review-pr/SKILL.md index 005a06f87..d8b1b43c5 100644 --- a/.claude/skills/review-pr/SKILL.md +++ b/.claude/skills/review-pr/SKILL.md @@ -97,6 +97,7 @@ If docs and live code disagree, verify the code/tests and report the drift. | [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. | | [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. | | [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. | +| [examples-policy.md](../precheck-pr/references/examples-policy.md) | The PR adds, copies, or renames Python under `examples/`; apply the canonical policy shared with `precheck-pr`. | ### Delivery and reviewer coordination @@ -152,6 +153,13 @@ contract they protect or use only the applicable evidence checks. Apply every category in [general-checks.md](references/process/general-checks.md) before lower-priority comments. +If the diff census contains an added, copied, or renamed Python path under +`examples/`, read and apply the canonical +[examples policy](../precheck-pr/references/examples-policy.md). Treat a new +model-specific Python example as blocking. Do not flag model-specific example +debt that the PR only modifies or removes, and do not run the rest of the +author-oriented `precheck-pr` workflow. + For each changed value or behavior, trace: ```text From 9b0df3909ae06207337cff817ce47e09cbbfe697 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=87=91=E6=97=AD?= <105263726+wjinxu@users.noreply.github.com> Date: Wed, 12 Aug 2026 13:33:34 +0800 Subject: [PATCH 06/41] [Feature] timestamps for TTS via a forced-aligner pooling (#4795) Signed-off-by: wjinxu <1299461899@qq.com> Co-authored-by: xurui Co-authored-by: Claude Opus 4.8 (1M context) --- .../text_to_speech/README.md | 16 +- .../qwen3_tts/word_timestamps.py | 172 -------- .../online_serving/text_to_speech/README.md | 7 + tests/config/test_forced_aligner_injection.py | 71 ++++ .../sched/test_omni_ar_scheduler_logprobs.py | 71 ++++ .../openai_api/test_serving_speech_stream.py | 60 +-- .../test_forced_aligner_stage.py | 117 ++++++ tests/utils/test_forced_aligner.py | 26 +- .../utils/test_qwen3_force_align_processor.py | 17 - vllm_omni/config/config_factory.py | 4 + vllm_omni/config/model.py | 3 + vllm_omni/config/stage_config.py | 9 +- vllm_omni/core/sched/omni_ar_scheduler.py | 16 +- .../core/sched/omni_generation_scheduler.py | 22 +- vllm_omni/core/sched/omni_scheduler_mixin.py | 25 ++ .../deploy/qwen3_tts_forced_aligner.yaml | 14 +- vllm_omni/engine/arg_utils.py | 8 + vllm_omni/engine/orchestrator.py | 10 +- vllm_omni/engine/stage_init_utils.py | 20 +- vllm_omni/entrypoints/cli/serve.py | 12 +- vllm_omni/entrypoints/openai/api_server.py | 6 +- .../entrypoints/openai/protocol/audio.py | 10 +- .../entrypoints/openai/serving_speech.py | 88 +++- .../openai/serving_speech_stream.py | 83 ++-- .../stage_input_processors/forced_aligner.py | 173 ++++++++ vllm_omni/utils/forced_aligner.py | 375 ++++++------------ .../utils/qwen3_force_align_processor.py | 33 +- vllm_omni/worker/gpu_ar_model_runner.py | 9 +- 28 files changed, 850 insertions(+), 627 deletions(-) delete mode 100644 examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py create mode 100644 tests/config/test_forced_aligner_injection.py create mode 100644 tests/model_executor/stage_input_processors/test_forced_aligner_stage.py create mode 100644 vllm_omni/model_executor/stage_input_processors/forced_aligner.py diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 56135b55e..d246f9e45 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -372,18 +372,10 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ Streaming requires `async_chunk: true` in the stage config. ### Word Timestamps -Generate a WAV offline and a JSON sidecar with word-level timestamps from -`Qwen/Qwen3-ForcedAligner-0.6B`: -```bash -python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \ - --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \ - --text "Hello world." \ - --output-dir /tmp/qwen3_tts_timestamps -``` -The script writes `qwen3_tts_word_timestamps.wav` and -`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit, -set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT. +Word-level timestamps are currently a serving-path feature: launch +`vllm-omni serve` with `--forced-aligner` and request `word_timestamps` +(see `examples/online_serving/text_to_speech/README.md`). An offline +example is not provided in this release. ### Batched decoding The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates): diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py deleted file mode 100644 index 581be2a7d..000000000 --- a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Offline Qwen3-TTS example with word-level timestamps. - -This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with -the same shared forced aligner used by the streaming server path, then writes -the WAV and JSON sidecar: - - python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \ - --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B - -On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0`` -to avoid FlashInfer sampler JIT during warmup. -""" - -from __future__ import annotations - -import asyncio -import json -import os -from pathlib import Path -from typing import Any - -os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") -# The in-process offline aligner consumes token_classify outputs. Keeping V1 -# multiprocessing off avoids msgspec serialization issues for token-wise logits. -os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") - -import numpy as np -import soundfile as sf -import torch -from end2end import _estimate_prompt_len -from vllm.utils.argparse_utils import FlexibleArgumentParser - -from vllm_omni import Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults -from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config - - -def _default_stage_config() -> str: - repo_root = Path(__file__).resolve().parents[4] - return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml") - - -def _build_custom_voice_input(args: Any) -> dict[str, Any]: - additional_information = { - "task_type": ["CustomVoice"], - "text": [args.text], - "language": [args.language], - "speaker": [args.speaker], - "instruct": [args.instructions], - "max_new_tokens": [args.max_new_tokens], - } - return { - "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model), - "additional_information": additional_information, - } - - -def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]: - audio_data = mm["audio"] - sr_raw = mm["sr"] - sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw - sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val) - audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data - return audio_tensor.float().detach().cpu().flatten(), sample_rate - - -def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes: - samples = audio.numpy() - pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype(" list[dict[str, Any]] | None: - aligner_config = build_forced_aligner_config(args) - if aligner_config is None: - raise ValueError("--forced-aligner or --forced-aligner-config is required") - timestamps = await align( - audio=_float_audio_to_pcm16_bytes(audio), - text=args.text, - sample_rate=sample_rate, - config=aligner_config, - language=args.language, - ) - if timestamps is None: - return None - return [ - { - "word": item.word, - "start_ms": item.start_ms, - "end_ms": item.end_ms, - } - for item in timestamps - ] - - -def main(args: Any) -> None: - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - omni_kwargs = vars(args).copy() - for key in ( - "forced_aligner", - "text", - "language", - "speaker", - "instructions", - "max_new_tokens", - "output_dir", - ): - omni_kwargs.pop(key, None) - omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config() - omni_kwargs["log_stats"] = args.log_stats - omni = Omni(**omni_kwargs) - - prompt = _build_custom_voice_input(args) - final_output = None - for stage_outputs in omni.generate([prompt]): - final_output = stage_outputs.request_output - if final_output is None: - raise RuntimeError("Qwen3-TTS did not produce an output.") - - mm = final_output.outputs[0].multimodal_output - audio, sample_rate = _audio_tensor_and_sample_rate(mm) - - wav_path = output_dir / "qwen3_tts_word_timestamps.wav" - # Align the in-memory tensor, not the saved WAV. The WAV is only a demo - # artifact so users can listen to the same audio referenced by the sidecar. - timestamps = asyncio.run(_run_alignment(args, audio, sample_rate)) - sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV") - - sidecar = { - "text": args.text, - "sample_rate": sample_rate, - "audio_path": str(wav_path), - "timestamps": timestamps, - } - json_path = output_dir / "qwen3_tts_word_timestamps.json" - json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") - - print(f"Saved audio: {wav_path}") - print(f"Saved timestamps: {json_path}") - print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2)) - - -def parse_args() -> Any: - parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example") - parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name") - parser.add_argument( - "--forced-aligner", - default=None, - help="Qwen3 forced aligner model path/name", - ) - parser.add_argument( - "--forced-aligner-config", - default=None, - help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)", - ) - parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML") - parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align") - parser.add_argument("--language", default="English", help="Qwen3-TTS language field") - parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name") - parser.add_argument("--instructions", default="", help="Optional speaking style instruction") - parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens") - parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar") - parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging") - nullify_stage_engine_defaults(parser) - return parser.parse_args() - - -if __name__ == "__main__": - main(parse_args()) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 2f9d84bf3..1ea191a37 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -605,6 +605,13 @@ python qwen3_tts/streaming_speech_client.py \ The client writes one PCM file per sentence and a matching `sentence_XXX_timestamps.json` sidecar. +Non-streaming requests can also ask for timestamps: pass +`"word_timestamps": true` to `POST /v1/audio/speech` and read the +`X-Word-Timestamps` response header (JSON list of `{word, start_ms, end_ms}`, +ASCII-escaped). Past 4 KB the header is replaced by +`X-Word-Timestamps-Omitted: oversize; bytes=; limit=4096` and the audio +still returns — use the WebSocket path for long transcripts. + To *see* the alignment instead of reading a JSON sidecar, run the word-timestamp Gradio demo (server must be launched with `--forced-aligner`): ```bash diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py new file mode 100644 index 000000000..73a7bcda2 --- /dev/null +++ b/tests/config/test_forced_aligner_injection.py @@ -0,0 +1,71 @@ +import pytest + +import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry) +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY +from vllm_omni.config.stage_config import ( + DeployConfig, + StageExecutionType, +) +from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( + POOLING_OUTPUT_DECODER_PATH, + TIMESTAMPS_MODALITY, +) +from vllm_omni.utils.forced_aligner import _DEFAULT_CONFIG_PATH, inject_forced_aligner_stage + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_inject_appends_pooling_aligner_stage(): + pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + deploy = DeployConfig() + n0 = len(pipeline.stages) + + ext_pipeline, ext_deploy = inject_forced_aligner_stage( + pipeline, deploy, {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"} + ) + + # A new pooling stage is appended; the original pipeline is unchanged. + assert len(pipeline.stages) == n0 + assert len(ext_pipeline.stages) == n0 + 1 + # The original deploy config is not mutated either (double-injection safe). + assert len(deploy.stages) == 0 + assert len(ext_deploy.stages) == 1 + aligner = ext_pipeline.stages[-1] + # Pooling stage reuses LLM_AR (+ runner=pooling), no dedicated LLM_POOLING. + assert aligner.execution_type == StageExecutionType.LLM_AR + assert aligner.final_output is True + assert aligner.final_output_type == TIMESTAMPS_MODALITY + assert aligner.input_sources == (n0 - 1,) + assert aligner.requires_multimodal_data is True + assert aligner.owns_tokenizer is True + + # The deploy stage carries the separate aligner checkpoint + decoder hook. + ds = ext_deploy.stages[-1] + assert ds.engine_extras["model"] == "/models/Qwen3-ForcedAligner-0.6B" + assert ds.engine_extras["runner"] == "pooling" + assert ds.engine_extras["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH + + +def test_inject_from_config_only(): + # --forced-aligner-config alone (model in the YAML) must also inject. + pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + deploy = DeployConfig() + n0 = len(pipeline.stages) + + ext_pipeline, ext_deploy = inject_forced_aligner_stage( + pipeline, deploy, {"forced_aligner": None, "forced_aligner_config": str(_DEFAULT_CONFIG_PATH)} + ) + + assert len(ext_pipeline.stages) == n0 + 1 + assert ext_pipeline.stages[-1].final_output_type == TIMESTAMPS_MODALITY + assert ext_deploy.stages[-1].engine_extras["model"] == "Qwen/Qwen3-ForcedAligner-0.6B" + + +def test_inject_noop_without_forced_aligner(): + pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + deploy = DeployConfig() + + ext_pipeline, ext_deploy = inject_forced_aligner_stage(pipeline, deploy, {}) + + assert ext_pipeline is pipeline + assert len(ext_deploy.stages) == 0 diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 2205ab32e..604904dd0 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -27,6 +27,7 @@ _MIXIN_UPDATE_HELPERS = ( "_attach_finished_request_sets", "_attach_scheduler_stats", "_capture_omni_connector_output", + "_maybe_decode_pooling_output", ) @@ -138,6 +139,7 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace: pending_stop_after_extraction=set(), waiting_for_transfer_free=set(), _new_prompt_len_snapshot={}, + _pooling_output_decoder=None, finished_req_ids=set(), finished_req_ids_dict=defaultdict(set), kv_cache_manager=SimpleNamespace(take_events=lambda: None), @@ -267,3 +269,72 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None: assert output_by_id["bad"].new_token_ids == [] assert output_by_id["good"].new_token_ids == [8] np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8]) + + +def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace: + return SimpleNamespace( + sampled_token_ids=[[]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[pooler_tensor], + num_nans_in_logits=None, + kv_connector_output=None, + cudagraph_stats=None, + req_id_to_index={"req": 0}, + routed_experts=None, + ) + + +def _make_pooling_stub(decoder) -> tuple[_Request, SimpleNamespace]: + request = _Request("req") + request.sampling_params = SimpleNamespace(num_logprobs=None) + request.pooling_params = SimpleNamespace() + scheduler = _make_scheduler_stub([request]) + scheduler._pooling_output_decoder = decoder + scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None)) + _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False)) + return request, scheduler + + +def test_pooling_decode_success_emits_decoded_payload() -> None: + import torch + + decoded = torch.tensor([[1, 2]], dtype=torch.int32) + request, scheduler = _make_pooling_stub(lambda logits, req, hf: decoded) + + scheduler_output = SimpleNamespace( + num_scheduled_tokens={"req": 1}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + outputs = OmniARScheduler.update_from_output( + scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4)) + ) + (output,) = outputs[0].outputs + + assert output.finish_reason is FinishReason.STOP + assert output.pooling_output is decoded + + +def test_pooling_decode_failure_finishes_request_with_error() -> None: + import torch + + def raising_decoder(logits, req, hf): + raise ValueError("boom") + + request, scheduler = _make_pooling_stub(raising_decoder) + + scheduler_output = SimpleNamespace( + num_scheduled_tokens={"req": 1}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + outputs = OmniARScheduler.update_from_output( + scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4)) + ) + (output,) = outputs[0].outputs + + # Decoder failure is a request error (500), not an empty success. + assert output.finish_reason is FinishReason.ERROR + assert output.pooling_output is None + assert "pooling output decode failed" in str(request.stop_reason) diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 44e84b346..37e3509f5 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -3,6 +3,7 @@ import base64 from types import SimpleNamespace import pytest +import torch from fastapi import FastAPI, WebSocket from pytest_mock import MockerFixture from starlette.testclient import TestClient @@ -11,11 +12,24 @@ from starlette.websockets import WebSocketDisconnect from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler -from vllm_omni.utils.forced_aligner import WordTimestamp +from vllm_omni.model_executor.stage_input_processors.forced_aligner import ALIGNER_WORDS_KEY pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _fake_aligner_res(pairs, words): + """Build a stand-in for the forced-aligner stage's pooling output. + + Mirrors what rides the generator in production: ``res.outputs.data`` is a + ``[n_words, 2]`` int32 tensor of ``[start_ms, end_ms]`` and the word strings + travel in ``additional_information``. Decoded by ``extract_word_timestamps``. + """ + return SimpleNamespace( + outputs=SimpleNamespace(data=torch.tensor(pairs, dtype=torch.int32), additional_information=None), + additional_information={ALIGNER_WORDS_KEY: list(words)}, + ) + + def _build_test_app( speech_service=None, *, @@ -28,9 +42,9 @@ def _build_test_app( speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav")) speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {})) - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield (chunk, 24000) if include_sample_rate else chunk @@ -235,7 +249,7 @@ class TestStreamingSpeechWebSocket: speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False async def mock_prepare_speech_generation(request): captured_requests.append(request) @@ -311,13 +325,13 @@ class TestStreamingSpeechWebSocket: assert error["type"] == "error" assert "without --forced-aligner" in error["message"] - def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture): + def test_word_timestamps_emit_pipeline_json_frame(self, mocker: MockerFixture): captured_requests = [] speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = SimpleNamespace(model="aligner") + speech_service.forced_aligner_enabled = True async def mock_prepare_speech_generation(request): captured_requests.append(request) @@ -328,19 +342,15 @@ class TestStreamingSpeechWebSocket: first_chunk = b"\x01" * 1000 second_chunk = b"\x02" * 1000 - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + # The forced-aligner stage rides the same generator: its pooling output + # is surfaced via the ``collect`` channel once the audio has streamed. + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): for chunk in (first_chunk, second_chunk): yield (chunk, 1000) if include_sample_rate else chunk + if collect is not None: + collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"]) speech_service._generate_pcm_chunks = mock_generate_pcm_chunks - # Sentence-level: aligner runs once over the whole sentence audio. - mock_align = mocker.AsyncMock( - return_value=[ - WordTimestamp("Hello", 0, 200), - WordTimestamp("world", 200, 900), - ] - ) - mocker.patch.object(streaming_speech_module, "forced_align", mock_align) app, _ = _build_test_app(speech_service) with TestClient(app) as client: @@ -406,7 +416,6 @@ class TestStreamingSpeechWebSocket: } assert captured_requests[0].word_timestamps is True - assert mock_align.await_count == 1 def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture): # The streaming layer forwards the aligner's (already monotonic, @@ -415,21 +424,16 @@ class TestStreamingSpeechWebSocket: speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = SimpleNamespace(model="aligner") + speech_service.forced_aligner_enabled = True speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): chunk = b"\x01" * 1000 yield (chunk, 1000) if include_sample_rate else chunk + if collect is not None: + collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"]) speech_service._generate_pcm_chunks = mock_generate_pcm_chunks - mock_align = mocker.AsyncMock( - return_value=[ - WordTimestamp("Hello", 0, 1000), - WordTimestamp("world", 1000, 1200), - ] - ) - mocker.patch.object(streaming_speech_module, "forced_align", mock_align) app, _ = _build_test_app(speech_service) with TestClient(app) as client: @@ -572,7 +576,7 @@ class TestStreamingSpeechWebSocket: speech_service._generate_pcm_chunks = mocker.AsyncMock() speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False app, _ = _build_test_app(speech_service) with TestClient(app) as client: @@ -602,7 +606,7 @@ class TestStreamingSpeechWebSocket: speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {})) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): yield b"\x01\x02" @@ -696,7 +700,7 @@ class TestStreamingSpeechWebSocket: speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {})) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): yield b"\x01\x02" diff --git a/tests/model_executor/stage_input_processors/test_forced_aligner_stage.py b/tests/model_executor/stage_input_processors/test_forced_aligner_stage.py new file mode 100644 index 000000000..9b1d5ec2a --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_forced_aligner_stage.py @@ -0,0 +1,117 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( + ALIGNER_WORDS_KEY, + _extract_text, + _extract_waveform, + code2wav2aligner, + decode_pooling_output, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeCompletion: + def __init__(self, multimodal_output): + self.multimodal_output = multimodal_output + + +class _FakeOutput: + def __init__(self, multimodal_output, finished=True): + self.finished = finished + self.outputs = [_FakeCompletion(multimodal_output)] + + +def _audio_mm(n=2400, sr=24000): + return {"audio": torch.zeros(n, dtype=torch.float32), "sr": torch.tensor(sr, dtype=torch.int32)} + + +def test_extract_waveform_from_audio_key(): + wav, sr = _extract_waveform(_audio_mm(n=100, sr=16000)) + assert isinstance(wav, np.ndarray) and wav.dtype == np.float32 + assert wav.shape == (100,) + assert sr == 16000 + + +def test_extract_waveform_falls_back_to_model_outputs_and_default_sr(): + wav, sr = _extract_waveform({"model_outputs": [torch.ones(50)]}) + assert wav.shape == (50,) + assert sr == 24000 # default when sr missing + + +def test_extract_text_unwraps_additional_information_list(): + prompt = {"additional_information": {"text": ["Hello world."]}} + assert _extract_text(prompt) == "Hello world." + + +def test_code2wav2aligner_builds_text_prompt_and_audio(): + prompt = {"additional_information": {"text": ["U.S.A test"], "language": ["English"]}} + out = code2wav2aligner([_FakeOutput(_audio_mm())], prompt) + + assert len(out) == 1 + item = out[0] + # text prompt: tokenization + audio features deferred to the stage preprocessor + assert "prompt_token_ids" not in item + # prompt built from official segmentation: "U.S.A" -> "USA" + assert "USAtest" in item["prompt"] + # audio carried as (np.float32, sr) tuple + audio, sr = item["multi_modal_data"]["audio"] + assert isinstance(audio, np.ndarray) and sr == 24000 + assert item["additional_information"]["aligner_words"] == [["USA", "test"]] + assert item["additional_information"]["language"] == ["English"] + + +def test_code2wav2aligner_skips_unfinished_and_empty(): + # unfinished -> skipped + assert code2wav2aligner([_FakeOutput(_audio_mm(), finished=False)], {}) == [] + # finished but empty waveform -> skipped + empty = _FakeOutput({"audio": torch.zeros(0), "sr": torch.tensor(24000)}) + assert code2wav2aligner([empty], {"additional_information": {"text": ["hi"]}}) == [] + + +def test_decode_pooling_output_produces_word_timestamps_ms(): + # 2 words -> 4 markers (id=99). classify_num=5, seg=200ms. + hf_config = SimpleNamespace( + timestamp_token_id=99, + timestamp_segment_time=200.0, + thinker_config=SimpleNamespace(classify_num=5), + ) + # prompt: marker positions at indices 2,3,5,6 + prompt_token_ids = [1, 1, 99, 99, 1, 99, 99] + request = SimpleNamespace( + prompt_token_ids=prompt_token_ids, + additional_information={ALIGNER_WORDS_KEY: [["hello", "world"]]}, + ) + logits = torch.zeros((7, 5), dtype=torch.float32) + logits[2, 0] = 1.0 # word0 start -> bin 0 + logits[3, 2] = 1.0 # word0 end -> bin 2 + logits[5, 2] = 1.0 # word1 start -> bin 2 + logits[6, 4] = 1.0 # word1 end -> bin 4 + + ms = decode_pooling_output(logits, request, hf_config) + + # decoder returns a bare [n_words, 2] int32 tensor (not a dict) + assert ms.tolist() == [[0, 400], [400, 800]] # bins x 200ms + + +def test_decode_pooling_output_handles_hf_config_without_thinker(): + # hf_config without thinker_config falls back to itself for classify_num. + hf_config = SimpleNamespace( + timestamp_token_id=99, + timestamp_segment_time=100.0, + classify_num=4, + ) + request = SimpleNamespace( + prompt_token_ids=[99, 99], + additional_information={ALIGNER_WORDS_KEY: [["hi"]]}, + ) + logits = torch.zeros((2, 4), dtype=torch.float32) + logits[0, 1] = 1.0 + logits[1, 3] = 1.0 + + ms = decode_pooling_output(logits, request, hf_config) + assert ms.tolist() == [[100, 300]] diff --git a/tests/utils/test_forced_aligner.py b/tests/utils/test_forced_aligner.py index 4964f3900..0e15f5b64 100644 --- a/tests/utils/test_forced_aligner.py +++ b/tests/utils/test_forced_aligner.py @@ -105,9 +105,7 @@ forced_aligner: """, encoding="utf-8", ) - args = type("Args", (), {"forced_aligner": None, "forced_aligner_config": str(cfg)})() - - out = forced_aligner.build_forced_aligner_config(args) + out = forced_aligner.build_forced_aligner_config(None, str(cfg)) assert out == forced_aligner.ForcedAlignerConfig( model="Qwen/Qwen3-ForcedAligner-0.6B", @@ -127,16 +125,7 @@ def test_build_config_cli_model_overrides_yaml(tmp_path): "forced_aligner:\n model: old\n gpu_memory_utilization: 0.2\n dtype: float16\n", encoding="utf-8", ) - args = type( - "Args", - (), - { - "forced_aligner": "new", - "forced_aligner_config": str(cfg), - }, - )() - - out = forced_aligner.build_forced_aligner_config(args) + out = forced_aligner.build_forced_aligner_config("new", str(cfg)) assert out is not None # --forced-aligner overrides the YAML model; gpu_memory_utilization/dtype @@ -148,16 +137,7 @@ def test_build_config_cli_model_overrides_yaml(tmp_path): def test_build_config_from_cli_model_uses_default_yaml(): - args = type( - "Args", - (), - { - "forced_aligner": "local-aligner", - "forced_aligner_config": None, - }, - )() - - out = forced_aligner.build_forced_aligner_config(args) + out = forced_aligner.build_forced_aligner_config("local-aligner", None) assert out is not None assert out.model == "local-aligner" diff --git a/tests/utils/test_qwen3_force_align_processor.py b/tests/utils/test_qwen3_force_align_processor.py index c9af7167e..4a928d0e5 100644 --- a/tests/utils/test_qwen3_force_align_processor.py +++ b/tests/utils/test_qwen3_force_align_processor.py @@ -46,20 +46,3 @@ def test_fix_timestamp_repairs_dip_onto_monotonic_sequence(): def test_fix_timestamp_passes_through_monotonic(): assert processor.fix_timestamp([0, 100, 100, 250]) == [0, 100, 100, 250] - - -def test_resolve_timestamp_token_id_defaults_to_marker_token(): - # Regression: the default must resolve the same marker that - # build_prompt inserts, not None (which would degrade every request). - seen = {} - - class FakeTokenizer: - def convert_tokens_to_ids(self, token): - seen["token"] = token - return 151705 - - tid = processor.resolve_timestamp_token_id(FakeTokenizer()) - - assert tid == 151705 - assert seen["token"] == processor.TIMESTAMP_TOKEN - assert processor.TIMESTAMP_TOKEN in processor.build_prompt(["hello", "world"]) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index ed9b05f1e..a1778eda3 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -409,6 +409,10 @@ class StageConfigFactory: if cli_async_chunk is not None: deploy_cfg.async_chunk = bool(cli_async_chunk) + from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage + + pipeline_cfg, deploy_cfg = inject_forced_aligner_stage(pipeline_cfg, deploy_cfg, cli_overrides) + stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides) # Overlay declarative parallel strategies (opt-in) before CLI overrides. diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index f830fc37e..9164e1352 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -129,6 +129,9 @@ class OmniModelConfig(ModelConfig): model_arch: str | None = None worker_type: str | None = None engine_output_type: str | None = None + # Optional dotted path of a per-stage pooling-output decoder applied + # worker-side before IPC. Read by the AR scheduler. + pooling_output_decoder: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None stage_connector_config: dict[str, Any] = field( diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 55c4a46cc..fa8e45295 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -340,6 +340,7 @@ class StageDeployConfig: output_connectors: dict[str, str] | None = None input_connectors: dict[str, str] | None = None default_sampling_params: dict[str, Any] | None = None + default_pooling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None # === Generic stage engine fields === @@ -509,6 +510,7 @@ _STAGE_RESERVED_KEYS = frozenset( "output_connectors", "input_connectors", "default_sampling_params", + "default_pooling_params", "engine_extras", "engine_args", "runtime", @@ -562,11 +564,14 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") + kwargs["default_pooling_params"] = stage_data.get("default_pooling_params") kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args) return StageDeployConfig(**kwargs) -_DEEP_MERGE_KEYS = frozenset({"default_sampling_params", "subtalker_sampling_params", "engine_extras", "engine_args"}) +_DEEP_MERGE_KEYS = frozenset( + {"default_sampling_params", "default_pooling_params", "subtalker_sampling_params", "engine_extras", "engine_args"} +) def _deep_merge_stage(base: dict, overlay: dict) -> dict: @@ -886,6 +891,8 @@ def _build_extras( sampling.update(ps.sampling_constraints) if sampling: extras["default_sampling_params"] = sampling + if ds is not None and ds.default_pooling_params: + extras["default_pooling_params"] = dict(ds.default_pooling_params) if ds is not None and ds.output_connectors: extras["output_connectors"] = dict(ds.output_connectors) if ds is not None and ds.input_connectors: diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 7cf3b5a6d..66da231a6 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -415,6 +415,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason = None routed_experts = None + # Decode the pooling output before stop handling so a decoder + # failure finishes the request with FinishReason.ERROR (500). + try: + pooling_output_payload = self._maybe_decode_pooling_output(request, pooler_output) + except Exception as exc: + logger.exception("[pooling] decoder hook failed for request %s", req_id) + pooling_output_payload = None + request.status = RequestStatus.FINISHED_ERROR + request.stop_reason = f"pooling output decode failed: {exc}" + request.resumable = False + # Check for stop and update request status. if new_token_ids: num_sampled_tokens = len(new_token_ids) @@ -428,7 +439,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) elif request.pooling_params and pooler_output is not None: # Pooling stops as soon as there is output. - request.status = RequestStatus.FINISHED_STOPPED + if request.status != RequestStatus.FINISHED_ERROR: + request.status = RequestStatus.FINISHED_STOPPED stopped = True # If criteria returns True, it means we must STOP the request. @@ -522,7 +534,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooler_output, + pooling_output=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, prefill_stats=request.take_prefill_stats(), diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index cf3678251..ae19a41a3 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -5,6 +5,7 @@ from collections import defaultdict from typing import Any from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.logger import init_logger from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import SchedulerOutput @@ -25,6 +26,8 @@ from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestDat from vllm_omni.core.sched.utils import omni_routed_experts_for_request from vllm_omni.outputs import OmniModelRunnerOutput +logger = init_logger(__name__) + class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): def __init__(self, *args, **kwargs): @@ -442,10 +445,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): is_segment_finished = False routed_experts = None + # Decode the pooling output before stop handling so a decoder + # failure finishes the request with FinishReason.ERROR (500). + try: + pooling_output_payload = self._maybe_decode_pooling_output(request, pooler_output) + except Exception as exc: + logger.exception("[pooling] decoder hook failed for request %s", req_id) + pooling_output_payload = None + request.status = RequestStatus.FINISHED_ERROR + request.stop_reason = f"pooling output decode failed: {exc}" + request.resumable = False + # One-shot generation request: finish after its current input unit # has been fully processed. if ( - request.status == RequestStatus.FINISHED_STOPPED + request.status == RequestStatus.FINISHED_ERROR + or request.status == RequestStatus.FINISHED_STOPPED or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens) or ( self.chunk_transfer_adapter is not None @@ -453,7 +468,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): and request.num_computed_tokens >= len(request.prompt_token_ids) ) ): - request.status = RequestStatus.FINISHED_STOPPED + if request.status != RequestStatus.FINISHED_ERROR: + request.status = RequestStatus.FINISHED_STOPPED # Optional: set a stop_reason for front-end clarity # (does not affect protocol) stopped = True @@ -508,7 +524,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooler_output, + pooling_output=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, prefill_stats=request.take_prefill_stats(), diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index d0a4cefab..e0df774c0 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -5,10 +5,12 @@ import time from collections.abc import Iterable from typing import Any +import torch from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import ( @@ -80,6 +82,29 @@ class OmniSchedulerMixin: else None ) self._latest_omni_connector_output = None + # Optional per-stage pooling-output decoder hook (dotted path in + # model_config); applied worker-side before IPC. + self._pooling_output_decoder = None + _decoder_path = getattr(model_config, "pooling_output_decoder", None) + if _decoder_path: + self._pooling_output_decoder = resolve_obj_by_qualname(str(_decoder_path)) + + def _maybe_decode_pooling_output(self, request: Request, pooler_output: Any) -> Any: + """Apply the stage's pooling-output decoder hook to the pooler tensor + before IPC, or pass it through unchanged when none is configured. + Decoder exceptions propagate; callers fail the request with + FinishReason.ERROR rather than emitting an empty success.""" + if self._pooling_output_decoder is None: + return pooler_output + if pooler_output is None or getattr(request, "pooling_params", None) is None: + return pooler_output + if not isinstance(pooler_output, torch.Tensor): + return pooler_output + return self._pooling_output_decoder( + pooler_output, + request, + self.vllm_config.model_config.hf_config, + ) def _free_input_coordinator_request(self, request_id: str) -> None: """Prune full-payload coordinator state for a completed request.""" diff --git a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml index 41145343d..2aecdbf9b 100644 --- a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml +++ b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml @@ -1,10 +1,12 @@ -# Reference deploy knobs for Qwen/Qwen3-ForcedAligner-0.6B, the shared forced -# aligner for streaming TTS word timestamps. Pass with --forced-aligner-config; -# --forced-aligner overrides the model path without editing this file. +# Default deploy knobs for Qwen/Qwen3-ForcedAligner-0.6B, the TTS word-timestamp +# aligner. These configure the pooling stage that --forced-aligner appends to +# the pipeline tail (see utils.forced_aligner.inject_forced_aligner_stage); they are +# NOT a standalone engine. Pass --forced-aligner-config to override this file; +# --forced-aligner sets just the model path. # -# Run as LLM(runner="pooling") with pooling_task="token_classify". Prompt tokens, -# chat template and word segmentation are Qwen-specific and live in -# qwen3_force_align_processor (not configurable here). +# The stage runs the model with runner="pooling" + pooling_task="token_classify". +# Prompt tokens, chat template and word segmentation are Qwen-specific and live +# in qwen3_force_align_processor (not configurable here). forced_aligner: model: Qwen/Qwen3-ForcedAligner-0.6B diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a66f39ce3..1d9eb803b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -191,6 +191,8 @@ class OmniEngineArgs(EngineArgs): quantization_config: Any | None = None force_cutlass_fp8: bool | None = None worker_type: str | None = None + # Dotted path of a per-stage pooling-output decoder applied worker-side. + pooling_output_decoder: str | None = None task_type: str | None = None worker_cls: str = None # type: ignore[assignment] # Upstream default is "auto"; omni resolves # in __post_init__ based on worker_type (ar/generation), so None is safe here. @@ -380,6 +382,7 @@ class OmniEngineArgs(EngineArgs): model_stage=self.model_stage, model_arch=self.model_arch, worker_type=self.worker_type, + pooling_output_decoder=self.pooling_output_decoder, engine_output_type=self.engine_output_type, hf_config_name=self.hf_config_name, custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, @@ -475,6 +478,11 @@ class OrchestratorArgs: # === Mode Switches (orchestrator reads, DeployConfig redistributes) === async_chunk: bool | None = None + # === Forced aligner (orchestrator injects a pooling stage; never a per-stage knob) === + forced_aligner: str | None = None + forced_aligner_config: str | None = None + forced_aligner_device: str | None = None + # === Observability === log_stats: bool = False enable_orch_monitor: bool = False diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index c539e9559..5646e82eb 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -1543,11 +1543,19 @@ class Orchestrator: return request processor = self._get_stage_input_processor(next_stage_id) + # A pooling stage is driven by PoolingParams; vLLM validates them against + # the stage's supported tasks, so advertise the model's pooling tasks (or + # the configured task) instead of generate. + if isinstance(params, PoolingParams): + model_cfg = next_pool.stage_vllm_config.model_config + supported_tasks = tuple(getattr(model_cfg, "supported_tasks", ()) or ()) or (params.task,) + else: + supported_tasks = ("generate",) request = processor.process_inputs( request_id=req_id, prompt=next_input, params=params, - supported_tasks=("generate",), + supported_tasks=supported_tasks, arrival_time=_time.time(), resumable=resumable, ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index b0a95d22c..c2e1ab5bc 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -20,6 +20,7 @@ from dataclasses import dataclass, fields, replace from typing import Any, Literal, cast from vllm.logger import init_logger +from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams from vllm.tokenizers import cached_tokenizer_from_config @@ -395,8 +396,19 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: final_output_type: str | None = stage_config.final_output_type default_sp = _to_dict(_get_attr_or_item(stage_config, "default_sampling_params", {})) - SPClass = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams - default_sampling_params: OmniSamplingParams = SPClass(**default_sp) + # A pooling stage carries its task via default_pooling_params, set where the + # stage is declared. + default_pp = _to_dict(_get_attr_or_item(stage_config, "default_pooling_params", {})) + # A pooling stage is an LLM stage run with runner="pooling" (vLLM's + # is_pooling_model signal); pick params by that signal, not execution_type. + is_pooling = str(engine_args.get("runner", "")).lower() == "pooling" + default_params: OmniSamplingParams | PoolingParams + if stage_type == "diffusion": + default_params = OmniDiffusionSamplingParams(**default_sp) + elif is_pooling: + default_params = PoolingParams(**default_pp) + else: # generative llm: ar / generation + default_params = SamplingParams(**default_sp) custom_process_input_func: Callable | None = None _cpif_path = _get_attr_or_item(stage_config, "custom_process_input_func") @@ -428,7 +440,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: engine_input_source=engine_input_source, final_output=final_output, final_output_type=final_output_type, - default_sampling_params=default_sampling_params, + default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, @@ -448,7 +460,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: engine_input_source=engine_input_source, final_output=final_output, final_output_type=final_output_type, - default_sampling_params=default_sampling_params, + default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 815ce4d6e..dfa222cd1 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -237,9 +237,8 @@ class OmniServeCommand(CLISubcommand): "VoiceDesign, or Base, while diffusion models may use it to select " "task-specific weights. If omitted, the model default is used.", ) - # Forced aligner / word timestamps. --forced-aligner is the opt-in - # toggle; heavier knobs (gpu_memory_utilization, dtype, max_model_len) - # live in the deploy YAML passed via --forced-aligner-config. + # Forced aligner / word timestamps. Either flag opts in (--forced-aligner-config + # alone works when its YAML sets the model); heavier knobs live in that YAML. omni_config_group.add_argument( "--forced-aligner", type=str, @@ -260,6 +259,13 @@ class OmniServeCommand(CLISubcommand): "flag, when set, overrides the YAML model field." ), ) + omni_config_group.add_argument( + "--forced-aligner-device", + type=str, + default=None, + help="Device(s) for the forced-aligner stage (e.g. '2'). Defaults to " + "sharing an existing stage's GPU when unset.", + ) omni_config_group.add_argument( "--deploy-config", type=str, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 940f62f2e..0ea0f3a3f 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1127,7 +1127,11 @@ async def omni_init_app_state( state.openai_serving_models, request_logger=request_logger, model_name=model_name, - forced_aligner_config=build_forced_aligner_config(args), + forced_aligner_enabled=build_forced_aligner_config( + getattr(args, "forced_aligner", None), + getattr(args, "forced_aligner_config", None), + ) + is not None, ) # Warm up speech pipeline (CUDA Graph capture, torch.compile) so the first diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 1564c4523..b2b55bc52 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -161,10 +161,12 @@ class OpenAICreateSpeechRequest(BaseModel): word_timestamps: bool = Field( default=False, description=( - "When true, the server runs a shared forced aligner alongside the streamed " - "audio and emits per-chunk word timestamps. Requires the server to be " - "launched with --forced-aligner pointing at an aligner model. No effect " - "when streaming is off." + "When true, non-streaming responses carry per-word timestamps in the " + "X-Word-Timestamps header (JSON list of {word, start_ms, end_ms}, " + "ASCII-escaped; replaced by X-Word-Timestamps-Omitted past 4 KB). " + "Requires the server to be launched with --forced-aligner (400 otherwise). " + "Not supported with stream=true; for streaming use the WebSocket " + "/v1/audio/speech/stream path." ), ) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 6b7ca1318..88e7624f1 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -423,7 +423,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def __init__(self, *args, **kwargs): self.model_name = kwargs.pop("model_name", None) - self.forced_aligner_config: Any | None = kwargs.pop("forced_aligner_config", None) + # True when the server was launched with --forced-aligner (a pooling + # aligner stage is appended to the pipeline). Gates word_timestamps. + self.forced_aligner_enabled: bool = bool(kwargs.pop("forced_aligner_enabled", False)) super().__init__(*args, **kwargs) self._init_speaker_storage() @@ -2186,6 +2188,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_start_s: float | None = None, include_sample_rate: bool = False, usage_acc: SpeechOutputTokenCounter | None = None, + collect: dict | None = None, ): """Generate audio chunks for streaming response. @@ -2218,6 +2221,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): usage_acc.observe(res) audio_output, audio_key = self._extract_audio_output(res) if audio_key is None: + # Stash the aligner's timestamps output for streaming callers. + if collect is not None and self._is_timestamps_output(res): + collect["aligner_res"] = res continue sr_raw = audio_output.get("sr") @@ -2406,6 +2412,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): data = json.dumps(payload, separators=(",", ":")) yield f"event: speech.audio.error\ndata: {data}\n\n" + @staticmethod + def _is_timestamps_output(res) -> bool: + """True when ``res`` is the forced-aligner stage's terminal timestamps output.""" + from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY + + return getattr(res, "final_output_type", None) == TIMESTAMPS_MODALITY + @staticmethod def _extract_audio_output(res) -> tuple[dict | None, str | None]: """Return (audio_output dict, audio key) or (None, None). @@ -3332,11 +3345,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): stage0_params.extra_args = {} stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed)) + # When word_timestamps is requested, also ask for the aligner stage's + # output so the orchestrator drives the request through the forced-aligner + # stage (final_stage_id extends to it). Harmless if no aligner stage exists. + output_modalities = ["audio"] + if getattr(request, "word_timestamps", False): + from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY + + output_modalities.append(TIMESTAMPS_MODALITY) + generator = self.engine_client.generate( prompt=prompt, request_id=request_id, sampling_params_list=sampling_params_list, - output_modalities=["audio"], + output_modalities=output_modalities, ) self._track_ref_audio_artifact_warmup( request_id, @@ -3345,17 +3367,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) return request_id, generator, tts_params - async def _generate_pcm_chunks(self, generator, request_id: str, *, include_sample_rate: bool = False): + async def _generate_pcm_chunks( + self, generator, request_id: str, *, include_sample_rate: bool = False, collect: dict | None = None + ): """Yield raw PCM byte chunks from the engine generator. Delegates to ``_generate_audio_chunks`` with ``response_format="pcm"``. Used by the WebSocket streaming handler and ``_iter_pcm_audio_bytes``. + ``collect`` (when given) receives the forced-aligner stage's pooling + output under ``"aligner_res"`` for downstream word-timestamp extraction. """ async for chunk in self._generate_audio_chunks( generator, request_id, response_format="pcm", include_sample_rate=include_sample_rate, + collect=collect, ): yield chunk @@ -3375,6 +3402,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id: str | None = None, usage_out: list[SpeechTokenUsage] | None = None, has_inline_ref_audio: bool | None = None, + collect: dict | None = None, ) -> tuple[bytes | str, str]: # ``usage_out`` is an opt-in output channel: when a list is passed, the # computed SpeechTokenUsage is appended to it. The return stays a @@ -3398,9 +3426,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Non-streaming is FINAL_ONLY, so the stage-0 output carries the full # token sequence; the counter records its length for output_tokens. usage_acc = SpeechOutputTokenCounter() + audio_res: OmniRequestOutput | None = None + aligner_res: OmniRequestOutput | None = None async for res in generator: final_output = res usage_acc.observe(res) + # The generator yields both the audio output (Code2Wav) and, with + # a forced-aligner stage, a timestamps output. Keep the audio res + # for the WAV and the aligner res for word timestamps. + if self._is_timestamps_output(res): + aligner_res = res + else: + _, audio_key = self._extract_audio_output(res) + if audio_key is not None: + audio_res = res if not is_moss: continue try: @@ -3422,10 +3461,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if final_output is None: raise ValueError("No output generated from the model.") - audio_output, audio_key = self._extract_audio_output(final_output) + # Extract audio from the audio-bearing res (not necessarily the last + # yielded one, which may be the aligner's timestamps output). + audio_source = audio_res if audio_res is not None else final_output + audio_output, audio_key = self._extract_audio_output(audio_source) if audio_key is None: raise ValueError("TTS model did not produce audio output.") + # Surface forced-aligner word timestamps to the caller (set as a + # response header) when requested and an aligner stage produced them. + if collect is not None and getattr(request, "word_timestamps", False): + from vllm_omni.utils.forced_aligner import extract_word_timestamps + + ts = ( + extract_word_timestamps(aligner_res, request.input, getattr(request, "language", None)) + if aligner_res is not None + else None + ) + if ts is not None: + collect["word_timestamps"] = ts + audio_tensor = audio_output[audio_key] sr_raw = audio_output.get("sr", 24000) sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw @@ -3695,6 +3750,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "/v1/audio/speech/stream path. Use session.config with " "stream_audio=true and response_format='pcm'." ) + if request.word_timestamps and not self.forced_aligner_enabled: + # Fail loud instead of silently returning 200 with no timestamps header. + return self.create_error_response( + "word_timestamps=true requires the server to be launched with --forced-aligner." + ) if request.is_raw_audio_stream(): response_format, error = self._validate_speech_streaming_request( @@ -3739,7 +3799,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): media_type="text/event-stream", ) - audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id) + collect: dict = {} + audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id, collect=collect) total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.info( "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d", @@ -3747,7 +3808,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): total_ms, len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)), ) - return Response(content=audio_bytes, media_type=media_type) + headers = {} + if collect.get("word_timestamps") is not None: + # Default ensure_ascii keeps the header latin-1 encodable (non-ASCII words \uXXXX-escaped). + ts_json = json.dumps(collect["word_timestamps"]) + # Cap at 4 KB: oversized headers turn into opaque 502s at common reverse-proxy defaults. + if len(ts_json) <= 4096: + headers["X-Word-Timestamps"] = ts_json + else: + # Marker header so clients can tell an oversized alignment from no alignment. + headers["X-Word-Timestamps-Omitted"] = f"oversize; bytes={len(ts_json)}; limit=4096" + logger.warning( + "X-Word-Timestamps header omitted: %d bytes exceeds the 4 KB budget " + "(use the WebSocket streaming path for long transcripts)", + len(ts_json), + ) + return Response(content=audio_bytes, media_type=media_type, headers=headers) except asyncio.CancelledError: total_ms = (time.perf_counter() - request_start_s) * 1000.0 diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 102cfda45..2b71ad6c7 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -61,8 +61,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech -from vllm_omni.utils.forced_aligner import ForcedAlignerLoadError -from vllm_omni.utils.forced_aligner import align as forced_align +from vllm_omni.utils.forced_aligner import extract_word_timestamps logger = init_logger(__name__) @@ -273,7 +272,7 @@ class OmniStreamingSpeechHandler: # Reject unmet word-timestamps preconditions early with a clear reason. if config.word_timestamps: - if self._speech_service.forced_aligner_config is None: + if not self._speech_service.forced_aligner_enabled: await self._send_error( websocket, "word_timestamps=true but the server was launched without " @@ -285,7 +284,8 @@ class OmniStreamingSpeechHandler: await self._send_error( websocket, "word_timestamps=true requires stream_audio=true and " - "response_format='pcm' (the aligner consumes raw PCM).", + "response_format='pcm' (timestamps ride the per-sentence " + "PCM audio.chunk stream).", ) return @@ -395,19 +395,20 @@ class OmniStreamingSpeechHandler: ) -> int: """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence. - Forward each PCM chunk live (``timestamps: null``) while buffering the - sentence audio, then run the forced aligner once over the whole - sentence and emit a final empty-audio ``audio.chunk`` with the word - timestamps. On aligner failure timestamps is ``null``; for silence it - is ``[]`` (audio always flows regardless). + Forward each PCM chunk live (``timestamps: null``). The forced-aligner + pipeline stage (appended when the server is launched with + ``--forced-aligner``) consumes the synthesized audio internally and its + pooling output rides the same generator, so once the audio finishes we + pull the word timestamps straight off that aligner output and emit a + final empty-audio ``audio.chunk`` carrying them. Timestamps is ``null`` + when the aligner produced none; audio always flows regardless. """ - aligner_config = self._speech_service.forced_aligner_config - assert aligner_config is not None # gated by the precondition check - - sentence_audio = bytearray() + audio_bytes_seen = 0 total_bytes = 0 sample_rate = _PCM_SAMPLE_RATE chunk_id = 0 + # Receives the aligner stage's pooling output from the generator. + collect: dict = {} async def send_chunk( chunk: bytes, @@ -437,62 +438,30 @@ class OmniStreamingSpeechHandler: generator, request_id, include_sample_rate=True, + collect=collect, ) ) as stream: async for chunk, chunk_sample_rate in stream: sample_rate = chunk_sample_rate - chunk_start_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) - sentence_audio.extend(chunk) - chunk_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + chunk_start_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + audio_bytes_seen += len(chunk) + chunk_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) total_bytes += len(chunk) # Audio first, timestamps after the whole sentence is aligned. await send_chunk(chunk, chunk_sample_rate, None, chunk_start_ms, chunk_end_ms) - # Single alignment pass over the full sentence, then emit timestamps. - # A load/config failure is permanent, so surface the reason once; audio - # has already streamed, so the trailing frame still carries null. - try: - timestamps_payload = await self._align_sentence( - audio=bytes(sentence_audio), - text=sentence_text, - sample_rate=sample_rate, - config=aligner_config, - language=language, - ) - except ForcedAlignerLoadError as exc: - await self._send_error(websocket, f"forced aligner unavailable: {exc}") - timestamps_payload = None - sentence_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + # Pull word timestamps off the aligner stage's pooling output (it rode + # the same generator); extract_word_timestamps re-segments the sentence + # text for the word strings when the aligner output doesn't carry them. + aligner_res = collect.get("aligner_res") + timestamps_payload = ( + extract_word_timestamps(aligner_res, sentence_text, language) if aligner_res is not None else None + ) + sentence_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) await send_chunk(b"", sample_rate, timestamps_payload, 0, sentence_end_ms) return total_bytes - @staticmethod - async def _align_sentence( - *, - audio: bytes, - text: str, - sample_rate: int, - config, - language: str | None = None, - ) -> list[dict] | None: - """Convert a sentence alignment into JSON word-timestamp dicts. - - Returns ``None`` on aligner failure, ``[]`` when it ran but produced no - tokens. Monotonic, non-overlapping bounds are guaranteed by the decoder. - ``language`` is forwarded to word segmentation. - """ - aligned = await forced_align( - audio=audio, - text=text, - sample_rate=sample_rate, - config=config, - language=language, - ) - if aligned is None: - return None - return [{"word": ts.word, "start_ms": ts.start_ms, "end_ms": ts.end_ms} for ts in aligned] - @staticmethod async def _send_error(websocket: WebSocket, message: str) -> None: """Send an error message to the client.""" diff --git a/vllm_omni/model_executor/stage_input_processors/forced_aligner.py b/vllm_omni/model_executor/stage_input_processors/forced_aligner.py new file mode 100644 index 000000000..7646c6f73 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/forced_aligner.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""stage input processor for a forced-aligner stage. + +Bridges a Code2Wav (audio) stage to a pooling forced-aligner stage: + + Talker (LLM_AR) -> Code2Wav (LLM_GENERATION) -> ForcedAligner (LLM_AR + runner=pooling) + +``code2wav2aligner`` is a synchronous ``custom_process_input_func`` (the same +shape as ``qwen3_tts.talker2code2wav``). It runs in the orchestrator process, +reads the synthesized waveform from the upstream Code2Wav output and the +ground-truth text from the original request, then builds the aligner prompt +(```` markers per word) + audio multimodal input. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from vllm.logger import init_logger + +from vllm_omni.model_executor.stage_input_processors.tts_utils import ( + extract_language_from_prompt, +) +from vllm_omni.utils.qwen3_force_align_processor import build_prompt, segment_words + +logger = init_logger(__name__) + +# Shared keys / names for the forced-aligner stage payloads, so the producer +# (this processor), the worker-side decode, and the serving consumer don't +# duplicate magic strings. +ALIGNER_WORDS_KEY = "aligner_words" +ALIGNER_DURATION_MS_KEY = "aligner_audio_duration_ms" +TIMESTAMPS_MODALITY = "timestamps" +# Dotted path of the pooling-output decoder hook (read by the AR scheduler for +# a stage whose engine_args set ``pooling_output_decoder``). +POOLING_OUTPUT_DECODER_PATH = "vllm_omni.model_executor.stage_input_processors.forced_aligner.decode_pooling_output" + + +def decode_pooling_output(logits: Any, request: Any, hf_config: Any) -> Any: + """Worker-side decoder for the forced-aligner pooling stage. + + Converts the aligner's ``[n_token, classify_num]`` logits into a bare + ``[n_words, 2]`` int32 tensor of per-word ``[start_ms, end_ms]`` (word + strings travel separately via ``additional_information``). + """ + from vllm_omni.engine.serialization import deserialize_additional_information + from vllm_omni.utils.forced_aligner import _decode_timestamps + + ts_id = int(getattr(hf_config, "timestamp_token_id")) + thinker_cfg = getattr(hf_config, "thinker_config", hf_config) + classify_num = int(getattr(thinker_cfg, "classify_num")) + seg_ms = float(getattr(hf_config, "timestamp_segment_time")) + + prompt_token_ids = list(getattr(request, "prompt_token_ids", []) or []) + positions = [i for i, t in enumerate(prompt_token_ids) if t == ts_id] + + info = deserialize_additional_information(getattr(request, "additional_information", None)) or {} + words_field = info.get(ALIGNER_WORDS_KEY) + words = words_field[0] if isinstance(words_field, list) and words_field else [] + + # Real waveform duration (stashed by code2wav2aligner) drives the decode clamp; fall back to the + # theoretical ceiling (clamp no-op) if absent. + duration_field = info.get(ALIGNER_DURATION_MS_KEY) + duration_ms = ( + float(duration_field[0]) + if isinstance(duration_field, list) and duration_field + else float(classify_num) * seg_ms + ) + + timestamps = _decode_timestamps( + logits=logits, + words=list(words), + timestamp_positions=positions, + classify_num=classify_num, + timestamp_segment_time_ms=seg_ms, + audio_duration_ms=duration_ms, + ) + # No failure sentinel: an empty [0, 2] tensor decodes downstream to timestamps: null. + ms = torch.tensor( + [[t.start_ms, t.end_ms] for t in timestamps], + dtype=torch.int32, + ).reshape(-1, 2) + return ms + + +def _extract_waveform(multimodal_output: dict[str, Any]) -> tuple[np.ndarray, int]: + """Pull a mono float32 waveform + sample rate from a Code2Wav output.""" + wav = multimodal_output.get("audio") + if wav is None: + wav = multimodal_output.get("model_outputs") + if isinstance(wav, list): + wav = wav[0] if wav else None + if isinstance(wav, torch.Tensor): + wav_np = wav.detach().cpu().to(torch.float32).flatten().numpy() + elif isinstance(wav, np.ndarray): + wav_np = wav.astype(np.float32).flatten() + else: + wav_np = np.zeros(0, dtype=np.float32) + + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + if isinstance(sr_raw, torch.Tensor): + sr = int(sr_raw.reshape(-1)[-1].item()) if sr_raw.numel() > 0 else 24000 + elif sr_raw is None: + sr = 24000 + else: + sr = int(sr_raw) + return wav_np, sr + + +def _extract_text(prompt: Any, index: int = 0) -> str: + """Read the ground-truth text the aligner needs from the original prompt.""" + if not isinstance(prompt, dict): + return "" + text = (prompt.get("additional_information") or {}).get("text") + if isinstance(text, list): + return str(text[index]) if index < len(text) else "" + return str(text) if text else "" + + +def code2wav2aligner( + source_outputs: list[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, + _streaming_context: Any | None = None, +) -> list[Any]: + """Build forced-aligner stage inputs from finished Code2Wav audio. + + Emits a **text prompt** (``prompt`` string + ``multi_modal_data`` audio). + Tokenization and audio feature extraction are deferred to the aligner + stage's input preprocessor (run by the orchestrator), so this function + needs neither the aligner tokenizer nor its mm processor and stays + unit-testable on CPU. + """ + from vllm_omni.inputs.data import OmniTextPrompt + + aligner_inputs: list[OmniTextPrompt] = [] + for i, src in enumerate(source_outputs): + if not getattr(src, "finished", False): + # Sentence-final alignment: only run once the audio is complete. + continue + mm = src.outputs[0].multimodal_output or {} + wav_np, sample_rate = _extract_waveform(mm) + if wav_np.size == 0: + logger.warning("code2wav2aligner: empty waveform for request; skipping alignment.") + continue + + text = _extract_text(prompt, index=i) + language_list = extract_language_from_prompt(prompt, index=i) + language = language_list[0] if language_list else None + words = segment_words(text, language) + prompt_str = build_prompt(words) + + additional_information: dict[str, Any] = { + "text": [text], + ALIGNER_WORDS_KEY: [words], + ALIGNER_DURATION_MS_KEY: [float(wav_np.size) / float(sample_rate) * 1000.0], + } + if language_list is not None: + additional_information["language"] = language_list + + aligner_inputs.append( + OmniTextPrompt( + prompt=prompt_str, + multi_modal_data={"audio": (wav_np, sample_rate)}, + additional_information=additional_information, + ) + ) + return aligner_inputs diff --git a/vllm_omni/utils/forced_aligner.py b/vllm_omni/utils/forced_aligner.py index 2a97df3d8..3b578ee51 100644 --- a/vllm_omni/utils/forced_aligner.py +++ b/vllm_omni/utils/forced_aligner.py @@ -1,30 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared forced aligner for streaming TTS word timestamps. +"""Forced-aligner config + word-timestamp decoding for TTS. -Hosts one lazy-loaded, in-process pooling ``vllm.LLM`` shared by the TTS -frontend. ``llm.encode`` is sync/blocking, so :func:`align` runs it in -``asyncio.to_thread``, serialized on a lock since the offline LLM is not -thread-safe. See :func:`align` for the return-value contract. +The aligner runs as a pooling stage appended to the pipeline by +``--forced-aligner`` (see :func:`inject_forced_aligner_stage`). Qwen-specific +word segmentation / prompt building / marker repair live in +:mod:`vllm_omni.utils.qwen3_force_align_processor`. """ from __future__ import annotations -import asyncio import logging -import os -import threading -from dataclasses import dataclass, field +from dataclasses import dataclass, replace from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import yaml from vllm_omni.utils import qwen3_force_align_processor as _processor -logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from vllm_omni.config.stage_config import DeployConfig, PipelineConfig +logger = logging.getLogger(__name__) _DEFAULT_CONFIG_PATH = Path(__file__).resolve().parents[1] / "deploy" / "qwen3_tts_forced_aligner.yaml" @@ -42,7 +41,8 @@ class WordTimestamp: @dataclass(frozen=True, slots=True) class ForcedAlignerConfig: - """Plain-data config for constructing the aligner ``vllm.LLM`` (lazy-loaded).""" + """Plain-data config for the forced-aligner pipeline stage; fields are + consumed by :func:`inject_forced_aligner_stage`.""" model: str runner: str | None = None @@ -52,7 +52,6 @@ class ForcedAlignerConfig: dtype: str | None = None max_model_len: int | None = None trust_remote_code: bool | None = None - extra_llm_kwargs: dict[str, Any] = field(default_factory=dict) def _load_forced_aligner_yaml(path: str | Path) -> dict[str, Any]: @@ -61,8 +60,8 @@ def _load_forced_aligner_yaml(path: str | Path) -> dict[str, Any]: return dict(raw.get("forced_aligner", raw)) -def build_forced_aligner_config(args: Any) -> ForcedAlignerConfig | None: - """Build a config from CLI args, or ``None`` when the feature is off. +def build_forced_aligner_config(model: str | None, config_path: str | None = None) -> ForcedAlignerConfig | None: + """Build a config from the CLI values, or ``None`` when the feature is off. Precedence (lowest to highest): the packaged default YAML (:data:`_DEFAULT_CONFIG_PATH`, Qwen deploy defaults) -> a user YAML passed @@ -70,254 +69,94 @@ def build_forced_aligner_config(args: Any) -> ForcedAlignerConfig | None: feature is off (returns ``None``) unless a model resolves from this chain. Per-field overrides such as ``gpu_memory_utilization`` live in the YAML. """ - config_path = getattr(args, "forced_aligner_config", None) config_data: dict[str, Any] = {} - model = getattr(args, "forced_aligner", None) if config_path or model: config_data.update(_load_forced_aligner_yaml(_DEFAULT_CONFIG_PATH)) if config_path: config_data.update(_load_forced_aligner_yaml(config_path)) if model: config_data["model"] = str(model) - model = config_data.get("model") - if not model: + if not config_data.get("model"): return None allowed = set(ForcedAlignerConfig.__dataclass_fields__) return ForcedAlignerConfig(**{k: v for k, v in config_data.items() if k in allowed}) -class ForcedAlignerLoadError(RuntimeError): - """The aligner model/config could not be loaded (permanent until restart). - - Distinguished from a per-request alignment failure so the streaming layer - can report a clear reason once rather than degrading every request to - ``timestamps: null``. +def inject_forced_aligner_stage( + pipeline: PipelineConfig, + deploy: DeployConfig, + cli_overrides: dict[str, Any], +) -> tuple[PipelineConfig, DeployConfig]: + """Append a forced-aligner pooling stage to the pipeline tail when + ``--forced-aligner`` or ``--forced-aligner-config`` resolves an aligner + model; no-op otherwise. The stage runs the aligner model with + ``runner="pooling"``, consuming the previous stage's audio and emitting a + terminal word-timestamps side-output. """ + if not (cli_overrides.get("forced_aligner") or cli_overrides.get("forced_aligner_config")): + return pipeline, deploy + from vllm_omni.config.stage_config import ( + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, + ) + from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( + POOLING_OUTPUT_DECODER_PATH, + TIMESTAMPS_MODALITY, + ) -# --- Singleton state --- -# A single LLM serves the whole API server. ``_lock`` guards the lazy -# constructor; ``_encode_lock`` serializes ``encode`` because the offline -# ``vllm.LLM`` is a sync batch interface that is not safe to call from multiple -# threads at once (concurrent WebSocket sessions otherwise race the engine). -_lock = threading.Lock() -_encode_lock = asyncio.Lock() -_llm: Any = None -_classify_num: int | None = None -_timestamp_token_id: int | None = None -_timestamp_segment_time_ms: float | None = None -_loaded_config: ForcedAlignerConfig | None = None + fa = build_forced_aligner_config( + cli_overrides.get("forced_aligner"), + cli_overrides.get("forced_aligner_config"), + ) + if fa is None: + return pipeline, deploy + _FA = "vllm_omni.model_executor.stage_input_processors.forced_aligner" + new_id = len(pipeline.stages) + aligner_ps = StagePipelineConfig( + stage_id=new_id, + model_stage="forced_aligner", + # Pooling stage = LLM_AR stage run with runner="pooling" (vLLM derives + # is_pooling_model from runner_type). + execution_type=StageExecutionType.LLM_AR, + input_sources=(new_id - 1,), + final_output=True, + final_output_type=TIMESTAMPS_MODALITY, + owns_tokenizer=True, + requires_multimodal_data=True, + model_arch=fa.architecture, + custom_process_input_func=f"{_FA}.code2wav2aligner", + ) + extended = replace(pipeline, stages=pipeline.stages + (aligner_ps,)) -async def align( - *, - audio: bytes, - text: str, - sample_rate: int, - config: ForcedAlignerConfig, - language: str | None = None, -) -> list[WordTimestamp] | None: - """Run one forced-alignment pass. - - Args: - audio: Signed-int16 little-endian mono PCM bytes. - text: Ground-truth text whose words to align. - sample_rate: Sample rate of ``audio`` in Hz. - config: Aligner config (same instance for every call across the - server's lifetime; reload requires a server restart). - language: Optional language hint for word segmentation. ``None`` / - ``"auto"`` use the space + Chinese-mixed path; ``"japanese"`` / - ``"korean"`` (or their codes) need ``qwen_asr`` installed for - faithful tokenisation, else they degrade to whitespace splitting. - - Returns: - List of :class:`WordTimestamp` on success (possibly empty for - silence / no aligned tokens), ``None`` on a per-request alignment - failure (decode error). - - Raises: - ForcedAlignerLoadError: the model/config could not be loaded. This is - permanent until restart, so callers surface it once instead of - degrading every request to ``None``. - """ - async with _encode_lock: - try: - await asyncio.to_thread(_ensure_loaded, config) - except Exception as exc: # noqa: BLE001 - logger.exception("Forced aligner failed to load") - raise ForcedAlignerLoadError(str(exc)) from exc - try: - return await asyncio.to_thread(_align_sync, audio, text, sample_rate, config, language) - except Exception: # noqa: BLE001 - logger.exception("Forced aligner failed for text=%r", text) - return None - - -def _align_sync( - audio: bytes, - text: str, - sample_rate: int, - config: ForcedAlignerConfig, - language: str | None = None, -) -> list[WordTimestamp]: - # Precondition: align() has already run _ensure_loaded(config) under - # _encode_lock and owns the load-failure -> ForcedAlignerLoadError mapping. - # Don't reload here, so a load error can't be swallowed as a per-request None. - audio_arr = _pcm_bytes_to_float32(audio) - if audio_arr.size == 0: - return [] - audio_duration_ms = (audio_arr.size / sample_rate) * 1000.0 - - # Segment once and reuse for both the prompt and the decode: the word - # units MUST match between the two or the markers drift out of sync. - words = _processor.segment_words(text, language) - prompt = _processor.build_prompt(words) - request = { - "prompt": prompt, - "multi_modal_data": {"audio": (audio_arr, sample_rate)}, + engine_extras: dict[str, Any] = { + "model": fa.model, + "runner": fa.runner or "pooling", + # Worker-side decoder hook: pooler logits -> word-timestamps tensor. + "pooling_output_decoder": POOLING_OUTPUT_DECODER_PATH, } - - from vllm.pooling_params import PoolingParams - - outputs = _llm.encode( # type: ignore[union-attr] - [request], - pooling_params=PoolingParams(), - pooling_task=config.pooling_task or "token_classify", - use_tqdm=False, + if fa.architecture: + engine_extras["hf_overrides"] = {"architectures": [fa.architecture]} + if fa.trust_remote_code is not None: + engine_extras["trust_remote_code"] = fa.trust_remote_code + if fa.dtype is not None: + engine_extras["dtype"] = fa.dtype + aligner_ds = StageDeployConfig( + stage_id=new_id, + devices=cli_overrides.get("forced_aligner_device"), + gpu_memory_utilization=fa.gpu_memory_utilization, + max_model_len=fa.max_model_len, + engine_extras=engine_extras, + default_pooling_params={"task": fa.pooling_task or "token_classify"}, + # Pooling is one-shot -> sync AR scheduler (the async one assumes + # token-by-token generation). + async_scheduling=False, ) - if not outputs: - return [] - - result = outputs[0] - logits = result.outputs.data # [n_token, classify_num] - prompt_token_ids = list(result.prompt_token_ids) - timestamp_positions = [i for i, tid in enumerate(prompt_token_ids) if tid == _timestamp_token_id] - if not timestamp_positions: - logger.warning( - "No <|timestamp|> tokens found in prompt for text=%r; aligner returned %d rows.", - text, - logits.shape[0] if hasattr(logits, "shape") else len(logits), - ) - return [] - - return _decode_timestamps( - logits=logits, - words=words, - timestamp_positions=timestamp_positions, - classify_num=_classify_num, - timestamp_segment_time_ms=_timestamp_segment_time_ms, - audio_duration_ms=audio_duration_ms, - ) - - -def _ensure_loaded(config: ForcedAlignerConfig) -> None: - """Lazy-load the singleton ``vllm.LLM`` under lock; idempotent.""" - global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config - - if _llm is not None: - if _loaded_config is not None and _loaded_config.model != config.model: - # Multiple configs from different requests — refuse rather - # than swap models silently. A server restart is required. - raise RuntimeError( - f"Forced aligner already loaded with config={_loaded_config!r}; " - f"cannot serve a request that asks for {config!r}. " - "Restart the server to change the aligner config." - ) - return - - with _lock: - if _llm is not None: - return # raced; another caller did the load - - # Lazy import: vllm pulls torch + CUDA, which we want to avoid at - # module import time. - from vllm import LLM - - logger.info( - "Loading forced aligner %s (gpu_memory_utilization=%s)", - config.model, - config.gpu_memory_utilization if config.gpu_memory_utilization is not None else "default", - ) - llm_kwargs: dict[str, Any] = {"model": config.model} - if config.runner is not None: - llm_kwargs["runner"] = config.runner - if config.architecture is not None: - llm_kwargs["hf_overrides"] = {"architectures": [config.architecture]} - if config.gpu_memory_utilization is not None: - llm_kwargs["gpu_memory_utilization"] = config.gpu_memory_utilization - if config.trust_remote_code is not None: - llm_kwargs["trust_remote_code"] = config.trust_remote_code - if config.dtype is not None: - llm_kwargs["dtype"] = config.dtype - if config.max_model_len is not None: - llm_kwargs["max_model_len"] = config.max_model_len - llm_kwargs.update(config.extra_llm_kwargs) - - # Force the aligner engine in-process: in multiprocessing mode the - # token_classify ``pooling_output`` fails to round-trip over the - # engine-core IPC on current vLLM. Hard-set (not setdefault) so a - # globally exported ``=1`` can't push it into a subprocess; save/restore - # to avoid leaking the override to anything built later in this process. - prev_mp = os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING") - os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" - try: - llm = LLM(**llm_kwargs) - finally: - if prev_mp is None: - os.environ.pop("VLLM_ENABLE_V1_MULTIPROCESSING", None) - else: - os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = prev_mp - - thinker_config = getattr(llm.llm_engine.model_config.hf_config, "thinker_config", None) - if thinker_config is None or not hasattr(thinker_config, "classify_num"): - raise RuntimeError( - "Loaded aligner has no thinker_config.classify_num; " - "expected a Qwen3ASRForcedAlignerForTokenClassification checkpoint." - ) - timestamp_segment_time_ms = getattr(llm.llm_engine.model_config.hf_config, "timestamp_segment_time", None) - if timestamp_segment_time_ms is None: - raise RuntimeError( - "Loaded aligner has no timestamp_segment_time; expected a Qwen3ASR forced aligner checkpoint." - ) - - tokenizer = llm.get_tokenizer() - timestamp_token_id = _processor.resolve_timestamp_token_id(tokenizer) - - # Publish in this order so a concurrent reader either sees - # _llm == None (will block on the lock) or sees a fully - # initialized aligner. - _classify_num = int(thinker_config.classify_num) - _timestamp_token_id = timestamp_token_id - _timestamp_segment_time_ms = float(timestamp_segment_time_ms) - _loaded_config = config - _llm = llm - - logger.info( - "Forced aligner ready: timestamp_token_id=%d, classify_num=%d, timestamp_segment_time_ms=%.1f", - timestamp_token_id, - _classify_num, - _timestamp_segment_time_ms, - ) - - -# --- pure helpers (testable without a GPU / vllm) --- -# -# Word segmentation, prompt building, timestamp repair and marker-token -# lookup are Qwen-specific and live in -# :mod:`vllm_omni.utils.qwen3_force_align_processor`. This module owns only -# the generic vLLM orchestration. - - -def _pcm_bytes_to_float32(audio: bytes) -> np.ndarray: - """Decode signed-int16 mono PCM bytes into a [-1, 1] float32 array.""" - if not audio: - return np.zeros(0, dtype=np.float32) - if len(audio) % 2 != 0: - # Drop a trailing odd byte rather than raise; keeps streaming - # robust against off-by-one chunk boundaries. - audio = audio[:-1] - pcm = np.frombuffer(audio, dtype=np.int16) - return (pcm.astype(np.float32) / 32768.0).copy() + deploy = replace(deploy, stages=list(deploy.stages) + [aligner_ds]) + logger.info("[stage_init] Injected forced-aligner stage %d (model=%s)", new_id, fa.model) + return extended, deploy def _decode_timestamps( @@ -382,18 +221,40 @@ def _decode_timestamps( return out -# Test hooks --------------------------------------------------------------- -# Tests need a way to reset module state without restarting Python. Not -# part of the public API; do not call in production code. +def extract_word_timestamps(res: Any, text: str, language: str | None = None) -> list[dict] | None: + """Read a forced-aligner stage's engine output into per-word timestamps. + Returns ``[{word, start_ms, end_ms}, ...]`` from the stage's pooling output + (``res.outputs.data`` is an int32 ``[n_words, 2]`` tensor). Words come from + re-segmenting ``text``; any mismatch or non-aligner input returns ``None``. + """ + if not text: + return None + try: + outs = getattr(res, "outputs", None) + if outs is None: + return None + # The aligner stage's res.outputs is a single PoolingOutput (not a list). + # Its .data is a [n_words, 2] int32 tensor of [start_ms, end_ms] per word. + out0 = outs[0] if isinstance(outs, (list, tuple)) else outs + ms = getattr(out0, "data", None) + if ms is None or not hasattr(ms, "tolist"): + return None + ms_list = ms.tolist() + if not ms_list or not isinstance(ms_list[0], (list, tuple)): + return None -def _reset_for_tests() -> None: - """Drop the cached aligner state so the next call reloads.""" - global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config - with _lock: - _llm = None - _classify_num = None - _timestamp_token_id = None - _timestamp_segment_time_ms = None - _loaded_config = None - _processor._reset_for_tests() + words = _processor.segment_words(text, language) + if len(words) != len(ms_list): + # Segmentation drifted from the engine side; pairing positionally would mislabel every word. + logger.warning( + "word/timestamp count mismatch (%d words vs %d intervals); returning no timestamps", + len(words), + len(ms_list), + ) + return None + + return [{"word": w, "start_ms": int(p[0]), "end_ms": int(p[1])} for w, p in zip(words, ms_list)] + except Exception: + logger.debug("failed to extract word timestamps from aligner output", exc_info=True) + return None diff --git a/vllm_omni/utils/qwen3_force_align_processor.py b/vllm_omni/utils/qwen3_force_align_processor.py index 66a084b44..e70d270db 100644 --- a/vllm_omni/utils/qwen3_force_align_processor.py +++ b/vllm_omni/utils/qwen3_force_align_processor.py @@ -12,15 +12,15 @@ """Qwen3 forced-aligner text/timestamp processor. This is the model-specific half of upstream's ``Qwen3ForceAlignProcessor``: -it turns text into the aligner's word units and prompt, repairs the predicted -timestamp bins, and resolves the marker token id. The generic "run a vLLM -pooling model" half lives in :mod:`vllm_omni.utils.forced_aligner`, which -calls into this module. +it turns text into the aligner's word units and prompt, and repairs the +predicted timestamp bins. It feeds the forced-aligner pipeline stage (input +glue in :mod:`vllm_omni.model_executor.stage_input_processors.forced_aligner`, +timestamp decoding in :mod:`vllm_omni.utils.forced_aligner`). Keeping the Qwen-specific pieces here marks the seam for the model-agnostic aligner the issue asks for: a different aligner family would supply its own processor exposing the same small surface — :func:`segment_words`, -:func:`build_prompt`, :func:`fix_timestamp`, :func:`resolve_timestamp_token_id`. +:func:`build_prompt`, :func:`fix_timestamp`. Word segmentation prefers qwen_asr's official ``Qwen3ForceAlignProcessor`` when installed (full multilingual fidelity, incl. Japanese/Korean) and @@ -260,26 +260,3 @@ def fix_timestamp(values: list[int]) -> list[int]: i = j return [int(r) for r in result] - - -def resolve_timestamp_token_id(tokenizer: Any, timestamp_token: str = TIMESTAMP_TOKEN) -> int: - """Look up the integer id of the timestamp special token.""" - convert = getattr(tokenizer, "convert_tokens_to_ids", None) - if not callable(convert): - raise RuntimeError("Aligner tokenizer has no convert_tokens_to_ids method.") - tid = convert(timestamp_token) - if isinstance(tid, list): - tid = tid[0] if tid else None - if tid is None or (isinstance(tid, int) and tid < 0): - raise RuntimeError( - f"Aligner tokenizer does not recognise {timestamp_token!r} (got id={tid}). " - "Check the model card; the marker token may use a different name." - ) - return int(tid) - - -def _reset_for_tests() -> None: - """Drop the cached official-processor probe so the next call re-checks.""" - global _official_processor, _official_processor_unavailable - _official_processor = None - _official_processor_unavailable = False diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b68c722ec..0f65489b7 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1360,7 +1360,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): hidden_states = model_output aux_hidden_states = None - hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) + # A pooling stage skips mm extraction + the omni prefix cache (both assume an AR + # [n_tok, hidden] output) and exits through the shared _pool() branch below. + if self.is_pooling_model: + multimodal_outputs = None + else: + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) hidden_states_cpu = None # Async-write pipeline (replaces the per-step blocking @@ -1369,7 +1374,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # the actual CPU scatter into ``hidden_states_cache`` / # ``mm_outputs_cache`` happens in ``drain_ready_async_writes`` # at the top of subsequent execute_model() calls. - if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: + if not self.is_pooling_model and self.omni_prefix_cache is not None and get_pp_group().is_last_rank: hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None # Some models (e.g. qwen3-tts-talker) opt out of full-hidden-state # prefix caching but the downstream pooler payload path still From 5be7541d01ccc4f34f8d2049cc9780ba67a7f26f Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Wed, 12 Aug 2026 16:25:48 +0800 Subject: [PATCH 07/41] [Model] Support NVIDIA-NemotronLabs-VoiceChat-11B offline speech-to-speech (#5842) Signed-off-by: Yuekai Zhang Co-authored-by: Claude Fable 5 --- .../nemotron_voicechat/end2end.py | 229 +++ mkdocs.yml | 1 + pyproject.toml | 13 + recipes/NVIDIA/NemotronLabs-VoiceChat.md | 113 ++ recipes/README.md | 1 + tests/config/test_omni_config.py | 1 + .../sched/test_omni_scheduling_coordinator.py | 3 + .../test_nemotron_voicechat_registration.py | 129 ++ .../models/test_nemotron_voicechat_shapes.py | 227 +++ .../test_nemotron_voicechat.py | 422 +++++ vllm_omni/config/omni_config.py | 3 + vllm_omni/config/pipeline_registry.py | 8 + vllm_omni/config/stage_config.py | 3 + .../core/sched/omni_scheduling_coordinator.py | 4 + vllm_omni/deploy/nemotron_labs_voicechat.yaml | 81 + .../nemotron_labs_voicechat_streaming.yaml | 84 + vllm_omni/engine/arg_utils.py | 22 + .../models/nemotron_voicechat/__init__.py | 19 + .../configuration_nemotron_voicechat.py | 211 +++ .../nemo_vendored/__init__.py | 34 + .../nemo_vendored/asr/__init__.py | 25 + .../nemo_vendored/asr/audio_preprocessing.py | 259 +++ .../nemo_vendored/asr/causal_convs.py | 154 ++ .../nemo_vendored/asr/conformer_encoder.py | 850 +++++++++ .../nemo_vendored/asr/conformer_modules.py | 372 ++++ .../nemo_vendored/asr/features.py | 633 +++++++ .../nemo_vendored/asr/multi_head_attention.py | 1133 ++++++++++++ .../nemo_vendored/asr/regularization_utils.py | 67 + .../nemo_vendored/asr/subsampling.py | 628 +++++++ .../nemo_vendored/compat.py | 644 +++++++ .../nemo_vendored/duplex_ear_tts.py | 1038 +++++++++++ .../nemo_vendored/ear_tts_commons.py | 260 +++ .../nemo_vendored/ear_tts_model.py | 1620 +++++++++++++++++ .../nemo_vendored/ear_tts_vae_codec.py | 874 +++++++++ .../nemo_vendored/fusion.py | 403 ++++ .../nemo_vendored/perception.py | 135 ++ .../nemo_vendored/tokenizer.py | 501 +++++ .../nemotron_voicechat_code2wav.py | 249 +++ .../nemotron_voicechat_talker.py | 490 +++++ .../nemotron_voicechat_thinker.py | 596 ++++++ .../models/nemotron_voicechat/pipeline.py | 85 + .../models/nemotron_voicechat/runtime_info.py | 30 + vllm_omni/model_executor/models/registry.py | 17 + .../nemotron_voicechat.py | 355 ++++ vllm_omni/worker/gpu_ar_model_runner.py | 5 + .../worker/gpu_generation_model_runner.py | 3 + 46 files changed, 13034 insertions(+) create mode 100644 examples/offline_inference/nemotron_voicechat/end2end.py create mode 100644 recipes/NVIDIA/NemotronLabs-VoiceChat.md create mode 100644 tests/model_executor/models/test_nemotron_voicechat_registration.py create mode 100644 tests/model_executor/models/test_nemotron_voicechat_shapes.py create mode 100644 tests/model_executor/stage_input_processors/test_nemotron_voicechat.py create mode 100644 vllm_omni/deploy/nemotron_labs_voicechat.yaml create mode 100644 vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/__init__.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py create mode 100644 vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py create mode 100644 vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py diff --git a/examples/offline_inference/nemotron_voicechat/end2end.py b/examples/offline_inference/nemotron_voicechat/end2end.py new file mode 100644 index 000000000..b020e5bad --- /dev/null +++ b/examples/offline_inference/nemotron_voicechat/end2end.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline speech-to-speech with NVIDIA-NemotronLabs-VoiceChat-11B. + +Feeds a 16 kHz user utterance (plus a spoken-style system prompt) through the +3-stage nemotron_voicechat pipeline and writes the agent's reply as text and a +22.05 kHz WAV. + +Usage: + python end2end.py \ + --checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \ + --wav /path/to/user_question_16k.wav \ + --output-dir results/nemotron_voicechat + +The system prompt defaults to the NeMo reference default. The text tokenizer +resolves from the ``nvidia/NVIDIA-Nemotron-Nano-9B-v2`` HF id (or the +``NEMOTRON_VOICECHAT_LLM_PATH`` env var for air-gapped runs). + +Sizing the input WAV: the timeline is frame-locked, so the reply budget IS the +input duration — the acoustic channel trails the text channel, and a reply +that has not finished when the frames run out is truncated SILENTLY (no error; +ASR of the WAV will just disagree with the emitted text). Leave generous +trailing silence after the question (a question ending at ~4.5 s truncated in +an 8 s WAV but completed cleanly in a 16 s one). The script warns when the +text channel is still emitting non-PAD tokens near the last frame. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import numpy as np +import soundfile as sf + +# Warn if the text channel is still speaking within this many final frames +# (12 frames = ~1 s) — the frame-locked reply likely got cut off. +TRUNCATION_GUARD_FRAMES = 12 + +# The NeMo reference offline script's default system prompt. +DEFAULT_SYSTEM_PROMPT = ( + "You are an AI voice assistant developed by NVIDIA. " + "Your name is NVIDIA Voice Chat. " + "Answer in a spoken, conversational style rather than a written one. " + "Do not repeat the same sentence over and over again. " + "Start the conversation by greeting the user." +) + + +def load_wav_16k_mono(path: str) -> np.ndarray: + wav, sr = sf.read(path, dtype="float32") + if wav.ndim > 1: + wav = wav.mean(axis=1) + if sr != 16000: + from vllm.multimodal.audio import resample_audio_scipy + + wav = resample_audio_scipy(wav, orig_sr=sr, target_sr=16000) + return wav.astype(np.float32) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", required=True, help="NVIDIA-NemotronLabs-VoiceChat-11B directory") + parser.add_argument("--wav", required=True, help="user speech (any sr; resampled to 16 kHz mono)") + parser.add_argument("--output-dir", default="results/nemotron_voicechat") + parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT) + parser.add_argument( + "--stage-configs-path", + default=None, + help=( + "Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml " + "for the async-chunk streaming pipeline (stages overlap; audio starts decoding " + "before the thinker finishes). Default: the bundled non-streaming yaml (bit-exact " + "parity path)." + ), + ) + args = parser.parse_args() + + if not Path(args.wav).is_file(): + raise FileNotFoundError(f"Input WAV not found: {args.wav}") + ckpt = Path(args.checkpoint) + raw_cfg = json.loads((ckpt / "config.json").read_text()) + stt_cfg = raw_cfg["model"]["stt"]["model"] + + from transformers import AutoTokenizer + from vllm.sampling_params import SamplingParams + + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import ( + compute_acoustic_frame_count, + ) + + tok_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get( + "pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + ) + tokenizer = AutoTokenizer.from_pretrained(tok_ref, trust_remote_code=False) + bos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("bos_token", "")) + eos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("eos_token", "")) + pad_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("pad_token", "")) + # NeMo encode_system_prompt: [bos] + tokens + [eos]. + prompt_ids = [bos_id] + tokenizer.encode(args.system_prompt, add_special_tokens=False) + [eos_id] + + wav = load_wav_16k_mono(args.wav) + n_frames = compute_acoustic_frame_count(stt_cfg, int(wav.shape[0])) + print(f"input: {wav.shape[0]} samples @16kHz -> {n_frames} acoustic frames; prompt {len(prompt_ids)} tokens") + + inputs = { + # vLLM prompt = system prompt + one placeholder position (acoustic + # frame 0); the thinker replaces every position with fused embeddings. + "prompt_token_ids": prompt_ids + [pad_id], + "additional_information": { + # Plain float list: raw ndarrays do not survive the engine-core + # message serialization. + "nvc_audio": wav.astype(np.float32).tolist(), + "nvc_sr": 16000, + "nvc_prompt_token_ids": prompt_ids, + "nvc_expected_frames": n_frames, + }, + } + sampling_params_list = [ + # thinker: greedy, frame-locked generation length. NO min_tokens: the + # tokenizer's eos_token is — the frame-locked PAD/silence + # token — and vLLM's min_tokens masks "EOS" to -inf, which would forbid + # silence and force the agent to babble through the whole utterance. + SamplingParams(temperature=0.0, max_tokens=n_frames, ignore_eos=True, seed=0), + # talker: placeholder loop, stops on the stage's stop token. 16383: + # the 1-token placeholder prompt + max_tokens must fit the stage's + # max_model_len (16384). + SamplingParams(temperature=0.0, max_tokens=16383, stop_token_ids=[1], detokenize=False, seed=0), + # code2wav: single decode step. + SamplingParams(temperature=0.0, max_tokens=1, detokenize=False, seed=0), + ] + + omni_kwargs: dict = {} + if args.stage_configs_path: + omni_kwargs["stage_configs_path"] = args.stage_configs_path + omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs) + outputs = omni.generate(inputs, sampling_params_list) + + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + stem = Path(args.wav).stem + + agent_text = None + audio = None + sr = 22050 + for stage_outputs in outputs: + # One entry per final stage (text from the thinker, audio from code2wav); + # mirrors the qwen2_5_omni offline example's collection loop. + output = getattr(stage_outputs, "request_output", stage_outputs) + final_type = getattr(stage_outputs, "final_output_type", None) + completions = getattr(output, "outputs", None) or [] + if not completions: + continue + completion = completions[0] + if final_type in (None, "text"): + text = getattr(completion, "text", None) + if text: + agent_text = text + token_ids = getattr(completion, "token_ids", None) + if token_ids: + (out_dir / f"{stem}_text_tokens.json").write_text(json.dumps(list(map(int, token_ids)))) + # Frame lock: the reply budget IS the input length. If the text + # channel is still mid-sentence when the frames run out, the + # spoken answer is silently truncated — warn instead. + tail = list(map(int, token_ids))[-TRUNCATION_GUARD_FRAMES:] + if any(t != pad_id for t in tail): + print( + f"WARNING: the text channel emits non-PAD tokens within the last " + f"{TRUNCATION_GUARD_FRAMES} frames ({TRUNCATION_GUARD_FRAMES * 80} ms) — the reply " + "likely ran out of acoustic frames and is cut off mid-sentence. " + "Pad the input WAV with more trailing silence and rerun." + ) + # Best-effort: dump the function-channel debug timeline if the + # engine surfaced the request info on this output. + info = getattr(output, "additional_information", None) + func_timeline = info.get("nvc_function_tokens") if isinstance(info, dict) else None + if func_timeline is not None: + values = func_timeline.tolist() if hasattr(func_timeline, "tolist") else list(func_timeline) + (out_dir / f"{stem}_function_tokens.json").write_text(json.dumps([int(v) for v in values])) + print(f"function-channel timeline: {len(values)} entries") + if final_type in (None, "audio"): + # multimodal_output is a MultimodalPayload (Mapping) — duck-type it. + mm = getattr(completion, "multimodal_output", None) + get = getattr(mm, "get", None) + if not callable(get): + continue + candidate = None + for key in ("audio", "model_outputs", "wav", "waveform"): + value = get(key) + if value is not None: + candidate = value + break + if isinstance(candidate, list): + candidate = candidate[0] if candidate else None + if candidate is not None: + arr = np.asarray( + candidate.float().cpu() if hasattr(candidate, "cpu") else candidate, dtype=np.float32 + ).reshape(-1) + if arr.size > 0: + audio = arr + sr_val = get("sr") + if sr_val is not None: + sr_item = sr_val[0] if isinstance(sr_val, list) else sr_val + sr = int(sr_item.item() if hasattr(sr_item, "item") else sr_item) + + if agent_text is not None: + # The frame-locked text channel emits PAD tokens on silent frames; + # strip them (NeMo's post-inference detokenizer does the same). + pad_token = stt_cfg.get("pad_token", "") + agent_text = " ".join(agent_text.replace(pad_token, " ").split()) + text_path = out_dir / f"{stem}_output.txt" + text_path.write_text(agent_text + "\n") + print(f"agent text -> {text_path}\n{agent_text}") + else: + print("WARNING: no agent text in outputs") + + if audio is not None and audio.size > 0: + wav_path = out_dir / f"{stem}_output.wav" + sf.write(wav_path, audio, sr) + print(f"agent audio -> {wav_path} ({audio.size / sr:.2f}s @ {sr} Hz)") + else: + raise RuntimeError("No audio produced by the code2wav stage.") + + +if __name__ == "__main__": + main() diff --git a/mkdocs.yml b/mkdocs.yml index fd97b9b37..242eb33ae 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -111,6 +111,7 @@ plugins: - "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building - "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building - "vllm_omni.experimental" # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building + - "vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored" # vendored NeMo code; upstream docstrings trip griffe strict warnings - mkdocstrings: handlers: python: diff --git a/pyproject.toml b/pyproject.toml index f9fbdfb2a..0e91bfae6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -206,6 +206,12 @@ ignore = [ # These files keep upstream conventions (star imports, long lines, math-symbol # argument names) so they stay easy to diff against upstream releases. "vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"] +# NemotronVoiceChat inference modules vendored from NVIDIA-NeMo/Speech (branch +# nemotron-labs-voicechat). Numerics-bearing code is kept byte-close to +# upstream to stay diffable and parity-safe, so upstream style (long lines, +# legacy typing, whitespace, unused-import shims) is exempted. TID251 and the +# F-correctness rules still apply. +"vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"] "vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines [tool.ruff.lint.flake8-tidy-imports.banned-api] @@ -314,6 +320,10 @@ extend-ignore-identifiers-re = [ ".*[Nn][Dd]3.*", ".*updator.*", "byt5.*", + # NemotronVoiceChat vendored triton kernel: dout/DOUT abbreviate d_out + # (output feature dim), not "doubt". + ".*dout.*", + ".*DOUT.*", ] extend-ignore-re = [ "\\.\\.\\.n,d->\\.\\.\\.nd", @@ -334,6 +344,9 @@ byt5 = "byt5" # "seeked" is the standard HTML media element DOM event name (fires when a # seek completes), used by the word-timestamps Gradio demo. Not a typo. seeked = "seeked" +# Ryan Leary: author credited in the NemotronVoiceChat vendored NeMo +# features.py copyright header. A surname, not a typo of "leery". +Leary = "Leary" # IndexTTS 2.5 language codes and speech-emotion-recognition event token. fo = "fo" ba = "ba" diff --git a/recipes/NVIDIA/NemotronLabs-VoiceChat.md b/recipes/NVIDIA/NemotronLabs-VoiceChat.md new file mode 100644 index 000000000..8faa268d2 --- /dev/null +++ b/recipes/NVIDIA/NemotronLabs-VoiceChat.md @@ -0,0 +1,113 @@ +# NemotronLabs VoiceChat 11B + +> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply +> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline. + +## Summary + +- Vendor: NVIDIA +- Model: [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B) +- Task: Speech-to-speech voice chat. The model runs a frame-locked 12.5 Hz + timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token + per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into + 31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio. +- Mode: Offline single-turn inference (batch=1). Online/duplex serving, + function calling, and batch>1 are documented follow-ups. +- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang) + +## When to use this recipe + +Use this recipe to run a single-turn voice-chat exchange with +`NVIDIA-NemotronLabs-VoiceChat-11B` on one GPU: you provide a user utterance as +a WAV file (any sample rate; it is resampled to 16 kHz mono) plus an optional +spoken-style system prompt, and get back the agent's reply as text and a +22.05 kHz WAV. The integration is NeMo-free at runtime — the perception +Conformer, EAR-TTS talker, and RVQ-VAE codec are vendored, dependency-stripped +NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed. + +## References + +- Offline example: + [`examples/offline_inference/nemotron_voicechat/end2end.py`](../../examples/offline_inference/nemotron_voicechat/end2end.py) +- Model modules (thinker / talker / code2wav / vendored NeMo): + [`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/) +- Staged pipeline config: + [`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml) +- Upstream model card: + [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B) +- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat` + +## Pipeline + +| stage | arch | dtype | role | +|---|---|---|---| +| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | fp32 | WAV + system prompt -> frame-locked text-token timeline (+ function channel) | +| 1 talker | `NemotronVoiceChatTalker` (LLM_AR) | fp32 | text timeline -> 31-quantizer RVQ code stacks (one per 80 ms frame) | +| 2 code2wav | `NemotronVoiceChatCode2Wav` (LLM_GENERATION) | fp32 | RVQ-VAE decode -> 22.05 kHz PCM | + +Stages 0/1 default to fp32 for exact parity with the NeMo reference +implementation (greedy decoding matches it token for token on the acceptance +fixture). The deploy yaml documents a ~2x-faster bf16 thinker option whose +output stayed within one word of the reference in testing. + +## Hardware Support + +## GPU + +### 1x H100 80GB + +#### Environment + +- OS: Linux +- Python: 3.12 +- vLLM version: 0.26.0 +- vLLM-Omni version or commit: this PR / current `main` + +#### Command + +```bash +# Tokenizer: the checkpoint ships no HF tokenizer; it resolves from the +# nvidia/NVIDIA-Nemotron-Nano-9B-v2 HF id automatically. For air-gapped runs, +# point NEMOTRON_VOICECHAT_LLM_PATH at a local snapshot of that repo instead. +python examples/offline_inference/nemotron_voicechat/end2end.py \ + --checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \ + --wav /path/to/user_question.wav \ + --output-dir results/nemotron_voicechat +``` + +#### Verification + +```bash +ls results/nemotron_voicechat +# _output.txt the agent reply as text +# _output.wav the agent reply as 22.05 kHz audio +# _text_tokens.json the frame-locked text-token timeline +``` + +The reply text should read as a coherent spoken-style answer to the question in +the input WAV, and the WAV should transcribe to (approximately) the same text +with any ASR model. + +#### Notes + +- Memory usage: the shipped yaml runs all three stages on one GPU + (`gpu_memory_utilization` 0.62 / 0.12 / 0.06); peak usage is dominated by the + fp32 thinker. The fp32 default has a hard floor of roughly 43 GB of thinker + weights alone (9B backbone + 587M `embed_tokens` + 587M `function_head` + + 0.6B Conformer), so 48 GB cards cannot run it — use the bf16 thinker option + documented in the deploy yaml on anything smaller than an 80 GB part. +- Input sizing: the timeline is frame-locked, so the reply budget IS the input + duration. The acoustic channel trails the text channel; if the WAV does not + carry enough trailing silence for the reply to finish, the spoken answer is + truncated silently. Leave generous trailing silence (a question ending at + ~4.5 s truncated in an 8 s WAV but completed cleanly in 16 s); the offline + example warns when the text channel is still speaking near the last frame. +- Key flags: sampling is greedy end to end. The thinker is frame-locked — + `max_tokens` equals the acoustic frame count with `ignore_eos=True`. Do NOT + set `min_tokens` on the thinker: the tokenizer's EOS token is also the + frame-locked PAD/silence token, so masking it forces the model to babble + instead of pausing. +- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder + token, and the stage context is 16384). +- Known limitations: batch=1 only; offline single turn (no online/duplex + serving yet); the function-call channel is decoded but not yet acted on. diff --git a/recipes/README.md b/recipes/README.md index ca287bf70..e010ade67 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -66,6 +66,7 @@ recipes/ | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB | | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB | | [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB | +| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB | | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 | | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 | | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB | diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 5676a17a6..1fb32239d 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -513,6 +513,7 @@ def test_sub_config_fields_match_structured_scopes(): "enable_prefix_caching", "disable_hybrid_kv_cache_manager", "mm_processor_cache_gb", + "mamba_ssm_cache_dtype", } assert {f.name for f in fields(OmniStageSchedulerConfig)} == { "max_num_seqs", diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index a93929306..444d21c34 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -120,6 +120,9 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase): ("MiMoAudioModel", "code2wav"), ("Qwen3TTSCode2Wav", "code2wav"), ("MiniCPMO45Code2Wav", "code2wav"), + # nemotron_voicechat: sync-path RVQ-VAE decode waits for the + # talker's full [frames, 31] code-stack payload. + ("NemotronVoiceChatCode2Wav", "code2wav"), ("CosyVoice3Model", "cosyvoice3_code2wav"), # audex: sync-path TTS decoder (async_chunk: false deploys) and # the TTA XCodec1 stage (always sync full-payload). diff --git a/tests/model_executor/models/test_nemotron_voicechat_registration.py b/tests/model_executor/models/test_nemotron_voicechat_registration.py new file mode 100644 index 000000000..09d85b271 --- /dev/null +++ b/tests/model_executor/models/test_nemotron_voicechat_registration.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights).""" + +import subprocess +from pathlib import Path + +import pytest + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.config.stage_config import StageExecutionType +from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import ( + NemotronVoiceChatConfig, +) +from vllm_omni.model_executor.models.registry import _OMNI_MODELS + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_ARCHS = ( + "NemotronVoiceChatThinkerForConditionalGeneration", + "NemotronVoiceChatTalkerForConditionalGeneration", + "NemotronVoiceChatCode2Wav", +) + + +def _minimal_nemo_dict() -> dict: + return { + "data": {"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050}, + "model": { + "inference_speaker_name": "Aria", + "stt": { + "model": { + "perception": {"encoder": {}, "preprocessor": {}, "modality_adapter": {}}, + "pretrained_llm": "nvidia/NVIDIA-Nemotron-Nano-9B-v2", + "use_function_head": True, + } + }, + "speech_generation": { + "model": { + "tts_config": {"backbone_config": {"hidden_size": 1152, "sliding_window": 7500}}, + "codec_config": {"num_quantizers": 31, "codebook_size": 1024}, + } + }, + }, + } + + +def test_pipeline_registered_with_three_stages() -> None: + pipeline = OMNI_PIPELINES["nemotron_voicechat"] + assert OMNI_PIPELINES["nemotron_labs_voicechat"] is pipeline + assert len(pipeline.stages) == 3 + thinker, talker, code2wav = pipeline.stages + assert thinker.execution_type == StageExecutionType.LLM_AR + assert thinker.final_output and thinker.final_output_type == "text" + assert thinker.owns_tokenizer + assert talker.execution_type == StageExecutionType.LLM_AR + assert talker.model_arch == "NemotronVoiceChatTalkerForConditionalGeneration" + assert talker.input_sources == (0,) + assert code2wav.execution_type == StageExecutionType.LLM_GENERATION + assert code2wav.final_output and code2wav.final_output_type == "audio" + assert code2wav.input_sources == (1,) + assert pipeline.validate() == [] + + +def test_architectures_registered() -> None: + for arch in _ARCHS: + assert arch in _OMNI_MODELS, f"{arch} missing from _OMNI_MODELS" + + +def test_full_payload_whitelists_cover_code2wav_only() -> None: + from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES + + assert ("NemotronVoiceChatCode2Wav", "code2wav") in _FULL_PAYLOAD_INPUT_STAGES + # The thinker->talker hop is token-path only: the talker must NOT wait on a + # connector payload or stage 1 hangs. + assert ("NemotronVoiceChatTalkerForConditionalGeneration", "talker") not in _FULL_PAYLOAD_INPUT_STAGES + + +def test_config_parses_nemo_layout_without_hub_access(monkeypatch) -> None: + # Provide the thinker text config explicitly so the test never touches HF. + cfg = NemotronVoiceChatConfig( + thinker_text_config={"model_type": "llama", "hidden_size": 4480, "num_hidden_layers": 56}, + **_minimal_nemo_dict(), + ) + assert cfg.get_text_config().hidden_size == 4480 + assert cfg.talker_config.hidden_size == 1152 + assert cfg.talker_config.sliding_window == 7500 + assert cfg.code2wav_config.num_quantizers == 31 + assert cfg.inference_speaker_name == "Aria" + + +@pytest.mark.parametrize( + "mutation", + ["drop_stt", "drop_speech_generation", "drop_perception", "drop_codec"], +) +def test_config_rejects_non_voicechat_checkpoints(mutation: str) -> None: + raw = _minimal_nemo_dict() + if mutation == "drop_stt": + del raw["model"]["stt"] + elif mutation == "drop_speech_generation": + del raw["model"]["speech_generation"] + elif mutation == "drop_perception": + del raw["model"]["stt"]["model"]["perception"] + else: + del raw["model"]["speech_generation"]["model"]["codec_config"] + with pytest.raises(ValueError, match="NemotronVoiceChat|missing"): + NemotronVoiceChatConfig( + thinker_text_config={"model_type": "llama", "hidden_size": 4480}, + **raw, + ) + + +def test_no_nemo_imports_in_nemotron_voicechat_tree() -> None: + """Dependency guard: the vendored tree must not import the nemo package. + + Scoped to the nemotron_voicechat model tree (other models may have their + own optional nemo integrations). + """ + repo_root = Path(__file__).resolve().parents[3] + targets = [ + str(repo_root / "vllm_omni/model_executor/models/nemotron_voicechat"), + str(repo_root / "vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py"), + ] + result = subprocess.run( + ["grep", "-rnE", r"^\s*(import nemo\b|from nemo(\.|\s))", *targets, "--include=*.py"], + capture_output=True, + text=True, + ) + assert result.returncode != 0, f"Found nemo imports in the nemotron_voicechat tree:\n{result.stdout}" diff --git a/tests/model_executor/models/test_nemotron_voicechat_shapes.py b/tests/model_executor/models/test_nemotron_voicechat_shapes.py new file mode 100644 index 000000000..8bd279bb8 --- /dev/null +++ b/tests/model_executor/models/test_nemotron_voicechat_shapes.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Golden shape test for the NemotronVoiceChat frame-locked timeline contract. + +Pins the exact prefill/off-by-one arithmetic against values verified with the +NeMo reference run on sample_general.wav (fp32, greedy): perception produced +196 acoustic frames for the 249,734-sample 16 kHz fixture (the NeMo dump's +250,760 was the loader's padded batch width; perception consumed the true +length); the system prompt tokenized to 56 ids ([bos] + 54 + [eos]); NeMo's +timeline T = 252 = 56 + 196; and our fp32 pipeline matched the reference text +channel 196/196 tokens. +""" + +import pytest + +from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import ( + compute_acoustic_frame_count, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Perception config subset matching the checkpoint (mel 10 ms hop, n_fft 512, +# dw-striding 8x CAUSAL subsampling, conv kernel 3). causal_downsampling is +# what the real checkpoint sets; omitting it here is exactly what let the +# non-causal padding formula pass review round 0 (helper vs itself). +_STT_CFG = { + "perception": { + "preprocessor": {"sample_rate": 16000, "window_stride": 0.01, "window_size": 0.025, "n_fft": 512}, + "encoder": { + "subsampling": "dw_striding", + "subsampling_factor": 8, + "subsampling_conv_kernel_size": 3, + "causal_downsampling": True, + }, + } +} + +# (num_16k_samples, expected 12.5 Hz frames), all verified against the REAL +# vendored AudioPerceptionModule built from the checkpoint config. The first +# row is the sample_general.wav acceptance fixture (196 frames -> the NeMo +# T=252 timeline). 71271 is the reviewer's 4.45 s repro that the non-causal +# formula undersized by one; 250760 is the NeMo loader's padded width, which +# genuinely yields 197 (not the fixture's 196). +_GOLDEN = [ + (249734, 196), + (250760, 197), + (71271, 57), + (16000, 14), + (16000 * 2, 26), + (16000 * 4, 51), + (16000 * 8, 101), +] + + +@pytest.mark.parametrize(("num_samples", "expected_frames"), _GOLDEN) +def test_acoustic_frame_count_golden(num_samples: int, expected_frames: int) -> None: + assert compute_acoustic_frame_count(_STT_CFG, num_samples) == expected_frames + + +def test_frame_count_matches_vendored_subsampling_geometry() -> None: + """Pin the helper against the real ConvSubsampling, not against itself. + + Builds the vendored subsampling module with the checkpoint geometry (both + causal and non-causal) and checks the helper's padding arithmetic against + the module's actual ``_left_padding + _right_padding`` and its + ``calc_length`` output over a spread of mel lengths. + """ + import torch + + from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.subsampling import ( + ConvSubsampling, + calc_length, + ) + + for causal in (True, False): + sub = ConvSubsampling( + subsampling="dw_striding", + subsampling_factor=8, + feat_in=128, + feat_out=8, + conv_channels=8, + is_causal=causal, + ) + expected_paddings = sub._left_padding + sub._right_padding + cfg = { + "perception": { + "preprocessor": _STT_CFG["perception"]["preprocessor"], + "encoder": { + "subsampling": "dw_striding", + "subsampling_factor": 8, + "subsampling_conv_kernel_size": 3, + "causal_downsampling": causal, + }, + } + } + for num_samples in (16000, 71271, 249734, 250760): + # Recompute the helper's mel length, then ask the module's own + # calc_length what the subsampled length should be. + mel_len = (num_samples + 512 - 512) // 160 + 1 + module_frames = int( + calc_length( + torch.tensor([mel_len]), + all_paddings=expected_paddings, + kernel_size=3, + stride=2, + ceil_mode=False, + repeat_num=3, + )[0] + ) + assert compute_acoustic_frame_count(cfg, num_samples) == module_frames, ( + f"helper diverged from vendored geometry (causal={causal}, samples={num_samples})" + ) + + +def test_timeline_max_model_len_validation() -> None: + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import ( + validate_timeline_fits_max_model_len, + ) + + # The acceptance fixture (56 + 1 + 196 = 253) fits the shipped 8192 limit. + validate_timeline_fits_max_model_len(56, 196, 8192) + # Exactly at the limit: the vLLM request occupies prompt + 1 (frame-0 + # placeholder) + frames tokens, so 57 + 8135 == 8192 is allowed... + validate_timeline_fits_max_model_len(56, 8135, 8192) + # ...and one more frame must be rejected, naming every offending length. + with pytest.raises(ValueError) as excinfo: + validate_timeline_fits_max_model_len(56, 8136, 8192) + message = str(excinfo.value) + for fragment in ( + "vllm_prefill_len=57", + "acoustic_frame_count=8136", + "vllm_total_len=8193", + "max_model_len=8192", + "logical_prompt_token_len=56", + "logical_total_len=8192", + ): + assert fragment in message + + +def _bare_thinker(): + import torch + from torch import nn + + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import ( + NemotronVoiceChatThinkerForConditionalGeneration, + ) + + thinker = NemotronVoiceChatThinkerForConditionalGeneration.__new__(NemotronVoiceChatThinkerForConditionalGeneration) + nn.Module.__init__(thinker) + thinker._use_function_head = True + thinker._dtype = torch.float32 + thinker.function_head = nn.Linear(8, 16, bias=False) + return thinker + + +def test_postprocess_appends_function_token_to_timeline(monkeypatch: pytest.MonkeyPatch) -> None: + import torch + + thinker = _bare_thinker() + hidden = torch.randn(3, 8) + expected = int(thinker.function_head(hidden[-1:]).argmax(dim=-1)) + + # Default: only the load-bearing prev-token; the cumulative debug timeline + # (a per-step host sync + growing concat) stays off. + update = thinker.postprocess(hidden) + assert int(update["nvc_prev_function_token"]) == expected + assert "nvc_function_tokens" not in update + + monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1") + # No existing timeline: starts one. + update = thinker.postprocess(hidden) + assert update["nvc_function_tokens"].tolist() == [expected] + # Existing timeline: appended, not replaced. + existing = torch.tensor([7, 8, 9], dtype=torch.long) + update = thinker.postprocess(hidden, nvc_function_tokens=existing) + assert update["nvc_function_tokens"].tolist() == [7, 8, 9, expected] + + +def test_postprocess_skips_intermediate_prefill_chunks(monkeypatch: pytest.MonkeyPatch) -> None: + import torch + + monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1") + thinker = _bare_thinker() + hidden = torch.randn(4, 8) + # Intermediate chunk: 0 computed + 4 rows < prompt_len 10 -> no token. + assert thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=0, _omni_prompt_len=10) == {} + # Final chunk reaches the last prefill position (acoustic frame 0). + update = thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=6, _omni_prompt_len=10) + assert "nvc_prev_function_token" in update + assert update["nvc_function_tokens"].numel() == 1 + + +def test_prefill_contract_arithmetic() -> None: + # vllm_prefill_len = logical_prompt_token_len + 1 (the +1 is acoustic + # frame 0); decode steps == acoustic_frame_count; NeMo timeline + # T == prompt + frames. Values from the verified sample_general.wav run. + logical_prompt_len = 56 + frames = compute_acoustic_frame_count(_STT_CFG, 249734) + vllm_prefill_len = logical_prompt_len + 1 + assert vllm_prefill_len == 57 + assert logical_prompt_len + frames == 252 # NeMo reference T + # The talker consumes the full timeline from t=1 and the producer trims + # the prompt region: rows P-1 onward of the (T-1)-row code stack. + talker_rows = (logical_prompt_len + frames) - 1 + trimmed = talker_rows - (logical_prompt_len - 1) + assert trimmed == frames + + +def test_torchaudio_mel_filterbank_matches_librosa() -> None: + """The perception frontend builds its mel filterbank with torchaudio. + + torchaudio's melscale_fbanks(norm="slaney", mel_scale="slaney") implements + the same slaney-convention math as librosa.filters.mel (the NeMo reference + dependency) up to float32 rounding. The 1e-6 tolerance is ~4x the observed + 2.6e-7 max diff; e2e this stays token- and WAV-identical to the reference + on the parity fixture (per-feature normalization absorbs it). + """ + librosa = pytest.importorskip("librosa") + torchaudio = pytest.importorskip("torchaudio") + + # The checkpoint's perception preprocessor config (16 kHz, n_fft=512, 128 mels). + ta = torchaudio.functional.melscale_fbanks( + n_freqs=257, f_min=0.0, f_max=8000.0, n_mels=128, sample_rate=16000, norm="slaney", mel_scale="slaney" + ).T.numpy() + reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney") + assert ta.shape == reference.shape + assert abs(ta - reference).max() < 1e-6 diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py new file mode 100644 index 000000000..dcdff5bbb --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import ( + _FULL_PAYLOAD_REPLACE_KEYS, + talker2code2wav_full_payload, + talker2code2wav_token_only, + thinker2talker_token_only, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_PAD = 12 + + +def _thinker_output(prompt_len: int, generated: list[int]): + return SimpleNamespace( + finished=True, + prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD], + additional_information=None, + outputs=[SimpleNamespace(token_ids=generated, cumulative_token_ids=None)], + ) + + +def test_thinker2talker_builds_pad_prompt_timeline() -> None: + generated = [7, 8, 9, 10] + inputs = thinker2talker_token_only([_thinker_output(prompt_len=5, generated=generated)]) + assert len(inputs) == 1 + prompt = inputs[0] + # 1-token placeholder prompt; the timeline rides additional_information. + assert prompt["prompt_token_ids"] == [0] + info = prompt["additional_information"] + assert info["nvc_logical_prompt_len"] == 5 + assert info["nvc_frame_count"] == 4 + assert info["nvc_text_timeline"] == [_PAD] * 5 + generated + + +def test_thinker2talker_skips_unfinished() -> None: + out = _thinker_output(3, [1]) + out.finished = False + assert thinker2talker_token_only([out]) == [] + + +def test_full_payload_trims_prompt_region_rows() -> None: + # Rows correspond to timeline steps t=1..T-1; with prompt_len P the NeMo + # trim keeps rows for t >= P, i.e. row index P-1 onward. + prompt_len = 4 + frames, q = 10, 31 + codes = torch.arange(frames * q, dtype=torch.long).reshape(frames, q) % 1024 + request = SimpleNamespace( + additional_information={ + "codes": {"audio": codes}, + "nvc_logical_prompt_len": prompt_len, + }, + additional_information_cpu=None, + ) + payload = talker2code2wav_full_payload(request=request, is_finished=True) + assert payload.codes.audio.shape == (frames - (prompt_len - 1), q) + assert torch.equal(payload.codes.audio, codes[prompt_len - 1 :]) + assert bool(payload.meta.finished) + + +def test_full_payload_without_codes_is_empty_finished() -> None: + # The sync producer must never emit a meta-only "pending" struct (it would + # be shipped as the real payload); with no codes it ships an empty finished + # payload regardless of the request's is_finished flag. + payload = talker2code2wav_full_payload(request=SimpleNamespace(), is_finished=False) + assert payload is not None + assert bool(payload.meta.finished) + assert payload.codes.audio.numel() == 0 + + +def test_full_payload_replace_semantics_declared() -> None: + # Full-sequence tensors must REPLACE in the accumulator, not concatenate. + # The AR runner flattens nested payloads before accumulation, so the + # flattened "codes.audio" key is the one that matters. + assert "codes.audio" in _FULL_PAYLOAD_REPLACE_KEYS + assert "codes" in _FULL_PAYLOAD_REPLACE_KEYS + + +def test_full_payload_missing_prompt_len_fails_explicitly() -> None: + codes = torch.zeros((5, 31), dtype=torch.long) + request = SimpleNamespace( + additional_information={"codes": {"audio": codes}}, + additional_information_cpu=None, + ) + with pytest.raises(ValueError, match="nvc_logical_prompt_len"): + talker2code2wav_full_payload(request=request, is_finished=True) + + +def test_code2wav_distinguishes_empty_payload_from_missing_info() -> None: + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import ( + NemotronVoiceChatCode2Wav, + ) + + extract = NemotronVoiceChatCode2Wav._codes_from_runtime_info + # Explicit empty payload (e.g. zero frames after the prompt trim): the + # empty tensor comes back as-is -> empty audio, never a placeholder-id + # decode. + empty = extract(None, {"codes": {"audio": torch.empty(0, dtype=torch.long)}}) + assert isinstance(empty, torch.Tensor) and empty.numel() == 0 + # Missing codes info entirely -> None (forward raises: placeholder + # input_ids are never decoded). + assert extract(None, {"meta": {}}) is None + assert extract(None, None) is None + # Normal payload passes through (nested and flattened key forms). + codes = torch.zeros((3, 31), dtype=torch.long) + assert extract(None, {"codes": {"audio": codes}}) is codes + assert extract(None, {"codes.audio": codes}) is codes + + +def test_validate_code_stack_negatives() -> None: + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import ( + validate_code_stack, + ) + + good = torch.zeros((4, 31), dtype=torch.long) + assert validate_code_stack(good, 31, 1024).shape == (4, 31) + with pytest.raises(ValueError, match="not divisible"): + validate_code_stack(torch.zeros(100, dtype=torch.long), 31, 1024) + with pytest.raises(ValueError, match="shaped"): + validate_code_stack(torch.zeros((4, 30), dtype=torch.long), 31, 1024) + bad = good.clone() + bad[1, 5] = 1024 + with pytest.raises(ValueError, match="out-of-range"): + validate_code_stack(bad, 31, 1024) + + +def test_tokenizer_resolution_failure_is_actionable(monkeypatch) -> None: + from transformers import AutoConfig + + from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import ( + NemotronVoiceChatConfig, + ) + + def _boom(*args, **kwargs): + raise OSError("offline") + + monkeypatch.setattr(AutoConfig, "from_pretrained", _boom) + with pytest.raises(RuntimeError, match="NEMOTRON_VOICECHAT_LLM_PATH"): + NemotronVoiceChatConfig( + data={"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050}, + model={ + "inference_speaker_name": "Aria", + "stt": {"model": {"perception": {"encoder": {}}, "pretrained_llm": "nvidia/nope"}}, + "speech_generation": {"model": {"tts_config": {}, "codec_config": {}}}, + }, + ) + + +def test_sanitize_tts_model_cfg_disables_pretrained_loads(monkeypatch) -> None: + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import ( + sanitize_tts_model_cfg, + ) + + original = { + "pretrained_lm_name": "nvidia/NVIDIA-Nemotron-Nano-9B-v2", + "pretrained_codec_model": "/some/external/codec.ckpt", + "tts_config": { + "pretrained_text_name": "some/hub-llm", + "backbone_type": "gemma3_text", + "cas_config": {"pretrained_tokenizer_name": "some/hub-tokenizer", "keep": 1}, + }, + } + monkeypatch.delenv("NEMOTRON_VOICECHAT_LLM_PATH", raising=False) + cfg = sanitize_tts_model_cfg(original) + # External weight-load triggers are gone; construction is config-only. + assert "pretrained_text_name" not in cfg["tts_config"] + assert "pretrained_codec_model" not in cfg + assert "pretrained_tokenizer_name" not in cfg["tts_config"]["cas_config"] + assert cfg["tts_config"]["cas_config"]["keep"] == 1 + # Tokenizer reference survives. + assert cfg["pretrained_lm_name"] == "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + # The shared stage config is never mutated (deep copy). + assert original["tts_config"]["pretrained_text_name"] == "some/hub-llm" + + # Air-gapped override reroutes the tokenizer reference. + monkeypatch.setenv("NEMOTRON_VOICECHAT_LLM_PATH", "/local/nemotron") + cfg = sanitize_tts_model_cfg(original) + assert cfg["pretrained_lm_name"] == "/local/nemotron" + + +def test_code2wav_token_only_sizes_placeholder() -> None: + talker_output = SimpleNamespace( + finished=True, + outputs=[SimpleNamespace(token_ids=[0] * 9, cumulative_token_ids=None)], + ) + inputs = talker2code2wav_token_only([talker_output]) + assert len(inputs) == 1 + assert len(inputs[0]["prompt_token_ids"]) == 9 + + +# ========================= +# Async-chunk (streaming) producers +# ========================= + + +def _fake_transfer_manager(codec_chunk_frames: int | None = None): + extra = {} + if codec_chunk_frames is not None: + extra["codec_chunk_frames"] = codec_chunk_frames + return SimpleNamespace( + request_payload={}, + connector=SimpleNamespace(config={"extra": extra}), + put_req_chunk={}, + ) + + +def _streaming_request(prompt_len: int, generated: list[int], info: dict | None = None): + return SimpleNamespace( + external_req_id="req-0", + prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD], + output_token_ids=generated, + additional_information=info, + ) + + +def test_thinker2talker_async_chunk_cumulative_timeline() -> None: + from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import ( + thinker2talker_async_chunk, + ) + + tm = _fake_transfer_manager() + req = _streaming_request(prompt_len=3, generated=[7]) + payload = thinker2talker_async_chunk(tm, {}, req, is_finished=False) + assert payload is not None + assert payload.ids.all == [_PAD] * 3 + [7] + assert payload.meta.num_processed_tokens == 3 + assert not bool(payload.meta.finished) + + # No new tokens since the last chunk -> skip. + assert thinker2talker_async_chunk(tm, {}, req, is_finished=False) is None + + # More tokens -> longer cumulative timeline; final chunk sets finished. + req.output_token_ids = [7, 8, 9] + payload = thinker2talker_async_chunk(tm, {}, req, is_finished=True) + assert payload.ids.all == [_PAD] * 3 + [7, 8, 9] + assert bool(payload.meta.finished) + + +def test_talker2code2wav_async_chunk_cadence_and_left_context() -> None: + from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import ( + talker2code2wav_async_chunk, + ) + + tm = _fake_transfer_manager(codec_chunk_frames=2) + prompt_len = 3 + info = {"nvc_logical_prompt_len": prompt_len} + req = _streaming_request(prompt_len=prompt_len, generated=[], info=info) + + def cumulative(rows: int) -> dict: + # Talker rows are timeline steps t=1..; prompt region occupies the + # first prompt_len-1 rows before acoustic frames start. + return {"codes": {"audio": torch.arange((prompt_len - 1 + rows) * 31).reshape(-1, 31)}} + + # Only prompt-region rows so far -> nothing to ship. + assert talker2code2wav_async_chunk(tm, cumulative(0), req, is_finished=False) is None + # One new frame < chunk size 2 -> hold. + assert talker2code2wav_async_chunk(tm, cumulative(1), req, is_finished=False) is None + # Two frames -> first chunk, no left context. + payload = talker2code2wav_async_chunk(tm, cumulative(2), req, is_finished=False) + assert payload is not None + assert payload.meta.left_context_size == 0 + assert payload.codes.audio.shape == (2, 31) + # Third frame + finished -> cumulative history with left context 2. + payload = talker2code2wav_async_chunk(tm, cumulative(3), req, is_finished=True) + assert payload.meta.left_context_size == 2 + assert payload.codes.audio.shape == (3, 31) + assert bool(payload.meta.finished) + + +def test_talker2code2wav_async_chunk_requires_prompt_len() -> None: + from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import ( + talker2code2wav_async_chunk, + ) + + tm = _fake_transfer_manager() + req = _streaming_request(prompt_len=3, generated=[], info=None) + with pytest.raises(ValueError, match="logical prompt length"): + talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req) + + +# ========================= +# Talker drain-per-wake (async scheduling review P1/P2) +# ========================= + + +class _FakeTTSModelOutput: + past_key_values = None + + +class _FakeTTS: + """Stub DuplexEARTTS: infer_codes_one_step returns a constant code row.""" + + def __init__(self, q: int = 31): + self.q = q + self.steps = 0 + + def set_init_inputs(self, speaker_name: str) -> None: + del speaker_name + + def get_init_inputs(self, B: int) -> dict: # noqa: N803 — mirrors the vendored DuplexEARTTS signature + return { + "code": torch.zeros((B, 2, self.q), dtype=torch.long), + "subword_ids": torch.zeros((B, 3), dtype=torch.long), + } + + def _get_generation_config(self, guidance_enabled: bool): + del guidance_enabled + return {} + + def tts_model(self, **kwargs): + del kwargs + return _FakeTTSModelOutput() + + def infer_codes_one_step(self, **kwargs): + del kwargs + self.steps += 1 + return torch.full((1, 1, self.q), self.steps, dtype=torch.long), None + + +def _bare_talker(): + from torch import nn + + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import ( + NemotronVoiceChatTalkerForConditionalGeneration, + ) + + talker = NemotronVoiceChatTalkerForConditionalGeneration.__new__(NemotronVoiceChatTalkerForConditionalGeneration) + nn.Module.__init__(talker) + talker.tts = _FakeTTS() + talker._sessions = {} + talker._hidden = 8 + talker._dtype = torch.float32 + talker._speaker_name = "Aria" + talker.config = SimpleNamespace(tts_cfg={"inference_guidance_enabled": False}) + return talker + + +def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict: + return { + "request_id": request_id, + "additional_information": { + "ids": {"all": list(timeline)}, + "meta": { + "num_processed_tokens": prompt_len, + "finished": torch.tensor(finished), + "request_id": request_id, + }, + }, + } + + +def test_talker_drains_all_received_positions_per_wake() -> None: + """P2: chunk 0's PAD prompt region must not cost one upstream chunk per step, + and P1: a coalesced chunk (2+ new positions, 1 wake) must be fully drained.""" + talker = _bare_talker() + pad, prompt_len = 12, 4 + ids = torch.zeros(1, dtype=torch.long) + + # Prefill wake: chunk 0 carries only the PAD prompt region. + info = _async_info([pad] * prompt_len, prompt_len, finished=False) + talker.preprocess(ids, None, _omni_is_prefill=True, **info) + session = talker._sessions["req-0"] + assert session["step"] == 1 and not session["sync_mode"] + + # First decode wake: drains ALL prompt-region steps (t=1..P-1) at once. + _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info) + assert session["step"] == prompt_len # t advanced 1 -> P in one wake + assert update["codes"]["audio"].shape == (prompt_len - 1, 31) + assert float(embeds[0, 0]) == 0.0 # not finished + + # Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save + # thread merged two thinker steps). Both must be consumed by this wake. + info2 = _async_info([pad] * prompt_len + [7, 8], prompt_len, finished=False) + _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2) + assert session["step"] == prompt_len + 2 + assert update["codes"]["audio"].shape == (prompt_len + 1, 31) + assert float(embeds[0, 0]) == 0.0 + + # Zero-progress wake (no new positions, not finished): plain CONTINUE, no + # fabricated frames, no codes update. + _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2) + assert session["step"] == prompt_len + 2 + assert update == {} + assert float(embeds[0, 0]) == 0.0 + + # Final chunk: one more token + finished marker -> drain + stop flag. + info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True) + _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3) + assert session["step"] == prompt_len + 3 + assert update["codes"]["audio"].shape == (prompt_len + 2, 31) + assert float(embeds[0, 0]) == 1.0 # stop + + +def test_talker_sync_mode_steps_once_per_wake() -> None: + """The verified sync parity path keeps its one-TTS-step-per-engine-step cadence.""" + talker = _bare_talker() + pad, prompt_len = 12, 3 + ids = torch.zeros(1, dtype=torch.long) + info = { + "request_id": "req-1", + "additional_information": { + "nvc_text_timeline": [pad] * prompt_len + [7, 8], + "nvc_logical_prompt_len": prompt_len, + }, + } + talker.preprocess(ids, None, _omni_is_prefill=True, **info) + session = talker._sessions["req-1"] + assert session["sync_mode"] and session["upstream_finished"] + for expected_step in (2, 3, 4): + _, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info) + assert session["step"] == expected_step + # Last position -> stop flag on the final step. + _, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info) + assert session["step"] == 5 + assert float(embeds[0, 0]) == 1.0 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 5ae83d1b9..4b5c73a39 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -128,6 +128,7 @@ class _CacheEngineOverrides(TypedDict, total=False): enable_prefix_caching: bool disable_hybrid_kv_cache_manager: bool mm_processor_cache_gb: float + mamba_ssm_cache_dtype: str class _SchedulerEngineOverrides(TypedDict, total=False): @@ -377,6 +378,8 @@ class OmniStageCacheConfig: enable_prefix_caching: bool | None = None disable_hybrid_kv_cache_manager: bool | None = None mm_processor_cache_gb: float | None = Field(default=None, ge=0.0) + # Hybrid-mamba SSM state dtype ("auto"/"float32"); vLLM CacheConfig field. + mamba_ssm_cache_dtype: str | None = None @config diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index ac2e9d38d..f02a3578e 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -88,6 +88,9 @@ from vllm_omni.model_executor.models.moss_tts.pipeline import ( MOSS_TTS_REALTIME_PIPELINE, ) from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE +from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import ( + NEMOTRON_VOICECHAT_PIPELINE, +) from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE from vllm_omni.model_executor.models.personaplex.pipeline import PERSONAPLEX_PIPELINE from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import ( @@ -118,6 +121,11 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "qwen2_5_omni": QWEN2_5_OMNI_PIPELINE, "qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, "personaplex": PERSONAPLEX_PIPELINE, + "nemotron_voicechat": NEMOTRON_VOICECHAT_PIPELINE, + # Alias: lets bare `vllm-omni serve ` + # auto-detect through the path-basename fallback (the checkpoint config.json + # has no model_type key). + "nemotron_labs_voicechat": NEMOTRON_VOICECHAT_PIPELINE, "qwen3_omni_moe": resolve_qwen3_omni_pipeline, "qwen3_tts": QWEN3_TTS_PIPELINE, "step_audio_2": STEP_AUDIO2_PIPELINE, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index fa8e45295..496aea208 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -357,6 +357,9 @@ class StageDeployConfig: async_scheduling: bool | None = None disable_hybrid_kv_cache_manager: bool | None = None mm_processor_cache_gb: float | None = None + # Hybrid-mamba stages (e.g. the NemotronVoiceChat thinker's NemotronH + # backbone) pin the SSM state dtype; projected onto vLLM CacheConfig. + mamba_ssm_cache_dtype: str | None = None # Generic compilation, profiling, tokenizer/config parsing, and model # loading controls. diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 424fed2d9..75f234d8c 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -51,6 +51,10 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset( ("MiniCPMO45Code2Wav", "code2wav"), # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1). ("CosyVoice3Model", "cosyvoice3_code2wav"), + # nemotron_voicechat: talker (Stage 1) -> code2wav (Stage 2). Stage 2 + # waits for the talker's full-payload code stacks; the thinker (Stage 0) + # -> talker (Stage 1) hop is token-path only and must NOT be listed. + ("NemotronVoiceChatCode2Wav", "code2wav"), # audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1). # The default deploy is async_chunk; this covers async_chunk: false. ("AudexCode2Wav", "audex_code2wav"), diff --git a/vllm_omni/deploy/nemotron_labs_voicechat.yaml b/vllm_omni/deploy/nemotron_labs_voicechat.yaml new file mode 100644 index 000000000..fe872caea --- /dev/null +++ b/vllm_omni/deploy/nemotron_labs_voicechat.yaml @@ -0,0 +1,81 @@ +# NemotronVoiceChat-11B offline speech-to-speech deployment. +# 3 stages on one GPU: thinker (NemotronH AR) -> talker (EAR-TTS AR) -> code2wav +# (RVQ-VAE decode). thinker->talker is token-path only; talker->code2wav ships +# the [frames, 31] code stack via the shared-memory connector. +pipeline: nemotron_voicechat +async_chunk: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + +stages: + - stage_id: 0 # thinker: speech -> frame-locked text (final text output) + max_num_seqs: 1 + gpu_memory_utilization: 0.62 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + # fp32 matches the NeMo reference exactly (196/196 token parity on the + # acceptance fixture). bf16 is ~2x faster and measured within one word of + # the reference on the same fixture — flip to bfloat16 if you prefer speed + # over exact parity. + dtype: float32 + # NemotronH mamba SSM state in fp32 (matches NVIDIA's own vLLM deployment + # of this model; bf16 SSM state accumulates drift over long sequences). + mamba_ssm_cache_dtype: float32 + devices: "0" + default_sampling_params: + # Greedy, frame-locked: the example sets max_tokens to the + # exact acoustic frame count per request; text EOS never stops generation. + temperature: 0.0 + ignore_eos: true + seed: 0 + + - stage_id: 1 # talker: text timeline -> 31-quantizer code stacks + max_num_seqs: 1 + gpu_memory_utilization: 0.12 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + skip_tokenizer_init: true + max_num_batched_tokens: 4096 + max_model_len: 16384 + # fp32 like the NeMo reference: the EAR-TTS MoG/codec-embedding path mixes + # fp32 buffers into the forward, and 1B params in fp32 are cheap. + dtype: float32 + devices: "0" + output_connectors: + to_stage_2: connector_of_shared_memory + default_sampling_params: + # Placeholder-token loop; the stage emits stop token 1 after the last + # timeline position. 16383 (not 16384): the request also carries one + # placeholder prompt token, and prompt + max_tokens must fit + # max_model_len. + temperature: 0.0 + max_tokens: 16383 + stop_token_ids: [1] + detokenize: false + seed: 0 + + - stage_id: 2 # code2wav: codes -> 22.05 kHz PCM (fp32 codec) + max_num_seqs: 1 + gpu_memory_utilization: 0.06 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + skip_tokenizer_init: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + detokenize: false + seed: 0 diff --git a/vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml b/vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml new file mode 100644 index 000000000..133bbee34 --- /dev/null +++ b/vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml @@ -0,0 +1,84 @@ +# NemotronVoiceChat-11B STREAMING speech-to-speech deployment (async-chunk). +# +# Same 3 stages as nemotron_labs_voicechat.yaml, but the stages overlap: +# the thinker streams its frame-locked text timeline to the talker per decode +# step (cumulative ids.all payloads), and the talker streams prompt-trimmed +# cumulative code stacks to code2wav every `codec_chunk_frames` frames. The +# codec re-decodes the full prefix per chunk and emits only the new samples +# (NeMo decode_one_audio_step semantics), so the concatenated stream is +# seam-free. First audio arrives after ~codec_chunk_frames thinker steps +# instead of after the whole pipeline completes. +# +# The non-streaming yaml remains the parity-acceptance path (bit-exact WAV); +# streamed audio is validated by ASR/WER instead (prefix decodes are not +# bit-identical to the one-shot decode of the final stack). +pipeline: nemotron_voicechat +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + # Streaming cadence for talker -> code2wav: emit a chunk once this many + # NEW frames (80 ms each) have accumulated; ~1 s of audio per chunk. + codec_chunk_frames: 13 + +stages: + - stage_id: 0 # thinker: speech -> frame-locked text (final text output) + max_num_seqs: 1 + gpu_memory_utilization: 0.62 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + dtype: float32 + mamba_ssm_cache_dtype: float32 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + ignore_eos: true + seed: 0 + + - stage_id: 1 # talker: text timeline -> 31-quantizer code stacks + max_num_seqs: 1 + gpu_memory_utilization: 0.12 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + skip_tokenizer_init: true + max_num_batched_tokens: 4096 + max_model_len: 16384 + dtype: float32 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + output_connectors: + to_stage_2: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + max_tokens: 16383 + stop_token_ids: [1] + detokenize: false + seed: 0 + + - stage_id: 2 # code2wav: codes -> 22.05 kHz PCM (fp32 codec) + max_num_seqs: 1 + gpu_memory_utilization: 0.06 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + skip_tokenizer_init: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + detokenize: false + seed: 0 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 1d9eb803b..a3aa96f60 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -30,6 +30,11 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = { # the HF config without manual overrides. "PersonaPlexTalkerForConditionalGeneration": "personaplex", "PersonaPlexCode2Wav": "personaplex", + # NemotronVoiceChat ships a NeMo-style config.json without model_type; + # create_model_config() patches model_type=nemotron_voicechat from the arch. + "NemotronVoiceChatThinkerForConditionalGeneration": "nemotron_voicechat", + "NemotronVoiceChatTalkerForConditionalGeneration": "nemotron_voicechat", + "NemotronVoiceChatCode2Wav": "nemotron_voicechat", "VoxCPM2TalkerForConditionalGeneration": "voxcpm2", } @@ -56,6 +61,9 @@ def _register_omni_hf_configs() -> None: MossTTSLocalConfig, MossTTSRealtimeConfig, ) + from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import ( + NemotronVoiceChatConfig, + ) from vllm_omni.model_executor.models.personaplex.configuration_personaplex import ( PersonaPlexConfig, ) @@ -87,6 +95,7 @@ def _register_omni_hf_configs() -> None: ("moss_tts_realtime", MossTTSRealtimeConfig), ("qwen3_tts", Qwen3TTSConfig), ("personaplex", PersonaPlexConfig), + ("nemotron_voicechat", NemotronVoiceChatConfig), ("cosyvoice3", CosyVoice3Config), ("glm_tts", GLMTTSConfig), ("omnivoice", OmniVoiceConfig), @@ -324,6 +333,19 @@ class OmniEngineArgs(EngineArgs): if model_type is not None: self._patch_empty_hf_config(model_type) + # NemotronVoiceChat: the checkpoint's only tokenizer subfolder is the + # auxiliary rnnt_tokenizer/ (ASR loss vocab), which the generic + # subfolder auto-detect below would wrongly pick up. The text tokenizer + # is the Nemotron backbone's, resolved from a local override or its + # HF id (cache-friendly). + if not self.tokenizer and self.model_arch in ( + "NemotronVoiceChatThinkerForConditionalGeneration", + "NemotronVoiceChatTalkerForConditionalGeneration", + "NemotronVoiceChatCode2Wav", + ): + self.tokenizer = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + logger.info("NemotronVoiceChat: using text tokenizer from %s", self.tokenizer) + # Auto-detect tokenizer for models that store it in a subdirectory # rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/). if not self.tokenizer and self.model: diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/__init__.py b/vllm_omni/model_executor/models/nemotron_voicechat/__init__.py new file mode 100644 index 000000000..f8046bfc5 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA NemotronLabs VoiceChat-11B: offline speech-to-speech 3-stage pipeline.""" + +from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import ( + NemotronVoiceChatCode2WavConfig, + NemotronVoiceChatConfig, + NemotronVoiceChatTalkerConfig, +) +from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import ( + NEMOTRON_VOICECHAT_PIPELINE, +) + +__all__ = [ + "NemotronVoiceChatConfig", + "NemotronVoiceChatTalkerConfig", + "NemotronVoiceChatCode2WavConfig", + "NEMOTRON_VOICECHAT_PIPELINE", +] diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py b/vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py new file mode 100644 index 000000000..1469f72df --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HF-compatible configs synthesized from the NeMo-style NemotronVoiceChat checkpoint. + +The ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` checkpoint ships a NeMo training +config (top-level ``data`` / ``model`` keys) instead of a transformers config: +there is no ``model_type`` or ``architectures``. vLLM-Omni's empty-config patching +(:meth:`OmniEngineArgs._patch_empty_hf_config`) injects +``model_type="nemotron_voicechat"`` and routes the raw dict here, where it is +parsed into per-stage sub-configs: + +* ``get_text_config()`` -> the NemotronH backbone config (thinker stage; the AR + runner reads ``hidden_size`` etc. from it). Resolved from the checkpoint's + ``model.stt.model.pretrained_llm`` HF id via the local HF cache, or from the + ``NEMOTRON_VOICECHAT_LLM_PATH`` env var pointing at a local snapshot. +* ``talker_config`` -> Gemma3-based EAR-TTS backbone dims (talker stage). +* ``code2wav_config`` -> RVQ-VAE codec dims (code2wav stage). + +The raw NeMo sections are preserved on the config (``stt_cfg`` / ``tts_cfg`` / +``nemo_data``) because the stage models consume them verbatim (perception layout, +fusion weights, TTS generation parameters, codec architecture). +""" + +from __future__ import annotations + +import os +from typing import Any + +from transformers import AutoConfig, PretrainedConfig + +_EXPECTED_STRUCTURE = ( + "Expected a NemotronVoiceChat checkpoint config.json with the NeMo layout: " + "top-level 'model' containing both 'stt' (with 'model.stt.model.perception' " + "and 'model.stt.model.pretrained_llm') and 'speech_generation' (with " + "'model.speech_generation.model.tts_config' and '...codec_config')." +) + +# Default HF id of the thinker text backbone; the checkpoint records the same id +# under model.stt.model.pretrained_llm. +_DEFAULT_LLM_ID = "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + +# Env override for air-gapped runs: local directory containing the Nemotron +# backbone's config.json (and optionally tokenizer files). +_LLM_PATH_ENV = "NEMOTRON_VOICECHAT_LLM_PATH" + + +class NemotronVoiceChatTalkerConfig(PretrainedConfig): + """Stage config for the EAR-TTS talker (28-layer Gemma3-style backbone).""" + + model_type = "nemotron_voicechat_talker" + + def __init__(self, **kwargs: Any) -> None: + backbone: dict[str, Any] = kwargs.pop("backbone_config", None) or {} + super().__init__(**kwargs) + self.backbone_config = backbone + self.hidden_size = int(backbone.get("hidden_size", 1152)) + self.num_hidden_layers = int(backbone.get("num_hidden_layers", 28)) + self.num_attention_heads = int(backbone.get("num_attention_heads", 16)) + self.head_dim = int(backbone.get("head_dim", 72)) + self.intermediate_size = int(backbone.get("intermediate_size", 4608)) + self.sliding_window = int(backbone.get("sliding_window", 7500)) + # The talker consumes/produces per-frame code stacks; vocab is nominal + # (codes are sampled by the vendored MOG head, not a vLLM sampler). + self.vocab_size = int(kwargs.get("vocab_size", 1024)) + self.max_position_embeddings = int(backbone.get("max_position_embeddings", 32768)) + + def get_text_config(self, **_: Any) -> NemotronVoiceChatTalkerConfig: + return self + + +class NemotronVoiceChatCode2WavConfig(PretrainedConfig): + """Stage config for the RVQ-VAE codec decoder.""" + + model_type = "nemotron_voicechat_code2wav" + + def __init__(self, **kwargs: Any) -> None: + codec: dict[str, Any] = kwargs.pop("codec_config", None) or {} + super().__init__(**kwargs) + self.codec_config = codec + self.num_quantizers = int(codec.get("num_quantizers", 31)) + self.codebook_size = int(codec.get("codebook_size", 1024)) + self.latent_size = int(codec.get("latent_size", 512)) + self.wav_to_token_ratio = int(codec.get("wav_to_token_ratio", 1764)) + self.sample_rate = int(kwargs.get("sample_rate", 22050)) + # The generation runner sizes buffers from hidden_size. + self.hidden_size = self.latent_size + self.num_attention_heads = 1 + self.num_hidden_layers = 1 + self.vocab_size = self.codebook_size + self.max_position_embeddings = int(kwargs.get("max_position_embeddings", 65536)) + + def get_text_config(self, **_: Any) -> NemotronVoiceChatCode2WavConfig: + return self + + +class NemotronVoiceChatConfig(PretrainedConfig): + """Unified config: parses the raw NeMo dict into per-stage sub-configs.""" + + model_type = "nemotron_voicechat" + + def __init__(self, **kwargs: Any) -> None: + nemo_model = kwargs.pop("model", None) + nemo_data = kwargs.pop("data", None) + thinker_text_config = kwargs.pop("thinker_text_config", None) + talker_config = kwargs.pop("talker_config", None) + code2wav_config = kwargs.pop("code2wav_config", None) + stt_cfg = kwargs.pop("stt_cfg", None) + tts_cfg = kwargs.pop("tts_cfg", None) + # Drop remaining NeMo training-only top-level sections. + for training_key in ("exp_manager", "trainer", "hf_export_dir", "name"): + kwargs.pop(training_key, None) + # transformers may call get_text_config() during super().__init__; + # give it a safe placeholder until the real sub-config is resolved. + self.thinker_text_config = None + super().__init__(**kwargs) + + if nemo_model is not None: + stt = nemo_model.get("stt") or {} + speech_generation = nemo_model.get("speech_generation") or {} + if not (isinstance(stt, dict) and isinstance(stt.get("model"), dict)) or not ( + isinstance(speech_generation, dict) and isinstance(speech_generation.get("model"), dict) + ): + raise ValueError( + "config.json is missing the NemotronVoiceChat NeMo signature keys " + "'model.stt.model' and/or 'model.speech_generation.model'. " + _EXPECTED_STRUCTURE + ) + stt_cfg = stt["model"] + tts_cfg = speech_generation["model"] + self.inference_speaker_name = nemo_model.get("inference_speaker_name") or "Aria" + self.tts_data = speech_generation.get("data") or {} + elif stt_cfg is None or tts_cfg is None: + # Bare construction (e.g. transformers plumbing); leave a skeleton + # that fails loudly if a stage model tries to use it. + self.inference_speaker_name = kwargs.get("inference_speaker_name", "Aria") + self.tts_data = {} + if stt_cfg is None or tts_cfg is None: + self.stt_cfg: dict[str, Any] = {} + self.tts_cfg: dict[str, Any] = {} + else: + if "perception" not in stt_cfg: + raise ValueError("'model.stt.model.perception' is missing from config.json. " + _EXPECTED_STRUCTURE) + if "tts_config" not in tts_cfg or "codec_config" not in tts_cfg: + raise ValueError( + "'model.speech_generation.model.tts_config'/'codec_config' are missing " + "from config.json. " + _EXPECTED_STRUCTURE + ) + self.stt_cfg = dict(stt_cfg) + self.tts_cfg = dict(tts_cfg) + self.nemo_data = dict(nemo_data) if isinstance(nemo_data, dict) else {} + self.frame_length = float(self.nemo_data.get("frame_length", 0.08)) + self.source_sample_rate = int(self.nemo_data.get("source_sample_rate", 16000)) + self.target_sample_rate = int(self.nemo_data.get("target_sample_rate", 22050)) + + self.thinker_text_config = self._resolve_thinker_text_config(thinker_text_config) + self.talker_config = self._coerce( + talker_config, + NemotronVoiceChatTalkerConfig, + defaults={"backbone_config": (self.tts_cfg.get("tts_config") or {}).get("backbone_config") or {}}, + ) + self.code2wav_config = self._coerce( + code2wav_config, + NemotronVoiceChatCode2WavConfig, + defaults={ + "codec_config": self.tts_cfg.get("codec_config") or {}, + "sample_rate": self.target_sample_rate, + }, + ) + + if not getattr(self, "architectures", None): + self.architectures = ["NemotronVoiceChatThinkerForConditionalGeneration"] + + @staticmethod + def _coerce(value: Any, cls: type[PretrainedConfig], defaults: dict[str, Any]) -> PretrainedConfig: + if isinstance(value, cls): + return value + if isinstance(value, PretrainedConfig): + value = value.to_dict() + if isinstance(value, dict): + merged = dict(defaults) + merged.update(value) + return cls(**merged) + return cls(**defaults) + + def _resolve_thinker_text_config(self, provided: Any) -> PretrainedConfig: + """NemotronH backbone config: explicit > env-local > HF cache/hub.""" + if isinstance(provided, PretrainedConfig): + return provided + if isinstance(provided, dict) and provided: + try: + return AutoConfig.for_model(**dict(provided)) + except (KeyError, ValueError): + return PretrainedConfig(**provided) + if not self.stt_cfg: + # Skeleton config (bare construction); resolved lazily by the stage. + return PretrainedConfig() + llm_ref = os.environ.get(_LLM_PATH_ENV) or self.stt_cfg.get("pretrained_llm") or _DEFAULT_LLM_ID + try: + # nemotron_h is a native transformers model type; trust_remote_code=True + # would chase the auto_map .py modules and break air-gapped dirs that + # ship only config.json. + return AutoConfig.from_pretrained(llm_ref, trust_remote_code=False) + except Exception as exc: + raise RuntimeError( + f"Failed to load the NemotronVoiceChat thinker backbone config from {llm_ref!r}. " + f"Make sure the HF cache contains 'nvidia/NVIDIA-Nemotron-Nano-9B-v2' or set " + f"{_LLM_PATH_ENV} to a local directory with its config.json." + ) from exc + + def get_text_config(self, **_: Any) -> PretrainedConfig: + text_config = self.__dict__.get("thinker_text_config") + return text_config if text_config is not None else self diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py new file mode 100644 index 000000000..d30a94542 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); dependency-stripped +# copies of the NeMo inference modules needed to run NVIDIA-NemotronLabs-VoiceChat-11B. + +from .duplex_ear_tts import DuplexEARTTS +from .ear_tts_commons import PreTrainedModel +from .ear_tts_vae_codec import RVQVAEModel +from .fusion import create_fusion_module +from .perception import AudioPerceptionModule, IdentityConnector +from .tokenizer import AutoTokenizer + +__all__ = [ + "AudioPerceptionModule", + "AutoTokenizer", + "DuplexEARTTS", + "IdentityConnector", + "PreTrainedModel", + "RVQVAEModel", + "create_fusion_module", +] diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py new file mode 100644 index 000000000..d1dda45b6 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); dependency-stripped +# copies of the NeMo ASR conformer modules needed by the perception encoder. + +from .audio_preprocessing import AudioToMelSpectrogramPreprocessor +from .conformer_encoder import ConformerEncoder + +__all__ = [ + "AudioToMelSpectrogramPreprocessor", + "ConformerEncoder", +] diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py new file mode 100644 index 000000000..e4343ef44 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/modules/audio_preprocessing.py; training-only code removed. +# Kept: AudioPreprocessor, AudioToMelSpectrogramPreprocessor (+ its config dataclass). +# Removed: AudioToMFCCPreprocessor, SpectrogramAugmentation, MaskedPatchAugmentation, +# CropOrPadSpectrogramAugmentation and numba/spec-augment imports (training-only). + +import math +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Optional + +import torch +from packaging import version + +from ..compat import ( + Exportable, + NeuralModule, + logging, + logging_mode, + typecheck, +) +from .features import FilterbankFeatures, FilterbankFeaturesTA + +try: + import torchaudio + import torchaudio.functional + import torchaudio.transforms + + TORCHAUDIO_VERSION = version.parse(torchaudio.__version__) + TORCHAUDIO_VERSION_MIN = version.parse("0.5") + + HAVE_TORCHAUDIO = True +except ModuleNotFoundError: + HAVE_TORCHAUDIO = False + +__all__ = [ + "AudioToMelSpectrogramPreprocessor", +] + + +class AudioPreprocessor(NeuralModule, ABC): + """ + An interface for Neural Modules that performs audio pre-processing, + transforming the wav files to features. + """ + + def __init__(self, win_length, hop_length): + super().__init__() + + self.win_length = win_length + self.hop_length = hop_length + + self.torch_windows = { + "hann": torch.hann_window, + "hamming": torch.hamming_window, + "blackman": torch.blackman_window, + "bartlett": torch.bartlett_window, + "ones": torch.ones, + None: torch.ones, + } + + # Normally, when you call to(dtype) on a torch.nn.Module, all + # floating point parameters and buffers will change to that + # dtype, rather than being float32. The AudioPreprocessor + # classes, uniquely, don't actually have any parameters or + # buffers from what I see. In addition, we want the input to + # the preprocessor to be float32, but need to create the + # output in appropriate precision. We have this empty tensor + # here just to detect which dtype tensor this module should + # output at the end of execution. + self.register_buffer("dtype_sentinel_tensor", torch.tensor((), dtype=torch.float32), persistent=False) + + @typecheck() + @torch.no_grad() + def forward(self, input_signal, length): + if input_signal.dtype != torch.float32: + logging.warning( + f"AudioPreprocessor received an input signal of dtype {input_signal.dtype}, rather than torch.float32. In sweeps across multiple datasets, we have found that the preprocessor is not robust to low precision mathematics. As such, it runs in float32. Your input will be cast to float32, but this is not necessarily enough to recovery full accuracy. For example, simply casting input_signal from torch.float32 to torch.bfloat16, then back to torch.float32 before running AudioPreprocessor causes drops in absolute WER of up to 0.1%. torch.bfloat16 simply does not have enough mantissa bits to represent enough values in the range [-1.0,+1.0] correctly.", + mode=logging_mode.ONCE, + ) + processed_signal, processed_length = self.get_features(input_signal.to(torch.float32), length) + processed_signal = processed_signal.to(self.dtype_sentinel_tensor.dtype) + return processed_signal, processed_length + + @abstractmethod + def get_features(self, input_signal, length): + # Called by forward(). Subclasses should implement this. + pass + + +class AudioToMelSpectrogramPreprocessor(AudioPreprocessor, Exportable): + """Featurizer module that converts wavs to mel spectrograms. + + Args: + sample_rate (int): Sample rate of the input audio data. + Defaults to 16000 + window_size (float): Size of window for fft in seconds + Defaults to 0.02 + window_stride (float): Stride of window for fft in seconds + Defaults to 0.01 + n_window_size (int): Size of window for fft in samples + Defaults to None. Use one of window_size or n_window_size. + n_window_stride (int): Stride of window for fft in samples + Defaults to None. Use one of window_stride or n_window_stride. + window (str): Windowing function for fft. can be one of ['hann', + 'hamming', 'blackman', 'bartlett'] + Defaults to "hann" + normalize (str): Can be one of ['per_feature', 'all_features']; all + other options disable feature normalization. 'all_features' + normalizes the entire spectrogram to be mean 0 with std 1. + 'pre_features' normalizes per channel / freq instead. + Defaults to "per_feature" + n_fft (int): Length of FT window. If None, it uses the smallest power + of 2 that is larger than n_window_size. + Defaults to None + preemph (float): Amount of pre emphasis to add to audio. Can be + disabled by passing None. + Defaults to 0.97 + features (int): Number of mel spectrogram freq bins to output. + Defaults to 64 + lowfreq (int): Lower bound on mel basis in Hz. + Defaults to 0 + highfreq (int): Lower bound on mel basis in Hz. + Defaults to None + log (bool): Log features. + Defaults to True + log_zero_guard_type(str): Need to avoid taking the log of zero. There + are two options: "add" or "clamp". + Defaults to "add". + log_zero_guard_value(float, or str): Add or clamp requires the number + to add with or clamp to. log_zero_guard_value can either be a float + or "tiny" or "eps". torch.finfo is used if "tiny" or "eps" is + passed. + Defaults to 2**-24. + dither (float): Amount of white-noise dithering. + Defaults to 1e-5 + pad_to (int): Ensures that the output size of the time dimension is + a multiple of pad_to. + Defaults to 16 + frame_splicing (int): Defaults to 1 + exact_pad (bool): If True, sets stft center to False and adds padding, such that num_frames = audio_length + // hop_length. Defaults to False. + pad_value (float): The value that shorter mels are padded with. + Defaults to 0 + mag_power (float): The power that the linear spectrogram is raised to + prior to multiplication with mel basis. + Defaults to 2 for a power spec + rng : Random number generator + nb_augmentation_prob (float) : Probability with which narrowband augmentation would be applied to + samples in the batch. + Defaults to 0.0 + nb_max_freq (int) : Frequency above which all frequencies will be masked for narrowband augmentation. + Defaults to 4000 + use_torchaudio: Whether to use the `torchaudio` implementation. + mel_norm: Normalization used for mel filterbank weights. + Defaults to 'slaney' (area normalization) + stft_exact_pad: Deprecated argument, kept for compatibility with older checkpoints. + stft_conv: Deprecated argument, kept for compatibility with older checkpoints. + """ + + def save_to(self, save_path: str): + pass + + @classmethod + def restore_from(cls, restore_path: str): + pass + + def __init__( + self, + sample_rate=16000, + window_size=0.02, + window_stride=0.01, + n_window_size=None, + n_window_stride=None, + window="hann", + normalize="per_feature", + n_fft=None, + preemph=0.97, + features=64, + lowfreq=0, + highfreq=None, + log=True, + log_zero_guard_type="add", + log_zero_guard_value=2**-24, + dither=1e-5, + pad_to=16, + frame_splicing=1, + exact_pad=False, + pad_value=0, + mag_power=2.0, + rng=None, + nb_augmentation_prob=0.0, + nb_max_freq=4000, + use_torchaudio: bool = False, + mel_norm="slaney", + stft_exact_pad=False, # Deprecated arguments; kept for config compatibility + stft_conv=False, # Deprecated arguments; kept for config compatibility + ): + super().__init__(n_window_size, n_window_stride) + + self._sample_rate = sample_rate + if window_size and n_window_size: + raise ValueError(f"{self} received both window_size and n_window_size. Only one should be specified.") + if window_stride and n_window_stride: + raise ValueError(f"{self} received both window_stride and n_window_stride. Only one should be specified.") + if window_size: + n_window_size = int(window_size * self._sample_rate) + if window_stride: + n_window_stride = int(window_stride * self._sample_rate) + + # Given the long and similar argument list, point to the class and instantiate it by reference + if not use_torchaudio: + featurizer_class = FilterbankFeatures + else: + featurizer_class = FilterbankFeaturesTA + self.featurizer = featurizer_class( + sample_rate=self._sample_rate, + n_window_size=n_window_size, + n_window_stride=n_window_stride, + window=window, + normalize=normalize, + n_fft=n_fft, + preemph=preemph, + nfilt=features, + lowfreq=lowfreq, + highfreq=highfreq, + log=log, + log_zero_guard_type=log_zero_guard_type, + log_zero_guard_value=log_zero_guard_value, + dither=dither, + pad_to=pad_to, + frame_splicing=frame_splicing, + exact_pad=exact_pad, + pad_value=pad_value, + mag_power=mag_power, + rng=rng, + nb_augmentation_prob=nb_augmentation_prob, + nb_max_freq=nb_max_freq, + mel_norm=mel_norm, + stft_exact_pad=stft_exact_pad, # Deprecated arguments; kept for config compatibility + stft_conv=stft_conv, # Deprecated arguments; kept for config compatibility + ) + + def get_features(self, input_signal, length): + return self.featurizer(input_signal, length) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py new file mode 100644 index 000000000..4d593895e --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/causal_convs.py; training-only code removed. + +from typing import Union + +import torch +import torch.nn.functional as F +from torch import nn + +__all__ = ["CausalConv2D", "CausalConv1D"] + + +class CausalConv2D(nn.Conv2d): + """ + A causal version of nn.Conv2d where each location in the 2D matrix would have no access to locations on its right or down + All arguments are the same as nn.Conv2d except padding which should be set as None + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: Union[str, int] = 0, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + if padding is not None: + raise ValueError("Argument padding should be set to None for CausalConv2D.") + self._left_padding = kernel_size - 1 + self._right_padding = stride - 1 + + padding = 0 + super(CausalConv2D, self).__init__( + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + groups, + bias, + padding_mode, + device, + dtype, + ) + + def forward( + self, + x, + ): + x = F.pad(x, pad=(self._left_padding, self._right_padding, self._left_padding, self._right_padding)) + x = super().forward(x) + return x + + +class CausalConv1D(nn.Conv1d): + """ + A causal version of nn.Conv1d where each step would have limited access to locations on its right or left + All arguments are the same as nn.Conv1d except padding. + + If padding is set None, then paddings are set automatically to make it a causal convolution where each location would not see any steps on its right. + + If padding is set as a list (size of 2), then padding[0] would be used as left padding and padding[1] as right padding. + It would make it possible to control the number of steps to be accessible on the right and left. + This mode is not supported when stride > 1. padding[0]+padding[1] should be equal to (kernel_size - 1). + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: Union[str, int] = 0, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + self.cache_drop_size = None + if padding is None: + self._left_padding = kernel_size - 1 + self._right_padding = stride - 1 + else: + if stride != 1 and padding != kernel_size - 1: + raise ValueError("No striding allowed for non-symmetric convolutions!") + if isinstance(padding, int): + self._left_padding = padding + self._right_padding = padding + elif isinstance(padding, list) and len(padding) == 2 and padding[0] + padding[1] == kernel_size - 1: + self._left_padding = padding[0] + self._right_padding = padding[1] + else: + raise ValueError(f"Invalid padding param: {padding}!") + + self._max_cache_len = self._left_padding + + super(CausalConv1D, self).__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + + def update_cache(self, x, cache=None): + if cache is None: + new_x = F.pad(x, pad=(self._left_padding, self._right_padding)) + next_cache = cache + else: + new_x = F.pad(x, pad=(0, self._right_padding)) + new_x = torch.cat([cache, new_x], dim=-1) + if self.cache_drop_size > 0: + next_cache = new_x[:, :, : -self.cache_drop_size] + else: + next_cache = new_x + next_cache = next_cache[:, :, -cache.size(-1) :] + return new_x, next_cache + + def forward(self, x, cache=None): + x, cache = self.update_cache(x, cache=cache) + x = super().forward(x) + if cache is None: + return x + else: + return x, cache diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py new file mode 100644 index 000000000..b6f8531a8 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py @@ -0,0 +1,850 @@ +# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/modules/conformer_encoder.py; training-only code removed. + +import math +import random +from collections import OrderedDict +from dataclasses import dataclass +from typing import List, Optional, Set, Tuple + +import torch +import torch.distributed +from omegaconf import DictConfig, ListConfig, open_dict +from torch import nn + +from ..compat import ( + AccessMixin, + CacheAwareStreamingConfig, + Exportable, + NeuralModule, + StreamingEncoder, + logging, + typecheck, +) +from .causal_convs import CausalConv1D +from .conformer_modules import ConformerLayer +from .multi_head_attention import ( + LocalAttRelPositionalEncoding, + MultiHeadAttention, + PositionalEncoding, + RelPositionalEncoding, + RelPositionMultiHeadAttention, + RelPositionMultiHeadAttentionLongformer, +) +from .regularization_utils import compute_stochastic_depth_drop_probs +from .subsampling import ( + ConvSubsampling, + StackingSubsampling, + SubsamplingReductionModule, +) + +__all__ = ["ConformerEncoder"] + + +class ConformerEncoder(NeuralModule, StreamingEncoder, Exportable, AccessMixin): + """ + The encoder for ASR model of Conformer. + Based on this paper: + 'Conformer: Convolution-augmented Transformer for Speech Recognition' by Anmol Gulati et al. + https://arxiv.org/abs/2005.08100 + + Args: + feat_in (int): the size of feature channels + n_layers (int): number of layers of ConformerBlock + d_model (int): the hidden size of the model + feat_out (int): the size of the output features + Defaults to -1 (means feat_out is d_model) + subsampling (str): the method of subsampling: + choices = ['vggnet', 'striding', 'dw-striding', 'stacking', 'stacking_norm'] + Defaults to striding. + subsampling_factor (int): the subsampling factor which should be power of 2 + Defaults to 4. + subsampling_conv_chunking_factor(int): optionally, force chunk inputs (helpful for large inputs) + Should be power of 2, 1 (auto-chunking, default), or -1 (no chunking) + subsampling_conv_channels (int): the size of the convolutions in the subsampling module + Defaults to -1 which would set it to d_model. + reduction (str, Optional): the method of reduction, choices=['pooling', 'striding']. If no value + is passed, then no reduction is performed and the models runs with the original 4x subsampling. + reduction_position (int, Optional): the index of the layer to apply reduction. If -1, apply reduction + at the end. + reduction_factor (int): the reduction factor which should be either 1 or a power of 2 + Defaults to 1. + ff_expansion_factor (int): the expansion factor in feed forward layers + Defaults to 4. + self_attention_model (str): the type of the attention layer and positional encoding. + + 'rel_pos': + relative positional embedding and Transformer-XL + 'rel_pos_local_attn': + relative positional embedding and Transformer-XL with local attention using + overlapping chunks. Attention context is determined by att_context_size parameter. + 'abs_pos': + absolute positional embedding and Transformer + + Default is rel_pos. + pos_emb_max_len (int): the maximum length of positional embeddings + Defaults to 5000 + n_heads (int): number of heads in multi-headed attention layers + Defaults to 4. + att_context_size (List[Union[List[int],int]]): specifies the context sizes on each side. + Each context size should be a list of two integers like `[100, 100]`. + A list of context sizes like `[[100,100]`, `[100,50]]` can also be passed. -1 means unlimited context. + Defaults to `[-1, -1]` + att_context_probs (List[float]): a list of probabilities of each one of the att_context_size + when a list of them is passed. If not specified, uniform distribution is being used. + Defaults to None + att_context_style (str): 'regular' or 'chunked_limited'. + Defaults to 'regular' + xscaling (bool): enables scaling the inputs to the multi-headed attention layers by `sqrt(d_model)`. + Defaults to True. + untie_biases (bool): whether to not share (untie) the bias weights between layers of Transformer-XL + Defaults to True. + conv_kernel_size (int): the size of the convolutions in the convolutional modules + Defaults to 31. + conv_norm_type (str): the type of the normalization in the convolutional modules + Defaults to 'batch_norm'. + conv_context_size (list): it can be"causal" or a list of two integers + while `conv_context_size[0]+conv_context_size[1]+1==conv_kernel_size`. + `None` means `[(conv_kernel_size-1)//2`, `(conv_kernel_size-1)//2]`, and 'causal' means + `[(conv_kernel_size-1), 0]`. + Defaults to None. + conv_dual_mode (bool): specifies if convolution should be dual mode when dual_offline mode is being used. + When enables, the left half of the convolution kernel would get masked in streaming cases. + Defaults to False. + use_bias (bool): Use bias in all Linear and Conv1d layers from each ConformerLayer to improve + activation flow and stabilize training of huge models. + Defaults to True. + dropout (float): the dropout rate used in all layers except the attention layers + Defaults to 0.1. + dropout_pre_encoder (float): the dropout rate used before the encoder + Defaults to 0.1. + dropout_emb (float): the dropout rate used for the positional embeddings + Defaults to 0.1. + dropout_att (float): the dropout rate used for the attention layer + Defaults to 0.0. + stochastic_depth_drop_prob (float): if non-zero, will randomly drop + layers during training. The higher this value, the more often layers + are dropped. Defaults to 0.0. + stochastic_depth_mode (str): can be either "linear" or "uniform". If + set to "uniform", all layers have the same probability of drop. If + set to "linear", the drop probability grows linearly from 0 for the + first layer to the desired value for the final layer. Defaults to + "linear". + stochastic_depth_start_layer (int): starting layer for stochastic depth. + All layers before this will never be dropped. Note that drop + probability will be adjusted accordingly if mode is "linear" when + start layer is > 1. Defaults to 1. + global_tokens (int): number of tokens to be used for global attention. + Only relevant if self_attention_model is 'rel_pos_local_attn'. + Defaults to 0. + global_tokens_spacing (int): how far apart the global tokens are + Defaults to 1. + global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate. + Defaults to False. + use_pytorch_sdpa (bool): use torch sdpa instead of manual attention. + Defaults to False. + use_pytorch_sdpa_backends (list[str]): list of backend names to use in sdpa. + None or empty list means all backends. e.g. ["MATH"] + Defaults to None. + bypass_pre_encode: if True, skip the pre-encoder module and the `audio_signal` should be pre-encoded + embeddings. The `audio_signal` input supports two formats depending on the `bypass_pre_encode` + boolean flag. This determines the required format of the input variable `audio_signal`. + Defaults to `bypass_pre_encode=False`. `bypass_pre_encode=True` is used for the cases + where frame-level, context-independent embeddings are needed to be saved or reused. + (e.g., speaker cache in streaming speaker diarization) + sync_max_audio_length (bool): when true, performs NCCL all_reduce to allocate the same amount of memory for + positional encoding buffers on all GPUs. Disabling this setting may help with deadlocks in certain + scenarios such as model parallelism, or generally when this module is not being ran on some GPUs + as a part of the training step. + """ + + def __init__( + self, + feat_in, + n_layers, + d_model, + feat_out=-1, + causal_downsampling=False, + subsampling="striding", + subsampling_factor=4, + subsampling_conv_chunking_factor=1, + subsampling_conv_channels=-1, + reduction=None, + reduction_position=None, + reduction_factor=1, + ff_expansion_factor=4, + self_attention_model="rel_pos", + n_heads=4, + att_context_size=None, + att_context_probs=None, + att_context_style="regular", + xscaling=True, + untie_biases=True, + pos_emb_max_len=5000, + conv_kernel_size=31, + conv_norm_type="batch_norm", + conv_context_size=None, + use_bias=True, + dropout=0.1, + dropout_pre_encoder=0.1, + dropout_emb=0.1, + dropout_att=0.0, + stochastic_depth_drop_prob: float = 0.0, + stochastic_depth_mode: str = "linear", + stochastic_depth_start_layer: int = 1, + global_tokens: int = 0, + global_tokens_spacing: int = 1, + global_attn_separate: bool = False, + use_pytorch_sdpa: bool = False, + use_pytorch_sdpa_backends=None, + sync_max_audio_length: bool = True, + ): + super().__init__() + d_ff = d_model * ff_expansion_factor + self.d_model = d_model + self.n_layers = n_layers + self._feat_in = feat_in + self.att_context_style = att_context_style + self.subsampling_factor = subsampling_factor + self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor + + self.self_attention_model = self_attention_model + self.global_tokens = global_tokens + self.global_attn_separate = global_attn_separate + self.global_tokens_spacing = global_tokens_spacing + self.use_pytorch_sdpa = use_pytorch_sdpa + if use_pytorch_sdpa_backends is None: + use_pytorch_sdpa_backends = [] + self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends + self.sync_max_audio_length = sync_max_audio_length + + # Setting up the att_context_size + ( + self.att_context_size_all, + self.att_context_size, + self.att_context_probs, + self.conv_context_size, + ) = self._calc_context_sizes( + att_context_style=att_context_style, + att_context_size=att_context_size, + att_context_probs=att_context_probs, + conv_context_size=conv_context_size, + conv_kernel_size=conv_kernel_size, + ) + + if xscaling: + self.xscale = math.sqrt(d_model) + else: + self.xscale = None + + # Subsampling + if subsampling_conv_channels == -1: + subsampling_conv_channels = d_model + if subsampling and subsampling_factor > 1: + if subsampling in ["stacking", "stacking_norm"]: + # stacking_norm has an extra layer norm after stacking comparing to stacking + self.pre_encode = StackingSubsampling( + subsampling_factor=subsampling_factor, + feat_in=feat_in, + feat_out=d_model, + norm=True if subsampling == "stacking_norm" else False, + ) + else: + self.pre_encode = ConvSubsampling( + subsampling=subsampling, + subsampling_factor=subsampling_factor, + feat_in=feat_in, + feat_out=d_model, + conv_channels=subsampling_conv_channels, + subsampling_conv_chunking_factor=subsampling_conv_chunking_factor, + activation=nn.ReLU(True), + is_causal=causal_downsampling, + ) + else: + self.pre_encode = nn.Linear(feat_in, d_model) + + # Reduction + if reduction and reduction_factor > 1: + assert reduction_position >= -1 and reduction_position < n_layers + self.reduction_subsampling = SubsamplingReductionModule( + reduction=reduction, + d_model=d_model, + reduction_factor=reduction_factor, + ) + self.reduction_position = reduction_position + else: + self.reduction_subsampling = None + self.reduction_position = None + + self._feat_out = d_model + + # Biases for relative positional encoding + if not untie_biases and self_attention_model == "rel_pos": + d_head = d_model // n_heads + pos_bias_u = nn.Parameter(torch.Tensor(n_heads, d_head)) + pos_bias_v = nn.Parameter(torch.Tensor(n_heads, d_head)) + nn.init.zeros_(pos_bias_u) + nn.init.zeros_(pos_bias_v) + else: + pos_bias_u = None + pos_bias_v = None + + # Positional encodings + self.pos_emb_max_len = pos_emb_max_len + if self_attention_model == "rel_pos": + self.pos_enc = RelPositionalEncoding( + d_model=d_model, + dropout_rate=dropout_pre_encoder, + max_len=pos_emb_max_len, + xscale=self.xscale, + dropout_rate_emb=dropout_emb, + ) + elif self_attention_model == "rel_pos_local_attn": + if max(att_context_size) <= 0: + raise ValueError("When using local attention, context size must be set > 0") + self.pos_enc = LocalAttRelPositionalEncoding( + att_context_size=att_context_size, + d_model=d_model, + dropout_rate=dropout, + max_len=pos_emb_max_len, + xscale=self.xscale, + dropout_rate_emb=dropout_emb, + ) + elif self_attention_model == "abs_pos": + pos_bias_u = None + pos_bias_v = None + self.pos_enc = PositionalEncoding( + d_model=d_model, dropout_rate=dropout_pre_encoder, max_len=pos_emb_max_len, xscale=self.xscale + ) + else: + raise ValueError(f"Not valid self_attention_model: '{self_attention_model}'!") + + self.layers = nn.ModuleList() + for i in range(n_layers): + layer = ConformerLayer( + d_model=d_model, + d_ff=d_ff, + self_attention_model=self_attention_model, + global_tokens=global_tokens, + global_tokens_spacing=global_tokens_spacing, + global_attn_separate=global_attn_separate, + n_heads=n_heads, + conv_kernel_size=conv_kernel_size, + conv_norm_type=conv_norm_type, + conv_context_size=self.conv_context_size, + dropout=dropout, + dropout_att=dropout_att, + pos_bias_u=pos_bias_u, + pos_bias_v=pos_bias_v, + att_context_size=self.att_context_size, + use_bias=use_bias, + use_pytorch_sdpa=self.use_pytorch_sdpa, + use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends, + ) + self.layers.append(layer) + + if feat_out > 0 and feat_out != self._feat_out: + self.out_proj = nn.Linear(self._feat_out, feat_out) + self._feat_out = feat_out + else: + self.out_proj = None + self._feat_out = d_model + self.set_max_audio_length(self.pos_emb_max_len) + self.use_pad_mask = True + + self.setup_streaming_params() + self.export_cache_support = False + + self.layer_drop_probs = compute_stochastic_depth_drop_probs( + len(self.layers), stochastic_depth_drop_prob, stochastic_depth_mode, stochastic_depth_start_layer + ) + # will be set in self.forward() if defined in AccessMixin config + self.interctc_capture_at_layers = None + + @typecheck() + def forward( + self, + audio_signal, + length, + cache_last_channel=None, + cache_last_time=None, + cache_last_channel_len=None, + bypass_pre_encode=False, + ): + """ + Forward function for the ConformerEncoder accepting an audio signal and its corresponding length. + The `audio_signal` input supports two formats depending on the `bypass_pre_encode` boolean flag. + This determines the required format of the input variable `audio_signal`: + (1) bypass_pre_encode = False (default): + `audio_signal` must be a tensor containing audio features. + Shape: (batch, self._feat_in, n_frames) + (2) bypass_pre_encode = True: + `audio_signal` must be a tensor containing pre-encoded embeddings. + Shape: (batch, n_frame, self.d_model) + """ + if not bypass_pre_encode and audio_signal.shape[-2] != self._feat_in: + raise ValueError( + f"If bypass_pre_encode is False, audio_signal should have shape " + f"(batch, {self._feat_in}, n_frame) but got last dimension {audio_signal.shape[-2]}." + ) + if bypass_pre_encode and audio_signal.shape[-1] != self.d_model: + raise ValueError( + f"If bypass_pre_encode is True, audio_signal should have shape " + f"(batch, n_frame, {self.d_model}) but got last dimension {audio_signal.shape[-1]}." + ) + + if bypass_pre_encode: + self.update_max_seq_length( + seq_length=audio_signal.size(2) * self.subsampling_factor, device=audio_signal.device + ) + else: + self.update_max_seq_length(seq_length=audio_signal.size(2), device=audio_signal.device) + return self.forward_internal( + audio_signal, + length, + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + bypass_pre_encode=bypass_pre_encode, + ) + + def forward_internal( + self, + audio_signal, + length, + cache_last_channel=None, + cache_last_time=None, + cache_last_channel_len=None, + bypass_pre_encode=False, + ): + """ + The `audio_signal` input supports two formats depending on the `bypass_pre_encode` boolean flag. + This determines the required format of the input variable `audio_signal`: + (1) bypass_pre_encode = False (default): + `audio_signal` must be a tensor containing audio features. + Shape: (batch, self._feat_in, n_frames) + (2) bypass_pre_encode = True: + `audio_signal` must be a tensor containing pre-encoded embeddings. + Shape: (batch, n_frame, self.d_model) + + `bypass_pre_encode=True` is used in cases where frame-level, context-independent embeddings are + needed to be saved or reused (e.g., speaker cache in streaming speaker diarization). + """ + if length is None: + length = audio_signal.new_full( + (audio_signal.size(0),), audio_signal.size(-1), dtype=torch.int64, device=audio_signal.device + ) + + # select a random att_context_size with the distribution specified by att_context_probs during training + # for non-validation cases like test, validation or inference, it uses the first mode in self.att_context_size + if self.training and len(self.att_context_size_all) > 1: + cur_att_context_size = random.choices(self.att_context_size_all, weights=self.att_context_probs)[0] + else: + cur_att_context_size = self.att_context_size + + if not bypass_pre_encode: + audio_signal = torch.transpose(audio_signal, 1, 2) + + if isinstance(self.pre_encode, nn.Linear): + audio_signal = self.pre_encode(audio_signal) + else: + audio_signal, length = self.pre_encode(x=audio_signal, lengths=length) + length = length.to(torch.int64) + # `self.streaming_cfg` is set by setup_streaming_cfg(), called in the init + if self.streaming_cfg.drop_extra_pre_encoded > 0 and cache_last_channel is not None: + audio_signal = audio_signal[:, self.streaming_cfg.drop_extra_pre_encoded :, :] + length = (length - self.streaming_cfg.drop_extra_pre_encoded).clamp(min=0) + + if self.reduction_position is not None and cache_last_channel is not None: + raise ValueError("Caching with reduction feature is not supported yet!") + + max_audio_length = audio_signal.size(1) + if cache_last_channel is not None: + cache_len = self.streaming_cfg.last_channel_cache_size + cache_keep_size = max_audio_length - self.streaming_cfg.cache_drop_size + max_audio_length = max_audio_length + cache_len + padding_length = length + cache_len + offset = torch.neg(cache_last_channel_len) + cache_len + else: + padding_length = length + cache_last_channel_next = None + cache_len = 0 + offset = None + + audio_signal, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len) + + # Create the self-attention and padding masks + pad_mask, att_mask = self._create_masks( + att_context_size=cur_att_context_size, + padding_length=padding_length, + max_audio_length=max_audio_length, + offset=offset, + device=audio_signal.device, + ) + + if cache_last_channel is not None: + pad_mask = pad_mask[:, cache_len:] + if att_mask is not None: + att_mask = att_mask[:, cache_len:] + # Convert caches from the tensor to list + cache_last_time_next = [] + cache_last_channel_next = [] + + for lth, (drop_prob, layer) in enumerate(zip(self.layer_drop_probs, self.layers)): + original_signal = audio_signal + if cache_last_channel is not None: + cache_last_channel_cur = cache_last_channel[lth] + cache_last_time_cur = cache_last_time[lth] + else: + cache_last_channel_cur = None + cache_last_time_cur = None + audio_signal = layer( + x=audio_signal, + att_mask=att_mask, + pos_emb=pos_emb, + pad_mask=pad_mask, + cache_last_channel=cache_last_channel_cur, + cache_last_time=cache_last_time_cur, + ) + + if cache_last_channel_cur is not None: + (audio_signal, cache_last_channel_cur, cache_last_time_cur) = audio_signal + cache_last_channel_next.append(cache_last_channel_cur) + cache_last_time_next.append(cache_last_time_cur) + + # applying stochastic depth logic from https://arxiv.org/abs/2102.03216 + if self.training and drop_prob > 0.0: + should_drop = torch.rand(1) < drop_prob + # adjusting to match expectation + if should_drop: + # that's not efficient, but it's hard to implement distributed + # version of dropping layers without deadlock or random seed meddling + # so multiplying the signal by 0 to ensure all weights get gradients + audio_signal = audio_signal * 0.0 + original_signal + else: + # not doing this operation if drop prob is 0 as it's identity in that case + audio_signal = (audio_signal - original_signal) / (1.0 - drop_prob) + original_signal + + if self.reduction_position == lth: + audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length) + max_audio_length = audio_signal.size(1) + # Don't update the audio_signal here because then it will again scale the audio_signal + # and cause an increase in the WER + _, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len) + pad_mask, att_mask = self._create_masks( + att_context_size=cur_att_context_size, + padding_length=length, + max_audio_length=max_audio_length, + offset=offset, + device=audio_signal.device, + ) + + # saving tensors if required for interctc loss + if self.is_access_enabled(getattr(self, "model_guid", None)): + if self.interctc_capture_at_layers is None: + self.interctc_capture_at_layers = self.access_cfg.get("interctc", {}).get("capture_layers", []) + if lth in self.interctc_capture_at_layers: + lth_audio_signal = audio_signal + if self.out_proj is not None: + lth_audio_signal = self.out_proj(audio_signal) + # shape is the same as the shape of audio_signal output, i.e. [B, D, T] + self.register_accessible_tensor( + name=f"interctc/layer_output_{lth}", tensor=torch.transpose(lth_audio_signal, 1, 2) + ) + self.register_accessible_tensor(name=f"interctc/layer_length_{lth}", tensor=length) + + if self.out_proj is not None: + audio_signal = self.out_proj(audio_signal) + + # Reduction + if self.reduction_position == -1: + audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length) + + audio_signal = torch.transpose(audio_signal, 1, 2) + length = length.to(dtype=torch.int64) + + if cache_last_channel is not None: + cache_last_channel_next = torch.stack(cache_last_channel_next, dim=0) + cache_last_time_next = torch.stack(cache_last_time_next, dim=0) + return ( + audio_signal, + length, + cache_last_channel_next, + cache_last_time_next, + torch.clamp(cache_last_channel_len + cache_keep_size, max=cache_len), + ) + else: + return audio_signal, length + + def update_max_seq_length(self, seq_length: int, device): + """ + Updates the maximum sequence length for the model. + + Args: + seq_length (int): New maximum sequence length. + device (torch.device): Device to use for computations. + """ + # Find global max audio length across all nodes + if self.sync_max_audio_length and torch.distributed.is_initialized(): + global_max_len = torch.tensor([seq_length], dtype=torch.float32, device=device) + + # Update across all ranks in the distributed system + torch.distributed.all_reduce(global_max_len, op=torch.distributed.ReduceOp.MAX) + + seq_length = global_max_len.int().item() + + if seq_length > self.max_audio_length: + self.set_max_audio_length(seq_length) + + def set_max_audio_length(self, max_audio_length): + """ + Sets maximum input length. + Pre-calculates internal seq_range mask. + + Args: + max_audio_length (int): New maximum sequence length. + """ + self.max_audio_length = max_audio_length + device = next(self.parameters()).device + dtype = next(self.parameters()).dtype + self.pos_enc.extend_pe(max_audio_length, device, dtype) + + def _create_masks(self, att_context_size, padding_length, max_audio_length, offset, device): + if self.self_attention_model != "rel_pos_local_attn": + att_mask = torch.ones(1, max_audio_length, max_audio_length, dtype=torch.bool, device=device) + + if self.att_context_style == "regular": + if att_context_size[0] >= 0: + att_mask = att_mask.triu(diagonal=-att_context_size[0]) + if att_context_size[1] >= 0: + att_mask = att_mask.tril(diagonal=att_context_size[1]) + elif self.att_context_style == "chunked_limited": + # When right context is unlimited, just the left side of the masking need to get updated + if att_context_size[1] == -1: + if att_context_size[0] >= 0: + att_mask = att_mask.triu(diagonal=-att_context_size[0]) + else: + chunk_size = att_context_size[1] + 1 + # left_chunks_num specifies the number of chunks to be visible by each chunk on the left side + if att_context_size[0] >= 0: + left_chunks_num = att_context_size[0] // chunk_size + else: + left_chunks_num = 10000 + + chunk_idx = torch.arange(0, max_audio_length, dtype=torch.int, device=att_mask.device) + chunk_idx = torch.div(chunk_idx, chunk_size, rounding_mode="trunc") + diff_chunks = chunk_idx.unsqueeze(1) - chunk_idx.unsqueeze(0) + chunked_limited_mask = torch.logical_and( + torch.le(diff_chunks, left_chunks_num), torch.ge(diff_chunks, 0) + ) + att_mask = torch.logical_and(att_mask, chunked_limited_mask.unsqueeze(0)) + else: + att_mask = None + + # pad_mask is the masking to be used to ignore paddings + pad_mask = torch.arange(0, max_audio_length, device=device).expand( + padding_length.size(0), -1 + ) < padding_length.unsqueeze(-1) + + if offset is not None: + pad_mask_off = torch.arange(0, max_audio_length, device=device).expand( + padding_length.size(0), -1 + ) >= offset.unsqueeze(-1) + pad_mask = pad_mask_off.logical_and(pad_mask) + + if att_mask is not None: + # pad_mask_for_att_mask is the mask which helps to ignore paddings + pad_mask_for_att_mask = pad_mask.unsqueeze(1).repeat([1, max_audio_length, 1]) + pad_mask_for_att_mask = torch.logical_and(pad_mask_for_att_mask, pad_mask_for_att_mask.transpose(1, 2)) + # att_mask is the masking to be used by the MHA layers to ignore the tokens not supposed to be visible + att_mask = att_mask[:, :max_audio_length, :max_audio_length] + # paddings should also get ignored, so pad_mask_for_att_mask is used to ignore their corresponding scores + att_mask = torch.logical_and(pad_mask_for_att_mask, att_mask.to(pad_mask_for_att_mask.device)) + att_mask = ~att_mask + + pad_mask = ~pad_mask + return pad_mask, att_mask + + def _calc_context_sizes( + self, att_context_size, att_context_probs, att_context_style, conv_context_size, conv_kernel_size + ): + # convert att_context_size to a standard list of lists + if att_context_size: + att_context_size_all = list(att_context_size) + if isinstance(att_context_size_all[0], int): + att_context_size_all = [att_context_size_all] + for i, att_cs in enumerate(att_context_size_all): + if isinstance(att_cs, ListConfig): + att_context_size_all[i] = list(att_cs) + if att_context_style == "chunked_limited": + if att_cs[0] > 0 and att_cs[0] % (att_cs[1] + 1) > 0: + raise ValueError(f"att_context_size[{i}][0] % (att_context_size[{i}][1] + 1) should be zero!") + if att_cs[1] < 0 and len(att_context_size_all) <= 1: + raise ValueError( + f"Right context (att_context_size[{i}][1]) can not be unlimited for chunked_limited style!" + ) + else: + att_context_size_all = [[-1, -1]] + + if att_context_probs: + if len(att_context_probs) != len(att_context_size_all): + raise ValueError("The size of the att_context_probs should be the same as att_context_size.") + att_context_probs = list(att_context_probs) + if sum(att_context_probs) != 1: + raise ValueError("The sum of numbers in att_context_probs should be equal to one to be a distribution.") + else: + att_context_probs = [1.0 / len(att_context_size_all)] * len(att_context_size_all) + + if conv_context_size is not None: + if isinstance(conv_context_size, ListConfig): + conv_context_size = list(conv_context_size) + if not isinstance(conv_context_size, list) and not isinstance(conv_context_size, str): + raise ValueError( + "Invalid conv_context_size! It should be the string 'causal' or a list of two integers." + ) + if conv_context_size == "causal": + conv_context_size = [conv_kernel_size - 1, 0] + else: + if conv_context_size[0] + conv_context_size[1] + 1 != conv_kernel_size: + raise ValueError(f"Invalid conv_context_size: {self.conv_context_size}!") + else: + conv_context_size = [(conv_kernel_size - 1) // 2, (conv_kernel_size - 1) // 2] + return att_context_size_all, att_context_size_all[0], att_context_probs, conv_context_size + + def setup_streaming_params( + self, + chunk_size: int = None, + shift_size: int = None, + left_chunks: int = None, + att_context_size: list = None, + max_context: int = 10000, + ): + """ + This function sets the needed values and parameters to perform streaming. + The configuration would be stored in self.streaming_cfg. + The streaming configuration is needed to simulate streaming inference. + + Args: + chunk_size (int): overrides the chunk size + shift_size (int): overrides the shift size for chunks + left_chunks (int): overrides the number of left chunks visible to each chunk + max_context (int): the value used for the cache size of last_channel layers + if left context is set to infinity (-1) + Defaults to -1 (means feat_out is d_model) + """ + streaming_cfg = CacheAwareStreamingConfig() + + # When att_context_size is not specified, it uses the default_att_context_size + if att_context_size is None: + att_context_size = self.att_context_size + + if chunk_size is not None: + if chunk_size < 1: + raise ValueError("chunk_size needs to be a number larger or equal to one.") + lookahead_steps = chunk_size - 1 + streaming_cfg.cache_drop_size = chunk_size - shift_size + elif self.att_context_style == "chunked_limited": + lookahead_steps = att_context_size[1] + streaming_cfg.cache_drop_size = 0 + elif self.att_context_style == "regular": + lookahead_steps = att_context_size[1] * self.n_layers + self.conv_context_size[1] * self.n_layers + streaming_cfg.cache_drop_size = lookahead_steps + else: + streaming_cfg.cache_drop_size = 0 + lookahead_steps = None + + if chunk_size is None: + streaming_cfg.last_channel_cache_size = att_context_size[0] if att_context_size[0] >= 0 else max_context + else: + if left_chunks is None: + streaming_cfg.last_channel_cache_size = att_context_size[0] if att_context_size[0] >= 0 else max_context + logging.warning( + f"left_chunks is not set. Setting it to default: {streaming_cfg.last_channel_cache_size}." + ) + else: + streaming_cfg.last_channel_cache_size = left_chunks * chunk_size + + if hasattr(self.pre_encode, "get_sampling_frames"): + sampling_frames = self.pre_encode.get_sampling_frames() + else: + sampling_frames = 0 + + if isinstance(sampling_frames, list): + streaming_cfg.chunk_size = [ + sampling_frames[0] + self.subsampling_factor * lookahead_steps, + sampling_frames[1] + self.subsampling_factor * lookahead_steps, + ] + else: + streaming_cfg.chunk_size = sampling_frames * (1 + lookahead_steps) + + if isinstance(sampling_frames, list): + streaming_cfg.shift_size = [ + sampling_frames[0] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size), + sampling_frames[1] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size), + ] + else: + streaming_cfg.shift_size = sampling_frames * (1 + lookahead_steps - streaming_cfg.cache_drop_size) + + if isinstance(streaming_cfg.shift_size, list): + streaming_cfg.valid_out_len = ( + streaming_cfg.shift_size[1] - sampling_frames[1] + ) // self.subsampling_factor + 1 + else: + streaming_cfg.valid_out_len = streaming_cfg.shift_size // self.subsampling_factor + + if hasattr(self.pre_encode, "get_streaming_cache_size"): + streaming_cfg.pre_encode_cache_size = self.pre_encode.get_streaming_cache_size() + else: + streaming_cfg.pre_encode_cache_size = 0 + + if isinstance(streaming_cfg.pre_encode_cache_size, list): + if streaming_cfg.pre_encode_cache_size[1] >= 1: + streaming_cfg.drop_extra_pre_encoded = ( + 1 + (streaming_cfg.pre_encode_cache_size[1] - 1) // self.subsampling_factor + ) + else: + streaming_cfg.drop_extra_pre_encoded = 0 + else: + streaming_cfg.drop_extra_pre_encoded = streaming_cfg.pre_encode_cache_size // self.subsampling_factor + + for m in self.layers.modules(): + if hasattr(m, "_max_cache_len"): + if isinstance(m, MultiHeadAttention): + m.cache_drop_size = streaming_cfg.cache_drop_size + if isinstance(m, CausalConv1D): + m.cache_drop_size = streaming_cfg.cache_drop_size + + self.streaming_cfg = streaming_cfg + + def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int): + """ + Update the conv_chunking_factor (int) + Default is 1 (auto) + Set it to -1 (disabled) or to a specific value (power of 2) if you OOM in the conv subsampling layers + + + Args: + subsampling_conv_chunking_factor (int) + """ + + if not hasattr(self.pre_encode, "change_subsampling_conv_chunking_factor"): + logging.info("Model pre_encoder doesn't have a change_subsampling_conv_chunking_factor method ") + return + + self.pre_encode.change_subsampling_conv_chunking_factor( + subsampling_conv_chunking_factor=subsampling_conv_chunking_factor + ) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py new file mode 100644 index 000000000..4875ff734 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py @@ -0,0 +1,372 @@ +# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/conformer_modules.py; training-only code removed. + +import torch +from torch import nn as nn +from torch.nn import LayerNorm + +from ..compat import AccessMixin, AttentionAdapterModuleMixin, FusedBatchNorm1d, Swish, activation_registry +from .causal_convs import CausalConv1D +from .multi_head_attention import ( + MultiHeadAttention, + RelPositionMultiHeadAttention, + RelPositionMultiHeadAttentionLongformer, +) + +__all__ = ["ConformerConvolution", "ConformerFeedForward", "ConformerLayer"] + + +class ConformerLayer(torch.nn.Module, AttentionAdapterModuleMixin, AccessMixin): + """A single block of the Conformer encoder. + + Args: + d_model (int): input dimension of MultiheadAttentionMechanism and PositionwiseFeedForward + d_ff (int): hidden dimension of PositionwiseFeedForward + self_attention_model (str): type of the attention layer and positional encoding + 'rel_pos': relative positional embedding and Transformer-XL + 'rel_pos_local_attn': relative positional embedding and Transformer-XL with local attention using + overlapping chunks. Attention context is determined by att_context_size parameter. + 'abs_pos': absolute positional embedding and Transformer + Default is rel_pos. + global_tokens (int): number of tokens to be used for global attention. + Only relevant if self_attention_model is 'rel_pos_local_attn'. + Defaults to 0. + global_tokens_spacing (int): how far apart the global tokens are + Defaults to 1. + global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate. + Defaults to False. + n_heads (int): number of heads for multi-head attention + conv_kernel_size (int): kernel size for depthwise convolution in convolution module + dropout (float): dropout probabilities for linear layers + dropout_att (float): dropout probabilities for attention distributions + use_bias (bool): Apply bias to all Linear and Conv1d layers from each ConformerLayer to improve activation flow and stabilize training of huge models. + Defaults to True. + """ + + def __init__( + self, + d_model, + d_ff, + self_attention_model="rel_pos", + global_tokens=0, + global_tokens_spacing=1, + global_attn_separate=False, + n_heads=4, + conv_kernel_size=31, + conv_norm_type="batch_norm", + conv_context_size=None, + dropout=0.1, + dropout_att=0.1, + pos_bias_u=None, + pos_bias_v=None, + att_context_size=[-1, -1], + use_bias=True, + use_pytorch_sdpa=False, + use_pytorch_sdpa_backends=None, + ): + super(ConformerLayer, self).__init__() + + self.use_pytorch_sdpa = use_pytorch_sdpa + if use_pytorch_sdpa_backends is None: + use_pytorch_sdpa_backends = [] + self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends + self.self_attention_model = self_attention_model + self.n_heads = n_heads + self.fc_factor = 0.5 + + # first feed forward module + self.norm_feed_forward1 = LayerNorm(d_model) + self.feed_forward1 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout, use_bias=use_bias) + + # convolution module + self.norm_conv = LayerNorm(d_model) + self.conv = ConformerConvolution( + d_model=d_model, + kernel_size=conv_kernel_size, + norm_type=conv_norm_type, + conv_context_size=conv_context_size, + use_bias=use_bias, + ) + + # multi-headed self-attention module + self.norm_self_att = LayerNorm(d_model) + MHA_max_cache_len = att_context_size[0] + + if self_attention_model == "rel_pos": + self.self_attn = RelPositionMultiHeadAttention( + n_head=n_heads, + n_feat=d_model, + dropout_rate=dropout_att, + pos_bias_u=pos_bias_u, + pos_bias_v=pos_bias_v, + max_cache_len=MHA_max_cache_len, + use_bias=use_bias, + use_pytorch_sdpa=self.use_pytorch_sdpa, + use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends, + ) + elif self_attention_model == "rel_pos_local_attn": + self.self_attn = RelPositionMultiHeadAttentionLongformer( + n_head=n_heads, + n_feat=d_model, + dropout_rate=dropout_att, + pos_bias_u=pos_bias_u, + pos_bias_v=pos_bias_v, + max_cache_len=MHA_max_cache_len, + att_context_size=att_context_size, + global_tokens=global_tokens, + global_tokens_spacing=global_tokens_spacing, + global_attn_separate=global_attn_separate, + use_bias=use_bias, + ) + elif self_attention_model == "abs_pos": + self.self_attn = MultiHeadAttention( + n_head=n_heads, + n_feat=d_model, + dropout_rate=dropout_att, + max_cache_len=MHA_max_cache_len, + use_bias=use_bias, + use_pytorch_sdpa=self.use_pytorch_sdpa, + use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends, + ) + else: + raise ValueError( + f"'{self_attention_model}' is not not a valid value for 'self_attention_model', " + f"valid values can be from ['rel_pos', 'rel_pos_local_attn', 'abs_pos']" + ) + + # second feed forward module + self.norm_feed_forward2 = LayerNorm(d_model) + self.feed_forward2 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout, use_bias=use_bias) + + self.dropout = nn.Dropout(dropout) + self.norm_out = LayerNorm(d_model) + + def forward(self, x, att_mask=None, pos_emb=None, pad_mask=None, cache_last_channel=None, cache_last_time=None): + """ + Args: + x (torch.Tensor): input signals (B, T, d_model) + att_mask (torch.Tensor): attention masks(B, T, T) + pos_emb (torch.Tensor): (L, 1, d_model) + pad_mask (torch.tensor): padding mask + cache_last_channel (torch.tensor) : cache for MHA layers (B, T_cache, d_model) + cache_last_time (torch.tensor) : cache for convolutional layers (B, d_model, T_cache) + Returns: + x (torch.Tensor): (B, T, d_model) + cache_last_channel (torch.tensor) : next cache for MHA layers (B, T_cache, d_model) + cache_last_time (torch.tensor) : next cache for convolutional layers (B, d_model, T_cache) + """ + residual = x + x = self.norm_feed_forward1(x) + x = self.feed_forward1(x) + residual = residual + self.dropout(x) * self.fc_factor + + x = self.norm_self_att(residual) + if self.self_attention_model == "rel_pos": + x = self.self_attn(query=x, key=x, value=x, mask=att_mask, pos_emb=pos_emb, cache=cache_last_channel) + elif self.self_attention_model == "rel_pos_local_attn": + x = self.self_attn(query=x, key=x, value=x, pad_mask=pad_mask, pos_emb=pos_emb, cache=cache_last_channel) + elif self.self_attention_model == "abs_pos": + x = self.self_attn(query=x, key=x, value=x, mask=att_mask, cache=cache_last_channel) + else: + x = None + + if x is not None and cache_last_channel is not None: + (x, cache_last_channel) = x + + residual = residual + self.dropout(x) + + if self.is_adapter_available(): + # Call the MHA adapters + pack_input = { + "x": residual, + "loc": "mha", + "att_mask": att_mask, + "pos_emb": pos_emb, + } + pack_input = self.forward_enabled_adapters(pack_input) + residual = pack_input["x"] + + x = self.norm_conv(residual) + x = self.conv(x, pad_mask=pad_mask, cache=cache_last_time) + if cache_last_time is not None: + (x, cache_last_time) = x + residual = residual + self.dropout(x) + + x = self.norm_feed_forward2(residual) + x = self.feed_forward2(x) + residual = residual + self.dropout(x) * self.fc_factor + + x = self.norm_out(residual) + + if self.is_adapter_available(): + # Call the adapters + pack_input = { + "x": x, + "loc": "post", + } + pack_input = self.forward_enabled_adapters(pack_input) + x = pack_input["x"] + + if self.is_access_enabled(getattr(self, "model_guid", None)) and self.access_cfg.get( + "save_encoder_tensors", False + ): + self.register_accessible_tensor(name="encoder", tensor=x) + if cache_last_channel is None: + return x + else: + return x, cache_last_channel, cache_last_time + + +class ConformerConvolution(nn.Module): + """The convolution module for the Conformer model. + Args: + d_model (int): hidden dimension + kernel_size (int): kernel size for depthwise convolution + pointwise_activation (str): name of the activation function to be used for the pointwise conv. + Note that Conformer uses a special key `glu_` which is treated as the original default from + the paper. + use_bias (bool): Use bias in all Linear and Conv1d layers improve activation flow and stabilize training of huge models. + Defaults to True + """ + + def __init__( + self, + d_model, + kernel_size, + norm_type="batch_norm", + conv_context_size=None, + pointwise_activation="glu_", + use_bias=True, + ): + super(ConformerConvolution, self).__init__() + assert (kernel_size - 1) % 2 == 0 + self.d_model = d_model + self.kernel_size = kernel_size + self.norm_type = norm_type + self.use_bias = use_bias + + if conv_context_size is None: + conv_context_size = (kernel_size - 1) // 2 + + if pointwise_activation in activation_registry: + self.pointwise_activation = activation_registry[pointwise_activation]() + dw_conv_input_dim = d_model * 2 + + if hasattr(self.pointwise_activation, "inplace"): + self.pointwise_activation.inplace = True + else: + self.pointwise_activation = pointwise_activation + dw_conv_input_dim = d_model + + self.pointwise_conv1 = nn.Conv1d( + in_channels=d_model, + out_channels=d_model * 2, + kernel_size=1, + stride=1, + padding=0, + bias=self.use_bias, + ) + + self.depthwise_conv = CausalConv1D( + in_channels=dw_conv_input_dim, + out_channels=dw_conv_input_dim, + kernel_size=kernel_size, + stride=1, + padding=conv_context_size, + groups=dw_conv_input_dim, + bias=self.use_bias, + ) + + if norm_type == "batch_norm": + self.batch_norm = nn.BatchNorm1d(dw_conv_input_dim) + elif norm_type == "instance_norm": + self.batch_norm = nn.InstanceNorm1d(dw_conv_input_dim) + elif norm_type == "layer_norm": + self.batch_norm = nn.LayerNorm(dw_conv_input_dim) + elif norm_type == "fused_batch_norm": + self.batch_norm = FusedBatchNorm1d(dw_conv_input_dim) + elif norm_type.startswith("group_norm"): + num_groups = int(norm_type.replace("group_norm", "")) + self.batch_norm = nn.GroupNorm(num_groups=num_groups, num_channels=d_model) + else: + raise ValueError(f"conv_norm_type={norm_type} is not valid!") + + self.activation = Swish() + self.pointwise_conv2 = nn.Conv1d( + in_channels=dw_conv_input_dim, + out_channels=d_model, + kernel_size=1, + stride=1, + padding=0, + bias=self.use_bias, + ) + + def forward(self, x, pad_mask=None, cache=None): + x = x.transpose(1, 2) + x = self.pointwise_conv1(x) + + # Compute the activation function or use GLU for original Conformer + if self.pointwise_activation == "glu_": + x = nn.functional.glu(x, dim=1) + else: + x = self.pointwise_activation(x) + + if pad_mask is not None: + x = x.masked_fill(pad_mask.unsqueeze(1), 0.0) + + x = self.depthwise_conv(x, cache=cache) + if cache is not None: + x, cache = x + + if self.norm_type == "layer_norm": + x = x.transpose(1, 2) + x = self.batch_norm(x) + x = x.transpose(1, 2) + else: + x = self.batch_norm(x) + + x = self.activation(x) + x = self.pointwise_conv2(x) + x = x.transpose(1, 2) + if cache is None: + return x + else: + return x, cache + + +class ConformerFeedForward(nn.Module): + """ + feed-forward module of Conformer model. + use_bias (bool): Apply bias to all Linear and Conv1d layers improve activation flow and stabilize training of huge models. + """ + + def __init__(self, d_model, d_ff, dropout, activation=Swish(), use_bias=True): + super(ConformerFeedForward, self).__init__() + self.d_model = d_model + self.d_ff = d_ff + self.use_bias = use_bias + self.linear1 = nn.Linear(d_model, d_ff, bias=self.use_bias) + self.activation = activation + self.dropout = nn.Dropout(p=dropout) + self.linear2 = nn.Linear(d_ff, d_model, bias=self.use_bias) + + def forward(self, x): + x = self.linear1(x) + x = self.activation(x) + x = self.dropout(x) + x = self.linear2(x) + return x diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py new file mode 100644 index 000000000..2261d739e --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py @@ -0,0 +1,633 @@ +# Copyright (c) 2018 Ryan Leary +# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# This file contains code artifacts adapted from https://github.com/ryanleary/patter +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/preprocessing/features.py; training-only code removed. +# Removed: WaveformFeaturizer / FeaturizerFactory (data-loading only, required perturb/segment). + +import math +import random +from typing import Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn + +from ..compat import logging + +try: + import torchaudio + + HAVE_TORCHAUDIO = True +except ModuleNotFoundError: + HAVE_TORCHAUDIO = False + + +CONSTANT = 1e-5 + + +def normalize_batch(x, seq_len, normalize_type): + x_mean = None + x_std = None + if normalize_type == "per_feature": + batch_size = x.shape[0] + max_time = x.shape[2] + + # When doing stream capture to a graph, item() is not allowed + # because it calls cudaStreamSynchronize(). Therefore, we are + # sacrificing some error checking when running with cuda graphs. + if ( + torch.cuda.is_available() + and not torch.cuda.is_current_stream_capturing() + and torch.any(seq_len == 1).item() + ): + raise ValueError( + "normalize_batch with `per_feature` normalize_type received a tensor of length 1. This will result " + "in torch.std() returning nan. Make sure your audio length has enough samples for a single " + "feature (ex. at least `hop_length` for Mel Spectrograms)." + ) + time_steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(batch_size, max_time) + valid_mask = time_steps < seq_len.unsqueeze(1) + x_mean_numerator = torch.where(valid_mask.unsqueeze(1), x, 0.0).sum(axis=2) + x_mean_denominator = valid_mask.sum(axis=1) + x_mean = x_mean_numerator / x_mean_denominator.unsqueeze(1) + + # Subtract 1 in the denominator to correct for the bias. + x_std = torch.sqrt( + torch.sum(torch.where(valid_mask.unsqueeze(1), x - x_mean.unsqueeze(2), 0.0) ** 2, axis=2) + / (x_mean_denominator.unsqueeze(1) - 1.0) + ) + # make sure x_std is not zero + x_std += CONSTANT + return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2), x_mean, x_std + elif normalize_type == "all_features": + x_mean = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device) + x_std = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device) + for i in range(x.shape[0]): + x_mean[i] = x[i, :, : seq_len[i].item()].mean() + x_std[i] = x[i, :, : seq_len[i].item()].std() + # make sure x_std is not zero + x_std += CONSTANT + return (x - x_mean.view(-1, 1, 1)) / x_std.view(-1, 1, 1), x_mean, x_std + elif "fixed_mean" in normalize_type and "fixed_std" in normalize_type: + x_mean = torch.tensor(normalize_type["fixed_mean"], device=x.device) + x_std = torch.tensor(normalize_type["fixed_std"], device=x.device) + return ( + (x - x_mean.view(x.shape[0], x.shape[1]).unsqueeze(2)) / x_std.view(x.shape[0], x.shape[1]).unsqueeze(2), + x_mean, + x_std, + ) + else: + return x, x_mean, x_std + + +def splice_frames(x, frame_splicing): + """Stacks frames together across feature dim + + input is batch_size, feature_dim, num_frames + output is batch_size, feature_dim*frame_splicing, num_frames + + """ + seq = [x] + for n in range(1, frame_splicing): + seq.append(torch.cat([x[:, :, :n], x[:, :, n:]], dim=2)) + return torch.cat(seq, dim=1) + + +@torch.jit.script_if_tracing +def make_seq_mask_like( + lengths: torch.Tensor, like: torch.Tensor, time_dim: int = -1, valid_ones: bool = True +) -> torch.Tensor: + """ + + Args: + lengths: Tensor with shape [B] containing the sequence length of each batch element + like: The mask will contain the same number of dimensions as this Tensor, and will have the same max + length in the time dimension of this Tensor. + time_dim: Time dimension of the `shape_tensor` and the resulting mask. Zero-based. + valid_ones: If True, valid tokens will contain value `1` and padding will be `0`. Else, invert. + + Returns: + A :class:`torch.Tensor` containing 1's and 0's for valid and invalid tokens, respectively, if `valid_ones`, else + vice-versa. Mask will have the same number of dimensions as `like`. Batch and time dimensions will match + the `like`. All other dimensions will be singletons. E.g., if `like.shape == [3, 4, 5]` and + `time_dim == -1', mask will have shape `[3, 1, 5]`. + """ + # Mask with shape [B, T] + mask = torch.arange(like.shape[time_dim], device=like.device).repeat(lengths.shape[0], 1).lt(lengths.view(-1, 1)) + # [B, T] -> [B, *, T] where * is any number of singleton dimensions to expand to like tensor + for _ in range(like.dim() - mask.dim()): + mask = mask.unsqueeze(1) + # If needed, transpose time dim + if time_dim != -1 and time_dim != mask.dim() - 1: + mask = mask.transpose(-1, time_dim) + # Maybe invert the padded vs. valid token values + if not valid_ones: + mask = ~mask + return mask + + +class FilterbankFeatures(nn.Module): + """Featurizer that converts wavs to Mel Spectrograms. + See AudioToMelSpectrogramPreprocessor for args. + """ + + def __init__( + self, + sample_rate=16000, + n_window_size=320, + n_window_stride=160, + window="hann", + normalize="per_feature", + n_fft=None, + preemph=0.97, + nfilt=64, + lowfreq=0, + highfreq=None, + log=True, + log_zero_guard_type="add", + log_zero_guard_value=2**-24, + dither=CONSTANT, + pad_to=16, + max_duration=16.7, + frame_splicing=1, + exact_pad=False, + pad_value=0, + mag_power=2.0, + use_grads=False, + rng=None, + nb_augmentation_prob=0.0, + nb_max_freq=4000, + mel_norm="slaney", + stft_exact_pad=False, # Deprecated arguments; kept for config compatibility + stft_conv=False, # Deprecated arguments; kept for config compatibility + ): + super().__init__() + if stft_conv or stft_exact_pad: + logging.warning( + "Using torch_stft is deprecated and has been removed. The values have been forcibly set to False " + "for FilterbankFeatures and AudioToMelSpectrogramPreprocessor. Please set exact_pad to True " + "as needed." + ) + if exact_pad and n_window_stride % 2 == 1: + raise NotImplementedError( + f"{self} received exact_pad == True, but hop_size was odd. If audio_length % hop_size == 0. Then the " + "returned spectrogram would not be of length audio_length // hop_size. Please use an even hop_size." + ) + self.log_zero_guard_value = log_zero_guard_value + if ( + n_window_size is None + or n_window_stride is None + or not isinstance(n_window_size, int) + or not isinstance(n_window_stride, int) + or n_window_size <= 0 + or n_window_stride <= 0 + ): + raise ValueError( + f"{self} got an invalid value for either n_window_size or n_window_stride. Both must be positive ints." + ) + logging.info(f"PADDING: {pad_to}") + + self.win_length = n_window_size + self.hop_length = n_window_stride + self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length)) + self.stft_pad_amount = (self.n_fft - self.hop_length) // 2 if exact_pad else None + self.exact_pad = exact_pad + + if exact_pad: + logging.info("STFT using exact pad") + torch_windows = { + "hann": torch.hann_window, + "hamming": torch.hamming_window, + "blackman": torch.blackman_window, + "bartlett": torch.bartlett_window, + "none": None, + } + window_fn = torch_windows.get(window, None) + window_tensor = window_fn(self.win_length, periodic=False) if window_fn else None + self.register_buffer("window", window_tensor) + + self.normalize = normalize + self.log = log + self.dither = dither + self.frame_splicing = frame_splicing + self.nfilt = nfilt + self.preemph = preemph + self.pad_to = pad_to + highfreq = highfreq or sample_rate / 2 + + # librosa.filters.mel replaced by the torchaudio equivalent (norm="slaney" + + # mel_scale="slaney" matches librosa's default HTK=False convention; filterbank + # max abs diff 2.6e-7, verified token- and WAV-identical on the parity fixture). + if not HAVE_TORCHAUDIO: + raise ValueError(f"Need to install torchaudio to instantiate a {self.__class__.__name__}") + filterbanks = ( + torchaudio.functional.melscale_fbanks( + n_freqs=self.n_fft // 2 + 1, + f_min=lowfreq, + f_max=highfreq, + n_mels=nfilt, + sample_rate=sample_rate, + norm=mel_norm, + mel_scale="slaney", + ) + .T.to(torch.float) + .unsqueeze(0) + ) + self.register_buffer("fb", filterbanks) + + # Calculate maximum sequence length + max_length = self.get_seq_len(torch.tensor(max_duration * sample_rate, dtype=torch.float)) + max_pad = pad_to - (max_length % pad_to) if pad_to > 0 else 0 + self.max_length = max_length + max_pad + self.pad_value = pad_value + self.mag_power = mag_power + + # We want to avoid taking the log of zero + # There are two options: either adding or clamping to a small value + if log_zero_guard_type not in ["add", "clamp"]: + raise ValueError( + f"{self} received {log_zero_guard_type} for the " + f"log_zero_guard_type parameter. It must be either 'add' or " + f"'clamp'." + ) + + self.use_grads = use_grads + if not use_grads: + self.forward = torch.no_grad()(self.forward) + self._rng = random.Random() if rng is None else rng + self.nb_augmentation_prob = nb_augmentation_prob + if self.nb_augmentation_prob > 0.0: + if nb_max_freq >= sample_rate / 2: + self.nb_augmentation_prob = 0.0 + else: + self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * n_fft) + + # log_zero_guard_value is the the small we want to use, we support + # an actual number, or "tiny", or "eps" + self.log_zero_guard_type = log_zero_guard_type + logging.debug(f"sr: {sample_rate}") + logging.debug(f"n_fft: {self.n_fft}") + logging.debug(f"win_length: {self.win_length}") + logging.debug(f"hop_length: {self.hop_length}") + logging.debug(f"n_mels: {nfilt}") + logging.debug(f"fmin: {lowfreq}") + logging.debug(f"fmax: {highfreq}") + logging.debug(f"using grads: {use_grads}") + logging.debug(f"nb_augmentation_prob: {nb_augmentation_prob}") + + def stft(self, x): + return torch.stft( + x, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + center=False if self.exact_pad else True, + window=self.window.to(dtype=torch.float, device=x.device), + return_complex=True, + ) + + def log_zero_guard_value_fn(self, x): + if isinstance(self.log_zero_guard_value, str): + if self.log_zero_guard_value == "tiny": + return torch.finfo(x.dtype).tiny + elif self.log_zero_guard_value == "eps": + return torch.finfo(x.dtype).eps + else: + raise ValueError( + f"{self} received {self.log_zero_guard_value} for the " + f"log_zero_guard_type parameter. It must be either a " + f"number, 'tiny', or 'eps'" + ) + else: + return self.log_zero_guard_value + + def get_seq_len(self, seq_len): + # Assuming that center is True is stft_pad_amount = 0 + pad_amount = self.stft_pad_amount * 2 if self.stft_pad_amount is not None else self.n_fft // 2 * 2 + seq_len = torch.floor_divide((seq_len + pad_amount - self.n_fft), self.hop_length) + 1 + return seq_len.to(dtype=torch.long) + + def forward(self, x, seq_len, linear_spec=False): + # Initialize debug dictionary + # debug_snapshots = {} + # debug_snapshots['step0_input'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu() + # debug_snapshots['input_len'] = seq_len[0].item() + + seq_len = self.get_seq_len(seq_len) + # debug_snapshots['output_seq_len'] = seq_len[0].item() + + if self.stft_pad_amount is not None: + x = torch.nn.functional.pad( + x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), "reflect" + ).squeeze(1) + # debug_snapshots['step1_after_reflect_pad'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu() + + # dither (only in training mode for eval determinism) + if self.training and self.dither > 0: + x += self.dither * torch.randn_like(x) + # debug_snapshots['step2_after_dither'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu() + + # do preemphasis + if self.preemph is not None: + x = torch.cat((x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1) + # debug_snapshots['step3_after_preemph'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu() + + # disable autocast to get full range of stft values + with torch.amp.autocast(x.device.type, enabled=False): + x = self.stft(x) + + # torch stft returns complex tensor (of shape [B,N,T]); so convert to magnitude + # guard is needed for sqrt if grads are passed through + guard = 0 if not self.use_grads else CONSTANT + x = torch.view_as_real(x) + x = torch.sqrt(x.pow(2).sum(-1) + guard) + # debug_snapshots['step4_after_stft_magnitude'] = x[0, :, :10].cpu() # First 10 frames, all freq bins + + if self.training and self.nb_augmentation_prob > 0.0: + for idx in range(x.shape[0]): + if self._rng.random() < self.nb_augmentation_prob: + x[idx, self._nb_max_fft_bin :, :] = 0.0 + + # get power spectrum + if self.mag_power != 1.0: + x = x.pow(self.mag_power) + # debug_snapshots['step5_after_power'] = x[0, :, :10].cpu() + + # return plain spectrogram if required + if linear_spec: + return x, seq_len + + # disable autocast, otherwise it might be automatically casted to fp16 + # on fp16 compatible GPUs and get NaN values for input value of 65520 + with torch.amp.autocast(x.device.type, enabled=False): + # dot with filterbank energies + x = torch.matmul(self.fb.to(x.dtype), x) + # debug_snapshots['step6_after_mel_filterbank'] = x[0, :, :10].cpu() + + # log features if required + if self.log: + if self.log_zero_guard_type == "add": + x = torch.log(x + self.log_zero_guard_value_fn(x)) + elif self.log_zero_guard_type == "clamp": + x = torch.log(torch.clamp(x, min=self.log_zero_guard_value_fn(x))) + else: + raise ValueError("log_zero_guard_type was not understood") + # debug_snapshots['step7_after_log'] = x[0, :, :10].cpu() + + # frame splicing if required + if self.frame_splicing > 1: + x = splice_frames(x, self.frame_splicing) + # debug_snapshots['step8_after_frame_splicing'] = x[0, :, :10].cpu() + + # normalize if required + if self.normalize: + x, _, _ = normalize_batch(x, seq_len, normalize_type=self.normalize) + # debug_snapshots['step9_after_normalize'] = x[0, :, :10].cpu() + + # mask to zero any values beyond seq_len in batch, pad to multiple of `pad_to` (for efficiency) + max_len = x.size(-1) + mask = torch.arange(max_len, device=x.device) + mask = mask.repeat(x.size(0), 1) >= seq_len.unsqueeze(1) + x = x.masked_fill(mask.unsqueeze(1).type(torch.bool).to(device=x.device), self.pad_value) + # debug_snapshots['step10_after_masking'] = x[0, :, :10].cpu() + del mask + + pad_to = self.pad_to + if pad_to == "max": + x = nn.functional.pad(x, (0, self.max_length - x.size(-1)), value=self.pad_value) + elif pad_to > 0: + pad_amt = x.size(-1) % pad_to + if pad_amt != 0: + x = nn.functional.pad(x, (0, pad_to - pad_amt), value=self.pad_value) + # debug_snapshots['step11_after_padding'] = x[0, :, :10].cpu() + + # Debug: Save all intermediate steps + # import os + + # # Add final output and config to debug snapshots + # debug_snapshots['final_output'] = x[0].cpu() + # debug_snapshots['final_shape'] = list(x.shape) + # debug_snapshots['config'] = { + # 'pad_to': pad_to, + # 'normalize': self.normalize, + # 'exact_pad': self.exact_pad, + # 'n_fft': self.n_fft, + # 'hop_length': self.hop_length, + # 'preemph': self.preemph, + # 'mag_power': self.mag_power, + # } + + # torch.save(debug_snapshots, debug_path) + # print(f"[Preprocessor Debug] Saved {len(debug_snapshots)} steps to {debug_path}") + # print(f" Input length: {debug_snapshots['input_len']}, Output seq_len: {seq_len[0].item()}, Final shape: {x.shape}") + + return x, seq_len + + +class FilterbankFeaturesTA(nn.Module): + """ + Exportable, `torchaudio`-based implementation of Mel Spectrogram extraction. + + See `AudioToMelSpectrogramPreprocessor` for args. + + """ + + def __init__( + self, + sample_rate: int = 16000, + n_window_size: int = 320, + n_window_stride: int = 160, + normalize: str | None = "per_feature", + nfilt: int = 64, + n_fft: int | None = None, + preemph: float = 0.97, + lowfreq: float = 0, + highfreq: float | None = None, + log: bool = True, + log_zero_guard_type: str = "add", + log_zero_guard_value: Union[float, str] = 2**-24, + dither: float = 1e-5, + window: str = "hann", + pad_to: int = 0, + pad_value: float = 0.0, + mel_norm="slaney", + # Seems like no one uses these options anymore. Don't convolute the code by supporting them. + use_grads: bool = False, # Deprecated arguments; kept for config compatibility + max_duration: float = 16.7, # Deprecated arguments; kept for config compatibility + frame_splicing: int = 1, # Deprecated arguments; kept for config compatibility + exact_pad: bool = False, # Deprecated arguments; kept for config compatibility + nb_augmentation_prob: float = 0.0, # Deprecated arguments; kept for config compatibility + nb_max_freq: int = 4000, # Deprecated arguments; kept for config compatibility + mag_power: float = 2.0, # Deprecated arguments; kept for config compatibility + rng: random.Random | None = None, # Deprecated arguments; kept for config compatibility + stft_exact_pad: bool = False, # Deprecated arguments; kept for config compatibility + stft_conv: bool = False, # Deprecated arguments; kept for config compatibility + ): + super().__init__() + if not HAVE_TORCHAUDIO: + raise ValueError(f"Need to install torchaudio to instantiate a {self.__class__.__name__}") + + # Make sure log zero guard is supported, if given as a string + supported_log_zero_guard_strings = {"eps", "tiny"} + if isinstance(log_zero_guard_value, str) and log_zero_guard_value not in supported_log_zero_guard_strings: + raise ValueError( + f"Log zero guard value must either be a float or a member of {supported_log_zero_guard_strings}" + ) + + # Copied from `AudioPreprocessor` due to the ad-hoc structuring of the Mel Spec extractor class + self.torch_windows = { + "hann": torch.hann_window, + "hamming": torch.hamming_window, + "blackman": torch.blackman_window, + "bartlett": torch.bartlett_window, + "ones": torch.ones, + None: torch.ones, + } + + # Ensure we can look up the window function + if window not in self.torch_windows: + raise ValueError(f"Got window value '{window}' but expected a member of {self.torch_windows.keys()}") + + self.win_length = n_window_size + self.hop_length = n_window_stride + self._sample_rate = sample_rate + self._normalize_strategy = normalize + self._use_log = log + self._preemphasis_value = preemph + self.log_zero_guard_type = log_zero_guard_type + self.log_zero_guard_value: Union[str, float] = log_zero_guard_value + self.dither = dither + self.pad_to = pad_to + self.pad_value = pad_value + self.n_fft = n_fft + self._mel_spec_extractor: torchaudio.transforms.MelSpectrogram = torchaudio.transforms.MelSpectrogram( + sample_rate=self._sample_rate, + win_length=self.win_length, + hop_length=self.hop_length, + n_mels=nfilt, + window_fn=self.torch_windows[window], + mel_scale="slaney", + norm=mel_norm, + n_fft=n_fft, + f_max=highfreq, + f_min=lowfreq, + wkwargs={"periodic": False}, + ) + + def _resolve_log_zero_guard_value(self, dtype: torch.dtype) -> float: + if isinstance(self.log_zero_guard_value, float): + return self.log_zero_guard_value + return getattr(torch.finfo(dtype), self.log_zero_guard_value) + + def _apply_dithering(self, signals: torch.Tensor) -> torch.Tensor: + if self.training and self.dither > 0.0: + noise = torch.randn_like(signals) * self.dither + signals = signals + noise + return signals + + def _apply_preemphasis(self, signals: torch.Tensor) -> torch.Tensor: + if self._preemphasis_value is not None: + padded = torch.nn.functional.pad(signals, (1, 0)) + signals = signals - self._preemphasis_value * padded[:, :-1] + return signals + + def _compute_output_lengths(self, input_lengths: torch.Tensor) -> torch.Tensor: + out_lengths = input_lengths.div(self.hop_length, rounding_mode="floor").add(1).long() + return out_lengths + + def _apply_pad_to(self, features: torch.Tensor) -> torch.Tensor: + # Only apply during training; else need to capture dynamic shape for exported models + if not self.training or self.pad_to == 0 or features.shape[-1] % self.pad_to == 0: + return features + pad_length = self.pad_to - (features.shape[-1] % self.pad_to) + return torch.nn.functional.pad(features, pad=(0, pad_length), value=self.pad_value) + + def _apply_log(self, features: torch.Tensor) -> torch.Tensor: + if self._use_log: + zero_guard = self._resolve_log_zero_guard_value(features.dtype) + if self.log_zero_guard_type == "add": + features = features + zero_guard + elif self.log_zero_guard_type == "clamp": + features = features.clamp(min=zero_guard) + else: + raise ValueError(f"Unsupported log zero guard type: '{self.log_zero_guard_type}'") + features = features.log() + return features + + def _extract_spectrograms(self, signals: torch.Tensor) -> torch.Tensor: + # Complex FFT needs to be done in single precision + with torch.amp.autocast("cuda", enabled=False): + features = self._mel_spec_extractor(waveform=signals) + return features + + def _apply_normalization(self, features: torch.Tensor, lengths: torch.Tensor, eps: float = 1e-5) -> torch.Tensor: + # For consistency, this function always does a masked fill even if not normalizing. + mask: torch.Tensor = make_seq_mask_like(lengths=lengths, like=features, time_dim=-1, valid_ones=False) + features = features.masked_fill(mask, 0.0) + # Maybe don't normalize + if self._normalize_strategy is None: + return features + # Use the log zero guard for the sqrt zero guard + guard_value = self._resolve_log_zero_guard_value(features.dtype) + if self._normalize_strategy == "per_feature" or self._normalize_strategy == "all_features": + # 'all_features' reduces over each sample; 'per_feature' reduces over each channel + reduce_dim = 2 + if self._normalize_strategy == "all_features": + reduce_dim = [1, 2] + # [B, D, T] -> [B, D, 1] or [B, 1, 1] + means = features.sum(dim=reduce_dim, keepdim=True).div(lengths.view(-1, 1, 1)) + stds = ( + features.sub(means) + .masked_fill(mask, 0.0) + .pow(2.0) + .sum(dim=reduce_dim, keepdim=True) # [B, D, T] -> [B, D, 1] or [B, 1, 1] + .div(lengths.view(-1, 1, 1) - 1) # assume biased estimator + .clamp(min=guard_value) # avoid sqrt(0) + .sqrt() + ) + features = (features - means) / (stds + eps) + else: + # Deprecating constant std/mean + raise ValueError(f"Unsupported norm type: '{self._normalize_strategy}") + features = features.masked_fill(mask, 0.0) + return features + + def forward(self, input_signal: torch.Tensor, length: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + feature_lengths = self._compute_output_lengths(input_lengths=length) + signals = self._apply_dithering(signals=input_signal) + signals = self._apply_preemphasis(signals=signals) + features = self._extract_spectrograms(signals=signals) + features = self._apply_log(features=features) + features = self._apply_normalization(features=features, lengths=feature_lengths) + features = self._apply_pad_to(features=features) + return features, feature_lengths diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py new file mode 100644 index 000000000..fcfc36f4b --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py @@ -0,0 +1,1133 @@ +# Copyright 2017 Johns Hopkins University (Shinji Watanabe) +# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/multi_head_attention.py; training-only code removed. + +""" +Part of this code is adopted from https://github.com/espnet/espnet +""" + +import math +from functools import lru_cache +from typing import List, Tuple + +import torch +import torch.nn as nn +import torch.nn.attention +import torch.nn.functional as F + +from ..compat import avoid_float16_autocast_context + +__all__ = [ + "RelPositionMultiHeadAttention", + "RelPositionalEncoding", + "PositionalEncoding", +] + +INF_VAL = 10000.0 + + +class MultiHeadAttention(nn.Module): + """Multi-Head Attention layer of Transformer. + Args: + n_head (int): number of heads + n_feat (int): size of the features + dropout_rate (float): dropout rate + use_bias (bool): whether to remove bias in linear and conv layers + use_pytorch_sdpa (bool): use torch sdpa instead of manual attention + use_pytorch_sdpa_backends list[str]: list of backend names to use in sdpa. None or empty list means all backends. e.g. ["MATH"] + """ + + def __init__( + self, + n_head, + n_feat, + dropout_rate, + max_cache_len=0, + use_bias=True, + use_pytorch_sdpa=False, + use_pytorch_sdpa_backends=None, + ): + """Construct an MultiHeadedAttention object.""" + super(MultiHeadAttention, self).__init__() + self.use_pytorch_sdpa = use_pytorch_sdpa + if self.use_pytorch_sdpa and use_pytorch_sdpa_backends: + use_pytorch_sdpa_backends = list( + map( + lambda backend_name: getattr(torch.nn.attention.SDPBackend, backend_name), + use_pytorch_sdpa_backends, + ) + ) + self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends + + self.cache_drop_size = None + self.use_bias = use_bias + self.dropout_rate = dropout_rate + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.s_d_k = math.sqrt(self.d_k) + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat, bias=use_bias) + self.linear_k = nn.Linear(n_feat, n_feat, bias=use_bias) + self.linear_v = nn.Linear(n_feat, n_feat, bias=use_bias) + self.linear_out = nn.Linear(n_feat, n_feat, bias=use_bias) + self.dropout = nn.Dropout(p=dropout_rate) + + self._max_cache_len = max_cache_len + + def forward_qkv(self, query, key, value): + """Transforms query, key and value. + Args: + query (torch.Tensor): (batch, time1, size) + key (torch.Tensor): (batch, time2, size) + value (torch.Tensor): (batch, time2, size) + returns: + q (torch.Tensor): (batch, head, time1, size) + k (torch.Tensor): (batch, head, time2, size) + v (torch.Tensor): (batch, head, time2, size) + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + return q, k, v + + def forward_attention(self, value, scores, mask): + """Compute attention context vector. + Args: + value (torch.Tensor): (batch, time2, size) + scores(torch.Tensor): (batch, time1, time2) + mask(torch.Tensor): (batch, time1, time2) + returns: + value (torch.Tensor): transformed `value` (batch, time2, d_model) weighted by the attention scores + """ + n_batch = value.size(0) + if mask is not None: + mask = mask.unsqueeze(1) # (batch, 1, time1, time2) + scores = scores.masked_fill(mask, -INF_VAL) + attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) # (batch, head, time1, time2) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = x.transpose(1, 2).reshape(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, query, key, value, mask, pos_emb=None, cache=None): + """Compute 'Scaled Dot Product Attention'. + Args: + query (torch.Tensor): (batch, time1, size) + key (torch.Tensor): (batch, time2, size) + value(torch.Tensor): (batch, time2, size) + mask (torch.Tensor): (batch, time1, time2) + cache (torch.Tensor) : (batch, time_cache, size) + + returns: + output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention + cache (torch.Tensor) : (batch, time_cache_next, size) + """ + key, value, query, cache = self.update_cache(key=key, value=value, query=query, cache=cache) + + if torch.is_autocast_enabled(): + query, key, value = query.to(torch.float32), key.to(torch.float32), value.to(torch.float32) + + # temporary until we solve this more gracefully + with avoid_float16_autocast_context(): + q, k, v = self.forward_qkv(query, key, value) + + if self.use_pytorch_sdpa: + n_batch = value.size(0) + + if mask is not None: + mask = ~mask.unsqueeze(1) + + dropout_rate = self.dropout_rate if self.training else 0 + if self.use_pytorch_sdpa_backends: + with torch.nn.attention.sdpa_kernel(self.use_pytorch_sdpa_backends): + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=mask, dropout_p=dropout_rate + ) + else: + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=mask, dropout_p=dropout_rate + ) + + # this IF block can be deleted when https://github.com/pytorch/pytorch/pull/131863 is in the stable version + if mask is not None: + all_masked_rows = torch.all(~mask, dim=-1) + all_masked_rows.unsqueeze_(-1) + out = out.masked_fill(all_masked_rows, 0.0) + + out = out.transpose(1, 2).reshape(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) + out = self.linear_out(out) # (batch, time1, d_model) + else: + scores = torch.matmul(q, k.transpose(-2, -1)) / self.s_d_k + out = self.forward_attention(v, scores, mask) + + if cache is None: + return out + else: + return out, cache + + def update_cache(self, key, value, query, cache): + if cache is not None: + key = value = torch.cat([cache, key], dim=1) + q_keep_size = query.shape[1] - self.cache_drop_size + cache = torch.cat([cache[:, q_keep_size:, :], query[:, :q_keep_size, :]], dim=1) + return key, value, query, cache + + +class RelPositionMultiHeadAttention(MultiHeadAttention): + """Multi-Head Attention layer of Transformer-XL with support of relative positional encoding. + Paper: https://arxiv.org/abs/1901.02860 + Args: + n_head (int): number of heads + n_feat (int): size of the features + dropout_rate (float): dropout rate + use_bias (bool): whether to apply bias in linear and conv layers of MultiHeadAttention + """ + + def __init__( + self, + n_head, + n_feat, + dropout_rate, + pos_bias_u, + pos_bias_v, + max_cache_len=0, + use_bias=True, + use_pytorch_sdpa=False, + use_pytorch_sdpa_backends=None, + ): + """Construct an RelPositionMultiHeadedAttention object.""" + super().__init__( + n_head=n_head, + n_feat=n_feat, + dropout_rate=dropout_rate, + max_cache_len=max_cache_len, + use_bias=use_bias, + use_pytorch_sdpa=use_pytorch_sdpa, + use_pytorch_sdpa_backends=use_pytorch_sdpa_backends, + ) + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable biases are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + if pos_bias_u is None or pos_bias_v is None: + self.pos_bias_u = nn.Parameter(torch.FloatTensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.FloatTensor(self.h, self.d_k)) + # nn.init.normal_(self.pos_bias_u, 0.0, 0.02) + # nn.init.normal_(self.pos_bias_v, 0.0, 0.02) + nn.init.zeros_(self.pos_bias_u) + nn.init.zeros_(self.pos_bias_v) + else: + self.pos_bias_u = pos_bias_u + self.pos_bias_v = pos_bias_v + + def rel_shift(self, x): + """Compute relative positional encoding. + Args: + x (torch.Tensor): (batch, nheads, time, 2*time-1) + """ + b, h, qlen, pos_len = x.size() # (b, h, t1, t2) + # need to add a column of zeros on the left side of last dimension to perform the relative shifting + x = torch.nn.functional.pad(x, pad=(1, 0)) # (b, h, t1, t2+1) + x = x.view(b, h, -1, qlen) # (b, h, t2+1, t1) + # need to drop the first row + x = x[:, :, 1:].view(b, h, qlen, pos_len) # (b, h, t1, t2) + return x + + def forward(self, query, key, value, mask, pos_emb, cache=None): + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + Args: + query (torch.Tensor): (batch, time1, size) + key (torch.Tensor): (batch, time2, size) + value(torch.Tensor): (batch, time2, size) + mask (torch.Tensor): (batch, time1, time2) + pos_emb (torch.Tensor) : (batch, time1, size) + cache (torch.Tensor) : (batch, time_cache, size) + + Returns: + output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention + cache (torch.Tensor) : (batch, time_cache_next, size) + """ + key, value, query, cache = self.update_cache(key=key, value=value, query=query, cache=cache) + + if torch.is_autocast_enabled(): + query, key, value = query.to(torch.float32), key.to(torch.float32), value.to(torch.float32) + + # temporary until we solve this more gracefully + with avoid_float16_autocast_context(): + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + n_batch_pos = pos_emb.size(0) + n_batch = value.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + + # compute matrix b and matrix d + # (batch, head, time1, time2) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + + if self.use_pytorch_sdpa: + scale_factor = 1 / math.sqrt(q_with_bias_u.size(-1)) + matrix_bd = matrix_bd[:, :, :, : k.size(-2)] * scale_factor + + if mask is not None: + mask = mask.unsqueeze(1) + matrix_bd.masked_fill_(mask, -INF_VAL) + + dropout_rate = self.dropout_rate if self.training else 0 + if self.use_pytorch_sdpa_backends: + with torch.nn.attention.sdpa_kernel(self.use_pytorch_sdpa_backends): + out = torch.nn.functional.scaled_dot_product_attention( + q_with_bias_u, k, v, attn_mask=matrix_bd, dropout_p=dropout_rate + ) + else: + out = torch.nn.functional.scaled_dot_product_attention( + q_with_bias_u, k, v, attn_mask=matrix_bd, dropout_p=dropout_rate + ) + + # this IF block can be deleted when https://github.com/pytorch/pytorch/pull/131863 is in the stable version + if mask is not None: + all_masked_rows = torch.all(mask, dim=-1) + all_masked_rows.unsqueeze_(-1) + all_masked_rows = all_masked_rows.expand(-1, out.size(1), -1, out.size(-1)) + out = out.masked_fill(all_masked_rows, 0.0) + + out = out.transpose(1, 2).reshape(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) + out = self.linear_out(out) # (batch, time1, d_model) + else: + # drops extra elements in the matrix_bd to match the matrix_ac's size + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)] + scores = (matrix_ac + matrix_bd) / self.s_d_k # (batch, head, time1, time2) + out = self.forward_attention(v, scores, mask) + + if cache is None: + return out + else: + return out, cache + + +class RelPositionMultiHeadAttentionLongformer(RelPositionMultiHeadAttention): + """Multi-Head Attention layer of Transformer-XL with sliding window local+global attention from Longformer. + Partially adapted from allenai (https://github.com/allenai/longformer/blob/master/longformer/sliding_chunks.py) + and huggingface (https://github.com/huggingface/transformers/blob/main/src/transformers/models/longformer/modeling_longformer.py) + Paper: https://arxiv.org/abs/1901.02860 (Transformer-XL), + https://arxiv.org/abs/2004.05150 (Longformer) + Args: + n_head (int): number of heads + n_feat (int): size of the features + dropout_rate (float): dropout rate + pos_bias_u (Tensor): the positional bias matrix U + pos_bias_v (Tensor): the positional bias matrix V + att_context_size (List[int]): List of 2 ints corresponding to left and right attention context sizes. + max_cache_len (int): the maximum size of cache + global_tokens (int): number of tokens to be used for global attention + global_tokens_spacing (int): how far apart the global tokens are + global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate + use_bias (bool): whether to apply bias in linear and conv layers of MultiHeadAttention + """ + + def __init__( + self, + n_head, + n_feat, + dropout_rate, + pos_bias_u, + pos_bias_v, + att_context_size, + max_cache_len=0, + global_tokens=0, + global_tokens_spacing=1, + global_attn_separate=False, + use_bias=True, + use_pytorch_sdpa=False, + use_pytorch_sdpa_backends=None, + ): + """Construct an RelPositionMultiHeadAttentionLongformer object.""" + super().__init__( + n_head=n_head, + n_feat=n_feat, + dropout_rate=dropout_rate, + pos_bias_u=pos_bias_u, + pos_bias_v=pos_bias_v, + max_cache_len=max_cache_len, + use_bias=use_bias, + use_pytorch_sdpa=use_pytorch_sdpa, + use_pytorch_sdpa_backends=use_pytorch_sdpa_backends, + ) + + if use_pytorch_sdpa: + raise NotImplementedError("Not implemented for Longformer yet") + + self.att_context_size = att_context_size + self.global_tokens = global_tokens + self.global_tokens_spacing = global_tokens_spacing + self.global_attn_separate = global_attn_separate + + if self.global_attn_separate: + self.global_q = nn.Linear(n_feat, n_feat, bias=use_bias) + self.global_k = nn.Linear(n_feat, n_feat, bias=use_bias) + self.global_v = nn.Linear(n_feat, n_feat, bias=use_bias) + + def forward(self, query, key, value, pad_mask, pos_emb, cache=None): + """Compute Scaled Dot Product Local Attention with rel. positional encoding. using overlapping chunks + Args: + query (torch.Tensor): (batch, time, size) + key (torch.Tensor): (batch, time, size) + value(torch.Tensor): (batch, time, size) + pad_mask (torch.Tensor): (batch, time) + pos_emb (torch.Tensor) : (batch, 2w + 1, size) + cache (torch.Tensor) : (batch, time_cache, size) + Returns: + output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention + cache (torch.Tensor) : (batch, time_cache_next, size) + """ + + key, value, query, cache = self.update_cache(key=key, value=value, query=query, cache=cache) + + if torch.is_autocast_enabled(): + query, key, value = query.to(torch.float32), key.to(torch.float32), value.to(torch.float32) + + # temporary until we solve this more gracefully + with avoid_float16_autocast_context(): + q, k, v = self.forward_qkv(query, key, value) + n_batch, _, T, _ = q.size() + + w = max(self.att_context_size[0], self.att_context_size[1]) + if w <= 0: + raise ValueError("When using local attention, context size must be set > 0") + pad_len = (2 * w - T % (2 * w)) % (2 * w) # pad time to 2w + q = F.pad(q, (0, 0, 0, pad_len)) # (batch, head, time, size) + k = F.pad(k, (0, 0, 0, pad_len)) # (batch, head, time, size) + v = F.pad(v, (0, 0, 0, pad_len)) # (batch, head, time, size) + mask = F.pad(pad_mask, (0, pad_len), value=1.0) + + q_with_bias_u = q + self.pos_bias_u.unsqueeze(1) # (batch, head, time, size) + q_with_bias_v = q + self.pos_bias_v.unsqueeze(1) # (batch, head, time, size) + + diagonal_matrix_ac = self.sliding_chunks_matmul_qk( + q_with_bias_u, k, w, padding_value=0.0 + ) # (batch, head, time, 2w + 1) + + # add relative positional embedding + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k).transpose(1, 2) + # (batch, head, 2w, size) + diagonal_matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + # (batch, head, time, 2w + 1) + + start_pos = w - self.att_context_size[0] + end_pos = w + self.att_context_size[1] + + diagonal_matrix_ac[:, :, :, : self.att_context_size[0]] += diagonal_matrix_bd[ + :, :, :, : self.att_context_size[0] + ] + diagonal_matrix_ac[:, :, :, -(self.att_context_size[1] + 1) :] += diagonal_matrix_bd[ + :, :, :, self.att_context_size[0] : + ] + scores = diagonal_matrix_ac / self.s_d_k + # (batch, head, time, 2w + 1) + + # mask invalid positions + scores[:, :, :, :start_pos] = -INF_VAL + scores[:, :, :, end_pos + 1 :] = -INF_VAL + + # This implementation is fast and takes very little memory because num_heads x hidden_size = 1 + # from (bsz x seq_len) to (bsz x num_heads x seqlen x hidden_size) + mask = mask.unsqueeze(dim=1).unsqueeze(dim=-1) + # cast to float/half then replace 1's with -inf + float_mask = mask.type_as(scores).masked_fill(mask, -INF_VAL) + ones = float_mask.new_ones(size=float_mask.size()) # tensor of ones + # diagonal mask with zeros everywhere and -inf inplace of padding + d_mask = self.sliding_chunks_matmul_qk(ones, float_mask, w, padding_value=0.0) + # (batch, head, time, 2w + 1) + + scores += d_mask + + if self.global_tokens > 0: + # create q, k, v for global attn + if self.global_attn_separate: + global_q = self.global_q(query).view(n_batch, -1, self.h, self.d_k) + global_k = self.global_k(key).view(n_batch, -1, self.h, self.d_k) + global_v = self.global_v(value).view(n_batch, -1, self.h, self.d_k) + global_q = global_q.transpose(1, 2) + global_k = global_k.transpose(1, 2) + global_v = global_v.transpose(1, 2) + global_q = F.pad(global_q, (0, 0, 0, pad_len)) # (batch, head, time, size) + global_k = F.pad(global_k, (0, 0, 0, pad_len)) # (batch, head, time, size) + global_v = F.pad(global_v, (0, 0, 0, pad_len)) # (batch, head, time, size) + else: + global_q, global_k, global_v = q, k, v + + global_q /= self.s_d_k + + # assign which tokens are global + is_index_global_attn = torch.zeros_like(pad_mask) + is_index_global_attn[ + :, : self.global_tokens * self.global_tokens_spacing : self.global_tokens_spacing + ] = 1.0 + + # compute global attn indices + ( + max_num_global_attn_indices, + is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero, + ) = self._get_global_attn_indices(is_index_global_attn=is_index_global_attn) + + # calculate global attn probs with global keys + # (batch, time, head, max_num_global_attn_indices) + global_key_attn = self._compute_global_key_attn( + query=global_q.transpose(1, 2), + key=global_k.transpose(1, 2), + max_num_global_attn_indices=max_num_global_attn_indices, + is_index_global_attn_nonzero=is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero, + ).transpose(1, 2) + + # concat to local_attn_probs + # (batch, time, head, max_num_global_attn_indices + 2*w) + scores = torch.cat((global_key_attn, scores), dim=-1) + + # free memory + del global_key_attn + + attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) + p_attn = self.dropout(attn) + # (batch, head, time, 2w + 1) + + if self.global_tokens > 0: + # compute sum of global and local attn + out = self._compute_attn_output_with_global_indices( + value=v, + attn_probs=p_attn, + max_num_global_attn_indices=max_num_global_attn_indices, + is_index_global_attn_nonzero=is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero, + w=w, + ) + else: + # compute local attn only + out = self.sliding_chunks_matmul_pv(p_attn, v, w) + + out = out.reshape(n_batch, -1, self.h * self.d_k)[:, :T] + + if self.global_tokens > 0: + out_global_to_all = self._compute_out_global_to_all( + query=global_q, + key=global_k, + value=global_v, + max_num_global_attn_indices=max_num_global_attn_indices, + is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero, + is_index_global_attn_nonzero=is_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero, + is_index_masked=mask, + ) + + # overwrite values with global attention + out[is_index_global_attn_nonzero] = out_global_to_all + + ret = self.linear_out(out) + + if cache is None: + return ret + else: + return ret, cache + + def _get_global_attn_indices(self, is_index_global_attn: torch.Tensor) -> Tuple: + """ + Compute global attention indices. + + Args: + is_index_global_attn (torch.Tensor): (batch, time) A boolean tensor indicating if an index is a global attention index. + + Returns: + max_num_global_attn_indices (int): Maximum number of global attention indices in the batch. + is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements). + is_local_index_global_attn_nonzero (tuple): Indices of non-padding values within global attention indices. + is_local_index_no_global_attn_nonzero (tuple): Indices of padding values within global attention indices. + """ + # Calculate the number of global attention indices in the batch + num_global_attn_indices = is_index_global_attn.long().sum(dim=1) + + # Find the maximum number of global attention indices in the batch + max_num_global_attn_indices = num_global_attn_indices.max() + + # Get the indices of global attention (non-zero elements) + is_index_global_attn_nonzero = is_index_global_attn.nonzero(as_tuple=True) + + # Create a helper tensor to find the local indices of global attention + is_local_index_global_attn = torch.arange( + max_num_global_attn_indices, device=is_index_global_attn.device + ) < num_global_attn_indices.unsqueeze(dim=-1) + + # Find the non-padding values within global attention indices + is_local_index_global_attn_nonzero = is_local_index_global_attn.nonzero(as_tuple=True) + + # Find the padding values within global attention indices + is_local_index_no_global_attn_nonzero = (is_local_index_global_attn == 0).nonzero(as_tuple=True) + + return ( + max_num_global_attn_indices, + is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero, + ) + + def _compute_global_key_attn( + self, + key: torch.Tensor, + query: torch.Tensor, + max_num_global_attn_indices: int, + is_index_global_attn_nonzero: tuple, + is_local_index_global_attn_nonzero: tuple, + is_local_index_no_global_attn_nonzero: tuple, + ) -> torch.Tensor: + """ + Compute the attention probabilities using only global key vectors. + + Args: + key (torch.Tensor): (batch, time, head, head_dim) The key vectors. + query (torch.Tensor): (batch, time, head, head_dim) The query vectors. + max_num_global_attn_indices (int): Maximum number of global attention indices in the batch. + is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements). + is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices. + is_local_index_no_global_attn_nonzero (tuple): Padding values within global attention indices. + + Returns: + attn_probs_from_global_key (torch.Tensor): (batch, time, head, max_num_global_attn_indices) The computed attention probabilities using only global key vectors. + """ + batch_size = key.shape[0] + + # create only global key vectors + key_only_global = key.new_zeros(batch_size, max_num_global_attn_indices, self.h, self.d_k) + + key_only_global[is_local_index_global_attn_nonzero] = key[is_index_global_attn_nonzero] + + # (batch_size, seq_len, head, max_num_global_attn_indices) + attn_probs_from_global_key = torch.einsum("blhd,bshd->blhs", (query, key_only_global)) + + # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets + attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3) + attn_probs_from_global_key[ + is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, : + ] = torch.finfo(attn_probs_from_global_key.dtype).min + attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3) + + return attn_probs_from_global_key + + def _compute_attn_output_with_global_indices( + self, + value: torch.Tensor, + attn_probs: torch.Tensor, + max_num_global_attn_indices: int, + is_index_global_attn_nonzero: tuple, + is_local_index_global_attn_nonzero: tuple, + w: int, + ) -> torch.Tensor: + """ + Compute the attention output with global indices. + + Args: + value (torch.Tensor): (batch, head, time, head_dim) The value vectors for global attention. + attn_probs (torch.Tensor): (batch, time, head, 2w) The attention probabilities. + max_num_global_attn_indices (int): Maximum number of global attention indices in the batch. + is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements). + is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices. + w (int): Local context size + Returns: + torch.Tensor: (batch, time, head x head_dim) The attention output of all tokens attending to global. + """ + batch_size, time = attn_probs.shape[0], attn_probs.shape[2] + + value = value.transpose(1, 2) + + # get value vectors for global only + value_vectors_only_global = value.new_zeros(batch_size, max_num_global_attn_indices, self.h, self.d_k) + value_vectors_only_global[is_local_index_global_attn_nonzero] = value[is_index_global_attn_nonzero] + + # cut local attn probs to global only + attn_probs_only_global = attn_probs.narrow(-1, 0, max_num_global_attn_indices) + # compute attn output only global + attn_output_only_global = torch.matmul( + attn_probs_only_global.clone(), value_vectors_only_global.transpose(1, 2).clone() + ).transpose(1, 2) + + # reshape attn probs + attn_probs_without_global = attn_probs.narrow( + -1, max_num_global_attn_indices, attn_probs.size(-1) - max_num_global_attn_indices + ).contiguous() + + # compute attn output with global + attn_output_without_global = self.sliding_chunks_matmul_pv(attn_probs_without_global, value.transpose(1, 2), w) + + return attn_output_only_global + attn_output_without_global + + def _compute_out_global_to_all( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + max_num_global_attn_indices: int, + is_local_index_global_attn_nonzero: tuple, + is_index_global_attn_nonzero: tuple, + is_local_index_no_global_attn_nonzero: tuple, + is_index_masked: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute the attention output of global tokens attending to all. + + Args: + query (torch.Tensor): (batch, head, time, head_dim) The queries for global attention. + key (torch.Tensor): (batch, head, time, head_dim) The keys for global attention. + value (torch.Tensor): (batch, head, time, head_dim) The values for global attention. + max_num_global_attn_indices (int): Maximum number of global attention indices in the batch. + is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices. + is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements). + is_local_index_no_global_attn_nonzero (tuple): Padding values within global attention indices. + is_index_masked (torch.Tensor): (batch, time) A boolean tensor indicating if an index is masked. + + Returns: + global_attn_output (torch.Tensor): (batch, max_num_global_attn_indices, head x head_dim) + The attention output of global tokens attending to all. + """ + + batch_size = key.shape[0] + seq_len = key.shape[2] + + global_k = key.reshape(batch_size * self.h, -1, self.d_k) + global_v = value.reshape(batch_size * self.h, -1, self.d_k) + + global_q = query.transpose(1, 2) + global_q_from_global = global_q.new_zeros(batch_size, max_num_global_attn_indices, self.h, self.d_k) + global_q_from_global[is_local_index_global_attn_nonzero] = global_q[is_index_global_attn_nonzero] + global_q_from_global = global_q_from_global.transpose(0, 1).reshape(batch_size * self.h, -1, self.d_k) + + # compute attn scores + global_attn_scores = torch.bmm(global_q_from_global, global_k.transpose(1, 2)) + global_attn_scores = global_attn_scores.view(batch_size, self.h, max_num_global_attn_indices, seq_len) + + # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets + global_attn_scores = global_attn_scores.transpose(1, 2) + global_attn_scores[is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, :] = ( + torch.finfo(global_attn_scores.dtype).min + ) + global_attn_scores = global_attn_scores.transpose(1, 2) + + global_attn_scores = global_attn_scores.masked_fill( + is_index_masked.transpose(2, 3), + torch.finfo(global_attn_scores.dtype).min, + ) + + global_attn_scores = global_attn_scores.view(batch_size * self.h, max_num_global_attn_indices, seq_len) + + # compute global attn probs + if self.training: + global_attn_probs_float = nn.functional.softmax(global_attn_scores, dim=-1, dtype=torch.float32) + else: + global_attn_probs_float = nn.functional.softmax(global_attn_scores, dim=-1) + + global_attn_probs = self.dropout(global_attn_probs_float) + + # global attn output + global_attn_output = torch.bmm(global_attn_probs, global_v) + global_attn_output = global_attn_output.view(batch_size, self.h, max_num_global_attn_indices, self.d_k) + + global_attn_output = global_attn_output[ + is_local_index_global_attn_nonzero[0], :, is_local_index_global_attn_nonzero[1] + ] + + global_attn_output = global_attn_output.reshape(global_attn_output.shape[0], -1) + + return global_attn_output + + # Longformer implementation for overlap case + # + def _skew(self, x: torch.Tensor, direction: List[int], padding_value: float) -> torch.Tensor: + """Convert diagonals into columns (or columns into diagonals depending on `direction` + + Args: + x (torch.Tensor): (batch x head, chunk_count, 2w, 2w) + direction (List[int]): padding directions + padding_value (float): value to pad with + + Returns: + output (torch.Tensor): (batch x head, chunk_count, 2w, 2w + 1) + + """ + x_padded = F.pad(x, direction, value=padding_value) + x_padded = x_padded.view(*x_padded.size()[:-2], x_padded.size(-1), x_padded.size(-2)) + return x_padded + + def _skew2(self, x: torch.Tensor, padding_value: float) -> torch.Tensor: + """Shift every row 1 step to right converting columns into diagonals + + Args: + x (torch.Tensor): (batch x head, chunks_count + 1, w, 2w + 1) + padding_value (float): value to pad with + + Returns: + output (torch.Tensor): (batch x head, chunks_count + 1, w, 3w) + """ + # X = B x C x M x L + B, C, M, L = x.size() + x = F.pad(x, (0, M + 1), value=padding_value) # B x C x M x (L+M+1) + x = x.view(B, C, -1) # B x C x ML+MM+M + x = x[:, :, :-M] # B x C x ML+MM + x = x.view(B, C, M, M + L) # B x C, M x L+M + x = x[:, :, :, :-1] + return x + + def _chunk_overlap(self, x: torch.Tensor, w: int) -> torch.Tensor: + """Convert into overlapping chunks. + + Args: + x (torch.Tensor): # (batch x head, time, size) + w (int): Chunk overlap size + + Returns: + output (torch.Tensor): # (batch x head, chunk_count, 2w, size) + """ + + # non-overlapping chunks of size = 2w + x = x.view(x.size(0), x.size(1) // (w * 2), w * 2, x.size(2)) + + # use `as_strided` to make the chunks overlap with an overlap size = w + chunk_size = list(x.size()) + chunk_size[1] = chunk_size[1] * 2 - 1 + + chunk_stride = list(x.stride()) + chunk_stride[1] = chunk_stride[1] // 2 + return x.as_strided(size=chunk_size, stride=chunk_stride) + + @lru_cache() + def _get_invalid_locations_mask(self, w: int, device: str): + diagonals_list = [] + for j in range(-w, 1): + diagonal_mask = torch.zeros(w, device="cpu", dtype=torch.uint8) + diagonal_mask[:-j] = 1 + diagonals_list.append(diagonal_mask) + + mask = torch.stack(diagonals_list, dim=-1) + mask = mask[None, None, :, :] + + ending_mask = mask.flip(dims=(2, 3)).bool().to(device) + return mask.bool().to(device), ending_mask + + def mask_invalid_locations( + self, + input_tensor: torch.Tensor, + w: int, + ): + """ + Mask locations invalid for the sliding window attention + + Args: + input_tensor (torch.Tensor): # (batch x head, time, size) + w (int): Chunk overlap size + """ + beginning_mask, ending_mask = self._get_invalid_locations_mask(w, input_tensor.device) + seq_len = input_tensor.size(2) + beginning_input = input_tensor[:, :, :w, : w + 1] + beginning_mask = beginning_mask[:, :, :seq_len].expand(beginning_input.size()) + beginning_input.masked_fill_(beginning_mask, -float("inf")) + + ending_input = input_tensor[:, :, -w:, -(w + 1) :] + ending_mask = ending_mask[:, :, -seq_len:].expand(ending_input.size()) + ending_input.masked_fill_(ending_mask, -float("inf")) + + def sliding_chunks_matmul_qk(self, q: torch.Tensor, k: torch.Tensor, w: int, padding_value: float) -> torch.Tensor: + """Matrix multiplication of query x key tensors using with a sliding window attention pattern. + This implementation splits the input into overlapping chunks of size 2w + with an overlap of size w + + Args: + q (torch.Tensor): (batch, head, time, size) + k (torch.Tensor): (batch, head, time, size) + w (int): Chunk overlap size + padding_value (float): Value to pad with + + Returns: + output (torch.Tensor): (batch, head, time, 2w + 1) + """ + bsz, num_heads, seqlen, head_dim = q.size() + assert seqlen % (w * 2) == 0 + assert q.size() == k.size() + + chunks_count = seqlen // w - 1 + + # group bsz and num_heads dimensions into one, then chunk seqlen into chunks of size w * 2 + q = q.reshape(bsz * num_heads, seqlen, head_dim) + k = k.reshape(bsz * num_heads, seqlen, head_dim) + + chunk_q = self._chunk_overlap(q, w) # (batch x head, chunk_count, 2w, size) + chunk_k = self._chunk_overlap(k, w) # (batch x head, chunk_count, 2w, size) + + # matrix multiplication + # bcxd: bsz*num_heads x chunks x 2w x head_dim + # bcyd: bsz*num_heads x chunks x 2w x head_dim + # bcxy: bsz*num_heads x chunks x 2w x 2w + chunk_attn = torch.einsum("bcxd,bcyd->bcxy", (chunk_q, chunk_k)) # multiply + # (batch x head, chunk_count, 2w, 2w) + + # convert diagonals into columns + diagonal_chunk_attn = self._skew(chunk_attn, direction=(0, 0, 0, 1), padding_value=padding_value) + # (batch x head, chunk_count, 2w, 2w + 1) + + # allocate space for the overall attention matrix where the chunks are combined. The last dimension + # has (w * 2 + 1) columns. The first (w) columns are the w lower triangles (attention from a word to + # w previous words). The following column is attention score from each word to itself, then + # followed by w columns for the upper triangle. + + diagonal_attn = diagonal_chunk_attn.new_empty((bsz * num_heads, chunks_count + 1, w, w * 2 + 1)) + # (batch x head, chunk_count + 1, w, 2w + 1) + + # copy parts from diagonal_chunk_attn into the compined matrix of attentions + # - copying the main diagonal and the upper triangle + diagonal_attn[:, :-1, :, w:] = diagonal_chunk_attn[:, :, :w, : w + 1] + diagonal_attn[:, -1, :, w:] = diagonal_chunk_attn[:, -1, w:, : w + 1] + # - copying the lower triangle + diagonal_attn[:, 1:, :, :w] = diagonal_chunk_attn[:, :, -(w + 1) : -1, w + 1 :] + diagonal_attn[:, 0, 1:w, 1:w] = diagonal_chunk_attn[:, 0, : w - 1, 1 - w :] + + # separate bsz and num_heads dimensions again + diagonal_attn = diagonal_attn.view(bsz, num_heads, seqlen, 2 * w + 1) + # (batch, head, time, 2w + 1) + + self.mask_invalid_locations(diagonal_attn, w) + + return diagonal_attn + + def sliding_chunks_matmul_pv(self, prob: torch.Tensor, v: torch.Tensor, w: int): + """Same as sliding_chunks_matmul_qk but for prob and value tensors. + + Args: + prob (torch.Tensor): (batch, head, time, size) + v (torch.Tensor): (batch, head, time, size) + w (int): Chunk overlap size + + Returns: + output (torch.Tensor): (batch, time, head, size) + """ + bsz, num_heads, seqlen, head_dim = v.size() + chunks_count = seqlen // w - 1 + # group bsz and num_heads dimensions into one, then chunk seqlen into chunks of size 2w + chunk_prob = prob.reshape(bsz * num_heads, seqlen // w, w, 2 * w + 1) + # (batch x head, chunks_count + 1, w, 2w + 1) + + # group bsz and num_heads dimensions into one + v = v.reshape(bsz * num_heads, seqlen, head_dim) + # (batch x head, time, size) + + # pad seqlen with w at the beginning of the sequence and another w at the end + padded_v = F.pad(v, (0, 0, w, w), value=-1) + # (batch x head, time + 2w, size) + + # chunk padded_v into chunks of size 3w and an overlap of size w + chunk_v_size = (bsz * num_heads, chunks_count + 1, 3 * w, head_dim) + chunk_v_stride = padded_v.stride() + chunk_v_stride = chunk_v_stride[0], w * chunk_v_stride[1], chunk_v_stride[1], chunk_v_stride[2] + chunk_v = padded_v.as_strided(size=chunk_v_size, stride=chunk_v_stride) + # (batch x head, chunks_count + 1, 3w, size) + + skewed_prob = self._skew2(chunk_prob, padding_value=0) + # (batch x head, chunks_count + 1, w, 3w) + + context = torch.einsum("bcwd,bcdh->bcwh", (skewed_prob, chunk_v)) + # (batch x head, chunks_count + 1, w, size) + + return context.view(bsz, num_heads, seqlen, head_dim).transpose(1, 2) + + +class PositionalEncoding(torch.nn.Module): + """Fixed sinusoidal positional encoding. + Args: + d_model (int): embedding dim + dropout_rate (float): dropout rate + max_len (int): maximum input length + xscale (bool): whether to scale the input by sqrt(d_model) + dropout_rate_emb (float): dropout rate for the positional embeddings + """ + + def __init__(self, d_model, dropout_rate, max_len=5000, xscale=None, dropout_rate_emb=0.0): + """Construct an PositionalEncoding object.""" + super(PositionalEncoding, self).__init__() + self.d_model = d_model + self.xscale = xscale + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.max_len = max_len + if dropout_rate_emb > 0: + self.dropout_emb = nn.Dropout(dropout_rate_emb) + else: + self.dropout_emb = None + + def create_pe(self, positions, dtype): + pos_length = positions.size(0) + pe = torch.zeros(pos_length, self.d_model, device=positions.device) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32, device=positions.device) + * -(math.log(INF_VAL) / self.d_model) + ) + pe[:, 0::2] = torch.sin(positions * div_term) + pe[:, 1::2] = torch.cos(positions * div_term) + pe = pe.unsqueeze(0).to(dtype) + if hasattr(self, "pe"): + self.pe = pe + else: + self.register_buffer("pe", pe, persistent=False) + + def extend_pe(self, length, device, dtype): + """Reset and extend the positional encodings if needed.""" + if hasattr(self, "pe") and self.pe.size(1) >= length: + return + positions = torch.arange(0, length, dtype=torch.float32, device=device).unsqueeze(1) + self.create_pe(positions=positions, dtype=dtype) + + def forward(self, x: torch.Tensor, cache_len=0): + """Adds positional encoding. + Args: + x (torch.Tensor): Input. Its shape is (batch, time, feature_size) + cache_len (int): the size of the cache which is used to shift positions + Returns: + x+pos_emb (torch.Tensor): Its shape is (batch, time, feature_size) + pos_emb (torch.Tensor): Its shape is (1, time, feature_size) + """ + input_len = x.size(1) + cache_len + if self.xscale: + x = x * self.xscale + pos_emb = self.pe[:, :input_len] + if self.dropout_emb: + pos_emb = self.dropout_emb(pos_emb) + x = x + pos_emb + return self.dropout(x), pos_emb + + +class RelPositionalEncoding(PositionalEncoding): + """Relative positional encoding for TransformerXL's layers + See : Appendix B in https://arxiv.org/abs/1901.02860 + Args: + d_model (int): embedding dim + dropout_rate (float): dropout rate + max_len (int): maximum input length + xscale (bool): whether to scale the input by sqrt(d_model) + dropout_rate_emb (float): dropout rate for the positional embeddings + """ + + def extend_pe(self, length, device, dtype): + """Reset and extend the positional encodings if needed.""" + needed_size = 2 * length - 1 + if hasattr(self, "pe") and self.pe.size(1) >= needed_size: + return + # positions would be from negative numbers to positive + # positive positions would be used for left positions and negative for right positions + positions = torch.arange(length - 1, -length, -1, dtype=torch.float32, device=device).unsqueeze(1) + self.create_pe(positions=positions, dtype=dtype) + + def forward(self, x, cache_len=0): + """Compute positional encoding. + Args: + x (torch.Tensor): Input. Its shape is (batch, time, feature_size) + cache_len (int): the size of the cache which is used to shift positions + Returns: + x (torch.Tensor): Its shape is (batch, time, feature_size) + pos_emb (torch.Tensor): Its shape is (1, time, feature_size) + """ + + if self.xscale: + x = x * self.xscale + + # center_pos would be the index of position 0 + # negative positions would be used for right and positive for left tokens + # for input of length L, 2*L-1 positions are needed, positions from (L-1) to -(L-1) + input_len = x.size(1) + cache_len + center_pos = self.pe.size(1) // 2 + 1 + start_pos = center_pos - input_len + end_pos = center_pos + input_len - 1 + pos_emb = self.pe[:, start_pos:end_pos] + if self.dropout_emb: + pos_emb = self.dropout_emb(pos_emb) + return self.dropout(x), pos_emb + + +class LocalAttRelPositionalEncoding(PositionalEncoding): + """Relative positional encoding for sliding window attention or chunked attention. + See above for relative positional encoding based on Transformer-XL paper + Args: + left_chunk_size (int): number of frames to in past chunks + chunk size (int): number of frames (max frames if using multimode) in current chunk + d_model (int): embedding dim + dropout_rate (float): dropout rate + max_len (int): maximum input length + xscale (bool): whether to scale the input by sqrt(d_model) + dropout_rate_emb (float): dropout rate for the positional embeddings + """ + + def __init__(self, att_context_size, **kwargs): + super(LocalAttRelPositionalEncoding, self).__init__(**kwargs) + self.left_context = att_context_size[0] + self.right_context = att_context_size[1] + + def extend_pe(self, length, device, dtype): + """Reset and extend the positional encodings only at the beginning""" + if hasattr(self, "pe"): + return + + positions = torch.arange( + self.left_context, -self.right_context - 1, -1, dtype=torch.float32, device=device + ).unsqueeze(1) + self.create_pe(positions=positions, dtype=dtype) + + def forward(self, x, cache_len=0): + """Compute positional encoding. + Args: + x (torch.Tensor): Input. Its shape is (batch, time, feature_size) + Returns: + x (torch.Tensor): Its shape is (batch, time, feature_size) + pos_emb (torch.Tensor): Its shape is (1, time, feature_size) + """ + + if self.xscale: + x = x * self.xscale + + end_pos = self.left_context + self.right_context + 1 + pos_emb = self.pe[:, :end_pos] + if self.dropout_emb: + pos_emb = self.dropout_emb(pos_emb) + return self.dropout(x), pos_emb diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py new file mode 100644 index 000000000..a4ccc14ae --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/utils/regularization_utils.py; training-only code removed. + +from typing import List + + +def compute_stochastic_depth_drop_probs( + num_layers: int, + stochastic_depth_drop_prob: float = 0.0, + stochastic_depth_mode: str = "linear", + stochastic_depth_start_layer: int = 1, +) -> List[float]: + """Computes drop probabilities for stochastic depth regularization technique. + The first layer is never dropped and the starting layer needs to be greater + or equal to 1. + + Args: + num_layers (int): number of layers in the network. + stochastic_depth_drop_prob (float): if non-zero, will randomly drop + layers during training. The higher this value, the more often layers + are dropped. Defaults to 0.0. + stochastic_depth_mode (str): can be either "linear" or "uniform". If + set to "uniform", all layers have the same probability of drop. If + set to "linear", the drop probability grows linearly from 0 for the + first layer to the desired value for the final layer. Defaults to + "linear". + stochastic_depth_start_layer (int): starting layer for stochastic depth. + All layers before this will never be dropped. Note that drop + probability will be adjusted accordingly if mode is "linear" when + start layer is > 1. Defaults to 1. + Returns: + List[float]: list of drop probabilities for all layers + """ + if not (0 <= stochastic_depth_drop_prob < 1.0): + raise ValueError("stochastic_depth_drop_prob has to be in [0, 1).") + if not (1 <= stochastic_depth_start_layer <= num_layers): + raise ValueError("stochastic_depth_start_layer has to be in [1, num layers].") + + # Layers before `stochastic_depth_start_layer` are never dropped + layer_drop_probs = [0.0] * stochastic_depth_start_layer + + # Layers starting with `stochastic_depth_start_layer` may be dropped + if (L := num_layers - stochastic_depth_start_layer) > 0: + if stochastic_depth_mode == "linear": + # we start with 1/L * drop_prob and and end with the desired drop probability. + layer_drop_probs += [l / L * stochastic_depth_drop_prob for l in range(1, L + 1)] + elif stochastic_depth_mode == "uniform": + layer_drop_probs += [stochastic_depth_drop_prob] * L + else: + raise ValueError( + f'stochastic_depth_mode has to be one of ["linear", "uniform"]. Current value: {stochastic_depth_mode}' + ) + return layer_drop_probs diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py new file mode 100644 index 000000000..1faffee7d --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py @@ -0,0 +1,628 @@ +# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/subsampling.py; training-only code removed. + +import math + +import torch +import torch.nn as nn +from torch.nn import LayerNorm + +from .causal_convs import CausalConv1D, CausalConv2D +from ..compat import logging + + +class StackingSubsampling(torch.nn.Module): + """Stacking subsampling which simply stacks consecutive frames to reduce the sampling rate + Args: + subsampling_factor (int): The subsampling factor + feat_in (int): size of the input features + feat_out (int): size of the output features + norm (bool): whether to use an MLP layer after the stacking along with normalization. default is False. + """ + + def __init__(self, subsampling_factor, feat_in, feat_out, norm=False): + super(StackingSubsampling, self).__init__() + self.subsampling_factor = subsampling_factor + self.proj_out = torch.nn.Linear(subsampling_factor * feat_in, feat_out) + if norm: + self.pre_norm = LayerNorm(feat_in) + else: + self.pre_norm = None + + def get_sampling_frames(self): + return self.subsampling_factor + + def get_streaming_cache_size(self): + return 0 + + def forward(self, x, lengths): + b, t, h = x.size() + pad_size = (self.subsampling_factor - (t % self.subsampling_factor)) % self.subsampling_factor + x = torch.nn.functional.pad(x, (0, 0, 0, pad_size)) + if self.pre_norm is not None: + x = self.pre_norm(x) + _, t, _ = x.size() + x = torch.reshape(x, (b, t // self.subsampling_factor, h * self.subsampling_factor)) + x = self.proj_out(x) + lengths = torch.div(lengths + pad_size, self.subsampling_factor, rounding_mode="floor") + return x, lengths + + +class ConvSubsampling(torch.nn.Module): + """Convolutional subsampling which supports VGGNet and striding approach introduced in: + VGGNet Subsampling: Transformer-transducer: end-to-end speech recognition with self-attention (https://arxiv.org/pdf/1910.12977.pdf) + Striding Subsampling: "Speech-Transformer: A No-Recurrence Sequence-to-Sequence Model for Speech Recognition" by Linhao Dong et al. (https://ieeexplore.ieee.org/document/8462506) + Args: + subsampling (str): The subsampling technique from {"vggnet", "striding", "dw-striding"} + subsampling_factor (int): The subsampling factor which should be a power of 2 + subsampling_conv_chunking_factor (int): Input chunking factor which can be -1 (no chunking) + 1 (auto) or a power of 2. Default is 1 + feat_in (int): size of the input features + feat_out (int): size of the output features + conv_channels (int): Number of channels for the convolution layers. + activation (Module): activation function, default is nn.ReLU() + """ + + def __init__( + self, + subsampling, + subsampling_factor, + feat_in, + feat_out, + conv_channels, + subsampling_conv_chunking_factor=1, + activation=nn.ReLU(), + is_causal=False, + ): + super(ConvSubsampling, self).__init__() + self._subsampling = subsampling + self._conv_channels = conv_channels + self._feat_in = feat_in + self._feat_out = feat_out + + if subsampling_factor % 2 != 0: + raise ValueError("Sampling factor should be a multiply of 2!") + self._sampling_num = int(math.log(subsampling_factor, 2)) + self.subsampling_factor = subsampling_factor + self.is_causal = is_causal + + if ( + subsampling_conv_chunking_factor != -1 + and subsampling_conv_chunking_factor != 1 + and subsampling_conv_chunking_factor % 2 != 0 + ): + raise ValueError("subsampling_conv_chunking_factor should be -1, 1, or a power of 2") + self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor + + in_channels = 1 + layers = [] + + if subsampling == "vggnet": + self._stride = 2 + self._kernel_size = 2 + self._ceil_mode = True + + self._left_padding = 0 + self._right_padding = 0 + + for i in range(self._sampling_num): + layers.append( + torch.nn.Conv2d( + in_channels=in_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1 + ) + ) + layers.append(activation) + layers.append( + torch.nn.Conv2d( + in_channels=conv_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1 + ) + ) + layers.append(activation) + layers.append( + torch.nn.MaxPool2d( + kernel_size=self._kernel_size, + stride=self._stride, + padding=self._left_padding, + ceil_mode=self._ceil_mode, + ) + ) + in_channels = conv_channels + + elif subsampling == "dw_striding": + self._stride = 2 + self._kernel_size = 3 + self._ceil_mode = False + + if self.is_causal: + self._left_padding = self._kernel_size - 1 + self._right_padding = self._stride - 1 + self._max_cache_len = subsampling_factor + 1 + else: + self._left_padding = (self._kernel_size - 1) // 2 + self._right_padding = (self._kernel_size - 1) // 2 + self._max_cache_len = 0 + + # Layer 1 + if self.is_causal: + layers.append( + CausalConv2D( + in_channels=in_channels, + out_channels=conv_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=None, + ) + ) + else: + layers.append( + torch.nn.Conv2d( + in_channels=in_channels, + out_channels=conv_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=self._left_padding, + ) + ) + in_channels = conv_channels + layers.append(activation) + + for i in range(self._sampling_num - 1): + if self.is_causal: + layers.append( + CausalConv2D( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=None, + groups=in_channels, + ) + ) + else: + layers.append( + torch.nn.Conv2d( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=self._left_padding, + groups=in_channels, + ) + ) + + layers.append( + torch.nn.Conv2d( + in_channels=in_channels, + out_channels=conv_channels, + kernel_size=1, + stride=1, + padding=0, + groups=1, + ) + ) + layers.append(activation) + in_channels = conv_channels + + elif subsampling == "striding": + self._stride = 2 + self._kernel_size = 3 + self._ceil_mode = False + + if self.is_causal: + self._left_padding = self._kernel_size - 1 + self._right_padding = self._stride - 1 + self._max_cache_len = subsampling_factor + 1 + else: + self._left_padding = (self._kernel_size - 1) // 2 + self._right_padding = (self._kernel_size - 1) // 2 + self._max_cache_len = 0 + + for i in range(self._sampling_num): + if self.is_causal: + layers.append( + CausalConv2D( + in_channels=in_channels, + out_channels=conv_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=None, + ) + ) + else: + layers.append( + torch.nn.Conv2d( + in_channels=in_channels, + out_channels=conv_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=self._left_padding, + ) + ) + layers.append(activation) + in_channels = conv_channels + + elif subsampling == "striding_conv1d": + in_channels = feat_in + + self._stride = 2 + self._kernel_size = 5 + self._ceil_mode = False + + if self.is_causal: + self._left_padding = self._kernel_size - 1 + self._right_padding = self._stride - 1 + self._max_cache_len = subsampling_factor + 1 + else: + self._left_padding = (self._kernel_size - 1) // 2 + self._right_padding = (self._kernel_size - 1) // 2 + self._max_cache_len = 0 + + for i in range(self._sampling_num): + if self.is_causal: + layers.append( + CausalConv1D( + in_channels=in_channels, + out_channels=feat_out if self._sampling_num == i + 1 else conv_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=None, + ) + ) + else: + layers.append( + torch.nn.Conv1d( + in_channels=in_channels, + out_channels=feat_out if self._sampling_num == i + 1 else conv_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=self._left_padding, + ) + ) + layers.append(activation) + in_channels = conv_channels + + elif subsampling == "dw_striding_conv1d": + in_channels = feat_in + + self._stride = 2 + self._kernel_size = 5 + self._ceil_mode = False + + self._left_padding = (self._kernel_size - 1) // 2 + self._right_padding = (self._kernel_size - 1) // 2 + + # Layer 1 + layers.extend( + [ + torch.nn.Conv1d( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=self._left_padding, + groups=in_channels, + ), + torch.nn.Conv1d( + in_channels=in_channels, + out_channels=feat_out if self._sampling_num == 1 else conv_channels, + kernel_size=1, + stride=1, + padding=0, + groups=1, + ), + ] + ) + in_channels = conv_channels + layers.append(activation) + + for i in range(self._sampling_num - 1): + layers.extend( + [ + torch.nn.Conv1d( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=self._kernel_size, + stride=self._stride, + padding=self._left_padding, + groups=in_channels, + ), + torch.nn.Conv1d( + in_channels=in_channels, + out_channels=feat_out if self._sampling_num == i + 2 else conv_channels, + kernel_size=1, + stride=1, + padding=0, + groups=1, + ), + ] + ) + layers.append(activation) + in_channels = conv_channels + + else: + raise ValueError(f"Not valid sub-sampling: {subsampling}!") + + if subsampling in ["vggnet", "dw_striding", "striding"]: + in_length = torch.tensor(feat_in, dtype=torch.float) + out_length = calc_length( + lengths=in_length, + all_paddings=self._left_padding + self._right_padding, + kernel_size=self._kernel_size, + stride=self._stride, + ceil_mode=self._ceil_mode, + repeat_num=self._sampling_num, + ) + self.out = torch.nn.Linear(conv_channels * int(out_length), feat_out) + self.conv2d_subsampling = True + elif subsampling in ["striding_conv1d", "dw_striding_conv1d"]: + self.out = None + self.conv2d_subsampling = False + else: + raise ValueError(f"Not valid sub-sampling: {subsampling}!") + + self.conv = torch.nn.Sequential(*layers) + + def get_sampling_frames(self): + return [1, self.subsampling_factor] + + def get_streaming_cache_size(self): + return [0, self.subsampling_factor + 1] + + def forward(self, x, lengths): + lengths = calc_length( + lengths, + all_paddings=self._left_padding + self._right_padding, + kernel_size=self._kernel_size, + stride=self._stride, + ceil_mode=self._ceil_mode, + repeat_num=self._sampling_num, + ) + + # Unsqueeze Channel Axis + if self.conv2d_subsampling: + x = x.unsqueeze(1) + # Transpose to Channel First mode + else: + x = x.transpose(1, 2) + + # split inputs if chunking_factor is set + if self.subsampling_conv_chunking_factor != -1 and self.conv2d_subsampling: + if self.subsampling_conv_chunking_factor == 1: + # if subsampling_conv_chunking_factor is 1, we split only if needed + # avoiding a bug / feature limiting indexing of tensors to 2**31 + # see https://github.com/pytorch/pytorch/issues/80020 + x_ceil = 2**31 / self._conv_channels * self._stride * self._stride + if torch.numel(x) > x_ceil: + need_to_split = True + else: + need_to_split = False + else: + # if subsampling_conv_chunking_factor > 1 we always split + need_to_split = True + + if need_to_split: + x, success = self.conv_split_by_batch(x) + if not success: # if unable to split by batch, try by channel + if self._subsampling == "dw_striding": + x = self.conv_split_by_channel(x) + else: + x = self.conv(x) # try anyway + else: + x = self.conv(x) + else: + x = self.conv(x) + + # Flatten Channel and Frequency Axes + if self.conv2d_subsampling: + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).reshape(b, t, -1)) + # Transpose to Channel Last mode + else: + x = x.transpose(1, 2) + + return x, lengths + + def reset_parameters(self): + # initialize weights + if self._subsampling == "dw_striding": + with torch.no_grad(): + # init conv + scale = 1.0 / self._kernel_size + dw_max = (self._kernel_size**2) ** -0.5 + pw_max = self._conv_channels**-0.5 + + torch.nn.init.uniform_(self.conv[0].weight, -scale, scale) + torch.nn.init.uniform_(self.conv[0].bias, -scale, scale) + + for idx in range(2, len(self.conv), 3): + torch.nn.init.uniform_(self.conv[idx].weight, -dw_max, dw_max) + torch.nn.init.uniform_(self.conv[idx].bias, -dw_max, dw_max) + torch.nn.init.uniform_(self.conv[idx + 1].weight, -pw_max, pw_max) + torch.nn.init.uniform_(self.conv[idx + 1].bias, -pw_max, pw_max) + + # init fc (80 * 64 = 5120 from https://github.com/kssteven418/Squeezeformer/blob/13c97d6cf92f2844d2cb3142b4c5bfa9ad1a8951/src/models/conformer_encoder.py#L487 + fc_scale = (self._feat_out * self._feat_in / self._sampling_num) ** -0.5 + torch.nn.init.uniform_(self.out.weight, -fc_scale, fc_scale) + torch.nn.init.uniform_(self.out.bias, -fc_scale, fc_scale) + + def conv_split_by_batch(self, x): + """Tries to split input by batch, run conv and concat results""" + b, _, _, _ = x.size() + if b == 1: # can't split if batch size is 1 + return x, False + + if self.subsampling_conv_chunking_factor > 1: + cf = self.subsampling_conv_chunking_factor + logging.debug(f"using manually set chunking factor: {cf}") + else: + # avoiding a bug / feature limiting indexing of tensors to 2**31 + # see https://github.com/pytorch/pytorch/issues/80020 + x_ceil = 2**31 / self._conv_channels * self._stride * self._stride + p = math.ceil(math.log(torch.numel(x) / x_ceil, 2)) + cf = 2**p + logging.debug(f"using auto set chunking factor: {cf}") + + new_batch_size = b // cf + if new_batch_size == 0: # input is too big + return x, False + + logging.debug(f"conv subsampling: using split batch size {new_batch_size}") + return torch.cat([self.conv(chunk) for chunk in torch.split(x, new_batch_size, 0)]), True + + def conv_split_by_channel(self, x): + """For dw convs, tries to split input by time, run conv and concat results""" + x = self.conv[0](x) # full conv2D + x = self.conv[1](x) # activation + + for i in range(self._sampling_num - 1): + _, c, t, _ = x.size() + + if self.subsampling_conv_chunking_factor > 1: + cf = self.subsampling_conv_chunking_factor + logging.debug(f"using manually set chunking factor: {cf}") + else: + # avoiding a bug / feature limiting indexing of tensors to 2**31 + # see https://github.com/pytorch/pytorch/issues/80020 + p = math.ceil(math.log(torch.numel(x) / 2**31, 2)) + cf = 2**p + logging.debug(f"using auto set chunking factor: {cf}") + + new_c = int(c // cf) + if new_c == 0: + logging.warning(f"chunking factor {cf} is too high; splitting down to one channel.") + new_c = 1 + + new_t = int(t // cf) + if new_t == 0: + logging.warning(f"chunking factor {cf} is too high; splitting down to one timestep.") + new_t = 1 + + logging.debug(f"conv dw subsampling: using split C size {new_c} and split T size {new_t}") + x = self.channel_chunked_conv(self.conv[i * 3 + 2], new_c, x) # conv2D, depthwise + + # splitting pointwise convs by time + x = torch.cat([self.conv[i * 3 + 3](chunk) for chunk in torch.split(x, new_t, 2)], 2) # conv2D, pointwise + x = self.conv[i * 3 + 4](x) # activation + return x + + def channel_chunked_conv(self, conv, chunk_size, x): + """Performs channel chunked convolution""" + + ind = 0 + out_chunks = [] + for chunk in torch.split(x, chunk_size, 1): + step = chunk.size()[1] + + if self.is_causal: + chunk = nn.functional.pad( + chunk, pad=(self._kernel_size - 1, self._stride - 1, self._kernel_size - 1, self._stride - 1) + ) + ch_out = nn.functional.conv2d( + chunk, + conv.weight[ind : ind + step, :, :, :], + bias=conv.bias[ind : ind + step], + stride=self._stride, + padding=0, + groups=step, + ) + else: + ch_out = nn.functional.conv2d( + chunk, + conv.weight[ind : ind + step, :, :, :], + bias=conv.bias[ind : ind + step], + stride=self._stride, + padding=self._left_padding, + groups=step, + ) + out_chunks.append(ch_out) + ind += step + + return torch.cat(out_chunks, 1) + + def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int): + if ( + subsampling_conv_chunking_factor != -1 + and subsampling_conv_chunking_factor != 1 + and subsampling_conv_chunking_factor % 2 != 0 + ): + raise ValueError("subsampling_conv_chunking_factor should be -1, 1, or a power of 2") + self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor + + +def calc_length(lengths, all_paddings, kernel_size, stride, ceil_mode, repeat_num=1): + """Calculates the output length of a Tensor passed through a convolution or max pooling layer""" + add_pad: float = all_paddings - kernel_size + one: float = 1.0 + for i in range(repeat_num): + lengths = torch.div(lengths.to(dtype=torch.float) + add_pad, stride) + one + if ceil_mode: + lengths = torch.ceil(lengths) + else: + lengths = torch.floor(lengths) + return lengths.to(dtype=torch.int) + + +class SubsamplingReductionModule(nn.Module): + """Downsamples the audio signal in time dimension.""" + + def __init__(self, reduction: str, d_model: int, reduction_factor: int = 2): + super().__init__() + + assert reduction in ["pooling", "striding"] + + self.reduction = reduction + self.d_model = d_model + self._sampling_num = int(math.log(reduction_factor, 2)) + + if reduction == "pooling": + self.reduction_enc = nn.MaxPool1d(kernel_size=reduction_factor) + self.padding = 0 + self.kernel_size = self.reduction_enc.kernel_size + self.stride = self.reduction_enc.stride + elif reduction == "striding": + self.reduction_enc = ConvSubsampling( + subsampling="striding", + subsampling_factor=reduction_factor, + feat_in=d_model, + feat_out=d_model, + conv_channels=d_model, + activation=nn.ReLU(), + is_causal=False, + ) + + def forward(self, x, lengths): + """Shapes: + - x: [B, T, C] + - lengths: [B] + """ + + if self.reduction == "striding": + x, lengths = self.reduction_enc(x=x, lengths=lengths) + else: + x = torch.transpose(x, 1, 2) # [B, C, T] + lengths = calc_length( + lengths=lengths, + all_paddings=self.padding, + kernel_size=self.kernel_size, + stride=self.stride, + ceil_mode=False, + repeat_num=self._sampling_num, + ) + x = self.reduction_enc(x) + x = torch.transpose(x, 1, 2) # [B, T, C] + + return x, lengths diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py new file mode 100644 index 000000000..f1c648acd --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py @@ -0,0 +1,644 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); compatibility shims +# replacing nemo.core / nemo.utils machinery for dependency-free inference. +# Each vendored function below is an exact copy from its annotated original source path; +# only imports were adjusted. Stub classes replace NeMo base classes whose functionality +# is not needed at inference time. + +import importlib +import logging as _stdlib_logging +import math +import warnings +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from typing import Dict, Optional + +import torch +import torch.nn as nn +from omegaconf import OmegaConf +from safetensors.torch import load_file +from transformers import AutoConfig, AutoModelForCausalLM + + +# ============================================================================== +# Logging shim (replaces the nemo.utils logging import) +# ============================================================================== +class logging_mode: + """Stub of nemo.utils.logging's dedup modes (used as `mode=logging_mode.ONCE`).""" + + EACH = 0 + ONCE = 1 + + +class _NemoLoggerShim: + """stdlib-logging based replacement for the nemo.utils logging module. + + NeMo's logger accepts an extra `mode=` kwarg (deduplication); it is accepted + and ignored here, with a simple set-based dedup for `mode=logging_mode.ONCE`. + """ + + def __init__(self, logger): + self._logger = logger + self._once_seen = set() + + def _log(self, level, msg, args, mode=None): + if mode == logging_mode.ONCE: + key = (level, str(msg)) + if key in self._once_seen: + return + self._once_seen.add(key) + self._logger.log(level, msg, *args) + + def debug(self, msg, *args, mode=None, **kwargs): + self._log(_stdlib_logging.DEBUG, msg, args, mode=mode) + + def info(self, msg, *args, mode=None, **kwargs): + self._log(_stdlib_logging.INFO, msg, args, mode=mode) + + def warning(self, msg, *args, mode=None, **kwargs): + self._log(_stdlib_logging.WARNING, msg, args, mode=mode) + + def error(self, msg, *args, mode=None, **kwargs): + self._log(_stdlib_logging.ERROR, msg, args, mode=mode) + + def critical(self, msg, *args, mode=None, **kwargs): + self._log(_stdlib_logging.CRITICAL, msg, args, mode=mode) + + +logging = _NemoLoggerShim(_stdlib_logging.getLogger("nemo_vendored")) + + +# ============================================================================== +# typecheck no-op (replaces nemo.core.classes.common.typecheck) +# ============================================================================== +class typecheck: + """No-op replacement for NeMo's typecheck decorator. + + Supports the two usages found in the vendored files: + @typecheck() + def forward(...) + + @typecheck.disable_checks() + def forward(...) + + with typecheck.disable_checks(): + ... + """ + + def __init__(self, *args, **kwargs): + pass + + def __call__(self, func): + return func + + @staticmethod + @contextmanager + def disable_checks(): + yield + + +# ============================================================================== +# Neural type stubs (replace nemo.core.neural_types) +# ============================================================================== + + +# ============================================================================== +# from_config_dict target resolver (replaces hydra-based NeMo Serialization) +# ============================================================================== +# Maps the `_target_` strings that appear in the NVIDIA-NemotronLabs-VoiceChat-11B +# checkpoint config.json to the vendored classes. Extend as needed. +_TARGET_REGISTRY = { + "nemo.collections.asr.modules.ConformerEncoder": (".asr.conformer_encoder", "ConformerEncoder"), + "nemo.collections.asr.modules.conformer_encoder.ConformerEncoder": ( + ".asr.conformer_encoder", + "ConformerEncoder", + ), + "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor": ( + ".asr.audio_preprocessing", + "AudioToMelSpectrogramPreprocessor", + ), + "nemo.collections.asr.modules.audio_preprocessing.AudioToMelSpectrogramPreprocessor": ( + ".asr.audio_preprocessing", + "AudioToMelSpectrogramPreprocessor", + ), + "nemo.collections.speechlm2.modules.perception.IdentityConnector": (".perception", "IdentityConnector"), + "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule": (".perception", "AudioPerceptionModule"), +} + +_PACKAGE = __name__.rsplit(".", 1)[0] # ...nemotron_voicechat.nemo_vendored + + +def from_config_dict(config): + """Instantiate a vendored class from an OmegaConf/dict config with a `_target_` key. + + Hydra-free replacement for nemo.core.classes.common.Serialization.from_config_dict. + Like hydra.utils.instantiate with the default `_convert_="none"`, all keys except + `_target_` are passed as keyword arguments (nested containers stay OmegaConf nodes). + """ + cfg = OmegaConf.create(config) if not OmegaConf.is_config(config) else config + if "_target_" not in cfg: + raise ValueError(f"from_config_dict requires a '_target_' key in the config; got keys: {list(cfg.keys())}") + target = cfg["_target_"] + if target not in _TARGET_REGISTRY: + raise ValueError( + f"Unknown _target_ '{target}' — not in the vendored nemo_vendored.compat._TARGET_REGISTRY. " + f"Known targets: {sorted(_TARGET_REGISTRY)}. Add a mapping if this class was vendored." + ) + module_name, class_name = _TARGET_REGISTRY[target] + cls = getattr(importlib.import_module(module_name, package=_PACKAGE), class_name) + kwargs = {str(k): cfg[k] for k in cfg if k != "_target_"} + return cls(**kwargs) + + +# ============================================================================== +# NeuralModule / Exportable / mixin stubs (replace nemo.core.classes.*) +# ============================================================================== +class NeuralModule(nn.Module): + """Slim replacement of nemo.core.classes.module.NeuralModule.""" + + @property + def num_weights(self): + num: int = 0 + for p in self.parameters(): + if p.requires_grad: + num += p.numel() + return num + + @staticmethod + def from_config_dict(config): + return from_config_dict(config) + + def unfreeze(self, partial: bool = False) -> None: + """Unfreeze all parameters for training. (exact copy of NeuralModule.unfreeze)""" + if partial and not hasattr(self, "_frozen_grad_map"): + raise ValueError("Cannot unfreeze partially without first freezing the module with `freeze()`") + for pname, param in self.named_parameters(): + if not partial: + param.requires_grad = True + elif pname in self._frozen_grad_map: + param.requires_grad = self._frozen_grad_map[pname] + if hasattr(self, "_frozen_grad_map"): + delattr(self, "_frozen_grad_map") + self.train() + + +class Exportable: + """Stub of nemo.core.classes.exportable.Exportable (ONNX export machinery removed).""" + + @property + def input_names(self): + return [] + + @property + def output_names(self): + return [] + + +class AccessMixin: + """Stub of nemo.core.classes.mixins.access_mixins.AccessMixin. + + Tensor-capture registry is training/interctc-only; at inference access is + always disabled so guarded blocks are skipped. + """ + + @property + def access_cfg(self): + return {} + + def is_access_enabled(self, guid=None): + return False + + def register_accessible_tensor(self, name, tensor): + pass + + +class AttentionAdapterModuleMixin: + """Stub of nemo.collections.asr.parts.submodules.adapters.attention_adapter_mixin. + Adapters (PEFT) are training-only; report none available.""" + + def is_adapter_available(self) -> bool: + return False + + +class AdapterModuleMixin(AttentionAdapterModuleMixin): + """Stub of nemo.core.classes.mixins.adapter_mixins.AdapterModuleMixin.""" + + pass + + +# ============================================================================== +# StreamingEncoder — exact copy of +# nemo/collections/asr/parts/mixins/streaming.py (imports adjusted only) +# ============================================================================== +from abc import ABC, abstractmethod # noqa: E402 + + +class StreamingEncoder(ABC): + @abstractmethod + def setup_streaming_params( + self, + max_look_ahead: int = 10000, + ): + """ + This function sets the needed values and parameters to perform streaming. The configuration (CacheAwareStreamingConfig) need to be stored in self.streaming_cfg. + The streaming configuration is needed to simulate streaming inference. It would set the following + """ + pass + + +# ============================================================================== +# CacheAwareStreamingConfig — exact copy of the dataclass from +# nemo/collections/asr/models/configs/asr_models_config.py +# ============================================================================== +@dataclass +class CacheAwareStreamingConfig: + chunk_size: int = 0 # the size of each chunk at each step, it can be a list of two integers to specify different chunk sizes for the first step and others + shift_size: int = 0 # the size of the shift in each step, it can be a list of two integers to specify different shift sizes for the first step and others + + cache_drop_size: int = 0 # the number of steps to drop from the cache + last_channel_cache_size: int = 0 # the size of the needed cache for last channel layers + + valid_out_len: int = ( + 0 # the number of the steps in the final output which are valid (have the same value as in the offline mode) + ) + + pre_encode_cache_size: int = 0 # the size of the needed cache for the pre-encoding part of the model to avoid caching inside the pre-encoding layers + drop_extra_pre_encoded: int = 0 # the number of steps to get dropped after the pre-encoding layer + + last_channel_num: int = 0 # number of the last channel layers (like MHA layers) which need caching in the model + last_time_num: int = 0 # number of the last time layers (like convolutions) which need caching in the model + + +# ============================================================================== +# avoid_float16_autocast_context / avoid_bfloat16_autocast_context — exact +# copies of NeMo nemo/utils/cast_utils.py +# ============================================================================== + + +def avoid_float16_autocast_context(): + """ + If the current autocast context is float16, cast it to bfloat16 + if available (unless we're in jit) or float32 + """ + + if torch.is_autocast_enabled() and torch.get_autocast_gpu_dtype() == torch.float16: + if torch.jit.is_scripting() or torch.jit.is_tracing(): + return torch.amp.autocast("cuda", dtype=torch.float32) + + if torch.cuda.is_bf16_supported(): + return torch.amp.autocast("cuda", dtype=torch.bfloat16) + else: + return torch.amp.autocast("cuda", dtype=torch.float32) + else: + return nullcontext() + + +# ============================================================================== +# fp32_precision — exact copy of nemo/collections/speechlm2/parts/precision.py +# ============================================================================== +@contextmanager +def fp32_precision(): + """ + Workaround for precision related issues when training with bf16-true PyTorch Lightning precision setting. + In bf16-true, PTL changes PyTorch's default dtype, which may break implicit assumptions for some models. + This context manager restores default float32 precision and runs the computation in float32 autocast context. + """ + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + try: + with torch.amp.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu", dtype=torch.float32): + yield + finally: + torch.set_default_dtype(default_dtype) + + +# ============================================================================== +# get_pad_id — exact copy of nemo/collections/speechlm2/data/utils.py +# ============================================================================== +def get_pad_id(tokenizer) -> int: + pad_id = tokenizer.pad + if pad_id is not None: + return pad_id + pad_id = tokenizer.unk_id + if pad_id is not None: + return pad_id + warnings.warn( + "The text tokenizer has no or tokens available, using ID 0 for padding (this may lead to silent bugs)." + ) + return 0 + + +# ============================================================================== +# resample — exact copies of resample/_get_sinc_resample_kernel/_apply_sinc_resample_kernel +# of nemo/collections/audio/parts/utils/resampling.py (itself vendored from +# torchaudio 2.6.0, BSD 2-Clause; see that file for the license text). +# nemo/collections/audio/parts/utils/transforms.py re-exports `resample` from there. +# ============================================================================== +def resample( + waveform: torch.Tensor, + orig_freq: int, + new_freq: int, + lowpass_filter_width: int = 6, + rolloff: float = 0.99, + resampling_method: str = "sinc_interp_hann", + beta: float | None = None, +) -> torch.Tensor: + r"""Resamples the waveform at the new frequency using bandlimited interpolation. :cite:`RESAMPLE`.""" + + if orig_freq <= 0.0 or new_freq <= 0.0: + raise ValueError("Original frequency and desired frequency should be positive") + + if orig_freq == new_freq: + return waveform + + gcd = math.gcd(int(orig_freq), int(new_freq)) + + kernel, width = _get_sinc_resample_kernel( + orig_freq, + new_freq, + gcd, + lowpass_filter_width, + rolloff, + resampling_method, + beta, + waveform.device, + waveform.dtype, + ) + resampled = _apply_sinc_resample_kernel(waveform, orig_freq, new_freq, gcd, kernel, width) + return resampled + + +def _get_sinc_resample_kernel( + orig_freq: int, + new_freq: int, + gcd: int, + lowpass_filter_width: int = 6, + rolloff: float = 0.99, + resampling_method: str = "sinc_interp_hann", + beta: float | None = None, + device: torch.device = "cpu", + dtype: torch.dtype | None = None, +): + if not (int(orig_freq) == orig_freq and int(new_freq) == new_freq): + raise Exception( + "Frequencies must be of integer type to ensure quality resampling computation. " + "To work around this, manually convert both frequencies to integer values " + "that maintain their resampling rate ratio before passing them into the function. " + "Example: To downsample a 44100 hz waveform by a factor of 8, use " + "`orig_freq=8` and `new_freq=1` instead of `orig_freq=44100` and `new_freq=5512.5`. " + "For more information, please refer to https://github.com/pytorch/audio/issues/1487." + ) + + if resampling_method not in ["sinc_interp_hann", "sinc_interp_kaiser"]: + raise ValueError("Invalid resampling method: {}".format(resampling_method)) + + orig_freq = int(orig_freq) // gcd + new_freq = int(new_freq) // gcd + + if lowpass_filter_width <= 0: + raise ValueError("Low pass filter width should be positive.") + base_freq = min(orig_freq, new_freq) + # This will perform antialiasing filtering by removing the highest frequencies. + # At first I thought I only needed this when downsampling, but when upsampling + # you will get edge artifacts without this, as the edge is equivalent to zero padding, + # which will add high freq artifacts. + base_freq *= rolloff + + # The key idea of the algorithm is that x(t) can be exactly reconstructed from x[i] (tensor) + # using the sinc interpolation formula: + # x(t) = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - t)) + # We can then sample the function x(t) with a different sample rate: + # y[j] = x(j / new_freq) + # or, + # y[j] = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - j / new_freq)) + + # We see here that y[j] is the convolution of x[i] with a specific filter, for which + # we take an FIR approximation, stopping when we see at least `lowpass_filter_width` zeros crossing. + # But y[j+1] is going to have a different set of weights and so on, until y[j + new_freq]. + # Indeed: + # y[j + new_freq] = sum_i x[i] sinc(pi * orig_freq * ((i / orig_freq - (j + new_freq) / new_freq)) + # = sum_i x[i] sinc(pi * orig_freq * ((i - orig_freq) / orig_freq - j / new_freq)) + # = sum_i x[i + orig_freq] sinc(pi * orig_freq * (i / orig_freq - j / new_freq)) + # so y[j+new_freq] uses the same filter as y[j], but on a shifted version of x by `orig_freq`. + # This will explain the F.conv1d after, with a stride of orig_freq. + width = math.ceil(lowpass_filter_width * orig_freq / base_freq) + # If orig_freq is still big after GCD reduction, most filters will be very unbalanced, i.e., + # they will have a lot of almost zero values to the left or to the right... + # There is probably a way to evaluate those filters more efficiently, but this is kept for + # future work. + idx_dtype = dtype if dtype is not None else torch.float64 + + idx = torch.arange(-width, width + orig_freq, dtype=idx_dtype, device=device)[None, None] / orig_freq + + t = torch.arange(0, -new_freq, -1, dtype=dtype, device=device)[:, None, None] / new_freq + idx + t *= base_freq + t = t.clamp_(-lowpass_filter_width, lowpass_filter_width) + + # we do not use built in torch windows here as we need to evaluate the window + # at specific positions, not over a regular grid. + if resampling_method == "sinc_interp_hann": + window = torch.cos(t * math.pi / lowpass_filter_width / 2) ** 2 + else: + # sinc_interp_kaiser + if beta is None: + beta = 14.769656459379492 + beta_tensor = torch.tensor(float(beta)) + window = torch.i0(beta_tensor * torch.sqrt(1 - (t / lowpass_filter_width) ** 2)) / torch.i0(beta_tensor) + + t *= math.pi + + scale = base_freq / orig_freq + kernels = torch.where(t == 0, torch.tensor(1.0).to(t), t.sin() / t) + kernels *= window * scale + + if dtype is None: + kernels = kernels.to(dtype=torch.float32) + + return kernels, width + + +def _apply_sinc_resample_kernel( + waveform: torch.Tensor, + orig_freq: int, + new_freq: int, + gcd: int, + kernel: torch.Tensor, + width: int, +): + if not waveform.is_floating_point(): + raise TypeError(f"Expected floating point type for waveform tensor, but received {waveform.dtype}.") + + orig_freq = int(orig_freq) // gcd + new_freq = int(new_freq) // gcd + + # pack batch + shape = waveform.size() + waveform = waveform.view(-1, shape[-1]) + + num_wavs, length = waveform.shape + waveform = torch.nn.functional.pad(waveform, (width, width + orig_freq)) + resampled = torch.nn.functional.conv1d(waveform[:, None], kernel, stride=orig_freq) + resampled = resampled.transpose(1, 2).reshape(num_wavs, -1) + target_length = torch.ceil(torch.as_tensor(new_freq * length / orig_freq)).long() + resampled = resampled[..., :target_length] + + # unpack batch + resampled = resampled.view(shape[:-1] + resampled.shape[-1:]) + return resampled + + +# ============================================================================== +# Swish — exact copy of nemo/collections/asr/parts/utils/activations.py +# ============================================================================== +class Swish(nn.SiLU): + """ + Swish activation function introduced in 'https://arxiv.org/abs/1710.05941' + Mathematically identical to SiLU. See note in nn.SiLU for references. + """ + + +# ============================================================================== +# activation_registry — exact copy of nemo/collections/common/parts/utils.py +# ============================================================================== +activation_registry = { + "identity": nn.Identity, + "hardtanh": nn.Hardtanh, + "relu": nn.ReLU, + "selu": nn.SELU, + "swish": nn.SiLU, + "silu": nn.SiLU, + "gelu": nn.GELU, +} + + +# ============================================================================== +# FusedBatchNorm1d — exact copy of nemo/collections/asr/parts/submodules/batchnorm.py +# ============================================================================== +class FusedBatchNorm1d(nn.Module): + """ + Fused BatchNorm to use in Conformer to improve accuracy in finetuning with TTS scenario + Drop-in replacement for BatchNorm1d with simple affine projection + """ + + def __init__(self, num_features: int): + """ + Args: + num_features: number of channels, see original BatchNorm1d documentation + """ + super().__init__() + self.num_features = num_features + self.weight = nn.Parameter(torch.ones(num_features)) + self.bias = nn.Parameter(torch.zeros(num_features)) + + def forward(self, x: torch.Tensor): + if x.dim() == 3: + return x * self.weight.unsqueeze(-1) + self.bias.unsqueeze(-1) + assert x.dim() == 2 + return x * self.weight + self.bias + + +# ============================================================================== +# load_pretrained_hf / set_model_dict_for_partial_init / load_checkpoint — +# exact copies of nemo/collections/speechlm2/parts/pretrained.py +# (only the functions used by the kept inference paths were vendored) +# ============================================================================== +def load_pretrained_hf( + model_path_or_name: str, pretrained_weights: bool = True, dtype=torch.float32, trust_remote_code: bool = True +): + """ + Load pretrained HuggingFace AutoModelForCausalLM. + + Setting ``pretrained_weights=False`` returns a model that has identical architecture with the checkpoint, + but is randomly initialized. + + Args: + model_path_or_name: Path or name of the model to load + pretrained_weights: Whether to load pretrained weights (True) or random init (False) + dtype: Data type for the model + trust_remote_code: Whether to trust remote code when loading model (needed for some models like Nemotron) + """ + if pretrained_weights: + return AutoModelForCausalLM.from_pretrained( + model_path_or_name, torch_dtype=dtype, trust_remote_code=trust_remote_code + ) + else: + config = AutoConfig.from_pretrained(model_path_or_name, trust_remote_code=trust_remote_code) + return AutoModelForCausalLM.from_config(config, torch_dtype=dtype, trust_remote_code=trust_remote_code) + + +def set_model_dict_for_partial_init( + pretrained_dict: Dict[str, torch.Tensor], model_dict: Dict[str, torch.Tensor] +) -> Dict[str, torch.Tensor]: + """ + Partially initialize a model's state dictionary with a pretrained state dictionary. + This function safely copies compatible layers from a pretrained model into a new model, + ignoring layers with mismatched shapes or missing keys. + + Steps: + 1. Remove layers from the pretrained dictionary if their shape does not match the target model. + 2. Keep only keys that exist in the target model. + 3. Update the model dictionary with the filtered pretrained weights. + + Args: + pretrained_dict (Dict[str, torch.Tensor]): + The state dictionary of the pretrained model. + model_dict (Dict[str, torch.Tensor]): + The state dictionary of the target model to be partially initialized. + + Returns: + Dict[str, torch.Tensor]: + The updated model state dictionary with compatible layers loaded from the pretrained dictionary. + """ + # 1. Remove layers where pretrained shape differs from model shape + for k, v in list(pretrained_dict.items()): + if k in model_dict and hasattr(model_dict[k], "numel") and v.numel() != model_dict[k].numel(): + del pretrained_dict[k] + logging.info(f" | > Layer with shape mismatch in the model definition: {k}") + + # 2. Keep only keys that exist in the target model + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} + + # 3. Update model dictionary with filtered pretrained layers + model_dict.update(pretrained_dict) + logging.info(f" | > {len(pretrained_dict)} / {len(model_dict)} layers are restored.") + + return model_dict + + +def load_checkpoint(checkpoint_path): + """ + Load a model checkpoint from disk. + + Supports loading checkpoints stored in either PyTorch (`.ckpt`, `.pt`) or + SafeTensors (`.safetensors`) formats. All parameters are loaded onto CPU + regardless of the original device. + + Args: + checkpoint_path (str): + Path to the checkpoint file. If the filename contains `.safetensors`, + it is loaded using the SafeTensors backend; otherwise, it is assumed + to be a PyTorch checkpoint containing a `state_dict` field. + + Returns: + dict: + A state dictionary mapping parameter names to tensors. + """ + if ".safetensors" in checkpoint_path: + checkpoint_state = load_file(checkpoint_path, device="cpu") + else: + checkpoint_state = torch.load(checkpoint_path, map_location="cpu")["state_dict"] + return checkpoint_state diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py new file mode 100644 index 000000000..7c674af5c --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py @@ -0,0 +1,1038 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/models/duplex_ear_tts.py; training-only code removed. +# Reduced for inference: base class changed LightningModule+HFHubMixin -> nn.Module; +# removed training_step/validation_step/test_step, epoch hooks, log_model_stats, +# run_evaluation_one_batch, configure_optimizers/backward/configure_model (FSDP/TP), +# oomptimizer_schema, and the lightning/peft/hydra/metrics imports they required. + +import os +import time +from collections import Counter +from contextlib import contextmanager + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from omegaconf import DictConfig + +from .compat import ( + fp32_precision, + get_pad_id, + load_checkpoint, + load_pretrained_hf, + logging, + set_model_dict_for_partial_init, +) +from .ear_tts_model import RVQEARTTSModel +from .ear_tts_vae_codec import RVQVAEModel +from .tokenizer import AutoTokenizer + + +class DuplexEARTTS(nn.Module): + def __init__(self, cfg: dict) -> None: + assert isinstance(cfg, dict), ( + "You must pass the config to DuplexEARTTS as a Python dict to support hyperparameter serialization " + f"in PTL checkpoints (we got: '{type(cfg)=}')." + ) + super().__init__() + # convert dict to config + cfg = DictConfig(cfg) + self.trainer_config = cfg.get("trainer", None) + self.data_cfg = cfg.data + self.cfg = cfg.model + self.target_sample_rate = cfg.data.target_sample_rate + self.source_sample_rate = cfg.data.source_sample_rate + self.normalize_text = cfg.data.get("normalize_text", False) + + # exp_manager is training-only; keep validation_save_path optional for inference configs. + if cfg.get("exp_manager", None) is not None: + self.validation_save_path = os.path.join(cfg.exp_manager.explicit_log_dir, "validation_logs") + else: + self.validation_save_path = None + + # move back text channel by x, in inference it advance the text channel prediction by x frames + self.advance_text_channel_by = self.cfg.get("advance_text_channel_by", None) + + # Load ForCausalLM + if self.cfg.tts_config.context_hidden_size is not None: + self.language_model = self._load_language_model(self.cfg) + self.embed_tokens = self._load_embed_tokens(self.cfg) + # delete llm because we use it only to get the embedding tokens + del self.language_model + + # get codec run precision + self.audio_codec_run_dtype = getattr(torch, self.cfg.get("audio_codec_run_dtype", "float32"), torch.float32) + + # Load tokenizer + self.tokenizer = AutoTokenizer( + self.cfg.pretrained_lm_name, + use_fast=True, + trust_remote_code=True, + bos_token=self.cfg.get("bos_token", None), + eos_token=self.cfg.get("eos_token", None), + pad_token=self.cfg.get("pad_token", None), + ) # Note that we are using fast tokenizer + + # Instantiate TTS model + self.tts_model = RVQEARTTSModel(DictConfig(self.cfg.tts_config), tokenizer=self.tokenizer) + # Load and initialize audio codec, and bind RVQ embeddings to the TTS model + setup_audio_codec(self) + + self._codebook_size = self.tts_model.config.codebook_size + + # compute samples per frame + self.source_samples_per_frame = int(self.source_sample_rate * cfg.data.frame_length) + + # get codec silence tokens + codec_silence_tokens = self.get_codec_silence_frame() + self.register_buffer("codec_silence_tokens", codec_silence_tokens) + + # cached for quicker audio decoding + self.register_buffer( + "_control_codes", + torch.tensor([self.speech_bos_id, self.speech_eos_id, self.speech_pad_id], device=self.device), + ) + + self._use_fsdp = False + self._use_tp = False + self.audio_prompt_latents = nn.ParameterDict() + + def get_codec_silence_frame(self): + # Generate long zero waveform (silence) + audio = torch.zeros(1, 10 * self.target_sample_rate).float().to(self.device) + audio_len = torch.tensor([audio.size(-1)]).long() + audio, audio_len = self.pad_audio_to_factor(audio, audio_len, self.target_samples_per_frame) + + with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): + sil_codes, _ = self.audio_codec.encode(audio.unsqueeze(1), audio_len) # [1, T, C] + sil_codes = sil_codes[0] # [T, C] + + # Convert each frame (C tokens) into a tuple + combos = [tuple(row.tolist()) for row in sil_codes] + + # Count frequencies + counter = Counter(combos) + + # Pick the most common combination + most_common_combo, freq = counter.most_common(1)[0] + + # Return as tensor [C] + return torch.tensor(most_common_combo, device=self.device, dtype=torch.long) + + def _load_embed_tokens(self, cfg) -> nn.Embedding: + """Load token embedding layer for RVQ-EAR-TTS.""" + if self.language_model: + assert callable(self.language_model.get_input_embeddings) + embed_tokens: nn.Embedding = self.language_model.get_input_embeddings() + else: + embed_tokens_state_dict = torch.load( + cfg.pretrained_lm_embedding_path, map_location="cpu", weights_only=True + ) + + # Create token embedding layer + vocab_size, hidden_size = embed_tokens_state_dict["weight"].size() + embed_tokens = nn.Embedding(vocab_size, hidden_size, dtype=torch.bfloat16) + embed_tokens.load_state_dict(embed_tokens_state_dict) + return embed_tokens + + def _load_language_model(self, cfg): + """Load language model for RVQ-EAR-TTS.""" + if cfg.pretrained_lm_name: + language_model = load_pretrained_hf( + self.cfg.pretrained_lm_name, pretrained_weights=True, trust_remote_code=True + ).eval() + else: + language_model = None + return language_model + + @property + def device(self): + return next(self.parameters()).device + + @property + def speech_vocab_size(self): + """Return the size of the audio codec codebook including extra speech BOS and EOS tokens.""" + if self.use_local_transformer and self.local_transformer_type == "nar": # add extra token for mask + return self._codebook_size + 4 + return self._codebook_size + 3 + + @property + def speech_bos_id(self) -> int: + """Indicates start of utterance generation (not start of inference!).""" + if self.cfg.get("custom_speech_bos_id", None): + return self.cfg.get("custom_speech_bos_id") + return self._codebook_size + 2 + + @property + def speech_eos_id(self) -> int: + """Indicates end of utterance generation.""" + if self.cfg.get("custom_speech_eos_id", None): + return self.cfg.get("custom_speech_eos_id") + return self._codebook_size + 1 + + @property + def speech_pad_id(self) -> int: + """Indicates start of inference (the very first frame).""" + if self.cfg.get("custom_speech_pad_id", None): + return self.cfg.get("custom_speech_pad_id") + return self._codebook_size + + @property + def text_vocab_size(self): + """Return the size of the text tokenizer.""" + return self.tokenizer.vocab_size + + @property + def text_bos_id(self) -> int: + return self.tokenizer.bos_id + + @property + def text_eos_id(self) -> int: + return self.tokenizer.eos_id + + @property + def text_pad_id(self) -> int: + """ + Text pad ID is used as a 'blank' for frames when the model is not speaking + and for frames where the model is speaking but has already predicted the + entire text channel's content. + + Example: + + flow: |---user---||-------assistant--------||-user-| + text channel: 0000000000 1xxxxxxx0000000000000002 000000 + + Where 0 indicates PAD ID, 1 indicates BOS ID, 2 indacates EOS ID, + and x indicates tokens corresponding to actual text + + """ + return get_pad_id(self.tokenizer) + + def pad_audio_to_factor(self, audio, audio_len, samples_per_frame, downsampling_factor: int = 1): + """ + Zero pad the end of the audio so that we do not have a partial end frame. + The output will be zero-padded to have an integer number of frames of + length `samples_per_frame * downsampling_factor`. + + Args: + audio: input time-domain signal (B, T) + audio_len: valid length for each example in the batch (B,) + samples_per_frame: number of samples per frame + downsampling_factor: how much each frame is downsampled in later processing + + Returns: + padded_audio: Padded time-domain signal (B, T') + padded_len: Adjusted valid lengths (B,) + """ + with fp32_precision(): + total_factor = samples_per_frame * downsampling_factor + padded_len = total_factor * torch.ceil(audio_len / total_factor).int() + max_len = padded_len.max().int().item() + num_padding = max_len - audio.shape[1] + padded_audio = F.pad(audio, (0, num_padding)) + return padded_audio, padded_len + + def prepare_inputs(self, batch: dict): + """ + Prepare inputs, extracting audio tokens and padding if needed. + """ + # check if audios has the same batch size + assert batch["source_audio"].size(0) == batch["target_audio"].size(0) + assert batch["audio_prompt"].size(0) == batch["target_audio"].size(0) + + target_audio = batch["target_audio"] + target_audio_lens = batch["target_audio_lens"] + target_text_tokens = batch["target_text_tokens"] + non_prompt_mask = batch["non_prompt_mask"] + aligned_attention_mask = batch["aligned_attention_mask"] + aligned_position_ids = batch["aligned_position_ids"] + + if self.training and (self.cfg.get("empty_turn_probability", 0.0) > 0): + # Randomly decide whether this batch gets emptied + if torch.rand(1).item() < self.cfg.empty_turn_probability: + # Zero out audio + target_audio = torch.zeros_like(target_audio) + + # Create mask for tokens we want to drop + # Keep BOS and EOS, drop the rest. + keep_mask = (target_text_tokens == self.text_bos_id) | (target_text_tokens == self.text_eos_id) + full_dropout_mask = ~keep_mask # True = positions to replace with PAD + + # Replace all non-BOS/EOS with PAD + target_text_tokens = torch.where( + full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens + ) + + # extract target audio codes + target_audio, target_audio_lens = self.pad_audio_to_factor( + target_audio, target_audio_lens, self.target_samples_per_frame, 1 + ) + with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): + target_codes, target_codes_lens = self.audio_codec.encode(target_audio.unsqueeze(1), target_audio_lens) + + with fp32_precision(): + target_len = target_codes.shape[1] + + # Pad or truncate sequence variables + def pad_or_truncate(x, pad_value=0): + if x.dim() == 2: # [B, T] + L = x.shape[1] + if L < target_len: + return F.pad(x, (0, target_len - L), value=pad_value) + else: + return x[:, :target_len] + return x # leave others for now + + target_text_tokens = pad_or_truncate(target_text_tokens, pad_value=self.text_pad_id) + non_prompt_mask = pad_or_truncate(non_prompt_mask, pad_value=0) + aligned_position_ids = pad_or_truncate(aligned_position_ids, pad_value=0) + + # Correct attention mask padding/truncation + B, H, L1, L2 = aligned_attention_mask.shape + new_len = target_len + if L1 < new_len or L2 < new_len: + pad_rows = new_len - L1 + pad_cols = new_len - L2 + aligned_attention_mask = F.pad(aligned_attention_mask, (0, pad_cols, 0, pad_rows)) + elif L1 > new_len or L2 > new_len: + aligned_attention_mask = aligned_attention_mask[:, :, :new_len, :new_len] + + # set the pad token for the first BOS frame + target_codes_aligned = target_codes.clone() + target_codes_aligned[:, 0] = self.speech_pad_id + + # set special token in the last audio prompt (it will works as a BOS token) + pos = non_prompt_mask.float().argmax(dim=1) # shape: [B] + row_idx = torch.arange(B, device=self.device) + # set the extra self.speech_pad_id at first 1 position in non_prompt_mask + target_codes_aligned[row_idx, pos] = self.speech_pad_id + + # EOS dropout to make the model more robust + if self.training and self.cfg.get("text_eos_dropout_prob", 0.0) > 0: + # Mask EOS positions + eos_mask = target_text_tokens == self.text_eos_id + + # Random dropout only on EOS positions + dropout_mask = torch.rand(eos_mask.sum(), device=target_text_tokens.device) < self.cfg.text_eos_dropout_prob + + # Scatter dropout decisions into [B, T] + full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool) + full_dropout_mask[eos_mask] = dropout_mask + + # Replace dropped EOS with PAD + target_text_tokens = torch.where( + full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens + ) + + # BOS dropout to make the model more robust + if self.training and self.cfg.get("text_bos_dropout_prob", 0.0) > 0: + # Mask BOS positions + bos_mask = target_text_tokens == self.text_bos_id + + # Random dropout only on BOS positions + dropout_mask = torch.rand(bos_mask.sum(), device=target_text_tokens.device) < self.cfg.text_bos_dropout_prob + + # Scatter dropout decisions into [B, T] + full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool) + full_dropout_mask[bos_mask] = dropout_mask + + # Replace dropped BOS with PAD + target_text_tokens = torch.where( + full_dropout_mask, + torch.full_like(target_text_tokens, self.text_pad_id), + target_text_tokens, + ) + + # BOS dropout to make the model more robust + if self.training and self.cfg.get("text_bos_dropout_prob", 0.0) > 0: + prob = self.cfg.text_bos_dropout_prob # e.g., 0.5 + + # Identify all BOS positions [B, T] + bos_mask = target_text_tokens == self.text_bos_id + + # Get indices of sequences that actually have a BOS token + # We need to know *where* the BOS tokens are to drop them. + # tensor of coordinates: [[batch_idx, seq_idx], ...] + bos_indices = torch.nonzero(bos_mask) + num_bos = bos_indices.shape[0] + + if num_bos > 0: + # Create a random dropout decision for each BOS instance + drop_decisions = torch.rand(num_bos, device=target_text_tokens.device) < prob + + # Ensure at least one is dropped + if drop_decisions.sum() == 0: + # Pick one random index from the available BOS locations to drop + force_idx = torch.randint(0, num_bos, (1,), device=target_text_tokens.device) + drop_decisions[force_idx] = True + + # 5. Apply the dropout + # We need to map the decisions back to the full tensor + # Create a mask of the same shape as target_text_tokens + full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool) + + # Set True only at the specific (batch, seq) coordinates we chose to drop + # bos_indices[:, 0] are batch indices, bos_indices[:, 1] are seq indices + full_dropout_mask[bos_indices[:, 0], bos_indices[:, 1]] = drop_decisions + + # 6. Replace dropped BOS with PAD + target_text_tokens = torch.where( + full_dropout_mask, + torch.full_like(target_text_tokens, self.text_pad_id), + target_text_tokens, + ) + + # shift text tokens + subword_ids = F.pad(target_text_tokens[:, 1:], [0, 1]) + # note that we are using a text mask where we are ignoring the desc + audio prompt but we are keeping 1 until the audio ends to support duplex + subword_mask = F.pad(non_prompt_mask[:, 1:], [0, 1]) + + # detach embedding as in eartts + if self.cfg.tts_config.context_hidden_size is not None: + context_hidden_state = self.embed_tokens(target_text_tokens).detach() + else: + context_hidden_state = None + + if self._use_tp: + tp_world_size = self.device_mesh["tensor_parallel"].size() + if (remainder := (target_text_tokens.shape[1] - 1) % tp_world_size) != 0: + target_text_tokens = target_text_tokens[:, :-remainder] + target_codes_aligned = target_codes_aligned[:, :-remainder] + target_codes_aligned = target_codes_aligned[:, :-remainder] + subword_ids = subword_ids[:, :-remainder] + subword_mask = subword_mask[:, :-remainder] + + return { + "code": target_codes_aligned, + "audio_mask": non_prompt_mask, # set audio_mask as non_prompt_mask to avoid the audio prompt in loss computation + "attention_mask": aligned_attention_mask, + "position_ids": aligned_position_ids, + "subword_ids": subword_ids, + "subword_mask": subword_mask, + "context_hidden_state": context_hidden_state, + "output_lens": target_codes_lens, + "non_prompt_mask": non_prompt_mask, + "target_text_tokens": target_text_tokens, + } + + def _get_generation_config(self, guidance_enabled: bool = False): + """Get default generation config for EAR-TTS.""" + return { + "num_iter": 8, + "guidance_scale": self.cfg.get("inference_guidance_scale", 0.5) if guidance_enabled else None, + "top_p_or_k": self.cfg.get("inference_top_p_or_k", 0.8), + "noise_scale": self.cfg.get("inference_noise_scale", 0.8), + "eos_threshold": -3.0, + } + + def get_audio_prompt_latent(self, name, B): + """ + Retrieve a cached audio prompt latent and adapt it to the requested batch size. + + This fetches a latent previously cached via `set_audio_prompt_latent()` and + ensures the returned tensor has batch size `B` by: + - returning as-is when the batch already matches, + - truncating when the cached batch is larger than `B`, + - expanding when the cached batch is smaller (commonly batch=1 warmup). + + Args: + name (str): + Key of the cached audio prompt latent to retrieve. + B (int): + Desired batch size. + + Returns: + Tensor: + Audio prompt latent with batch dimension equal to `B`, moved to `self.device`. + Shape: [B, ..., D] + + Raises: + KeyError: + If `name` does not exist in `self.audio_prompt_latents`. + """ + if not hasattr(self, "audio_prompt_latents") or name not in self.audio_prompt_latents: + raise KeyError(f"Unknown audio prompt latent '{name}'. Call set_audio_prompt_latent(...) first.") + + audio_prompt_latent = self.audio_prompt_latents[name] # cached on CPU + + if audio_prompt_latent.shape[0] == B: + out = audio_prompt_latent + elif audio_prompt_latent.shape[0] >= B: + out = audio_prompt_latent[:B] + else: + out = audio_prompt_latent[:1].expand(B, *audio_prompt_latent.shape[1:]) + + return out.to(self.device) + + def set_init_inputs(self, speaker_audio=None, speaker_audio_lens=None, system_prompt=None, speaker_name=None): + """ + Registers and prepares initial input buffers for text/audio prompt and context, to warm up AR inference. + + Args: + speaker_audio (torch.Tensor): Batch of prompt audio, (B, T). + speaker_audio_lens (torch.Tensor): Lengths for each sample in speaker_audio, (B,). + system_prompt (str, optional): System prompt for context. + speaker_name (str, optional): Name of a pre-baked speaker latent in audio_prompt_latents. + When provided, speaker_audio/speaker_audio_lens are ignored and silent audio is used + as the prompt carrier; the actual voice identity comes from the cached latent. + + Returns: + dict: Dictionary of input tensors to be passed to inference, with registered buffers. + """ + # compute prompt audio size and slice it + with fp32_precision(): + # compute the exact number of samples for the prompt duration + prompt_audio_size = int( + ((self.data_cfg.audio_prompt_duration * self.target_sample_rate) // self.target_samples_per_frame) + * self.target_samples_per_frame + ) + + if speaker_name is not None: + logging.info( + f"set_init_inputs: using pre-baked latent for speaker '{speaker_name}' (silent carrier audio)" + ) + speaker_audio = torch.zeros((1, prompt_audio_size), device=self.device, dtype=torch.float32) + speaker_audio_lens = torch.LongTensor([speaker_audio.shape[1]]).to(self.device) + + B, T = speaker_audio.shape + device = speaker_audio.device + dtype = speaker_audio.dtype + + # allocate result + prompt_audio = torch.zeros(B, prompt_audio_size, device=device, dtype=dtype) + + # process each example independently + for b in range(B): + valid_len = min(speaker_audio_lens[b].item(), T) + + # handle empty + if valid_len <= 0: + continue + + # valid (non-padded) segment + valid_segment = speaker_audio[b, :valid_len] + + if valid_len >= prompt_audio_size: + # enough valid audio → crop from start (no silence) + prompt_audio[b] = valid_segment[:prompt_audio_size] + else: + # too short → repeat and crop + repeat_factor = (prompt_audio_size + valid_len - 1) // valid_len # ceil division + expanded = valid_segment.repeat(repeat_factor) + prompt_audio[b] = expanded[:prompt_audio_size] + + # add a silence in the end to smooth the transition between prompt and audio tokens + prompt_audio[:, -int(self.target_samples_per_frame * 2) :] = 0 + + # get prompt audio size + with fp32_precision(): + prompt_audio_text_pad_size = int(prompt_audio_size // self.target_samples_per_frame) + + # create a eos token id + if system_prompt is not None and self.cfg.get("use_system_prompt", None) and system_prompt != "": + text_prompt = torch.as_tensor( + [self.tokenizer.bos] + self.tokenizer.text_to_ids(system_prompt) + [self.tokenizer.eos], + dtype=torch.long, + device=self.device, + ) + else: + text_prompt = torch.tensor([self.tokenizer.eos], dtype=torch.long, device=self.device) + + # create a padding tensor + prompt_audio_text_pad = ( + torch.ones(prompt_audio_text_pad_size, device=self.device, dtype=text_prompt.dtype) * self.text_pad_id + ) + prompt_audio_text_pad[-1] = self.tokenizer.eos + + # Prepend an initial text EOS token followed by padding tokens that match + # the number of audio-prompt frames (in text-token units). + target_text_tokens = torch.cat([text_prompt, prompt_audio_text_pad.to(text_prompt.dtype)]) + + # create pad audio for the description + pad_size = text_prompt.size(-1) * self.target_samples_per_frame + pad_audio = ( + torch.zeros(pad_size, device=prompt_audio.device, dtype=prompt_audio.dtype) + .unsqueeze(0) + .repeat(prompt_audio.size(0), 1) + ) + + # repeat to reaches the batch size + target_text_tokens = target_text_tokens.unsqueeze(0).repeat(prompt_audio.size(0), 1) + target_audio = torch.cat([pad_audio, prompt_audio], dim=1) + + # extract code codes + target_audio_len = torch.tensor( + [target_audio.size(-1)] * target_audio.size(0), dtype=torch.long, device=self.device + ) + with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): + code, _ = self.audio_codec.encode(target_audio.unsqueeze(1), target_audio_len) + + # get context hidden + if self.cfg.tts_config.context_hidden_size is not None: + context_hidden_state = self.embed_tokens(target_text_tokens) + else: + context_hidden_state = None + + # create masks + # non_prompt_mask is all zeros, because all processed is prompt + non_prompt_mask = torch.zeros_like(target_text_tokens) + non_prompt_mask[:, -2:] = 1 # set last valid prompt frame as 1 to allow the addition of BOS in the right place + subword_mask = torch.zeros_like( + target_text_tokens + ) # subword_mask is almost all zeros because on the warmup there is only the prompt + subword_mask[:, -3:] = ( + 1 # -3 because of the it start right after the first valid prompt token and it is shifted by 1 + ) + + # set the pad token for the first BOS frame + code[:, 0] = self.speech_pad_id + + # shift subword_ids + subword_ids = F.pad(target_text_tokens[:, 1:], [0, 1], value=0.0) + + # set special token in the last audio prompt (it will works as a BOS token) + pos = non_prompt_mask.float().argmax(dim=1) # shape: [B] + row_idx = torch.arange(B, device=self.device) + # set the extra self.speech_pad_id at first 1 position in non_prompt_mask + code[row_idx, pos] = self.speech_pad_id + + init_inputs = { + "code": code[:, :-1], + "audio_mask": non_prompt_mask.bool()[ + :, :-1 + ], # set audio_mask as non_prompt_mask to avoid the audio prompt in loss computation + "context_hidden_state": context_hidden_state[:, :-1] if context_hidden_state is not None else None, + "subword_ids": subword_ids[:, :-1], + "subword_mask": subword_mask.bool()[:, :-1], + "non_prompt_mask": non_prompt_mask.bool()[:, :-1], + } + + if speaker_name is not None: + init_inputs["audio_prompt_latent"] = self.get_audio_prompt_latent(speaker_name, B=1) + + self._init_input_cache = {} + for k, v in init_inputs.items(): + if v is None: + self._init_input_cache[k] = None + else: + # clone → removes inference tensor flag + # detach → ensures no graph refs + self._init_input_cache[k] = v.detach().clone() + + return init_inputs + + def get_init_inputs( + self, + B: int, + init_inputs_names=[ + "code", + "audio_mask", + "context_hidden_state", + "subword_ids", + "subword_mask", + "non_prompt_mask", + ], + ): + """ + Returns a dictionary of initial inputs for inference, using registered buffers. + + Args: + B (int): Required batch size. + init_inputs_names (List[str], optional): Names of input buffers to fetch. + + Returns: + dict: Each key is name from init_inputs_names, and value is tensor of appropriate shape (B, ...). + + Notes: + Expands batch-1 buffers to B if necessary. + """ + if init_inputs_names is None: + init_inputs_names = [ + "code", + "audio_mask", + "context_hidden_state", + "subword_ids", + "subword_mask", + "non_prompt_mask", + ] + + if self._init_input_cache.get("audio_prompt_latent", None) is not None: + init_inputs_names = list(init_inputs_names) + ["audio_prompt_latent"] + + init_inputs = {} + for name in init_inputs_names: + buf = self._init_input_cache.get(name, None) + + if buf is None: + init_inputs[name] = None + continue + + # Batch already matches + if buf.shape[0] == B: + init_inputs[name] = buf + elif buf.shape[0] >= B: + init_inputs[name] = buf[:B] + else: + # assume batch=1 warmup → expand + init_inputs[name] = buf[:1].expand(B, *buf.shape[1:]) + + return init_inputs + + @torch.inference_mode() + def infer_codes_one_step( + self, + current_subword_id, + prev_subword_id, + current_subword_mask, + prev_audio_tokens, + past_key_values, + guidance_enabled=True, + generation_config=None, + ignore_eos_flag_stop=True, + request_id=None, + ): + """ + Runs a single autoregressive prediction step to infer audio codec codes. + + Args: + current_subword_id (torch.Tensor): Current text token IDs, shape (B, 1). + prev_subword_id (torch.Tensor): Previous text token IDs, shape (B, 1). + current_subword_mask (torch.Tensor): Current mask, shape (B, 1). + prev_audio_tokens (torch.Tensor): Previously generated audio tokens, shape (B, 1, C). + past_key_values: Key-value cache for transformer decoder state. + guidance_enabled (bool, optional): Enables classifier-free guidance. + generation_config (dict, optional): Generation hyperparameters. + ignore_eos_flag_stop (bool): If True, ignore EOS flag for stopping. + + Returns: + Tuple[torch.Tensor, Any]: + - Predicted audio codec token(s), shape (B, 1, C) + - Updated past_key_values for the next step. + """ + + if self.cfg.tts_config.context_hidden_size is not None: + # get context_hidden_state it is always one step behind current_subword_id + # for the first step uses the last step from warmup + context_hidden_state = self.embed_tokens(prev_subword_id) + else: + context_hidden_state = None + + # force silence as next token + if self.cfg.get("inference_force_speech_silence_on_eos", True): + silence_codes = self.codec_silence_tokens.view(1, 1, -1).expand(prev_audio_tokens.shape) + prev_audio_tokens = torch.where( + current_subword_id.unsqueeze(-1) == self.text_eos_id, + silence_codes, # silence + prev_audio_tokens, # keep original + ) + + # get subword_ids + inputs = { + "code": prev_audio_tokens, + "context_hidden_state": context_hidden_state, + "subword_ids": current_subword_id, + "subword_mask": current_subword_mask, + "past_key_values": past_key_values, + "use_cache": True, + "guidance_enabled": guidance_enabled, + "generation_config": generation_config, + "ignore_eos_flag_stop": ignore_eos_flag_stop, + } + # request_id is only used by vLLM backend — skip for native PyTorch inference + if request_id is not None: + inputs["request_id"] = request_id + + outputs = self.tts_model(**inputs) + + return outputs["codes"], outputs["past_key_values"] + + @torch.inference_mode() + def decode_one_audio_step(self, gen_audio_codes_history, number_prev_tokens=None): + """ + Decodes one step of generated audio codec tokens to raw waveform. + + Args: + gen_audio_codes_history (torch.Tensor): Audio tokens history, shape (B, T, C). + number_prev_tokens (int, optional): Number of previous tokens to decode, for incremental decoding. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + - audio_pred_cur_step: Latest decoded waveform chunk, shape (B, wav_to_token_ratio). + - audio_len: Lengths (number of samples), shape (B,). + """ + with fp32_precision(), torch.no_grad(): + if number_prev_tokens: + gen_audio_codes_history = gen_audio_codes_history[:, -number_prev_tokens:] + + gen_audio_codes_history = replace_control_speech_codes( + gen_audio_codes_history, self._control_codes, self.codec_silence_tokens + ) + gen_audio_codes_lens = torch.tensor( + [gen_audio_codes_history.size(1)] * gen_audio_codes_history.size(0), device=self.device + ) + audio_pred, audio_len = self.audio_codec.decode(gen_audio_codes_history, gen_audio_codes_lens) + + # return only the current/lastest audio chunk + audio_pred_cur_step = audio_pred.squeeze(1)[:, -self.audio_codec.config.wav_to_token_ratio :] + audio_len[:] = self.audio_codec.config.wav_to_token_ratio + return audio_pred_cur_step, audio_len + + @torch.inference_mode() + def offline_inference( + self, + next_subword_ids: torch.Tensor, + init_inputs: dict, + formatter: str = "", + guidance_enabled: bool = True, + generation_config: dict = None, + incremental_audio_decoding: bool = False, + ) -> dict[str, torch.Tensor]: + """ + Runs offline autoregressive inference for the Duplex EAR-TTS speech decoder. + + This method performs **text-to-speech (TTS)** generation: given subword/text + tokens and prompt-initialization states, it autoregressively generates + audio codec tokens and decodes them into a waveform. + + Args: + next_subword_ids (torch.Tensor): + Conditioning subword/text token IDs for the speech decoder. + Shape: (B, T_text). + + init_inputs (dict): + Dictionary of prompt-dependent initial states produced by + ``get_init_inputs()``. May include: + + • "code" — initial audio tokens (e.g., prompt audio) + • "audio_mask" — mask for prompt audio positions + • "context_hidden_state" — decoder hidden state at t = 0 + • "subword_ids" — prompt text tokens + • "subword_mask" — mask for prompt text + • "non_prompt_mask" — mask marking positions to be generated + + ``get_init_inputs()`` automatically expands batch-1 buffers to + batch size B. + + formatter (str, optional): + Optional formatter identifier used to customize the prompt structure. + + guidance_enabled (bool, optional): + Whether classifier-free guidance (CFG) is enabled. + If enabled and ``generation_config`` is ``None``, guidance parameters + are taken from ``_get_generation_config()``. + + generation_config (dict, optional): + Settings controlling autoregressive generation, including sampling + strategy, noise scale, refinement iterations, and EOS rules. + If ``None``, defaults are taken from + ``_get_generation_config(guidance_enabled)``. + + incremental_audio_decoding (bool, optional): + If True, codec-to-waveform decoding is performed incrementally during + autoregressive generation. + If False, waveform decoding occurs only after all audio tokens are produced. + + Returns: + dict[str, torch.Tensor]: + Contains: + + • **"audio"**: + Generated waveform of shape ``(B, T_audio)``, obtained via + ``audio_pred.squeeze(1)``. + + • **"audio_len"**: + Length of each generated waveform in samples, shape ``(B,)``. + """ + B = next_subword_ids.size(0) + + if generation_config is None: + generation_config = self._get_generation_config(guidance_enabled) + logging.info(f"Doing inference using the following config: {generation_config} !") + + init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": guidance_enabled}) + + # warmup the model and generate the very first audio token + outputs = self.tts_model(**init_inputs) + + if self.cfg.get("inference_skip_first_code_prediction_on_init", True): + # use the last token on init, because we are shifthing it in the model forward, so we dont really need to compute it + code = init_inputs["code"][:, -1:] + else: + code, _, _ = self.tts_model.generate_step(outputs.hidden_states[:, -1:], **generation_config) + + past_key_values = outputs["past_key_values"] + + # use the text tokens to stop generation + max_steps = next_subword_ids.size(-1) + # create variable to store the audios + gen_audio_codes = torch.zeros( + B, max_steps, self.tts_model.config.num_quantizers, device=self.device, dtype=torch.long + ) + + # init subwork as all ones + subword_mask = torch.ones(B, max_steps, device=self.device, dtype=torch.bool) + # get first context subword_id, that is the last subword_ids from the warmup + first_context_subword_id = init_inputs["subword_ids"][:, -1].unsqueeze(-1) + + # initialize variables used to save the output audio + audio_pred = None + audio_pred_len = torch.zeros(B, device=self.device, dtype=torch.long) + + for i in range(max_steps): + step_start = time.time() + # current subword id is always seem + current_subword_id = next_subword_ids[:, i].unsqueeze(-1) + + if i == 0: + prev_subword_id = first_context_subword_id + else: + prev_subword_id = next_subword_ids[:, i - 1].unsqueeze(-1) + + # create subword_mask + current_subword_mask = subword_mask[:, i].unsqueeze(-1) + + code, past_key_values = self.infer_codes_one_step( + current_subword_id=current_subword_id, + prev_subword_id=prev_subword_id, + current_subword_mask=current_subword_mask, + prev_audio_tokens=code, + past_key_values=past_key_values, + guidance_enabled=guidance_enabled, + generation_config=generation_config, + ignore_eos_flag_stop=True, + ) + + # cache audio tokens + gen_audio_codes[:, i] = code.squeeze(1) + + if incremental_audio_decoding: + audio_pred_i, audio_pred_i_len = self.decode_one_audio_step( + gen_audio_codes[:, : i + 1], + number_prev_tokens=self.cfg.get("inference_codec_decoding_prev_tokens_number", None), + ) + if audio_pred is None: + audio_pred = audio_pred_i + else: + audio_pred = torch.cat([audio_pred, audio_pred_i], dim=1) + audio_pred_len += audio_pred_i_len + + step_time = time.time() - step_start + logging.info(f"Autoregressive inference step: {i} of {max_steps} take around {step_time}s") + + if not incremental_audio_decoding: + gen_audio_codes_lens = torch.tensor([gen_audio_codes.shape[1]] * gen_audio_codes.shape[0]).to(self.device) + # decode audio. Note that it is not necessary because the prompt is removed, so no special token should be on the output, but lets do it for safety + gen_audio_codes = replace_control_speech_codes( + gen_audio_codes, self._control_codes, self.codec_silence_tokens + ) + with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): + audio_pred, audio_pred_len = self.audio_codec.decode(gen_audio_codes, gen_audio_codes_lens) + + return audio_pred.squeeze(1), audio_pred_len + + def maybe_recreate_cached_audio_prompt_latents_structure(self, state_dict): + """ + Recreate audio_prompt_latents ParameterDict structure + from keys present in the provided state_dict. + """ + for k, tensor in state_dict.items(): + if "audio_prompt_latents." in k: + name = k.split("audio_prompt_latents.")[1] + if name not in self.audio_prompt_latents: + self.audio_prompt_latents[name] = nn.Parameter( + torch.zeros_like(tensor), + requires_grad=False, + ) + + def load_state_dict(self, state_dict, strict: bool = True): + self.maybe_recreate_cached_audio_prompt_latents_structure(state_dict) + try: + return super().load_state_dict(state_dict, strict=strict) + except RuntimeError: + logging.info("Error loading model state_dict !! Retrying with partial initialization!") + model_dict = set_model_dict_for_partial_init(state_dict, self.state_dict()) + return super().load_state_dict(model_dict, strict=False) + + +@contextmanager +def ensures_target_precision(target_dtype): + """ + Workaround for precision related issues when training with bf16-true PyTorch Lightning precision setting. + In bf16-true, PTL changes PyTorch's default dtype, which may break implicit assumptions for some models. + This context manager restores default float32 precision and runs the computation in float32 autocast context. + """ + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(target_dtype) + try: + with torch.amp.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu", dtype=target_dtype): + yield + finally: + torch.set_default_dtype(default_dtype) + + +def replace_control_speech_codes( + speech_codes: torch.Tensor, control_codes: torch.Tensor, silence_tokens: torch.Tensor = None +) -> torch.Tensor: + """ + Replaces control codes (speech BOS, EOS, etc) in `speech_codes` with the first frame which is + assumed to consist of 'valid' codes representing silence. + """ + if silence_tokens is not None: + # Expand to [B, 1, 74] + silence_tokens_expanded = silence_tokens.unsqueeze(0).unsqueeze(1).expand(speech_codes.shape[0], 1, -1) + return torch.where(torch.isin(speech_codes, control_codes), silence_tokens_expanded, speech_codes) + + if torch.isin(speech_codes[:, :1], control_codes).any(): + return torch.where(torch.isin(speech_codes, control_codes), torch.zeros_like(speech_codes[:, :1]), speech_codes) + else: + return torch.where(torch.isin(speech_codes, control_codes), speech_codes[:, :1], speech_codes) + + +def setup_audio_codec(model): + """ + Instantiates the RVQ audio codec and injects codec embeddings into the TTS model. + + This function is responsible only for: + - Instantiating the codec model (`RVQVAEModel`). + - Loading pretrained codec weights (if configured). + - Freezing codec parameters. + - Registering RVQ embeddings inside the TTS model via `set_rvq_embs`. + + Args: + model: Model instance of DuplexEARTTS + """ + with ensures_target_precision(model.audio_codec_run_dtype): + model.audio_codec = RVQVAEModel(DictConfig(model.cfg.codec_config)) + # load pretrained codec checkpoint + if model.cfg.get("pretrained_codec_model", None): + checkpoint_state = load_checkpoint(model.cfg.pretrained_codec_model) + checkpoint_state = set_model_dict_for_partial_init(checkpoint_state, model.audio_codec.state_dict()) + model.audio_codec.load_state_dict(checkpoint_state, strict=True) + + for p in model.audio_codec.parameters(): + p.requires_grad = False + + model.audio_codec.eval() + + assert callable(model.tts_model.set_rvq_embs) + + model.tts_model.set_rvq_embs(torch.stack([x.detach() for x in model.audio_codec.prvq.mus_list], 0)) + model.tts_model.rvq_embs = model.tts_model.rvq_embs.to(next(model.tts_model.parameters()).dtype) + # compute target fps + model.target_fps = model.target_sample_rate / model.audio_codec.config.wav_to_token_ratio + model.target_samples_per_frame = model.audio_codec.config.wav_to_token_ratio diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py new file mode 100644 index 000000000..1b8b96169 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_commons.py; training-only code removed. + +import argparse +import glob +import importlib.machinery +import json +import os +import re +import shutil +import subprocess +import sys +from collections.abc import Mapping +from typing import Any + +from omegaconf import DictConfig +from safetensors import safe_open +from torch import nn + +from .compat import logging + +# ============================================================================== +# Constants +# ============================================================================== +PYTHON_CONFIG_GETTER_NAME = "get_config" +CHECKPOINT_FORMAT = "checkpoint_{}/ema.safetensors" +CONFIG_NAME = "config.json" +GIT_HASH_NAME = "githash" +SCRIPT_PLACEHOLDER = "[[[<<>>]]]" + + +# ============================================================================== +# Configuration Class and Utilities +# ============================================================================== +def get_config_from_file(config_path: str) -> DictConfig: + """ + Loads a configuration from a JSON or Python file. + + - For JSON files (`*.json`), it parses the file directly. + - For Python files (`*.py`), it imports the file as a module and calls a + `get_config()` function within it. + - It also supports a special syntax `path/to/config.py:config_name` to select + a specific configuration from a Python file that returns a dictionary of configs. + + Args: + config_path (str): The path to the configuration file. + + Returns: + Config: The loaded configuration object. + + Raises: + AssertionError: If the file path is invalid, does not exist, or is not in + the expected format. + """ + + match = re.search(r".+\.((json)|(py)|(py:.+))$", config_path) + assert match, f"Only Python (*.py) or JSON (*.json) files are supported, but got {config_path}." + + py_config_name: str | None = None + if not (config_path.endswith(".py") or config_path.endswith(".json")): + config_path_split = config_path.split(":") + config_path = ":".join(config_path_split[:-1]) + py_config_name = config_path_split[-1] + + assert os.path.isfile(config_path), f"Configuration file not found at: {config_path}" + + if config_path.endswith(".json"): + with open(config_path) as f: + config = json.load(f) + else: + config_module = importlib.machinery.SourceFileLoader("_config", config_path).load_module() + assert hasattr(config_module, PYTHON_CONFIG_GETTER_NAME), ( + f"Python config file must define a `{PYTHON_CONFIG_GETTER_NAME}` function." + ) + config = getattr(config_module, PYTHON_CONFIG_GETTER_NAME)(py_config_name) + assert isinstance(config, Mapping), f"`{PYTHON_CONFIG_GETTER_NAME}` must return a dictionary-like object." + cfg = DictConfig(config) + return cfg + + +def get_config() -> DictConfig: + """ + Parses command-line arguments to load the main configuration for a training run. + + This function implements a hierarchical configuration loading strategy: + 1. It checks if a `config.json` exists in the specified `--workdir`. If so, it loads it. + 2. If a `--config` argument is also provided, it uses that file to update the + configuration loaded from the work directory. + 3. If no config exists in the work directory, it requires the `--config` argument + to be provided as the base configuration. + + This allows for resuming training from a work directory while also being able to + override specific parameters for a new run. + + Returns: + Config: The final, consolidated configuration object. + """ + parser = argparse.ArgumentParser(description="Load training configuration.") + parser.add_argument("-c", "--config", type=str, default=None, help="Path to a Python or JSON configuration file.") + parser.add_argument("-w", "--workdir", type=str, required=True, help="Work directory to save logs and checkpoints.") + + args = parser.parse_args() + workdir_path = args.workdir + config_save_path = os.path.join(workdir_path, CONFIG_NAME) + + if os.path.exists(config_save_path): + logging.info(f"Resuming from work directory. Loading configuration from {config_save_path}.") + cfg = get_config_from_file(config_save_path) + if args.config and args.config != config_save_path: + logging.info(f"Updating loaded configuration with parameters from {args.config}.") + override_cfg = get_config_from_file(args.config) + cfg.update(override_cfg) + else: + assert args.config is not None, "A configuration file must be specified via `-c` or `--config` for a new run." + logging.info(f"Starting a new run. Loading configuration from {args.config}.") + cfg = get_config_from_file(args.config) + cfg.workdir_path = workdir_path + + return cfg + + +def get_config_from_dir(workdir_path: str) -> DictConfig: + """ + A simple utility to load the configuration directly from a work directory. + + Args: + workdir_path (str): The path to the work directory containing a `config.json`. + + Returns: + DictConfig: The loaded configuration object. + """ + config_save_path = os.path.join(workdir_path, CONFIG_NAME) + cfg = get_config_from_file(config_save_path) + cfg.workdir_path = workdir_path + return cfg + + +# ============================================================================== +# Base Model Classes +# ============================================================================== + + +class PreTrainedModel(nn.Module): + config_class = DictConfig + + """ + A base class for models to handle loading from pretrained checkpoints. + + This class provides a common interface for initializing a model and loading + weights from a saved checkpoint, following a pattern similar to libraries + like Hugging Face's Transformers. + + Args: + config (DictConfig | dict[str, Any]): A configuration object containing model hyperparameters. + """ + + def __init__(self, config: DictConfig | dict[str, Any], *args, **kwargs): + super().__init__() + self.config = config if isinstance(config, self.config_class) else self.config_class(config) + + @classmethod + def from_pretrained( + cls, + pretrained_dir: str, + cfg: DictConfig | dict[str, Any] | None = None, + checkpoint_regex: str = "checkpoint_*/ema.safetensors", + strict: bool = False, + **model_kwargs, + ) -> "PreTrainedModel": + """ + Loads a pretrained model from a directory. + + This method first loads the configuration file from the specified directory, + initializes the model with this configuration, and then loads the weights + from the latest checkpoint file found in that directory. + + Args: + cls (type): The model class to instantiate. + pretrained_dir (str): The directory containing the pretrained model + config and checkpoint files. + cfg (DictConfig | dict[str, Any] | None, optional): An optional config object to override + the loaded config. Defaults to None. + checkpoint_regex (str, optional): A regex pattern to find the checkpoint + file. Defaults to "checkpoint_*/ema.safetensors". + strict (bool, optional): Whether to strictly enforce that the keys in + the checkpoint match the keys of the model. + Defaults to False. + **model_kwargs: Additional keyword arguments to pass to the model's + constructor. + + Returns: + PreTrainedModel: An instance of the model with loaded weights. + """ + pretrained_cfg = get_config_from_dir(pretrained_dir).model + if cfg is not None: + pretrained_cfg.update(cfg) + logging.info(f"The loaded config of the pretrained model is updated to: {pretrained_cfg}") + model = cls( + pretrained_cfg, + **model_kwargs, + ) + model_state_dict = {} + with safe_open(latest_checkpoint_path(pretrained_dir, checkpoint_regex), framework="pt", device="cpu") as f: + for key in f.keys(): + model_state_dict[key] = f.get_tensor(key) + model.load_state_dict(model_state_dict, strict=strict) + return model + + +# ============================================================================== +# IO and Checkpointing Utilities +# ============================================================================== + + +def latest_checkpoint_path(dir_path: str, regex: str | None = None) -> str: + """ + Finds the path of the latest checkpoint file or directory in a directory. + + The latest checkpoint is determined by sorting the filenames alphanumerically + and picking the last one. This assumes a naming convention like `checkpoint_1000.pt`, + `checkpoint_2000.pt`, etc. + + Args: + dir_path (str): The directory to search for checkpoints. + regex (str | None, optional): A glob pattern to match checkpoint files. If None, + a default pattern is used. Defaults to None. + + Returns: + str: The full path to the latest checkpoint file. + + Raises: + AssertionError: If no files matching the regex are found in the directory. + """ + if regex is None: + regex = CHECKPOINT_FORMAT.format("*") + + f_list = glob.glob(os.path.join(dir_path, regex)) + if not f_list: + raise FileNotFoundError(f"No checkpoint files or directories found in {dir_path} matching '{regex}'") + + # Sort files based on the integer values in their names + f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f)))) + + latest_path = f_list[-1] + logging.info(f"Latest checkpoint '{os.path.relpath(latest_path, start=dir_path)}' found in '{dir_path}'.") + return latest_path diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py new file mode 100644 index 000000000..e83e78cb6 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py @@ -0,0 +1,1620 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_model.py; training-only code removed. +# NOTE: DuplexEARTTS (duplex_ear_tts.py) imports RVQEARTTSModel from THIS file, not from +# rvq_ear_tts_model.py — both variants are vendored to match upstream exactly. + +import json +import math +import os +from dataclasses import dataclass, fields +from typing import Any + +import torch +import transformers +from omegaconf import DictConfig, OmegaConf +from torch import Tensor, nn +from torch.nn import functional as F +from transformers import AutoConfig, AutoModel, AutoModelForTextEncoding, Cache +from transformers.generation.logits_process import TopKLogitsWarper, TopPLogitsWarper + +from .compat import fp32_precision, logging, set_model_dict_for_partial_init +from .tokenizer import AutoTokenizer + +# ============================================================================== +# MLP module and Norm +# ============================================================================== + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()) + # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16) + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + +class MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + ): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + self.act_fn = nn.GELU(approximate="tanh") + + def forward(self, x: Tensor) -> Tensor: + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class MLPLayer(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + eps: float = 1e-6, + ): + super().__init__() + self.pre_norm = RMSNorm(hidden_size, eps=eps) + self.mlp = MLP(hidden_size, intermediate_size) + self.post_norm = RMSNorm(hidden_size, eps=eps) + + def forward(self, x: Tensor) -> Tensor: + y = self.pre_norm(x) + y = self.mlp(y) + y = self.post_norm(y) + x = x + y + return x + + +# ============================================================================== +# Triton-accelerated and Fallback Functions +# ============================================================================== + +TRITON_IMPORTED = False +try: + import triton + import triton.language as tl + + TRITON_IMPORTED = True +except ImportError: + TRITON_IMPORTED = False + +USE_TRITON = TRITON_IMPORTED and torch.cuda.is_available() + +if USE_TRITON: + logging.info("Triton available & CUDA detected. Using Triton kernel for batch_matmul.") + + @triton.jit + def batch_matmul_kernel( + x_ptr, + w_ptr, + y_ptr, + result_ptr, + b, + d_in, + d_out, + n, + BLOCK_SIZE_DIN: tl.constexpr, + BLOCK_SIZE_DOUT: tl.constexpr, + ): + batch_id = tl.program_id(axis=0) + dout_block_id = tl.program_id(axis=1) + + if batch_id >= b: + return + + idx = tl.load(y_ptr + batch_id) + + x_offset = x_ptr + batch_id * d_in + w_offset = w_ptr + idx * d_out * d_in + + dout_offsets = dout_block_id * BLOCK_SIZE_DOUT + tl.arange(0, BLOCK_SIZE_DOUT) + dout_mask = dout_offsets < d_out + + result_block = tl.zeros([BLOCK_SIZE_DOUT], dtype=tl.float32) + + for din_start in range(0, d_in, BLOCK_SIZE_DIN): + din_offsets = din_start + tl.arange(0, BLOCK_SIZE_DIN) + din_mask = din_offsets < d_in + + x_i = tl.load(x_offset + din_offsets, mask=din_mask, other=0.0) + + w_i_block = tl.load( + w_offset + dout_offsets[:, None] * d_in + din_offsets[None, :], + mask=(dout_mask[:, None] & din_mask[None, :]), + other=0.0, + ) + + result_block += tl.sum(w_i_block * x_i[None, :], axis=1) + + result_offset = result_ptr + batch_id * d_out + dout_offsets + tl.store(result_offset, result_block, mask=dout_mask) + + def batch_matmul_triton(x, w, y, BLOCK_SIZE_DIN: int = 16, BLOCK_SIZE_DOUT: int = 64): + assert x.is_contiguous() and w.is_contiguous() and y.is_contiguous() + + b, d_in = x.shape + n, d_out, _ = w.shape + result = torch.empty(b, d_out, device=x.device, dtype=torch.float32) + + batch_matmul_kernel[lambda meta: (b, triton.cdiv(d_out, meta["BLOCK_SIZE_DOUT"]))]( + x.float(), + w.float(), + y, + result, + b, + d_in, + d_out, + n, + BLOCK_SIZE_DIN=BLOCK_SIZE_DIN, + BLOCK_SIZE_DOUT=BLOCK_SIZE_DOUT, + ) + + return result.to(dtype=x.dtype) + + batch_matmul = batch_matmul_triton + +else: + logging.info("Using PyTorch fallback (Triton unavailable or no CUDA).") + + # Fallback to PyTorch implementation if Triton is not available + def batch_matmul_pytorch(x: Tensor, w: Tensor, y: Tensor, *args, **kwargs) -> Tensor: + """ + Performs a batched matrix multiplication using PyTorch's native functions. + + This function serves as a fallback when Triton is not available. It achieves + the same result by gathering the appropriate weight matrices and using `torch.bmm`. + + Args: + x (Tensor): The input tensor of shape `[batch_size, d_in]`. + w (Tensor): The weight tensor of shape `[num_weights, d_out, d_in]`. + y (Tensor): The index tensor of shape `[batch_size]`. + + Returns: + Tensor: The result of the multiplication, shape `[batch_size, d_out]`. + """ + # w[y] gathers the weight matrices for each item in the batch. + # x.unsqueeze(2) reshapes x to [batch_size, d_in, 1] for bmm. + # The result is squeezed to remove the trailing dimension of size 1. + return torch.bmm(w[y], x.unsqueeze(2)).squeeze(2) + + batch_matmul = batch_matmul_pytorch + + +# ============================================================================== +# Core Mathematical and Masking Functions +# ============================================================================== + + +def gumbel_like(tensor: Tensor, eps: float = 1e-8) -> Tensor: + """ + Generates a tensor of Gumbel noise with the same shape as the input tensor. + + This is used for the Gumbel-Max trick, a technique to sample from a categorical + distribution in a differentiable way (using a straight-through estimator). + + Args: + tensor (torch.Tensor): The input tensor to match the shape of. + eps (float): A small epsilon value for numerical stability. + + Returns: + torch.Tensor: A tensor containing Gumbel noise. + """ + # Sample from a uniform distribution + u = torch.rand_like(tensor) + # Apply the inverse CDF of the Gumbel distribution + return -torch.log(-torch.log(u + eps) + eps) + + +def sequence_mask(lengths: Tensor, max_length: Tensor | int | None = None) -> Tensor: + """ + Creates a boolean mask from a 1D tensor of sequence lengths. + + This function is useful for masking out padding in sequences. Given a tensor + of lengths, it produces a 2D boolean tensor where `mask[i, j]` is `True` if + `j < lengths[i]` and `False` otherwise. + + Args: + lengths (Long Tensor): A 1D tensor of integer lengths. Shape: `[batch_size]`. + max_length (Long Tensor | int | None, optional): The maximum length of the mask. If None, + it is inferred from the maximum value + in `lengths`. Defaults to None. + + Returns: + Tensor: The boolean mask. Shape: `[batch_size, max_length]`. + """ + if max_length is None: + max_length = lengths.max() + + # Create a range tensor from 0 to max_length - 1 + x = torch.arange(max_length, dtype=lengths.dtype, device=lengths.device) # type: ignore[arg-type] + + # Compare each length with the range tensor to create the mask via broadcasting + return x.unsqueeze(0) < lengths.unsqueeze(1) + + +def get_masking_rate(rate: Tensor, exponent: float = 3.0) -> Tensor: + """ + Converts a desired token keep rate to a masking rate using a power function. + + This function is part of a scheduling strategy for masking, where the effective + masking rate changes non-linearly with the desired keep rate. This function is + its own inverse. + + Args: + rate (Tensor): The desired rate of tokens to keep (0 to 1). + exponent (float, optional): The exponent for the transformation. Defaults to 3.0. + + Returns: + Tensor: The corresponding masking rate. + """ + return (1 - rate.pow(exponent)).pow(1 / exponent) + + +# Alias the function for clarity in the inverse context +get_rate = get_masking_rate + + +def get_mask( + code_mask: Tensor, + num_masking: Tensor, + unmasking: bool = False, + validate: bool = False, +) -> Tensor: + """ + Adjusts a boolean mask by masking or unmasking tokens from the end. + + This function operates on a `code_mask` where `True` values represent valid + tokens and are assumed to be contiguous at the start of the sequence. It + calculates a new mask by decreasing (masking) or increasing (unmasking) + the number of `True` values. + + Args: + code_mask (Tensor): The input boolean mask. Shape: `[..., depth]`. + num_masking (Tensor): The number of tokens to mask or unmask. + Shape matching `code_mask`'s batch dimensions. + unmasking (bool, optional): If `True`, increases the number of valid + tokens (unmasking). Defaults to `False`. + validate (bool, optional): If `True`, asserts that the input `code_mask` + is contiguous. This adds a slight overhead and + is mainly for debugging. Defaults to `False`. + + Returns: + Tensor: A new boolean mask with the adjusted length of valid tokens. + Shape is identical to `code_mask`. + """ + depth = code_mask.size(-1) + num_valid = code_mask.sum(dim=-1, dtype=torch.long) + + if validate: + # Reconstruct the expected contiguous mask and assert equality. + expected_mask = sequence_mask(num_valid.view(-1), depth).view_as(code_mask) + assert torch.equal(code_mask, expected_mask), ( + "Input `code_mask` must have contiguous `True` values at the beginning." + ) + + # Calculate the target number of valid tokens. + if not unmasking: + # Masking: reduce the number of valid tokens, ensuring it's not negative. + num_to_keep = (num_valid - num_masking).clamp_min(0) + else: + # Unmasking: increase the number of valid tokens, capped by total depth. + num_to_keep = (num_valid + num_masking).clamp_max(depth) + + # Generate the new mask using the final number of tokens to keep. + return sequence_mask(num_to_keep.view(-1), depth).view_as(code_mask) + + +# ============================================================================== +# Model and Vocabulary Utilities +# ============================================================================== + + +@dataclass +class RVQEARTTSOutput: + """ + Output type for the RVQEARTTSModel, providing a structured way to return model outputs. + This class allows accessing outputs by attribute, key, or index. + """ + + loss: Tensor | None = None + lm_loss: Tensor | None = None + c_loss: Tensor | None = None + k_loss: Tensor | None = None + + hidden_states: Tensor | None = None + past_key_values: Tensor | None = None + audio_prompt_latent: Tensor | None = None + + codes: Tensor | None = None + lm_logits: Tensor | None = None + eos_flag: Tensor | None = None + + def __getitem__(self, item: str | int): + """Allows for accessing attributes by key or index.""" + if isinstance(item, str): + return getattr(self, item) + else: + # Access fields in the order they are defined in the dataclass + return getattr(self, fields(self)[item].name) + + +def find_and_delete_module(parent_module: nn.Module, target_module: nn.Module, parent_name: str) -> str | None: + """ + Recursively searches for a specific module instance and deletes it from its parent. + + This is useful for dynamically modifying a model's architecture, such as replacing + an existing embedding layer with a custom one. + + Args: + parent_module (nn.Module): The module to search within. + target_module (nn.Module): The exact module instance to find and delete. + parent_name (str): The initial name of the parent module for constructing the path. + + Returns: + str | None: The full dotted name of the deleted attribute if found, otherwise None. + """ + # Iterate over all direct children of the parent module + for name, module in parent_module.named_children(): + # Use the 'is' operator to check for object identity, not just value equality + if module is target_module: + # If found, delete the attribute from the parent and return its name + delattr(parent_module, name) + return f"{parent_name}.{name}" + + # If not found, recurse into the child module + found_path = find_and_delete_module(module, target_module, parent_name=f"{parent_name}.{name}") + if found_path: + return found_path + return None + + +def build_vocabs( + tokenizer: AutoTokenizer, vocab_dir: str | None = None +) -> tuple[dict[int, tuple[int, ...]], dict[str, int], int]: + """ + Builds or loads a character-level vocabulary derived from a subword tokenizer. + + This function creates a mapping from each subword in a pretrained tokenizer to a + sequence of character IDs. It follows a modern practice of using a directory + to save and load vocabulary files, making the process more robust and extensible. + + The primary source of truth is the `char_vocab.json` file. If it exists, it's + loaded. Otherwise, it's created from the pretrained tokenizer and saved. + + Args: + tokenizer (AutoTokenizer): The pretrained Hugging Face tokenizer class. + vocab_dir (str | None, optional): The directory to save or load the character + vocabulary from. Defaults to None. + + Returns: + tuple[dict[int, tuple[int, ...]], dict[str, int], int]: A tuple containing: + - A mapping from subword IDs to tuples of character IDs. + - The character-to-ID vocabulary dictionary. + - The ID for the subword padding token. + """ + + def _build_char_vocab() -> dict[str, int]: + # Find all single-character tokens in the original tokenizer's vocabulary + single_chars = { + subword: subword_id for subword, subword_id in tokenizer.tokenizer.vocab.items() if len(subword) == 1 + } + # Create a new, dense character vocabulary sorted by the original token ID + sorted_chars = sorted(single_chars.keys(), key=lambda k: single_chars[k]) + char_vocab = {char: i for i, char in enumerate(sorted_chars)} + return char_vocab + + # 1. Load or build the character vocabulary + if vocab_dir: + from filelock import FileLock + + char_vocab_file = os.path.join(vocab_dir, "char_vocab.json") + os.makedirs(vocab_dir, exist_ok=True) + with FileLock(char_vocab_file + ".lock", timeout=60): + if not os.path.exists(char_vocab_file): + char_vocab = _build_char_vocab() + + logging.info(f"Saving character vocabulary to {char_vocab_file}") + with open(char_vocab_file, "w", encoding="utf-8") as f: + json.dump(char_vocab, f, ensure_ascii=False, indent=2) + + # All processes can now safely load the file. + logging.info(f"Loading character vocabulary from {char_vocab_file}") + with open(char_vocab_file, encoding="utf-8") as f: + char_vocab = json.load(f) + else: + # No cache directory provided, build in memory. + logging.info("Building character vocabulary from tokenizer.") + char_vocab = _build_char_vocab() + + # 2. Reconstruct the subword-to-character mapping on the fly + subword_id_to_char_ids = { + subword_id: tuple(char_vocab[char] for char in subword if char in char_vocab) + for subword, subword_id in tokenizer.tokenizer.vocab.items() + } + # Filter out subwords that contain characters not in our character vocabulary + subword_id_to_char_ids = {k: v for k, v in subword_id_to_char_ids.items() if v} + + # 3. Define a padding index for subwords + subword_padding_idx = len(tokenizer.vocab) + # The padding subword maps to a new character padding ID + subword_id_to_char_ids[subword_padding_idx] = (len(char_vocab),) + return subword_id_to_char_ids, char_vocab, subword_padding_idx + + +# @torch.compile +@torch._dynamo.disable +def depthsum_encoding_step( + embs: Tensor, + r: Tensor, + code: Tensor, + depth_str: int = 0, + k: int = 72, +) -> Tensor: + for i in range(depth_str, depth_str + k): + idx_sel = ( + embs[i].pow(2).sum(-1) # [g?, v] + - 2 + * (r.unsqueeze(-2) @ embs[i].transpose(-1, -2)).squeeze(-2) # [b, ?, g?, h] , [g?, h, v] -> [b, ?, g?, v] + ).argmin(-1) + + emb_i = F.embedding(idx_sel, embs[i]) + r = r - emb_i + + code[..., i] = idx_sel + + return code + + +class MoGHead(nn.Module): + """ + A Mixture of Gaussians (MoG) prediction head. + + This module takes a hidden state and predicts the parameters for a mixture of + Gaussian distributions. It's suitable for modeling continuous, multi-modal data. + + Args: + hidden_size (int): The dimensionality of the input hidden state. + intermediate_size (int): The dimensionality of the MLP layers. + out_size (int): The dimensionality of the output vectors (the mean of each Gaussian). + num_layers (int): The number of MLP layers in the stack. + num_predictions (int): The number of Gaussian components in the mixture. + low_rank (int | None): The dimensionality used for compressing the hidden states. + min_log_std (float): The minimum value for the logarithm of the standard deviation. + eps (float): A small epsilon value for the RMSNorm layers. + """ + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + out_size: int, + num_layers: int, + num_predictions: int, + low_rank: int | None = 64, + min_log_std: float = -4.0, + eps: float = 1e-6, + ): + super().__init__() + self.out_size = out_size + self.low_rank = low_rank + self.num_predictions = num_predictions + self.min_log_std = min_log_std + + self.mlp_stack = nn.Sequential( + *[MLPLayer(hidden_size, intermediate_size, eps=eps) for _ in range(num_layers)], + RMSNorm(hidden_size, eps=eps), + ) + + if low_rank is None: + self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) # Predicts mixture weights + self.proj_mus = nn.Linear(hidden_size, num_predictions * out_size, bias=False) # Predicts means + self.proj_logs = nn.Linear(hidden_size, 1, bias=False) # Predicts log standard deviations + else: + assert low_rank < out_size + self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) # Predicts mixture weights + self.proj_mus = nn.Linear(hidden_size, num_predictions * low_rank, bias=False) # Predicts means + self.proj_logs = nn.Linear(hidden_size, 1, bias=False) # Predicts log standard deviations + self.proj_else = nn.Linear(hidden_size, out_size, bias=False) + self.low_mat = nn.Parameter(torch.randn(num_predictions, out_size, low_rank) * (low_rank**-0.5)) + + def infer(self, x: Tensor, guidance_scale: float = 0.0, top_p_or_k: float | int = 1.0) -> tuple[Tensor, Tensor]: + """ + Performs inference by sampling from the predicted mixture distribution. + + Args: + x (Tensor): The input hidden state. + guidance_scale (float): The weight for classifier-free guidance. + top_p_or_k (float | int): The value for top-p (nucleus) or top-k sampling of the mixture components. + + Returns: + tuple[Tensor, Tensor]: A tuple containing the mean of the chosen component, + and the log standard deviations. + """ + b, t, _ = x.size() + n, d = self.num_predictions, self.low_rank or self.out_size + x = self.mlp_stack(x) + if guidance_scale > 0: + b //= 2 + x_cond, x_uncond = x.chunk(2, dim=0) + x = x_cond + guidance_scale * (x_cond - x_uncond) + + logits = self.proj_logits(x) + + # Apply top-p or top-k filtering to the mixture logits + if top_p_or_k is not None: + logits = ( + TopPLogitsWarper(top_p_or_k)( + None, + logits.view(-1, n), + ).view_as(logits) + if isinstance(top_p_or_k, float) + else TopKLogitsWarper(top_p_or_k)( + None, + logits.view(-1, n), + ).view_as(logits) + ) + + # Sample a mixture component using the Gumbel-Max trick + with fp32_precision(): + mixture_indices = (F.log_softmax(logits, dim=-1) + gumbel_like(logits)).argmax(-1) + + # Select the mean corresponding to the sampled component + mu = batch_matmul( + x.view(b * t, -1), + self.proj_mus.weight.detach().view(n, d, -1), + mixture_indices.view(b * t), + ).view(b, t, d) + if self.proj_mus.bias is not None: + mu += self.proj_mus.bias.detach().view(n, d)[mixture_indices] + + if self.low_rank: + assert math.log2(d).is_integer() and math.log2(self.out_size).is_integer() + mu = batch_matmul( + mu.view(b * t, -1), + self.low_mat.detach().view(n, self.out_size, -1), + mixture_indices.view(b * t), + BLOCK_SIZE_DIN=d, + BLOCK_SIZE_DOUT=self.out_size, + ).view(b, t, self.out_size) + + mu_res = self.proj_else(x) + else: + mu_res = torch.zeros((b, t, d), device=x.device) + + logs = self.proj_logs(x).clamp_min(self.min_log_std) + return mu * torch.exp(logs) + mu_res, logs + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """ + Performs a forward pass for training. + + Args: + x (Tensor): The input hidden state. + + Returns: + tuple[Tensor, Tensor, Tensor]: A tuple containing the mixture logits, + the means for all components, and the + log standard deviations. + """ + b, t, _ = x.size() + d = self.low_rank or self.out_size + x = self.mlp_stack(x) + logits = self.proj_logits(x) + mus = self.proj_mus(x).view(b, t, self.num_predictions, d) + logs = self.proj_logs(x).clamp_min(self.min_log_std) + + if self.low_rank: + mu_res = self.proj_else(x) + else: + mu_res = torch.zeros((b, t, d), device=x.device) + return logits, mus, mu_res, logs + + def dist(self, mus: Tensor, mu: Tensor) -> Tensor: + """ + mus: [b, t, n, d] + mu: [b, t, d] + + return: [b, t, n] + """ + if self.low_rank is None: + return (mus - mu.unsqueeze(-2)).pow(2).sum(-1) + else: + low_mat_sq = self.low_mat.transpose(-1, -2) @ self.low_mat + x, y = mus, mu + wx_sq = ( + x + * torch.einsum( + "btni,nij->btnj", + x, + low_mat_sq.to(x), + ) + ).sum(-1) # [b, t, n] + y_sq = y.pow(2).sum(-1, keepdim=True) # [b, t, 1] + xwy = (x * torch.einsum("bti,nij->btnj", y, self.low_mat.to(y))).sum( + -1 + ) # [b, t, n, d_l], [n, d_i, d_l], [b, t, d_i] -> [b, t, n] + + dist = wx_sq + y_sq - 2 * xwy + return torch.abs(dist) + + +class SubwordFlagEmbedding(nn.Module): + """ + Adds a small continuation embedding for subwords (tokens without word-boundary marker). + Automatically adds a custom padding token at index vocab_size. + Ignores special tokens (starting with '<') when computing continuation flags. + """ + + def __init__(self, tokenizer: AutoTokenizer, d_model: int): + super().__init__() + + self.tokenizer = tokenizer + self.vocab_size = self.tokenizer.vocab_size + self.d_model = d_model + + # Custom pad token at vocab_size + self.pad_id = self.vocab_size + # register pad_id as a tensor buffer to avoid device issues + self.register_buffer("pad_tensor", torch.tensor(self.pad_id, dtype=torch.long)) + + # Precompute continuation flags + tokens = [self.tokenizer.ids_to_tokens(i) for i in range(self.vocab_size)] + cont_flags = [ + 1 if not (tok.startswith("Ġ") or tok.startswith("▁") or tok.startswith("<")) else 0 for tok in tokens + ] + cont_flags.append(0) # for the custom pad token + self.register_buffer("is_continuation", torch.tensor(cont_flags, dtype=torch.long)) + + # Continuation embedding + init_std = self.d_model**-0.5 + self.cont_emb = nn.Embedding(2, self.d_model) + nn.init.normal_(self.cont_emb.weight, mean=0.0, std=init_std) + self.cont_emb.weight.data[0].zero_() + + def forward(self, subword_embeds: torch.Tensor, token_ids: torch.LongTensor): + # Replace OOV token IDs with pad_id safely + token_ids_clamped = torch.where(token_ids >= self.vocab_size, self.pad_tensor, token_ids) + # Continuation flags + cont_flags = self.is_continuation[token_ids_clamped] + # Add continuation embedding + cont_emb = self.cont_emb(cont_flags) + return subword_embeds + cont_emb + + +class BOSEOSEmbedding(nn.Module): + """ + Adds independent embeddings for BOS and EOS tokens using a single embedding table. + Index 0 = regular token (ignored), 1 = BOS, 2 = EOS. + Compatible with Hugging Face tokenizers that may or may not have BOS/EOS. + """ + + def __init__(self, tokenizer: AutoTokenizer, d_model: int): + super().__init__() + + self.tokenizer = tokenizer + # vocab size that includes special tokens + vocab_dict = self.tokenizer.tokenizer.get_vocab() + self.vocab_size = max(vocab_dict.values()) + self.d_model = d_model + + # Custom pad token for OOVs + self.pad_id = self.vocab_size + self.register_buffer("pad_tensor", torch.tensor(self.pad_id, dtype=torch.long)) + + # Identify BOS and EOS tokens (may be None) + tokens = [self.tokenizer.ids_to_tokens(i) for i in range(self.vocab_size)] + + special_flags = [] + for tok in tokens: + if self.tokenizer.bos_token is not None and tok == self.tokenizer.bos_token: + special_flags.append(1) + elif self.tokenizer.eos_token is not None and tok == self.tokenizer.eos_token: + special_flags.append(2) + else: + special_flags.append(0) + special_flags.append(0) # for custom pad token + self.register_buffer("special_flags", torch.tensor(special_flags, dtype=torch.long)) + # Embedding table: 0 = regular, 1 = BOS, 2 = EOS + init_std = self.d_model**-0.5 + self.special_emb = nn.Embedding(3, d_model) + nn.init.normal_(self.special_emb.weight, mean=0.0, std=init_std) + self.special_emb.weight.data[0].zero_() # regular tokens ignored + + def forward(self, token_embeds: torch.Tensor, token_ids: torch.LongTensor): + """ + token_embeds: (B, T, d_model) + token_ids: (B, T) + """ + # Clamp OOVs to custom pad token + safe_ids = torch.where(token_ids >= self.vocab_size, self.pad_tensor, token_ids) + + # Lookup flags (0=regular, 1=BOS, 2=EOS) + flags = self.special_flags[safe_ids] + return token_embeds + self.special_emb(flags) + + +class CharAwareSubwordEncoder(nn.Module): + """ + An encoder that creates subword embeddings from character-level embeddings. + + This module replaces a standard subword embedding layer. It breaks down each + subword into its constituent characters, embeds the characters, and then + aggregates these character embeddings (e.g., via mean pooling) to form the + final subword representation. This allows the model to handle rare or out-of-vocabulary + subwords more gracefully. + + Args: + out_size (int): The dimensionality of the output embedding vectors. + tokenizer (AutoTokenizer): The Hugging Face tokenizer class. + vocab_dir (str | None): Directory to save/load the character vocabulary. + backbone_type (str | None): The type of backbone model from Hugging Face (e.g., "t5gemma"). + backbone_model_class (str | None): The class name of the backbone model if not using AutoModel. + backbone_config_class (str | None): The class name of the backbone config. + backbone_config (DictConfig | None): A configuration for the backbone model. + """ + + def __init__( + self, + out_size: int, + tokenizer: AutoTokenizer, + vocab_dir: str | None = None, + backbone_type: str | None = "t5gemma", + backbone_model_class: str | None = None, + backbone_config_class: str | None = None, + backbone_config: DictConfig | None = None, + use_subword_flag_emb: bool = True, + use_bos_eos_emb: bool = True, + use_cumulative_word_emb: bool = False, + ): + super().__init__() + + # 1. Build or load the character vocabulary + self.subword_id_to_char_ids, self.char_vocab, self.subword_padding_idx = build_vocabs( + tokenizer, + vocab_dir, + ) + + self.char_padding_idx = len(self.char_vocab) + self.use_subword_flag_emb = use_subword_flag_emb + self.use_bos_eos_emb = use_bos_eos_emb + + # 2. Initialize the backbone model + if backbone_type: + config = AutoConfig.for_model( + backbone_type, **(OmegaConf.to_container(backbone_config, resolve=True) if backbone_config else {}) + ) + self.backbone = AutoModelForTextEncoding.from_config(config) + else: + assert backbone_model_class and backbone_config_class + config_class = getattr(transformers, backbone_config_class) + model_class = getattr(transformers, backbone_model_class) + config = config_class(**(OmegaConf.to_container(backbone_config, resolve=True) if backbone_config else {})) + self.backbone = model_class(config) + + self.hidden_size = self.backbone.get_input_embeddings().weight.size(-1) + + # 3. Delete the original subword embedding layer and replace it with our character embedding layer + find_and_delete_module(self.backbone, self.backbone.get_input_embeddings(), "backbone") + self.embed_tokens = nn.Embedding(len(self.char_vocab) + 1, self.hidden_size, padding_idx=self.char_padding_idx) + self.proj_embedding = nn.Linear(self.hidden_size, out_size, bias=False) + + if self.use_subword_flag_emb: + self.subword_flag_emb = SubwordFlagEmbedding(tokenizer, self.hidden_size) + + if self.use_bos_eos_emb: + self.bos_eos_emb = BOSEOSEmbedding(tokenizer, self.hidden_size) + + def prepare_inputs(self, subword_ids: Tensor, padding_mask: Tensor) -> tuple[Tensor, Tensor]: + """ + Converts a batch of subword IDs into a padded batch of character IDs. + + Args: + subword_ids (Tensor): A tensor of subword IDs. Shape: `[batch, seq_len]`. + padding_mask (Tensor): A boolean mask indicating valid (non-padding) subwords. + + Returns: + tuple[Tensor, Tensor]: A tuple containing: + - Padded character IDs. Shape: `[num_valid_subwords, max_char_len]`. + - Lengths of each character sequence. Shape: `[num_valid_subwords]`. + """ + device = subword_ids.device + # Select only the valid subword IDs + subword_id_list = torch.masked_select(subword_ids, padding_mask).cpu().tolist() + # Map each subword ID to its sequence of character IDs + char_id_list = [list(self.subword_id_to_char_ids.get(x, ())) for x in subword_id_list] + + char_lengths = torch.tensor([len(x) for x in char_id_list], dtype=torch.long, device=device) + batch_size = char_lengths.size(0) + max_len = int(char_lengths.max().item()) if batch_size > 0 else 0 + + # Create a padded tensor for the character IDs + char_ids = torch.full((batch_size, max_len), self.char_padding_idx, dtype=torch.long, device=device) + for i, char_seq in enumerate(char_id_list): + char_ids[i, : len(char_seq)] = torch.tensor(char_seq, dtype=torch.long, device=device) + + return char_ids, char_lengths + + def forward(self, subword_ids: Tensor, subword_mask: Tensor | None = None) -> Tensor: + """ + Performs the forward pass to get character-aware subword embeddings. + + Args: + subword_ids (Tensor): A tensor of subword IDs. Shape: `[batch, seq_len]`. + subword_mask (Tensor | None): A boolean mask for padding. Defaults to None. + + Returns: + Tensor: The final subword embeddings. Shape: `[batch, seq_len, hidden_size]`. + """ + if subword_mask is None: + subword_mask = torch.ones_like(subword_ids, dtype=torch.bool) + + # 1. Convert subword IDs to character IDs + char_ids, char_lengths = self.prepare_inputs(subword_ids, subword_mask) + + # char_mask = sequence_mask(char_lengths).float() + char_mask = sequence_mask(char_lengths) + + # 2. Get character embeddings and pass them through the backbone + char_embeds = self.embed_tokens(char_ids) + + # The backbone model should be able to accept `inputs_embeds` + char_hidden_states = self.backbone(inputs_embeds=char_embeds, attention_mask=char_mask).last_hidden_state + + # 3. Aggregate character embeddings to form subword embeddings (mean pooling) + # We mask the padding characters before summing to get a correct mean. + masked_sum = (char_hidden_states * char_mask.unsqueeze(-1)).sum(dim=1) + # Avoid division by zero for empty sequences + mean_emb = masked_sum / (char_lengths.unsqueeze(-1).clamp(min=1)) + + # 4. Scatter the aggregated embeddings back to the original subword sequence shape + out_emb = self.proj_embedding(mean_emb) + subword_embeds = torch.zeros( + subword_ids.shape + (out_emb.size(-1),), device=subword_ids.device, dtype=out_emb.dtype + ) + subword_embeds[subword_mask] = out_emb + + if self.use_subword_flag_emb: + subword_embeds = self.subword_flag_emb(subword_embeds, subword_ids) + + if self.use_bos_eos_emb: + subword_embeds = self.bos_eos_emb(subword_embeds, subword_ids) + + return subword_embeds + + +class GatedProjectedSumRMSNorm(nn.Module): + def __init__(self, audio_dim, text_dim, hidden_dim, final_norm=True, num_codebooks=31, init_residual_scale=0.5): + super().__init__() + self.num_codebooks = num_codebooks + + self.audio_proj = nn.Linear(audio_dim, hidden_dim) + self.text_proj = nn.Linear(text_dim, hidden_dim) + + nn.init.normal_(self.audio_proj.weight, mean=0.0, std=0.015) + nn.init.zeros_(self.audio_proj.bias) + nn.init.normal_(self.text_proj.weight, mean=0.0, std=0.015) + nn.init.zeros_(self.text_proj.bias) + + # FP32 gate params + self.gate = nn.Parameter(torch.zeros(hidden_dim, dtype=torch.float32)) + self.residual_scale = nn.Parameter(torch.tensor(init_residual_scale, dtype=torch.float32)) + + self.final_norm = RMSNorm(hidden_dim) if final_norm else nn.Identity() + + def forward(self, audio_emb, text_emb): + audio_emb = audio_emb / self.num_codebooks + + # projections run in model dtype (BF16) + audio_h = self.audio_proj(audio_emb) + text_h = self.text_proj(text_emb) + + dtype = audio_h.dtype + + with fp32_precision(): + gate = torch.sigmoid(self.gate) # FP32 + res = torch.sigmoid(self.residual_scale) # FP32 + + h = gate.to(dtype) * audio_h + (1 - gate).to(dtype) * text_h + h = res.to(dtype) * h + h = self.final_norm(h.float()).to(dtype) + + return h + + +class RVQEARTTSModel(nn.Module): + """ + Main RVQEARTTS model for training and inference. + + This model integrates a character-aware text encoder with a transformer backbone + and a Mixture-of-Gaussians (MoG) prediction head. + + The architecture is based on the Streaming TTS model proposed in + "Audio Flamingo 3" (https://arxiv.org/abs/2507.08128), with several improvements: + + 1. Gated fusion of text and audio representations + (`GatedProjectedSumRMSNorm`). + + 2. Subword-aware embeddings for improved pronunciation of multi-token words + (`SubwordFlagEmbedding`). + + 3. Custom BOS and EOS embeddings for duplex interaction support, + enabling interruption-aware generation (`BOSEOSEmbedding`). + + Args: + config (DictConfig | dict[str, Any]): The configuration object for the model. + """ + + config_class = DictConfig + rvq_embs: Tensor + + def __init__(self, config: DictConfig | dict[str, Any], tokenizer: AutoTokenizer = None): + super().__init__() + self.config = config + + # Backbone module + if self.config.get("pretrained_text_name", None): + # Load pretrained backbone from huggingface + from .compat import load_pretrained_hf + + llm = load_pretrained_hf(self.config.pretrained_text_name, pretrained_weights=True).train() + self.backbone = llm.model # fetch PretrainedBaseModel from model "ForCausalLM" + else: + if self.config.get("backbone_type", None) is None: + assert ( + self.config.get("backbone_model_class", None) is not None + and self.config.get("backbone_config_class", None) is not None + ) + backbone_config = getattr(transformers, self.config.backbone_config_class)( + **( + OmegaConf.to_container(self.config.backbone_config, resolve=True) + if self.config.backbone_config + else {} + ), + ) + self.backbone = getattr(transformers, self.config.backbone_model_class)(backbone_config) + else: + backbone_config = AutoConfig.for_model( + self.config.backbone_type, + **( + OmegaConf.to_container(self.config.backbone_config, resolve=True) + if self.config.backbone_config + else {} + ), + ) + self.backbone = AutoModel.from_config(backbone_config) + + self.hidden_size = self.backbone.get_input_embeddings().weight.size(-1) + find_and_delete_module(self.backbone, self.backbone.get_input_embeddings(), "backbone") + + # Embedding and projection layers + self.bos_emb = nn.Parameter(torch.randn(self.hidden_size)) + self.null_emb = nn.Parameter(torch.randn(self.hidden_size)) + if self.config.random_target_masking: + self.embed_target_mask = nn.Embedding(self.config.num_quantizers, self.hidden_size) + + self.embed_code = nn.Linear(self.config.latent_size, self.hidden_size, bias=False) + + self.embed_context = ( + nn.Linear(self.config.context_hidden_size, self.hidden_size, bias=False) + if self.config.context_hidden_size + else None + ) + + cas_config = ( + {k: v for k, v in self.config.cas_config.items() if k != "pretrained_tokenizer_name"} + if self.config.cas_config + else {} + ) + self.embed_subword = ( + CharAwareSubwordEncoder( + tokenizer=tokenizer, + out_size=self.hidden_size, + use_subword_flag_emb=self.config.use_subword_flag_emb, + use_bos_eos_emb=self.config.use_bos_eos_emb, + **cas_config, + ) + if self.config.cas_config + else None + ) + + if self.config.use_gated_fusion_for_text_audio: + self.gated_fusion_audio_text = GatedProjectedSumRMSNorm( + self.hidden_size, self.hidden_size, self.hidden_size, self.config.num_quantizers + ) + + if self.config.get("audio_prompt_encoder_config", None): + # Dedicated projection for audio prompt (pre-BOS) + ape_cfg = OmegaConf.to_container(self.config.audio_prompt_encoder_config, resolve=True) + model_type = ape_cfg.pop("type") # remove "type" from kwargs + ape_cfg = AutoConfig.for_model(model_type, **ape_cfg) + self.audio_prompt_encoder = AutoModel.from_config(ape_cfg) + + if self.config.get("use_audio_prompt_frozen_projection", False): + with fp32_precision(): + U, _ = torch.linalg.qr(torch.randn(self.hidden_size, self.hidden_size)) + V, _ = torch.linalg.qr(torch.randn(self.hidden_size, self.hidden_size)) + smin, smax = 0.4, 2.5 + s = smin + (smax - smin) * torch.rand(self.hidden_size) + W = U @ torch.diag(s) @ V.T + self.register_buffer("audio_prompt_projection_W", W) # register as buffer to avoid weight update + + # Prediction Heads + if not self.config.disable_eos_prediction: + self.lm_head = nn.Linear(self.hidden_size, 2, bias=False) + + self.mog_head = MoGHead( + hidden_size=self.hidden_size, + out_size=self.config.latent_size, + **self.config.mog_head_config, + ) + + def set_rvq_embs(self, rvq_embs: Tensor): + self.register_buffer("rvq_embs", rvq_embs.detach().clone()) + + def depthsum_embedding(self, code: Tensor) -> Tensor: + """ + code: [b, t, d] + rvq_embs: [d, v, h] + + ret: [b, t, h] + """ + b, t, d = code.size() + _, v, h = self.rvq_embs.size() + device = code.device + + ret = torch.zeros((b, t, h), device=device) + embs = F.pad(self.rvq_embs, [0, 0, 0, 1]) + for i in range(d): + emb = embs[i] + ret = ret + F.embedding(code[..., i], emb) + return ret + + def prepare_training_inputs(self, code: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Prepares masked and dropped-out versions of the code for training.""" + b, t, d = code.size() + device = code.device + + src_rate = torch.rand((b, t), device=device) * self.config.max_training_rate + src_masking_rate = get_masking_rate(src_rate, self.config.exponent) + src_num_masking = torch.ceil(src_masking_rate * self.config.num_quantizers).long() + + src_code_mask = torch.ones((b, t, d), dtype=torch.bool, device=device) + src_code_mask = get_mask(src_code_mask, src_num_masking) + src_masked_code = code * src_code_mask + (torch.zeros_like(code) + self.config.codebook_size) * (~src_code_mask) + + if self.config.random_target_masking: + tgt_rate = src_rate + (1.0 - src_rate) * torch.rand((b, t), device=device) + tgt_masking_rate = get_masking_rate(tgt_rate, self.config.exponent) + tgt_num_masking = torch.floor(tgt_masking_rate * self.config.num_quantizers).long() + + tgt_code_mask = torch.ones((b, t, d), dtype=torch.bool, device=device) + tgt_code_mask = get_mask(tgt_code_mask, tgt_num_masking) + tgt_masked_code = code * tgt_code_mask + (torch.zeros_like(code) + self.config.codebook_size) * ( + ~tgt_code_mask + ) + else: + tgt_code_mask = torch.ones((b, t, d), dtype=torch.bool, device=device) + tgt_masked_code = code + + dropout_mask = torch.where( + torch.rand((b, t, 1), device=device) < self.config.quantizer_dropout, + (torch.randint(0, self.config.num_quantizers + 1, (b, t, 1), device=device)), + self.config.num_quantizers, + ) > torch.arange(d, dtype=torch.long, device=device) + dropped_code = code * dropout_mask + (torch.zeros_like(code) + self.config.codebook_size) * (~dropout_mask) + + return src_masked_code, src_code_mask, tgt_masked_code, tgt_code_mask, dropped_code + + def _prepare_conditioning( + self, + context_hidden_state: Tensor | None, + subword_ids: Tensor | None, + subword_mask: Tensor | None, + uncond_dec_flag: Tensor, + ) -> Tensor: + """Computes the final conditioning tensor by combining all sources.""" + cond = torch.zeros((1, 1, self.hidden_size), device=uncond_dec_flag.device) + + if self.embed_context is not None and context_hidden_state is not None: + cond = cond + self.embed_context(context_hidden_state) + + if self.embed_subword is not None and subword_ids is not None: + # Infer subword mask from context if not provided + if subword_mask is None and context_hidden_state is not None: + subword_mask = torch.any(context_hidden_state != 0, dim=-1) + # at least one value should be true, otherwise we can completely skip it to avoid errors + if subword_mask is not None and subword_mask.any(): + cond = cond + self.embed_subword(subword_ids, subword_mask) + + # Replace with null embedding for unconditional generation + cond = torch.where(uncond_dec_flag, self.null_emb, cond) + return cond + + def _compute_losses( + self, + code: Tensor, + lm_logits: Tensor, + mog_logits: Tensor, + mog_mus: Tensor, + mog_mu_res: Tensor, + mog_logs: Tensor, + src_code_mask: Tensor, + tgt_code_mask: Tensor, + audio_mask: Tensor, + ) -> tuple[Tensor, Tensor, Tensor]: + """Helper to compute all losses for the training step.""" + with torch.autocast(code.device.type, enabled=False): + # 1. LM Loss (predicting discrete tokens) + if not self.config.disable_eos_prediction: + eos_mask = (~audio_mask) & F.pad(audio_mask[:, :-1], [1, 0]) + lm_mask = eos_mask | audio_mask + lm_target = torch.where(eos_mask, 1, 0) + lm_loss = ( + F.cross_entropy(lm_logits.transpose(1, 2), lm_target, reduction="none") * lm_mask + ).sum() / lm_mask.sum().clamp_min(1) + else: + lm_loss = 0.0 + + # 2. Continuous & KL Losses (for the MoG head) + target_mask = (~src_code_mask & tgt_code_mask) & audio_mask.unsqueeze(-1) + reduced_target_mask = target_mask.any(dim=-1) + + cont_code_target = self.depthsum_embedding( + code * target_mask + (torch.zeros_like(code) + self.config.codebook_size) * (~target_mask) + ) + mog_logits = mog_logits.float() + mog_mus = mog_mus.float() + mog_mu_res = mog_mu_res.float() + mog_logs = mog_logs.float() + with fp32_precision(): + # Log probability of the true code under each Gaussian component + logp_code = ( + -0.5 * math.log(2 * math.pi) - mog_logs + ) * self.config.latent_size - 0.5 * self.mog_head.dist( + mog_mus, (cont_code_target - mog_mu_res) * torch.exp(-mog_logs) + ) + + # Compute posterior q(k|c) + q_kc = ( + torch.softmax( + logp_code, + -1, + ) + * (1 - self.config.label_smoothing) + + self.config.label_smoothing / self.mog_head.num_predictions + ).detach() + log_q_kc = torch.log(q_kc + 1e-8).detach() + + # Continuous Loss (negative log-likelihood) + c_loss = (-(q_kc * logp_code).sum(-1) * reduced_target_mask).sum() / target_mask.sum().clamp_min(1) + + # KL Divergence Loss + k_loss = ( + (q_kc * (log_q_kc - F.log_softmax(mog_logits, -1))).sum(-1) * reduced_target_mask + ).sum() / target_mask.sum().clamp_min(1) + + return lm_loss, c_loss, k_loss + + def forward( + self, + code: Tensor, + attention_mask: Tensor | None = None, + position_ids: Tensor | None = None, + context_hidden_state: Tensor | None = None, + subword_ids: Tensor | None = None, + subword_mask: Tensor | None = None, + audio_mask: Tensor | None = None, + non_prompt_mask: Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool = False, + training: bool | None = None, + guidance_enabled: bool = False, + generation_config: dict[str, Any] | None = None, + teacher_forcing_inference: bool = False, + ignore_eos_flag_stop: bool = False, + asr_speech_tokens_emb: Tensor | None = None, + audio_prompt_latent: Tensor | None = None, + ) -> RVQEARTTSOutput: + """ + Performs a forward pass handling training, generation, or single-step inference. + + Args: + code (Tensor): Input audio codes. For training, this is the ground truth. + For generation, this is the previously generated code token. + attention_mask (Tensor | None): Attention mask for the backbone transformer. + position_ids (Tensor | None): Position ids for the backbone transformer. + context_hidden_state (Tensor | None): Conditioning from a language model. + subword_ids (Tensor | None): Subword token IDs for conditioning. + subword_mask (Tensor | None): Mask for subword IDs. + audio_mask (Tensor | None): Mask for valid audio positions (for training and inference initialization). + past_key_values (Cache | None): Cache for past key-values for fast decoding. + use_cache (bool): If True, returns the updated `past_key_values`. + training (bool | None): Explicitly set training mode. If `None`, uses `self.training`. + guidance_enabled (bool): If True, duplicates inputs internally to run both + conditional and unconditional passes + generation_config (dict[str, Any] | None): If provided, triggers an iterative code generation. + + Returns: + RVQEARTTSOutput: A dataclass containing losses (for training) or generated outputs + and the cache (for inference). + """ + + # Determine operating mode. + if training is None: + training = self.training + + if audio_mask is not None: + if training: + (src_masked_code, src_code_mask, tgt_masked_code, tgt_code_mask, dropped_code) = ( + self.prepare_training_inputs(code) + ) + uncond_dec_flag = torch.rand(code.size(0), 1, 1, device=code.device) < self.config.p_uncond + else: + dropped_code = code + uncond_dec_flag = torch.zeros(code.size(0), 1, 1, device=code.device, dtype=torch.bool) + + # code_embeds = ( + # self.embed_code(self.depthsum_embedding(F.pad(dropped_code[:, :-1], [0, 0, 1, 0]))) + # + (audio_mask & (~F.pad(audio_mask[:, :-1], [1, 0]))).unsqueeze(-1) * self.bos_emb + # ) + + # Shifted code + shifted_code = F.pad(dropped_code[:, :-1], [0, 0, 1, 0]) + + # Base embeddings + code_embed = self.depthsum_embedding(shifted_code) + + # BOS mask + bos_mask = audio_mask & (~F.pad(audio_mask[:, :-1], [1, 0])) # [B, T] + bos_mask = bos_mask.unsqueeze(-1) # [B, T, 1] + + # Mask for tokens BEFORE BOS + pre_bos_mask = bos_mask.cumsum(dim=1) == 0 # [B, T, 1] + + # Apply projection to model size + code_embed = self.embed_code(code_embed) + + # Choose projection + if self.config.get("audio_prompt_encoder_config", None): + # Dedicated projection for audio prompt (pre-BOS) + if audio_prompt_latent is None: + prompt_attn_mask = pre_bos_mask.squeeze(-1).long() + audio_prompt_latent = self.audio_prompt_encoder( + inputs_embeds=code_embed, + attention_mask=prompt_attn_mask, + return_dict=True, + ).last_hidden_state + + code_embed = torch.where( + pre_bos_mask, + audio_prompt_latent, + code_embed, + ) + + if self.config.get("use_audio_prompt_frozen_projection", False): + if audio_prompt_latent is None: + W = self.audio_prompt_projection_W.to(code_embed.device, code_embed.dtype) + audio_prompt_latent = torch.nn.functional.linear(code_embed, W.T) + + code_embed = torch.where( + pre_bos_mask, + audio_prompt_latent, + code_embed, + ) + + # Add BOS embedding + code_embeds = code_embed + bos_mask * self.bos_emb + + else: # Inference + code_embeds = self.embed_code(self.depthsum_embedding(code)) + uncond_dec_flag = torch.zeros(code.size(0), 1, 1, device=code.device, dtype=torch.bool) + + if guidance_enabled: + assert not training, "Classifier-free guidance can only be used when `training` is False." + code_embeds = torch.cat([code_embeds] * 2, 0) + if attention_mask is not None: + attention_mask = torch.cat([attention_mask] * 2, 0) + if position_ids is not None: + position_ids = torch.cat([position_ids] * 2, 0) + if context_hidden_state is not None: + context_hidden_state = torch.cat([context_hidden_state] * 2, 0) + if subword_ids is not None: + subword_ids = torch.cat([subword_ids] * 2, 0) + if subword_mask is not None: + subword_mask = torch.cat([subword_mask] * 2, 0) + if asr_speech_tokens_emb is not None: + asr_speech_tokens_emb = torch.cat([asr_speech_tokens_emb] * 2, 0) + + uncond_dec_flag = torch.cat([uncond_dec_flag, torch.ones_like(uncond_dec_flag)], 0) + + # Prepare conditioning + cond = self._prepare_conditioning(context_hidden_state, subword_ids, subword_mask, uncond_dec_flag) + + if asr_speech_tokens_emb is not None: + cond = cond + asr_speech_tokens_emb + + if self.config.use_gated_fusion_for_text_audio: + inputs_embeds = self.gated_fusion_audio_text(code_embeds, cond) + else: + inputs_embeds = code_embeds + cond + + # Main backbone pass + backbone_outputs = self.backbone( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + ) + hidden_states = backbone_outputs.last_hidden_state + + if audio_mask is not None and training: + # --- Training-specific loss computation --- + if not self.config.disable_eos_prediction: + lm_logits = self.lm_head(hidden_states) + else: + lm_logits = None + + mog_input_embeds = self.embed_code(self.depthsum_embedding(src_masked_code)) + if self.config.random_target_masking: + mog_input_embeds = mog_input_embeds + self.embed_target_mask((tgt_code_mask.sum(-1) - 1).clamp_min(0)) + mog_input_embeds = mog_input_embeds + hidden_states + mog_logits, mog_mus, mog_mu_res, mog_logs = self.mog_head(mog_input_embeds) + + lm_loss, c_loss, k_loss = self._compute_losses( + code, lm_logits, mog_logits, mog_mus, mog_mu_res, mog_logs, src_code_mask, tgt_code_mask, audio_mask + ) + total_loss = lm_loss + c_loss + k_loss + + return RVQEARTTSOutput( + loss=total_loss, + lm_loss=lm_loss, + c_loss=c_loss, + k_loss=k_loss, + hidden_states=hidden_states, + audio_prompt_latent=audio_prompt_latent, + ) + else: # Inference + if not generation_config: + return RVQEARTTSOutput( + hidden_states=hidden_states, + past_key_values=backbone_outputs.past_key_values, + audio_prompt_latent=audio_prompt_latent, + ) + else: + if teacher_forcing_inference: + generated_codes, lm_logits, eos_flag = self.generate_teacher_forcing( + hidden_states, generation_config + ) + else: + generated_codes, lm_logits, eos_flag = self.generate_step( + hidden_states, ignore_eos_flag_stop=ignore_eos_flag_stop, **generation_config + ) + return RVQEARTTSOutput( + past_key_values=backbone_outputs.past_key_values, + codes=generated_codes, + lm_logits=lm_logits, + eos_flag=eos_flag, + hidden_states=hidden_states, + ) + + @torch.no_grad() + def generate_teacher_forcing(self, hidden_states: Tensor, generation_config: dict): + """ + Teacher-forcing wrapper for generate_step, processing all frames in parallel + using a per-frame loop internally. + + Args: + hidden_states: [B, T, H] hidden states + generation_config: kwargs for self.generate_step() + + Returns: + generated_codes: [B, T, ...] generated codes per frame + lm_logits: [B, T, vocab_size] language model logits + eos_flag: [B, T] boolean tensor indicating EOS + """ + B, T, H = hidden_states.shape + + # Preallocate caches + generated_codes_cache = [] + lm_logits_cache = [] + eos_flag_cache = [] + + # Iterate over time steps (frames) + for t in range(T): + # extract one frame (as the original generate_step expects) + frame_hidden = hidden_states[:, t, :] # [B, H] + + # call original generate_step + generated_codes, lm_logits, eos_flag = self.generate_step( + frame_hidden.unsqueeze(1), + **generation_config, # keep batch dim + frame dim + ) + if generated_codes is not None: + # store in cache + generated_codes_cache.append(generated_codes) + lm_logits_cache.append(lm_logits) + eos_flag_cache.append(eos_flag) + + # Stack results along time dimension + generated_codes = torch.stack(generated_codes_cache, dim=1) # [B, T, ...] + if not self.config.disable_eos_prediction: + lm_logits = torch.stack(lm_logits_cache, dim=1) # [B, T, vocab_size] + eos_flag = torch.stack(eos_flag_cache, dim=1) # [B, T] + else: + lm_logits = None + eos_flag = None + + return generated_codes, lm_logits, eos_flag + + @torch.no_grad() + def generate_step( + self, + hidden_states: Tensor, + num_iter: int, + guidance_scale: list[float] | float | None = None, + top_p_or_k: list[float | int] | float | int | None = None, + noise_scale: list[float] | float | None = None, + exponent: float | None = None, + eos_threshold: float | None = None, + ignore_eos_flag_stop: bool = False, + ) -> tuple[Tensor | None, Tensor, Tensor]: + """ + Performs the iterative unmasking process for a single generation step. + + This function takes the hidden state from the backbone transformer and generates + codes through an iterative unmasking process. + + Args: + hidden_states (Tensor): The hidden states from the backbone. If using CFG, + this should be the combined [uncond, cond] tensor. + num_iter (int): The number of unmasking iterations. + guidance_scale (list[float] | float | None): The scale for Classifier-Free Guidance. + top_p_or_k (list[float | int] | float | int | None): The value for top-p or top-k sampling. + noise_scale (list[float] | float | None): The scale of noise to add during MoG sampling. + exponent (float | None): The exponent for the masking schedule. + eos_threshold (float | None): The threshold for EOS prediction. + + Returns: + tuple[Tensor | None, Tensor, Tensor]: A tuple containing: + - the generated codes. + - The logits from `lm_head`. + - The EOS flag. + """ + # 1. Preparation + if guidance_scale is not None: + if not isinstance(guidance_scale, list): + guidance_scale = [guidance_scale] * (1 + num_iter) # includes one step for `lm_head` + assert len(guidance_scale) == 1 + num_iter + if top_p_or_k is not None: + if not isinstance(top_p_or_k, list): + top_p_or_k = [top_p_or_k] * (1 + num_iter) # includes one step for `lm_head` + assert len(top_p_or_k) == 1 + num_iter + if noise_scale is not None: + if not isinstance(noise_scale, list): + noise_scale = [noise_scale] * num_iter + assert len(noise_scale) == num_iter + if exponent is None: + exponent = self.config.exponent + + if guidance_scale is not None: + # The effective batch size is halved + hidden_states, uncond_hidden_states = hidden_states.chunk(2, dim=0) + else: + uncond_hidden_states = hidden_states[:0, :0, :0] + + b, t, _ = hidden_states.size() + d = self.config.num_quantizers + device = hidden_states.device + + # 2. Predict the discrete part of the code + if not self.config.disable_eos_prediction: + if guidance_scale is not None: + lm_logits = self.lm_head(hidden_states + guidance_scale[0] * (hidden_states - uncond_hidden_states)) + else: + lm_logits = self.lm_head(hidden_states) + if top_p_or_k is not None: + lm_logits = ( + TopPLogitsWarper(top_p_or_k[0])( + None, + lm_logits.view(-1, lm_logits.size(-1)), + ).view_as(lm_logits) + if isinstance(top_p_or_k[0], float) + else TopKLogitsWarper(top_p_or_k[0])( + None, + lm_logits.view(-1, lm_logits.size(-1)), + ).view_as(lm_logits) + ) + + with fp32_precision(): + lm_logits = F.log_softmax(lm_logits, -1) + + if eos_threshold is not None: + eos_flag = lm_logits[..., -1] > eos_threshold + else: + eos_flag = lm_logits.argmax(-1) == 1 + + if torch.all(eos_flag) and not ignore_eos_flag_stop: + return None, lm_logits, eos_flag + else: + lm_logits = None + eos_flag = None + + # Initialize the full code tensor + code = torch.zeros((b, t, d), dtype=torch.long, device=device) + self.config.codebook_size + + # 3. Set up the iterative denoising schedule for the continuous part + rates = torch.linspace(0.0, 1.0, num_iter + 1, device=device)[:-1].unsqueeze(-1) + masking_rates = get_masking_rate(rates, exponent=exponent) + num_maskings = torch.ceil(masking_rates * self.config.num_quantizers).long() + + ks = num_maskings - F.pad(num_maskings[1:], [0, 0, 0, 1]) + + # 4. Iteratively unmask the continuous part of the code + cnt = 0 + for i, k in enumerate(ks): + if torch.all(k == 0): + continue + + # Prepare input for the MoG head + guidance_scale_i = guidance_scale[i] if guidance_scale is not None else 0.0 + top_p_or_k_i = top_p_or_k[i] if top_p_or_k is not None else 1.0 + noise_scale_i = noise_scale[i] if noise_scale is not None else 1.0 + + mog_input_embeds = self.embed_code(self.depthsum_embedding(code)) + if self.config.random_target_masking: + mog_input_embeds += self.embed_target_mask(cnt + k - 1) + if guidance_scale_i > 0.0: + mog_input_embeds = torch.cat( + [mog_input_embeds + hidden_states, mog_input_embeds + uncond_hidden_states], 0 + ) + else: + mog_input_embeds += hidden_states + + mog_mu, mog_logs = self.mog_head.infer( + mog_input_embeds, + guidance_scale=guidance_scale_i, + top_p_or_k=top_p_or_k_i, + ) + z = mog_mu + torch.exp(mog_logs) * torch.randn_like(mog_mu) * noise_scale_i + code = depthsum_encoding_step(self.rvq_embs, z, code, cnt, k[0].item()) + cnt += k[0].item() + return code, lm_logits, eos_flag + + def load_state_dict(self, state_dict, strict: bool = True): + try: + super().load_state_dict(state_dict, strict=strict) + except RuntimeError: + logging.info("Error loading model state_dict !! Retrying with partial initialization!") + model_dict = set_model_dict_for_partial_init(state_dict, self.state_dict()) + super().load_state_dict(model_dict, strict=False) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py new file mode 100644 index 000000000..9385f0c75 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py @@ -0,0 +1,874 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_vae_codec.py; training-only code removed. + +import functools +import math +from collections.abc import Callable +from contextlib import contextmanager +from typing import Any, Concatenate + +import torch +from omegaconf import DictConfig +from torch import Tensor, nn +from torch.nn import functional as F + + +@contextmanager +def disable_tf32(): + prev = torch.backends.cudnn.allow_tf32 + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.backends.cudnn.allow_tf32 = prev + + +# ============================================================================== +# Utility Functions +# ============================================================================== + + +def zero_module(module: nn.Module) -> nn.Module: + """ + Zeros out the parameters of a PyTorch module in-place. + + This is a utility function that iterates through all parameters of a given + `nn.Module` and sets their values to zero. This is often used for specific + initialization strategies, for example in diffusion models where some layers + are initialized to zero. + + From: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/nn.py#L68 + + Args: + module (nn.Module): The PyTorch module to be zeroed. + + Returns: + nn.Module: The same module with its parameters zeroed. + """ + for p in module.parameters(): + # p.detach().zero_() performs the operation in-place without tracking it in autograd + p.detach().zero_() + return module + + +def sequence_mask(lengths: Tensor, max_length: int | None = None) -> Tensor: + """ + Creates a boolean mask from a 1D tensor of sequence lengths. + + This function is useful for masking out padding in sequences. Given a tensor + of lengths, it produces a 2D boolean tensor where `mask[i, j]` is `True` if + `j < lengths[i]` and `False` otherwise. + + Example: + >>> lengths = torch.tensor([1, 3, 2]) + >>> sequence_mask(lengths) + tensor([[ True, False, False], + [ True, True, True], + [ True, True, False]]) + + Args: + lengths (Tensor): A 1D tensor of integer lengths. Shape: `[batch_size]`. + max_length (int | None, optional): The maximum length of the mask. If None, + it is inferred from the maximum value + in `lengths`. Defaults to None. + + Returns: + Tensor: The boolean mask. Shape: `[batch_size, max_length]`. + """ + if max_length is None: + # If max_length is not provided, use the longest sequence length in the batch. + max_length = int(lengths.max().item()) + + # Create a range tensor from 0 to max_length - 1 + x = torch.arange(max_length, dtype=lengths.dtype, device=lengths.device) + + # Compare each length with the range tensor to create the mask. + # `x.unsqueeze(0)` is `[1, max_length]` + # `lengths.unsqueeze(1)` is `[batch_size, 1]` + # Broadcasting takes care of the comparison. + return x.unsqueeze(0) < lengths.unsqueeze(1) + + +# ============================================================================== +# Signal Processing Functions +# ============================================================================== + + +def spectrogram( + wav: Tensor, + n_fft: int, + hop_length: int, + win_length: int, + window_fn: Callable[Concatenate[int, ...], Tensor] = torch.hann_window, +) -> Tensor: + """ + Computes the Short-Time Fourier Transform (STFT) of a waveform with manual padding. + + This implementation manually applies zero padding before computing the STFT. + This is done to center the analysis window at the beginning of the signal + without using the `center=True` argument in `torch.stft`, giving more control. + + Args: + wav (Tensor): The input audio waveform. + Shape: [batch_size?, time_steps], where batch_size? is an + optional batch dimension. + n_fft (int): The size of the FFT. + hop_length (int): The number of samples between adjacent STFT columns. + win_length (int): The size of the window function. + window_fn (function, optional): The window function to apply. + Defaults to torch.hann_window. + + Returns: + Tensor: The complex-valued spectrogram. + Shape: [batch_size?, n_fft // 2 + 1, num_frames] + """ + # Calculate the padding required on the left and right sides to center the frames. + pad_size_l = (n_fft - hop_length) // 2 + pad_size_r = (n_fft - hop_length) - pad_size_l + + # Use a torch.autocast context to perform STFT in float32 for precision. + with torch.autocast(device_type=wav.device.type, enabled=False): + # Apply reflection padding to the waveform. + wav = F.pad(wav.float(), (pad_size_l, pad_size_r)) + + # Create the window tensor on the same device as the waveform. + window = window_fn(win_length, dtype=torch.float, device=wav.device) + + # Compute the STFT. + # `center=False` because we have already manually padded the signal. + spec = torch.stft( + wav, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=window, + center=False, + normalized=False, + onesided=True, + return_complex=True, + ) + return spec + + +def spec_to_wav( + spec: Tensor, + n_fft: int, + hop_length: int, + win_length: int, + window_fn: Callable[Concatenate[int, ...], Tensor] = torch.hann_window, + constrain_value_range: bool = False, +) -> Tensor: + """ + Converts a spectrogram back into a waveform using the overlap-add method. + This function is an approximate inverse of the `spectrogram` function. + + Args: + spec (Tensor): The input complex-valued spectrogram. + Shape: [batch_size?, dim, time_steps], where batch_size? + is an optional batch dimension. + n_fft (int): The size of the FFT used to create the spectrogram. + hop_length (int): The number of samples between frames in the original signal. + win_length (int): The size of the window function used in the original signal. + window_fn (function, optional): The window function used. Currently only + `torch.hann_window` is supported. + constrain_value_range (bool, optional): If True, constrains the IFFT values + to be within the range of the window. + This ensures that the output values + remain within the range of -1.0 to 1.0. + Defaults to False. + + Returns: + Tensor: The reconstructed waveform. + Shape: [batch_size?, time_steps] + + Raises: + ValueError: If a window function other than `torch.hann_window` is provided. + """ + with torch.autocast(device_type=spec.device.type, enabled=False): + if window_fn != torch.hann_window: + raise ValueError(f"`window_fn` should be 'torch.hann_window', but got '{window_fn}'.") + + # Calculate padding and number of frames + pad = (win_length - hop_length) // 2 + T = spec.size(-1) + window = window_fn(win_length, device=spec.device) + + # 1. Inverse FFT + # Convert from frequency domain back to time domain for each frame. + ifft = torch.fft.irfft(spec, n=n_fft, dim=-2, norm="backward") + window_unsqz = window.unsqueeze(-1) + + # 2. Optionally constrain values + if constrain_value_range: + ifft = torch.where( + ifft >= 0, + torch.minimum(ifft, window_unsqz), + torch.maximum(ifft, -window_unsqz), + ) + + # 3. Apply window to the IFFT result + ifft = ifft * window_unsqz + + # 4. Overlap and Add + # Use `torch.nn.functional.fold` to perform the overlap-add operation efficiently. + # This reconstructs the continuous signal from the windowed frames. + output_size = (T - 1) * hop_length + win_length + wav = F.fold( + ifft, + output_size=(1, output_size), + kernel_size=(1, win_length), + stride=(1, hop_length), + )[..., 0, 0, pad:-pad] + + # 5. Calculate the window envelope for normalization + # This is necessary to correct for the energy added by overlapping windows. + window_sq = window.square().expand(T, -1).transpose(0, 1) + window_envelope = F.fold( + window_sq, + output_size=(1, output_size), + kernel_size=(1, win_length), + stride=(1, hop_length), + ).squeeze()[pad:-pad] + + # 6. Normalize the waveform + # Divide by the window envelope to get the final reconstructed signal. + assert (window_envelope > 1e-11).all(), "Window envelope has zero values, cannot normalize." + wav = wav / window_envelope + + return wav + + +# ============================================================================== +# Basic Modules +# ============================================================================== + + +class CausalConv1dCache: + """ + A cache for managing states in causal 1D convolutions. + + This class is used during autoregressive inference to store and update the + tail of the input to a causal convolution, which is used as padding for the + next time step. This avoids re-computing the entire sequence at each step. + """ + + def __init__(self) -> None: + self.cache: dict[int | str, Tensor] = {} + + def __getitem__(self, layer_id: int | str) -> Tensor: + """Retrieves the cached tensor for a given layer.""" + return self.cache[layer_id] + + def update( + self, + states: Tensor, + layer_id: int | str, + padding: int, + padding_value: int = 0, + flush: bool = False, + ) -> Tensor: + """ + Updates the cache for a specific layer and returns the padded input. + + Args: + states (Tensor): The new input tensor for the current time step. + layer_id (int | str): An identifier for the convolutional layer. + padding (int): The amount of left padding required by the convolution. + padding_value (int, optional): The value to use for initial padding. Defaults to 0. + flush (bool, optional): If True, the cache for this layer is deleted + after use. Defaults to False. + + Returns: + Tensor: The input states concatenated with the cached padding. + """ + device = states.device + dtype = states.dtype + b, c, t = states.size() + + if layer_id not in self.cache: + # Initialize cache with zero padding if it's the first time step + padding_tensor = torch.zeros((b, c, padding), dtype=dtype, device=device) + padding_value + else: + padding_tensor = self.cache[layer_id] + assert padding_tensor.size(2) == padding + + # Concatenate the cached padding with the new states + padded_states = torch.cat([padding_tensor, states], dim=2) + # Update the cache with the tail of the new padded states + self.cache[layer_id] = padded_states[:, :, -padding:] + + if flush: + del self.cache[layer_id] + + return padded_states + + +class LayerNormNd(nn.Module): + """ + A LayerNorm module that works for N-dimensional inputs. + + This implementation normalizes over the channel dimension (dim=1), which is + a common setup for convolutional networks. + + Args: + channels (int): The number of channels of the input tensor. + eps (float, optional): A value added to the denominator for numerical + stability. Defaults to 1e-6. + elementwise_affine (bool, optional): If True, this module has learnable + affine parameters (weight and bias). + Defaults to True. + bias (bool, optional): If True, this module has a learnable bias. + Defaults to True. + """ + + def __init__(self, channels: int, eps=1e-6, elementwise_affine: bool = True, bias: bool = True): + super().__init__() + self.channels = channels + self.eps = eps + + self.weight = nn.Parameter(torch.ones((channels,)), requires_grad=elementwise_affine) + self.bias = nn.Parameter(torch.zeros((channels,)), requires_grad=elementwise_affine and bias) + + def forward(self, x: Tensor) -> Tensor: + # Calculate mean and reciprocal standard deviation over the channel dimension + mean = x.mean(1, keepdim=True) + x_shift = x - mean + # Using rsqrt for potentially better performance + x_rstd = torch.rsqrt(x_shift.pow(2).mean(1, keepdim=True) + self.eps) + + # Reshape weight and bias to be broadcastable with the input tensor + shape = [-1 if i == 1 else 1 for i in range(x.ndim)] + + # Apply normalization and affine transformation + return (x_shift * x_rstd) * self.weight.view(shape) + self.bias.view(shape) + + +class ConvNeXt1d(nn.Module): + """ + A 1D ConvNeXt block adapted for causal convolutions on audio signals. + + This block is a core component of modern convolutional architectures, featuring + a depthwise convolution, layer normalization, and pointwise convolutions to + expand and contract the channel dimension, similar to an inverted bottleneck. + + Implementation adapted from: https://github.com/charactr-platform/vocos + + Args: + dim (int): Number of input and output channels. + intermediate_dim (int): Dimensionality of the intermediate (expanded) layer. + kernel_size (int): The kernel size for the causal depthwise convolution. + identity_init (bool, optional): If True, the final pointwise convolution + is initialized to zero, making the block + an identity function at the start of training. + Defaults to False. + layer_idx (int, optional): An index for this layer, used for caching. Defaults to 0. + """ + + def __init__( + self, + dim: int, + intermediate_dim: int, + kernel_size: int, + identity_init: bool = False, + layer_idx: int = 0, + ): + super().__init__() + self.layer_idx = layer_idx + self.kernel_size = kernel_size + + # Depthwise convolution + self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, groups=dim) + + self.norm = LayerNormNd(dim) + self.pwconv1 = nn.Conv1d(dim, intermediate_dim, 1) # Pointwise/1x1 conv for expansion + self.act = nn.GELU() + + # Pointwise/1x1 conv for projection + if identity_init: + self.pwconv2 = zero_module(nn.Conv1d(intermediate_dim, dim, 1)) + else: + self.pwconv2 = nn.Conv1d(intermediate_dim, dim, 1) + + def forward(self, x: Tensor, cache: CausalConv1dCache | None = None, flush: bool = False) -> Tensor: + residual = x + + # Apply causal padding, either through a cache or manually + if cache is not None: + x = cache.update(x, self.layer_idx, self.kernel_size - 1, flush=flush) + else: + x = F.pad(x, [self.kernel_size - 1, 0]) # Left padding for causality + + # Main ConvNeXt path + x = self.dwconv(x) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + + # Add residual connection + x = residual + x + return x + + +class PreTrainedEMAVariance(nn.Module): + """ + Exponential Moving Average of Variance + """ + + def __init__(self, initial_value: float = 1.0): + super().__init__() + self.variance = nn.Parameter( + torch.tensor(initial_value), + requires_grad=False, + ) + + def forward(self) -> Tensor: + return self.variance + + +class PreTrainedProbabilisticVQ(nn.Module): + def __init__( + self, + channels: int, + num_mixtures: int, + depth: int = 1, + ): + super().__init__() + self.channels = channels + self.num_mixtures = num_mixtures + self.depth = depth + + self.mus_list = nn.ParameterList( + [ + nn.Parameter( + F.normalize(torch.randn(num_mixtures, channels), p=2.0, dim=1) * ((depth - i) / depth), + requires_grad=False, + ) + for i in range(depth) + ] + ) + self._variance_list = nn.ModuleList([PreTrainedEMAVariance() for _ in range(depth)]) + + def encode(self, z: Tensor, return_z_q: bool = False) -> list[Tensor] | tuple[list[Tensor], Tensor]: + r = z + ids_sel = [] + for i in range(self.depth): + mus = self.mus_list[i] + idx_sel = self._dist_sq(r, mus).argmin(-1) # [b, ?, h], [v, h] -> [b, ?] + r = r - F.embedding(idx_sel, mus) + ids_sel.append(idx_sel) + if return_z_q: + return ids_sel, z - r + return ids_sel + + def decode(self, ids_sel: list[Tensor]) -> Tensor: + z = torch.zeros((*ids_sel[0].size(), self.channels), device=ids_sel[0].device) + for i in range(len(ids_sel)): + mus = self.mus_list[i] + z = z + F.embedding(ids_sel[i], mus) + return z # [b, ?, h] + + def _dist_sq(self, z: Tensor, mus: Tensor) -> Tensor: + """ + z: [b, ?, d?, h] + mus: [d?, v, h] + """ + return ( + z.pow(2).sum(-1, keepdim=True) # [b, ?, d?, 1] + + mus.pow(2).sum(-1) # [d?, v] + - 2 * (z.unsqueeze(-2) @ mus.transpose(-1, -2)).squeeze(-2) # [b, ?, d?, h] , [d?, h, v] -> [b, ?, d?, v] + ) + + +class Wav2Latent(nn.Module): + """ + An encoder model that transforms a raw waveform into a latent representation. + + This model first converts the waveform to a spectrogram, then processes it + through a series of ConvNeXt blocks and downsampling convolutional layers + to produce a compressed latent tensor. + + Args: + latent_size (int): The number of channels in the final latent representation. + n_fft (int): The FFT size for the initial spectrogram transformation. + hop_length (int): The hop length for the STFT. + base_hidden_size (int): The base number of channels for the hidden layers. + channel_mult (tuple[int, ...]): A tuple of multipliers for the hidden + size at each stage of downsampling. + rates (tuple[int, ...]): A tuple of downsampling factors (strides) for + the convolutional layers. + num_blocks (int): The number of ConvNeXt blocks per stage. + kernel_size (int): The kernel size for the ConvNeXt blocks. + groups (int): The number of groups for the downsampling convolutions. + """ + + def __init__( + self, + latent_size: int = 1024, + n_fft: int = 32, + hop_length: int = 8, + base_hidden_size: int = 384, + channel_mult: tuple[int, ...] = (1, 2, 4), + rates: tuple[int, ...] = (8, 8, 8), + num_blocks: int = 3, + kernel_size: int = 7, + groups: int = 1, + ): + super().__init__() + self.n_fft = n_fft + self.hop_length = hop_length + + # Initial projection from spectrogram to hidden size + layers: list[nn.Module] = [nn.Conv1d(n_fft + 2, base_hidden_size * channel_mult[0], 1, bias=False)] + + # Downsampling stages + for i in range(len(channel_mult)): + ch_mult, rate = channel_mult[i], rates[i] + hidden_size = base_hidden_size * ch_mult + # Add ConvNeXt blocks for this stage + for j in range(num_blocks): + layers.append(ConvNeXt1d(hidden_size, hidden_size * 4, kernel_size, True, layer_idx=i * num_blocks + j)) + # Add downsampling convolution + next_hidden_size = base_hidden_size * channel_mult[i + 1] if i < len(channel_mult) - 1 else latent_size + layers.append( + nn.Conv1d(hidden_size, next_hidden_size, kernel_size=rate, stride=rate, bias=False, groups=groups) + ) + + self.layers = nn.ModuleList(layers) + + def forward(self, x: Tensor, cache=None, flush: bool = False) -> Tensor: + if cache is not None: + raise NotImplementedError("Caching is not implemented for the encoder.") + + # Convert waveform to spectrogram (magnitude and phase) + with torch.autocast(device_type=x.device.type, enabled=False): + spec = spectrogram(x.squeeze(1), n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.n_fft) + # Split complex spectrogram into real and imaginary, then treat as magnitude and phase + mag, ph = torch.view_as_real(spec).chunk(2, dim=-1) + x = torch.cat([mag, ph], 1).squeeze(-1) + + # Pass through the network + for layer in self.layers: + if isinstance(layer, ConvNeXt1d): + x = layer(x, cache=cache, flush=flush) + else: + x = layer(x) + + # Transpose to [batch, time, channels] for compatibility with transformers + x = x.transpose(-1, -2) + return x + + +class Latent2Wav(nn.Module): + """ + A decoder (vocoder) model that transforms a latent representation back into a raw waveform. + + This model processes a latent tensor through a series of ConvNeXt blocks and + upsampling transposed convolutional layers to produce a spectrogram, which is + then converted back to a waveform using an inverse STFT. + + Args: + latent_size (int): The number of channels in the input latent representation. + n_fft (int): The FFT size for the final spectrogram reconstruction. + hop_length (int): The hop length for the ISTFT. + base_hidden_size (int): The base number of channels for the hidden layers. + channel_mult (tuple[int, ...]): A tuple of multipliers for the hidden + size at each stage of upsampling. + rates (tuple[int, ...]): A tuple of upsampling factors (strides) for + the transposed convolutional layers. + num_blocks (int): The number of ConvNeXt blocks per stage. + kernel_size (int): The kernel size for the ConvNeXt blocks. + groups (int): The number of groups for the upsampling convolutions. + """ + + def __init__( + self, + latent_size: int = 1024, + n_fft: int = 32, + hop_length: int = 8, + base_hidden_size: int = 384, + channel_mult: tuple[int, ...] = (4, 2, 1), + rates: tuple[int, ...] = (8, 8, 8), + num_blocks: int = 3, + kernel_size: int = 7, + groups=1, + ): + super().__init__() + self.n_fft = n_fft + self.hop_length = hop_length + self.spec_cache_idx = (len(channel_mult)) * num_blocks + + layers: list[nn.Module] = [] + + # Upsampling stages + for i in range(len(channel_mult)): + ch_mult, rate = channel_mult[i], rates[i] + hidden_size = base_hidden_size * ch_mult + # Add upsampling transposed convolution + in_size = base_hidden_size * channel_mult[i - 1] if i != 0 else latent_size + layers.append( + nn.ConvTranspose1d(in_size, hidden_size, kernel_size=rate, stride=rate, bias=False, groups=groups) + ) + # Add ConvNeXt blocks for this stage + for j in range(num_blocks): + layers.append(ConvNeXt1d(hidden_size, hidden_size * 4, kernel_size, True, layer_idx=i * num_blocks + j)) + + # Final projection to spectrogram dimensions (magnitude + phase) + layers.append(nn.Conv1d(hidden_size, n_fft + 2, 1, bias=False)) + self.layers = nn.ModuleList(layers) + + def forward(self, x: Tensor, cache=None, flush: bool = False, constrain_value_range: bool = True) -> Tensor: + # Transpose input from [batch, time, channels] to [batch, channels, time] + x = x.transpose(-1, -2) + + # Pass through the network + for layer in self.layers: + if isinstance(layer, ConvNeXt1d): + x = layer(x, cache=cache, flush=flush) + else: + x = layer(x) + + # Convert network output to a complex spectrogram and then to a waveform + with torch.autocast(device_type=x.device.type, enabled=False): + max_mag = 100.0 + # Split output into magnitude and phase components + mag, ph = x.float().chunk(2, dim=1) + # Safeguard to prevent excessively large magnitudes + mag = max_mag * torch.exp(-F.softplus(-mag + math.log(max_mag))) + + # Reconstruct the complex spectrogram from magnitude and phase + # The DC and Nyquist components are real, so their phase is applied via cosine. + mag_dc, mag_mid, mag_nyquist = mag.split([1, mag.size(1) - 2, 1], dim=1) + ph_dc, ph_mid, ph_nyquist = torch.cos(ph).split([1, ph.size(1) - 2, 1], dim=1) + ph_imag = torch.sin(ph[:, 1:-1, :]) + + spec_real = mag_mid * ph_mid + spec_imag = mag_mid * ph_imag + + spec = torch.cat([mag_dc * ph_dc, spec_real + 1j * spec_imag, mag_nyquist * ph_nyquist], 1) + + # Handle caching for autoregressive generation of the spectrogram + if cache is not None: + half_spec_padding = math.ceil(((self.n_fft - self.hop_length) // 2) / self.hop_length) + spec = cache.update(spec, self.spec_cache_idx, padding=half_spec_padding * 2, flush=flush) + if flush: + spec = F.pad(spec, [0, half_spec_padding]) + + # Convert spectrogram to waveform + x = spec_to_wav( + spec, self.n_fft, self.hop_length, self.n_fft, constrain_value_range=constrain_value_range + ).unsqueeze(1) + + if cache is not None: + # Trim the output to remove the padded region from the start + half_wav_padding = half_spec_padding * self.hop_length + x = x[:, :, half_wav_padding:-half_wav_padding] + + return x + + +class RVQVAEModel(nn.Module): + """ + Residual Vector-Quantized Variational Autoencoder (RVQ-VAE) model. + + This model learns a discrete representation of audio by encoding a waveform + into a latent space and then quantizing the latents into discrete codes. + It consists of an encoder, a quantizer, and a decoder. + + Args: + config (DictConfig | dict[str, Any]): A configuration object with model hyperparameters. + """ + + config_class: type[DictConfig] = DictConfig + + def __init__(self, config: DictConfig | dict[str, Any]): + super().__init__() + self.config = config + + self.encoder = Wav2Latent( + latent_size=self.config.latent_size, + n_fft=self.config.n_fft, + hop_length=self.config.hop_length, + base_hidden_size=self.config.base_hidden_size, + channel_mult=self.config.channel_mult, + rates=self.config.rates, + num_blocks=self.config.num_blocks, + kernel_size=self.config.kernel_size, + groups=self.config.groups, + ) + + # Layers for quantization + self.prvq = PreTrainedProbabilisticVQ( + channels=self.config.latent_size, + num_mixtures=self.config.codebook_size, + depth=self.config.num_quantizers, + ) + + self.decoder = Latent2Wav( + latent_size=self.config.latent_size, + n_fft=self.config.n_fft, + hop_length=self.config.hop_length, + base_hidden_size=self.config.base_hidden_size, + channel_mult=tuple(reversed(self.config.channel_mult)), + rates=tuple(reversed(self.config.rates)), + num_blocks=self.config.num_blocks, + kernel_size=self.config.kernel_size, + groups=self.config.groups, + ) + + for p in self.parameters(): + p.requires_grad = False + + def ae_encode(self, x: Tensor, cache: CausalConv1dCache | None = None, flush: bool = False) -> Tensor: + """ + Runs the encoder part of the autoencoder. + + Args: + x (Tensor): Input waveform. Shape: `[batch, 1, time]`. + cache (CausalConv1dCache | None): Not implemented for the encoder. + flush (bool): Not implemented for the encoder. + + Returns: + Tensor: The continuous latent representation. Shape: `[batch, time', channels]`. + """ + assert x.size(1) == 1 and x.dim() == 3, "Input must be a batch of mono audio." + assert x.size(2) % self.config.wav_to_token_ratio == 0, ( + f"Input audio length ({x.size(2)}) must be divisible by the model's " + f"wav_to_token_ratio ({self.config.wav_to_token_ratio}). " + f"Please pad the input to a compatible length." + ) + + if cache is not None: + raise NotImplementedError("Caching is not supported for the encoder.") + + return self.encoder(x, cache=cache, flush=flush) + + def ae_decode( + self, + x: Tensor, + constrain_value_range: bool = True, + cache: CausalConv1dCache | None = None, + flush: bool = False, + ) -> Tensor: + """ + Runs the decoder part of the autoencoder. + + Args: + x (Tensor): The (de-quantized) latent representation. Shape: `[batch, time', channels]`. + constrain_value_range (bool): If True, constrains the output of the ISTFT. + cache (CausalConv1dCache | None): Cache for autoregressive generation. + flush (bool): If True, flushes the cache. + + Returns: + Tensor: The reconstructed waveform. Shape: `[batch, 1, time]`. + """ + return self.decoder(x, constrain_value_range=constrain_value_range, cache=cache, flush=flush) + + def encode(self, x: Tensor, x_len: Tensor) -> tuple[Tensor, Tensor]: + """ + Encodes a waveform into discrete codes. + + Args: + x (Tensor): Input waveform. Shape: `[batch, 1, time]`. + x_len (Tensor): The original lengths of the waveforms in the batch. + + Returns: + tuple[Tensor, Tensor]: A tuple containing: + - The discrete codes. Shape: `[batch, time', n_quantizers]`. + - The lengths of the code sequences. + """ + with disable_tf32(): + z_e = self.ae_encode(x) + code_len = x_len // self.config.wav_to_token_ratio + return self.quantize(z_e), code_len + + def decode( + self, + code: Tensor, + code_len: Tensor | None = None, + constrain_value_range: bool = True, + cache: CausalConv1dCache | None = None, + flush: bool = False, + ) -> tuple[Tensor, Tensor | None]: + """ + Decodes discrete codes back into a waveform. + + Args: + code (Tensor): The discrete codes. Shape: `[batch, time', n_quantizers]`. + code_len (Tensor | None): The lengths of the code sequences. + constrain_value_range (bool): If True, constrains the output of the ISTFT. + cache (CausalConv1dCache | None): Cache for autoregressive generation. + flush (bool): If True, flushes the cache. + + Returns: + tuple[Tensor, Tensor | None]: A tuple containing: + - The reconstructed waveform. Shape: `[batch, 1, time]`. + - The lengths of the reconstructed waveforms. + """ + with disable_tf32(): + z_q = self.dequantize(code) + x_hat = self.ae_decode(z_q, constrain_value_range=constrain_value_range, cache=cache, flush=flush) + wav_len = code_len * self.config.wav_to_token_ratio if code_len is not None else None + return x_hat, wav_len + + def quantize(self, z: Tensor) -> Tensor: + """ + Quantizes a continuous latent tensor into discrete codes. + + Args: + z (Tensor): The continuous latent tensor from the encoder. + Shape: `[batch, time, channels]`. + + Returns: + Tensor: The quantized codes. Shape: `[batch, time, n_quantizers]`. + """ + with disable_tf32(): + ids_sel = self.prvq.encode(z, return_z_q=False) + return torch.stack(ids_sel, -1) + + def dequantize(self, code: Tensor) -> Tensor: + """ + De-quantizes discrete codes back into a continuous latent tensor. + + Args: + code (Tensor): The quantized codes. Shape: `[batch, time, n_quantizers]`. + + Returns: + Tensor: The de-quantized continuous latent tensor. + Shape: `[batch, time, latent_size]`. + """ + ids_sel = [x.squeeze(-1) for x in torch.split(code, 1, -1)] + return self.prvq.decode(ids_sel) + + def forward(self, x: Tensor, constrain_value_range: bool = False) -> Tensor: + """ + Performs a full autoencoding pass: encode, quantize, dequantize, and decode. + + Args: + x (Tensor): The input waveform. Shape: `[batch, 1, time]`. + constrain_value_range (bool): If True, constrains the output of the ISTFT. + + Returns: + Tensor: The reconstructed waveform. Shape: `[batch, 1, time]`. + """ + + with torch.no_grad(): + z_e = self.ae_encode(x) + code = self.quantize(z_e) + z_d = self.dequantize(code) + x_hat = self.ae_decode(z_d, constrain_value_range=constrain_value_range) + return x_hat diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py new file mode 100644 index 000000000..2cf7ba8f4 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/parts/fusion.py; training-only code removed. + +""" +Fusion modules for combining multi-modal embeddings in duplex speech models. + +Supported fusion methods: +- "add" or None: Simple additive fusion with per-modality weights +- "concat": Concatenate (after LayerNorm) and project back to hidden_dim +- "gated_simple": Per-timestep, per-dimension gates with softmax normalization +- "gated_gmu": Gated Multimodal Unit (Arevalo et al., 2017) with per-modality sigmoid gates + +For gated methods, the channel embeddings, tie_and_roll, and per-modality weights +are bypassed - the gating mechanism learns the optimal combination. +""" + +import torch +import torch.nn as nn +from typing import Optional, Literal, Union + +from .compat import logging + +FuseMethod = Literal["add", "concat", "gated_simple", "gated_gmu"] + + +class FusionModule(nn.Module): + """Base class for fusion modules.""" + + def __init__(self, hidden_dim: int): + super().__init__() + self.hidden_dim = hidden_dim + + def forward( + self, + agent_text_embeds: torch.Tensor, + user_audio_embeds: torch.Tensor, + user_text_embeds: torch.Tensor | None = None, + function_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Fuse multi-modal embeddings. + + Args: + agent_text_embeds: Agent text embeddings (B, T, D) + user_audio_embeds: User audio embeddings (B, T, D) + user_text_embeds: User text/ASR embeddings (B, T, D) or None + function_embeds: Function calling channel embeddings (B, T, D) or None + + Returns: + Fused embeddings (B, T, D) + """ + raise NotImplementedError + + +class AddFusion(FusionModule): + """ + Simple additive fusion with optional per-modality weights. + + This is the default behavior matching the original implementation. + """ + + def __init__( + self, + hidden_dim: int, + agent_text_weight: float = 1.0, + user_audio_weight: float = 1.0, + user_text_weight: float = 1.0, + function_weight: float = 1.0, + ): + super().__init__(hidden_dim) + self.agent_text_weight = agent_text_weight + self.user_audio_weight = user_audio_weight + self.user_text_weight = user_text_weight + self.function_weight = function_weight + + def forward( + self, + agent_text_embeds: torch.Tensor, + user_audio_embeds: torch.Tensor, + user_text_embeds: torch.Tensor | None = None, + function_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + output = agent_text_embeds * self.agent_text_weight + output = output + user_audio_embeds * self.user_audio_weight + if user_text_embeds is not None: + output = output + user_text_embeds * self.user_text_weight + if function_embeds is not None: + output = output + function_embeds * self.function_weight + return output + + +class ConcatFusion(FusionModule): + """ + Concatenate embeddings (after LayerNorm) and project back to hidden_dim. + + Supports optional user_text_embeds and function_embeds; zeros are used when + not provided. Uses 4 modalities (agent, audio, user_text, function) so that + function channel can be optionally present. + The linear layer learns to handle this case appropriately. + """ + + def __init__( + self, + hidden_dim: int, + agent_text_weight: float = 1.0, + user_audio_weight: float = 1.0, + user_text_weight: float = 1.0, + function_weight: float = 1.0, + num_modalities: int = 4, + ): + super().__init__(hidden_dim) + self.agent_text_weight = agent_text_weight + self.user_audio_weight = user_audio_weight + self.user_text_weight = user_text_weight + self.function_weight = function_weight + self.num_modalities = num_modalities + + # LayerNorm for each modality + self.ln_agent = nn.LayerNorm(hidden_dim) + self.ln_audio = nn.LayerNorm(hidden_dim) + self.ln_text = nn.LayerNorm(hidden_dim) + self.ln_function = nn.LayerNorm(hidden_dim) + + # Projection layer (4 modalities; pad with zeros when optional channels missing) + self.proj = nn.Linear(hidden_dim * num_modalities, hidden_dim, bias=False) + + def forward( + self, + agent_text_embeds: torch.Tensor, + user_audio_embeds: torch.Tensor, + user_text_embeds: torch.Tensor | None = None, + function_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + # Apply layernorm and weights + agent_normed = self.ln_agent(agent_text_embeds) * self.agent_text_weight + audio_normed = self.ln_audio(user_audio_embeds) * self.user_audio_weight + + if user_text_embeds is not None: + text_normed = self.ln_text(user_text_embeds) * self.user_text_weight + else: + # Zero-pad when ASR embeddings are not available + text_normed = torch.zeros_like(agent_text_embeds) + + if function_embeds is not None: + function_normed = self.ln_function(function_embeds) * self.function_weight + else: + function_normed = torch.zeros_like(agent_text_embeds) + + # Concatenate along hidden dim and project + concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1) + return self.proj(concat) + + +class GatedFusionSimple(FusionModule): + """ + Learns a softmax-normalized gate over 4 streams at each timestep (agent, audio, user_text, function). + + Gate input: LayerNorm(each stream) → concatenate → linear → softmax + Output: weighted sum of LayerNorm'd embeddings where weights are input-dependent. + Optional channels (user_text, function) use zeros when not provided; their gate is masked to 0. + + Note: This is PER-TIMESTEP gating (one scalar weight per modality per timestep). + For PER-DIMENSION gating, use GatedFusionGMU. + """ + + def __init__(self, hidden_dim: int, num_modalities: int = 4): + super().__init__(hidden_dim) + self.num_modalities = num_modalities + + # LayerNorm for each stream (normalizes distributions before gate computation) + self.ln_agent = nn.LayerNorm(hidden_dim) + self.ln_audio = nn.LayerNorm(hidden_dim) + self.ln_text = nn.LayerNorm(hidden_dim) + self.ln_function = nn.LayerNorm(hidden_dim) + + # Gate network: outputs one gate value per modality per timestep + # Input: concatenated normalized modalities (B, T, D*4) + # Output: gates (B, T, 4) - one scalar per modality + self.gate_net = nn.Linear(hidden_dim * num_modalities, num_modalities, bias=True) + + # Initialize for uniform weighting + self._init_uniform() + + def _init_uniform(self): + """Initialize gate network for uniform initial weighting.""" + nn.init.zeros_(self.gate_net.weight) + nn.init.zeros_(self.gate_net.bias) + + def forward( + self, + agent_text_embeds: torch.Tensor, + user_audio_embeds: torch.Tensor, + user_text_embeds: torch.Tensor | None = None, + function_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + B, T, D = agent_text_embeds.shape + + has_user_text = user_text_embeds is not None + has_function = function_embeds is not None + text_input = user_text_embeds if has_user_text else torch.zeros_like(agent_text_embeds) + function_input = function_embeds if has_function else torch.zeros_like(agent_text_embeds) + + # Apply LayerNorm to each stream before gate computation + agent_normed = self.ln_agent(agent_text_embeds) + audio_normed = self.ln_audio(user_audio_embeds) + text_normed = self.ln_text(text_input) + function_normed = self.ln_function(function_input) + + # Concatenate normalized inputs for gate computation + concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1) + + # Compute gates: (B, T, D*4) -> (B, T, 4) + gates = self.gate_net(concat) # (B, T, num_modalities) + + # Mask out optional channels when not present (before softmax) + if not has_user_text: + gates = gates.clone() + gates[:, :, 2] = float("-inf") + if not has_function: + gates = gates.clone() + gates[:, :, 3] = float("-inf") + + # Softmax over modalities for normalization + gates = torch.softmax(gates, dim=-1) # (B, T, 4) + + # Stack normalized modalities for weighted sum: (B, T, 4, D) + stacked = torch.stack([agent_normed, audio_normed, text_normed, function_normed], dim=2) + + # Apply gates (broadcast) and sum over modalities + fused = (stacked * gates.unsqueeze(-1)).sum(dim=2) # (B, T, D) + + return fused + + +class GatedFusionGMU(FusionModule): + """ + Gated Multimodal Unit inspired by Arevalo et al. (2017). + + PER-DIMENSION gating: each dimension can attend to different modalities. + Uses softmax over modalities for numerically stable normalization. + Supports 4 modalities (agent, audio, user_text, function); optional channels use zeros and are masked. + + All streams are LayerNorm'd before gate computation and fusion. + Gate output shape: (B, T, 4, D) with softmax over the modality dimension. + """ + + def __init__(self, hidden_dim: int, num_modalities: int = 4): + super().__init__(hidden_dim) + self.num_modalities = num_modalities + + # LayerNorm for each stream (normalizes distributions before gate computation) + self.ln_agent = nn.LayerNorm(hidden_dim) + self.ln_audio = nn.LayerNorm(hidden_dim) + self.ln_text = nn.LayerNorm(hidden_dim) + self.ln_function = nn.LayerNorm(hidden_dim) + + # Per-modality gate projections (output logits, not activations) + self.gate_agent = nn.Linear(hidden_dim * num_modalities, hidden_dim) + self.gate_audio = nn.Linear(hidden_dim * num_modalities, hidden_dim) + self.gate_text = nn.Linear(hidden_dim * num_modalities, hidden_dim) + self.gate_function = nn.Linear(hidden_dim * num_modalities, hidden_dim) + + # Initialize for uniform weighting (all logits = 0 → equal softmax weights) + self._init_uniform() + + def _init_uniform(self): + """Initialize gate networks for uniform initial weighting.""" + for gate in [self.gate_agent, self.gate_audio, self.gate_text, self.gate_function]: + nn.init.zeros_(gate.weight) + nn.init.zeros_(gate.bias) + + def forward( + self, + agent_text_embeds: torch.Tensor, + user_audio_embeds: torch.Tensor, + user_text_embeds: torch.Tensor | None = None, + function_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + B, T, D = agent_text_embeds.shape + + has_user_text = user_text_embeds is not None + has_function = function_embeds is not None + text_input = user_text_embeds if has_user_text else torch.zeros_like(agent_text_embeds) + function_input = function_embeds if has_function else torch.zeros_like(agent_text_embeds) + + # Apply LayerNorm to each stream before gate computation + agent_normed = self.ln_agent(agent_text_embeds) + audio_normed = self.ln_audio(user_audio_embeds) + text_normed = self.ln_text(text_input) + function_normed = self.ln_function(function_input) + + # Concatenate normalized inputs for gate computation + concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1) + + # Compute per-modality gate logits: each is (B, T, D) + logits_agent = self.gate_agent(concat) + logits_audio = self.gate_audio(concat) + logits_text = self.gate_text(concat) + logits_function = self.gate_function(concat) + + # Stack logits: (B, T, 4, D) + stacked_logits = torch.stack([logits_agent, logits_audio, logits_text, logits_function], dim=2) + + # Mask out optional channels when not present (before softmax) + if not has_user_text: + stacked_logits = stacked_logits.clone() + stacked_logits[:, :, 2, :] = float("-inf") + if not has_function: + stacked_logits = stacked_logits.clone() + stacked_logits[:, :, 3, :] = float("-inf") + + # Softmax over modalities (dim=2) for stable normalization + gates = torch.softmax(stacked_logits, dim=2) # (B, T, 4, D) + + # Stack normalized embeddings: (B, T, 4, D) + stacked_embeds = torch.stack([agent_normed, audio_normed, text_normed, function_normed], dim=2) + + # Fuse with element-wise gating + fused = (gates * stacked_embeds).sum(dim=2) # (B, T, D) + + return fused + + +def create_fusion_module( + fuse_method: Union[str, FuseMethod] | None, + hidden_dim: int, + agent_text_weight: float = 1.0, + user_audio_weight: float = 1.0, + user_text_weight: float = 1.0, + function_weight: float = 1.0, +) -> FusionModule: + """ + Factory function to create the appropriate fusion module. + + Args: + fuse_method: One of None, "add", "concat", "gated_simple", "gated_gmu" + hidden_dim: Hidden dimension of the embeddings + agent_text_weight: Weight for agent text (used by add/concat methods only) + user_audio_weight: Weight for user audio (used by add/concat methods only) + user_text_weight: Weight for user text (used by add/concat methods only) + function_weight: Weight for function channel (used by add/concat only; optional channel) + + Returns: + Instantiated fusion module + + Notes: + - All fusion modules accept optional function_embeds; when None, that channel + contributes nothing (zeros + mask for gated methods). + - For gated methods (gated_simple, gated_gmu), the weight parameters are + ignored as the gating mechanism learns the optimal combination. + - When using gated methods, you should also bypass channel embeddings, + tie_and_roll, and any manual embedding weight parameters in the model. + """ + if fuse_method is None or fuse_method == "add": + logging.info( + f"Using AddFusion with weights: agent={agent_text_weight}, " + f"audio={user_audio_weight}, text={user_text_weight}, function={function_weight}" + ) + return AddFusion( + hidden_dim=hidden_dim, + agent_text_weight=agent_text_weight, + user_audio_weight=user_audio_weight, + user_text_weight=user_text_weight, + function_weight=function_weight, + ) + elif fuse_method == "concat": + logging.info( + f"Using ConcatFusion with weights: agent={agent_text_weight}, " + f"audio={user_audio_weight}, text={user_text_weight}, function={function_weight}" + ) + return ConcatFusion( + hidden_dim=hidden_dim, + agent_text_weight=agent_text_weight, + user_audio_weight=user_audio_weight, + user_text_weight=user_text_weight, + function_weight=function_weight, + ) + elif fuse_method == "gated_simple": + logging.info( + f"Using GatedFusionSimple with hidden_dim={hidden_dim} (learned gating, optional function channel)" + ) + return GatedFusionSimple(hidden_dim=hidden_dim) + elif fuse_method == "gated_gmu": + logging.info(f"Using GatedFusionGMU with hidden_dim={hidden_dim} (learned gating, optional function channel)") + return GatedFusionGMU(hidden_dim=hidden_dim) + else: + raise ValueError( + f"Unknown fuse_method: '{fuse_method}'. Expected one of: None, 'add', 'concat', 'gated_simple', 'gated_gmu'" + ) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py new file mode 100644 index 000000000..36df3a0d3 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/perception.py; training-only code removed. + +import torch +from omegaconf import DictConfig +from torch import nn + +from .compat import ( + Exportable, + NeuralModule, + typecheck, +) + + +class AudioPerceptionModule(NeuralModule, Exportable): + """Audio perception module that consists of audio encoder(s) and modality adapter.""" + + def __init__(self, cfg: DictConfig): + super().__init__() + # Initialize components + self.cfg = cfg + self.preprocessor = self.from_config_dict(cfg.preprocessor) + self.encoder = self.from_config_dict(cfg.encoder) + + if "spec_augment" in cfg and cfg.spec_augment is not None: + self.spec_augmentation = self.from_config_dict(cfg.spec_augment) + else: + self.spec_augmentation = None + self.modality_adapter = self.from_config_dict(cfg.modality_adapter) + if "output_dim" not in cfg.modality_adapter and "d_model" in cfg.modality_adapter: # e.g., conformer encoder + self.proj = nn.Linear(cfg.modality_adapter.d_model, cfg.output_dim) + else: + self.proj = nn.Identity() + + self.modality_adapter_quantizer_levels = cfg.get("modality_adapter_quantizer_levels", None) + if self.modality_adapter_quantizer_levels: + # FiniteScalarQuantizer (nemo.collections.tts.modules.audio_codec_modules) was not + # vendored: `modality_adapter_quantizer_levels` is unset in the + # NVIDIA-NemotronLabs-VoiceChat-11B checkpoint config.json. + raise RuntimeError( + "modality_adapter_quantizer_levels is set in the perception config, but " + "FiniteScalarQuantizer was not vendored (it is unused by the " + "NVIDIA-NemotronLabs-VoiceChat-11B checkpoint). Vendor " + "nemo.collections.tts.modules.audio_codec_modules.FiniteScalarQuantizer to enable it." + ) + + def maybe_preprocess_audio( + self, + input_signal=None, + input_signal_length=None, + processed_signal=None, + processed_signal_length=None, + ): + has_input_signal = input_signal is not None and input_signal_length is not None + has_processed_signal = processed_signal is not None and processed_signal_length is not None + if (has_input_signal ^ has_processed_signal) is False: + raise ValueError( + f"{self.__class__} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive " + " with ``processed_signal`` and ``processed_signal_len`` arguments." + ) + + if not has_processed_signal: + processed_signal, processed_signal_length = self.preprocessor( + input_signal=input_signal, + length=input_signal_length, + ) + return processed_signal, processed_signal_length + + # disable type checks to avoid type-check errors when using Conformer as modality adapter + @typecheck.disable_checks() + def forward( + self, + input_signal=None, + input_signal_length=None, + processed_signal=None, + processed_signal_length=None, + return_encoder_emb=False, + sample_id=None, + ): + processed_signal, processed_signal_length = self.maybe_preprocess_audio( + input_signal, input_signal_length, processed_signal, processed_signal_length + ) + + # Spec augment is not applied during evaluation/testing + if self.spec_augmentation is not None and self.training: + processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_length) + + encoder_emb, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length) + + if self.modality_adapter_quantizer_levels is not None: + encoded = self.modality_adapter_quantizer_bottleneck(encoder_emb.transpose(1, 2)) + encoded, _ = self.modality_adapter_vector_quantizer(inputs=encoded.transpose(1, 2), input_len=None) + encoded = self.modality_adapter_quantizer_projection(encoded.transpose(1, 2)).transpose(1, 2) + else: + encoded = encoder_emb + + encoded, encoded_len = self.modality_adapter(audio_signal=encoded, length=encoded_len) + + # encoded_orig = encoded.clone() + + # b, c, t -> b, t, c + encoded = self.proj(encoded.transpose(1, 2)) + + if return_encoder_emb: + return encoded, encoded_len, encoder_emb.transpose(1, 2) + else: + return encoded, encoded_len + + +class IdentityConnector(NeuralModule, Exportable): + """User to pass encoder's representations as-is to the LLM.""" + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__() + + def forward(self, audio_signal, length=None, *args, **kwargs): + return audio_signal, length diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py new file mode 100644 index 000000000..ed89746b4 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py @@ -0,0 +1,501 @@ +# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) +# nemo/collections/common/tokenizers/tokenizer_spec.py and +# nemo/collections/common/tokenizers/huggingface/auto_tokenizer.py; training-only code removed. +# Only the imports were adjusted (nemo.utils logging -> stdlib shim); the class +# bodies are exact copies. + +from abc import ABC, abstractmethod +from collections import OrderedDict +from typing import List, Optional + +from transformers import AutoTokenizer as AUTOTOKENIZER + +from .compat import logging + +__all__ = ["TokenizerSpec", "AutoTokenizer"] + + +class TokenizerSpec(ABC): + """ + Inherit this class to implement a new tokenizer. + """ + + @abstractmethod + def text_to_tokens(self, text): + """Converts text into a list of tokens.""" + pass + + @abstractmethod + def tokens_to_ids(self, tokens): + """Converts a list of tokens to their corresponding IDs.""" + pass + + @abstractmethod + def ids_to_tokens(self, ids): + """Converts a list of token IDs back to tokens.""" + pass + + @abstractmethod + def text_to_ids(self, text): + """Converts text directly to token IDs.""" + pass + + def add_special_tokens(self, special_tokens: List[str]): + """Adds special tokens (eos, pad, cls...) to vocab.""" + raise NotImplementedError("To be implemented") + + def apply_chat_template(self, *args, **kwargs): + """Applies chat template and tokenizes results""" + raise NotImplementedError("To be implemented") + + @property + def name(self): + """name of the class""" + return type(self).__name__ + + @property + def cls(self): + """Property alias to match MegatronTokenizer; returns cls_id if available.""" + if hasattr(self, "cls_id"): + return self.cls_id + raise AttributeError(f"{type(self).__name__} has no attribute 'cls' or 'cls_id'") + + @property + def sep(self): + """Property alias to match MegatronTokenizer; returns sep_id if available.""" + if hasattr(self, "sep_id"): + return self.sep_id + raise AttributeError(f"{type(self).__name__} has no attribute 'sep' or 'sep_id'") + + @property + def pad(self): + """Property alias to match MegatronTokenizer; returns pad_id if available.""" + if hasattr(self, "pad_id"): + return self.pad_id + raise AttributeError(f"{type(self).__name__} has no attribute 'pad' or 'pad_id'") + + @property + def eod(self): + """Property alias to match MegatronTokenizer; returns eod_id if available.""" + if hasattr(self, "eod_id"): + return self.eod_id + if hasattr(self, "eos_id"): + # Default to end-of-sentence id if end-of-document is not defined. + return self.eos_id + raise AttributeError(f"{type(self).__name__} has no attribute 'eod', 'eod_id', 'eos', or 'eos_id'") + + @property + def bos(self): + """Property alias to match MegatronTokenizer; returns bos_id if available.""" + if hasattr(self, "bos_id"): + return self.bos_id + raise AttributeError(f"{type(self).__name__} has no attribute 'bos' or 'bos_id'") + + @property + def eos(self): + """Property alias to match MegatronTokenizer; returns eos_id if available.""" + if hasattr(self, "eos_id"): + return self.eos_id + raise AttributeError(f"{type(self).__name__} has no attribute 'eos' or 'eos_id'") + + @property + def mask(self): + """Property alias to match MegatronTokenizer; returns mask_id if available.""" + if hasattr(self, "mask_id"): + return self.mask_id + raise AttributeError(f"{type(self).__name__} has no attribute 'mask' or 'mask_id'") + + +class AutoTokenizer(TokenizerSpec): + """ + Wrapper of HuggingFace AutoTokenizer + https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoTokenizer. + + """ + + def __init__( + self, + pretrained_model_name: str, + vocab_file: str | None = None, + merges_file: str | None = None, + mask_token: str | None = None, + bos_token: str | None = None, + eos_token: str | None = None, + pad_token: str | None = None, + sep_token: str | None = None, + cls_token: str | None = None, + unk_token: str | None = None, + additional_special_tokens: List | None = [], + use_fast: bool | None = True, + trust_remote_code: bool | None = False, + include_special_tokens: bool = False, + ): + """ + Args: + pretrained_model_name: corresponds to HuggingFace-AutoTokenizer's 'pretrained_model_name_or_path' input + argument. For more details please refer to the documentation of the `from_pretrained` method here: + https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoTokenizer. + The list of all supported models can be found here: https://huggingface.co/models + vocab_file: path to file with vocabulary which consists + of characters separated by newlines. + mask_token: mask token + bos_token: the beginning of sequence token + eos_token: the end of sequence token. Usually equal to sep_token + pad_token: token to use for padding + sep_token: token used for separating sequences + cls_token: class token. Usually equal to bos_token + unk_token: token to use for unknown tokens + additional_special_tokens: list of other tokens beside standard special tokens (bos, eos, pad, etc.). For + example, sentinel tokens for T5 (, , etc.) + use_fast: whether to use fast HuggingFace tokenizer + include_special_tokens: when True, converting text to ids will include special tokens / prompt tokens (if + any), yielding self.tokenizer(text).input_ids + """ + try: + self._initialize_tokenizer(pretrained_model_name, vocab_file, merges_file, use_fast, trust_remote_code) + assert self.tokenizer, "tokenizer not initialized" + except Exception: + try: + self._initialize_tokenizer( + pretrained_model_name, vocab_file, merges_file, not use_fast, trust_remote_code + ) + assert self.tokenizer, "tokenizer not initialized" + except Exception as e: + raise ValueError( + f"Unable to instantiate HuggingFace AUTOTOKENIZER for {pretrained_model_name}. Exception: {e}" + ) + + self.include_special_tokens = include_special_tokens + self.original_vocab_size = len(self.tokenizer) + special_tokens_dict = {} + + # # setting special tokens, by default the default model's special tokens will be preserved + # # unless passes new values to the special tokens + if unk_token is not None: + special_tokens_dict["unk_token"] = unk_token + if mask_token is not None: + special_tokens_dict["mask_token"] = mask_token + if pad_token is not None: + special_tokens_dict["pad_token"] = pad_token + + # if the model does not have eos_token but has sep_token, + # set eos_token = sep_token, and vice versa + if sep_token is not None: + special_tokens_dict["sep_token"] = sep_token + elif self.tokenizer.sep_token is None and self.tokenizer.eos_token: + special_tokens_dict["sep_token"] = self.tokenizer.eos_token + if eos_token is not None: + special_tokens_dict["eos_token"] = eos_token + elif self.tokenizer.eos_token is None and self.tokenizer.sep_token: + special_tokens_dict["eos_token"] = self.tokenizer.sep_token + + # if the model does not have bos_token but has cls_token, + # set bos_token = cls_token, and vice versa + if bos_token is not None: + special_tokens_dict["bos_token"] = bos_token + elif self.tokenizer.bos_token is None and self.tokenizer.cls_token: + special_tokens_dict["bos_token"] = self.tokenizer.cls_token + if cls_token is not None: + special_tokens_dict["cls_token"] = cls_token + elif self.tokenizer.cls_token is None and self.tokenizer.bos_token: + special_tokens_dict["cls_token"] = self.tokenizer.bos_token + + # add additional special tokens (not standard special tokens such as bos, eod, sep) + if additional_special_tokens is not None: + special_tokens_dict["additional_special_tokens"] = additional_special_tokens + + new_tokens_in_vocab = [] + for token in [mask_token, bos_token, eos_token, pad_token, sep_token, cls_token, unk_token]: + if token is not None and token not in self.tokenizer.get_vocab(): + new_tokens_in_vocab.append(token) + for token in additional_special_tokens: + if token is not None and token not in self.tokenizer.get_vocab(): + new_tokens_in_vocab.append(token) + + if len(new_tokens_in_vocab) > 0: + """ + Special tokens that were not previously included in the tokenizer's vocabulary file will be added to + the vocabulary and, as a result, the model should be resized, for example: + + # define your model + pretrained_model_name = 'roberta-base' + model = nemo_nlp.modules.get_lm_model(pretrained_model_name=pretrained_model_name) + + # define pretrained tokenizer + tokenizer_default = nemo_nlp.modules.get_tokenizer(tokenizer_name=pretrained_model_name) + + special_tokens = {'bos_token': '', + 'cls_token': '', + 'additional_special_tokens': ['', '']} + tokenizer_default.add_special_tokens(special_tokens_dict=special_tokens) + + # resize your model so that the embeddings for newly added tokens are updated during training/finetuning + model.resize_token_embeddings(tokenizer_default.vocab_size) + + See NLP_Tokenizers.ipynb for more details. + """ + logging.warning( + f"{new_tokens_in_vocab} \n will be added to the vocabulary.\n" + f"Please resize your model accordingly, " + f"see NLP_Tokenizers.ipynb for more details." + ) + self.add_special_tokens(special_tokens_dict) + self.space_sensitive = self.text_to_tokens("x y") != self.text_to_tokens("x") + self.text_to_tokens("y") + self._inv_vocab_dict = {} + + def _initialize_tokenizer( + self, + pretrained_model_name: str, + vocab_file: str | None = None, + merges_file: str | None = None, + use_fast: bool | None = False, + trust_remote_code: bool | None = False, + ): + # this logic deals with different huggingface tokenizers having different positional args + if vocab_file is None: + self.tokenizer = AUTOTOKENIZER.from_pretrained( + pretrained_model_name_or_path=pretrained_model_name, + use_fast=use_fast, + trust_remote_code=trust_remote_code, + ) + elif merges_file is None: + self.tokenizer = AUTOTOKENIZER.from_pretrained( + pretrained_model_name_or_path=pretrained_model_name, + vocab_file=vocab_file, + use_fast=use_fast, + trust_remote_code=trust_remote_code, + ) + else: + self.tokenizer = AUTOTOKENIZER.from_pretrained( + pretrained_model_name_or_path=pretrained_model_name, + vocab_file=vocab_file, + merges_file=merges_file, + use_fast=use_fast, + trust_remote_code=trust_remote_code, + ) + + @property + def vocab_size(self): + """ + Returns the size of the tokenizer's vocabulary. + + Returns: + int: The number of tokens in the vocabulary. + """ + return len(self.tokenizer) + + def add_special_tokens(self, special_tokens_dict: dict) -> int: + """ + Adds a dictionary of special tokens (eos, pad, cls...). If special tokens are NOT in the vocabulary, they are + added to it (indexed starting from the last index of the current vocabulary). + + Args: + special_tokens_dict: dict of string. Keys should be in the list of predefined special attributes: + [``bos_token``, ``eos_token``, ``unk_token``, ``sep_token``, ``pad_token``, ``cls_token``, + ``mask_token``, ``additional_special_tokens``]. + Tokens are only added if they are not already in the vocabulary. + + Returns: + Number of tokens added to the vocabulary. + """ + num_tokens_added = self.tokenizer.add_special_tokens(special_tokens_dict) + + if num_tokens_added > 0: + logging.info(f"{num_tokens_added} special tokens added, resize your model accordingly.") + for k in self.tokenizer.SPECIAL_TOKENS_ATTRIBUTES: + setattr(self, k, getattr(self.tokenizer, k, None)) + return num_tokens_added + + def text_to_tokens(self, text): + """ + Converts text into a list of tokens. + + Args: + text (str): Input text to be tokenized. + + Returns: + List[str]: List of tokens. + """ + tokens = self.tokenizer.tokenize(text) + return tokens + + def tokens_to_ids(self, tokens): + """ + Converts a list of tokens to their corresponding IDs. + + Args: + tokens (List[str]): List of tokens to convert. + + Returns: + List[int]: List of token IDs. + """ + ids = self.tokenizer.convert_tokens_to_ids(tokens) + return ids + + def ids_to_tokens(self, ids): + """ + Converts a list of token IDs back to tokens. + + Args: + ids (List[int]): List of token IDs to convert. + + Returns: + List[str]: List of tokens. + """ + tokens = self.tokenizer.convert_ids_to_tokens(ids) + return tokens + + def text_to_ids(self, text): + """ + Converts text directly to token IDs. + + Args: + text (str): Input text to be converted to IDs. + + Returns: + List[int]: List of token IDs. If include_special_tokens is True, will include special tokens from the + tokenizer's configuration. + """ + if self.include_special_tokens: + return self.tokenizer(text).input_ids + tokens = self.text_to_tokens(text) + ids = self.tokens_to_ids(tokens) + return ids + + def apply_chat_template(self, *args, **kwargs): + """Applies chat template and tokenizes results""" + return self.tokenizer.apply_chat_template(*args, **kwargs) + + @property + def vocab(self): + """ + Returns the vocabulary as a list where the index corresponds to the token ID. + + Returns: + List[str]: List of tokens in the vocabulary. + """ + id2vocab = {v: k for k, v in self.tokenizer.vocab.items()} + return [id2vocab[i] for i in range(len(id2vocab))] + + @property + def pad_id(self): + """ + Gets the ID of the padding token. + + Returns: + int or None: The ID of the padding token if it exists, None otherwise. + """ + if getattr(self, "pad_token") is None: + return None + return self.tokens_to_ids([getattr(self, "pad_token")])[0] + + @property + def bos_id(self): + """ + Gets the ID of the beginning-of-sequence token. + + Returns: + int or None: The ID of the BOS token if it exists, None otherwise. + """ + if getattr(self, "bos_token") is None: + return None + return self.tokens_to_ids([getattr(self, "bos_token")])[0] + + @property + def eos_id(self): + """ + Gets the ID of the end-of-sequence token. + + Returns: + int or None: The ID of the EOS token if it exists, None otherwise. + """ + if getattr(self, "eos_token") is None: + return None + return self.tokens_to_ids([getattr(self, "eos_token")])[0] + + @property + def eod(self): + """ + Gets the ID of the end-of-document token (same as EOS token). Required for megatron-core compatibility. + + Returns: + int: The ID of the EOD/EOS token. + """ + return self.tokens_to_ids([getattr(self, "eos_token")])[0] + + @property + def sep_id(self): + """ + Gets the ID of the separator token. + + Returns: + int or None: The ID of the separator token if it exists, None otherwise. + """ + if getattr(self, "sep_token") is None: + return None + return self.tokens_to_ids([getattr(self, "sep_token")])[0] + + @property + def cls_id(self): + """ + Gets the ID of the classifier token. + + Returns: + int or None: The ID of the classifier token if it exists, None otherwise. + """ + if getattr(self, "cls_token") is None: + return None + return self.tokens_to_ids([getattr(self, "cls_token")])[0] + + @property + def unk_id(self): + """ + Gets the ID of the unknown token. + + Returns: + int or None: The ID of the unknown token if it exists, None otherwise. + """ + if getattr(self, "unk_token") is None: + return None + return self.tokens_to_ids([getattr(self, "unk_token")])[0] + + @property + def mask_id(self): + """ + Gets the ID of the mask token. + + Returns: + int or None: The ID of the mask token if it exists, None otherwise. + """ + if getattr(self, "mask_token") is None: + return None + return self.tokens_to_ids([getattr(self, "mask_token")])[0] + + @property + def name(self): + """ + Returns the name of the underlying HuggingFace tokenizer class. + + Returns: + str: Name of the tokenizer class. + """ + return type(self.tokenizer).__name__ + + def save_pretrained(self, save_directory: str): + """Saves tokenizer's vocabulary and other artifacts to the specified directory""" + return self.tokenizer.save_pretrained(save_directory) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py new file mode 100644 index 000000000..3ceaf93de --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-2 Code2Wav for NemotronVoiceChat: RVQ-VAE codec decode to 22.05 kHz PCM. + +``LLM_GENERATION`` stage consuming the talker's per-frame 31-quantizer code +stacks and decoding them with the vendored NeMo RVQ-VAE codec. The decode runs +in fp32 (the NeMo reference wraps codec decode in an fp32 context; bf16 collapses +the reconstruction). + +Input contract (mirrors ``PersonaPlexCode2Wav``): the connector payload carries +``{"codes": {"audio": Tensor[frames, num_quantizers]}}`` per request; the +scheduler-side ``input_ids`` are placeholders. Codes must lie in +``[0, codebook_size)`` — out-of-range codes are a hard error (an early bf16-drift +signal upstream), never clamped. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +def validate_code_stack(codes: torch.Tensor, num_quantizers: int, codebook_size: int) -> torch.Tensor: + """Validate and reshape a code payload to ``[frames, num_quantizers]``. + + Out-of-range codes are a hard error (an early numeric-drift signal + upstream), never clamped. + """ + q = num_quantizers + if codes.dim() == 1: + if codes.numel() % q != 0: + raise ValueError( + f"NemotronVoiceChat code2wav received a flat code sequence of length " + f"{codes.numel()} which is not divisible by num_quantizers={q}." + ) + codes = codes.reshape(-1, q) + if codes.dim() != 2 or codes.shape[-1] != q: + raise ValueError(f"NemotronVoiceChat code2wav expects codes shaped [frames, {q}], got {tuple(codes.shape)}.") + codes = codes.to(dtype=torch.long) + bad = (codes < 0) | (codes >= codebook_size) + if bool(bad.any()): + bad_vals = codes[bad] + raise ValueError( + "NemotronVoiceChat code2wav received out-of-range codec codes " + f"(min={int(bad_vals.min())}, max={int(bad_vals.max())}, count={int(bad.sum())}); " + f"valid range is [0, {codebook_size}). Refusing to clamp — this usually " + "signals numeric drift or a corrupted stage payload upstream." + ) + return codes + + +class NemotronVoiceChatCode2Wav(nn.Module): + """RVQ-VAE codec decoder stage (GenerationModelRunner).""" + + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + codec_cfg = getattr(self.config, "code2wav_config", None) + self._num_quantizers = int(getattr(codec_cfg, "num_quantizers", 31)) + self._codebook_size = int(getattr(codec_cfg, "codebook_size", 1024)) + self._sample_rate = int(getattr(codec_cfg, "sample_rate", 22050)) + # Samples per 12.5 Hz frame (streaming chunks slice new audio by frame). + self._wav_per_frame = int(getattr(codec_cfg, "wav_to_token_ratio", 1764)) + + # Runner-facing capability flags (same contract as PersonaPlexCode2Wav). + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + # Constructed in load_weights (owns the tts_model.audio_codec.* subtree). + self.audio_codec: nn.Module | None = None + + # ------------------------------------------------------------------ + # Runner-facing placeholder hooks. + # ------------------------------------------------------------------ + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None: + return None + + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: + return [{"meta": {"nemotron_voicechat_dummy_profile": True}} for _ in range(num_reqs)] + + # ------------------------------------------------------------------ + # Decode. + # ------------------------------------------------------------------ + def _codes_from_runtime_info(self, runtime_info: Mapping[str, Any] | None) -> torch.Tensor | None: + """Extract the codes payload; None means "no codes info at all". + + An EMPTY tensor is returned as-is: it is an explicit empty payload + (e.g. zero acoustic frames after the prompt trim) and must yield empty + audio, NOT a fallback decode of the placeholder input_ids. + """ + if not isinstance(runtime_info, Mapping): + return None + if "codes" in runtime_info: + codes = runtime_info["codes"] + if isinstance(codes, Mapping): + codes = codes.get("audio") + else: + codes = runtime_info.get("codes.audio") + if isinstance(codes, list | tuple): + codes = torch.as_tensor(codes, dtype=torch.long) + if isinstance(codes, torch.Tensor): + return codes + return None + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + infos = runtime_additional_information or [] + num_req = max(len(infos), 1) + audios: list[torch.Tensor] = [empty] * num_req + srs = [sr_tensor] * num_req + + if self.audio_codec is None: + raise RuntimeError("NemotronVoiceChatCode2Wav.forward called before load_weights().") + device = next(self.audio_codec.parameters()).device + + for i in range(num_req): + runtime_info = infos[i] if i < len(infos) else None + meta = runtime_info.get("meta") if isinstance(runtime_info, Mapping) else None + if isinstance(meta, Mapping) and meta.get("nemotron_voicechat_dummy_profile") is True: + continue + codes = self._codes_from_runtime_info(runtime_info) + if codes is None: + # No input_ids fallback (unlike PersonaPlexCode2Wav): this + # pipeline's scheduler-side prompt is all-zero placeholders, so + # decoding it would produce plausible-sounding garbage whenever + # n_frames % num_quantizers == 0. A missing payload is a bug. + raise ValueError( + "NemotronVoiceChat code2wav request carries no 'codes' payload in its " + "runtime info; the talker stage must ship the [frames, num_quantizers] " + "code stack (placeholder input_ids are never decoded)." + ) + codes = validate_code_stack(codes, self._num_quantizers, self._codebook_size).to(device=device) + frames = codes.shape[0] + if frames == 0: + continue + # Async-chunk streaming: the payload carries the CUMULATIVE trimmed + # code history plus meta.left_context_size = frames whose audio was + # already emitted by earlier chunks. Mirroring NeMo's + # decode_one_audio_step (unbounded window), we re-decode the whole + # prefix and slice off only the new tail — seam-free by + # construction. Sync full-payload chunks carry no left context. + left_context_frames = 0 + if isinstance(meta, Mapping) and meta.get("left_context_size") is not None: + left_context_frames = max(int(meta["left_context_size"]), 0) + elif isinstance(runtime_info, Mapping) and runtime_info.get("meta.left_context_size") is not None: + left_context_frames = max(int(runtime_info["meta.left_context_size"]), 0) + if left_context_frames >= frames: + continue # terminal empty chunk: all frames already emitted + lens = torch.tensor([frames], device=device, dtype=torch.long) + # NeMo decodes the codec strictly in fp32. + with torch.autocast(device_type=device.type, enabled=False): + wav, wav_len = self.audio_codec.decode(codes.unsqueeze(0), lens) + wav = wav.reshape(-1)[: int(wav_len.reshape(-1)[0])] + if left_context_frames > 0: + wav = wav[left_context_frames * self._wav_per_frame :] + audios[i] = wav.detach().to(dtype=torch.float32).cpu() + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios, "sr": srs}, + ) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + raise TypeError(f"NemotronVoiceChatCode2Wav expected OmniOutput, got {type(model_outputs)}") + + # ------------------------------------------------------------------ + # Weights. + # ------------------------------------------------------------------ + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Build the vendored RVQ-VAE codec and load ``tts_model.audio_codec.*``. + + The vLLM safetensors iterator streams the full unified checkpoint; only + the codec subtree is materialized here (everything else is skipped + without copying). The codec stays fp32. + """ + from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.ear_tts_vae_codec import ( + RVQVAEModel, + ) + + codec_cfg = dict(getattr(self.config, "tts_cfg", {}).get("codec_config") or {}) + if not codec_cfg: + raise ValueError( + "NemotronVoiceChat checkpoint config lacks " + "'model.speech_generation.model.codec_config'; cannot build the RVQ-VAE codec." + ) + from omegaconf import DictConfig + + codec = RVQVAEModel(DictConfig(codec_cfg)) + + prefix = "tts_model.audio_codec." + state: dict[str, torch.Tensor] = {} + for name, tensor in weights: + if name.startswith(prefix): + state[name.removeprefix(prefix)] = tensor.to(dtype=torch.float32) + if not state: + raise ValueError( + f"No '{prefix}*' tensors found in the checkpoint; the NemotronVoiceChat " + "unified model.safetensors is required for the code2wav stage." + ) + missing, unexpected = codec.load_state_dict(state, strict=False) + if missing: + raise ValueError(f"NemotronVoiceChat code2wav is missing codec weights: {sorted(missing)[:10]} ...") + if unexpected: + logger.warning( + "NemotronVoiceChat code2wav ignoring %d unexpected codec tensors (e.g. %s)", + len(unexpected), + sorted(unexpected)[:5], + ) + device = self.vllm_config.device_config.device + codec = codec.to(device=device, dtype=torch.float32).eval() + self.audio_codec = codec + logger.info( + "NemotronVoiceChat code2wav loaded RVQ-VAE codec: %d tensors, num_quantizers=%d, sr=%d", + len(state), + self._num_quantizers, + self._sample_rate, + ) + return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py new file mode 100644 index 000000000..f4f98ab00 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py @@ -0,0 +1,490 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-1 talker for NemotronVoiceChat: frame-locked text -> 31-quantizer codes. + +``LLM_AR`` stage driven per frame by the vLLM engine, with the vendored NeMo +EAR-TTS model (28-layer Gemma3-style backbone + MoG head) executed inside the +per-request ``preprocess`` hook. Classifier-free guidance (internal batch +doubling) and MoG sampling happen inside the vendored forward exactly as in +NeMo — vLLM's sampler never touches the audio codes. + +Execution contract (mirrors the NeMo ``NemotronVoiceChat.offline_inference`` +TTS half): + +* The request's ``additional_information`` carries the full frame-aligned text + timeline (built by ``thinker2talker_token_only``), INCLUDING the PAD prompt + region — the talker's KV state at the first acoustic frame depends on having + processed it (NeMo trims prompt-region codes only after the loop). +* vLLM prompt = one placeholder token (timeline position 0; NeMo's loop starts + at t=1). Each decode step t runs ``infer_codes_one_step`` with + ``current_subword_id=timeline[t]``, ``prev_subword_id=timeline[t-1]`` (the + first step uses the warmup's last init subword, as in NeMo), and the previous + step's code stack; the resulting ``[1, 31]`` codes accumulate under + ``("codes","audio")`` in the request state for the full-payload producer. +* Warmup: ``set_init_inputs(speaker_name=...)`` (pre-baked "Aria" latent) + + one ``tts_model(**init_inputs)`` forward seeds ``past_key_values`` and the + initial code, exactly as in the reference. +* Completion: after consuming timeline position ``T-1`` the stage emits a stop + placeholder token (the deploy yaml sets ``stop_token_ids: [1]``); codes for + positions ``1..T-1`` are shipped and the producer trims the prompt region. + +Heavyweight per-request state (HF ``past_key_values`` etc.) lives in a +model-side session dict keyed by request id (PersonaPlex pattern), cleaned up +via ``on_requests_finished``. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import ( + merge_runtime_info, + require_request_id, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + +_CONTINUE_TOKEN = 0 +_STOP_TOKEN = 1 + + +def sanitize_tts_model_cfg(tts_model_cfg: dict[str, Any]) -> dict[str, Any]: + """Make the vendored DuplexEARTTS construction config-only. + + Every parameter comes from the unified checkpoint streamed in + ``load_weights``; the vendored constructors must never download or load + external pretrained weights. Returns a deep copy so the shared stage + config is never mutated: + + * ``tts_config.pretrained_text_name`` would make ``RVQEARTTSModel`` + download a full LLM backbone — dropped (the backbone builds from + ``backbone_type``/``backbone_config``). + * ``pretrained_codec_model`` would make ``setup_audio_codec`` load an + external codec checkpoint at construction — dropped (the codec weights + arrive with the ``tts_model.audio_codec.*`` subtree, and the RVQ + embedding binding is overwritten by the same ``load_state_dict``). + * ``tts_config.cas_config.pretrained_tokenizer_name`` is dropped, matching + NeMo's own inference override (the CAS vocab is baked into the weights). + * ``pretrained_lm_name`` survives ONLY as the tokenizer reference (the + context-LM load path is disabled by ``context_hidden_size: null``); the + ``NEMOTRON_VOICECHAT_LLM_PATH`` override keeps it air-gap friendly. + """ + import copy + import os + + cfg = copy.deepcopy(tts_model_cfg) + cfg.pop("pretrained_codec_model", None) + tts_config = cfg.get("tts_config") + if isinstance(tts_config, dict): + tts_config.pop("pretrained_text_name", None) + cas_cfg = tts_config.get("cas_config") + if isinstance(cas_cfg, dict): + cas_cfg.pop("pretrained_tokenizer_name", None) + llm_override = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") + if llm_override: + cfg["pretrained_lm_name"] = llm_override + return cfg + + +class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): + """vLLM AR stage wrapping the vendored NeMo EAR-TTS per-frame step.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + talker_cfg = getattr(self.config, "talker_config", None) + self._hidden = int(getattr(talker_cfg, "hidden_size", 1152)) + self._vocab = max(int(getattr(talker_cfg, "vocab_size", 1024)), 2) + self._dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16) + self._speaker_name = getattr(self.config, "inference_speaker_name", "Aria") or "Aria" + + # Omni AR runner contract flags. + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = False + self.requires_full_prefix_cached_hidden_states = False + + # Stateful per-request execution is only validated at batch size 1 + # (the shipped offline scope); fail fast on silent multi-request use. + max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1)) + if max_num_seqs != 1: + raise NotImplementedError( + f"NemotronVoiceChat talker supports max_num_seqs=1 only (got {max_num_seqs}); " + "the vendored EAR-TTS session state is not batch-safe." + ) + + # Vendored DuplexEARTTS; constructed in load_weights. + self.tts: nn.Module | None = None + # request_id -> mutable session state (past_key_values, prev code, ...). + self._sessions: dict[str, dict[str, Any]] = {} + + # ------------------------------------------------------------------ + # vLLM model protocol (placeholder token loop). + # ------------------------------------------------------------------ + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return torch.zeros((input_ids.shape[0], self._hidden), device=input_ids.device, dtype=self._dtype) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor: + # The per-frame compute happens in preprocess; the "hidden state" only + # transports the stop flag (column 0) to compute_logits. + if inputs_embeds is not None: + return inputs_embeds + assert input_ids is not None + return self.embed_input_ids(input_ids) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + hs = hidden_states + flag = hs[..., 0].to(torch.float32) # 1.0 => finished + logits = torch.full((hs.shape[0], self._vocab), -1.0e4, device=hs.device, dtype=torch.float32) + logits[:, _CONTINUE_TOKEN] = (1.0 - flag) * 10.0 + logits[:, _STOP_TOKEN] = flag * 10.0 + return logits + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + hidden = model_outputs + info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] + codes_list: list[torch.Tensor] = [] + prompt_len: int | None = None + for info in info_dicts: + if not isinstance(info, dict): + continue + codes = info.get("codes", {}) + audio = codes.get("audio") if isinstance(codes, dict) else None + if isinstance(audio, torch.Tensor): + codes_list.append(audio) + meta = info.get("meta") + reported = meta.get("nvc_logical_prompt_len") if isinstance(meta, dict) else None + if reported is None: + reported = info.get("meta.nvc_logical_prompt_len") + if reported is not None: + prompt_len = int(reported) + if not codes_list: + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + outputs: dict[str, Any] = {"codes": {"audio": torch.cat(codes_list, dim=0)}} + # The prompt length must ride the WIRE payload (multimodal_outputs): + # the async-chunk dispatcher hands exactly this dict to the + # talker2code2wav producer, while request.additional_information is + # overwritten concurrently by arriving thinker chunks. + if prompt_len is not None: + # 1-D on purpose: the wire-payload builder indexes element.shape[0]. + outputs["meta"] = {"nvc_logical_prompt_len": torch.tensor([prompt_len], dtype=torch.long)} + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs=outputs, + ) + + # ------------------------------------------------------------------ + # Per-frame vendored TTS step. + # ------------------------------------------------------------------ + def _try_timeline(self, info: dict[str, Any]) -> torch.Tensor | None: + """Extract the frame-locked text timeline from either payload shape. + + Sync mode ships ``nvc_text_timeline`` (+ ``nvc_logical_prompt_len``) in + additional_information; async-chunk mode replaces the request info with + the connector payload each chunk, carrying the CUMULATIVE timeline under + ``ids.all`` (and the logical prompt length under + ``meta.num_processed_tokens``). Returns None when neither is present. + """ + timeline = info.get("nvc_text_timeline") + if timeline is not None and info.get("nvc_logical_prompt_len") is None: + raise ValueError( + "NemotronVoiceChat talker got 'nvc_text_timeline' without " + "'nvc_logical_prompt_len'; the code2wav producer needs the prompt " + "length to trim prompt-region codes." + ) + if timeline is None: + ids = info.get("ids") + timeline = ids.get("all") if isinstance(ids, dict) else None + if timeline is None: + timeline = info.get("ids.all") + if timeline is None: + return None + if isinstance(timeline, torch.Tensor): + return timeline.reshape(-1).to(torch.long) + return torch.as_tensor(list(timeline), dtype=torch.long) + + @staticmethod + def _logical_prompt_len(info: dict[str, Any]) -> int | None: + """Logical prompt length from either payload shape (sync/async).""" + value = info.get("nvc_logical_prompt_len") + if value is None: + meta = info.get("meta") + value = meta.get("num_processed_tokens") if isinstance(meta, dict) else None + if value is None: + value = info.get("meta.num_processed_tokens") + return int(value) if value is not None else None + + @staticmethod + def _upstream_finished(info: dict[str, Any]) -> bool: + """True when the async producer marked its final chunk (meta.finished).""" + meta = info.get("meta") + flag = meta.get("finished") if isinstance(meta, dict) else None + if flag is None: + flag = info.get("meta.finished") + if isinstance(flag, torch.Tensor): + return bool(flag.reshape(-1)[0].item()) if flag.numel() else False + return bool(flag) + + def _timeline(self, info: dict[str, Any]) -> torch.Tensor: + timeline = self._try_timeline(info) + if timeline is None: + # Explicit failure — no prompt_token_ids fallback: an implicit + # timeline would silently skip the prompt-region trim downstream. + raise ValueError( + "NemotronVoiceChat talker request is missing its timeline metadata " + "('nvc_text_timeline' in additional_information, or 'ids.all' on the " + "async-chunk payload); the stage cannot synthesize speech without " + "the thinker's frame-locked tokens." + ) + return timeline + + def _init_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]: + assert self.tts is not None + timeline = self._timeline(info).to(device) + if timeline.numel() < 2: + raise ValueError( + f"NemotronVoiceChat talker timeline for request {request_id} has " + f"{timeline.numel()} positions; need at least 2 (prompt + one frame)." + ) + # Warmup: pre-baked speaker latent + init prompt forward (NeMo + # offline_inference lines around get_init_inputs/first forward). + self.tts.set_init_inputs(speaker_name=self._speaker_name) + init_inputs = self.tts.get_init_inputs(B=1) + generation_config = self.tts._get_generation_config(self._guidance_enabled()) + init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": self._guidance_enabled()}) + with torch.inference_mode(): + outputs = self.tts.tts_model(**init_inputs) + prompt_len = self._logical_prompt_len(info) + if prompt_len is None: + raise ValueError( + "NemotronVoiceChat talker request is missing its logical prompt length " + "('nvc_logical_prompt_len' in additional_information, or async " + "meta.num_processed_tokens); the code2wav producer cannot trim the " + "prompt-region codes without it." + ) + session = { + "timeline": timeline, + "prompt_len": prompt_len, + "step": 1, # NeMo's loop starts at t=1 + "past_key_values": outputs.past_key_values, + "code": init_inputs["code"][:, -1:], + "first_context_subword_id": init_inputs["subword_ids"][:, -1].unsqueeze(-1), + "tts_initialized": False, + "generation_config": generation_config, + "codes_rows": [], + # Sync mode ships the whole timeline up front via + # nvc_text_timeline; async-chunk mode grows it across chunks and + # the last chunk sets meta.finished. + "sync_mode": info.get("nvc_text_timeline") is not None, + "upstream_finished": self._upstream_finished(info) or info.get("nvc_text_timeline") is not None, + } + self._sessions[request_id] = session + return session + + def _refresh_session_timeline(self, session: dict[str, Any], info: dict[str, Any], device: torch.device) -> None: + """Adopt a longer timeline from the latest async-chunk payload.""" + latest = self._try_timeline(info) + if latest is not None and latest.numel() > int(session["timeline"].numel()): + session["timeline"] = latest.to(device) + if self._upstream_finished(info): + session["upstream_finished"] = True + + def _guidance_enabled(self) -> bool: + return bool(getattr(self.config, "tts_cfg", {}).get("inference_guidance_enabled", True)) + + def _step_session(self, session: dict[str, Any]) -> tuple[torch.Tensor, bool]: + """One NeMo TTS step; returns (codes_row [1, Q], finished).""" + assert self.tts is not None + timeline: torch.Tensor = session["timeline"] + t = int(session["step"]) + total = int(timeline.numel()) + if t >= total: + # The drain loop in preprocess never enters here; a direct call + # past the received timeline is a programming error. Hard-fail + # rather than fabricating a frame — a repeated/guessed subword + # would silently desync the audio from the text channel. + raise RuntimeError( + f"NemotronVoiceChat talker stepped past the received timeline: step {t} " + f"needs timeline position {t} but only {total} positions have arrived " + f"(upstream_finished={session['upstream_finished']})." + ) + current_subword_id = timeline[t].reshape(1, 1) + if not session["tts_initialized"]: + prev_subword_id = session["first_context_subword_id"] + session["tts_initialized"] = True + else: + prev_subword_id = timeline[t - 1].reshape(1, 1) + current_subword_mask = torch.ones(1, 1, device=timeline.device, dtype=torch.bool) + with torch.inference_mode(): + code, past_key_values = self.tts.infer_codes_one_step( + current_subword_id=current_subword_id, + prev_subword_id=prev_subword_id, + current_subword_mask=current_subword_mask, + prev_audio_tokens=session["code"], + past_key_values=session["past_key_values"], + guidance_enabled=self._guidance_enabled(), + generation_config=session["generation_config"], + ignore_eos_flag_stop=True, + ) + session["code"] = code + session["past_key_values"] = past_key_values + session["step"] = t + 1 + # Async mode: exhausting the received timeline only ends the request + # once the upstream marked its final chunk; otherwise more entries are + # in flight and the scheduler parks us until they arrive. + finished = (t + 1) >= int(session["timeline"].numel()) and bool(session["upstream_finished"]) + return code.reshape(1, -1).to(torch.long), finished + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + info = merge_runtime_info(info_dict) + device = input_ids.device + span = int(input_ids.shape[0]) + is_prefill_raw = info.get("_omni_is_prefill") + is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1 + request_id = require_request_id(info, "talker") + + embeds = torch.zeros((span, self._hidden), device=device, dtype=self._dtype) + if is_prefill or request_id not in self._sessions: + # The 1-token vLLM prefill is the session boundary: run the vendored + # warmup here so the first decode step is a pure NeMo t=1 step. + self._init_session(request_id, info, device) + return input_ids, embeds, {} + + session = self._sessions[request_id] + self._refresh_session_timeline(session, info, device) + # Sync mode: one NeMo TTS step per engine step (the verified parity + # path, unchanged). Async-chunk mode: DRAIN every received timeline + # position on each wake. The engine tokens are placeholders + # (CONTINUE/STOP) and the codes ride the cumulative payload, so + # multiple TTS steps per engine step are safe — and necessary: + # * the save_async background thread can coalesce thinker steps into + # one chunk (fewer wakes than new positions), and + # * the prompt-region PAD steps (logical_prompt_len - 1 of them) all + # become available with chunk 0 and must not cost one upstream + # chunk each. + finished = False + steps_run = 0 + while int(session["step"]) < int(session["timeline"].numel()): + codes_row, finished = self._step_session(session) + session["codes_rows"].append(codes_row.cpu()) + steps_run += 1 + if session["sync_mode"]: + break + if steps_run == 0: + # Zero-progress wake (duplicate/terminal-marker chunk): no TTS step + # was run, so emit a plain CONTINUE placeholder — never a + # fabricated frame. The request finishes only when the upstream is + # exhausted AND the timeline is fully consumed. + finished = bool(session["upstream_finished"]) + if finished: + embeds[:, 0] = 1.0 + return input_ids, embeds, {} + cumulative = torch.cat(session["codes_rows"], dim=0) + if finished: + embeds[:, 0] = 1.0 # stop flag -> compute_logits emits the stop token + info_update: dict[str, Any] = { + "codes": {"audio": cumulative}, + # nvc_logical_prompt_len rides EVERY step's payload: the request's + # additional_information is overwritten both by arriving thinker + # chunks and by these model updates, so the async code2wav producer + # reads the prompt length from the per-step multimodal_output + # instead of racing the request-level dict. + "meta": {"nvc_tts_step": int(session["step"]), "nvc_logical_prompt_len": int(session["prompt_len"])}, + } + return input_ids, embeds, info_update + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + for request_id in finished_req_ids: + self._sessions.pop(str(request_id), None) + + # ------------------------------------------------------------------ + # Weights. + # ------------------------------------------------------------------ + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Build the vendored DuplexEARTTS and load the ``tts_model.*`` subtree. + + The backbone/MoG run in the engine dtype (bf16 by default); the codec + submodule inside DuplexEARTTS stays fp32 (it is not used by this stage — + decode happens in code2wav — but keeping its dtype canonical avoids + surprises if warmup paths touch it). + """ + from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.duplex_ear_tts import ( + DuplexEARTTS, + ) + + tts_model_cfg = dict(getattr(self.config, "tts_cfg", {}) or {}) + if not tts_model_cfg: + raise ValueError( + "NemotronVoiceChat checkpoint config lacks 'model.speech_generation.model'; cannot build the talker." + ) + tts_model_cfg = sanitize_tts_model_cfg(tts_model_cfg) + tts_data = dict(getattr(self.config, "tts_data", {}) or {}) + tts_data.setdefault("source_sample_rate", 22050) + tts_data.setdefault("target_sample_rate", int(getattr(self.config, "target_sample_rate", 22050))) + tts_data.setdefault("frame_length", float(getattr(self.config, "frame_length", 0.08))) + # DuplexEARTTS consumes the whole speech_generation section layout + # ({"data": ..., "model": ...}), matching NeMo's constructor call. + tts = DuplexEARTTS({"data": tts_data, "model": tts_model_cfg}) + + prefix = "tts_model." + state: dict[str, torch.Tensor] = {} + for name, tensor in weights: + if not name.startswith(prefix): + continue + target = name.removeprefix(prefix) + if target.startswith("audio_codec."): + state[target] = tensor.to(dtype=torch.float32) + else: + state[target] = tensor.to(dtype=self._dtype) + if not state: + raise ValueError( + "No 'tts_model.*' tensors found in the checkpoint; the NemotronVoiceChat " + "unified model.safetensors is required for the talker stage." + ) + missing, unexpected = tts.load_state_dict(state, strict=False) + real_missing = [m for m in missing if not m.startswith("audio_codec.")] + if real_missing: + raise ValueError(f"NemotronVoiceChat talker is missing TTS weights: {sorted(real_missing)[:10]} ...") + if unexpected: + logger.warning( + "NemotronVoiceChat talker ignoring %d unexpected TTS tensors (e.g. %s)", + len(unexpected), + sorted(unexpected)[:5], + ) + device = self.vllm_config.device_config.device + tts = tts.to(device=device) + tts.tts_model.to(dtype=self._dtype) + if hasattr(tts, "audio_codec") and tts.audio_codec is not None: + tts.audio_codec.to(dtype=torch.float32) + self.tts = tts.eval() + logger.info( + "NemotronVoiceChat talker loaded EAR-TTS: %d tensors, speaker=%s, guidance=%s", + len(state), + self._speaker_name, + self._guidance_enabled(), + ) + return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py new file mode 100644 index 000000000..9f237a885 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py @@ -0,0 +1,596 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-0 thinker for NemotronVoiceChat: 16 kHz speech -> frame-locked text. + +``LLM_AR`` stage combining: + +* the vendored NeMo FastConformer perception (mel + 24-layer encoder + linear + projection to the LLM width, 12.5 Hz frames); +* the additive channel fusion from the NeMo duplex STT model + (``fused[t] = E(text_prev) * w_text + timeline[t] * w_audio + + E(func_prev) * w_func`` with the checkpoint's AddFusion weights); +* the NemotronH (hybrid Mamba) backbone on vLLM paged execution via + ``init_vllm_registered_model`` — this is where vLLM earns its keep; +* the separate ``lm_head`` (text channel, sampled by the vLLM engine) and + ``function_head`` (greedy, maintained internally per request). + +Timeline contract (verified against NeMo ``_init_inference``/``_step_zero``/ +``_step_inference``; the golden shape test guards it): + +* The vLLM prompt covers ``logical_prompt_token_len + 1`` positions: the system + prompt tokens (``[bos] + tokens + [eos]``) fused with a PAD text channel + (NeMo's ``_get_bos_embedding`` returns the PAD embedding, so position 0 is + uniform with the rest of the prompt region), plus acoustic frame 0 as the + final prefill position. The first engine-sampled token is therefore NeMo's + ``gen_text[prompt_len]`` — the first acoustic-frame prediction. +* Each decode step t consumes conformer frame t (held in per-request state) + fused with the previously sampled text token and the previous greedy + function token. Text EOS never terminates generation (``ignore_eos``); the + example sets ``max_tokens = acoustic_frame_count``. Never set ``min_tokens``: + the tokenizer's eos_token is — the frame-locked PAD/silence + token — and min_tokens masks "EOS" to -inf, forbidding silence entirely. +* The function channel exists for numerical correctness (fusion weight 2.0) + even though tool calling is out of scope; the greedy token timeline is + exposed for debugging via the request info. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid +from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix + +from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import ( + merge_runtime_info, + require_request_id, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + +_LM_ARCHITECTURE = "NemotronHForCausalLM" + + +def validate_timeline_fits_max_model_len( + logical_prompt_token_len: int, + acoustic_frame_count: int, + max_model_len: int, +) -> None: + """Reject audio whose frame-locked timeline exceeds the stage limit. + + The vLLM request occupies ``logical_prompt_token_len + 1`` prefill + positions (the +1 is the acoustic-frame-0 placeholder) plus + ``acoustic_frame_count`` sampled tokens, so the binding sequence size is + one MORE than NeMo's logical timeline (``prompt + frames``). + """ + logical_total_len = logical_prompt_token_len + acoustic_frame_count + vllm_prefill_len = logical_prompt_token_len + 1 + vllm_total_len = vllm_prefill_len + acoustic_frame_count + if vllm_total_len > max_model_len: + raise ValueError( + "NemotronVoiceChat input exceeds the thinker stage's max_model_len: " + f"vllm_prefill_len={vllm_prefill_len} + " + f"acoustic_frame_count={acoustic_frame_count} = " + f"vllm_total_len={vllm_total_len} > max_model_len={max_model_len} " + f"(logical_prompt_token_len={logical_prompt_token_len}, " + f"logical_total_len={logical_total_len}). Use shorter audio or raise the " + "stage's max_model_len in the deploy yaml." + ) + + +def compute_acoustic_frame_count(nemo_stt_cfg: dict[str, Any], num_samples: int) -> int: + """Exact 12.5 Hz frame count for ``num_samples`` of 16 kHz audio. + + Replicates the vendored mel featurizer length arithmetic followed by the + dw-striding conv subsampling length recurrence (NOT ``duration / 0.08``). + The thinker validates this against the real perception output at prefill, + and the golden shape test pins it against the NeMo reference. + """ + from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.subsampling import ( + calc_length, + ) + + perception = nemo_stt_cfg.get("perception") or {} + pre = perception.get("preprocessor") or {} + encoder = perception.get("encoder") or {} + sample_rate = int(pre.get("sample_rate", 16000)) + window_stride = float(pre.get("window_stride", 0.01)) + window_size = float(pre.get("window_size", 0.025)) + n_fft = int(pre.get("n_fft") or 2 ** int(np.ceil(np.log2(window_size * sample_rate)))) + hop = int(window_stride * sample_rate) + pad_amount = n_fft // 2 * 2 # torch.stft center padding on both sides + mel_len = (num_samples + pad_amount - n_fft) // hop + 1 + subsampling_factor = int(encoder.get("subsampling", "dw_striding") and encoder.get("subsampling_factor", 8)) + conv_kernel = int(encoder.get("subsampling_conv_kernel_size", 3) or 3) + stride = 2 + repeat = int(np.log2(subsampling_factor)) + # Mirror the vendored ConvSubsampling padding exactly: causal downsampling + # (this checkpoint sets encoder.causal_downsampling=true) pads + # left=kernel-1 plus right=stride-1, i.e. one more column than the + # non-causal symmetric 2*(kernel//2) — which shifts the output length by + # one for roughly half of all input durations. + if bool(encoder.get("causal_downsampling", False)): + all_paddings = (conv_kernel - 1) + (stride - 1) + else: + all_paddings = 2 * (conv_kernel // 2) + out = calc_length( + torch.tensor([mel_len]), + all_paddings=all_paddings, + kernel_size=conv_kernel, + stride=stride, + ceil_mode=False, + repeat_num=repeat, + ) + return int(out[0]) + + +class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, IsHybrid): + """vLLM-native thinker (perception + fusion + NemotronH + dual heads). + + Declares the hybrid-Mamba interfaces (``IsHybrid``/``HasInnerState``) so + vLLM configures the mamba state cache for the wrapped NemotronH backbone; + the state shape/dtype classmethods mirror ``NemotronHForCausalLM`` but read + the NemotronH sub-config (``get_text_config()``) instead of the top-level + checkpoint config. + """ + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig) -> tuple[torch.dtype, torch.dtype]: + from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM + + return NemotronHForCausalLM.get_mamba_state_dtype_from_config(vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig) -> tuple[tuple[int, int], tuple[int, int, int]]: + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator + + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config.get_text_config() + intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim + return MambaStateShapeCalculator.mamba2_state_shape( + intermediate_size=intermediate_size, + tp_world_size=parallel_config.tensor_parallel_size, + n_groups=hf_config.n_groups, + num_heads=hf_config.mamba_num_heads, + head_dim=hf_config.mamba_head_dim, + state_size=hf_config.ssm_state_size, + conv_kernel=hf_config.conv_kernel, + num_spec=vllm_config.num_speculative_tokens, + ) + + @classmethod + def get_mamba_state_copy_func(cls): + from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM + + return NemotronHForCausalLM.get_mamba_state_copy_func() + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + self.config = config + stt_cfg: dict[str, Any] = dict(getattr(config, "stt_cfg", {}) or {}) + if not stt_cfg: + raise ValueError("NemotronVoiceChat checkpoint config lacks 'model.stt.model'; cannot build the thinker.") + self.stt_cfg = stt_cfg + text_cfg = config.get_text_config() + self._hidden = int(getattr(text_cfg, "hidden_size", 4480)) + self._vocab = int(getattr(text_cfg, "vocab_size", 131072)) + self._dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16) + + # NemotronH backbone (+ lm_head) on vLLM native layers. + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=text_cfg, + prefix=maybe_prefix(prefix, "language_model"), + architectures=[_LM_ARCHITECTURE], + ) + + # Shared channel embedding + separate function head (both distinct + # checkpoints subtrees from the backbone). + self.embed_tokens = nn.Embedding(self._vocab, self._hidden) + self.function_head = nn.Linear(self._hidden, self._vocab, bias=False) + + # Vendored perception (mel + FastConformer + projection). + from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.perception import ( + AudioPerceptionModule, + ) + + perception_cfg = stt_cfg.get("perception") + if not perception_cfg: + raise ValueError("'model.stt.model.perception' missing from the checkpoint config.") + from omegaconf import DictConfig + + self.perception = AudioPerceptionModule(DictConfig(perception_cfg)) + + # AddFusion weights (the checkpoint uses fuse_method="add" defaults; + # the same duplex_*_weight keys double as the fusion weights). + use_function_head = bool(stt_cfg.get("use_function_head", True)) + self._w_text = float(stt_cfg.get("duplex_text_channel_weight", 1.0)) + self._w_audio = float(stt_cfg.get("duplex_user_channel_weight", 1.0)) + self._w_func = float(stt_cfg.get("duplex_function_channel_weight", 1.0)) if use_function_head else 0.0 + self._use_function_head = use_function_head + fuse_method = stt_cfg.get("fuse_method", "add") or "add" + if fuse_method != "add": + raise NotImplementedError( + f"NemotronVoiceChat thinker only implements fuse_method='add' " + f"(this checkpoint uses it); got {fuse_method!r}." + ) + if bool(stt_cfg.get("predict_user_text", False)): + raise NotImplementedError( + "NemotronVoiceChat thinker does not implement the user-ASR channel " + "(predict_user_text=True); this checkpoint has it disabled." + ) + + self._special_ids: dict[str, int] | None = None + + # Omni AR runner contract. + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True # per-step greedy function-channel token + self.requires_full_prefix_cached_hidden_states = False + # The per-step function token is a scalar; no GPU-resident buffers needed. + self.gpu_resident_buffer_keys: set[tuple[str, str]] = set() + + # Stateful per-request execution is only validated at batch size 1 + # (the shipped offline scope); fail fast on silent multi-request use. + max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1)) + if max_num_seqs != 1: + raise NotImplementedError( + f"NemotronVoiceChat thinker supports max_num_seqs=1 only (got {max_num_seqs}); " + "multi-request batching of the per-request fusion state is not implemented." + ) + + # request_id -> per-request state (audio frames, prefill embeds, counters). + self._sessions: dict[str, dict[str, Any]] = {} + + # ------------------------------------------------------------------ + # Special token ids (from the checkpoint's token strings). + # ------------------------------------------------------------------ + def _resolve_special_ids(self) -> dict[str, int]: + if self._special_ids is not None: + return self._special_ids + from transformers import AutoTokenizer + + ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or self.stt_cfg.get( + "pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + ) + tok = AutoTokenizer.from_pretrained(ref, trust_remote_code=False) + ids = { + "bos": tok.convert_tokens_to_ids(self.stt_cfg.get("bos_token", "")), + "eos": tok.convert_tokens_to_ids(self.stt_cfg.get("eos_token", "")), + "pad": tok.convert_tokens_to_ids(self.stt_cfg.get("pad_token", "")), + } + for name, value in ids.items(): + if value is None or value < 0: + raise ValueError(f"NemotronVoiceChat could not resolve the {name} token id from {ref!r}.") + self._special_ids = ids + return ids + + # ------------------------------------------------------------------ + # vLLM model protocol. + # ------------------------------------------------------------------ + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + # Placeholder; preprocess replaces every position with fused embeds. + return torch.zeros((input_ids.shape[0], self._hidden), device=input_ids.device, dtype=self._dtype) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor: + return self.language_model( + input_ids=None, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + # vLLM >= 0.25 dropped the sampling_metadata parameter. + return self.language_model.compute_logits(hidden_states) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs={}) + + # ------------------------------------------------------------------ + # Fusion helpers. + # ------------------------------------------------------------------ + def _fuse( + self, + text_ids: torch.Tensor, + timeline_embeds: torch.Tensor, + func_ids: torch.Tensor, + ) -> torch.Tensor: + """AddFusion over the three channels; shapes [T] ids, [T, H] timeline.""" + text_e = self.embed_tokens(text_ids) + fused = text_e * self._w_text + timeline_embeds.to(text_e.dtype) * self._w_audio + if self._use_function_head: + fused = fused + self.embed_tokens(func_ids) * self._w_func + return fused + + # ------------------------------------------------------------------ + # Per-request session. + # ------------------------------------------------------------------ + def _load_audio(self, info: dict[str, Any], device: torch.device) -> torch.Tensor: + audio = info.get("nvc_audio") + if audio is None: + raise ValueError( + "NemotronVoiceChat thinker request is missing its audio: pass " + "additional_information={'nvc_audio': <16 kHz mono float waveform>, " + "'nvc_sr': 16000, 'nvc_prompt_token_ids': [...]}." + ) + wav = torch.as_tensor(np.asarray(audio), dtype=torch.float32, device=device).reshape(-1) + if wav.numel() == 0: + raise ValueError("NemotronVoiceChat thinker received empty audio input.") + sr = int(info.get("nvc_sr", 16000)) + expected_sr = int(getattr(self.config, "source_sample_rate", 16000)) + if sr != expected_sr: + raise ValueError( + f"NemotronVoiceChat thinker expects {expected_sr} Hz audio, got {sr} Hz — " + "resample client-side (the offline example does this)." + ) + return wav + + def _init_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]: + pad_id = self._resolve_special_ids()["pad"] + wav = self._load_audio(info, device) + + prompt_ids = info.get("nvc_prompt_token_ids") + if prompt_ids is None: + raise ValueError( + "NemotronVoiceChat thinker request is missing 'nvc_prompt_token_ids' " + "(the system prompt token ids, [bos] + tokens + [eos])." + ) + prompt_ids = torch.as_tensor(list(prompt_ids), dtype=torch.long, device=device) + if prompt_ids.numel() == 0: + raise ValueError("NemotronVoiceChat thinker requires a non-empty system prompt.") + + # Perception: mel + FastConformer + projection -> [N, hidden] frames. + with torch.inference_mode(): + encoded, encoded_len = self.perception( + input_signal=wav.unsqueeze(0).to(next(self.perception.parameters()).dtype), + input_signal_length=torch.tensor([wav.numel()], device=device), + ) + n_frames = int(encoded_len.reshape(-1)[0]) + audio_embeds = encoded[0, :n_frames].to(self._dtype) # [N, H] + expected = info.get("nvc_expected_frames") + if expected is not None and int(expected) != n_frames: + raise ValueError( + f"NemotronVoiceChat frame-count mismatch: the request was sized for " + f"{int(expected)} acoustic frames but perception produced {n_frames}. " + "The example must size max_tokens with compute_acoustic_frame_count() (never min_tokens)." + ) + + p = int(prompt_ids.numel()) + validate_timeline_fits_max_model_len( + logical_prompt_token_len=p, + acoustic_frame_count=n_frames, + max_model_len=int(self.vllm_config.model_config.max_model_len), + ) + + # Fused prefill embeds for positions 0..P (P prompt slots + frame 0). + # All prompt positions use the PAD text channel (NeMo's + # _get_bos_embedding returns the PAD embedding) and PAD function channel. + timeline = torch.cat([self.embed_tokens(prompt_ids).to(self._dtype), audio_embeds[0:1]], dim=0) # [P+1, H] + pad_ids = torch.full((p + 1,), pad_id, dtype=torch.long, device=device) + prefill = self._fuse(pad_ids, timeline, pad_ids) # [P+1, H] + + session = { + "audio_embeds": audio_embeds, + "prefill_embeds": prefill, + "prompt_len": p, + "n_frames": n_frames, + "func_token": int(pad_id), + "decode_step": 0, + } + self._sessions[request_id] = session + return session + + # ------------------------------------------------------------------ + # Omni AR hooks. + # ------------------------------------------------------------------ + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + info = merge_runtime_info(info_dict) + device = input_ids.device + span = int(input_ids.shape[0]) + is_prefill_raw = info.get("_omni_is_prefill") + is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1 + request_id = require_request_id(info, "thinker") + + if is_prefill or request_id not in self._sessions: + session = self._sessions.get(request_id) or self._init_session(request_id, info, device) + offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0)) + prefill: torch.Tensor = session["prefill_embeds"] + if offset + span > prefill.shape[0]: + raise ValueError( + f"NemotronVoiceChat thinker prefill span [{offset}, {offset + span}) exceeds " + f"the fused prompt length {prefill.shape[0]}; the request prompt must be " + "the system prompt tokens plus exactly one placeholder position." + ) + return input_ids, prefill[offset : offset + span].to(self._dtype), {} + + session = self._sessions[request_id] + pad_id = self._resolve_special_ids()["pad"] + + # Function channel: greedy argmax over the PREVIOUS position's hidden, + # computed by postprocess right after each forward. At the first decode + # step this is the prefill's last position (acoustic frame 0), matching + # NeMo's gen_function[prompt_len]. The cumulative debug timeline + # (nvc_function_tokens) is owned by postprocess, which also covers the + # final acoustic frame (no later preprocess exists for it). + if self._use_function_head: + prev_func = info.get("nvc_prev_function_token") + if prev_func is not None: + token = int(prev_func.reshape(-1)[0]) if isinstance(prev_func, torch.Tensor) else int(prev_func) + session["func_token"] = token + + # Position-addressed frame index: decode of vLLM position P+k consumes + # acoustic frame k (frame 0 was the last prefill position). Deriving it + # from _omni_num_computed_tokens keeps preprocess idempotent — if the + # runner ever re-enters the same decode position (preemption or + # recompute), the audio pointer stays aligned with the text channel + # instead of silently advancing. Falls back to a call counter only if + # the runner did not inject the key. + computed = info.get("_omni_num_computed_tokens") + if computed is not None: + frame_idx = int(computed) - int(session["prompt_len"]) + else: + frame_idx = int(session["decode_step"]) + 1 + session["decode_step"] = frame_idx + n_frames = int(session["n_frames"]) + if frame_idx < 1: + raise ValueError( + f"NemotronVoiceChat thinker derived decode frame index {frame_idx} " + f"(computed_tokens={computed}, prompt_len={session['prompt_len']}); " + "decode must start at acoustic frame 1." + ) + if frame_idx >= n_frames: + raise ValueError( + f"NemotronVoiceChat thinker was asked for decode step {frame_idx} but only " + f"{n_frames} acoustic frames exist; max_tokens must equal the frame count." + ) + text_prev = input_ids.reshape(-1)[:1].to(torch.long) + func_prev = torch.full_like(text_prev, session["func_token"] if self._use_function_head else pad_id) + frame = session["audio_embeds"][frame_idx : frame_idx + 1] + fused = self._fuse(text_prev, frame, func_prev) # [1, H] + + info_update: dict[str, Any] = { + "meta": {"nvc_frame": frame_idx}, + "nvc_text_pad_id": pad_id, + } + return input_ids, fused, info_update + + def postprocess(self, hidden_states: torch.Tensor, **kwargs: Any) -> dict[str, Any]: + """Greedy function-channel token from this step's last hidden. + + Runs after EVERY forward (prefill and each decode step), so the final + acoustic frame's token is covered too — NeMo produces gen_function for + every non-prompt position. ``nvc_prev_function_token`` is load-bearing + (it feeds the next step's fusion at weight 2.0). The cumulative debug + timeline (``nvc_function_tokens``, one entry per acoustic frame) costs + a host sync plus a growing concat per step, so it is only accumulated + when ``NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE=1``. + On multi-chunk prefill, only the final chunk (whose forward covers the + last prefill position, acoustic frame 0) contributes a token. + """ + if hidden_states is None or hidden_states.numel() == 0 or not self._use_function_head: + return {} + if kwargs.get("_omni_is_prefill"): + computed = int(kwargs.get("_omni_num_computed_tokens", 0) or 0) + prompt_len = int(kwargs.get("_omni_prompt_len", 0) or 0) + if prompt_len and computed + int(hidden_states.shape[0]) < prompt_len: + return {} # intermediate prefill chunk: frame 0 not reached yet + with torch.inference_mode(): + func_logits = self.function_head(hidden_states[-1, :].reshape(1, -1).to(self._dtype)) + token = func_logits.argmax(dim=-1).reshape(()).detach() + update: dict[str, Any] = {"nvc_prev_function_token": token} + if os.environ.get("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "0") == "1": + existing = kwargs.get("nvc_function_tokens") + if isinstance(existing, torch.Tensor) and existing.numel() > 0: + update["nvc_function_tokens"] = torch.cat([existing.reshape(-1).cpu(), token.reshape(1).cpu()]) + else: + update["nvc_function_tokens"] = token.reshape(1).cpu() + return update + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + for request_id in finished_req_ids: + self._sessions.pop(str(request_id), None) + + # ------------------------------------------------------------------ + # Weights. + # ------------------------------------------------------------------ + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Route the unified checkpoint into the thinker's components. + + * ``stt_model.llm.*`` -> vLLM NemotronH (renamed to ``backbone.*``; the + separate ``stt_model.embed_tokens`` doubles as ``backbone.embeddings`` + since the checkpoint keeps a single shared text embedding). + * ``stt_model.lm_head.weight`` -> the backbone's ``lm_head``. + * ``stt_model.embed_tokens/function_head`` -> the fusion embedding / + function head. + * ``stt_model.perception.*`` -> the vendored Conformer. + * everything else (``rnnt_*``, ``tts_model.*``) is skipped lazily. + """ + loaded: set[str] = set() + lm_weights: list[tuple[str, torch.Tensor]] = [] + perception_state: dict[str, torch.Tensor] = {} + embed_weight: torch.Tensor | None = None + + for name, tensor in weights: + if name.startswith("stt_model.llm."): + lm_weights.append(("backbone." + name.removeprefix("stt_model.llm."), tensor)) + elif name == "stt_model.lm_head.weight": + lm_weights.append(("lm_head.weight", tensor)) + elif name == "stt_model.embed_tokens.weight": + embed_weight = tensor + elif name == "stt_model.function_head.weight": + with torch.no_grad(): + self.function_head.weight.copy_(tensor) + loaded.add("function_head.weight") + elif name.startswith("stt_model.perception."): + perception_state[name.removeprefix("stt_model.perception.")] = tensor + + if embed_weight is None: + raise ValueError("Checkpoint is missing 'stt_model.embed_tokens.weight'.") + if not lm_weights: + raise ValueError("Checkpoint is missing the 'stt_model.llm.*' backbone subtree.") + if not perception_state: + raise ValueError("Checkpoint is missing the 'stt_model.perception.*' subtree.") + if "function_head.weight" not in loaded and self._use_function_head: + raise ValueError("Checkpoint is missing 'stt_model.function_head.weight'.") + + with torch.no_grad(): + self.embed_tokens.weight.copy_(embed_weight) + loaded.add("embed_tokens.weight") + + # The vLLM NemotronH model expects an input embedding; the checkpoint's + # shared embed_tokens fills that slot (inputs_embeds bypass it at runtime). + lm_weights.append(("backbone.embeddings.weight", embed_weight)) + lm_loaded = self.language_model.load_weights(iter(lm_weights)) + loaded |= {f"language_model.{n}" for n in lm_loaded} + + missing, unexpected = self.perception.load_state_dict(perception_state, strict=False) + if missing: + raise ValueError(f"NemotronVoiceChat perception is missing weights: {sorted(missing)[:10]} ...") + if unexpected: + logger.warning( + "NemotronVoiceChat perception ignoring %d unexpected tensors (e.g. %s)", + len(unexpected), + sorted(unexpected)[:5], + ) + device = self.vllm_config.device_config.device + self.perception.to(device=device, dtype=self._dtype).eval() + self.embed_tokens.to(device=device, dtype=self._dtype) + self.function_head.to(device=device, dtype=self._dtype) + loaded |= {f"perception.{n}" for n in perception_state} + logger.info( + "NemotronVoiceChat thinker loaded: %d backbone tensors, %d perception tensors, " + "fusion weights (text=%.1f, audio=%.1f, function=%.1f)", + len(lm_weights), + len(perception_state), + self._w_text, + self._w_audio, + self._w_func, + ) + return loaded diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py new file mode 100644 index 000000000..4f4834740 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NemotronVoiceChat pipeline: thinker (speech -> frame-locked text) -> talker +(text timeline -> 31-quantizer RVQ code stacks) -> code2wav (codes -> 22.05 kHz PCM). + +``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` is an 11B full-duplex speech-to-speech +model. Offline it is served as a qwen3-omni-style 3-stage pipeline: + +* Stage 0 (``thinker``, LLM_AR): vendored FastConformer perception + additive + channel fusion feeding the NemotronH backbone on vLLM; emits the frame-locked + agent text channel (also the final text output) and carries the full frame + timeline + metadata to the talker as a latent payload. +* Stage 1 (``talker``, LLM_AR): the vendored Gemma3-based EAR-TTS model stepped + per frame (classifier-free guidance + MoG sampling run inside the vendored + forward); emits one 31-code stack per frame as latents. +* Stage 2 (``code2wav``, LLM_GENERATION): the vendored RVQ-VAE codec decoder + (fp32) producing the final 22.05 kHz waveform. + +The STT->TTS coupling is one-directional and token-ids-only (verified against +the NeMo reference), which is what makes the offline cascade exact. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.nemotron_voicechat" + +NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig( + model_type="nemotron_voicechat", + model_arch="NemotronVoiceChatThinkerForConditionalGeneration", + # Named after the alias key so the deploy-yaml stem substring-matches the + # checkpoint directory name (NVIDIA-NemotronLabs-VoiceChat-11B) during + # bare-path auto-detection. + default_deploy_config_name="nemotron_labs_voicechat.yaml", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + final_output=True, + final_output_type="text", + engine_output_type="latent", + # thinker -> talker sync mode is token-path only: the frame-locked + # text timeline rides the talker's prompt_token_ids (built by + # thinker2talker_token_only on the talker stage); no connector + # payload and no full-payload wait for stage 1. With + # ``async_chunk: true`` in the deploy yaml, the timeline instead + # streams incrementally over the connector (cumulative ids.all per + # thinker step) so the talker starts before the thinker finishes. + async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="talker", + execution_type=StageExecutionType.LLM_AR, + model_arch="NemotronVoiceChatTalkerForConditionalGeneration", + hf_config_name="talker_config", + input_sources=(0,), + engine_output_type="latent", + custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", + async_chunk_process_next_stage_input_func=f"{_PROC}.talker2code2wav_async_chunk", + sync_process_input_func=f"{_PROC}.thinker2talker_token_only", + sampling_constraints={"detokenize": False}, + ), + StagePipelineConfig( + stage_id=2, + model_stage="code2wav", + execution_type=StageExecutionType.LLM_GENERATION, + model_arch="NemotronVoiceChatCode2Wav", + hf_config_name="code2wav_config", + input_sources=(1,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + sync_process_input_func=f"{_PROC}.talker2code2wav_token_only", + sampling_constraints={"detokenize": False}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py new file mode 100644 index 000000000..7f5a897d2 --- /dev/null +++ b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Runtime-info helpers shared by the thinker and talker AR stages.""" + +from __future__ import annotations + +from typing import Any + + +def merge_runtime_info(info_dict: dict[str, Any]) -> dict[str, Any]: + """Flatten ``additional_information`` into the info dict (top-level wins).""" + additional = info_dict.get("additional_information") + if isinstance(additional, dict): + merged = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional.items(): + merged.setdefault(k, v) + return merged + return info_dict + + +def require_request_id(info: dict[str, Any], stage: str) -> str: + """Resolve the request id from the runtime info (or its meta), or fail.""" + request_id = info.get("request_id") + if request_id is None: + meta = info.get("meta") + if isinstance(meta, dict): + request_id = meta.get("request_id") + if request_id is None: + raise RuntimeError(f"NemotronVoiceChat {stage} requires a request_id in the runtime info.") + return str(request_id) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 3c79ba34c..498f9eb85 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -176,6 +176,23 @@ _OMNI_MODELS = { "personaplex_code2wav", "PersonaPlexCode2Wav", ), + ## nemotron_voicechat (NVIDIA-NemotronLabs-VoiceChat-11B, full-duplex S2S; + ## 3-stage offline speech->speech pipeline: thinker/talker/code2wav) + "NemotronVoiceChatThinkerForConditionalGeneration": ( + "nemotron_voicechat", + "nemotron_voicechat_thinker", + "NemotronVoiceChatThinkerForConditionalGeneration", + ), + "NemotronVoiceChatTalkerForConditionalGeneration": ( + "nemotron_voicechat", + "nemotron_voicechat_talker", + "NemotronVoiceChatTalkerForConditionalGeneration", + ), + "NemotronVoiceChatCode2Wav": ( + "nemotron_voicechat", + "nemotron_voicechat_code2wav", + "NemotronVoiceChatCode2Wav", + ), ## higgs-audio v2 "HiggsAudioV2ForConditionalGeneration": ( "higgs_audio_v2", diff --git a/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py new file mode 100644 index 000000000..48de5a380 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage transfer processors for the NemotronVoiceChat 3-stage pipeline. + +thinker (0) -> talker (1): token path only. The thinker's frame-locked text +timeline becomes the talker's ``prompt_token_ids``: +``[text_pad_id] * logical_prompt_token_len + generated frame tokens``. The +talker must process the PAD prompt region too (its KV state at the first +acoustic frame depends on it — NeMo trims prompt-region codes only after the +loop), so the prompt slots are reconstructed here, and the timeline metadata +rides ``additional_information``. + +talker (1) -> code2wav (2): PersonaPlex-style full payload. The talker +accumulates one 31-quantizer code stack per frame under ``("codes","audio")``; +the producer ships the prompt-trimmed ``[frames, 31]`` stack once the request +finishes, and the sync placeholder builder sizes stage-2's dummy prompt. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm_omni.data_entry_keys import ( + CodesStruct, + IdsStruct, + MetaStruct, + OmniPayloadStruct, +) + +# Full-sequence tensors are re-shipped whole; the full-payload accumulator must +# replace (not concatenate) them if the producer fires more than once. The AR +# runner flattens nested payloads before accumulation, so the effective key is +# the flattened "codes.audio" (the unflattened root is kept for robustness). +_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes", "codes.audio", "meta", "meta.nvc_logical_prompt_len"}) + +# in the Nemotron-Nano-9B-v2 vocab; the checkpoint config's +# stt pad_token. The thinker also reports it in its latent metadata, which +# takes precedence when present. +_DEFAULT_TEXT_PAD_ID = 12 + + +def _info_get(info: Any, key: str) -> Any: + if isinstance(info, dict): + if key in info: + return info[key] + meta = info.get("meta") + if isinstance(meta, dict): + return meta.get(key) + return None + + +def thinker2talker_token_only( + source_outputs: list, + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list: + """Build the talker's engine inputs from the thinker's finished outputs. + + ``prompt_token_ids`` = the full frame-aligned text timeline including the + PAD prompt region. The thinker's vLLM prompt covers + ``logical_prompt_token_len + 1`` positions (the +1 is acoustic frame 0), so + ``logical_prompt_token_len = len(thinker prompt) - 1`` and the generated + tokens are exactly the acoustic-frame text channel. + """ + from vllm_omni.inputs.data import OmniTokensPrompt + + del prompt, _requires_multimodal_data + inputs: list = [] + for thinker_output in source_outputs: + if not getattr(thinker_output, "finished", False): + continue + output = thinker_output.outputs[0] + generated = list(getattr(output, "cumulative_token_ids", None) or getattr(output, "token_ids", None) or []) + prompt_ids = list(getattr(thinker_output, "prompt_token_ids", None) or []) + logical_prompt_len = max(len(prompt_ids) - 1, 0) + pad_id = _DEFAULT_TEXT_PAD_ID + info = getattr(thinker_output, "additional_information", None) + reported_pad = _info_get(info, "nvc_text_pad_id") + if reported_pad is not None: + pad_id = int(reported_pad) + timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated] + inputs.append( + OmniTokensPrompt( + # The vLLM-side prompt is one placeholder (timeline position 0; + # the NeMo TTS loop starts at t=1). Each decode step consumes + # one timeline position from additional_information, so the + # talker needs len(timeline) - 1 decode steps. + prompt_token_ids=[0], + additional_information={ + "nvc_text_timeline": timeline, + "nvc_logical_prompt_len": logical_prompt_len, + "nvc_frame_count": len(generated), + "nvc_text_pad_id": pad_id, + }, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + return inputs + + +def thinker2talker_async_chunk( + transfer_manager: Any, + multimodal_output: Any = None, + request: Any = None, + is_finished: bool = False, + **kwargs: Any, +) -> OmniPayloadStruct | None: + """Streaming producer: ship the CUMULATIVE text timeline each thinker step. + + The AR receive path REPLACES the downstream request's + additional_information with each chunk payload, so every chunk must carry + the whole timeline so far: ``[pad] * logical_prompt_len + sampled tokens`` + under ``ids.all`` (the talker adopts any longer timeline it sees). The + timeline is a few hundred ints, so cumulative re-ship is cheap and keeps + the payload self-contained. ``meta.num_processed_tokens`` carries the + logical prompt length for the code2wav producer's prompt-region trim. + """ + del multimodal_output, kwargs + request_id = request.external_req_id + prompt_ids = list(getattr(request, "prompt_token_ids", None) or []) + # vLLM prompt = logical prompt + 1 placeholder (acoustic frame 0). + logical_prompt_len = max(len(prompt_ids) - 1, 0) + generated = list(getattr(request, "output_token_ids", None) or []) + pad_id = _DEFAULT_TEXT_PAD_ID + reported_pad = _info_get(getattr(request, "additional_information", None), "nvc_text_pad_id") + if reported_pad is not None: + pad_id = int(reported_pad) + timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated] + + # Skip no-progress wakeups (same length as the last emitted chunk). + state = transfer_manager.request_payload.get(request_id) + last_len = int(state.get("nvc_timeline_len", -1)) if isinstance(state, dict) else -1 + if len(timeline) <= last_len and not is_finished: + return None + transfer_manager.request_payload[request_id] = {"nvc_timeline_len": len(timeline)} + + return OmniPayloadStruct( + ids=IdsStruct(all=timeline), + meta=MetaStruct( + finished=torch.tensor(bool(is_finished), dtype=torch.bool), + num_processed_tokens=logical_prompt_len, + # The talker's vLLM prompt stays a single placeholder token. + next_stage_prompt_len=1, + ), + ) + + +def talker2code2wav_async_chunk( + transfer_manager: Any, + multimodal_output: Any = None, + request: Any = None, + is_finished: bool = False, + **kwargs: Any, +) -> OmniPayloadStruct | None: + """Streaming producer: ship prompt-trimmed CUMULATIVE code stacks in chunks. + + Follows the NeMo incremental-decode recipe (``decode_one_audio_step``): + each chunk carries the FULL trimmed code history and + ``meta.left_context_size`` = frames the codec already emitted, so the + code2wav stage re-decodes the whole prefix and slices off only the new + samples — no seams, no crossfade, and the concatenated stream equals the + prefix-decode of the final stack. Chunk cadence comes from the connector + config key ``codec_chunk_frames`` (default 13 frames ~= 1 s of audio). + """ + del kwargs + request_id = request.external_req_id + + codes = None + if isinstance(multimodal_output, dict): + nested = multimodal_output.get("codes") + codes = nested.get("audio") if isinstance(nested, dict) else None + if codes is None: + codes = multimodal_output.get("codes.audio") + if not isinstance(codes, torch.Tensor) or codes.numel() == 0: + if not is_finished: + return None + return _empty_finished_payload() + if codes.ndim == 1: + codes = codes.reshape(1, -1) + + # Prompt-region trim (rows are NeMo timeline steps t=1..; keep t >= P). + # Read the prompt length from the PER-STEP model payload first: the + # request-level additional_information is overwritten both by arriving + # thinker chunks and by talker info updates, so it races. Cache the first + # sighting in the transfer-manager state as a fallback. + state = transfer_manager.request_payload.get(request_id) + prompt_len = _info_get(multimodal_output if isinstance(multimodal_output, dict) else None, "nvc_logical_prompt_len") + if prompt_len is None and isinstance(multimodal_output, dict): + prompt_len = multimodal_output.get("meta.nvc_logical_prompt_len") + if prompt_len is None: + info = getattr(request, "additional_information", None) + prompt_len = _info_get(info, "nvc_logical_prompt_len") + if prompt_len is None: + prompt_len = _info_get(info, "num_processed_tokens") + if prompt_len is None and isinstance(state, dict): + prompt_len = state.get("nvc_prompt_len") + if prompt_len is None: + raise ValueError( + "NemotronVoiceChat talker request is missing its logical prompt length " + "('nvc_logical_prompt_len' or async meta.num_processed_tokens); cannot " + "trim the prompt-region codes for streaming code2wav." + ) + trimmed = codes[max(int(prompt_len) - 1, 0) :] + + frames_sent = int(state.get("nvc_frames_sent", 0)) if isinstance(state, dict) else 0 + new_frames = int(trimmed.shape[0]) - frames_sent + if new_frames <= 0 and not is_finished: + return None + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + chunk_frames = int(cfg.get("codec_chunk_frames", 13)) + if new_frames < chunk_frames and not is_finished: + return None + if new_frames <= 0 and is_finished and frames_sent > 0: + # Everything already shipped; emit a terminal empty chunk so the + # code2wav request observes meta.finished. + return _empty_finished_payload() + + transfer_manager.request_payload[request_id] = { + "nvc_frames_sent": int(trimmed.shape[0]), + "nvc_prompt_len": int(prompt_len), + } + return OmniPayloadStruct( + codes=CodesStruct(audio=trimmed.detach().to(dtype=torch.long, device="cpu")), + meta=MetaStruct( + left_context_size=frames_sent, + finished=torch.tensor(bool(is_finished), dtype=torch.bool), + ), + ) + + +def _empty_finished_payload() -> OmniPayloadStruct: + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + + +def talker2code2wav_full_payload( + transfer_manager: Any = None, + pooling_output: Any = None, + request: Any = None, + is_finished: bool = False, + **kwargs: Any, +) -> OmniPayloadStruct | None: + """Producer: ship the talker's prompt-trimmed ``[frames, 31]`` code stack. + + Codes accumulate per frame under ``("codes","audio")`` in the request's + additional_information; the prompt-region rows were never emitted by the + talker (it only stores codes for generated positions), so no extra trim is + needed here. Ships once, at request finish, with replace semantics. + """ + # NOTE: no is_finished gating — in sync full-payload mode the producer only + # runs at flush time (request already finished), and the flush-time request + # object may still report is_finished()=False. Returning a meta-only + # "pending" struct here would be SENT as the real payload and starve the + # code2wav stage of its codes. + del transfer_manager, is_finished + + def _codes_from(src: Any) -> torch.Tensor | None: + if not isinstance(src, dict): + return None + nested = src.get("codes") + audio = nested.get("audio") if isinstance(nested, dict) else None + return audio if audio is not None else src.get("codes.audio") + + audio = None + for source in ( + getattr(request, "additional_information", None), + getattr(request, "additional_information_cpu", None), + pooling_output, + kwargs.get("multimodal_output"), + ): + audio = _codes_from(source) + if audio is not None: + break + if not isinstance(audio, torch.Tensor) or audio.numel() == 0: + import logging + + logging.getLogger(__name__).error( + "nemotron_voicechat talker full-payload producer found no codes: " + "pooling_output keys=%s, request.additional_information keys=%s", + sorted(pooling_output.keys()) if isinstance(pooling_output, dict) else type(pooling_output).__name__, + sorted(getattr(request, "additional_information", None) or {}) + if isinstance(getattr(request, "additional_information", None), dict) + else type(getattr(request, "additional_information", None)).__name__, + ) + return _empty_finished_payload() + if audio.ndim == 1: + audio = audio.reshape(1, -1) + # Rows correspond to NeMo timeline steps t = 1..T-1. Trim the prompt + # region: NeMo keeps gen_codes rows t = P..T-1 (P = logical prompt len), + # i.e. row indices P-1 onward here. P >= 1 always (a system prompt is + # required), so this never drops acoustic frames. + prompt_len = _info_get(getattr(request, "additional_information", None), "nvc_logical_prompt_len") + if prompt_len is None: + prompt_len = _info_get(getattr(request, "additional_information_cpu", None), "nvc_logical_prompt_len") + if prompt_len is None: + # Explicit failure, never a silent trim=0 (untrimmed prompt-region codes + # would ship to the codec). The connector layer logs producer errors. + raise ValueError( + "NemotronVoiceChat talker request is missing 'nvc_logical_prompt_len' in its " + "additional_information; cannot trim the prompt-region codes for code2wav." + ) + audio = audio[max(int(prompt_len) - 1, 0) :] + if audio.shape[0] == 0: + return _empty_finished_payload() + # 2-D [frames, Q] payload: the chunk/full-payload transfer adapters route + # >=2-D tensors through the payload channel (1-D would be coerced into + # prompt token ids and dropped from the info merge). + return OmniPayloadStruct( + codes=CodesStruct(audio=audio.detach().to(dtype=torch.long, device="cpu")), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + + +def talker2code2wav_token_only( + source_outputs: list, + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list: + """Sync placeholder builder for stage 2 (codes arrive via the connector).""" + from vllm_omni.inputs.data import OmniTokensPrompt + + del prompt, _requires_multimodal_data + inputs: list = [] + for talker_output in source_outputs: + if not getattr(talker_output, "finished", False): + continue + output = talker_output.outputs[0] + token_ids = getattr(output, "cumulative_token_ids", None) or getattr(output, "token_ids", None) or [] + n_frames = len(token_ids) + inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * max(n_frames, 1), + additional_information=None, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + return inputs + + +__all__ = [ + "thinker2talker_token_only", + "thinker2talker_async_chunk", + "talker2code2wav_full_payload", + "talker2code2wav_async_chunk", + "talker2code2wav_token_only", +] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 0f65489b7..b3c388e3a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -339,6 +339,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): "NemotronHForCausalLM", "DyninOmniForConditionalGeneration", "IndexTTS2TalkerForConditionalGeneration", + # nemotron_voicechat: the talker (stage 1) is the full-payload + # producer for code2wav (stage 2); the thinker (stage 0) only + # produces over the connector in async-chunk (streaming) mode. + "NemotronVoiceChatTalkerForConditionalGeneration", + "NemotronVoiceChatThinkerForConditionalGeneration", } # The stage-level ``model_arch`` override may be blank so the class # resolves from the checkpoint's own ``architectures`` (e.g. the Audex diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 766e84c45..7cf26d005 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -69,6 +69,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin "CosyVoice3Model", "DyninOmniForConditionalGeneration", "IndexTTS2S2MelDecoder", + # nemotron_voicechat: code2wav (stage 2) consumes the talker's + # full-payload code stacks. + "NemotronVoiceChatCode2Wav", } if ( getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS From 31384cb39cf107809e64338519de55e119361221 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:30:55 +0800 Subject: [PATCH 08/41] [CI/Build] Add nightly DockerHub publish and cleanup to release pipeline (#6048) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] [Your Email] Co-authored-by: Hongsheng Liu --- .buildkite/release/release-pipeline.yml | 24 ++- .../release/scripts/cleanup-nightly-builds.sh | 147 ++++++++++++++++++ .../release/scripts/publish-release-images.sh | 93 +++++++---- 3 files changed, 235 insertions(+), 29 deletions(-) create mode 100644 .buildkite/release/scripts/cleanup-nightly-builds.sh diff --git a/.buildkite/release/release-pipeline.yml b/.buildkite/release/release-pipeline.yml index 351e105cc..80425b928 100644 --- a/.buildkite/release/release-pipeline.yml +++ b/.buildkite/release/release-pipeline.yml @@ -40,9 +40,11 @@ steps: commands: - "bash .buildkite/release/scripts/publish-release-wheels-pypi.sh" + # Auto-skip when NIGHTLY=1 (scheduled); otherwise require manual unblock. - block: "Unblock to build release Docker images" depends_on: ~ key: block-build-release-images + if: build.env("NIGHTLY") != "1" - group: "Build release Docker images" key: "build-release-images" @@ -63,7 +65,7 @@ steps: --progress plain \ -f docker/Dockerfile.cuda . - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)" - # re-tag to default image tag and push, just in case arm64 build fails + # Fallback default tag (no arch suffix) if arm64 build fails. - "docker tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT" - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT" env: @@ -86,6 +88,26 @@ steps: env: DOCKER_BUILDKIT: "1" + # NIGHTLY=1 only: promote ECR images to DockerHub nightly tags, then prune. + - label: "Publish nightly multi-arch image to DockerHub" + depends_on: + - build-release-image-x86 + - build-release-image-arm64 + if: build.env("NIGHTLY") == "1" + key: publish-nightly-images-dockerhub + agents: + queue: small_cpu_queue_release + commands: + - "bash .buildkite/release/scripts/publish-release-images.sh nightly" + - "bash .buildkite/release/scripts/cleanup-nightly-builds.sh" + plugins: + - docker-login#v3.0.0: + username: vllmbot + password-env: DOCKERHUB_TOKEN + env: + DOCKER_BUILDKIT: "1" + DOCKERHUB_USERNAME: "vllmbot" + - block: "Publish release images to DockerHub" key: block-publish-release-images depends_on: diff --git a/.buildkite/release/scripts/cleanup-nightly-builds.sh b/.buildkite/release/scripts/cleanup-nightly-builds.sh new file mode 100644 index 000000000..7d1838d89 --- /dev/null +++ b/.buildkite/release/scripts/cleanup-nightly-builds.sh @@ -0,0 +1,147 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# +# Clean up old nightly builds from DockerHub (vllm/vllm-omni), keeping the +# newest 14 tags with the nightly- prefix. + +set -ex + +TAG_PREFIX="nightly-" +REPO="vllm/vllm-omni" +REPO_API_URL="https://hub.docker.com/v2/repositories/${REPO}/tags" + +echo "Cleaning up tags with prefix: $TAG_PREFIX in repository: $REPO" + +if [ -z "$DOCKERHUB_TOKEN" ]; then + echo "Error: DOCKERHUB_TOKEN environment variable is not set" + exit 1 +fi + +if [ -z "$DOCKERHUB_USERNAME" ]; then + echo "Error: DOCKERHUB_USERNAME environment variable is not set" + exit 1 +fi + +echo "Getting DockerHub bearer token..." +set +x +BEARER_TOKEN=$(curl -s -X POST \ + -H "Content-Type: application/json" \ + -d "{\"username\": \"$DOCKERHUB_USERNAME\", \"password\": \"$DOCKERHUB_TOKEN\"}" \ + "https://hub.docker.com/v2/users/login" | jq -r '.token') +set -x + +if [ -z "$BEARER_TOKEN" ] || [ "$BEARER_TOKEN" = "null" ]; then + echo "Error: Failed to get DockerHub bearer token" + exit 1 +fi + +get_all_tags() { + local page=1 + local all_tags="" + # Hub `name` is a partial match (not a strict prefix). Use it to narrow the + # list, then enforce startswith() client-side. Exhaust API pages; do not + # stop merely because a page has no matching tags after filtering. + local name_query + name_query=$(printf '%s' "$TAG_PREFIX" | jq -sRr @uri) + + while true; do + set +x + local response + response=$(curl -s -H "Authorization: Bearer $BEARER_TOKEN" \ + "$REPO_API_URL?page=$page&page_size=100&name=${name_query}") + set -x + + local page_count + page_count=$(echo "$response" | jq -r '.results | length') + if [ -z "$page_count" ] || [ "$page_count" = "0" ] || [ "$page_count" = "null" ]; then + break + fi + + local tags + tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" \ + '.results[] | select(.name | startswith($prefix)) | "\(.last_updated)|\(.name)"') + + if [ -n "$tags" ]; then + all_tags="$all_tags$tags"$'\n' + fi + page=$((page + 1)) + done + + echo "$all_tags" | sort -r | cut -d'|' -f2 +} + +delete_tag() { + local tag_name="$1" + echo "Deleting tag: $tag_name" + + local delete_url="https://hub.docker.com/v2/repositories/${REPO}/tags/$tag_name" + set +x + local response http_code body + response=$(curl -s -w "\n%{http_code}" -X DELETE \ + -H "Authorization: Bearer $BEARER_TOKEN" "$delete_url") + set -x + http_code=$(printf '%s\n' "$response" | tail -n1) + body=$(printf '%s\n' "$response" | sed '$d') + + if [ "$http_code" -ge 200 ] && [ "$http_code" -lt 300 ]; then + echo "Successfully deleted tag: $tag_name (HTTP ${http_code})" + return 0 + fi + + local detail="" + if echo "$body" | jq -e '.detail' > /dev/null 2>&1; then + detail=$(echo "$body" | jq -r '.detail') + elif [ -n "$body" ]; then + detail="$body" + else + detail="(empty response body)" + fi + echo "Error: Failed to delete tag $tag_name (HTTP ${http_code}): ${detail}" + return 1 +} + +echo "Fetching all tags from DockerHub..." +all_tags=$(get_all_tags) + +if [ -z "$all_tags" ]; then + echo "No tags found to clean up" + exit 0 +fi + +total_tags=$(echo "$all_tags" | wc -l) +echo "Found $total_tags tags" + +tags_to_keep=14 +tags_to_delete=$((total_tags - tags_to_keep)) + +if [ $tags_to_delete -le 0 ]; then + echo "No tags need to be deleted (only $total_tags tags found, keeping $tags_to_keep)" + exit 0 +fi + +echo "Will delete $tags_to_delete old tags, keeping the newest $tags_to_keep" + +tags_to_delete_list=$(echo "$all_tags" | tail -n +$((tags_to_keep + 1))) + +if [ -z "$tags_to_delete_list" ]; then + echo "No tags to delete" + exit 0 +fi + +echo "Deleting old tags..." +delete_failures=0 +while IFS= read -r tag; do + if [ -n "$tag" ]; then + if ! delete_tag "$tag"; then + delete_failures=$((delete_failures + 1)) + fi + sleep 1 + fi +done <<< "$tags_to_delete_list" + +if [ "$delete_failures" -gt 0 ]; then + echo "Cleanup finished with ${delete_failures} deletion failure(s)" + exit 1 +fi + +echo "Cleanup completed successfully" diff --git a/.buildkite/release/scripts/publish-release-images.sh b/.buildkite/release/scripts/publish-release-images.sh index e616604ce..56705613b 100755 --- a/.buildkite/release/scripts/publish-release-images.sh +++ b/.buildkite/release/scripts/publish-release-images.sh @@ -1,50 +1,87 @@ #!/bin/bash # SPDX-License-Identifier: Apache-2.0 # -# Publish release Docker images from ECR to DockerHub. -# Pulls per-arch images, tags with `latest` and versioned tags, pushes them, -# then creates and pushes a multi-arch manifest. +# Publish Docker images from ECR staging to DockerHub. +# +# Usage: +# publish-release-images.sh # release mode (default) +# publish-release-images.sh release +# publish-release-images.sh nightly +# +# Release tags: latest, v +# Nightly tags: nightly, nightly- set -euo pipefail -RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//') -if [ -z "${RELEASE_VERSION}" ]; then - echo "ERROR: release-version metadata not set" - exit 1 -fi +MODE="${1:-release}" -COMMIT="$BUILDKITE_COMMIT" DOCKERHUB_REPO="vllm/vllm-omni" ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-omni-release-repo" +COMMIT="$BUILDKITE_COMMIT" +PRIMARY_TAGS=() +EXTRA_MANIFEST_TAGS=() + +case "${MODE}" in + release) + RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//') + if [ -z "${RELEASE_VERSION}" ]; then + echo "ERROR: release-version metadata not set" + exit 1 + fi + PRIMARY_TAGS=("latest" "v${RELEASE_VERSION}") + ;; + nightly) + # Commit-pinned manifest reuses nightly-{x86_64,aarch64} arch images. + PRIMARY_TAGS=("nightly") + EXTRA_MANIFEST_TAGS=("nightly-${COMMIT}") + ;; + *) + echo "Usage: $0 {release|nightly}" + exit 2 + ;; +esac echo "========================================" -echo "Publishing release images v${RELEASE_VERSION}" +echo "Publishing ${MODE} images to ${DOCKERHUB_REPO}" echo " Commit: ${COMMIT}" +echo " Primary tags: ${PRIMARY_TAGS[*]}" +if [ ${#EXTRA_MANIFEST_TAGS[@]} -gt 0 ]; then + echo " Extra manifests: ${EXTRA_MANIFEST_TAGS[*]}" +fi echo "========================================" -# Login to ECR to pull staging images aws ecr-public get-login-password --region us-east-1 | \ docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 -docker pull "${ECR_REPO}:${COMMIT}-x86_64" -docker pull "${ECR_REPO}:${COMMIT}-aarch64" +SRC_X86="${ECR_REPO}:${COMMIT}-x86_64" +SRC_ARM="${ECR_REPO}:${COMMIT}-aarch64" -docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:latest-x86_64" -docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" -docker push "${DOCKERHUB_REPO}:latest-x86_64" -docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" +docker pull "${SRC_X86}" +docker pull "${SRC_ARM}" -docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:latest-aarch64" -docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64" -docker push "${DOCKERHUB_REPO}:latest-aarch64" -docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64" +publish_manifest() { + local tag="$1" + local arch_tag_base="$2" + docker manifest rm "${DOCKERHUB_REPO}:${tag}" || true + docker manifest create \ + "${DOCKERHUB_REPO}:${tag}" \ + "${DOCKERHUB_REPO}:${arch_tag_base}-x86_64" \ + "${DOCKERHUB_REPO}:${arch_tag_base}-aarch64" + docker manifest push "${DOCKERHUB_REPO}:${tag}" +} -docker manifest rm "${DOCKERHUB_REPO}:latest" || true -docker manifest rm "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" || true -docker manifest create "${DOCKERHUB_REPO}:latest" "${DOCKERHUB_REPO}:latest-x86_64" "${DOCKERHUB_REPO}:latest-aarch64" -docker manifest create "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64" -docker manifest push "${DOCKERHUB_REPO}:latest" -docker manifest push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" +for tag in "${PRIMARY_TAGS[@]}"; do + docker tag "${SRC_X86}" "${DOCKERHUB_REPO}:${tag}-x86_64" + docker tag "${SRC_ARM}" "${DOCKERHUB_REPO}:${tag}-aarch64" + docker push "${DOCKERHUB_REPO}:${tag}-x86_64" + docker push "${DOCKERHUB_REPO}:${tag}-aarch64" + publish_manifest "${tag}" "${tag}" +done + +ARCH_BASE="${PRIMARY_TAGS[0]}" +for tag in "${EXTRA_MANIFEST_TAGS[@]}"; do + publish_manifest "${tag}" "${ARCH_BASE}" +done echo "" -echo "Successfully published release images for v${RELEASE_VERSION}" +echo "Successfully published ${MODE} images" From 7cd826d28c7e986bbcdd740f099b4eb256b52c4f Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Wed, 12 Aug 2026 18:36:31 +0800 Subject: [PATCH 09/41] =?UTF-8?q?[Bugfix][MiniCPM-o]=20Align=20Daily-Omni?= =?UTF-8?q?=20offline=20loading=20and=20duplex=20soft-=E2=80=A6=20(#6095)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: amy-why-3459 --- .../test_daily_omni_local_source.py | 92 +++++++ .../benchmarks/test_daily_omni_pack_modes.py | 26 +- .../minicpmo_4_5/test_minicpmo_4_5.py | 9 +- .../test_minicpmo_4_5_duplex_expansion.py | 18 +- .../data_modules/daily_omni_dataset.py | 247 +++++++++++++----- vllm_omni/benchmarks/patch/patch.py | 32 ++- 6 files changed, 327 insertions(+), 97 deletions(-) create mode 100644 tests/benchmarks/test_daily_omni_local_source.py diff --git a/tests/benchmarks/test_daily_omni_local_source.py b/tests/benchmarks/test_daily_omni_local_source.py new file mode 100644 index 000000000..35045a04b --- /dev/null +++ b/tests/benchmarks/test_daily_omni_local_source.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Offline Daily-Omni sources: local mirrors must not fall back to the Hub id.""" + +import json +import os +import tarfile +from pathlib import Path + +import pytest + +from vllm_omni.benchmarks.data_modules.daily_omni_dataset import ( + daily_omni_local_qa_json, + daily_omni_local_videos_dir, + resolve_daily_omni_local_root, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark] + + +def _write_qa_json(root: Path) -> Path: + qa = root / "qa.json" + qa.write_text(json.dumps([{"Question": "q", "Answer": "A", "video_id": "vid0"}]), encoding="utf-8") + return qa + + +def test_resolve_local_root_returns_none_for_hub_id() -> None: + assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None + assert resolve_daily_omni_local_root("") is None + assert resolve_daily_omni_local_root(None) is None + + +def test_plain_local_mirror(tmp_path: Path) -> None: + _write_qa_json(tmp_path) + (tmp_path / "Videos" / "vid0").mkdir(parents=True) + + root = resolve_daily_omni_local_root(str(tmp_path)) + assert root == tmp_path.resolve() + assert daily_omni_local_qa_json(root) == (tmp_path / "qa.json").resolve() + assert daily_omni_local_videos_dir(root) == root / "Videos" + + +def test_hf_cache_dir_resolves_to_snapshot(tmp_path: Path) -> None: + cache_dir = tmp_path / "datasets--liarliar--Daily-Omni" + snapshot = cache_dir / "snapshots" / "deadbeef" + snapshot.mkdir(parents=True) + _write_qa_json(snapshot) + (cache_dir / "refs").mkdir() + (cache_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8") + + root = resolve_daily_omni_local_root(str(cache_dir)) + assert root == snapshot.resolve() + assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve() + + +def test_hf_cache_dir_without_refs_picks_newest_snapshot(tmp_path: Path) -> None: + cache_dir = tmp_path / "datasets--liarliar--Daily-Omni" + older = cache_dir / "snapshots" / "aaa" + newer = cache_dir / "snapshots" / "bbb" + older.mkdir(parents=True) + newer.mkdir(parents=True) + os.utime(older, (1, 1)) + os.utime(newer, (10, 10)) + + assert resolve_daily_omni_local_root(str(cache_dir)) == newer.resolve() + + +def test_videos_tar_is_extracted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("HF_HOME", str(tmp_path / "hf_home")) + root = tmp_path / "mirror" + (root / "Videos" / "vid0").mkdir(parents=True) + (root / "Videos" / "vid0" / "vid0_video.mp4").write_bytes(b"fake") + with tarfile.open(root / "Videos.tar", "w") as tf: + tf.add(root / "Videos", arcname="Videos") + # Only the tarball is left in the mirror. + for child in (root / "Videos").rglob("*"): + if child.is_file(): + child.unlink() + (root / "Videos" / "vid0").rmdir() + (root / "Videos").rmdir() + + videos_dir = daily_omni_local_videos_dir(root) + assert videos_dir is not None + assert (videos_dir / "vid0" / "vid0_video.mp4").is_file() + + # Second call reuses the extraction instead of redoing it. + assert daily_omni_local_videos_dir(root) == videos_dir + + +def test_missing_qa_json_reports_none(tmp_path: Path) -> None: + root = resolve_daily_omni_local_root(str(tmp_path)) + assert daily_omni_local_qa_json(root) is None + assert daily_omni_local_videos_dir(root) is None diff --git a/tests/benchmarks/test_daily_omni_pack_modes.py b/tests/benchmarks/test_daily_omni_pack_modes.py index 46ca369e5..0fd231857 100644 --- a/tests/benchmarks/test_daily_omni_pack_modes.py +++ b/tests/benchmarks/test_daily_omni_pack_modes.py @@ -2,9 +2,8 @@ The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches OpenBMB's reported Daily-Omni score when frames and 1s audio segments arrive strictly -interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and an empty MiniCPM -system message (not the Qwen system prompt). These tests pin that contract so a refactor -cannot silently regress the score. +interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and no Qwen system prompt. +These tests pin that contract so a refactor cannot silently regress the score. vllm stubs are installed by tests/benchmarks/conftest.py before collection. """ @@ -99,22 +98,13 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) -> # --------------------------------------------------------------------------- -@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (1000, 64), (7, 7), (3, 5), (1, 1)]) -def test_uniform_sample_indices_matches_openbmb_midpoint(n: int, k: int) -> None: - """Pin OpenBMB midpoint-bin sampling (not np.linspace endpoints).""" +@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (7, 7), (3, 5), (1, 1)]) +def test_uniform_sample_indices_matches_numpy_linspace(n: int, k: int) -> None: got = _uniform_sample_indices(n, k) if n <= k: assert got == list(range(n)) else: - gap = n / k - expected = [int(i * gap + gap / 2) for i in range(k)] - assert got == expected - # Endpoints differ from linspace for typical Daily-Omni downsample sizes. - linspace = [int(i) for i in np.linspace(0, n - 1, k, dtype=int)] - if n == 1000 and k == 64: - assert got[0] == 7 and got[-1] == 992 - assert linspace[0] == 0 and linspace[-1] == 999 - assert got != linspace + assert got == [int(i) for i in np.linspace(0, n - 1, k, dtype=int)] def test_uniform_sample_indices_non_positive_k() -> None: @@ -223,16 +213,14 @@ def test_minicpm_interleave_missing_local_video_returns_none(qa_json, tmp_path) # --------------------------------------------------------------------------- -def test_minicpm_interleave_keeps_empty_system_message(monkeypatch, qa_json, video_dir) -> None: - """Official MiniCPM omni mode still sends role=system with empty content.""" +def test_minicpm_interleave_omits_qwen_system_message(monkeypatch, qa_json, video_dir) -> None: ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all") _stub_extract(monkeypatch, ds, []) parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None) messages = ds._build_daily_omni_openai_messages(parts, "q?", {"A": "a"}) - assert [m["role"] for m in messages] == ["system", "user"] - assert messages[0]["content"][0]["text"] == "" + assert [m["role"] for m in messages] == ["user"] def test_qwen_pack_mode_keeps_system_message(qa_json, video_dir) -> None: diff --git a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py index 18591ee6e..50474a644 100644 --- a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py +++ b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py @@ -3,8 +3,8 @@ Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches ~78% overall accuracy on Daily-Omni (``minicpm-interleave``, ``temperature=0``, -text modalities with ``use_tts_template``, and server ``--interleave-mm-strings`` -+ 1fps / 128-frame media-io kwargs, also pinned in ``minicpmo_4_5.yaml``). +text modalities, and server ``--interleave-mm-strings`` + 1fps / 128-frame +media-io kwargs, also pinned in ``minicpmo_4_5.yaml``). """ from __future__ import annotations @@ -40,10 +40,7 @@ _MAX_SEED_TTS_MEAN_WER = 0.05 # Match the validated Daily-Omni client body from daily_omni_bench.sh. _DAILY_EXTRA_BODY = { "modalities": ["text"], - "chat_template_kwargs": { - "enable_thinking": False, - "use_tts_template": True, - }, + "chat_template_kwargs": {"enable_thinking": False}, } _SEED_EXTRA_BODY = { "modalities": ["text", "audio"], diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py index 874f09502..68204fe28 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py @@ -53,7 +53,11 @@ def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_ assert result["admission"]["overflow_error_code"] == "resource_exhausted" -@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +# CUDA-only for now: this contract asserts the model speaks while the user is +# still talking, which only holds when the duplex pipeline sustains real-time +# throughput. The current NPU stack runs several times slower than real time, +# so it never reaches a mid-stream decision point. +@hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -> None: input_wav = validated_soft_interrupt_wav() @@ -70,13 +74,8 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) - timeout_s=180.0, require_audio=True, no_realtime_pacing=False, - # Single-GPU co-location cannot reliably finish the fixture's - # two-response soft-interrupt handoff before the WAV ends. - # model-policy + deterministic sampling checks streaming audio - # lifecycle instead of the diagnostic two-response contract. - validation_mode="model-policy", - temperature=0.0, - min_responses=1, + validation_mode="response-required", + min_responses=2, min_audio_deltas_per_response=2, input_sha256=SOFT_INTERRUPT_SHA256, expect_followup_response_substring=None, @@ -88,5 +87,4 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) - assert result["error_count"] == 0 assert result["response_lifecycle_ok"] is True assert result["response_audio_contract_ok"] is True - assert result["response_before_final_commit"] is True - assert result["enough_responses"] is True + assert result["followup_response_transcript_ok"] is True diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py index 95d58e4b3..512e86538 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -110,47 +110,17 @@ def _daily_omni_find_videos_root_in_extract(tmp: Path) -> Path: ) -def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: - """Download ``Videos.tar`` from the Hugging Face dataset repo and return the ``Videos`` root. +def _extract_daily_omni_videos_tar(tar_path: Path, cache_key: str) -> Path: + """Extract ``Videos.tar`` once into the shared media cache and return the ``Videos`` root. - The returned path matches ``--daily-omni-video-dir`` (directory containing ``{{video_id}}/``). - - Cached under ``HF_HOME`` / ``vllm_omni/daily_omni_media/``. Reuses extraction when the - tarball fingerprint matches. - - Raises: - ImportError: if ``huggingface_hub`` is not installed. - FileNotFoundError / RuntimeError: if the archive is missing or malformed. + Cached under ``HF_HOME`` / ``vllm_omni/daily_omni_media/``. Reuses a previous + extraction when the tarball fingerprint matches. """ - rid = (repo_id or "").strip() - if not rid: - raise ValueError("repo_id is required to download Daily-Omni Videos.tar") - - try: - from huggingface_hub import hf_hub_download - except ImportError as e: - raise ImportError( - "Daily-Omni Hub media download requires huggingface_hub. " - "Install it (e.g. with vLLM) or provide --daily-omni-video-dir with a local extract." - ) from e - - safe = rid.replace("/", "__").replace("\\", "_") + safe = cache_key.replace("/", "__").replace("\\", "_") staging_root = _daily_omni_hf_cache_root() / "vllm_omni" / "daily_omni_media" / safe videos_dir = staging_root / "Videos" marker = staging_root / ".videos_extracted" - tar_path: Path | None = None - for fname in ("Videos.tar", "videos.tar"): - try: - tar_path = Path(hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) - break - except Exception: - continue - if tar_path is None or not tar_path.is_file(): - raise FileNotFoundError( - f"Could not download Videos.tar from Hugging Face dataset {rid!r} (tried Videos.tar / videos.tar)." - ) - fp = _daily_omni_tar_fingerprint(tar_path) if marker.is_file() and videos_dir.is_dir(): try: @@ -167,7 +137,7 @@ def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: shutil.rmtree(work, ignore_errors=True) work.mkdir(parents=True) try: - logger.info("Extracting Daily-Omni Videos.tar from %s (repo=%s)", tar_path, rid) + logger.info("Extracting Daily-Omni Videos.tar from %s (cache_key=%s)", tar_path, cache_key) with tarfile.open(tar_path, "r:*") as tf: tf.extractall(path=work, filter="data") found = _daily_omni_find_videos_root_in_extract(work) @@ -178,10 +148,106 @@ def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: shutil.rmtree(work, ignore_errors=True) marker.write_text(fp, encoding="utf-8") - logger.info("Daily-Omni Hub media ready at %s", videos_dir) + logger.info("Daily-Omni media ready at %s", videos_dir) return videos_dir +def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: + """Download ``Videos.tar`` from the Hugging Face dataset repo and return the ``Videos`` root. + + The returned path matches ``--daily-omni-video-dir`` (directory containing ``{{video_id}}/``). + + Raises: + ImportError: if ``huggingface_hub`` is not installed. + FileNotFoundError / RuntimeError: if the archive is missing or malformed. + """ + rid = (repo_id or "").strip() + if not rid: + raise ValueError("repo_id is required to download Daily-Omni Videos.tar") + + try: + from huggingface_hub import hf_hub_download + except ImportError as e: + raise ImportError( + "Daily-Omni Hub media download requires huggingface_hub. " + "Install it (e.g. with vLLM) or provide --daily-omni-video-dir with a local extract." + ) from e + + tar_path: Path | None = None + for fname in ("Videos.tar", "videos.tar"): + try: + tar_path = Path(hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) + break + except Exception: + continue + if tar_path is None or not tar_path.is_file(): + raise FileNotFoundError( + f"Could not download Videos.tar from Hugging Face dataset {rid!r} (tried Videos.tar / videos.tar)." + ) + + return _extract_daily_omni_videos_tar(tar_path, rid) + + +def _resolve_hf_cache_snapshot(path: Path) -> Path: + """Map a HF hub cache dir (``datasets--org--name``) to its checked-out snapshot dir. + + Any other directory is returned unchanged. + """ + snapshots = path / "snapshots" + if not snapshots.is_dir(): + return path + ref = path / "refs" / "main" + if ref.is_file(): + try: + target = snapshots / ref.read_text(encoding="utf-8").strip() + except OSError: + target = snapshots + if target.is_dir(): + return target + revisions = [p for p in snapshots.iterdir() if p.is_dir()] + if revisions: + return max(revisions, key=lambda p: p.stat().st_mtime) + return path + + +def resolve_daily_omni_local_root(dataset_path: str | None) -> Path | None: + """Return the local Daily-Omni mirror root for ``dataset_path``, or ``None`` if not local. + + Accepts a plain directory (``qa.json`` + ``Videos/`` or ``Videos.tar``, i.e. the layout of the + Hub repo) as well as a HuggingFace hub cache directory (``datasets--liarliar--Daily-Omni``), + which is resolved to its snapshot. Hub ids such as ``liarliar/Daily-Omni`` return ``None``. + """ + raw = (dataset_path or "").strip() + if not raw: + return None + path = Path(raw).expanduser() + if not path.is_dir(): + return None + return _resolve_hf_cache_snapshot(path.resolve()) + + +def daily_omni_local_qa_json(root: Path) -> Path | None: + """Locate ``qa.json`` inside a local Daily-Omni mirror (``None`` when absent).""" + for name in ("qa.json", "QA.json"): + candidate = root / name + if candidate.is_file(): + return candidate + return None + + +def daily_omni_local_videos_dir(root: Path) -> Path | None: + """Return the ``Videos`` root of a local mirror, extracting ``Videos.tar`` when needed.""" + for name in ("Videos", "videos"): + candidate = root / name + if candidate.is_dir(): + return candidate + for name in ("Videos.tar", "videos.tar"): + tar_path = root / name + if tar_path.is_file(): + return _extract_daily_omni_videos_tar(tar_path, f"local__{root.name}") + return None + + class _ListDatasetIterator: """Simple iterator wrapper around a list to mimic HuggingFace streaming dataset behavior.""" @@ -222,24 +288,92 @@ DAILY_OMNI_SYSTEM_TEXT = ( "capable of perceiving auditory and visual inputs, as well as generating text and speech." ) -# MiniCPM-o ``get_sys_prompt(mode="omni")`` without ref-audio is an empty system string -# (still sent as role=system; see ``_build_daily_omni_openai_messages``). +# MiniCPM-o ``get_sys_prompt(mode="omni")`` without ref-audio is an empty system string. MINICPM_OMNI_SYSTEM_TEXT = "" def _uniform_sample_indices(n: int, k: int) -> list[int]: - """Evenly pick ``k`` indices from ``0..n-1`` (OpenBMB ``uniform_sample`` midpoint bins). - - Matches OmniEvalKit / MiniCPM ``uniform_sample``: midpoint of each equal-width bin, - not ``np.linspace`` endpoints. See - https://github.com/OpenBMB/OmniEvalKit/blob/1adac0577258d539efc03f16d8a0f4b3f7df6c19/o_e_Kit/utils/utils.py#L50-L54 - """ + """Evenly pick ``k`` indices from ``0..n-1`` (MiniCPM ``uniform_sample`` on index lists).""" if k <= 0: return [] if n <= k: return list(range(n)) - gap = n / k - return [int(i * gap + gap / 2) for i in range(k)] + # Match numpy linspace(..., dtype=int) used upstream. + try: + import numpy as np + + return [int(i) for i in np.linspace(0, n - 1, k, dtype=int)] + except ImportError: + if k == 1: + return [0] + return [int(round(i * (n - 1) / (k - 1))) for i in range(k)] + + +def _probe_video_frames_and_fps(video_path: Path) -> tuple[int, float]: + """Frame count and average FPS, replacing ``len(vr)`` / ``vr.get_avg_fps()``. + + PyAV is used instead of ``decord`` because ``decord`` ships x86_64-only Linux wheels + and is unmaintained, while ``av`` is already a first-class requirement and has + aarch64 wheels (Ascend hosts). + """ + import av + + with av.open(str(video_path)) as container: + stream = container.streams.video[0] + avg_fps = float(stream.average_rate) if stream.average_rate else 0.0 + num_frames = stream.frames or 0 + if num_frames <= 0: + # Containers such as raw webm leave ``frames`` at 0; derive it from duration. + duration: float | None = None + if stream.duration is not None and stream.time_base is not None: + duration = float(stream.duration * stream.time_base) + elif container.duration is not None: + duration = container.duration / av.time_base + if duration and avg_fps: + num_frames = int(math.floor(duration * avg_fps)) + if num_frames <= 0: + # Last resort for streams with no usable metadata at all. + stream.thread_type = "AUTO" + num_frames = sum(1 for _ in container.decode(stream)) + if num_frames <= 0: + raise ValueError(f"No decodable video frames in {video_path}") + return num_frames, avg_fps or 1.0 + + +def _decode_video_frames_rgb(video_path: Path, frame_idx: list[int]) -> list[Any]: + """Decode presentation-order ``frame_idx`` into RGB ``PIL.Image``s. + + Mirrors ``decord.VideoReader.get_batch(frame_idx)``: one image per requested index, + in the requested order, duplicates repeated. Decoding is a single forward pass + (no seeking) so the mapping from index to picture matches decord's indexing. + """ + import av + + wanted = sorted({int(i) for i in frame_idx}) + if not wanted: + return [] + + decoded: dict[int, Any] = {} + with av.open(str(video_path)) as container: + stream = container.streams.video[0] + stream.thread_type = "AUTO" + cursor = 0 + for position, frame in enumerate(container.decode(stream)): + if wanted[cursor] > position: + continue + image = frame.to_image() + # ``<=`` also absorbs indices that overshoot the real frame count. + while cursor < len(wanted) and wanted[cursor] <= position: + decoded[wanted[cursor]] = image + cursor += 1 + if cursor >= len(wanted): + break + + if not decoded: + raise ValueError(f"Decoded no frames from {video_path} for indices {wanted[:8]}") + # Metadata may promise more frames than the stream actually decodes; clamp to the last one. + last_image = decoded[max(decoded)] + return [decoded.get(int(i), last_image) for i in frame_idx] def _numpy_to_wav_bytes(audio_np: Any, sample_rate: int = _MINICPM_AUDIO_SR) -> bytes: @@ -1028,28 +1162,24 @@ class DailyOmniDataset(BenchmarkDataset): ) -> tuple[list[Any], list[Any]]: """Port of MiniCPM ``get_video_frame_audio_segments`` (stack_frames=1, 1fps).""" import numpy as np - from decord import VideoReader, cpu - from PIL import Image from vllm.multimodal.media.audio import load_audio - vr = VideoReader(str(video_path), ctx=cpu(0)) - avg_fps = float(vr.get_avg_fps()) or 1.0 - duration = len(vr) / avg_fps + num_video_frames, avg_fps = _probe_video_frames_and_fps(video_path) + duration = num_video_frames / avg_fps num_seconds = max(1, int(math.ceil(duration))) second_timestamps = list(range(num_seconds)) if duration > max_num_frames: timestamps = [round(i * 0.1, 1) for i in range(int(duration / 0.1))] - frame_idx = [min(int(ts * avg_fps), len(vr) - 1) for ts in timestamps] + frame_idx = [min(int(ts * avg_fps), num_video_frames - 1) for ts in timestamps] pick = _uniform_sample_indices(len(frame_idx), max_num_frames) frame_idx = [frame_idx[i] for i in pick] timestamps = [timestamps[i] for i in pick] else: - frame_idx = [min(int(i * avg_fps), len(vr) - 1) for i in range(num_seconds)] + frame_idx = [min(int(i * avg_fps), num_video_frames - 1) for i in range(num_seconds)] timestamps = second_timestamps - video = vr.get_batch(frame_idx).asnumpy() - frames = [Image.fromarray(v.astype("uint8")).convert("RGB") for v in video] + frames = _decode_video_frames_rgb(video_path, frame_idx) audio_segments: list[Any] = [] if include_audio: @@ -1127,10 +1257,7 @@ class DailyOmniDataset(BenchmarkDataset): user_content: list[dict[str, Any]] = [*mm_list, {"type": "text", "text": user_text}] system_text = self._system_text_for_pack_mode() messages: list[dict[str, Any]] = [] - # MiniCPM-o ``get_sys_prompt(mode="omni")`` still emits role=system with empty - # content (no ref-audio). Truthiness would drop that role and change chat-template - # control tokens; keep the empty system message for minicpm-interleave. - if self.pack_mode == "minicpm-interleave" or system_text: + if system_text: messages.append({"role": "system", "content": [{"type": "text", "text": system_text}]}) messages.append({"role": "user", "content": user_content}) return messages diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 70f11c0a5..3f83a1b38 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -14,6 +14,7 @@ import wave from collections.abc import Iterable from dataclasses import dataclass, replace from datetime import datetime +from pathlib import Path from typing import Any, Literal from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit @@ -38,7 +39,13 @@ from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats -from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset, DailyOmniSampleRequest +from vllm_omni.benchmarks.data_modules.daily_omni_dataset import ( + DailyOmniDataset, + DailyOmniSampleRequest, + daily_omni_local_qa_json, + daily_omni_local_videos_dir, + resolve_daily_omni_local_root, +) from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT, @@ -282,6 +289,26 @@ def get_samples(args, tokenizer): if isinstance(qa_json, str): qa_json = qa_json.strip() or None + # A local mirror of the dataset repo (plain directory or HF hub cache dir) is used + # directly: no Hub round-trip, so offline / air-gapped runs work without extra flags. + local_root: Path | None = None + if qa_json is None: + local_root = resolve_daily_omni_local_root(getattr(args, "dataset_path", None)) or ( + resolve_daily_omni_local_root(getattr(args, "hf_name", None)) + ) + if local_root is not None: + local_qa = daily_omni_local_qa_json(local_root) + if local_qa is not None: + qa_json = str(local_qa) + if video_dir is None: + video_dir = daily_omni_local_videos_dir(local_root) + logger.info("Using local Daily-Omni mirror: root=%s", local_root) + else: + logger.info( + "Local Daily-Omni path %s has no qa.json; loading it with `datasets` instead", + local_root, + ) + if qa_json is not None: logger.info( "Loading Daily-Omni dataset: qa_json=%s, video_dir=%s (Hub not used for QA)", @@ -304,7 +331,8 @@ def get_samples(args, tokenizer): repo_id = _daily_omni_repo_from_args(args) if args.dataset_name == "daily-omni": if repo_id is None: - repo_id = _DEFAULT_DAILY_OMNI_REPO + # Prefer an on-disk copy over the Hub id so offline runs keep working. + repo_id = str(local_root) if local_root is not None else _DEFAULT_DAILY_OMNI_REPO elif repo_id is None: raise ValueError( "Daily-Omni with --dataset-name hf requires " From 38a2500bdf887cd0d23c1c1616c59faec6ffd317 Mon Sep 17 00:00:00 2001 From: Wenjia Li Date: Wed, 12 Aug 2026 19:12:54 +0800 Subject: [PATCH 10/41] [Perf] MiniMax H3 Qwen3-VL support fused RMSNorm on NPU (#5915) Signed-off-by: Wenjia Li Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- .../minimax_h3/test_minimax_h3_rmsnorm.py | 26 +++++++++++++ .../diffusion/models/minimax_h3/encoder.py | 38 +++++++++++-------- 2 files changed, 49 insertions(+), 15 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py new file mode 100644 index 000000000..551fd68a5 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm + + +def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None: + """Qwen keeps BF16 gamma while native fallback accumulates in FP32.""" + eps = 1e-6 + norm = MiniMaxH3Qwen3VLRMSNorm(hidden_size=4, eps=eps, dtype=torch.bfloat16) + norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16)) + x = torch.tensor( + [[1.0, -2.0, 3.0, -4.0], [-5.0, 6.0, -7.0, 8.0]], + dtype=torch.bfloat16, + ) + + x_fp32 = x.float() + expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype) + + assert isinstance(norm, RMSNorm) + assert norm.weight.dtype == torch.bfloat16 + assert set(norm.state_dict()) == {"weight"} + torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0) diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 15098b098..2e07ace96 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -39,6 +39,8 @@ import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger +from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm + MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50 MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 @@ -291,18 +293,22 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): # --------------------------------------------------------------------------- -class MiniMaxH3Qwen3VLRMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps +class MiniMaxH3Qwen3VLRMSNorm(DiffusionRMSNorm): + """Qwen3-VL RMSNorm using the common implementation. - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - 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) + The model-specific name keeps checkpoint weight keys stable, while the + common RMSNorm dispatches to torch_npu.npu_rms_norm on Ascend. Gamma uses + the checkpoint/model dtype, but native fallbacks keep the RMS variance + reduction and scaling in float32. + """ + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + ) -> None: + super().__init__(hidden_size, eps=eps, dtype=dtype) def _rotate_half(x: torch.Tensor) -> torch.Tensor: @@ -685,8 +691,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): dtype=dtype, input_is_parallel=True, ) - self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) + self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) def forward( self, @@ -760,8 +766,10 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): self.hidden_size = config.hidden_size self.self_attn = MiniMaxH3Qwen3VLTextAttention(group, config, dtype) self.mlp = MiniMaxH3Qwen3VLTextMLP(group, config, dtype) - self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype) + self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm( + config.hidden_size, eps=config.rms_norm_eps, dtype=dtype + ) def forward( self, From fcf1085395940df1f9486e2f1b6f11581d8002c9 Mon Sep 17 00:00:00 2001 From: Shi Boao Bowie Date: Wed, 12 Aug 2026 22:58:23 +0800 Subject: [PATCH 11/41] [Refactor][OutputProcessor 2/8]: OmniRequestOutput should inherit RequestOutput and no nested wrap-up (#5146) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Boao Shi Signed-off-by: bowieshi Signed-off-by: princepride Signed-off-by: R2-Y Signed-off-by: amy-why-3459 Signed-off-by: Shi Boao Bowie Co-authored-by: 汪志鹏 Co-authored-by: ruirui(rein) yang <73573651+R2-Y@users.noreply.github.com> Co-authored-by: amy-why-3459 Co-authored-by: Gao Han Co-authored-by: Hongsheng Liu --- benchmarks/diffusion/quantization_quality.py | 21 +- benchmarks/glm_image/vllm-omni/inference.py | 2 +- .../offline_inference/aura_omni/end2end.py | 2 +- .../offline_inference/covo_audio/end2end.py | 2 +- .../image_to_image/image_edit.py | 2 +- .../offline_inference/dynin_omni/end2end.py | 2 +- .../hunyuan_image3/end2end.py | 2 +- .../image_to_image/image_edit.py | 2 +- .../image_to_video/image_to_video.py | 4 +- .../run_mammothmoda2_image_summarize.py | 117 ++++++ .../offline_inference/mimo_audio/end2end.py | 2 +- .../ming_flash_omni/end2end.py | 2 +- .../offline_inference/qwen2_5_omni/end2end.py | 2 +- .../offline_inference/qwen3_omni/end2end.py | 2 +- .../qwen3_omni/end2end_async_chunk.py | 2 +- .../speech_to_video/speech_to_video.py | 5 +- .../offline_inference/step_audio2/end2end.py | 2 +- .../text_to_audio/text_to_audio.py | 5 +- .../text_to_image/gradio_demo.py | 2 +- .../text_to_image/text_to_image.py | 2 +- .../text_to_speech/cosyvoice3/end2end.py | 2 +- .../text_to_speech/fish_speech/end2end.py | 4 +- .../text_to_speech/moss_tts_nano/end2end.py | 2 +- .../text_to_speech/omnivoice/end2end.py | 2 +- .../text_to_speech/qwen3_tts/end2end.py | 4 +- .../qwen3_tts/word_timestamps.py | 172 ++++++++ .../text_to_speech/soulxsinger/benchmark.py | 2 +- .../text_to_speech/soulxsinger/end2end.py | 2 +- .../text_to_video/text_to_video.py | 6 +- .../x_to_video_audio/x_to_video_audio.py | 2 +- .../batching/test_diffusion_batching.py | 2 +- tests/diffusion/cache/test_teacache.py | 4 +- .../distributed/test_sequence_parallel.py | 2 +- .../distributed/test_tensor_parallel.py | 4 +- tests/diffusion/lora/test_diffusion_lora.py | 4 +- .../offloader/test_diffusion_cpu_offload.py | 4 +- .../test_diffusion_layerwise_offload.py | 23 +- .../quantization/test_gguf_diffusion.py | 4 +- .../quantization/test_quantization_fp8.py | 18 +- .../quantization/test_quantization_quality.py | 24 +- .../test_trajectory_similarity_tool.py | 8 +- .../test_diffusion_streaming_output.py | 2 +- .../test_bagel_mooncake_connector.py | 5 +- .../test_bagel_shared_memory_connector.py | 5 +- tests/e2e/accuracy/test_hunyuan_image3.py | 4 +- .../accuracy/test_videogen_modelopt_quant.py | 4 +- .../comfyui/test_comfyui_integration.py | 30 +- .../e2e/features/helpers/verl_omni_server.py | 2 +- .../test_audiox_model_expansion.py | 4 +- tests/e2e/offline_inference/test_bagel.py | 2 +- .../test_cosyvoice3_expansion.py | 5 +- .../test_dynin_omni_expansion.py | 8 +- ...est_flux1_dev_autoround_w4a16_expansion.py | 2 +- ...est_glm_image_autoround_w4a16_expansion.py | 4 +- .../offline_inference/test_hunyuanimage3.py | 2 +- .../test_hunyuanimage3_i2t_expansion.py | 2 +- .../test_mammoth_moda2_expansion.py | 2 +- .../test_ming_tts_expansion.py | 18 +- .../test_moss_tts_nano_expansion.py | 4 +- .../test_omnivoice_expansion.py | 2 +- ...st_qwen_image_autoround_w4a16_expansion.py | 2 +- .../test_stable_audio_expansion.py | 4 +- .../test_step_audio2_expansion.py | 8 +- .../test_minicpmo_realtime_duplex_drivers.py | 6 +- tests/engine/test_async_omni_engine_abort.py | 2 +- tests/engine/test_omni_request_output.py | 394 +++++++++++------- .../openai_api/test_duplex_handler.py | 8 +- .../openai_api/test_image_server.py | 16 +- .../openai_api/test_serving_speech.py | 12 +- .../openai_api/test_serving_video_stream.py | 4 +- tests/entrypoints/test_async_omni_duplex.py | 10 +- tests/entrypoints/test_omni_entrypoints.py | 19 +- .../entrypoints/test_stream_finish_reason.py | 12 +- tests/helpers/media.py | 20 +- tests/helpers/runtime.py | 6 +- vllm_omni/diffusion/utils/image_output.py | 9 +- .../omni_connectors/utils/serialization.py | 41 +- vllm_omni/engine/orchestrator.py | 4 +- vllm_omni/entrypoints/omni_base.py | 14 +- vllm_omni/entrypoints/openai/api_server.py | 6 - .../openai/serving_audio_generate.py | 3 - vllm_omni/entrypoints/openai/serving_chat.py | 38 +- .../entrypoints/openai/serving_speech.py | 2 +- vllm_omni/entrypoints/openai/serving_video.py | 26 +- .../openai/serving_video_output_stream.py | 13 +- .../entrypoints/openai/video_stream_base.py | 4 +- vllm_omni/outputs/__init__.py | 305 ++++++-------- ...compare_diffusion_trajectory_similarity.py | 7 +- vllm_omni/utils/audio.py | 4 +- 89 files changed, 944 insertions(+), 639 deletions(-) create mode 100644 examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py create mode 100644 examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py diff --git a/benchmarks/diffusion/quantization_quality.py b/benchmarks/diffusion/quantization_quality.py index b47fae41f..ac0fe9d70 100644 --- a/benchmarks/diffusion/quantization_quality.py +++ b/benchmarks/diffusion/quantization_quality.py @@ -158,7 +158,6 @@ def _build_omni_kwargs(args, quantization=None): def _generate_image(omni, args, prompt, seed): """Generate a single image and return (PIL.Image, time_seconds, memory_gib).""" from vllm_omni.inputs.data import OmniDiffusionSamplingParams - from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed) @@ -176,7 +175,7 @@ def _generate_image(omni, args, prompt, seed): elapsed = time.perf_counter() - start peak_mem = torch.accelerator.max_memory_allocated() / (1024**3) - req_out = OmniRequestOutput.unwrap_result(outputs) + req_out = outputs[0] if isinstance(outputs, list) else outputs if not req_out.images: raise ValueError("Could not extract image output from result.") img = req_out.images[0] @@ -206,15 +205,15 @@ def _generate_video(omni, args, prompt, seed): elapsed = time.perf_counter() - start peak_mem = torch.accelerator.max_memory_allocated() / (1024**3) - first = outputs[0] - if hasattr(first, "request_output") and isinstance(first.request_output, list): - inner = first.request_output[0] - if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"): - frames = inner.images[0] if inner.images else None - else: - frames = inner - elif hasattr(first, "images") and first.images: - frames = first.images[0] + if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput): + first = outputs[0] + elif isinstance(outputs, OmniRequestOutput): + first = outputs + else: + raise ValueError("Could not extract video frames from output.") + + if hasattr(first, "images"): + frames = first.images[0] if first.images else None else: raise ValueError("Could not extract video frames from output.") diff --git a/benchmarks/glm_image/vllm-omni/inference.py b/benchmarks/glm_image/vllm-omni/inference.py index 5729da071..584ed5c63 100644 --- a/benchmarks/glm_image/vllm-omni/inference.py +++ b/benchmarks/glm_image/vllm-omni/inference.py @@ -300,7 +300,7 @@ def benchmark(args: argparse.Namespace) -> None: output_idx = 0 for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True): if stage_outputs.final_output_type == "image": - request_output = stage_outputs.request_output + request_output = stage_outputs request_id = getattr(request_output, "request_id", "") images = getattr(request_output, "images", []) diff --git a/examples/offline_inference/aura_omni/end2end.py b/examples/offline_inference/aura_omni/end2end.py index c08287515..bac385d8e 100644 --- a/examples/offline_inference/aura_omni/end2end.py +++ b/examples/offline_inference/aura_omni/end2end.py @@ -131,7 +131,7 @@ def build_sampling_params() -> list[SamplingParams]: def save_stage_output(stage_outputs, output_dir: str) -> None: - output = stage_outputs.request_output + output = stage_outputs request_id = output.request_id if stage_outputs.final_output_type == "text": text_output = output.outputs[0].text diff --git a/examples/offline_inference/covo_audio/end2end.py b/examples/offline_inference/covo_audio/end2end.py index ce728373d..5b7ac5fef 100644 --- a/examples/offline_inference/covo_audio/end2end.py +++ b/examples/offline_inference/covo_audio/end2end.py @@ -104,7 +104,7 @@ def main(args): os.makedirs(output_dir, exist_ok=True) for stage_outputs in omni_outputs: - output = stage_outputs.request_output + output = stage_outputs if stage_outputs.final_output_type == "text": request_id = output.request_id text_output = output.outputs[0].text diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py index 7c91c0323..5f4311050 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -228,7 +228,7 @@ async def main(): if not outputs: raise ValueError("No output produced from omni.generate()") - first_out = outputs[0].request_output + first_out = outputs[0] req_out: OmniRequestOutput = first_out # Verify trajectory data (from custom pipeline) diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py index 3442bd9bc..0cab46fee 100644 --- a/examples/offline_inference/dynin_omni/end2end.py +++ b/examples/offline_inference/dynin_omni/end2end.py @@ -741,7 +741,7 @@ def make_mmu_prompt( def iter_mm_outputs(outputs: list[Any]): for omni_out in outputs: - req_out = getattr(omni_out, "request_output", None) + req_out = omni_out req_list = req_out if isinstance(req_out, list) else [req_out] for item in req_list: if item is None: diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 62836f963..9a2350156 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -303,7 +303,7 @@ def main(): ) img_idx = 0 for req_output in omni_outputs: - ro = getattr(req_output, "request_output", None) + ro = req_output stage_id = getattr(req_output, "stage_id", None) # AR stage text — each CompletionOutput.text is already a delta when diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 10cf91088..ae84353f7 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -621,7 +621,7 @@ def main(): images = getattr(output, "images", None) if images: break - req_out = getattr(output, "request_output", None) + req_out = output images = getattr(req_out, "images", None) if req_out is not None else None if images: break diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index c34e9101b..222502e8a 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -646,8 +646,8 @@ def main(): if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) - if frames.is_pipeline_output and frames.request_output is not None: - inner_output = frames.request_output + if frames.is_pipeline_output and frames is not None: + inner_output = frames if isinstance(inner_output, OmniRequestOutput): if inner_output.multimodal_output and "audio" in inner_output.multimodal_output: audio = inner_output.multimodal_output["audio"] diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py new file mode 100644 index 000000000..5363948dd --- /dev/null +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py @@ -0,0 +1,117 @@ +""" +Offline inference example: MammothModa2 image summarization (single AR stage). + +Example: + uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \ + --model path/to/MammothModa2-Preview \ + --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \ + --image /path/to/input.jpg \ + --question "Please summarize the content of this image." +""" + +from __future__ import annotations + +import argparse +import os + +from PIL import Image +from vllm import SamplingParams +from vllm.multimodal.image import convert_image_mode + +from vllm_omni import Omni + +DEFAULT_SYSTEM = "You are a helpful assistant." +DEFAULT_QUESTION = "Please summarize the content of this image." + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).") + parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.") + parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.") + parser.add_argument("--image", type=str, required=True, help="Path to input image.") + parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.") + parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.") + parser.add_argument( + "--max-tokens", + type=int, + default=512, + help="Max new tokens to generate.", + ) + parser.add_argument("--temperature", type=float, default=0.2) + parser.add_argument("--top-p", type=float, default=0.9) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) + return parser.parse_args() + + +def build_prompt(system: str, question: str) -> str: + return ( + f"<|im_start|>system\n{system}<|im_end|>\n" + "<|im_start|>user\n" + "<|vision_start|><|image_pad|><|vision_end|>" + f"{question}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + + +def main() -> None: + args = parse_args() + + if not os.path.exists(args.image): + raise FileNotFoundError(f"Image file not found: {args.image}") + + pil_image = Image.open(args.image) + image_data = convert_image_mode(pil_image, "RGB") + prompt = build_prompt(args.system, args.question) + + omni = Omni( + model=args.model, + deploy_config=args.deploy_config, + trust_remote_code=args.trust_remote_code, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + ) + try: + sp = SamplingParams( + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=-1, + max_tokens=int(args.max_tokens), + seed=int(args.seed), + detokenize=True, + ) + # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None]. + # Consume it inside the try block so the worker isn't closed early. + outputs = list( + omni.generate( + [ + { + "prompt": prompt, + "multi_modal_data": {"image": image_data}, + "additional_information": {"omni_task": ["chat"]}, + } + ], + [sp], + ) + ) + finally: + omni.close() + + lines: list[str] = [] + for stage_outputs in outputs: + ro = stage_outputs + text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro) + lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n") + lines.append("answer:\n") + lines.append(text.strip() + "\n") + lines.append("\n") + + print("\n".join(lines)) + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index 4ba8f483a..e1772838a 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/examples/offline_inference/mimo_audio/end2end.py @@ -297,7 +297,7 @@ def main(args): os.makedirs(output_dir, exist_ok=True) for stage_outputs in omni_outputs: - output = stage_outputs.request_output + output = stage_outputs if stage_outputs.final_output_type == "text": request_id = output.request_id text_output = output.outputs[0].text diff --git a/examples/offline_inference/ming_flash_omni/end2end.py b/examples/offline_inference/ming_flash_omni/end2end.py index 9fe907ae5..d41a8045b 100644 --- a/examples/offline_inference/ming_flash_omni/end2end.py +++ b/examples/offline_inference/ming_flash_omni/end2end.py @@ -347,7 +347,7 @@ def main(args): omni.start_profile(stages=args.profiler_stages) for stage_outputs in omni.generate(prompts, sampling_params_list): - output = stage_outputs.request_output + output = stage_outputs if stage_outputs.final_output_type == "text": request_id = output.request_id text_output = output.outputs[0].text diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index d8769b5f7..6dbafdfc1 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -392,7 +392,7 @@ def main(args): total_requests = len(prompts) processed_count = 0 for stage_outputs in omni_generator: - output = stage_outputs.request_output + output = stage_outputs if stage_outputs.final_output_type == "text": request_id = output.request_id text_output = output.outputs[0].text diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index ae719e1c3..9434b3058 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -367,7 +367,7 @@ def main(args): print(f"query type: {args.query_type}") for stage_outputs in omni_generator: - output = stage_outputs.request_output + output = stage_outputs if stage_outputs.final_output_type == "text": request_id = output.request_id text_output = output.outputs[0].text diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py index 7d051ac4a..c224586c4 100644 --- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py +++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py @@ -237,7 +237,7 @@ async def run_single_request( sampling_params_list=sampling_params_list, output_modalities=output_modalities, ): - output = omni_output.request_output + output = omni_output if omni_output.final_output_type == "text": if stage_0_first_output_ts is None: stage_0_first_output_ts = time.perf_counter() diff --git a/examples/offline_inference/speech_to_video/speech_to_video.py b/examples/offline_inference/speech_to_video/speech_to_video.py index bfa6d961a..964cd4874 100644 --- a/examples/offline_inference/speech_to_video/speech_to_video.py +++ b/examples/offline_inference/speech_to_video/speech_to_video.py @@ -26,7 +26,6 @@ import torch from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -352,8 +351,8 @@ def main(): else: print("[Profiler] No valid profiling data returned.") - # Extract output from result - output = OmniRequestOutput.unwrap_result(result) + # omni.generate() returns a list for sync calls; unwrap single-result list. + output = result[0] if isinstance(result, list) else result if not output.images: raise ValueError("No video frames found in OmniRequestOutput.") diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py index e0d1836a8..f175db750 100644 --- a/examples/offline_inference/step_audio2/end2end.py +++ b/examples/offline_inference/step_audio2/end2end.py @@ -345,7 +345,7 @@ def main(args): # Process outputs from each stage for stage_idx, stage_outputs in enumerate(omni_outputs): # request_output may be a single object or a list - raw_output = stage_outputs.request_output + raw_output = stage_outputs outputs_list = raw_output if isinstance(raw_output, list) else [raw_output] if stage_outputs.final_output_type == "text": diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 3ad88a6f3..eff9c9d1a 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -43,6 +43,7 @@ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_extras import get_extra_body_params, get_model_class_name +from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m") @@ -388,9 +389,9 @@ def main(): raise ValueError("No output generated from omni.generate()") output = outputs[0] - if not hasattr(output, "request_output") or not output.request_output: + if not isinstance(output, OmniRequestOutput) or not output: raise ValueError("No request_output found in OmniRequestOutput") - request_output = output.request_output + request_output = output if not hasattr(request_output, "multimodal_output"): raise ValueError("No multimodal_output found in request_output") diff --git a/examples/offline_inference/text_to_image/gradio_demo.py b/examples/offline_inference/text_to_image/gradio_demo.py index 09e011240..bf60c3064 100644 --- a/examples/offline_inference/text_to_image/gradio_demo.py +++ b/examples/offline_inference/text_to_image/gradio_demo.py @@ -115,7 +115,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks: ) images_outputs = [] for output in outputs: - req_out = output.request_output + req_out = output if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): raise ValueError("Invalid request_output structure or missing 'images' key") images = req_out.images diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 2cadb49e6..360c726fd 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -656,7 +656,7 @@ def main(): images = getattr(output, "images", None) if images: break - req_out = getattr(output, "request_output", None) + req_out = output images = getattr(req_out, "images", None) if req_out is not None else None if images: break diff --git a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py index ccfb1aaef..62917de07 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -168,7 +168,7 @@ def run_e2e(): print(f"Received {len(outputs)} outputs.") for i, output in enumerate(outputs): try: - ro = output.request_output + ro = output if ro is None: print("No request_output found.") continue diff --git a/examples/offline_inference/text_to_speech/fish_speech/end2end.py b/examples/offline_inference/text_to_speech/fish_speech/end2end.py index 5661cccd3..18b4d79f8 100644 --- a/examples/offline_inference/text_to_speech/fish_speech/end2end.py +++ b/examples/offline_inference/text_to_speech/fish_speech/end2end.py @@ -179,7 +179,7 @@ def main(args): t_start = time.perf_counter() for stage_outputs in omni.generate(inputs): - request_output = stage_outputs.request_output + request_output = stage_outputs if request_output is None or not request_output.outputs: continue _save_wav( @@ -215,7 +215,7 @@ async def main_streaming(args): sample_rate = None async for stage_output in omni.generate(prompt, request_id=request_id): - mm = stage_output.request_output.outputs[0].multimodal_output + mm = stage_output.outputs[0].multimodal_output if not stage_output.finished: t_now = time.perf_counter() sample_rate, n = _collect_audio_chunk(mm, all_audio_chunks, sample_rate) diff --git a/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py b/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py index 91291743e..979770865 100644 --- a/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py +++ b/examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py @@ -137,7 +137,7 @@ def main(args) -> None: params_list = sampling_params for stage_outputs in omni.generate(inputs, params_list): - for i, req_output in enumerate(stage_outputs.request_output): + for i, req_output in enumerate(stage_outputs): for j, out in enumerate(req_output.outputs): mm = out.multimodal_output if mm is None: diff --git a/examples/offline_inference/text_to_speech/omnivoice/end2end.py b/examples/offline_inference/text_to_speech/omnivoice/end2end.py index a6bdc7b4e..f1699d0b8 100644 --- a/examples/offline_inference/text_to_speech/omnivoice/end2end.py +++ b/examples/offline_inference/text_to_speech/omnivoice/end2end.py @@ -135,7 +135,7 @@ def run_e2e(): print(f"Received {len(outputs)} outputs.") for i, output in enumerate(outputs): try: - ro = output.request_output + ro = output if ro is None: print("No request_output found.") continue diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py index 8150338d5..7c0ce125e 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py @@ -391,7 +391,7 @@ def main(args): for batch_start in range(0, len(inputs), batch_size): batch = inputs[batch_start : batch_start + batch_size] for stage_outputs in omni.generate(batch): - output = stage_outputs.request_output + output = stage_outputs _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output) @@ -412,7 +412,7 @@ async def main_streaming(args): t_prev = t_start chunk_idx = 0 async for stage_output in omni.generate(prompt, request_id=request_id): - mm = stage_output.request_output.outputs[0].multimodal_output + mm = stage_output.outputs[0].multimodal_output if not stage_output.finished: t_now = time.perf_counter() audio = mm.get("audio") diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py new file mode 100644 index 000000000..62f5e2caa --- /dev/null +++ b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py @@ -0,0 +1,172 @@ +"""Offline Qwen3-TTS example with word-level timestamps. + +This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with +the same shared forced aligner used by the streaming server path, then writes +the WAV and JSON sidecar: + + python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \ + --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B + +On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0`` +to avoid FlashInfer sampler JIT during warmup. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from pathlib import Path +from typing import Any + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") +# The in-process offline aligner consumes token_classify outputs. Keeping V1 +# multiprocessing off avoids msgspec serialization issues for token-wise logits. +os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + +import numpy as np +import soundfile as sf +import torch +from end2end import _estimate_prompt_len +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import Omni +from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config + + +def _default_stage_config() -> str: + repo_root = Path(__file__).resolve().parents[4] + return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml") + + +def _build_custom_voice_input(args: Any) -> dict[str, Any]: + additional_information = { + "task_type": ["CustomVoice"], + "text": [args.text], + "language": [args.language], + "speaker": [args.speaker], + "instruct": [args.instructions], + "max_new_tokens": [args.max_new_tokens], + } + return { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model), + "additional_information": additional_information, + } + + +def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]: + audio_data = mm["audio"] + sr_raw = mm["sr"] + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val) + audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data + return audio_tensor.float().detach().cpu().flatten(), sample_rate + + +def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes: + samples = audio.numpy() + pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype(" list[dict[str, Any]] | None: + aligner_config = build_forced_aligner_config(args) + if aligner_config is None: + raise ValueError("--forced-aligner or --forced-aligner-config is required") + timestamps = await align( + audio=_float_audio_to_pcm16_bytes(audio), + text=args.text, + sample_rate=sample_rate, + config=aligner_config, + language=args.language, + ) + if timestamps is None: + return None + return [ + { + "word": item.word, + "start_ms": item.start_ms, + "end_ms": item.end_ms, + } + for item in timestamps + ] + + +def main(args: Any) -> None: + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + omni_kwargs = vars(args).copy() + for key in ( + "forced_aligner", + "text", + "language", + "speaker", + "instructions", + "max_new_tokens", + "output_dir", + ): + omni_kwargs.pop(key, None) + omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config() + omni_kwargs["log_stats"] = args.log_stats + omni = Omni(**omni_kwargs) + + prompt = _build_custom_voice_input(args) + final_output = None + for stage_outputs in omni.generate([prompt]): + final_output = stage_outputs + if final_output is None: + raise RuntimeError("Qwen3-TTS did not produce an output.") + + mm = final_output.outputs[0].multimodal_output + audio, sample_rate = _audio_tensor_and_sample_rate(mm) + + wav_path = output_dir / "qwen3_tts_word_timestamps.wav" + # Align the in-memory tensor, not the saved WAV. The WAV is only a demo + # artifact so users can listen to the same audio referenced by the sidecar. + timestamps = asyncio.run(_run_alignment(args, audio, sample_rate)) + sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV") + + sidecar = { + "text": args.text, + "sample_rate": sample_rate, + "audio_path": str(wav_path), + "timestamps": timestamps, + } + json_path = output_dir / "qwen3_tts_word_timestamps.json" + json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + print(f"Saved audio: {wav_path}") + print(f"Saved timestamps: {json_path}") + print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2)) + + +def parse_args() -> Any: + parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example") + parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name") + parser.add_argument( + "--forced-aligner", + default=None, + help="Qwen3 forced aligner model path/name", + ) + parser.add_argument( + "--forced-aligner-config", + default=None, + help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)", + ) + parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML") + parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align") + parser.add_argument("--language", default="English", help="Qwen3-TTS language field") + parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name") + parser.add_argument("--instructions", default="", help="Optional speaking style instruction") + parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens") + parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar") + parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging") + nullify_stage_engine_defaults(parser) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py index 97ed7eb97..14dd5e145 100644 --- a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py +++ b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py @@ -38,7 +38,7 @@ def _stage_durations_from_output(outputs) -> dict[str, float]: durations = getattr(omni_out, "stage_durations", None) or {} if durations: return dict(durations) - ro = getattr(omni_out, "request_output", None) + ro = omni_out if ro is not None: return dict(getattr(ro, "stage_durations", None) or {}) return {} diff --git a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py index 116b8da4b..a1bb112af 100644 --- a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py +++ b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py @@ -99,7 +99,7 @@ def extract_audio(outputs) -> tuple[np.ndarray, int]: out = outputs[0] mm = getattr(out, "multimodal_output", None) if mm is None: - ro = getattr(out, "request_output", None) + ro = out if ro is not None: mm = getattr(ro, "multimodal_output", None) if mm is None and ro.outputs: diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 383887f5a..2e272b7a7 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -390,7 +390,7 @@ def _extract_peak_memory_mb(result: Any) -> float: return 0.0 val = getattr(result, "peak_memory_mb", 0.0) if not val: - inner = getattr(result, "request_output", None) + inner = result if isinstance(inner, list): inner = inner[0] if inner else None val = getattr(inner, "peak_memory_mb", 0.0) @@ -575,8 +575,8 @@ def main(): if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) - if frames.is_pipeline_output and frames.request_output is not None: - inner_output = frames.request_output + if frames.is_pipeline_output and frames is not None: + inner_output = frames if isinstance(inner_output, OmniRequestOutput): if inner_output.multimodal_output and "audio" in inner_output.multimodal_output: audio = inner_output.multimodal_output["audio"] diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 0efbdc37d..9512ac65c 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py @@ -138,7 +138,7 @@ def _extract_peak_memory_mb(result: Any) -> float: return 0.0 val = getattr(result, "peak_memory_mb", 0.0) if not val: - inner = getattr(result, "request_output", None) + inner = result if isinstance(inner, list): inner = inner[0] if inner else None val = getattr(inner, "peak_memory_mb", 0.0) diff --git a/tests/diffusion/batching/test_diffusion_batching.py b/tests/diffusion/batching/test_diffusion_batching.py index fa96b54a4..60136ce4d 100644 --- a/tests/diffusion/batching/test_diffusion_batching.py +++ b/tests/diffusion/batching/test_diffusion_batching.py @@ -125,7 +125,7 @@ def _extract_images(output: OmniRequestOutput) -> list: return output.images # When the output comes from the orchestrator pipeline, images may be # nested inside request_output. - inner = getattr(output, "request_output", None) + inner = output if inner is not None and hasattr(inner, "images") and inner.images: return inner.images return [] diff --git a/tests/diffusion/cache/test_teacache.py b/tests/diffusion/cache/test_teacache.py index 8152792fc..62e0d6698 100644 --- a/tests/diffusion/cache/test_teacache.py +++ b/tests/diffusion/cache/test_teacache.py @@ -56,10 +56,10 @@ def test_teacache(model_name: str): # Extract images from request_output['images'] first_output = outputs[0] assert first_output.final_output_type == "image" - if not hasattr(first_output, "request_output") or not first_output.request_output: + if not isinstance(first_output, OmniRequestOutput) or not first_output: raise ValueError("No request_output found in OmniRequestOutput") - req_out = first_output.request_output + req_out = first_output if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): raise ValueError("Invalid request_output structure or missing 'images' key") diff --git a/tests/diffusion/distributed/test_sequence_parallel.py b/tests/diffusion/distributed/test_sequence_parallel.py index c02667b93..d43d18c49 100644 --- a/tests/diffusion/distributed/test_sequence_parallel.py +++ b/tests/diffusion/distributed/test_sequence_parallel.py @@ -131,7 +131,7 @@ def _run_inference( elapsed_ms = (time.time() - start) * 1000 return InferenceResult( - images=outputs[0].request_output.images, + images=outputs[0].images, elapsed_ms=elapsed_ms, ) finally: diff --git a/tests/diffusion/distributed/test_tensor_parallel.py b/tests/diffusion/distributed/test_tensor_parallel.py index 9f39112ae..c50c4336d 100644 --- a/tests/diffusion/distributed/test_tensor_parallel.py +++ b/tests/diffusion/distributed/test_tensor_parallel.py @@ -49,10 +49,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: def _extract_single_image(outputs) -> Image.Image: first_output = outputs[0] assert first_output.final_output_type == "image" - if not hasattr(first_output, "request_output") or not first_output.request_output: + if not isinstance(first_output, OmniRequestOutput) or not first_output: raise ValueError("No request_output found in OmniRequestOutput") - req_out = first_output.request_output + req_out = first_output if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): raise ValueError("Invalid request_output structure or missing 'images' key") diff --git a/tests/diffusion/lora/test_diffusion_lora.py b/tests/diffusion/lora/test_diffusion_lora.py index 75e122a04..40a79a5e2 100644 --- a/tests/diffusion/lora/test_diffusion_lora.py +++ b/tests/diffusion/lora/test_diffusion_lora.py @@ -29,10 +29,10 @@ def test_diffusion_model(tmp_path: Path, omni_runner): raise ValueError("Empty outputs from Omni.generate()") first_output = outputs[0] assert first_output.final_output_type == "image" - if not hasattr(first_output, "request_output") or not first_output.request_output: + if not isinstance(first_output, OmniRequestOutput) or not first_output: raise ValueError("No request_output found in OmniRequestOutput") - req_out = first_output.request_output + req_out = first_output if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): raise ValueError("Invalid request_output structure or missing 'images' key") return req_out.images diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 549a98142..810f0dd52 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -122,8 +122,8 @@ def test_cpu_offload_diffusion_model(model_name: str): print(f"No offload peak memory: {no_offload_peak_memory} MB") if model_name == "stabilityai/stable-audio-open-1.0": - audio_offload = output_offload[0].request_output.multimodal_output.get("audio") - audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio") + audio_offload = output_offload[0].multimodal_output.get("audio") + audio_no_offload = output_no_offload[0].multimodal_output.get("audio") check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2) # Set platform-specific VRAM saving thresholds to account diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index 342c9833b..26944b202 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -20,18 +20,10 @@ IMAGE_VIDEO_MODELS = { MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} -LAYERWISE_MODEL_CASES = [ - pytest.param( - "stabilityai/stable-audio-open-1.0", - marks=[pytest.mark.full_model, pytest.mark.diffusion], - id="stable_audio", - ), - pytest.param( - "riverclouds/qwen_image_random", - marks=[pytest.mark.core_model, pytest.mark.diffusion], - id="qwen_image_random", - ), -] +MODEL_MARKS = { + "riverclouds/qwen_image_random": pytest.mark.core_model, + "stabilityai/stable-audio-open-1.0": pytest.mark.full_model, +} AUDIO_MODEL_PARAMS = { "runner_params": {}, @@ -101,8 +93,9 @@ def run_inference( return peak, output +@pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325"}) -@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES) +@pytest.mark.parametrize("model_name", list(MODELS.keys())) def test_layerwise_offload_diffusion_model(model_name: str): """Test that layerwise offloading reduces GPU memory usage. @@ -134,8 +127,8 @@ def test_layerwise_offload_diffusion_model(model_name: str): print(f"No offload peak memory: {no_offload_peak_memory} MB") if model_name == "stabilityai/stable-audio-open-1.0": - audio_offload = output_offload[0].request_output.multimodal_output.get("audio") - audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio") + audio_offload = output_offload[0].multimodal_output.get("audio") + audio_no_offload = output_no_offload[0].multimodal_output.get("audio") # Match the sibling cpu-offload test's tolerance: layerwise offload moves # blocks across the PCIe bus on a side stream, which can perturb cuBLAS # algorithm selection and produce ~ULP-level drift larger than 1e-3. diff --git a/tests/diffusion/quantization/test_gguf_diffusion.py b/tests/diffusion/quantization/test_gguf_diffusion.py index fa87e8a7c..243282b0c 100644 --- a/tests/diffusion/quantization/test_gguf_diffusion.py +++ b/tests/diffusion/quantization/test_gguf_diffusion.py @@ -94,8 +94,8 @@ def _generate_single_stage_image( if hasattr(first_output, "images") and first_output.images: images = first_output.images else: - assert hasattr(first_output, "request_output") and first_output.request_output - request_output = first_output.request_output + assert isinstance(first_output, OmniRequestOutput) and first_output + request_output = first_output if isinstance(request_output, list): req_out = request_output[0] else: diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 58f91742d..581a7498f 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -88,8 +88,8 @@ def _generate_single_stage_image( if hasattr(first_output, "images") and first_output.images: images = first_output.images else: - assert hasattr(first_output, "request_output") and first_output.request_output - request_output = first_output.request_output + assert isinstance(first_output, OmniRequestOutput) and first_output + request_output = first_output if isinstance(request_output, list): req_out = request_output[0] else: @@ -147,10 +147,10 @@ def _generate_single_stage_video( first = outputs[0] - # Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output). + # Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput). frames: Any = None - if hasattr(first, "request_output") and isinstance(first.request_output, list): - inner = first.request_output[0] + if isinstance(first, OmniRequestOutput) and isinstance(first, list): + inner = first[0] if isinstance(inner, OmniRequestOutput) and inner.images: frames = inner.images[0] if frames is None and hasattr(first, "images") and first.images: @@ -234,8 +234,8 @@ def _generate_bagel_image( if images := getattr(req_output, "images", None): generated_image = images[0] break - if hasattr(req_output, "request_output") and req_output.request_output: - stage_outputs = req_output.request_output + if isinstance(req_output, OmniRequestOutput) and req_output: + stage_outputs = req_output if not isinstance(stage_outputs, list): stage_outputs = [stage_outputs] for stage_out in stage_outputs: @@ -250,8 +250,8 @@ def _generate_bagel_image( # Check LLM stage output — should have finish_reason=stop (not length) for req_output in omni_outputs: - if hasattr(req_output, "request_output") and req_output.request_output: - stage_outputs = req_output.request_output + if isinstance(req_output, OmniRequestOutput) and req_output: + stage_outputs = req_output if not isinstance(stage_outputs, list): stage_outputs = [stage_outputs] for stage_out in stage_outputs: diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index b704bca1d..fdcef37b5 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -310,7 +310,7 @@ def _generate_image(omni, config: QualityTestConfig): first = outputs[0] if hasattr(first, "images") and first.images: return first.images[0], peak_mem - inner = first.request_output + inner = first if inner is not None and hasattr(inner, "images") and inner.images: return inner.images[0], peak_mem raise ValueError("Could not extract image from output.") @@ -341,15 +341,15 @@ def _generate_video(omni, config: QualityTestConfig): ) peak_mem = torch.accelerator.max_memory_allocated() / (1024**3) - first = outputs[0] - if hasattr(first, "request_output") and isinstance(first.request_output, list): - inner = first.request_output[0] - if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"): - frames = inner.images[0] if inner.images else None - else: - frames = inner - elif hasattr(first, "images") and first.images: - frames = first.images[0] + if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput): + first = outputs[0] + elif isinstance(outputs, OmniRequestOutput): + first = outputs + else: + raise ValueError("Could not extract video frames from output.") + + if hasattr(first, "images"): + frames = first.images[0] if first.images else None else: raise ValueError("Could not extract video frames from output.") @@ -444,11 +444,11 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch image = Image.new("RGB", (2, 2)) inner = OmniRequestOutput.from_diffusion(request_id="req", images=[image]) - outer = OmniRequestOutput( + outer = OmniRequestOutput.from_stage_output( + inner, request_id="req", stage_id=0, final_output_type="image", - request_output=inner, finished=True, ) diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py index 963ed5685..1561764fa 100644 --- a/tests/diffusion/quantization/test_trajectory_similarity_tool.py +++ b/tests/diffusion/quantization/test_trajectory_similarity_tool.py @@ -73,11 +73,9 @@ def test_run_summary_reports_worker_peak_memory(): assert summary["per_run_peak_memory_mb"] == [100.0, 150.0] -def test_request_peak_memory_prefers_inner_diffusion_output(): - inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})() - outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})() - - assert _request_peak_memory_mb(outer) == 321.0 +def test_request_peak_memory_reads_peak_memory_mb_directly(): + result = type("Output", (), {"peak_memory_mb": 123.0})() + assert _request_peak_memory_mb(result) == 123.0 def test_metric_guidance_describes_thresholds(): diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py index 824902ad4..d84f49b31 100644 --- a/tests/diffusion/test_diffusion_streaming_output.py +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -273,7 +273,7 @@ class TestPipelineStreamingOutputToEntrypoint: await _wait_for(lambda: len(pipeline.requests) == 1) assert pipeline.requests[0].request_id.startswith("req-omni-") assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1] - assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True] + assert [output.finished for output in outputs] == [False, True] finally: await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client) diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py index fef63f4f8..57997905c 100644 --- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py @@ -25,6 +25,7 @@ from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.omni import Omni +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")] @@ -95,8 +96,8 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None: for req_output in omni_outputs: if images := getattr(req_output, "images", None): return images[0] - if hasattr(req_output, "request_output") and req_output.request_output: - stage_out = req_output.request_output + if isinstance(req_output, OmniRequestOutput) and req_output: + stage_out = req_output if hasattr(stage_out, "images") and stage_out.images: return stage_out.images[0] return None diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py index 0545e9e3f..b97161335 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py @@ -24,6 +24,7 @@ from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.omni import Omni +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")] @@ -122,8 +123,8 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None: for req_output in omni_outputs: if images := getattr(req_output, "images", None): return images[0] - if hasattr(req_output, "request_output") and req_output.request_output: - stage_out = req_output.request_output + if isinstance(req_output, OmniRequestOutput) and req_output: + stage_out = req_output if hasattr(stage_out, "images") and stage_out.images: return stage_out.images[0] return None diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index f28ac16b9..eed0ef6e8 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -328,7 +328,7 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag images = None cot_text = "" for out in outputs: - ro = getattr(out, "request_output", None) + ro = out if ro and getattr(ro, "outputs", None): cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs) if not cot_text: @@ -467,7 +467,7 @@ def _extract_image(outputs) -> Image.Image: assert outputs, "Pipeline produced no outputs" for output in outputs: images = getattr(output, "images", None) - request_output = getattr(output, "request_output", None) + request_output = output if not images and request_output is not None: images = getattr(request_output, "images", None) if images: diff --git a/tests/e2e/accuracy/test_videogen_modelopt_quant.py b/tests/e2e/accuracy/test_videogen_modelopt_quant.py index e91c15175..3faa6737d 100644 --- a/tests/e2e/accuracy/test_videogen_modelopt_quant.py +++ b/tests/e2e/accuracy/test_videogen_modelopt_quant.py @@ -141,8 +141,8 @@ def _extract_frames(result) -> list[Image.Image]: if isinstance(result, list): result = result[0] if result else None - if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None: - result = result.request_output + if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result is not None: + result = result images = getattr(result, "images", None) if not images: raise AssertionError("Pipeline output had no images") diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py index e8d68b159..5dcb9c449 100644 --- a/tests/e2e/features/comfyui/test_comfyui_integration.py +++ b/tests/e2e/features/comfyui/test_comfyui_integration.py @@ -171,11 +171,11 @@ def _build_text_output(text: str = "This is a test response.") -> OmniRequestOut metrics=None, lora_request=None, ) - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + request_output, request_id="test_req_text", finished=True, final_output_type="text", - request_output=request_output, ) @@ -200,11 +200,11 @@ def _build_audio_chat_output(num_samples: int = 24000) -> OmniRequestOutput: metrics=None, lora_request=None, ) - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + request_output, request_id="test_req_audio_chat", finished=True, final_output_type="audio", - request_output=request_output, ) @@ -240,11 +240,11 @@ def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput: images=[_build_image_output(color="blue")], finished=True, ) - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + request_output, request_id="test_req_img_chat", - finished=True, + images=[_build_image_output(color="blue")], final_output_type="image", - request_output=request_output, ) @@ -515,8 +515,19 @@ def mock_async_omni( @pytest.fixture -def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni): +def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni, tmp_path): """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni.""" + # Override the STORAGE_MANAGER path to a writable temp directory before + # forking the server subprocess. The default /tmp/storage may be owned + # by a different user (e.g. from a previous CI run) and cause + # PermissionError when the server tries to save generated video files. + from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER + + old_storage_path = STORAGE_MANAGER.storage_path + new_storage_path = tmp_path / "omni_storage" + new_storage_path.mkdir(exist_ok=True) + STORAGE_MANAGER.storage_path = str(new_storage_path) + parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") cmd = OmniServeCommand() @@ -555,6 +566,9 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni yield f"http://127.0.0.1:{port}/v1" + # Restore the original storage path after the test. + STORAGE_MANAGER.storage_path = old_storage_path + if server_process.is_alive(): server_process.terminate() server_process.join(timeout=10) diff --git a/tests/e2e/features/helpers/verl_omni_server.py b/tests/e2e/features/helpers/verl_omni_server.py index d53e0936f..a38a3a12e 100644 --- a/tests/e2e/features/helpers/verl_omni_server.py +++ b/tests/e2e/features/helpers/verl_omni_server.py @@ -294,7 +294,7 @@ class vLLMOmniHttpServerLocal: "global_steps": self.global_steps, } - req_output = final_res.request_output + req_output = final_res if req_output is not None and hasattr(req_output, "finish_reason"): finish_reason = req_output.finish_reason or "stop" else: diff --git a/tests/e2e/offline_inference/test_audiox_model_expansion.py b/tests/e2e/offline_inference/test_audiox_model_expansion.py index d4b2f5223..af944b144 100644 --- a/tests/e2e/offline_inference/test_audiox_model_expansion.py +++ b/tests/e2e/offline_inference/test_audiox_model_expansion.py @@ -53,9 +53,9 @@ def test_audiox_model(omni_runner: OmniRunner) -> None: assert outputs is not None first_output = outputs[0] assert first_output.final_output_type == "audio" - assert hasattr(first_output, "request_output") and first_output.request_output + assert isinstance(first_output, OmniRequestOutput) and first_output - req_out = first_output.request_output + req_out = first_output assert isinstance(req_out, OmniRequestOutput) assert req_out.final_output_type == "audio" assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output diff --git a/tests/e2e/offline_inference/test_bagel.py b/tests/e2e/offline_inference/test_bagel.py index 6d495a156..9ba7bdac4 100644 --- a/tests/e2e/offline_inference/test_bagel.py +++ b/tests/e2e/offline_inference/test_bagel.py @@ -55,7 +55,7 @@ pytestmark = [ def _extract_text(omni_outputs: list) -> str: """Extract generated text from OmniRequestOutput list.""" for req_output in omni_outputs: - ro = getattr(req_output, "request_output", None) + ro = req_output if ro and getattr(ro, "outputs", None): return "".join(getattr(o, "text", "") or "" for o in ro.outputs) return "" diff --git a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py index 49079266d..fa1bdfead 100644 --- a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py +++ b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py @@ -28,7 +28,6 @@ from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer -from vllm_omni.outputs import OmniRequestOutput from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" @@ -176,12 +175,12 @@ def test_cosyvoice3_offline_reference_zero_shot(omni_runner: OmniRunner) -> None assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})" # Code2wav is ``final_output`` (stage 1); ``Omni.generate`` yields its ``request_output`` / completions. - pipeline_out = OmniRequestOutput.unwrap_result(outputs[0]) + pipeline_out = outputs[0] assert pipeline_out.stage_id == 1, f"expected final stage 1, got {pipeline_out.stage_id}" assert pipeline_out.final_output_type == "audio", ( f"expected audio final_output_type, got {pipeline_out.final_output_type!r}" ) - ro = pipeline_out.request_output + ro = pipeline_out assert ro is not None, "stage 1 should include request_output" completions = getattr(ro, "outputs", None) or [] assert completions, "stage 1 request_output should include at least one completion" diff --git a/tests/e2e/offline_inference/test_dynin_omni_expansion.py b/tests/e2e/offline_inference/test_dynin_omni_expansion.py index 19a847332..922518557 100644 --- a/tests/e2e/offline_inference/test_dynin_omni_expansion.py +++ b/tests/e2e/offline_inference/test_dynin_omni_expansion.py @@ -175,7 +175,7 @@ def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None: # Prefer the latest finished chunk to avoid picking an intermediate stream output. for stage_output in reversed(matched): - if _is_finished_request_output(getattr(stage_output, "request_output", None)): + if _is_finished_request_output(stage_output): return stage_output return matched[-1] @@ -205,7 +205,7 @@ def _to_token_list(value: Any) -> list[int]: def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str: - request_output = getattr(stage_output, "request_output", None) + request_output = stage_output if request_output is None: return "" req_list = request_output if isinstance(request_output, list) else [request_output] @@ -240,7 +240,7 @@ def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str: def _extract_audio(stage_output: Any) -> Any | None: - request_output = getattr(stage_output, "request_output", None) + request_output = stage_output if request_output is None: return None req_list = request_output if isinstance(request_output, list) else [request_output] @@ -256,7 +256,7 @@ def _extract_audio(stage_output: Any) -> Any | None: def _extract_image(stage_output: Any) -> Any | None: - request_output = getattr(stage_output, "request_output", None) + request_output = stage_output if request_output is None: return None req_list = request_output if isinstance(request_output, list) else [request_output] diff --git a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py index 4e9e30712..ec48a49f7 100644 --- a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py @@ -54,7 +54,7 @@ def _sampling_params() -> OmniDiffusionSamplingParams: def _first_request_images(outputs) -> list: first_output = outputs[0] assert first_output.final_output_type == "image" - req_out = first_output.request_output + req_out = first_output assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") return req_out.images diff --git a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py index 2ea4dc641..28eaab9b3 100644 --- a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py @@ -161,7 +161,7 @@ def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunn first_output = outputs[0] assert first_output.final_output_type == "image" - req_out = first_output.request_output + req_out = first_output assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") images = req_out.images @@ -224,7 +224,7 @@ def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunne first_output = outputs[0] assert first_output.final_output_type == "image" - req_out = first_output.request_output + req_out = first_output assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") images = req_out.images diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py index ebd338fe8..97179d963 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3.py @@ -288,7 +288,7 @@ def _extract_generated_image(outputs: list[object]) -> Image.Image: if images := getattr(first_output, "images", None): return images[0] - request_output = getattr(first_output, "request_output", None) + request_output = first_output if request_output is not None and (images := getattr(request_output, "images", None)): return images[0] diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py index 008f0d2a9..4bdb42c34 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py @@ -74,7 +74,7 @@ def test_i2t_generates_text(omni: Omni) -> None: outputs = omni.generate(prompts=[prompt_dict]) assert outputs, "No outputs returned from Omni.generate()" - request_output = outputs[0].request_output + request_output = outputs[0] assert request_output.outputs, "No completion outputs" completion = request_output.outputs[0] diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py index 04308c4a7..ec1545f6a 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py @@ -172,7 +172,7 @@ def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner): found_image = False for out in outputs: - ro_list = getattr(out, "request_output", out) + ro_list = out if not isinstance(ro_list, list): ro_list = [ro_list] for ro in ro_list: diff --git a/tests/e2e/offline_inference/test_ming_tts_expansion.py b/tests/e2e/offline_inference/test_ming_tts_expansion.py index 82da74227..8dcefd182 100644 --- a/tests/e2e/offline_inference/test_ming_tts_expansion.py +++ b/tests/e2e/offline_inference/test_ming_tts_expansion.py @@ -137,14 +137,11 @@ def _extract_final_audio_outputs(outputs): if getattr(item, "final_output_type", None) == "audio": final_outputs.append(item) continue - request_output = getattr(item, "request_output", None) - if request_output is None: - continue - multimodal_output = getattr(request_output, "multimodal_output", None) + multimodal_output = getattr(item, "multimodal_output", None) if isinstance(multimodal_output, Mapping): final_outputs.append(item) continue - completions = getattr(request_output, "outputs", None) or [] + completions = getattr(item, "outputs", None) or [] if any(isinstance(getattr(completion, "multimodal_output", None), Mapping) for completion in completions): final_outputs.append(item) return final_outputs @@ -155,16 +152,11 @@ def _extract_multimodal_output(output) -> Mapping[str, Any]: if isinstance(multimodal_output, Mapping): return multimodal_output - request_output = getattr(output, "request_output", None) - if request_output is not None: - multimodal_output = getattr(request_output, "multimodal_output", None) + completions = getattr(output, "outputs", None) or [] + for completion in completions: + multimodal_output = getattr(completion, "multimodal_output", None) if isinstance(multimodal_output, Mapping): return multimodal_output - completions = getattr(request_output, "outputs", None) or [] - for completion in completions: - multimodal_output = getattr(completion, "multimodal_output", None) - if isinstance(multimodal_output, Mapping): - return multimodal_output raise AssertionError("No multimodal audio output found in Ming generate results") diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py index d3e7fcb0f..b894fe4ea 100644 --- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py @@ -108,7 +108,7 @@ def _build_request( def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: """Run a single request and return (waveform, sample_rate).""" for stage_outputs in omni.generate(request, DEFAULT_SAMPLING): - req_output = stage_outputs.request_output + req_output = stage_outputs if req_output is not None: mm = req_output.outputs[0].multimodal_output assert mm is not None, "Expected multimodal_output to be non-None" @@ -167,7 +167,7 @@ def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None: results = [] # Single-stage model (num_stages=1): one sampling param for all requests. for stage_outputs in omni_runner.omni.generate(requests, [DEFAULT_SAMPLING]): - req_output = stage_outputs.request_output + req_output = stage_outputs if req_output is not None: mm = req_output.outputs[0].multimodal_output assert mm is not None diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py index e3a138112..5b4c7a935 100644 --- a/tests/e2e/offline_inference/test_omnivoice_expansion.py +++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py @@ -58,7 +58,7 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None: # Check final output has audio final_output = outputs[-1] - ro = final_output.request_output + ro = final_output assert ro is not None, "No request_output" mm = getattr(ro, "multimodal_output", None) diff --git a/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py index 8d7d31882..e7b4d061e 100644 --- a/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py @@ -42,7 +42,7 @@ def _sampling_params() -> OmniDiffusionSamplingParams: def _first_request_images(outputs) -> list: first_output = outputs[0] assert first_output.final_output_type == "image" - req_out = first_output.request_output + req_out = first_output assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") return req_out.images diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index fc3b5596d..00ecd09d9 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py @@ -62,9 +62,9 @@ def generate_stable_audio_short_clip( # Audio-output diffusion pipelines (those with ``support_audio_output = True``) now have # ``final_output_type="audio"`` set on the outer stage metadata as well as the inner request. assert first_output.final_output_type == "audio" - assert hasattr(first_output, "request_output") and first_output.request_output + assert isinstance(first_output, OmniRequestOutput) and first_output - req_out = first_output.request_output + req_out = first_output assert isinstance(req_out, OmniRequestOutput) assert req_out.final_output_type == "audio" assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output diff --git a/tests/e2e/offline_inference/test_step_audio2_expansion.py b/tests/e2e/offline_inference/test_step_audio2_expansion.py index cac4fbe52..483396b23 100644 --- a/tests/e2e/offline_inference/test_step_audio2_expansion.py +++ b/tests/e2e/offline_inference/test_step_audio2_expansion.py @@ -63,8 +63,8 @@ def _assert_non_empty_text_output(outputs: list[OmniRequestOutput]) -> None: text_output = next((o for o in outputs if o.final_output_type == "text"), None) assert text_output is not None - assert text_output.request_output is not None - text_content = text_output.request_output.outputs[0].text + assert text_output is not None + text_content = text_output.outputs[0].text assert text_content is not None assert len(text_content.strip()) > 0 @@ -81,8 +81,8 @@ def test_audio_to_text_and_audio(omni_runner: OmniRunner) -> None: audio_output = next((o for o in outputs if o.final_output_type == "audio"), None) if audio_output is not None: - assert audio_output.request_output is not None - audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio") + assert audio_output is not None + audio_tensor = audio_output.outputs[0].multimodal_output.get("audio") if audio_tensor is not None: assert audio_tensor.numel() > 0 diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index 87d635807..b2e40c58f 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -65,7 +65,7 @@ def test_realtime_duplex_multi_session_script_is_directly_executable(): cwd=MULTI_DEMO_PATH.parents[3], capture_output=True, text=True, - timeout=20, + timeout=60, check=False, ) @@ -78,7 +78,7 @@ def test_realtime_duplex_demo_pair_script_is_directly_executable(): cwd=PAIR_DEMO_PATH.parents[3], capture_output=True, text=True, - timeout=20, + timeout=60, check=False, ) @@ -91,7 +91,7 @@ def test_realtime_duplex_soft_interrupt_script_is_directly_executable(): cwd=SOFT_INTERRUPT_DEMO_PATH.parents[3], capture_output=True, text=True, - timeout=20, + timeout=60, check=False, ) diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py index dfe00732f..40f89df84 100644 --- a/tests/engine/test_async_omni_engine_abort.py +++ b/tests/engine/test_async_omni_engine_abort.py @@ -50,7 +50,7 @@ async def generate( output_modalities=["text"], ): stage_id = omni_output.stage_id - out = omni_output.request_output + out = omni_output if stage_id == 0: num_tokens = sum(len(output.token_ids) for output in out.outputs) count = num_tokens diff --git a/tests/engine/test_omni_request_output.py b/tests/engine/test_omni_request_output.py index 8da19a498..fa25bba9e 100644 --- a/tests/engine/test_omni_request_output.py +++ b/tests/engine/test_omni_request_output.py @@ -4,171 +4,269 @@ import pytest from PIL import Image -from pytest_mock import MockerFixture +from vllm.outputs import CompletionOutput, RequestOutput from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -class TestOmniRequestOutput: - """Tests for OmniRequestOutput class.""" +# ── helpers ──────────────────────────────────────────────────────────────── - def test_from_diffusion(self): - """Test creating output from diffusion model.""" - images = [Image.new("RGB", (64, 64), color="red")] - output = OmniRequestOutput.from_diffusion( - request_id="test-123", - images=images, - prompt="a cat", + +def _make_text_request_output( + request_id: str = "req-1", + text: str = "hello world", + token_ids: list[int] | None = None, + finished: bool = True, +) -> RequestOutput: + """Build a minimal vLLM RequestOutput for testing.""" + return RequestOutput( + request_id=request_id, + prompt="test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=text, + token_ids=token_ids or [10, 11, 12], + cumulative_logprob=-1.0, + logprobs=None, + finish_reason="stop" if finished else None, + stop_reason=None, + ) + ], + finished=finished, + metrics=None, + lora_request=None, + ) + + +# ── from_stage_output ────────────────────────────────────────────────────── + + +class TestFromStageOutput: + """Tests for OmniRequestOutput.from_stage_output().""" + + def test_copies_request_output_content(self): + """Content from a vLLM RequestOutput is flattened onto the new object.""" + source = _make_text_request_output(text="copied text") + out = OmniRequestOutput.from_stage_output(source, request_id="req-1", stage_id=0, final_output_type="text") + + assert out.prompt == "test prompt" + assert out.prompt_token_ids == [1, 2, 3] + assert len(out.outputs) == 1 + assert out.outputs[0].text == "copied text" + assert out.finished is True + assert out.stage_id == 0 + assert out.final_output_type == "text" + + def test_copies_omni_content(self): + """Diffusion fields are copied when source is another OmniRequestOutput.""" + image = Image.new("RGB", (64, 64), color="blue") + inner = OmniRequestOutput.from_diffusion(request_id="req-diff", images=[image], prompt="a cat") + outer = OmniRequestOutput.from_stage_output(inner, request_id="req-diff", stage_id=0, final_output_type="image") + + assert len(outer.images) == 1 + assert outer.images[0] is image + assert outer.prompt == "a cat" + assert outer.is_diffusion_output + + def test_with_duck_typed_source(self): + """A duck-typed source (SimpleNamespace-style) copies correctly.""" + from types import SimpleNamespace + + source = SimpleNamespace( + outputs=[], + prompt="duck prompt", + prompt_token_ids=[7, 8, 9], + finished=False, + ) + out = OmniRequestOutput.from_stage_output(source, request_id="duck-1", stage_id=2) + + assert out.prompt == "duck prompt" + assert out.prompt_token_ids == [7, 8, 9] + assert out.finished is False + assert out.stage_id == 2 + + def test_request_id_always_copied_from_source(self): + """request_id is always copied from source, even when passed in kwargs.""" + source = _make_text_request_output(request_id="source-id") + out = OmniRequestOutput.from_stage_output(source, request_id="kwargs-id", stage_id=1) + # Source wins — request_id is content, not a control field. + assert out.request_id == "source-id" + + def test_passes_control_fields(self): + """Control fields (metrics, stage_durations, peak_memory_mb) are set.""" + source = _make_text_request_output() + out = OmniRequestOutput.from_stage_output( + source, + request_id="r1", + stage_id=3, + metrics={"ttft": 0.5}, + stage_durations={"prefill": 0.12}, + peak_memory_mb=1024.0, + ) + + assert out.metrics == {"ttft": 0.5} + assert out.stage_durations == {"prefill": 0.12} + assert out.peak_memory_mb == 1024.0 + + def test_finished_field_copied(self): + """The *finished* flag is copied from the source when present.""" + source = _make_text_request_output(finished=False) + out = OmniRequestOutput.from_stage_output(source, request_id="r1") + assert out.finished is False + + source2 = _make_text_request_output(finished=True) + out2 = OmniRequestOutput.from_stage_output(source2, request_id="r2") + assert out2.finished is True + + def test_outputs_field_copied(self): + """The *outputs* list is copied from the source.""" + source = _make_text_request_output(text="hello", token_ids=[1, 2, 3]) + out = OmniRequestOutput.from_stage_output(source, request_id="r1") + + assert len(out.outputs) == 1 + assert out.outputs[0].text == "hello" + assert out.outputs[0].token_ids == [1, 2, 3] + + +# ── dataclass safety ──────────────────────────────────────────────────────── + + +class TestDataclassSafety: + def test_from_stage_output_is_classmethod(self): + """from_stage_output must be a classmethod.""" + assert callable(OmniRequestOutput.from_stage_output) + + +# ── msgpack round-trip ───────────────────────────────────────────────────── + + +class TestMsgpackRoundTrip: + """Tests that OmniRequestOutput survives msgpack encode → decode without + infinite recursion or data loss.""" + + def test_round_trip_pipeline_output(self): + """A pipeline output round-trips without data loss.""" + from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackDecoder, + OmniMsgpackEncoder, + ) + + source = _make_text_request_output(text="round trip") + out = OmniRequestOutput.from_stage_output( + source, + request_id="rt-1", + stage_id=0, + final_output_type="text", + ) + + encoder = OmniMsgpackEncoder() + decoder = OmniMsgpackDecoder() + + encoded = encoder.encode(out) + decoded = decoder.decode(encoded) + + assert isinstance(decoded, OmniRequestOutput) + assert decoded.request_id == "req-1" + assert decoded.stage_id == 0 + assert decoded.final_output_type == "text" + assert decoded.finished is True + assert decoded.outputs[0].text == "round trip" + assert decoded.prompt == "test prompt" + assert decoded.prompt_token_ids == [1, 2, 3] + + def test_round_trip_diffusion_output(self): + """A diffusion output round-trips without data loss.""" + from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackDecoder, + OmniMsgpackEncoder, + ) + + image = Image.new("RGB", (64, 64), color="red") + out = OmniRequestOutput.from_diffusion( + request_id="diff-1", + images=[image], + prompt="a red image", metrics={"steps": 50}, - ) - assert output.request_id == "test-123" - assert output.images == images - assert output.prompt == "a cat" - assert output.metrics == {"steps": 50} - assert output.is_diffusion_output - assert output.num_images == 1 - - def test_from_pipeline(self, mocker: MockerFixture): - """Test creating output from pipeline stage.""" - mock_request_output = mocker.Mock() - mock_request_output.request_id = "pipeline-123" - mock_request_output.prompt_token_ids = [1, 2, 3] - mock_request_output.outputs = [mocker.Mock()] - mock_request_output.encoder_prompt_token_ids = None - mock_request_output.prompt_logprobs = None - mock_request_output.num_cached_tokens = 10 - mock_request_output.kv_transfer_params = None - mock_request_output.multimodal_output = {"image": mocker.Mock()} - - output = OmniRequestOutput.from_pipeline( - stage_id=0, - final_output_type="text", - request_output=mock_request_output, + multimodal_output={"audio": None}, ) - assert output.request_id == "pipeline-123" - assert output.stage_id == 0 - assert output.final_output_type == "text" - assert output.is_pipeline_output + encoder = OmniMsgpackEncoder() + decoder = OmniMsgpackDecoder() - def test_prompt_token_ids_property(self, mocker: MockerFixture): - """Test prompt_token_ids property for streaming compatibility.""" - mock_request_output = mocker.Mock() - mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5] + encoded = encoder.encode(out) + decoded = decoder.decode(encoded) - output = OmniRequestOutput.from_pipeline( - stage_id=0, - final_output_type="text", - request_output=mock_request_output, + assert isinstance(decoded, OmniRequestOutput) + assert decoded.request_id == "diff-1" + assert decoded.prompt == "a red image" + assert decoded.num_images == 1 + assert decoded.images[0].size == (64, 64) + assert decoded.metrics == {"steps": 50} + + def test_round_trip_does_not_recurse(self): + """Encoding a pipeline output must not trigger infinite recursion.""" + from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackEncoder, ) - assert output.prompt_token_ids == [1, 2, 3, 4, 5] + source = _make_text_request_output(text="no recursion") + out = OmniRequestOutput.from_stage_output(source, request_id="rec-1", stage_id=0) - def test_prompt_token_ids_none_when_no_request_output(self): - """Test prompt_token_ids returns None when no request_output.""" - output = OmniRequestOutput.from_diffusion( - request_id="test-123", - images=[], - prompt="a cat", - ) - assert output.prompt_token_ids is None + encoder = OmniMsgpackEncoder() + encoded = encoder.encode(out) + assert isinstance(encoded, bytes) + assert len(encoded) > 0 - def test_outputs_property(self, mocker: MockerFixture): - """Test outputs property for chat completion compatibility.""" - mock_output = mocker.Mock() - mock_request_output = mocker.Mock() - mock_request_output.outputs = [mock_output] - - output = OmniRequestOutput.from_pipeline( - stage_id=0, - final_output_type="text", - request_output=mock_request_output, + def test_legacy_wire_format_deserialized(self): + """When the wire format contains a nested ``request_output`` key + (pre-inheritance legacy), it is merged via _copy_content_from.""" + from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackDecoder, ) - assert output.outputs == [mock_output] + # Build a legacy-style dict: OmniRequestOutput fields + a nested + # request_output dict containing RequestOutput-like content. + legacy_dict = { + "request_id": "legacy-1", + "finished": True, + "stage_id": 1, + "final_output_type": "text", + "images": [], + "_multimodal_output": {}, + "_custom_output": {}, + "stage_durations": {}, + "peak_memory_mb": 0.0, + # --- nested request_output key (legacy) --- + "request_output": { + "request_id": "legacy-1", + "prompt": "legacy prompt", + "prompt_token_ids": [1, 2], + "prompt_logprobs": None, + "outputs": [], + "finished": True, + "metrics": {}, + "lora_request": None, + "encoder_prompt": None, + "encoder_prompt_token_ids": None, + "num_cached_tokens": None, + "kv_transfer_params": None, + }, + } - def test_outputs_empty_when_no_request_output(self): - """Test outputs returns empty list when no request_output.""" - output = OmniRequestOutput.from_diffusion( - request_id="test-123", - images=[], - prompt="a cat", - ) - assert output.outputs == [] + # Manually run the decoder's post_process to test legacy path. + # The decoder's _decode_omni_request_output pops "request_output" + # and merges it via _copy_content_from. + decoder = OmniMsgpackDecoder() + result = decoder._post_process(legacy_dict) - def test_encoder_prompt_token_ids_property(self, mocker: MockerFixture): - """Test encoder_prompt_token_ids property.""" - mock_request_output = mocker.Mock() - mock_request_output.encoder_prompt_token_ids = [10, 20, 30] - - output = OmniRequestOutput.from_pipeline( - stage_id=0, - final_output_type="text", - request_output=mock_request_output, - ) - - assert output.encoder_prompt_token_ids == [10, 20, 30] - - def test_num_cached_tokens_property(self, mocker: MockerFixture): - """Test num_cached_tokens property.""" - mock_request_output = mocker.Mock() - mock_request_output.num_cached_tokens = 42 - - output = OmniRequestOutput.from_pipeline( - stage_id=0, - final_output_type="text", - request_output=mock_request_output, - ) - - assert output.num_cached_tokens == 42 - - def test_multimodal_output_property(self, mocker: MockerFixture): - """Test multimodal_output property.""" - mock_request_output = mocker.Mock() - mock_audio = mocker.Mock() - expected_output = {"audio": mock_audio} - mock_request_output.outputs = [] - mock_request_output.multimodal_output = expected_output - - output = OmniRequestOutput.from_pipeline( - stage_id=0, - final_output_type="audio", - request_output=mock_request_output, - ) - - assert output.multimodal_output is expected_output - - def test_to_dict_diffusion(self): - """Test to_dict for diffusion output.""" - output = OmniRequestOutput.from_diffusion( - request_id="test-123", - images=[Image.new("RGB", (64, 64), color="red")], - prompt="a cat", - metrics={"steps": 50}, - ) - result = output.to_dict() - - assert result["request_id"] == "test-123" - assert result["finished"] is True - assert result["final_output_type"] == "image" - assert result["num_images"] == 1 - assert result["prompt"] == "a cat" - - def test_to_dict_pipeline(self, mocker: MockerFixture): - """Test to_dict for pipeline output.""" - mock_request_output = mocker.Mock() - mock_request_output.request_id = "pipeline-123" - - output = OmniRequestOutput.from_pipeline( - stage_id=0, - final_output_type="text", - request_output=mock_request_output, - ) - result = output.to_dict() - - assert result["request_id"] == "pipeline-123" - assert result["finished"] is True - assert result["final_output_type"] == "text" - assert result["stage_id"] == 0 + assert isinstance(result, OmniRequestOutput) + assert result.request_id == "legacy-1" + assert result.finished is True + # Content merged from legacy nested key + assert result.prompt == "legacy prompt" + assert result.prompt_token_ids == [1, 2] diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index 1ee454213..a37d9d2b1 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -2156,12 +2156,12 @@ def test_direct_listen_decision_survives_inner_completion_metadata(): ) assert decision is not None output = attach_duplex_output_decision( - OmniRequestOutput( + OmniRequestOutput.from_stage_output( + inner_output, request_id="duplex-direct-listen", finished=True, stage_id=0, final_output_type=decision.final_output_type, - request_output=inner_output, metrics={ "stage_metrics": { "0": { @@ -3049,12 +3049,12 @@ async def test_minicpmo_auto_response_listen_without_response_does_not_defer_com ) assert decision is not None listen_output = attach_duplex_output_decision( - OmniRequestOutput( + OmniRequestOutput.from_stage_output( + inner_output, request_id=request_id, finished=True, stage_id=0, final_output_type=decision.final_output_type, - request_output=inner_output, ), decision, ) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 19240c39c..2136d84f8 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -119,7 +119,6 @@ class MockGenerationResult: def __init__(self, images): self.images = images - self.request_output = SimpleNamespace(images=images) self.stage_durations = {} self.peak_memory_mb = 0.0 @@ -137,10 +136,7 @@ class MockStageResult: outputs = [SimpleNamespace(text=text, index=0)] else: outputs = [] - self.request_output = SimpleNamespace( - outputs=outputs, - images=self.images, - ) + self.outputs = outputs self.stage_durations = {} self.peak_memory_mb = 0.0 @@ -2080,22 +2076,20 @@ def test_extract_images_from_result(): assert all(isinstance(img, Image.Image) for img in images) assert all(img.size == (64, 64) for img in images) - # Test dict path: result.request_output["images"] + # Test result with "images" attribute set in __init__ class DictRequestOutput: def __init__(self): - self.request_output = {"images": [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]} + self.images = [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)] result = DictRequestOutput() images = _extract_images_from_result(result) assert len(images) == 1 assert isinstance(images[0], Image.Image) - # Test attribute path: result.request_output.images + # Test result with "images" attribute from an inner object class AttrRequestOutput: def __init__(self): - self.request_output = type( - "obj", (), {"images": [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]} - )() + self.images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)] result = AttrRequestOutput() images = _extract_images_from_result(result) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 32de8ac62..bcb89ff88 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -180,10 +180,10 @@ def create_mock_audio_output_for_test( audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples)) mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor) - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + mock_request_output, stage_id=0, final_output_type="audio", - request_output=mock_request_output, ) @@ -2292,10 +2292,10 @@ class TestStreamingResponse: self.prompt_logprobs = None self.kv_transfer_params = None - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + MockRequestOutput(audio_tensor=chunk), stage_id=0, final_output_type="audio", - request_output=MockRequestOutput(audio_tensor=chunk), finished=finished, ) @@ -3441,10 +3441,10 @@ class TestWAVStreaming: self.prompt_logprobs = None self.kv_transfer_params = None - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + MockRequestOutput(audio_tensor=chunk), stage_id=0, final_output_type="audio", - request_output=MockRequestOutput(audio_tensor=chunk), finished=finished, ) diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index fc8d83da6..bc80fe0db 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -47,7 +47,7 @@ def _text_result(text: str) -> OmniRequestOutput: output.text = text request_output = RequestOutput() request_output.outputs = [output] - return OmniRequestOutput(final_output_type="text", request_output=request_output) + return OmniRequestOutput.from_stage_output(request_output, final_output_type="text") def _audio_result(audio_data: Any) -> OmniRequestOutput: @@ -61,7 +61,7 @@ def _audio_result(audio_data: Any) -> OmniRequestOutput: output.multimodal_output = {"audio": audio_data} request_output = RequestOutput() request_output.outputs = [output] - return OmniRequestOutput(final_output_type="audio", request_output=request_output) + return OmniRequestOutput.from_stage_output(request_output, final_output_type="audio") class MockWebSocket: diff --git a/tests/entrypoints/test_async_omni_duplex.py b/tests/entrypoints/test_async_omni_duplex.py index 62f3f49ad..46bf668b3 100644 --- a/tests/entrypoints/test_async_omni_duplex.py +++ b/tests/entrypoints/test_async_omni_duplex.py @@ -480,11 +480,9 @@ def test_async_omni_duplex_direct_output_prefers_outer_control_metadata(): OmniRequestOutput( request_id="duplex-sid-e0-stage0", stage_id=0, - request_output=SimpleNamespace( - outputs=[], - multimodal_output=SimpleNamespace(kind="processed-payload"), - _custom_output={"inner": "completion"}, - ), + outputs=[], + _multimodal_output={"kind": "processed-payload"}, + _custom_output={"inner": "completion"}, ), decision, ) @@ -557,7 +555,7 @@ async def test_async_omni_duplex_collect_wraps_raw_response_stage_output(): assert outputs[0].finished is False assert outputs[0].stage_id == 1 assert outputs[0].final_output_type == "audio" - assert outputs[0].request_output is raw_stage1_output + assert outputs[0].peak_memory_mb == raw_stage1_output.peak_memory_mb def test_async_omni_duplex_request_info_includes_response_stage(): diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index d7f424d6e..161b3fe2a 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -72,8 +72,8 @@ def make_output_msg( final_output_type=final_output_type, images=images or [], stage_durations={}, + outputs=[SimpleNamespace(text=payload, index=0)], ) - engine_output.payload = payload return OutputMessage( request_id=request_id, stage_id=stage_id, @@ -474,7 +474,7 @@ async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.Mo app.shutdown() assert [output.stage_id for output in outputs] == [2] - assert [output.request_output.payload for output in outputs] == ["final"] + assert [output.outputs[0].text for output in outputs] == ["final"] assert "req-1" not in app.request_states @@ -492,7 +492,7 @@ async def test_async_omni_accepts_multiple_final_stage_streams(monkeypatch: pyte app.shutdown() assert [output.stage_id for output in outputs] == [0, 0, 0, 2, 2, 2] - assert [output.request_output.payload for output in outputs] == [ + assert [output.outputs[0].text for output in outputs] == [ "req-1-stage0-0", "req-1-stage0-1", "req-1-stage0-2", @@ -517,7 +517,7 @@ async def test_async_omni_stops_on_final_stage_finished(monkeypatch: pytest.Monk finally: app.shutdown() - assert [output.request_output.payload for output in outputs] == [ + assert [output.outputs[0].text for output in outputs] == [ "req-1-stage0", "req-1-stage2-final", ] @@ -544,7 +544,7 @@ async def test_async_omni_diffusion_only_yields_single_image_output(monkeypatch: assert outputs[0].stage_id == 0 assert outputs[0].final_output_type == "image" assert outputs[0].images == ["req-1-image"] - assert outputs[0].request_output.payload == "req-1-diffusion-final" + assert outputs[0].outputs[0].text == "req-1-diffusion-final" @pytest.mark.asyncio @@ -565,7 +565,7 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc assert [output.stage_id for output in outputs] == [0, 0, 0, 1] assert [output.final_output_type for output in outputs] == ["text", "text", "text", "image"] - assert [output.request_output.payload for output in outputs] == [ + assert [output.outputs[0].text for output in outputs] == [ "req-1-text-0", "req-1-text-1", "req-1-text-2", @@ -695,7 +695,7 @@ def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkey assert len(outputs) == 4 assert [output.stage_id for output in outputs] == [0, 2, 0, 2] - assert [output.request_output.payload for output in outputs] == [ + assert [output.outputs[0].text for output in outputs] == [ f"{engine.submitted[0]['request_id']}-stage0-0", f"{engine.submitted[0]['request_id']}-stage2-final", f"{engine.submitted[1]['request_id']}-stage0-0", @@ -739,7 +739,7 @@ def test_omni_generate_diffusion_only_yields_single_image_per_request(monkeypatc assert len(outputs) == 2 assert [output.stage_id for output in outputs] == [0, 0] assert [output.final_output_type for output in outputs] == ["image", "image"] - assert [output.request_output.payload for output in outputs] == [ + assert [output.outputs[0].text for output in outputs] == [ f"{engine.submitted[0]['request_id']}-diffusion-final", f"{engine.submitted[1]['request_id']}-diffusion-final", ] @@ -767,7 +767,7 @@ def test_omni_generate_llm_diffusion_yields_final_text_then_image_per_request( assert len(outputs) == 4 assert [output.stage_id for output in outputs] == [0, 1, 0, 1] assert [output.final_output_type for output in outputs] == ["text", "image", "text", "image"] - assert [output.request_output.payload for output in outputs] == [ + assert [output.outputs[0].text for output in outputs] == [ f"{engine.submitted[0]['request_id']}-text-0", f"{engine.submitted[0]['request_id']}-image-final", f"{engine.submitted[1]['request_id']}-text-0", @@ -1046,7 +1046,6 @@ def _enqueue_stage_error( if kill_engine: engine._alive = False engine_output = OmniRequestOutput.from_error(msg["request_id"], error_text) - engine_output.payload = "" engine.output_q.put_nowait( OutputMessage( request_id=msg["request_id"], diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 3d6052045..b02936838 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -69,10 +69,10 @@ def _make_text_omni_output( ], finished=finish_reason is not None, ) - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + res, request_id=request_id, final_output_type="text", - request_output=res, finished=finish_reason is not None, ) @@ -108,12 +108,12 @@ def _make_audio_omni_output( outputs=[completion], finished=True, ) - return OmniRequestOutput( + return OmniRequestOutput.from_stage_output( + res, request_id=request_id, stage_id=stage_id, replica_id=replica_id, final_output_type="audio", - request_output=res, finished=True, ) @@ -151,7 +151,7 @@ def _build_serving_chat(): ) instance._create_audio_choice = MagicMock( side_effect=lambda omni_res, role, request, stream=False: _mock_audio_choices( - index=omni_res.request_output.outputs[0].index, + index=omni_res.outputs[0].index, role=role, ) ) @@ -489,7 +489,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): request = _make_request(modalities=["text", "audio"]) empty_audio = _make_audio_omni_output() - empty_audio.request_output.outputs[0].multimodal_output = {"audio": []} + empty_audio.outputs[0].multimodal_output = {"audio": []} def create_audio_choice(omni_res, role, request, stream=False): return OmniOpenAIServingChat._create_audio_choice(serving_chat, omni_res, role, request, stream=stream) diff --git a/tests/helpers/media.py b/tests/helpers/media.py index 93e136f57..8eb82fc1f 100644 --- a/tests/helpers/media.py +++ b/tests/helpers/media.py @@ -26,10 +26,28 @@ from PIL import Image logger = logging.getLogger(__name__) +_synthetic_media_fallback_dir: Path | None = None + + def _resolve_synthetic_media_cache_dir(cache_dir: Path | str | None) -> Path: if cache_dir is not None: return Path(cache_dir).expanduser().resolve() - return Path(tempfile.gettempdir()) / "vllm_omni_test_synthetic_media" + + default = Path(tempfile.gettempdir()) / "vllm_omni_test_synthetic_media" + try: + default.mkdir(parents=True, exist_ok=True) + # Verify write access: the directory may exist but belong to + # another user (e.g. a previous CI job), causing PermissionError + # later when individual files are saved. + canary = default / ".write_test" + canary.touch() + canary.unlink() + return default + except (PermissionError, OSError): + global _synthetic_media_fallback_dir + if _synthetic_media_fallback_dir is None: + _synthetic_media_fallback_dir = Path(tempfile.mkdtemp(prefix="vllm_omni_test_synthetic_media_")) + return _synthetic_media_fallback_dir def _np_array_from_mp4_bytes(video_bytes: bytes) -> np.ndarray: diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 962075df2..82fcc308d 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -2938,9 +2938,9 @@ class OmniRunnerHandler: audio_content = None for stage_output in outputs: if getattr(stage_output, "final_output_type", None) == "text": - text_content = stage_output.request_output.outputs[0].text + text_content = stage_output.outputs[0].text if getattr(stage_output, "final_output_type", None) == "audio": - audio_content = stage_output.request_output.outputs[0].multimodal_output["audio"] + audio_content = stage_output.outputs[0].multimodal_output["audio"] result.audio_content = audio_content result.text_content = text_content result.success = True @@ -3080,7 +3080,7 @@ class OmniRunnerHandler: mm_out: dict[str, Any] | None = None for stage_out in outputs: if getattr(stage_out, "final_output_type", None) == "audio": - mm_out = stage_out.request_output.outputs[0].multimodal_output + mm_out = stage_out.outputs[0].multimodal_output break if mm_out is None: raise AssertionError("No audio output from pipeline") diff --git a/vllm_omni/diffusion/utils/image_output.py b/vllm_omni/diffusion/utils/image_output.py index 5f8c9903a..10d586894 100644 --- a/vllm_omni/diffusion/utils/image_output.py +++ b/vllm_omni/diffusion/utils/image_output.py @@ -51,8 +51,13 @@ def _iter_multimodal_image_payloads(output: OmniRequestOutput | RequestOutput) - if isinstance(output, OmniRequestOutput): if output.multimodal_output is not None: yield from _image_values_from_mapping_like(output.multimodal_output) - if output.request_output is not None: - yield from _iter_request_output_payloads(output.request_output) + # OmniRequestOutput IS the RequestOutput now; iterate completion + # outputs directly for multimodal payloads that the property + # (which returns only the first match) may have missed. + for completion in getattr(output, "outputs", []): + mm = getattr(completion, "multimodal_output", None) + if mm is not None: + yield from _image_values_from_mapping_like(mm) else: yield from _iter_request_output_payloads(output) diff --git a/vllm_omni/distributed/omni_connectors/utils/serialization.py b/vllm_omni/distributed/omni_connectors/utils/serialization.py index 1166ede19..60fc725d6 100644 --- a/vllm_omni/distributed/omni_connectors/utils/serialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/serialization.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping -from dataclasses import asdict, is_dataclass +from dataclasses import asdict, fields, is_dataclass from typing import Any import msgspec @@ -59,7 +59,10 @@ class OmniMsgpackEncoder: if isinstance(obj, Image.Image): return self._encode_pil_image(obj) - # RequestOutput (not a dataclass, needs special handling) + # RequestOutput (not a dataclass, needs special handling). + # Note: OmniRequestOutput subclasses RequestOutput but is a dataclass, + # so msgspec encodes it natively (field map) and it never reaches + # this hook. if isinstance(obj, RequestOutput): return self._encode_request_output(obj) @@ -148,7 +151,6 @@ class OmniMsgpackEncoder: "encoder_prompt": obj.encoder_prompt, "encoder_prompt_token_ids": obj.encoder_prompt_token_ids, "num_cached_tokens": obj.num_cached_tokens, - "multi_modal_placeholders": getattr(obj, "multi_modal_placeholders", None), "kv_transfer_params": obj.kv_transfer_params, } # Handle multimodal_output attribute (MultimodalPayload or dict) @@ -252,23 +254,27 @@ class OmniMsgpackDecoder: def _decode_omni_request_output(self, obj: dict[str, Any]) -> Any: """Decode dict to OmniRequestOutput. - OmniRequestOutput is a dataclass, so we can use msgspec.convert - or construct it directly. + By the time this runs, ``_post_process`` has already reconstructed + nested values bottom-up (CompletionOutput, tensors, images, ...), so + the object is built directly from the known dataclass fields. """ from vllm_omni.outputs import OmniRequestOutput + data = dict(obj) + # Legacy wire format (pre-inheritance): the stage output was nested + # under "request_output". Merge its content into the new object. + legacy_inner = data.pop("request_output", None) + field_names = {f.name for f in fields(OmniRequestOutput)} + known = {k: v for k, v in data.items() if k in field_names} try: - # Use msgspec.convert for dataclass reconstruction - return msgspec.convert(obj, OmniRequestOutput) + out = OmniRequestOutput(**known) except Exception: - try: - # Fallback: construct directly if msgspec.convert fails - # (e.g., if some fields are missing or have wrong types) - return OmniRequestOutput(**obj) - except Exception: - # If both attempts fail, return dict as-is (defensive fallback) - # This should rarely happen if _is_omni_request_output is correct - return obj + # Defensive fallback; should rarely happen if + # _is_omni_request_output is correct. + return obj + if legacy_inner is not None and not isinstance(legacy_inner, dict): + out._copy_content_from(legacy_inner) + return out def _decode_tensor(self, obj: dict[str, Any]) -> torch.Tensor: """Decode dict to torch.Tensor.""" @@ -314,19 +320,16 @@ class OmniMsgpackDecoder: We construct it manually using only the known __init__ parameters to avoid triggering the "Ignoring extra arguments" warning in vllm. Fields that are not part of RequestOutput.__init__ (e.g. - multi_modal_placeholders, multimodal_output) are extracted first and + multimodal_output) are extracted first and then restored as dynamic attributes after construction. """ # Extract dynamically-added / non-init fields before constructing so # they are not passed as unknown **kwargs to RequestOutput.__init__. mm_output = obj.pop("multimodal_output", None) - multi_modal_placeholders = obj.pop("multi_modal_placeholders", None) ro = RequestOutput(**obj) # Restore dynamic attributes that are not part of __init__. - if multi_modal_placeholders is not None: - setattr(ro, "multi_modal_placeholders", multi_modal_placeholders) if mm_output is not None: setattr(ro, "multimodal_output", mm_output) return ro diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 5646e82eb..901bdb980 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -1633,12 +1633,12 @@ class Orchestrator: from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision engine_output = attach_duplex_output_decision( - OmniRequestOutput( + OmniRequestOutput.from_stage_output( + output, request_id=req_id, finished=True, stage_id=stage_id, final_output_type=decision.final_output_type, - request_output=output, ), decision, ) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 587e0e73c..88e5a5edf 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -628,7 +628,6 @@ class OmniBase(PDDisaggregationMixin): self.prom_metrics.set_running(running) self.prom_metrics.set_waiting(max(0, total - running)) - images = getattr(engine_outputs, "images", []) if output_type == "image" else [] response_metrics: dict[str, Any] = {} stage_metrics: dict[str, dict[str, Any]] = {} rid_key = str(req_id) @@ -653,19 +652,16 @@ class OmniBase(PDDisaggregationMixin): response_metrics["final_output_type"] = current_stage_metrics["final_output_type"] response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"] response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"] - return OmniRequestOutput( + # Generation content (outputs, prompt, images, trajectory_*, ...) is + # copied from engine_outputs onto the returned object by + # OmniRequestOutput.from_stage_output(). + return OmniRequestOutput.from_stage_output( + engine_outputs, request_id=req_id or "", finished=finished, stage_id=stage_id, replica_id=result.replica_id, final_output_type=output_type, - request_output=engine_outputs, - images=images, - trajectory_latents=getattr(engine_outputs, "trajectory_latents", None), - trajectory_timesteps=getattr(engine_outputs, "trajectory_timesteps", None), - trajectory_log_probs=getattr(engine_outputs, "trajectory_log_probs", None), - trajectory_decoded=getattr(engine_outputs, "trajectory_decoded", None), - _custom_output=getattr(engine_outputs, "_custom_output", {}), metrics=response_metrics, stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 0ea0f3a3f..0a8daa055 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -2649,12 +2649,6 @@ def _extract_images_from_result(result: Any) -> list[Any]: images = [] if hasattr(result, "images") and result.images: images = result.images - elif hasattr(result, "request_output"): - request_output = result.request_output - if isinstance(request_output, dict) and request_output.get("images"): - images = request_output["images"] - elif hasattr(request_output, "images") and request_output.images: - images = request_output.images # Handle when generate more than one image if images and isinstance(images[0], np.ndarray) and images[0].shape[0] > 1 and images[0].ndim == 5: # Unwrap batch: (N, T, H, W, C) -> [img1, img2, ...] diff --git a/vllm_omni/entrypoints/openai/serving_audio_generate.py b/vllm_omni/entrypoints/openai/serving_audio_generate.py index f8fedd33f..594b8c137 100644 --- a/vllm_omni/entrypoints/openai/serving_audio_generate.py +++ b/vllm_omni/entrypoints/openai/serving_audio_generate.py @@ -119,9 +119,6 @@ class OmniOpenAIServingAudioGenerate(OpenAIServing, AudioMixin): audio_output = None if hasattr(final_output, "multimodal_output") and final_output.multimodal_output: audio_output = final_output.multimodal_output - if not audio_output and hasattr(final_output, "request_output"): - if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"): - audio_output = final_output.request_output.multimodal_output # Check for audio data using either "audio" or "model_outputs" key audio_key = None diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index ce019a970..423ca6d82 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -1468,7 +1468,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): try: async for omni_res in result_generator: final_output_type = omni_res.final_output_type - res = omni_res.request_output + res = omni_res if final_output_type not in first_iteration_dict: logger.warning(f"final output type: {final_output_type} is not needed by the request") continue @@ -2338,7 +2338,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): for omni_outputs in final_outputs: choices_data = [] - if omni_outputs.request_output is not None and not getattr(omni_outputs.request_output, "finished", False): + if not omni_outputs.finished: continue # Filter outputs based on requested modalites @@ -2347,7 +2347,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): continue if omni_outputs.final_output_type == "text": - if omni_outputs.request_output is not None: + if omni_outputs.stage_id is not None or omni_outputs.outputs: ( choices_data, usage, @@ -2362,7 +2362,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): role, reasoning_parser, ) - final_res = omni_outputs.request_output + final_res = omni_outputs else: # Diffusion pipeline text output (e.g. single-stage # img2text / text2text) — no AR request_output, so build @@ -2463,7 +2463,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): role: str, reasoning_parser: ReasoningParser | None = None, ): - final_res = omni_outputs.request_output + final_res = omni_outputs if self.tool_call_id_type == "kimi_k2": history_tool_call_cnt = get_history_tool_calls_cnt(conversation) else: @@ -2711,7 +2711,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): self, omni_outputs: OmniRequestOutput, role: str, request: ChatCompletionRequest, stream: bool = False ) -> list[ChatCompletionResponseChoice] | list[ChatCompletionResponseStreamChoice] | ErrorResponse: choices: list[ChatCompletionResponseChoice] = [] - final_res = omni_outputs.request_output + final_res = omni_outputs # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421 # The attribute is attached dynamically when stage audio arrives; fall @@ -2834,7 +2834,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): from PIL import Image choices: list[ChatCompletionResponseChoice] = [] - final_res = omni_outputs.request_output # Handle profiling data stage_durations = omni_outputs.stage_durations @@ -2846,10 +2845,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # First check omni_outputs.images directly (for diffusion mode via from_diffusion) if omni_outputs.images: images = omni_outputs.images - # Fall back to request_output for pipeline mode - # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput - elif final_res is not None and final_res.outputs: - completion_output = final_res.outputs[0] + # Fall back to completion outputs for pipeline mode (multimodal_output + # is attached to CompletionOutput by AR stages). + elif omni_outputs.outputs: + completion_output = omni_outputs.outputs[0] + final_res = omni_outputs if hasattr(completion_output, "multimodal_output") and completion_output.multimodal_output: image_data = completion_output.multimodal_output.get("image") if image_data is not None: @@ -3311,12 +3311,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): request_id=request_id, ) - images = getattr(result.request_output, "images", []) + images = getattr(result, "images", []) stage_durations = result.stage_durations peak_memory_mb = result.peak_memory_mb cot_output = None - req_out = getattr(result, "request_output", None) + req_out = result if req_out: prompt_obj = getattr(req_out, "prompt", None) if isinstance(prompt_obj, dict): @@ -3326,7 +3326,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if isinstance(ar_text, str) and ar_text.strip(): cot_output = ar_text - req_out = getattr(result, "request_output", None) + req_out = result if req_out: prompt_obj = getattr(req_out, "prompt", None) if isinstance(prompt_obj, dict): @@ -3336,7 +3336,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if isinstance(ar_text, str) and ar_text.strip(): cot_output = ar_text - req_out = getattr(result, "request_output", None) + req_out = result if req_out: prompt_obj = getattr(req_out, "prompt", None) if isinstance(prompt_obj, dict): @@ -3368,8 +3368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): stage_id = getattr(output, "stage_id", None) metrics = getattr(output, "metrics", None) if return_stage_metrics else None if final_output_type == "text" and stage_id == 0: - request_output = output.request_output - for completion in request_output.outputs: + for completion in output.outputs: text = completion.text or "" if not text: continue @@ -3382,7 +3381,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) yield f"data: {chunk.model_dump_json()}\n\n" elif final_output_type == "image": - images = self._flatten_diffusion_images(getattr(output.request_output, "images", [])) + images = self._flatten_diffusion_images(getattr(output, "images", [])) if not images: raise RuntimeError("Streaming image edit produced an empty final image output.") image_data = [ @@ -3686,8 +3685,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Image output path (text2img / img2img) final_output_type = getattr(result, "final_output_type", "image") - # Handle nested OmniRequestOutput structure where images might be in request_output - images = getattr(result.request_output, "images", []) + images = getattr(result, "images", []) multimodal_output = getattr(result, "multimodal_output", {}) or {} stage_durations = result.stage_durations peak_memory_mb = result.peak_memory_mb diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 88e7624f1..a662ae7de 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -2429,7 +2429,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): mm = getattr(res, "multimodal_output", None) ro = None if not mm: - ro = getattr(res, "request_output", None) + ro = res mm = getattr(ro, "multimodal_output", None) if ro else None if not mm: # MultimodalOutputProcessor attaches mm_accumulated on per-completion outputs. diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 3a25283c8..9d511ba6d 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -505,16 +505,6 @@ class OmniOpenAIServingVideo: videos = None if hasattr(result, "images") and result.images: videos = result.images - elif hasattr(result, "request_output"): - request_output = result.request_output - if isinstance(request_output, dict) and request_output.get("images"): - videos = request_output["images"] - elif hasattr(request_output, "images") and request_output.images: - videos = request_output.images - else: - request_multimodal_output = getattr(request_output, "multimodal_output", None) - if isinstance(request_multimodal_output, Mapping): - videos = request_multimodal_output.get("video") if videos is None: multimodal_output = getattr(result, "multimodal_output", None) if isinstance(multimodal_output, Mapping): @@ -534,16 +524,6 @@ class OmniOpenAIServingVideo: multimodal_output = getattr(result, "multimodal_output", None) if isinstance(multimodal_output, Mapping): audio = multimodal_output.get("audio") - elif hasattr(result, "request_output"): - request_output = result.request_output - if isinstance(request_output, dict) and request_output.get("multimodal_output"): - mm_output = request_output.get("multimodal_output") or {} - if isinstance(mm_output, Mapping): - audio = mm_output.get("audio") - else: - request_multimodal_output = getattr(request_output, "multimodal_output", None) - if isinstance(request_multimodal_output, Mapping): - audio = request_multimodal_output.get("audio") if audio is None: return [None] * expected_count @@ -599,7 +579,7 @@ class OmniOpenAIServingVideo: if isinstance(multimodal_output, Mapping): return dict(multimodal_output) - request_output = getattr(result, "request_output", None) + request_output = result if isinstance(request_output, dict): multimodal_output = request_output.get("multimodal_output") if multimodal_output is None: @@ -716,7 +696,7 @@ class OmniOpenAIServingVideo: except (TypeError, ValueError): pass - request_output = getattr(result, "request_output", None) + request_output = result if isinstance(request_output, dict): mm = request_output.get("multimodal_output") or {} if isinstance(mm, Mapping): @@ -755,7 +735,7 @@ class OmniOpenAIServingVideo: if sample_rate is not None: return sample_rate - request_output = getattr(result, "request_output", None) + request_output = result if isinstance(request_output, dict): multimodal_output = request_output.get("multimodal_output") or {} if isinstance(multimodal_output, Mapping): diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index 1923028e4..770890faf 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -671,17 +671,8 @@ class OmniStreamingVideoOutputHandler: if result.images: # Video frames may be list[list[Image.Image]] despite images: list[Image.Image]. videos = result.images # type: ignore[assignment] - elif result.request_output is not None: - request_output = result.request_output - if isinstance(request_output, dict) and request_output.get("images"): - videos = request_output["images"] - elif isinstance(request_output, OmniRequestOutput): - if request_output.images: - # Same nested-frames mismatch as result.images above. - videos = request_output.images # type: ignore[assignment] - else: - request_multimodal_output = request_output.multimodal_output - videos = request_multimodal_output.get("video") + # OmniRequestOutput IS the RequestOutput now — images and + # multimodal_output are checked directly on result. if videos is None: videos = result.multimodal_output.get("video") diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index 1694704f3..f975adbef 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -870,7 +870,7 @@ class OmniStreamingVideoHandler: @staticmethod def _get_audio_data(result: OmniRequestOutput): """Navigate OmniRequestOutput → multimodal_output['audio']. None on miss.""" - request_output = getattr(result, "request_output", None) + request_output = result if request_output is None: return None outputs = getattr(request_output, "outputs", None) @@ -1000,7 +1000,7 @@ class OmniStreamingVideoHandler: if result.final_output_type != "text": return "", previous_text - request_output = getattr(result, "request_output", None) + request_output = result if request_output is None: return "", previous_text diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index d0e41dbf0..105d0222d 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -3,7 +3,7 @@ from typing import Any import torch from PIL import Image -from vllm.outputs import RequestOutput +from vllm.outputs import CompletionOutput, RequestOutput from vllm.v1.outputs import ModelRunnerOutput from vllm_omni.inputs.data import OmniPromptType @@ -71,12 +71,55 @@ class OmniModelRunnerOutput(ModelRunnerOutput): ) +# RequestOutput attributes that hold generation content. When wrapping a +# stage output, these are copied onto the OmniRequestOutput itself so that +# the object IS the RequestOutput instead of holding one. +_REQUEST_OUTPUT_CONTENT_ATTRS = ( + "request_id", + "prompt", + "prompt_token_ids", + "prompt_logprobs", + "outputs", + "finished", + "lora_request", + "encoder_prompt", + "encoder_prompt_token_ids", + "num_cached_tokens", + "kv_transfer_params", +) + +# Omni-specific content copied when the wrapped stage output is itself an +# OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline). +_OMNI_CONTENT_ATTRS = ( + "images", + "latents", + "trajectory_latents", + "trajectory_timesteps", + "trajectory_log_probs", + "trajectory_decoded", + "_multimodal_output", + "_custom_output", +) + + @dataclass -class OmniRequestOutput: +class OmniRequestOutput(RequestOutput): """Unified request output for both pipeline stages and diffusion models. + Extends vLLM's ``RequestOutput`` so that omni outputs can flow directly + through vLLM serving codepaths (which expect ``prompt_token_ids``, + ``outputs``, etc. as real attributes). The inherited fields store the + LLM generation content; omni-specific fields store pipeline/diffusion + extras. + + Note: ``RequestOutput`` is a plain class (not a dataclass), so all of its + attributes are redeclared below as dataclass fields with defaults — the + dataclass-generated ``__init__`` replaces ``RequestOutput.__init__`` and + must set them itself. + This class handles outputs from: - 1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output) + 1. Multi-stage LLM pipelines (with stage_id, final_output_type, and the + inherited RequestOutput fields carrying the stage's generation content) 2. Diffusion models (with images, prompt, metrics) Attributes: @@ -85,31 +128,39 @@ class OmniRequestOutput: stage_id: Identifier of the stage that produced this output (pipeline mode) replica_id: Identifier of the stage replica that produced this output final_output_type: Type of output ("text", "image", "audio", "latents") - request_output: The underlying RequestOutput from the stage (pipeline mode) images: List of generated PIL images (diffusion mode) - prompt: The prompt used for generation (diffusion mode) + prompt: The prompt used for generation latents: Optional tensor of latent representations (diffusion mode) - metrics: Optional dictionary of generation metrics + metrics: Generation metrics. A plain dict for omni outputs; may carry + vLLM's request stats object when copied from a raw RequestOutput. """ + # --- Inherited RequestOutput attributes (redeclared as fields) --- request_id: str = "" + prompt: OmniPromptType | None = None + prompt_token_ids: list[int] | None = None + prompt_logprobs: Any = None + outputs: list[CompletionOutput] = field(default_factory=list) finished: bool = True + metrics: Any = field(default_factory=dict) + lora_request: Any = None + encoder_prompt: str | None = None + encoder_prompt_token_ids: list[int] | None = None + num_cached_tokens: int | None = None + kv_transfer_params: dict[str, Any] | None = None - # Pipeline stage fields + # --- Pipeline stage fields --- stage_id: int | None = None replica_id: int | None = None final_output_type: str = "text" - request_output: RequestOutput | None = None - # Diffusion model fields + # --- Diffusion model fields --- images: list[Image.Image] = field(default_factory=list) - prompt: OmniPromptType | None = None latents: torch.Tensor | None = None trajectory_latents: torch.Tensor | None = None trajectory_timesteps: torch.Tensor | None = None trajectory_log_probs: torch.Tensor | None = None trajectory_decoded: list | None = None - metrics: dict[str, Any] = field(default_factory=dict) _multimodal_output: dict[str, Any] = field(default_factory=dict) _custom_output: dict[str, Any] = field(default_factory=dict) @@ -124,6 +175,64 @@ class OmniRequestOutput: error_status_code: int | None = None error_type: str | None = None + def _copy_content_from(self, source: RequestOutput) -> None: + """Copy generation content from a stage output into this object. + + *source* should be a vLLM ``RequestOutput`` or a subclass such as + ``OmniRequestOutput``. Control fields (``stage_id``, + ``final_output_type``, ``stage_durations``, ``peak_memory_mb``, + ``error``, ...) are intentionally not copied. + """ + # RequestOutput attributes — guaranteed on any RequestOutput. + # Only copy attributes that actually exist on the source, so that + # duck-typed mocks can omit optional fields without overwriting the + # dataclass defaults (e.g. ``outputs=[]``, ``finished=True``). + for name in _REQUEST_OUTPUT_CONTENT_ATTRS: + if hasattr(source, name): + setattr(self, name, getattr(source, name)) + + # Omni-specific fields — only present when the source is itself an + # OmniRequestOutput (e.g. a diffusion stage inside a pipeline). + if isinstance(source, OmniRequestOutput): + for name in _OMNI_CONTENT_ATTRS: + if hasattr(source, name): + setattr(self, name, getattr(source, name)) + + # Propagate multimodal_output from duck-typed sources (e.g. test mocks + # that set multimodal_output directly on a non-OmniRequestOutput). + if not isinstance(source, OmniRequestOutput): + src_mm = getattr(source, "multimodal_output", None) + if src_mm: + self._multimodal_output = src_mm + + @classmethod + def from_stage_output(cls, source: RequestOutput, **kwargs: Any) -> "OmniRequestOutput": + """Create an OmniRequestOutput from a stage's raw output. + + Copies generation content (``outputs``, ``prompt``, ``prompt_token_ids``, + ``finished``, ``images``, ``latents``, etc.) from *source* onto the + returned object. *source* may be a vLLM ``RequestOutput``, another + ``OmniRequestOutput`` (which inherits from ``RequestOutput``). + + This is the **preferred** way to construct an ``OmniRequestOutput`` + that wraps a stage result. + + Args: + source: The stage output whose content is copied onto the new object. + **kwargs: Passed through to the dataclass constructor (``request_id``, + ``stage_id``, ``final_output_type``, ``metrics``, + ``stage_durations``, ``peak_memory_mb``, ``finished``, etc.). + Typed as ``Any`` because the exact set of valid keys is the + dataclass field list, which is validated by ``cls(**kwargs)`` + at call time. + + Returns: + A new ``OmniRequestOutput`` with the stage's content flattened onto it. + """ + obj = cls(**kwargs) + obj._copy_content_from(source) + return obj + @classmethod def from_error( cls, @@ -150,33 +259,6 @@ class OmniRequestOutput: error_type=error_type, ) - @classmethod - def from_pipeline( - cls, - stage_id: int, - final_output_type: str, - request_output: RequestOutput, - replica_id: int | None = None, - ) -> "OmniRequestOutput": - """Create output from pipeline stage. - - Args: - stage_id: Stage identifier - final_output_type: Type of output - request_output: The stage's output - - Returns: - OmniRequestOutput configured for pipeline mode - """ - return cls( - request_id=getattr(request_output, "request_id", ""), - stage_id=stage_id, - replica_id=replica_id, - final_output_type=final_output_type, - request_output=request_output, - finished=True, - ) - @classmethod def from_diffusion( cls, @@ -236,35 +318,21 @@ class OmniRequestOutput: @property def multimodal_output(self) -> Any: - """Return multimodal output from the underlying request output or local field. + """Return the multimodal output payload. - For pipeline outputs, this checks completion outputs first, then request_output. - For diffusion outputs, this returns the local _multimodal_output field. + Checks completion outputs first (where multimodal_output is attached + by AR stages), then the local _multimodal_output field. Returns either a MultimodalPayload (Phase 3+) or a plain dict (legacy). """ - if self.request_output is None: - return self._multimodal_output - - # Check completion outputs first (where multimodal_output is attached) - for output in getattr(self.request_output, "outputs", []): + for output in self.outputs: if mm := getattr(output, "multimodal_output", None): return mm - if mm := getattr(self.request_output, "multimodal_output", None): - return mm return self._multimodal_output @property def custom_output(self) -> dict[str, Any]: - """Return custom output data from diffusion pipelines. - - For diffusion outputs, returns the local _custom_output field. - For pipeline outputs with an inner OmniRequestOutput, forwards - the custom_output from the inner request output. - """ - if self.request_output is not None: - if isinstance(self.request_output, OmniRequestOutput): - return self.request_output._custom_output + """Return custom output data from diffusion pipelines.""" return self._custom_output @custom_output.setter @@ -276,58 +344,6 @@ class OmniRequestOutput: """Return the number of generated images.""" return len(self.images) - # Pass-through properties keep vLLM serving codepaths compatible with - # OmniRequestOutput for pipeline outputs (Issue #345). - @property - def prompt_token_ids(self) -> list[int] | None: - """Return prompt token IDs from the underlying request output. - - This property is required for compatibility with vLLM's streaming - chat completion generator which checks res.prompt_token_ids. - """ - if self.request_output is not None: - return getattr(self.request_output, "prompt_token_ids", None) - return None - - @property - def outputs(self) -> list[Any]: - """Return outputs from the underlying request output. - - This property is required for compatibility with vLLM's streaming - and non-streaming chat completion generators. - """ - if self.request_output is not None: - return getattr(self.request_output, "outputs", []) - return [] - - @property - def encoder_prompt_token_ids(self) -> list[int] | None: - """Return encoder prompt token IDs from the underlying request output.""" - if self.request_output is not None: - return getattr(self.request_output, "encoder_prompt_token_ids", None) - return None - - @property - def prompt_logprobs(self) -> Any: - """Return prompt logprobs from the underlying request output.""" - if self.request_output is not None: - return getattr(self.request_output, "prompt_logprobs", None) - return None - - @property - def num_cached_tokens(self) -> int | None: - """Return number of cached tokens from the underlying request output.""" - if self.request_output is not None: - return getattr(self.request_output, "num_cached_tokens", None) - return None - - @property - def kv_transfer_params(self) -> Any: - """Return KV transfer params from the underlying request output.""" - if self.request_output is not None: - return getattr(self.request_output, "kv_transfer_params", None) - return None - @property def is_diffusion_output(self) -> bool: """Check if this is a diffusion model output.""" @@ -336,73 +352,7 @@ class OmniRequestOutput: @property def is_pipeline_output(self) -> bool: """Check if this is a pipeline stage output.""" - return self.stage_id is not None and self.request_output is not None - - def unwrap(self) -> "OmniRequestOutput": - """Unwrap nested OmniRequestOutput to get the innermost result. - - This helper handles the common pattern where pipeline outputs may wrap - other OmniRequestOutput instances. It recursively unwraps until it reaches - the final output with actual content (images, text, etc.). - - Returns: - The innermost OmniRequestOutput containing the actual generation results. - - Example: - ```python - result = omni.generate(...) - output = OmniRequestOutput.unwrap_result(result) - if output.images: - # Access images directly - video_frames = output.images - ``` - """ - current = self - # Unwrap nested pipeline outputs - while current.is_pipeline_output and current.request_output is not None: - if isinstance(current.request_output, OmniRequestOutput): - current = current.request_output - else: - break - return current - - @staticmethod - def unwrap_result(result: Any) -> "OmniRequestOutput": - """Unwrap result from omni.generate() to get the final OmniRequestOutput. - - This static helper handles the full unwrapping pattern including: - 1. Extracting from list if needed - 2. Type validation - 3. Recursive unwrapping of nested pipeline outputs - - Args: - result: The result from omni.generate() - may be a list or OmniRequestOutput - - Returns: - The innermost OmniRequestOutput with actual content - - Raises: - ValueError: If result is not an OmniRequestOutput or list containing one - - Example: - ```python - result = omni.generate(...) - output = OmniRequestOutput.unwrap_result(result) - # output is guaranteed to be the final OmniRequestOutput - ``` - """ - # Handle list wrapper - if isinstance(result, list): - if not result: - raise ValueError("Result list is empty") - result = result[0] - - # Validate type - if not isinstance(result, OmniRequestOutput): - raise ValueError(f"Expected OmniRequestOutput, got {type(result)}") - - # Unwrap nested outputs - return result.unwrap() + return self.stage_id is not None def to_dict(self) -> dict[str, Any]: """Convert to dictionary for JSON serialization.""" @@ -441,7 +391,8 @@ class OmniRequestOutput: f"finished={self.finished}", f"stage_id={self.stage_id}", f"final_output_type={self.final_output_type!r}", - f"request_output={self.request_output}", + f"prompt_token_ids={self.prompt_token_ids}", + f"outputs={self.outputs}", f"images={images_repr}", f"prompt={self.prompt!r}", f"latents={self.latents}", diff --git a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py index 77b191738..4d5c69df5 100644 --- a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py +++ b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py @@ -259,17 +259,14 @@ def _extract_inner_output(outputs: Any) -> Any: if len(outputs) != 1: raise ValueError(f"Expected one request output, got {len(outputs)}.") outputs = outputs[0] - inner = getattr(outputs, "request_output", None) + inner = outputs if inner is not None: return inner return outputs def _request_peak_memory_mb(result: Any) -> float | None: - diffusion_output = getattr(result, "request_output", None) - peak_memory_mb = getattr(diffusion_output, "peak_memory_mb", None) - if peak_memory_mb is None: - peak_memory_mb = getattr(result, "peak_memory_mb", None) + peak_memory_mb = getattr(result, "peak_memory_mb", None) if peak_memory_mb is None: return None return float(peak_memory_mb) diff --git a/vllm_omni/utils/audio.py b/vllm_omni/utils/audio.py index bc2498efd..aafd0240b 100644 --- a/vllm_omni/utils/audio.py +++ b/vllm_omni/utils/audio.py @@ -79,7 +79,7 @@ def audio_chunk_pcm_bytes(omni_res: OmniRequestOutput) -> int: sample rather than recording a wrong byte count. """ try: - final_res = omni_res.request_output + final_res = omni_res mm_output = final_res.outputs[0].multimodal_output audio_data = mm_output.get("audio") if isinstance(audio_data, list): @@ -101,7 +101,7 @@ def audio_chunk_pcm_bytes(omni_res: OmniRequestOutput) -> int: def audio_chunk_sample_rate(omni_res: OmniRequestOutput) -> int: """Resolve audio sample rate for the request's audio stream.""" try: - mm_output = omni_res.request_output.outputs[0].multimodal_output + mm_output = omni_res.outputs[0].multimodal_output except Exception: return _metric_defs.DEFAULT_AUDIO_SAMPLE_RATE return _metric_defs.resolve_audio_sample_rate(mm_output) From 208c6776a89f9dfa1598457196d71538bb9b24ac Mon Sep 17 00:00:00 2001 From: Juan Pablo Zuluaga <46724788+JuanPZuluaga@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:02:02 +0200 Subject: [PATCH 12/41] [Perf][TTS] glm_tts: hoist loop-invariant text embed, RoPE, mask and CFG batch out of Euler loop (#5068) Signed-off-by: JuanPZuluaga Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../models/glm_tts/glm_tts_dit.py | 21 +++++-- .../models/glm_tts/glm_tts_dit_wrapper.py | 58 +++++++++++++------ 2 files changed, 58 insertions(+), 21 deletions(-) diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py index 666998e9d..03b241035 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py @@ -332,7 +332,7 @@ class DiTAttention(nn.Module): self, x: torch.Tensor, padding_mask: torch.Tensor | None = None, - rope=None, + rope: tuple[torch.Tensor, torch.Tensor | float] | None = None, attn_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, seq_len = x.shape[0], x.shape[1] @@ -406,7 +406,7 @@ class DiTBlock(nn.Module): x: torch.Tensor, t: torch.Tensor, padding_mask: torch.Tensor | None = None, - rope=None, + rope: tuple[torch.Tensor, torch.Tensor | float] | None = None, attn_mask: torch.Tensor | None = None, ) -> torch.Tensor: # Pre-norm & modulation for attention @@ -575,6 +575,8 @@ class GLMTTSDiT(nn.Module): is_causal: bool = False, attn_mask: torch.Tensor | None = None, block_pattern: list[int] | None = None, + text_embed: torch.Tensor | None = None, + rope: tuple[torch.Tensor, torch.Tensor | float] | None = None, ) -> torch.Tensor: """Forward pass. @@ -588,6 +590,9 @@ class GLMTTSDiT(nn.Module): is_causal: Use block-causal attention attn_mask: Custom attention mask (overrides is_causal) block_pattern: Token-level block sizes for block-causal mask + text_embed: Precomputed text embedding (must already be in model + dtype); skips the internal embedding + ConvNeXt pass + rope: Precomputed rotary embedding tuple for ``seq_len`` Returns: Predicted mel-spectrogram [B, T, mel_dim] @@ -610,13 +615,21 @@ class GLMTTSDiT(nn.Module): time_emb = torch.cat([time_emb, spkr_emb], dim=-1) # Text embedding - ensure cast to model dtype - text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype) + if text_embed is None: + text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype) + else: + # A precomputed text_embed replaces `text` entirely, so its length must + # match middle_point; a mismatch would apply the wrong rotary/mask silently. + assert text_embed.shape[1] == seq_len, ( + f"precomputed text_embed len {text_embed.shape[1]} != seq_len {seq_len}" + ) # Input projection x = self.emb_concator(middle_point, condition, text_embed, drop_audio_cond=False) # Rotary embeddings - rope = self.rotary_embed.forward_from_seq_len(seq_len) + if rope is None: + rope = self.rotary_embed.forward_from_seq_len(seq_len) # Build block-causal mask when is_causal and no explicit attn_mask if attn_mask is None and is_causal and block_pattern is not None: diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py index a7c26af30..d4b391fbe 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py @@ -857,6 +857,40 @@ class GLMTTSDiTForGeneration(nn.Module): if self.t_scheduler == "cosine": t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + # Hoist loop-invariant DiT inputs (text embedding, RoPE table, + # block-causal mask, CFG double-batch conditioning) out of the Euler + # loop; they depend only on the conditioning and seq_len, not on the + # per-step state x or the timestep. + seq_len = int(mel_cond.shape[1]) + # Must match the dtype forward() casts with internally (both are + # proj_out.weight.dtype); keep them in sync if that derivation changes. + dit_dtype = self.dit.proj_out.weight.dtype + # `device` is speech_tokens.device; in this call path the sampled state x + # lives on the same device, so the hoisted mask matches the old in-forward one. + attn_mask = ( + self.dit._create_block_causal_mask(block_pattern, seq_len, device, dit_dtype) + if block_pattern is not None + else None + ) + rope = self.dit.rotary_embed.forward_from_seq_len(seq_len) + + if self.inference_cfg_rate > 0: + condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0) + if self.speech_token_cfg: + text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0) + else: + text_b = speech_tokens.expand(2, -1) + padding_b = padding_mask.expand(2, -1) + if spkr_embedding_dit is not None and self.spkr_emb_adaln: + spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0) + else: + spkr_b = None + text_lens_b = text_lens.expand(2) if text_lens is not None else None + text_embed_b = self.dit.text_emb_layer(text_b, seq_len, text_lens=text_lens_b).to(dit_dtype) + else: + spkr_uncond = spkr_embedding_dit if (spkr_embedding_dit is not None and self.spkr_emb_adaln) else None + text_embed_u = self.dit.text_emb_layer(speech_tokens, seq_len, text_lens=text_lens).to(dit_dtype) + for step in range(n_timesteps): cache_step = step + 1 if last_step_cache is not None and cache_step in last_step_cache: @@ -872,24 +906,12 @@ class GLMTTSDiTForGeneration(nn.Module): t_current = t_span[step] dt = t_span[step + 1] - t_current - use_causal = block_pattern is not None x_dit = x.to(mel_cond.dtype) if self.inference_cfg_rate > 0: # Batch conditional + unconditional into a single forward middle_b = x_dit.expand(2, -1, -1) - condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0) - if self.speech_token_cfg: - text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0) - else: - text_b = speech_tokens.expand(2, -1) time_b = t_current.unsqueeze(0).expand(2) - padding_b = padding_mask.expand(2, -1) - if spkr_embedding_dit is not None and self.spkr_emb_adaln: - spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0) - else: - spkr_b = None - text_lens_b = text_lens.expand(2) if text_lens is not None else None pred_b = self.dit( middle_point=middle_b, condition=condition_b, @@ -897,9 +919,10 @@ class GLMTTSDiTForGeneration(nn.Module): time_step=time_b, padding_mask=padding_b, spkr_emb=spkr_b, - is_causal=use_causal, - block_pattern=block_pattern, + attn_mask=attn_mask, text_lens=text_lens_b, + text_embed=text_embed_b, + rope=rope, ).to(torch.float32) dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0) dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt @@ -910,10 +933,11 @@ class GLMTTSDiTForGeneration(nn.Module): text=speech_tokens, time_step=t_current.unsqueeze(0), padding_mask=padding_mask, - spkr_emb=spkr_embedding_dit if spkr_embedding_dit is not None and self.spkr_emb_adaln else None, - is_causal=use_causal, - block_pattern=block_pattern, + spkr_emb=spkr_uncond, + attn_mask=attn_mask, text_lens=text_lens, + text_embed=text_embed_u, + rope=rope, ).to(torch.float32) x = x + dt * dphi_dt From e677360168d7d4f512bcebd50fbb1cbfc422d133 Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Wed, 12 Aug 2026 19:04:07 +0200 Subject: [PATCH 13/41] Add cosmos-guardrail dependency and update error message (#6107) Signed-off-by: Maciej Bala --- requirements/common.txt | 2 ++ vllm_omni/diffusion/models/cosmos3/guardrails.py | 7 +++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 243a6ce7e..13c7d35a9 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -34,3 +34,5 @@ pydantic-settings>=2.0.0 # base image does not ship it, so declare it explicitly. The loader also guards # this import so a missing gguf no longer breaks all diffusion model loading. gguf>=0.10.0 +# Required for running the Cosmos3 model with guardrails enabled (default behavior). +cosmos-guardrail>=0.3.1 diff --git a/vllm_omni/diffusion/models/cosmos3/guardrails.py b/vllm_omni/diffusion/models/cosmos3/guardrails.py index 3c6166e85..3065f58ac 100644 --- a/vllm_omni/diffusion/models/cosmos3/guardrails.py +++ b/vllm_omni/diffusion/models/cosmos3/guardrails.py @@ -65,11 +65,10 @@ except ImportError: # ``model_config["guardrails"]`` is False. def __init__(self, *args, **kwargs): raise ValueError( - f"You have disabled the safety checker for {self.__class__}. This is in violation of the " + f"The guardrails are enabled but not functional for {self.__class__}. This is most likely due to a" + f"missing cosmos-guardrail package. Please install cosmos-guardrail to be compliant with the " "[NVIDIA Open Model License Agreement]" - "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). " - f"Please ensure that you are compliant with the license agreement." - f"Please install cosmos-guardrail package to enable safety checks." + "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license)." ) From 5358f5db01e512dc20fcc29a1d122e82c8546780 Mon Sep 17 00:00:00 2001 From: Juan Pablo Zuluaga <46724788+JuanPZuluaga@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:15:35 +0200 Subject: [PATCH 14/41] [Perf][TTS] omnivoice hot path: D2H batching, mask caching, cached text embeddings (#5174) Signed-off-by: JuanPZuluaga Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../models/omnivoice/omnivoice_generator.py | 168 +++++++++++++----- .../transformers_utils/configs/omnivoice.py | 2 + 2 files changed, 125 insertions(+), 45 deletions(-) diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py index 7eacfb9bb..087cba96f 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py @@ -242,6 +242,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge # --------------------------------------------------------------------------- +# Subclass keeps .weight name + ctor shape so the state_dict loader stays unchanged. class OmniVoiceRMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float = 1e-6): super().__init__() @@ -315,10 +316,9 @@ class OmniVoiceAttention(nn.Module): k = k.permute(0, 2, 1, 3) v = v.permute(0, 2, 1, 3) - # Convert [B, 1, S, S] bool mask to float mask for SDPA - # (0.0 where attend, -inf where masked) - sdpa_mask = None - if attention_mask is not None: + # Caller passes a float mask; materialize float form if a bool slips through. + sdpa_mask = attention_mask + if sdpa_mask is not None and sdpa_mask.dtype == torch.bool: sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf")) out = F.scaled_dot_product_attention( @@ -400,24 +400,48 @@ def _precompute_rope( theta: float = 1000000.0, device: torch.device | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Precompute RoPE cos/sin tensors.""" + """Precompute RoPE cos/sin pre-doubled to full head_dim so the hot path skips a per-call cat.""" inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) t = torch.arange(max_seq_len, device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) - cos = freqs.cos() - sin = freqs.sin() + cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1) + sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1) return cos, sin def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: - """Apply rotary position embedding. x shape: (B, S, H, D).""" + """Apply rotary position embedding. x shape: (B, S, H, D). cos/sin are pre-doubled to width D.""" seq_len = x.shape[1] - cos = cos[:seq_len].unsqueeze(0).unsqueeze(2) # (1, S, 1, D/2) + cos = cos[:seq_len].unsqueeze(0).unsqueeze(2) # (1, S, 1, D) sin = sin[:seq_len].unsqueeze(0).unsqueeze(2) x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] rotated = torch.cat([-x2, x1], dim=-1) - return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1) + return x * cos + rotated * sin + + +# --------------------------------------------------------------------------- +# TF32 opt-in (process-wide; default off) +# --------------------------------------------------------------------------- + +_TF32_ENABLED = False + + +def _maybe_enable_tf32() -> None: + """Enable TF32 matmuls process-wide (idempotent). Not bit-identical; opt-in via config.enable_tf32.""" + global _TF32_ENABLED + if _TF32_ENABLED or not torch.cuda.is_available(): + return + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision("high") + _TF32_ENABLED = True + logger.info( + "OmniVoice TF32 enabled process-wide: matmul.allow_tf32=%s cudnn.allow_tf32=%s float32_matmul_precision=%s", + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.get_float32_matmul_precision(), + ) # --------------------------------------------------------------------------- @@ -425,6 +449,19 @@ def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) # --------------------------------------------------------------------------- +def _additive_float_mask(mask: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Convert a boolean attention mask to its additive float form. + + ``True`` (attend) maps to ``0.0`` and ``False`` (masked) to ``-inf``. A bool + mask must never be copied straight into a float buffer: the implicit cast + maps True/False to 1.0/0.0, which leaves masked positions at 0.0 and so + silently *unmasks* them. + """ + if mask.dtype != torch.bool: + return mask + return torch.zeros_like(mask, dtype=dtype).masked_fill_(~mask, float("-inf")) + + class _OmniVoiceCUDAGraphForward: """Pre-captures CUDA graphs for predefined sequence-length buckets. @@ -476,11 +513,10 @@ class _OmniVoiceCUDAGraphForward: mask_padded[:, :S] = audio_mask if attention_mask is not None: - attn_padded = torch.zeros( - two_b, - 1, - bucket, - bucket, + # Callers normalize to the additive float form first, so pad with -inf. + attn_padded = torch.full( + (two_b, 1, bucket, bucket), + float("-inf"), dtype=attention_mask.dtype, device=attention_mask.device, ) @@ -564,7 +600,8 @@ class _OmniVoiceCUDAGraphForward: key = (two_b, bucket) dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device) dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device) - dummy_attn = torch.ones(two_b, 1, bucket, bucket, dtype=torch.bool, device=device) + # Capture with a float mask to match what forward() feeds at replay time. + dummy_attn = torch.zeros(two_b, 1, bucket, bucket, dtype=torch.float32, device=device) self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn) logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs)) @@ -586,6 +623,11 @@ class _OmniVoiceCUDAGraphForward: two_b = input_ids.shape[0] bucket = self._find_bucket(seq_len) if two_b == 2 else None + # Graphs are captured with (and their static buffers hold) the additive + # float mask, so normalize here, before padding or any copy_ into them. + if attention_mask is not None: + attention_mask = _additive_float_mask(attention_mask) + if bucket is None: # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket). # Lock prevents concurrent threads from double-capturing the same key. @@ -655,6 +697,10 @@ class OmniVoiceGenerator(nn.Module): super().__init__() self.config = config + # Opt-in TF32; must run before any CUDA-graph capture so captured kernels honour it. + if getattr(config, "enable_tf32", False): + _maybe_enable_tf32() + # Text embedding (shared with LLM) self.text_embedding = nn.Embedding(config.llm_vocab_size, config.llm_hidden_size) @@ -703,6 +749,8 @@ class OmniVoiceGenerator(nn.Module): self, input_ids: torch.Tensor, audio_mask: torch.Tensor, + text_embeds: torch.Tensor | None = None, + audio_mask_3d: torch.Tensor | None = None, ) -> torch.Tensor: """Prepare mixed text+audio embeddings. @@ -710,30 +758,38 @@ class OmniVoiceGenerator(nn.Module): input_ids: [B, 8, S] - text tokens replicated across codebooks, audio positions have per-codebook token IDs audio_mask: [B, S] - True for audio positions, False for text + text_embeds: optional cached [B, S, H] text-position embeddings + audio_mask_3d: optional cached [B, S, 1] audio_mask.unsqueeze(-1) Returns: embeddings: [B, S, hidden_size] """ - # Text embeddings from first codebook row (all rows identical for text) - text_embeds = self.text_embedding(input_ids[:, 0, :]) + # Cached across the denoising loop since text ids don't change. + if text_embeds is None: + text_embeds = self.text_embedding(input_ids[:, 0, :]) + if audio_mask_3d is None: + audio_mask_3d = audio_mask.unsqueeze(-1) # Audio embeddings: offset per codebook, then sum across codebooks shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1) audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1) # Merge: audio where audio_mask=True, text elsewhere - return torch.where(audio_mask.unsqueeze(-1), audio_embeds, text_embeds) + return torch.where(audio_mask_3d, audio_embeds, text_embeds) def _transformer_forward( self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor | None = None, + cos: torch.Tensor | None = None, + sin: torch.Tensor | None = None, ) -> torch.Tensor: """Run through transformer layers. Args: inputs_embeds: [B, S, hidden_size] attention_mask: [B, 1, S, S] or None + cos, sin: optional precomputed RoPE tables in target dtype Returns: hidden_states: [B, S, hidden_size] @@ -743,8 +799,15 @@ class OmniVoiceGenerator(nn.Module): self._ensure_rope(seq_len, device) hidden_states = inputs_embeds - cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype) - sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype) + if cos is None or sin is None: + cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype) + sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype) + + # Safety: convert bool mask if caller hasn't (e.g. external paths beyond forward()). + if attention_mask is not None and attention_mask.dtype == torch.bool: + attention_mask = torch.zeros_like(attention_mask, dtype=hidden_states.dtype).masked_fill_( + ~attention_mask, float("-inf") + ) for layer in self.layers: hidden_states = layer( @@ -859,21 +922,37 @@ class OmniVoiceGenerator(nn.Module): layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1) - # Compute c_lens for extracting target region from full sequence - c_lens = [] - for i in range(B): - # Conditional sequence length = number of non-padding positions - c_len = attention_mask[i, 0, 0].sum().item() - c_lens.append(int(c_len)) + # Single D2H pull for all conditional lengths instead of B per-item .item() syncs. + c_lens = attention_mask[:B, 0, 0].sum(dim=-1).tolist() + + # Materialize the SDPA float mask once so the captured graph (and eager path) skip per-layer conversion. + sdpa_attn_mask = torch.zeros_like(attention_mask, dtype=torch.float32).masked_fill_( + ~attention_mask, float("-inf") + ) + + use_cuda_graph = self._cuda_graph_fwd is not None and input_ids.is_cuda + if not use_cuda_graph: + # Eager-path-only constants (the cuda-graph captures its own). + text_embeds_cached = self.text_embedding(input_ids[:, 0, :]) + audio_mask_3d = audio_mask.unsqueeze(-1) + self._ensure_rope(input_ids.shape[-1], device) + target_dtype = text_embeds_cached.dtype + cos = self._rope_cos.to(device=device, dtype=target_dtype) + sin = self._rope_sin.to(device=device, dtype=target_dtype) # Main iterative loop for step in range(num_step): - if self._cuda_graph_fwd is not None and input_ids.is_cuda: - batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, attention_mask).to(torch.float32) + if use_cuda_graph: + # Float mask skips per-layer conversion; fp32 cast deferred to the per-item slices below. + batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, sdpa_attn_mask) else: - inputs_embeds = self._prepare_embeddings(input_ids, audio_mask) - hidden_states = self._transformer_forward(inputs_embeds, attention_mask) - batch_logits = self._get_logits(hidden_states).to(torch.float32) + # Eager fallback reuses hoisted constants (text embeds, sdpa mask, cos/sin). + inputs_embeds = self._prepare_embeddings( + input_ids, audio_mask, text_embeds=text_embeds_cached, audio_mask_3d=audio_mask_3d + ) + hidden_states = self._transformer_forward(inputs_embeds, sdpa_attn_mask, cos=cos, sin=sin) + # fp32 cast deferred to the per-item slices below. + batch_logits = self._get_logits(hidden_states) # batch_logits: [2*B, 8, S, 1025] for i in range(B): @@ -884,16 +963,17 @@ class OmniVoiceGenerator(nn.Module): c_len = c_lens[i] t_len = target_lens[i] - # Extract logits for target region - c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :] # [1, 8, T, 1025] - u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :] # [1, 8, T, 1025] + # Extract logits for target region; upcast only the slices we actually consume. + c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :].to(torch.float32) + u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :].to(torch.float32) - # Classifier-free guidance + # Classifier-free guidance. Fuse the chain: the two inner + # log_softmax normalizers are per-position scalars that the final + # shift-invariant log_softmax cancels, so guide on the raw logits + # with a single softmax: log_softmax((1+s)*c - s*u). Exact. if guidance_scale != 0: - c_log_probs = F.log_softmax(c_logits, dim=-1) - u_log_probs = F.log_softmax(u_logits, dim=-1) - log_probs = torch.log_softmax( - c_log_probs + guidance_scale * (c_log_probs - u_log_probs), + log_probs = F.log_softmax( + (1.0 + guidance_scale) * c_logits - guidance_scale * u_logits, dim=-1, ) else: @@ -922,15 +1002,13 @@ class OmniVoiceGenerator(nn.Module): sample_tokens = tokens[i : i + 1, :, :t_len] scores.masked_fill_(sample_tokens != mask_id, -float("inf")) - # Select top-k positions to unmask + # Select top-k positions to unmask. .flatten() on this non-contiguous view already copies. _, topk_idx = torch.topk(scores.flatten(), k) - flat_tokens = sample_tokens.flatten().clone() + flat_tokens = sample_tokens.flatten() flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx] sample_tokens.copy_(flat_tokens.view_as(sample_tokens)) - # Update tokens and batch inputs for next iteration - tokens[i : i + 1, :, :t_len] = sample_tokens - input_ids = input_ids.clone() + # Mirror update into both cond and uncond input_ids halves for the next step. input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0) input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0) diff --git a/vllm_omni/transformers_utils/configs/omnivoice.py b/vllm_omni/transformers_utils/configs/omnivoice.py index 1bd7207b6..ac203bf04 100644 --- a/vllm_omni/transformers_utils/configs/omnivoice.py +++ b/vllm_omni/transformers_utils/configs/omnivoice.py @@ -90,6 +90,8 @@ class OmniVoiceConfig(PretrainedConfig): "cuda_graph_capture_sizes", [128, 192, 256, 320, 384, 448, 512, 640, 768, 1024], ) + # TF32 matmuls: not bit-identical; opt-in (matches vLLM default-off), set OMNIVOICE_TF32=1 to enable. + self.enable_tf32 = getattr(self, "enable_tf32", os.environ.get("OMNIVOICE_TF32", "0") != "0") AutoConfig.register("omnivoice", OmniVoiceConfig) From 72a02b492a6957acc4659fde09be85c00ad0f90b Mon Sep 17 00:00:00 2001 From: Zhenghan Song Date: Thu, 13 Aug 2026 03:20:12 +0800 Subject: [PATCH 15/41] [Bugfix] Fix diffusion TTS adapter lookup in speech serving (#6121) Signed-off-by: HaningZS --- .../openai_api/test_serving_speech.py | 21 +++++++++++++++++-- .../entrypoints/openai/serving_speech.py | 6 ++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index bcb89ff88..94d84d041 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -680,7 +680,7 @@ class TestSpeechAPI: @pytest.mark.asyncio async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture): - """Test that extra_params are correctly applied to sampling_params_list in diffusion mode.""" + """Test public diffusion speech success and extra_params propagation.""" # Mock the engine client mock_engine = mocker.MagicMock() @@ -704,7 +704,11 @@ class TestSpeechAPI: req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"}) - await server._create_diffusion_speech(req) + response = await server.create_speech(req) + + assert response.status_code == 200 + assert response.media_type == "audio/wav" + assert response.body == b"dummy" # Verify generate was called mock_engine.generate.assert_called_once() @@ -737,6 +741,19 @@ class TestTTSMethods: yield server server.shutdown() + def test_diffusion_audio_encode_speed_skips_ar_adapter(self, mocker: MockerFixture): + """Pure-diffusion speech keeps generic speed handling outside AR adapters.""" + server = OmniOpenAIServingSpeech.for_diffusion( + diffusion_engine=mocker.MagicMock(), + model_name="test-model", + ) + resolve_adapter = mocker.patch.object(serving_speech_module, "resolve_adapter") + + speed = server._audio_encode_speed(OpenAICreateSpeechRequest(input="Hello", speed=1.25)) + + assert speed == 1.25 + resolve_adapter.assert_not_called() + def test_is_tts_detection_no_stage(self, speech_server): """Test TTS model detection when no TTS stage exists.""" # Fixture creates server with stage_configs = [] -> _is_tts should be False diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index a662ae7de..109777282 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -502,11 +502,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _get_tts_adapter(self): """Return the per-model serving adapter for the current ``_tts_model_type``. + Pure-diffusion speech uses its dedicated request path and does not + resolve adapters for AR-stage TTS models. + Resolved lazily (rebuilt if ``_tts_model_type`` changed since the cached instance was built) so callers that set ``_tts_model_type`` after construction still dispatch to the matching adapter. In production ``_tts_model_type`` is fixed at init, so the cached instance is reused. """ + if self._diffusion_mode: + return None + adapter_cls = resolve_adapter(self._tts_model_type) if adapter_cls is None: self._adapter = None From 12f81681793441632ffb7d68b200b9d61821e1f8 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 12 Aug 2026 16:03:49 -0400 Subject: [PATCH 16/41] [Tests] Enable tiny model testing for Qwen-Image Edit and EditPlus (#5656) Signed-off-by: Nick Cao --- .../diffusion/diff_model_builders.py | 60 +++++++++++++++---- tests/model_tests/diffusion/model_settings.py | 10 ++++ tests/model_tests/diffusion/test_alignment.py | 2 - 3 files changed, 58 insertions(+), 14 deletions(-) diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 9353d54f6..e47ea1895 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -13,10 +13,13 @@ def _shrink_dit_rope_config( target_head_dim: int = 32, target_heads: int = 4, default_axes_dims_rope: list[int] | None = None, + joint_attention_dim: int | None = None, ) -> dict: config["num_layers"] = num_layers if num_single_layers is not None: config["num_single_layers"] = num_single_layers + if joint_attention_dim is not None: + config["joint_attention_dim"] = joint_attention_dim axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope) factor = config["attention_head_dim"] / target_head_dim config["attention_head_dim"] = target_head_dim @@ -50,22 +53,55 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict: return config -def tiny_qwen_image_builder() -> str: - def shrink_text_encoder(config: dict) -> dict: - config["num_hidden_layers"] = 2 - config["intermediate_size"] = 64 - config["text_config"]["num_hidden_layers"] = 2 - config["text_config"]["intermediate_size"] = 64 - config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2] - config["vision_config"]["depth"] = 2 - config["vision_config"]["intermediate_size"] = 64 - config["vision_config"]["fullatt_block_indexes"] = [0, 1] - return config +def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict: + text_config = config["text_config"] + old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"] + text_config["num_hidden_layers"] = 2 + text_config["intermediate_size"] = 64 + text_config["hidden_size"] = hidden_size + text_config["num_attention_heads"] = 2 + text_config["num_key_value_heads"] = 2 + text_config["layer_types"] = text_config["layer_types"][:2] + factor = old_head_dim / (hidden_size / 2) + mrope_section = text_config["rope_scaling"]["mrope_section"] + text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section] + config["num_hidden_layers"] = 2 + config["intermediate_size"] = 64 + config["hidden_size"] = hidden_size + config["num_attention_heads"] = 2 + config["num_key_value_heads"] = 2 + config["vision_config"]["depth"] = 2 + config["vision_config"]["intermediate_size"] = 64 + config["vision_config"]["fullatt_block_indexes"] = [0, 1] + config["vision_config"]["out_hidden_size"] = hidden_size + return config + +def _shrink_qwen_transformer_config(config: dict, joint_attention_dim: int = 64) -> dict: + return _shrink_dit_rope_config(config, joint_attention_dim=joint_attention_dim) + + +def tiny_qwen_image_builder() -> str: return build_tiny_from_configs( "QwenImagePipeline", "Qwen/Qwen-Image", - transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config}, + transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config}, + ) + + +def tiny_qwen_image_edit_builder() -> str: + return build_tiny_from_configs( + "QwenImageEditPipeline", + "Qwen/Qwen-Image-Edit", + transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config}, + ) + + +def tiny_qwen_image_edit_plus_builder() -> str: + return build_tiny_from_configs( + "QwenImageEditPlusPipeline", + "Qwen/Qwen-Image-Edit-2511", + transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config}, ) diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 10500c688..cbb37729a 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -74,4 +74,14 @@ DIFFUSION_TEST_SETTINGS = { builder=diff_model_builders.tiny_flux2_builder, supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE], ), + "QwenImageEditPipeline": DiffusionModelTestOpts( + model="Qwen/Qwen-Image-Edit", + builder=diff_model_builders.tiny_qwen_image_edit_builder, + supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE], + ), + "QwenImageEditPlusPipeline": DiffusionModelTestOpts( + model="Qwen/Qwen-Image-Edit-2511", + builder=diff_model_builders.tiny_qwen_image_edit_plus_builder, + supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE], + ), } diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 17d1eb698..be064f603 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -21,8 +21,6 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] # if you are adding a new model and see the tests below fail, please follow the pattern # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS. EXCLUDED_MODELS = [ - "QwenImageEditPipeline", - "QwenImageEditPlusPipeline", "QwenImageLayeredPipeline", "GlmImagePipeline", "ZImagePipeline", From 9a59475015102608dc8655c967e817a66581e8ef Mon Sep 17 00:00:00 2001 From: mjZhaoElaine <38148457+mjZhaoElaine@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:52:01 -0700 Subject: [PATCH 17/41] [CI/Build] Fix ReadTheDocs build under RTD's seeded pip 23.1 (#6129) Signed-off-by: Mengjie Zhao Co-authored-by: Claude Opus 5 --- .readthedocs.yml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/.readthedocs.yml b/.readthedocs.yml index 2b5b81bbc..be2996355 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -11,10 +11,15 @@ build: post_checkout: - git fetch --unshallow || true pre_install: + # `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1 + # rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2). + - python -m pip install --upgrade pip # Pre-install CPU torch so the subsequent `pip install .[docs]` does not # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's - # 1 GB memory cap and time out the mkdocs build. - - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch + # 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url` + # replaces PyPI, so an sdist fallback cannot reach its build deps. The + # quotes keep YAML from reading `:all:` as a mapping key. + - "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch" mkdocs: configuration: mkdocs.yml From c57178c0e18df7b622fe6e3ee8357495305876b5 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 13 Aug 2026 08:41:59 +0800 Subject: [PATCH 18/41] [Dependency] Upgrade Cache-DiT to 1.5.0 (#6065) Signed-off-by: Hongsheng Liu --- .buildkite/cuda/test-nightly.yml | 10 +++++----- .buildkite/npu/test-npu-nightly.yml | 2 +- recipes/Tencent/HunyuanImage-3.0-Instruct.md | 2 +- requirements/common.txt | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index a1c0c4862..d3246a98f 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -345,7 +345,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -361,7 +361,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json @@ -377,7 +377,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json @@ -393,7 +393,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -409,7 +409,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json diff --git a/.buildkite/npu/test-npu-nightly.yml b/.buildkite/npu/test-npu-nightly.yml index 7f277258c..050cef4f7 100644 --- a/.buildkite/npu/test-npu-nightly.yml +++ b/.buildkite/npu/test-npu-nightly.yml @@ -115,7 +115,7 @@ steps: DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md index a9a63d5ce..4614df382 100644 --- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md +++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md @@ -82,7 +82,7 @@ example and the DiT sections below. - Optional environment variables: ```bash -export CACHE_DIT_VERSION=1.3.0 +export CACHE_DIT_VERSION=1.5.0 ``` HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA` diff --git a/requirements/common.txt b/requirements/common.txt index 13c7d35a9..e0233bd8d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -16,7 +16,7 @@ diffusers==0.38.0 safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 -cache-dit==1.3.0 +cache-dit==1.5.0 tqdm>=4.66.0 torchsde>=0.2.6 openai-whisper>=20250625 From 8b8d3f56db4a1e39c635db73fec96112918d8083 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 13 Aug 2026 08:42:37 +0800 Subject: [PATCH 19/41] [Misc] Consolidate LingBot text-to-video runner (#6076) Signed-off-by: Hongsheng Liu --- .../lingbot_video/dense_pipeline_parity.py | 4 +- .../offline_inference/text_to_video.md | 20 +++ .../text_to_video/text_to_video.md | 20 ++- .../text_to_video/text_to_video.py | 74 +++++++--- .../text_to_video/text_to_video_lingbot.py | 131 ------------------ recipes/Robbyant/LingBot-Video.md | 6 +- .../test_text_to_video_prompts.py | 91 ++++++++++++ tests/model_extras/test_model_extras.py | 3 + vllm_omni/model_extras/__init__.py | 2 + vllm_omni/model_extras/registry.py | 16 ++- 10 files changed, 208 insertions(+), 159 deletions(-) delete mode 100644 examples/offline_inference/text_to_video/text_to_video_lingbot.py create mode 100644 tests/examples/offline_inference/test_text_to_video_prompts.py diff --git a/benchmarks/lingbot_video/dense_pipeline_parity.py b/benchmarks/lingbot_video/dense_pipeline_parity.py index 0c0c1d7a0..cc781a95b 100644 --- a/benchmarks/lingbot_video/dense_pipeline_parity.py +++ b/benchmarks/lingbot_video/dense_pipeline_parity.py @@ -167,9 +167,11 @@ def main() -> None: ] native_cmd = [ sys.executable, - "examples/offline_inference/text_to_video/text_to_video_lingbot.py", + "examples/offline_inference/text_to_video/text_to_video.py", "--model", args.model, + "--model-class-name", + "LingBotVideoPipeline", "--prompt", args.prompt, "--output", diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index d32a1f562..1a2082f82 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) | |---|---|---|---|---|---| | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) | | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | @@ -35,6 +36,25 @@ python text_to_video.py \ --output t2v_out.mp4 ``` +### LingBot-Video + +The shared runner recognizes both official LingBot checkpoint IDs and selects +the native `LingBotVideoPipeline` automatically: + +```bash +python text_to_video.py \ + --model robbyant/lingbot-video-dense-1.3b \ + --prompt "a robotic arm picks up a red block" \ + --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \ + --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \ + --output lingbot_t2v.mp4 +``` + +Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local +paths without `lingbot` in the directory name, add +`--model-class-name LingBotVideoPipeline`. Model-only options use +`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`. + ### LTX-2 ```bash diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 50b207526..339c7dac9 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod |---|---|---|---|---|---| | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) | | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) | | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | @@ -57,18 +58,31 @@ python text_to_video.py \ ### LingBot-Video -The dedicated LingBot example explicitly selects the video output modality: +The shared runner recognizes the official dense and MoE checkpoint IDs, selects +`LingBotVideoPipeline`, and builds the canonical video request envelope: ```bash -python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ +python text_to_video.py \ --model robbyant/lingbot-video-dense-1.3b \ + --prompt "a robotic arm picks up a red block" \ --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \ + --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \ --output lingbot_t2v.mp4 ``` Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. Requested frame counts are rounded upward to the causal -VAE's `4n+1` grid. +VAE's `4n+1` grid. For a local checkpoint path whose name does not contain +`lingbot`, pass `--model-class-name LingBotVideoPipeline`. + +LingBot-only options such as `batch_cfg` and `output_type` use the shared +model-extra channel, for example: + +```bash +python text_to_video.py \ + --model robbyant/lingbot-video-dense-1.3b \ + --extra-body '{"batch_cfg": true, "output_type": "np"}' +``` ### LTX-2 diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 2e272b7a7..129ca860a 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id -from vllm_omni.model_extras import get_extra_body_params, get_model_class_name +from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -80,6 +80,17 @@ _MODEL_PRESETS = { "fps": 16, "output": "helios_output.mp4", }, + "lingbot": { + "model_class_name": "LingBotVideoPipeline", + "height": 192, + "width": 320, + "num_frames": 9, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "fps": 24, + "flow_shift": 3.0, + "output": "lingbot_video_output.mp4", + }, "ltx2": { "height": 512, "width": 768, @@ -110,27 +121,51 @@ _MODEL_PRESETS = { def _detect_preset(model: str, model_class_name: str | None = None) -> dict: model_lower = model.lower() class_lower = (model_class_name or "").lower() - if "distilled" in class_lower or "distilled" in model_lower: - return _MODEL_PRESETS["ltx2_distilled"] - if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower: - return _MODEL_PRESETS["ltx23"] - if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower: + if "lingbot" in model_lower or "lingbotvideo" in class_lower: + return _MODEL_PRESETS["lingbot"] + if "ltx" in class_lower or "ltx" in model_lower: + if "distilled" in class_lower or "distilled" in model_lower: + return _MODEL_PRESETS["ltx2_distilled"] + if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower: + return _MODEL_PRESETS["ltx23"] return _MODEL_PRESETS["ltx2"] - if "vace" in model_lower: + if "vace" in model_lower or "vace" in class_lower: return _MODEL_PRESETS["vace"] # Edge must be matched before the generic cosmos branch (its "cosmos" substring would # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset). if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower): return _MODEL_PRESETS["cosmos3_edge"] - if "cosmos" in model_lower: + if "cosmos" in model_lower or "cosmos" in class_lower: return _MODEL_PRESETS["cosmos"] - if "hunyuan" in model_lower: + if "hunyuan" in model_lower or "hunyuan" in class_lower: return _MODEL_PRESETS["hunyuan"] - if "helios" in model_lower: + if "helios" in model_lower or "helios" in class_lower: return _MODEL_PRESETS["helios"] return _MODEL_PRESETS["wan"] +def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]: + """Build the canonical request envelope for the shared T2V example.""" + result: dict[str, Any] = { + "prompt": prompt, + "modalities": ["video"], + } + if negative_prompt is not None: + result["negative_prompt"] = negative_prompt + return result + + +def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor: + """Normalize decoded frames according to the pipeline's output contract.""" + if not tensor.is_floating_point(): + return tensor + if source_range == "negative_one_to_one": + return tensor.clamp(-1, 1) * 0.5 + 0.5 + if source_range == "zero_to_one": + return tensor.clamp(0, 1) + raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}") + + def parse_profiler_config(value: str) -> dict[str, Any]: try: config = json.loads(value) @@ -160,7 +195,7 @@ def parse_extra_body(value: str) -> dict[str, Any]: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate a video from a text prompt. " - "Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models." + "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models." ) parser.add_argument( "--model", @@ -408,6 +443,7 @@ def main(): for key, default_val in preset.items(): if getattr(args, key.replace("-", "_"), None) is None: setattr(args, key.replace("-", "_"), default_val) + model_class_name = args.model_class_name generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) # Cache-dit config (Wan2.2 only) @@ -480,6 +516,7 @@ def main(): omni = Omni(**omni_kwargs) model_class_name = get_model_class_name(omni) or model_class_name declared_extra_body_params = get_extra_body_params(model_class_name) + output_tensor_range = get_output_tensor_range(model_class_name) if profiler_enabled: print("[Profiler] Starting profiling...") @@ -513,12 +550,11 @@ def main(): lora_path=lora_path, ) - prompt_dict = {"prompt": args.prompt} - if args.negative_prompt is not None: - prompt_dict["negative_prompt"] = args.negative_prompt - elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]): + negative_prompt = args.negative_prompt + if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")): # Preserve the historical empty-prompt behavior for non-LTX examples. - prompt_dict["negative_prompt"] = "" + negative_prompt = "" + prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt) extra_args = {} if lora_request: @@ -630,8 +666,7 @@ def main(): frame_tensor = frame_tensor[0] if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4): frame_tensor = frame_tensor.permute(1, 2, 0) - if frame_tensor.is_floating_point(): - frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5 + frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range) return frame_tensor.float().numpy() if isinstance(frame, np.ndarray): frame_array = frame @@ -683,8 +718,7 @@ def main(): video_tensor = video_tensor[0] elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): video_tensor = video_tensor.permute(1, 2, 3, 0) - if video_tensor.is_floating_point(): - video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5 + video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range) video_array = video_tensor.float().numpy() elif isinstance(frames, np.ndarray): video_array = frames diff --git a/examples/offline_inference/text_to_video/text_to_video_lingbot.py b/examples/offline_inference/text_to_video/text_to_video_lingbot.py deleted file mode 100644 index 108e2dc1f..000000000 --- a/examples/offline_inference/text_to_video/text_to_video_lingbot.py +++ /dev/null @@ -1,131 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import argparse -import time -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.") - parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b") - parser.add_argument("--prompt", default="a robotic arm picks up a red block") - parser.add_argument("--negative-prompt", default=None) - parser.add_argument("--output", default="lingbot_video_output.mp4") - parser.add_argument("--height", type=int, default=192) - parser.add_argument("--width", type=int, default=320) - parser.add_argument("--num-frames", type=int, default=9) - parser.add_argument("--num-inference-steps", type=int, default=2) - parser.add_argument("--guidance-scale", type=float, default=3.0) - parser.add_argument("--flow-shift", type=float, default=3.0) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--fps", type=int, default=24) - parser.add_argument("--output-type", choices=["pt", "np"], default="pt") - parser.add_argument("--batch-cfg", action="store_true") - return parser.parse_args() - - -def _peak_memory_mb(output: Any) -> float: - result = output[0] if isinstance(output, list) and output else output - return float(getattr(result, "peak_memory_mb", 0.0) or 0.0) - - -def _extract_video(output: Any) -> Any: - if isinstance(output, list): - output = output[0] if output else None - if isinstance(output, OmniRequestOutput): - if output.final_output_type != "image": - raise ValueError(f"Unexpected output type {output.final_output_type!r}.") - if not output.images: - raise ValueError("No video frames found in OmniRequestOutput.") - return output.images[0] if len(output.images) == 1 else output.images - return output - - -def _to_frame_array(video: Any) -> np.ndarray: - if isinstance(video, torch.Tensor): - tensor = video.detach().cpu() - if tensor.dim() == 5: - tensor = tensor[0] - if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4): - if tensor.shape[0] in (3, 4): - tensor = tensor.permute(1, 2, 3, 0) - elif tensor.shape[1] in (3, 4): - tensor = tensor.permute(0, 2, 3, 1) - if tensor.is_floating_point(): - if float(tensor.min()) < 0.0: - tensor = tensor.clamp(-1, 1) * 0.5 + 0.5 - else: - tensor = tensor.clamp(0, 1) - return tensor.float().numpy() - if isinstance(video, np.ndarray): - array = video[0] if video.ndim == 5 else video - if np.issubdtype(array.dtype, np.integer): - return array.astype(np.float32) / 255.0 - return array.astype(np.float32, copy=False) - if isinstance(video, list): - return np.stack([_to_frame_array(frame) for frame in video], axis=0) - return np.asarray(video, dtype=np.float32) - - -def main() -> None: - args = parse_args() - generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - omni = Omni( - model=args.model, - model_class_name="LingBotVideoPipeline", - flow_shift=args.flow_shift, - parallel_config=DiffusionParallelConfig(), - ) - - prompt: dict[str, Any] = { - "prompt": args.prompt, - "modalities": ["video"], - "num_frames": args.num_frames, - } - if args.negative_prompt is not None: - prompt["negative_prompt"] = args.negative_prompt - - sampling_params = OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - num_frames=args.num_frames, - num_inference_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - generator=generator, - output_type=args.output_type, - extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift}, - ) - - start = time.perf_counter() - output = omni.generate(prompt, sampling_params) - elapsed = time.perf_counter() - start - print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)") - peak_mb = _peak_memory_mb(output) - if peak_mb: - print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") - - video = _to_frame_array(_extract_video(output)) - output_path = Path(args.output) - output_path.parent.mkdir(parents=True, exist_ok=True) - from diffusers.utils import export_to_video - - export_to_video(list(video), str(output_path), fps=args.fps) - print(f"Saved generated video to {output_path}") - omni.close() - - -if __name__ == "__main__": - main() diff --git a/recipes/Robbyant/LingBot-Video.md b/recipes/Robbyant/LingBot-Video.md index 289d36990..29d112217 100644 --- a/recipes/Robbyant/LingBot-Video.md +++ b/recipes/Robbyant/LingBot-Video.md @@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends. - Dense checkpoint: - MoE checkpoint: - Upstream project: -- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py) +- Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py) - Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md) ## Hardware Support @@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights. ```bash CUDA_VISIBLE_DEVICES=0 \ -python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ +python examples/offline_inference/text_to_video/text_to_video.py \ --model robbyant/lingbot-video-dense-1.3b \ --prompt "a robotic arm picks up a red block" \ --output lingbot_t2v.mp4 \ @@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ ```bash CUDA_VISIBLE_DEVICES=0 \ -python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ +python examples/offline_inference/text_to_video/text_to_video.py \ --model robbyant/lingbot-video-moe-30b-a3b \ --prompt "a robotic arm picks up a red block" \ --output lingbot_moe_t2v.mp4 \ diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py new file mode 100644 index 000000000..ff7244d65 --- /dev/null +++ b/tests/examples/offline_inference/test_text_to_video_prompts.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest +import torch + +from examples.offline_inference.text_to_video.text_to_video import ( + _MODEL_PRESETS, + _detect_preset, + _normalize_float_tensor, + build_text_to_video_prompt, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("negative_prompt", "expected_negative_prompt"), + [(None, None), ("", ""), ("blurry", "blurry")], +) +def test_text_to_video_builds_canonical_prompt( + negative_prompt: str | None, + expected_negative_prompt: str | None, +) -> None: + result = build_text_to_video_prompt("a robot waves", negative_prompt) + + assert result["prompt"] == "a robot waves" + assert result["modalities"] == ["video"] + if expected_negative_prompt is None: + assert "negative_prompt" not in result + else: + assert result["negative_prompt"] == expected_negative_prompt + + +@pytest.mark.parametrize( + ("model", "model_class_name", "preset_name"), + [ + ("robbyant/lingbot-video-dense-1.3b", None, "lingbot"), + ("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"), + ("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"), + ("BestWishYsh/Helios-Distilled", None, "helios"), + ("/models/cosmos-edge", "Cosmos3EdgeVFMTransformer", "cosmos3_edge"), + ("/models/vace", "WanVACEPipeline", "vace"), + ], +) +def test_detect_preset_is_scoped_to_model_family( + model: str, + model_class_name: str | None, + preset_name: str, +) -> None: + assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name] + + +def test_lingbot_preset_matches_lingbot_defaults() -> None: + assert _MODEL_PRESETS["lingbot"] == { + "model_class_name": "LingBotVideoPipeline", + "height": 192, + "width": 320, + "num_frames": 9, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "fps": 24, + "flow_shift": 3.0, + "output": "lingbot_video_output.mp4", + } + + +def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None: + frames = torch.tensor([0.0, 0.25, 1.0, 1.5]) + + assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0])) + + +def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None: + frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5]) + + assert torch.equal( + _normalize_float_tensor(frames, "negative_one_to_one"), + torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]), + ) + + +def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None: + frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])] + + normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames] + + assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0])) + assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0])) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index fde3ea878..d57671adc 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -17,6 +17,7 @@ from vllm_omni.model_extras import ( build_x_to_text_prompt, get_extra_body_params, get_extra_output_params, + get_output_tensor_range, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, ) @@ -134,6 +135,8 @@ def test_lingbot_extra_registry_declares_request_params() -> None: } ) assert get_extra_output_params("LingBotVideoPipeline") == frozenset() + assert get_output_tensor_range("LingBotVideoPipeline") == "zero_to_one" + assert get_output_tensor_range("WanPipeline") == "negative_one_to_one" assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index 1a3d71870..b11ceadec 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -9,6 +9,7 @@ from vllm_omni.model_extras.registry import ( get_extra_body_params, get_extra_output_params, get_model_class_name, + get_output_tensor_range, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, ) @@ -21,6 +22,7 @@ __all__ = [ "get_extra_body_params", "get_extra_output_params", "get_model_class_name", + "get_output_tensor_range", "get_x_to_text_model_family", "should_init_extra_args_for_non_diffusion_stages", ] diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 4c41c5920..8d7ee1e29 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -4,7 +4,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping -from typing import Any +from typing import Any, Literal from PIL import Image @@ -93,6 +93,7 @@ ImageToVideoPromptBuilder = Callable[ dict[str, Any], ] XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]] +OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"] def default_x_to_text_prompt( @@ -193,6 +194,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { }, "LingBotVideoPipeline": { "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS, + "output_tensor_range": "zero_to_one", # Shared T2I/I2V envelopes select the output modality. LingBot's # pipeline owns model-specific validation and normalization. }, @@ -262,6 +264,18 @@ def get_extra_output_params(model_class_name: str | None) -> frozenset[str]: return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset() +def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange: + """Return the declared range for floating-point tensor outputs. + + The default preserves the shared examples' historical handling. Pipelines + that already return normalized tensors declare ``zero_to_one`` explicitly. + """ + spec = _get_spec(model_class_name) + if spec is None: + return "negative_one_to_one" + return spec.get("output_tensor_range", "negative_one_to_one") + + def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool: spec = _get_spec(model_class_name) return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) From 877beeaa4957e8caab0daad31930cb274cb601cb Mon Sep 17 00:00:00 2001 From: liwener Date: Thu, 13 Aug 2026 09:36:17 +0800 Subject: [PATCH 20/41] [Perf][MiniCPM-O-4.5]Add CUDA Graph For CFM DiT estimator (#6082) Signed-off-by: liwener Co-authored-by: amy-why-3459 --- .../minicpmo_4_5/test_cuda_graph_wrapper.py | 250 +++++++++++++++++- vllm_omni/deploy/minicpmo_4_5.yaml | 2 + vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 2 + vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 2 + vllm_omni/deploy/minicpmo_4_5_8x4090.yaml | 2 + .../models/minicpmo_4_5/batched_token2wav.py | 32 ++- .../models/minicpmo_4_5/cuda_graph_wrapper.py | 125 +++++++++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 5 + 8 files changed, 418 insertions(+), 2 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 16e04fc09..7ad35b86d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -1,21 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import lru_cache from types import SimpleNamespace from unittest.mock import Mock import pytest import torch import torch.nn as nn +from vllm.platforms import current_platform from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( HiFTGenerator, ) from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + CFMGraphWrapper, HiFTGraphWrapper, ) -pytestmark = [pytest.mark.core_model] +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] class _F0Predictor(nn.Module): @@ -142,3 +145,248 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) wrapper._capture.assert_not_called() wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) assert result is wrapper.decode_fn.return_value + + +# --------------------------------------------------------------------------- +# CFMGraphWrapper tests +# --------------------------------------------------------------------------- + + +class _MiniDiT(nn.Module): + """Minimal DiT-like module with blocks_forward_chunk for CFM graph testing. + + Mimics the upstream cosyvoice2 DiT's cache semantics: + - CausalConv1d.forward_chunk: cat([cnn_cache, x], dim=time) when cache is not None + - Attention.forward_chunk: cat([k, k_cache], dim=seq) when att_cache is not None + """ + + def __init__(self, in_dim: int = 16, hidden: int = 8, depth: int = 2) -> None: + super().__init__() + self.in_proj = nn.Linear(in_dim, hidden) + self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)]) + self.final_layer = nn.Linear(hidden, in_dim) + + def blocks_forward_chunk( + self, + x: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + cnn_cache: torch.Tensor | None = None, + att_cache: torch.Tensor | None = None, + cnn_cache_buffer: torch.Tensor | None = None, + att_cache_buffer: torch.Tensor | None = None, + ) -> torch.Tensor: + x = x.transpose(1, 2) + x = self.in_proj(x) + for b_idx in range(len(self.blocks)): + # Simulate CausalConv1d: cnn_cache contributes to the first frames. + # Upstream: if cnn_cache[b_idx] is None, creates zeros internally (same as zeros). + cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None + if cnn_b is not None: + x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2) + # Simulate Attention: att_cache contributes a bias when non-empty. + # Upstream: if att_cache[b_idx] is None, skips cat entirely (different path). + att_b = att_cache[b_idx] if att_cache is not None else None + if att_b is not None and att_b.shape[3] > 0: + x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1) + x = self.blocks[b_idx](x) + x = x + t + cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous() + dt = x.shape[1] + att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1) + x = self.final_layer(x) + x = x.transpose(1, 2) + return x + + +def _cfm_inputs( + batch_size: int, chunk_size: int, old_att_len: int, *, device: str = "cuda" +) -> tuple[torch.Tensor, ...]: + depth = 2 + hidden = 8 + estimator_input = torch.randn(batch_size, 16, chunk_size, device=device) + time_emb = torch.randn(batch_size, 1, hidden, device=device) + cnn_cache = torch.randn(depth, batch_size, hidden, 2, device=device) + att_cache = torch.randn(depth, batch_size, 1, old_att_len, hidden, device=device) + cnn_out = torch.empty(depth, batch_size, hidden, 2, device=device) + att_out = torch.empty(depth, batch_size, 1, old_att_len + chunk_size, hidden, device=device) + return estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32) + + with torch.inference_mode(): + for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)): + inputs = _cfm_inputs(2, chunk_size, old_att_len) + + eager_inputs = tuple(v.clone() for v in inputs) + with torch.no_grad(): + eager_result = estimator.blocks_forward_chunk( + eager_inputs[0], + eager_inputs[1], + None, + eager_inputs[2], + eager_inputs[3], + eager_inputs[4], + eager_inputs[5], + ) + + wrapper.replay(*inputs) + replay_inputs = tuple(v.clone() for v in inputs) + graph_result, graph_cnn, graph_att = wrapper.replay(*replay_inputs) + + torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5) + + wrapper._capture_graph.cache_clear() + + +def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = object.__new__(CFMGraphWrapper) + wrapper.max_graphs = max_graphs + wrapper.graph_fn = Mock(return_value=torch.tensor([42.0])) + wrapper.device = torch.device("cuda") + + @lru_cache(maxsize=max_graphs) + def _capture_graph(key: tuple): + return wrapper._do_capture(key) + + wrapper._capture_graph = _capture_graph + wrapper._do_capture = Mock(return_value=None) + return wrapper + + +def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _cfm_mock_wrapper(monkeypatch) + wrapper._capture_graph = Mock(return_value=None) + + inputs = _cfm_inputs(2, 10, 0) + wrapper.replay(*inputs) + + wrapper._capture_graph.assert_called_once() + wrapper.graph_fn.assert_called_once() + + +def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _cfm_mock_wrapper(monkeypatch) + wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None) + + inputs = _cfm_inputs(2, 10, 0) + result = wrapper.replay(*inputs) + + wrapper.graph_fn.assert_called_once() + assert result[0] is wrapper.graph_fn.return_value + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None: + """Verify that LRU cache evicts oldest entries when full.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2) + + with torch.inference_mode(): + inputs_a = _cfm_inputs(2, 10, 0) + wrapper.replay(*inputs_a) + assert wrapper._capture_graph.cache_info().currsize == 1 + + inputs_b = _cfm_inputs(2, 12, 0) + wrapper.replay(*inputs_b) + assert wrapper._capture_graph.cache_info().currsize == 2 + + wrapper.replay(*inputs_a) + assert wrapper._capture_graph.cache_info().currsize == 2 + + inputs_c = _cfm_inputs(2, 14, 0) + wrapper.replay(*inputs_c) + assert wrapper._capture_graph.cache_info().currsize == 2 + + wrapper.replay(*inputs_b) + hits_before = wrapper._capture_graph.cache_info().hits + wrapper.replay(*inputs_b) + hits_after = wrapper._capture_graph.cache_info().hits + assert hits_after == hits_before + 1 + + wrapper._capture_graph.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_none_cache_parity_between_graph_and_eager( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify that None→zeros (graph path) matches [None]*depth (eager path). + + setup_batch calls _decode_cfm with cnn_cache=None on every request's + prompt pass. The graph path replaces None with zeros; the eager path + passes [None]*depth. Both must produce identical outputs. + """ + torch.accelerator.synchronize() + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32) + + depth = len(estimator.blocks) + hidden = 8 + batch_size = 2 + chunk_size = 20 + + estimator_input = torch.randn(batch_size, 16, chunk_size, device="cuda") + time_emb = torch.randn(batch_size, 1, hidden, device="cuda") + cnn_out = torch.empty(depth, batch_size, hidden, 2, device="cuda") + att_out = torch.empty(depth, batch_size, 1, chunk_size, hidden, device="cuda") + + # Eager path: cnn_cache=[None]*depth, att_cache=[None]*depth + eager_input = estimator_input.clone() + eager_time = time_emb.clone() + eager_cnn_out = torch.empty_like(cnn_out) + eager_att_out = torch.empty_like(att_out) + with torch.no_grad(): + eager_result = estimator.blocks_forward_chunk( + eager_input, + eager_time, + None, + [None] * depth, + [None] * depth, + eager_cnn_out, + eager_att_out, + ) + + # Graph path: cnn_cache=zeros, att_cache=zero-length (simulating _estimator_step's None→zeros) + zero_cnn = torch.zeros_like(cnn_out) + zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + wrapper.replay(estimator_input, time_emb, zero_cnn, zero_att, cnn_out, att_out) + graph_input = estimator_input.clone() + graph_time = time_emb.clone() + graph_cnn_out = torch.empty_like(cnn_out) + graph_att_out = torch.empty_like(att_out) + graph_zero_cnn = torch.zeros_like(cnn_out) + graph_zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + graph_result, graph_cnn, graph_att = wrapper.replay( + graph_input, + graph_time, + graph_zero_cnn, + graph_zero_att, + graph_cnn_out, + graph_att_out, + ) + + torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 668954f5d..bc34bfe9f 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -19,6 +19,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index e44140b2d..2a4abed42 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -22,6 +22,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 4933d0c74..b34673910 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -22,6 +22,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 69fa6d5b6..813a2e2eb 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -33,6 +33,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 651284fe0..41223e864 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -14,7 +14,7 @@ import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger -from .cuda_graph_wrapper import HiFTGraphWrapper +from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper logger = init_logger(__name__) @@ -73,6 +73,7 @@ class BatchedToken2Wav(nn.Module): *, connector_config: Mapping[str, int] | None = None, hift_graph_config: Mapping[str, Any] | None = None, + cfm_graph_config: Mapping[str, Any] | None = None, ): super().__init__() self._token2wav = token2wav @@ -135,6 +136,22 @@ class BatchedToken2Wav(nn.Module): with torch.inference_mode(), _autocast_disabled(hift_parameter.device): self.hift_graph_wrapper.capture() logger.info("HiFT CUDA Graph captured successfully") + self._cfm_graph_wrapper: CFMGraphWrapper | None = None + cfm_graph_cfg = dict(cfm_graph_config or {}) + if bool(cfm_graph_cfg.get("enabled", False)): + flow_parameter = next(self.flow.parameters(), None) + if flow_parameter is not None and flow_parameter.device.type == "cuda": + estimator = self.flow.decoder.estimator + self._cfm_graph_wrapper = CFMGraphWrapper( + graph_fn=estimator.blocks_forward_chunk, + max_graphs=int(cfm_graph_cfg.get("max_graphs", 32)), + ) + logger.info("CFM CUDA Graph enabled (max_graphs=%d)", int(cfm_graph_cfg.get("max_graphs", 32))) + else: + logger.info( + "CFM CUDA Graph is disabled on device type %s", + flow_parameter.device.type if flow_parameter is not None else "unknown", + ) self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} def _hift_inference( @@ -265,6 +282,19 @@ class BatchedToken2Wav(nn.Module): speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) + if self._cfm_graph_wrapper is not None: + if cnn_cache is None: + cnn_cache = torch.zeros_like(cnn_out) + if att_cache is None: + att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + return self._cfm_graph_wrapper.replay( + estimator_input, + time_embedding, + cnn_cache, + att_cache, + cnn_out, + att_out, + ) old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) result = estimator.blocks_forward_chunk( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index 5f6c965d5..97015855e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import lru_cache + import torch from torch.cuda import CUDAGraph from vllm.logger import init_logger @@ -133,3 +135,126 @@ class HiFTGraphWrapper: cache_source = static_cache_source_output[:batch_size].clone() speech = self.finalize_fn(static_magnitude_output[:batch_size], static_phase_output[:batch_size]).clone() return speech, cache_source + + +def _tensor_signature(value: torch.Tensor) -> tuple: + return tuple(value.shape), str(value.dtype), str(value.device) + + +_DTYPE_MAP = { + "torch.float32": torch.float32, + "torch.float16": torch.float16, + "torch.bfloat16": torch.bfloat16, + "torch.float64": torch.float64, +} + + +def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]: + """Rebuild zero tensors from a cache key (shape, dtype, device tuples).""" + tensors = [] + for shape, dtype_str, device_str in key[1:]: + dtype = _DTYPE_MAP.get(dtype_str, torch.float32) + device = torch.device(device_str) + tensors.append(torch.zeros(shape, dtype=dtype, device=device)) + return tuple(tensors) + + +class CFMGraphWrapper: + """Per-shape CUDA graph capture/replay for the CFM DiT estimator. + + Captures one blocks_forward_chunk call (in_proj -> DiT blocks -> final_layer) + as the graph target. The 10-step Euler loop stays in Python, replaying + the graph 10 times per decode. + + Uses functools.lru_cache for automatic LRU eviction. Cache misses + trigger capture; capture failures fall back to eager. Outputs are + cloned after replay to prevent streaming cache corruption. + """ + + def __init__(self, graph_fn, *, max_graphs: int = 32) -> None: + self.graph_fn = graph_fn + self.max_graphs = int(max_graphs) + self.device = next(graph_fn.__self__.parameters()).device + + @lru_cache(maxsize=self.max_graphs) + def _capture_graph(key: tuple): + return self._capture(key) + + self._capture_graph = _capture_graph + + def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor: + return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5]) + + def _capture(self, key: tuple) -> tuple | None: + """Capture a CUDA graph for the given key. Returns None on failure.""" + static_inputs = _tensors_from_key(key) + + current_stream = torch.cuda.current_stream(self.device) + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(current_stream) + with torch.cuda.stream(warmup_stream), torch.no_grad(): + for _ in range(3): + self._call_graph_fn(static_inputs) + current_stream.wait_stream(warmup_stream) + + try: + graph = CUDAGraph() + with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_output = self._call_graph_fn(static_inputs) + except Exception: + logger.warning( + "CFM graph capture failed for shape=%s; using eager", + key, + exc_info=True, + ) + return None + + logger.info( + "Captured CFM CUDA Graph for shape %s (cache=%d/%d, hits=%d, misses=%d)", + key, + self._capture_graph.cache_info().currsize, + self.max_graphs, + self._capture_graph.cache_info().hits, + self._capture_graph.cache_info().misses, + ) + return (static_inputs, static_output, graph) + + def replay( + self, + estimator_input: torch.Tensor, + time_emb: torch.Tensor, + cnn_cache: torch.Tensor, + att_cache: torch.Tensor, + cnn_out: torch.Tensor, + att_out: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out) + + if torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda": + with torch.no_grad(): + result = self._call_graph_fn(inputs) + return result, inputs[4], inputs[5] + + key = ("estimator_step",) + tuple(_tensor_signature(v) for v in inputs) + entry = self._capture_graph(key) + + if entry is None: + logger.debug( + "CFM graph eager fallback for shape=%s (hits=%d, misses=%d)", + key, + self._capture_graph.cache_info().hits, + self._capture_graph.cache_info().misses, + ) + with torch.no_grad(): + result = self._call_graph_fn(inputs) + return result, inputs[4], inputs[5] + + static_inputs, static_output, graph = entry + for static, current in zip(static_inputs, inputs, strict=True): + static.copy_(current) + graph.replay() + return ( + static_output.detach().clone(), + static_inputs[4].detach().clone(), + static_inputs[5].detach().clone(), + ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index d6d3e3d47..686b39c5b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -162,6 +162,10 @@ class MiniCPMO45Code2Wav(nn.Module): "enabled": bool(extra.get("enable_hift_graph", False)), "capture_batch_sizes": capture_batch_sizes, } + self._cfm_graph_config = { + "enabled": bool(extra.get("enable_cfm_graph", False)), + "max_graphs": int(extra.get("cfm_max_graphs", 32)), + } self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) if self._min_batch_size < 1: raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") @@ -802,4 +806,5 @@ class MiniCPMO45Code2Wav(nn.Module): trt_stepper=trt_stepper, connector_config=self._connector_config, hift_graph_config=self._hift_graph_config, + cfm_graph_config=self._cfm_graph_config, ) From d77a35a32b25d1460ac5c72f66f4d9fd554399a2 Mon Sep 17 00:00:00 2001 From: NATURE Date: Thu, 13 Aug 2026 11:23:47 +0800 Subject: [PATCH 21/41] [Bugfix] Scope Bagel FP8 config to diffusion stage (#6085) Signed-off-by: natureofnature Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com> --- .../diffusion/quantization/test_quantization_fp8.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 581a7498f..617083398 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -189,7 +189,7 @@ def _generate_single_stage_video( def _generate_bagel_image( - quantization_config: str | None = None, + diffusion_quantization_config: str | None = None, num_inference_steps: int = 15, ) -> tuple[Any, float]: """Generate an image with BAGEL (multi-stage: LLM + Diffusion). @@ -202,8 +202,8 @@ def _generate_bagel_image( "stage_configs_path": config_path, "stage_init_timeout": 300, } - if quantization_config: - omni_kwargs["quantization_config"] = quantization_config + if diffusion_quantization_config: + omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config model_name = omni_kwargs.pop("model") with OmniRunner(model_name, **omni_kwargs) as runner: @@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory(): @hardware_test(res={"cuda": "H100"}) def test_bagel_fp8_generates_image(): - """BAGEL with FP8 quantization_config generates a valid image. + """BAGEL with diffusion-stage FP8 generates a valid image. FP8 should only apply to the diffusion stage (Stage-1), not the LLM stage (Stage-0). We verify this by checking: 1. Image is generated successfully 2. LLM stage finish_reason is 'stop' (not 'length' from garbled output) """ - image, _ = _generate_bagel_image(quantization_config="fp8") + image, _ = _generate_bagel_image(diffusion_quantization_config="fp8") image.save("test_bagel_fp8.png") @hardware_test(res={"cuda": "H100"}) def test_bagel_bf16_generates_image(): """BAGEL without quantization generates a valid image (baseline).""" - image, _ = _generate_bagel_image(quantization_config=None) + image, _ = _generate_bagel_image(diffusion_quantization_config=None) image.save("test_bagel_bf16.png") From c98517ce7fab78221e7806b56dd862a5fb69eae4 Mon Sep 17 00:00:00 2001 From: vraiti Date: Thu, 13 Aug 2026 02:59:13 -0400 Subject: [PATCH 22/41] [Metrics] Add Diffusion Metrics (#4755) Signed-off-by: vraiti Co-authored-by: Claude Opus 4.6 --- docs/design/metrics.md | 2 +- docs/usage/metrics.md | 76 ++++++--- .../test_diffusion_output_formatter.py | 1 + tests/metrics/test_modality.py | 150 +++++++++++++++++- vllm_omni/diffusion/output_formatter.py | 1 + vllm_omni/entrypoints/omni_base.py | 1 + vllm_omni/metrics/definitions.py | 9 ++ vllm_omni/metrics/modality.py | 135 ++++++++++++---- vllm_omni/metrics/stats.py | 17 +- 9 files changed, 328 insertions(+), 64 deletions(-) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 6e4946f4e..7ee35813f 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -1,6 +1,6 @@ # Prometheus Metrics Design -This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up. +This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. ## Objectives diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index 68a38043e..184a76eef 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -1,23 +1,23 @@ # Production Metrics -vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR. +vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. ```bash vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats curl http://localhost:8000/metrics ``` -**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate. +**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm_omni:*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate. ## Metric Namespaces | Prefix | Source | Present when | |--------|--------|--------------| -| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent | +| `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent | | `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage | | `http_` / `process_` | Uvicorn / Python runtime | Always | -## Pipeline-Level Metrics (`vllm:omni_`) +## Pipeline-Level Metrics (`vllm_omni:`) Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline. @@ -25,42 +25,66 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the | Metric | Type | Labels | Description | |--------|------|--------|-------------| -| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) | -| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue | -| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | +| `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) | +| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue | +| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | ### Latency | Metric | Type | Labels | Description | |--------|------|--------|-------------| -| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds | +| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds | -## Audio Modality Metrics (`vllm:omni_`) +## Audio Modality Metrics (`vllm_omni:`) Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label. | Metric | Type | Extra label | Description | |--------|------|-------------|-------------| -| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | -| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | -| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` | -| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` | -| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | -| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | -| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | +| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | +| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | +| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` | +| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` | +| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | +| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | +| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition. -## Cross-Stage Transfer Metrics (`vllm:omni_`) +## Diffusion Metrics (`vllm_omni:`) + +Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`. + +| Metric | Type | Description | +|--------|------|-------------| +| `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds | +| `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) | +| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds | +| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing (VAE decode) time per request in seconds | + +`diffusion_exec_s` uses `SECONDS_BUCKETS` (50 ms → 300 s); the other three use `SECONDS_FAST_BUCKETS` (1 ms → 60 s). `diffusion_exec_per_step_s` is skipped when `num_inference_steps` is not set. + +### Example PromQL + +```promql +# P99 DiT exec time by stage +histogram_quantile(0.99, rate(vllm_omni:diffusion_exec_s_bucket[5m])) + +# Ratio of preprocess to total exec wall time +rate(vllm_omni:diffusion_preprocess_s_sum[5m]) + / rate(vllm_omni:diffusion_exec_s_sum[5m]) +``` + +## Cross-Stage Transfer Metrics (`vllm_omni:`) Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed. | Metric | Type | Description | |--------|------|-------------| -| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | -| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | -| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | -| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | +| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | +| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | +| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | +| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | ## vLLM Engine Metrics (`vllm:`) @@ -83,9 +107,9 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl | Metric group | Multi-stage LLM (Qwen3-Omni) | |---|---| -| `vllm:omni_` request tracking + latency | With `--log-stats` | -| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage | -| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages | +| `vllm_omni:` request tracking + latency | With `--log-stats` | +| `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage | +| `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages | | `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` | | `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` | @@ -94,5 +118,5 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl - All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values. - Counters use the `_total` suffix (auto-appended by `prometheus_client`). - Sizes use the `_bytes` suffix. -- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`. -- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label. +- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm_omni:`. +- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm_omni:audio_ttfp_s` for audio) rather than a single metric with a `modality` label. diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py index eea53210c..dc21351ff 100644 --- a/tests/diffusion/test_diffusion_output_formatter.py +++ b/tests/diffusion/test_diffusion_output_formatter.py @@ -95,6 +95,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata( "resolution": 512, "width": None, "height": None, + "num_inference_steps": 1, } diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py index ba1785639..315594233 100644 --- a/tests/metrics/test_modality.py +++ b/tests/metrics/test_modality.py @@ -42,6 +42,10 @@ _EXPECTED_FAMILIES = [ defs.AUDIO_UNDERRUN_S, defs.AUDIO_CONTINUITY_OK_METRIC, defs.AUDIO_SKIPPED_REQUESTS_METRIC, + defs.DIFFUSION_EXEC_S, + defs.DIFFUSION_EXEC_PER_STEP_S, + defs.DIFFUSION_PREPROCESS_S, + defs.DIFFUSION_POSTPROCESS_S, ] @@ -55,6 +59,10 @@ class TestRegistration: mod.observe_audio_underrun("s", "r", 0.01) mod.inc_audio_continuity_ok("s", "r", 100) mod.inc_audio_skipped("s", "r", "malformed_codec") + mod.observe_diffusion_exec("s", "r", 0.5) + mod.observe_diffusion_exec_per_step("s", "r", 0.01) + mod.observe_diffusion_preprocess("s", "r", 0.01) + mod.observe_diffusion_postprocess("s", "r", 0.2) out = generate_latest(REGISTRY).decode() for name in _EXPECTED_FAMILIES: @@ -156,6 +164,50 @@ class TestAudio: assert unknown == 1.0 +# --------------------------------------------------------------------------- +# Diffusion observe API +# --------------------------------------------------------------------------- + + +class TestDiffusion: + def test_diffusion_exec_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "dit_exec", "0" + mod.observe_diffusion_exec(stage, replica, 1.25) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.DIFFUSION_EXEC_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 1.25 + + def test_diffusion_preprocess_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "dit_pre", "0" + mod.observe_diffusion_preprocess(stage, replica, 0.03) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.DIFFUSION_PREPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == pytest.approx(0.03) + + def test_diffusion_postprocess_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "dit_post", "0" + mod.observe_diffusion_postprocess(stage, replica, 0.8) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.DIFFUSION_POSTPROCESS_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == 0.8 + + def test_diffusion_exec_per_step_observed(self, mod: OmniModalityMetrics) -> None: + stage, replica = "dit_per_step", "0" + mod.observe_diffusion_exec_per_step(stage, replica, 0.078) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.DIFFUSION_EXEC_PER_STEP_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == pytest.approx(0.078) + + def test_log_stats_false_skips(self) -> None: + silent = OmniModalityMetrics(model_name="silent-model", log_stats=False) + silent.observe_diffusion_exec("s", "r", 1.0) + silent.observe_diffusion_exec_per_step("s", "r", 0.02) + silent.observe_diffusion_preprocess("s", "r", 0.01) + silent.observe_diffusion_postprocess("s", "r", 0.5) + out = generate_latest(REGISTRY).decode() + assert 'model_name="silent-model"' not in out + + # --------------------------------------------------------------------------- # Stub-driven routing tests for the finalize / streaming helpers. # --------------------------------------------------------------------------- @@ -183,6 +235,18 @@ class _StubModMetrics: def inc_audio_skipped(self, s, r, reason): self.calls.append(("inc_audio_skipped", s, r, reason)) + def observe_diffusion_exec(self, s, r, v): + self.calls.append(("observe_diffusion_exec", s, r, v)) + + def observe_diffusion_exec_per_step(self, s, r, v): + self.calls.append(("observe_diffusion_exec_per_step", s, r, v)) + + def observe_diffusion_preprocess(self, s, r, v): + self.calls.append(("observe_diffusion_preprocess", s, r, v)) + + def observe_diffusion_postprocess(self, s, r, v): + self.calls.append(("observe_diffusion_postprocess", s, r, v)) + class _Bag: def __init__(self, **kw): @@ -239,8 +303,88 @@ class TestObserveModalityAtFinalize: ) assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls - def test_non_audio_output_type_skipped(self): - # Image / video / text output types are out of scope for this module. + def test_diffusion_path_full(self): + stub = _StubModMetrics() + stage_metrics = _Bag( + stage_gen_time_ms=1000.0, + diffusion_metrics={ + "preprocess_time_s": 0.05, + "diffusion_engine_exec_time_s": 1.2, + "postprocess_time_s": 0.3, + "image_num": 1, + "resolution": 1024, + "num_inference_steps": 50, + }, + ) + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=stage_metrics, + engine_outputs=_Bag(), + ) + assert ("observe_diffusion_preprocess", "2", "0", 0.05) in stub.calls + assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls + assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls + assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls + + def test_diffusion_per_step_skipped_without_num_inference_steps(self): + stub = _StubModMetrics() + stage_metrics = _Bag( + stage_gen_time_ms=1000.0, + diffusion_metrics={ + "diffusion_engine_exec_time_s": 1.2, + "preprocess_time_s": 0.05, + "postprocess_time_s": 0.3, + }, + ) + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=stage_metrics, + engine_outputs=_Bag(), + ) + assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls + assert not any(c[0] == "observe_diffusion_exec_per_step" for c in stub.calls) + + def test_diffusion_metrics_none_skipped(self): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=_Bag(stage_gen_time_ms=500.0, diffusion_metrics=None), + engine_outputs=_Bag(), + ) + assert not any(c[0].startswith("observe_diffusion") for c in stub.calls) + + def test_diffusion_zero_values_observed(self): + stub = _StubModMetrics() + stage_metrics = _Bag( + stage_gen_time_ms=500.0, + diffusion_metrics={ + "preprocess_time_s": 0.0, + "diffusion_engine_exec_time_s": 0.0, + "postprocess_time_s": 0.0, + }, + ) + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=stage_metrics, + engine_outputs=_Bag(), + ) + assert ("observe_diffusion_preprocess", "2", "0", 0.0) in stub.calls + assert ("observe_diffusion_exec", "2", "0", 0.0) in stub.calls + assert ("observe_diffusion_postprocess", "2", "0", 0.0) in stub.calls + + def test_non_audio_output_type_skips_audio_but_still_checks_diffusion(self): stub = _StubModMetrics() for output_type in ("text", "image", "video"): observe_modality_at_finalize( @@ -251,7 +395,7 @@ class TestObserveModalityAtFinalize: stage_metrics=_Bag(stage_gen_time_ms=100.0), engine_outputs=_Bag(), ) - assert stub.calls == [] + assert not any(c[0].startswith("observe_audio") or c[0].startswith("inc_audio") for c in stub.calls) def test_replica_id_none_skipped(self): stub = _StubModMetrics() diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index d5b03443c..4e68bfd25 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -138,6 +138,7 @@ def format_diffusion_outputs( ), "width": request.sampling_params.width, "height": request.sampling_params.height, + "num_inference_steps": request.sampling_params.num_inference_steps, } # Only the primary payload determines the response type. Some image models diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 88e5a5edf..b787c44a4 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -568,6 +568,7 @@ class OmniBase(PDDisaggregationMixin): msg_id = id(result) consumed = self._consumed_metric_message_ids(req_id) if msg_id not in consumed: + metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs) metrics.on_stage_metrics(stage_id, req_id, _m, output_type) consumed.add(msg_id) diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 23a1364a4..b5d5217f5 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -134,6 +134,15 @@ AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS +# ============================================================================ +# Diffusion family (per-stage + per-replica diffusion timing breakdowns) +# ============================================================================ +DIFFUSION_EXEC_S = METRIC_PREFIX + "diffusion_exec_s" +DIFFUSION_EXEC_PER_STEP_S = METRIC_PREFIX + "diffusion_exec_per_step_s" +DIFFUSION_PREPROCESS_S = METRIC_PREFIX + "diffusion_preprocess_s" +DIFFUSION_POSTPROCESS_S = METRIC_PREFIX + "diffusion_postprocess_s" + + # ============================================================================ # Cross-stage Transfer family (per-physical-hop TX/RX/in-flight timings) # ============================================================================ diff --git a/vllm_omni/metrics/modality.py b/vllm_omni/metrics/modality.py index fd455ad36..b83912ed3 100644 --- a/vllm_omni/metrics/modality.py +++ b/vllm_omni/metrics/modality.py @@ -78,6 +78,35 @@ _audio_skipped_family = Counter( ) +# ---------------------------------------------------------------------------- +# Diffusion family +# ---------------------------------------------------------------------------- +_diffusion_exec_family = Histogram( + defs.DIFFUSION_EXEC_S, + "DiT forward pass execution time per request in seconds.", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_diffusion_exec_per_step_family = Histogram( + defs.DIFFUSION_EXEC_PER_STEP_S, + "DiT forward pass execution time per denoising step in seconds.", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) +_diffusion_preprocess_family = Histogram( + defs.DIFFUSION_PREPROCESS_S, + "Diffusion input preprocessing time per request in seconds.", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) +_diffusion_postprocess_family = Histogram( + defs.DIFFUSION_POSTPROCESS_S, + "Diffusion output postprocessing (VAE decode) time per request in seconds.", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) + + class OmniModalityMetrics: """Per-modality observe API. Stage/replica are passed at observe time because a single OmniModalityMetrics instance per pipeline serves all @@ -139,6 +168,44 @@ class OmniModalityMetrics: reason=reason or "unknown", ).inc() + # ---- Diffusion -------------------------------------------------------- + + def observe_diffusion_exec(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats: + return + _diffusion_exec_family.labels( + model_name=self._model_name, + stage=stage, + replica=replica, + ).observe(seconds) + + def observe_diffusion_exec_per_step(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats: + return + _diffusion_exec_per_step_family.labels( + model_name=self._model_name, + stage=stage, + replica=replica, + ).observe(seconds) + + def observe_diffusion_preprocess(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats: + return + _diffusion_preprocess_family.labels( + model_name=self._model_name, + stage=stage, + replica=replica, + ).observe(seconds) + + def observe_diffusion_postprocess(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats: + return + _diffusion_postprocess_family.labels( + model_name=self._model_name, + stage=stage, + replica=replica, + ).observe(seconds) + def observe_modality_at_finalize( mod_metrics: OmniModalityMetrics, @@ -161,43 +228,49 @@ def observe_modality_at_finalize( """ if replica_id is None or stage_metrics is None or output_type is None: return - if output_type != "audio": - return stage_label = str(stage_id) replica_label = str(replica_id) gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0 mm_out = extract_mm_output(engine_outputs) - sample_rate = defs.resolve_audio_sample_rate(mm_out) - # `stage_metrics.audio_generated_frames` is the legacy per-chunk - # accumulator field on StageRequestStats. No production path currently - # fills it, so the fallback below is the live source — but we leave the - # field lookup in place in case the accumulator gets re-wired upstream. - n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0) - if n_frames == 0: - n_frames = count_audio_frames(mm_out) - mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames) - duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0 - if duration_s > 0: - mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s) - mod_metrics.observe_audio_rtf( - stage_label, - replica_label, - defs.compute_audio_rtf(gen_time_s, duration_s), - ) - else: - # Request completed (finish_reason ∈ {stop, length} — error paths - # don't reach finalize) but no audio samples were produced. Covers - # silent `return None` skips in the talker→code2wav stage - # processors and the `parsed.append((0,0))` malformed-length path - # in qwen3-tts code2wav. raise-paths surface via the upstream - # vllm:request_success_total{finished_reason="error"} channel and - # never reach this branch. - mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") - # audio_underrun / continuity are emitted from the streaming path in - # observe_audio_streaming_finalize; finalize is too late for the - # per-chunk timeline they need. + if output_type == "audio": + sample_rate = defs.resolve_audio_sample_rate(mm_out) + n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0) + if n_frames == 0: + n_frames = count_audio_frames(mm_out) + mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames) + duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0 + if duration_s > 0: + mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s) + mod_metrics.observe_audio_rtf( + stage_label, + replica_label, + defs.compute_audio_rtf(gen_time_s, duration_s), + ) + else: + mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") + + dm = getattr(stage_metrics, "diffusion_metrics", None) + if dm: + _key_map = { + "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec, + "preprocess_time_s": mod_metrics.observe_diffusion_preprocess, + "postprocess_time_s": mod_metrics.observe_diffusion_postprocess, + } + for key, observe_fn in _key_map.items(): + val = dm.get(key) + if val is not None: + observe_fn(stage_label, replica_label, float(val)) + + exec_time = dm.get("diffusion_engine_exec_time_s") + num_steps = dm.get("num_inference_steps") + if exec_time is not None and num_steps and num_steps > 0: + mod_metrics.observe_diffusion_exec_per_step( + stage_label, + replica_label, + float(exec_time) / int(num_steps), + ) def observe_audio_first_packet( diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 3d08adae0..95198ef6a 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -44,7 +44,7 @@ class StageRequestStats: final_output_type: str | None = None request_id: str | None = None postprocess_time_ms: float = 0.0 - diffusion_metrics: dict[str, int] = None + diffusion_metrics: dict[str, float] = None audio_generated_frames: int = 0 audio_sample_rate: int = 0 audio_duration_s: float = 0.0 @@ -579,7 +579,7 @@ class OrchestratorAggregator: if final_output_type is not None: stats.final_output_type = final_output_type stats.diffusion_metrics = ( - {k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} + {k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} if req_id in self.diffusion_metrics else None ) @@ -644,8 +644,19 @@ class OrchestratorAggregator: if isinstance(diffusion_metrics, list): diffusion_metrics = diffusion_metrics[0] if diffusion_metrics: + _MS_TO_S = { + "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s", + "preprocess_time_ms": "preprocess_time_s", + "postprocess_time_ms": "postprocess_time_s", + "diffusion_engine_total_time_ms": "diffusion_engine_total_time_s", + } for key, value in diffusion_metrics.items(): - self.diffusion_metrics[req_id][key] += value + if value is None: + continue + if key in _MS_TO_S: + self.diffusion_metrics[req_id][_MS_TO_S[key]] = float(value) / 1000.0 + else: + self.diffusion_metrics[req_id][key] = value def on_forward( self, From bbb709cfb56358500cbd89a9cca0b3a24e66dbe5 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 13 Aug 2026 16:56:26 +0800 Subject: [PATCH 23/41] [BugFix] Use MOSS-TTS-Local official sample params (#6156) Signed-off-by: Canlin Guo --- vllm_omni/deploy/moss_tts_local.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 2ed937e0a..5f0bf9a14 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -31,12 +31,12 @@ stages: output_connectors: to_stage_1: shm default_sampling_params: - temperature: 1.0 - top_p: 0.95 - top_k: 50 + temperature: 1.7 + top_p: 0.8 + top_k: 25 max_tokens: 4096 seed: 42 - repetition_penalty: 1.1 + repetition_penalty: 1.0 - stage_id: 1 max_num_seqs: 64 From bbe6ccc512a404a2df8c977ea29003002f2683e8 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 13 Aug 2026 17:04:06 +0800 Subject: [PATCH 24/41] [Doc] Refresh community documentation (#6141) Signed-off-by: Hongsheng Liu --- docs/.nav.yml | 3 ++- docs/community/contact_us.md | 10 +++++++--- docs/community/governance.md | 26 +++----------------------- docs/community/meetings.md | 8 ++++++++ docs/community/meetups.md | 1 - docs/community/volunteers.md | 12 ------------ docs/contributing/README.md | 4 ++-- docs/usage/faq.md | 6 +++--- mkdocs.yml | 2 ++ 9 files changed, 27 insertions(+), 45 deletions(-) create mode 100644 docs/community/meetings.md delete mode 100644 docs/community/meetups.md delete mode 100644 docs/community/volunteers.md diff --git a/docs/.nav.yml b/docs/.nav.yml index bdc11ec7c..730c43be6 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -163,7 +163,8 @@ nav: - CLI Reference: cli - Community: - Governance: community/governance.md - - community/* + - Contact Us: community/contact_us.md + - Community Meetings: community/meetings.md - Slack: https://slack.vllm.ai - Blog: https://blog.vllm.ai - Forum: https://discuss.vllm.ai diff --git a/docs/community/contact_us.md b/docs/community/contact_us.md index 09c7815a0..cae175f4c 100644 --- a/docs/community/contact_us.md +++ b/docs/community/contact_us.md @@ -1,5 +1,9 @@ # Contact Us -- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues) -- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/) -- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature +- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues). +- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/). +- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md). + +## Security Disclosures + +Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request. diff --git a/docs/community/governance.md b/docs/community/governance.md index 63d38cd20..ae12ef038 100644 --- a/docs/community/governance.md +++ b/docs/community/governance.md @@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis - [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp - [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support -## Diffusion Workload Division +## Reviewer Routing -This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review. +For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents). -### Runtime-related - -- Diffusion models: @RuixiangMa, @wtomin, @xuechendi - -### Optimization - -- Parallel: @RuixiangMa, @wtomin -- Attention backend: @ZJY0516 -- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang -- Offload: @yuanheng-zhao -- Intel GPU support: @xuechendi -- Quantization: @lishunyang12, @david6666666 -- Kernel fusion / communication-computation: @ZJY0516 - -### Docs & Test - -- Diffusion docs: @RuixiangMa, @wtomin - -## Meetings - -Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project. +`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights. ## Committer Nomination Process diff --git a/docs/community/meetings.md b/docs/community/meetings.md new file mode 100644 index 000000000..7f812a599 --- /dev/null +++ b/docs/community/meetings.md @@ -0,0 +1,8 @@ +# Community Meetings + +vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth: + +- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0) +- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0) + +Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/). diff --git a/docs/community/meetups.md b/docs/community/meetups.md deleted file mode 100644 index 3374fe711..000000000 --- a/docs/community/meetups.md +++ /dev/null @@ -1 +0,0 @@ -# Meetups diff --git a/docs/community/volunteers.md b/docs/community/volunteers.md deleted file mode 100644 index 2c25485ea..000000000 --- a/docs/community/volunteers.md +++ /dev/null @@ -1,12 +0,0 @@ -# Volunteers for Bugfix and CI - -We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing. - -If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below. - -| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 | -|----------|----------|----------|----------|----------|----------| -| Conw729 | yinpeiqi | tzhouam | SamitHuang | gcanlin | natureofnature | -| david6666666 | R2-Y | hsliuustc0106 | Gaohan123 | ZJY0516 | qibaoyuan | - -We kindly welcome more contributors to fix bugs and contribute new features! diff --git a/docs/contributing/README.md b/docs/contributing/README.md index f86a18f2e..2224c667e 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -3,7 +3,7 @@ Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing. !!! note - We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting). + vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. See [Community Meetings](../community/meetings.md) for current schedules, access details, agendas, and past notes. ## Getting Started @@ -122,7 +122,7 @@ Repository skills guide the agent, but they do not replace contributor judgment, If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible. !!! important - If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label. + Do not report suspected security vulnerabilities through a public issue, pull request, discussion, or Slack channel. Follow the [security disclosure instructions](../community/contact_us.md#security-disclosures) to arrange a private report. ## Pull Requests & Code Reviews diff --git a/docs/usage/faq.md b/docs/usage/faq.md index 0539e158b..87b63998e 100644 --- a/docs/usage/faq.md +++ b/docs/usage/faq.md @@ -10,12 +10,12 @@ A: Refer to [GPU memory calculation and configuration](../configuration/gpu_memo > Q: I encounter some bugs or CI problems, which is urgent. How can I solve it? -A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help. +A: Search the current [issues](https://github.com/vllm-project/vllm-omni/issues) first. If no existing issue matches a bug or CI failure, open a [bug report](https://github.com/vllm-project/vllm-omni/issues/new?template=400-bug-report.yml) with a minimal reproducer, environment details, and relevant logs. For technical or usage questions, use the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or the [vLLM Forum](https://discuss.vllm.ai/). > Q: Does vLLM-Omni support AWQ or any other quantization? -A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057). +A: AWQ is available through the inherited vLLM quantization registry, and vLLM-Omni supports several additional methods. Actual support and validation vary by model, component, hardware, and checkpoint. See the [quantization overview](../user_guide/quantization/overview.md) for validated combinations and method-specific guides. > Q: Does vLLM-Omni support multimodal streaming input and output? -A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned! +A: Yes, for specific models and serving paths. Current capabilities include [streaming video input](../serving/video_stream_api.md), [streaming diffusion output](../user_guide/diffusion/execution_modes.md#streaming-output), and experimental full-duplex real-time audio described in the [project overview](../README.md). Check the model and API documentation because input and output support is model-specific. diff --git a/mkdocs.yml b/mkdocs.yml index 242eb33ae..1a3b5109f 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -73,6 +73,8 @@ plugins: - awesome-nav - redirects: redirect_maps: + community/meetups.md: community/meetings.md + community/volunteers.md: community/contact_us.md user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md From f41451f64c8b2fc20f8296f4c2a12c7d3bb16181 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Thu, 13 Aug 2026 20:01:01 +0800 Subject: [PATCH 25/41] [NPU] Upgrade to v0.27.0 (#6096) Signed-off-by: Weiming Liao --- docker/Dockerfile.npu | 2 +- docker/Dockerfile.npu.a3 | 2 +- docker/Dockerfile.npu.ci | 2 +- docker/Dockerfile.npu.ci.a3 | 2 +- .../npu/worker/npu_ar_model_runner.py | 7 +++-- .../npu/worker/npu_generation_model_runner.py | 11 ++++--- .../platforms/npu/worker/npu_model_runner.py | 31 ++++++++++--------- 7 files changed, 31 insertions(+), 26 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 05a780542..eaeef693c 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.26.0 +ARG VLLM_ASCEND_TAG=v0.27.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index d226a505f..2dc6334d9 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.26.0-a3 +ARG VLLM_ASCEND_TAG=v0.27.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 848a52f24..f5af5fb31 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.26.0 +ARG VLLM_ASCEND_TAG=v0.27.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 index 7ff709ee8..0529cb2da 100644 --- a/docker/Dockerfile.npu.ci.a3 +++ b/docker/Dockerfile.npu.ci.a3 @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.26.0-a3 +ARG VLLM_ASCEND_TAG=v0.27.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 0ff7144a9..63f3912f7 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -730,7 +730,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0 # Run forward pass - clear_kv_metadata = self.speculative_config is None + defer_kv_connector_finalize = self.speculative_config is not None and ( + get_pp_group().is_last_rank or self.broadcast_pp_output + ) with ( record_function_or_nullcontext("forward"), set_ascend_forward_context( @@ -744,13 +746,12 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): model_instance=self.model, skip_compiled=has_encoder_input, has_sinks=self._has_sinks, - input_ids=input_ids, eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ), self.maybe_get_kv_connector_output( scheduler_output, **( - {"defer_finalize": not clear_kv_metadata} + {"defer_finalize": defer_kv_connector_finalize} ), ) as kv_connector_output, ): diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index b728b3cc3..94c6f8916 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -381,7 +381,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0 # Run forward pass - clear_kv_metadata = self.speculative_config is None + defer_kv_connector_finalize = self.speculative_config is not None and ( + get_pp_group().is_last_rank or self.broadcast_pp_output + ) with ( record_function_or_nullcontext("forward"), set_ascend_forward_context( @@ -395,12 +397,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin model_instance=self.model, skip_compiled=has_encoder_input, has_sinks=self._has_sinks, - input_ids=input_ids, eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ), self.maybe_get_kv_connector_output( scheduler_output, - **({"defer_finalize": not clear_kv_metadata}), + **({"defer_finalize": defer_kv_connector_finalize}), ) as kv_connector_output, ): if self.cache_config.mamba_cache_mode == "align": @@ -654,9 +655,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin num_active_loras: int = 0, profile_seq_lens: int | None = None, profile_cpp: bool = False, + randomize_inputs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now - assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode() # If cudagraph_mode.decode_mode() == FULL and # cudagraph_mode.separate_routine(). This means that we are using # different graphs and/or modes for mixed prefill-decode batches vs. @@ -917,7 +919,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin batch_descriptor=batch_desc, model_instance=self.model, has_sinks=self._has_sinks, - input_ids=input_ids, eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ): outputs = self.model( diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index bc8e4f0a2..375d36047 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -125,9 +125,10 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): num_active_loras: int = 0, profile_seq_lens: int | None = None, profile_cpp: bool = False, + randomize_inputs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now - assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode() # If cudagraph_mode.decode_mode() == FULL and # cudagraph_mode.separate_routine(). This means that we are using # different graphs and/or modes for mixed prefill-decode batches vs. @@ -367,19 +368,21 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"): return self.drafter.model.compute_logits(hidden_states[dummy_indices]) - with set_ascend_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens_padded, - num_tokens_across_dp=num_tokens_across_dp, - in_profile_run=is_profile, - num_actual_tokens=num_tokens_padded, - aclgraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_desc, - model_instance=self.model, - has_sinks=self._has_sinks, - input_ids=input_ids, - eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, + with ( + self.maybe_randomize_inputs(input_ids, inputs_embeds, randomize_inputs=randomize_inputs), + set_ascend_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, + in_profile_run=is_profile, + num_actual_tokens=num_tokens_padded, + aclgraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, + model_instance=self.model, + has_sinks=self._has_sinks, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, + ), ): # ---------------------------------------Omni-new---------------------------------------------- if getattr(self.model, "talker", None) is not None and self.has_talker_mtp: From 4b66f62e10f1b765dc6dc6baa5bbfe633234ec25 Mon Sep 17 00:00:00 2001 From: Yukim1 <121286183+zwhzzz0821@users.noreply.github.com> Date: Thu, 13 Aug 2026 21:29:03 +0800 Subject: [PATCH 26/41] [Refactor] Remove internal stage_configs_path plumbing (#5741) Signed-off-by: zwhzzz0821 <2831474076@qq.com> Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- .../offline_inference/covo_audio/end2end.py | 6 +- .../offline_inference/dynin_omni/end2end.py | 6 +- .../hunyuan_image3/end2end.py | 19 +-- .../offline_inference/minicpmo/end2end.py | 6 - .../offline_inference/qwen3_omni/end2end.py | 4 +- .../offline_inference/step_audio2/end2end.py | 22 +-- .../text_to_image/text_to_image.py | 12 +- .../text_to_speech/fish_speech/end2end.py | 13 +- .../text_to_speech/glm_tts/end2end.py | 2 +- .../text_to_speech/qwen3_tts/end2end.py | 4 +- tests/config/test_config_factory.py | 52 +------ .../quantization/test_quantization_fp8.py | 2 +- .../test_bagel_mooncake_connector.py | 2 +- .../test_bagel_shared_memory_connector.py | 4 +- .../test_audiox_model_expansion.py | 2 +- .../offline_inference/test_bagel_expansion.py | 2 +- .../test_cosyvoice3_expansion.py | 2 +- .../test_ernie_image_expansion.py | 2 +- .../test_glm_tts_expansion.py | 2 +- .../test_hunyuanimage3_i2t_expansion.py | 2 +- .../test_longcat_video_avatar.py | 1 - .../test_mammoth_moda2_expansion.py | 2 +- .../offline_inference/test_minicpmo_4_5.py | 2 +- .../test_moss_tts_nano_expansion.py | 2 +- .../test_omnivoice_expansion.py | 6 +- .../test_sensenova_u1_expansion.py | 2 +- .../test_sensenova_u1_img2img_expansion.py | 6 +- .../test_sensenova_u1_text2img_expansion.py | 6 +- .../test_soulxsinger_expansion.py | 4 +- .../test_step_audio2_expansion.py | 2 +- .../test_voxtral_tts_expansion.py | 2 +- tests/engine/test_arg_utils.py | 139 ------------------ tests/engine/test_async_omni_engine_abort.py | 2 +- tests/engine/test_single_stage_mode.py | 4 +- tests/entrypoints/test_async_omni.py | 2 +- .../test_async_omni_diffusion_config.py | 6 +- tests/entrypoints/test_serve.py | 2 +- tests/entrypoints/test_utils.py | 16 +- tests/helpers/runtime.py | 13 +- tools/baselines/omni_hash_baseline.py | 2 +- vllm_omni/config/omni_config.py | 1 - vllm_omni/config/stage_config.py | 43 +----- vllm_omni/engine/arg_utils.py | 6 - vllm_omni/engine/async_omni_engine.py | 93 +----------- vllm_omni/entrypoints/async_omni.py | 5 +- vllm_omni/entrypoints/cli/serve.py | 2 - vllm_omni/entrypoints/openai/api_server.py | 4 +- vllm_omni/entrypoints/utils.py | 89 ++--------- 48 files changed, 93 insertions(+), 537 deletions(-) diff --git a/examples/offline_inference/covo_audio/end2end.py b/examples/offline_inference/covo_audio/end2end.py index 5b7ac5fef..b5050f62c 100644 --- a/examples/offline_inference/covo_audio/end2end.py +++ b/examples/offline_inference/covo_audio/end2end.py @@ -59,7 +59,7 @@ def main(args): omni = Omni( model=args.model_name, - stage_configs_path=args.stage_configs_path, + deploy_config=args.deploy_config, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, batch_timeout=args.batch_timeout, @@ -167,10 +167,10 @@ def parse_args(): help="Sampling rate for audio loading (default: 16000).", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, default=None, - help="Path to stage configs YAML file.", + help="Path to a deploy config YAML file.", ) parser.add_argument( "--log-stats", diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py index 0cab46fee..74d33fe01 100644 --- a/examples/offline_inference/dynin_omni/end2end.py +++ b/examples/offline_inference/dynin_omni/end2end.py @@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace: parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES) parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.") parser.add_argument( - "--stage-config-path", + "--deploy-config", type=str, default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"), help="Path to deploy config yaml.", @@ -1394,8 +1394,8 @@ def main() -> None: from vllm_omni.entrypoints.omni import Omni - stage_config_path = str(Path(args.stage_config_path).expanduser()) - omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype) + deploy_config_path = str(Path(args.deploy_config).expanduser()) + omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype) sampling_params_list = [ SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False) for _ in range(omni.num_stages) diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 9a2350156..fbe4289be 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -73,7 +73,6 @@ def parse_args(): ) parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.") parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.") - parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.") parser.add_argument( "--stream", action="store_true", @@ -146,13 +145,7 @@ def main(): else: bot_task = args.bot_task - if args.deploy_config is not None and args.stage_configs_path is not None: - raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.") - - deploy_config = args.deploy_config - stage_configs_path = args.stage_configs_path - if deploy_config is None and stage_configs_path is None: - deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality] + deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality] omni_kwargs = { "model": args.model, @@ -168,10 +161,7 @@ def main(): if additional_config is not None: omni_kwargs["additional_config"] = additional_config - if deploy_config is not None: - omni_kwargs["deploy_config"] = deploy_config - else: - omni_kwargs["stage_configs_path"] = stage_configs_path + omni_kwargs["deploy_config"] = deploy_config omni = Omni(**omni_kwargs) @@ -272,10 +262,7 @@ def main(): print(f" Modality: {args.modality}") print(f" Prompt task: {task}") print(f" Bot task: {bot_task}") - if deploy_config is not None: - print(f" Deploy config: {deploy_config}") - else: - print(f" Stage config: {stage_configs_path}") + print(f" Deploy config: {deploy_config}") print(f" Num stages: {omni.num_stages}") if args.modality in ("text2img", "img2img"): print(f" Inference steps: {args.steps}") diff --git a/examples/offline_inference/minicpmo/end2end.py b/examples/offline_inference/minicpmo/end2end.py index e9c9019c0..857f1f00f 100644 --- a/examples/offline_inference/minicpmo/end2end.py +++ b/examples/offline_inference/minicpmo/end2end.py @@ -452,12 +452,6 @@ def parse_args(): default=None, help="Path to a .txt file with one prompt per line (preferred).", ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="Path to a stage configs file (deprecated; prefer --deploy-config).", - ) parser.add_argument( "--deploy-config", type=str, diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 9434b3058..faa200bc9 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -484,10 +484,10 @@ def parse_args(): help="Path to a .txt file with one prompt per line (preferred).", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, default=None, - help="Path to a stage configs file.", + help="Path to a deploy config YAML.", ) parser.add_argument( "--video-path", diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py index f175db750..418a7f129 100644 --- a/examples/offline_inference/step_audio2/end2end.py +++ b/examples/offline_inference/step_audio2/end2end.py @@ -232,15 +232,8 @@ def main(args): raise ValueError(f"Unknown query type: {args.query_type}") # Initialize vLLM-Omni with Step-Audio2 - # Resolve deploy config path. Keep explicit stage_configs_path for legacy - # custom configs, but use bundled deploy configs by default. - deploy_config_path = None - stage_config_path = None - if args.stage_configs_path and args.deploy_config: - raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive") - if args.stage_configs_path: - stage_config_path = args.stage_configs_path - elif args.deploy_config: + # Resolve an explicit deploy config or use the bundled task default. + if args.deploy_config: deploy_config_path = args.deploy_config else: configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy" @@ -262,10 +255,7 @@ def main(args): "ray_address": args.ray_address, "trust_remote_code": True, } - if stage_config_path is not None: - omni_kwargs["stage_configs_path"] = stage_config_path - else: - omni_kwargs["deploy_config"] = deploy_config_path + omni_kwargs["deploy_config"] = deploy_config_path omni_llm = Omni(**omni_kwargs) @@ -418,12 +408,6 @@ def parse_args(): " - HuggingFace ID: stepfun-ai/Step-Audio-2-mini" ), ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="Path to a legacy stage config YAML file (default: use bundled deploy config)", - ) parser.add_argument( "--deploy-config", type=str, diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 360c726fd..2d1f9823a 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -92,12 +92,6 @@ def parse_args() -> argparse.Namespace: "meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, " "stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.", ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.", - ) parser.add_argument( "--deploy-config", type=str, @@ -475,8 +469,6 @@ def main(): omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager - if args.stage_configs_path: - omni_kwargs["stage_configs_path"] = args.stage_configs_path if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config if use_nextstep: @@ -515,8 +507,8 @@ def main(): print(f" Image size: {args.width}x{args.height}") if args.lora_path: print(f" LoRA: scale={args.lora_scale}") - if args.stage_configs_path: - print(f" stage-configs-path: {args.stage_configs_path}") + if args.deploy_config: + print(f" deploy-config: {args.deploy_config}") print(f"{'=' * 60}\n") # Build LoRA request when --lora-path is set diff --git a/examples/offline_inference/text_to_speech/fish_speech/end2end.py b/examples/offline_inference/text_to_speech/fish_speech/end2end.py index 18b4d79f8..8ce4664b9 100644 --- a/examples/offline_inference/text_to_speech/fish_speech/end2end.py +++ b/examples/offline_inference/text_to_speech/fish_speech/end2end.py @@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser logger = logging.getLogger(__name__) DEFAULT_MODEL = "fishaudio/s2-pro" -DEFAULT_STAGE_CONFIG = os.path.join( +DEFAULT_DEPLOY_CONFIG = os.path.join( os.path.dirname(__file__), "..", "..", @@ -152,10 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict: log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, ) - if args.deploy_config: - omni_kwargs["deploy_config"] = args.deploy_config - else: - omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG + omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG return omni_kwargs @@ -278,12 +275,6 @@ def parse_args(): default=None, help="Deploy config YAML path (under vllm_omni/deploy/).", ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="Path to stage configs YAML (legacy, prefer --deploy-config).", - ) parser.add_argument( "--output-dir", default="output_audio", diff --git a/examples/offline_inference/text_to_speech/glm_tts/end2end.py b/examples/offline_inference/text_to_speech/glm_tts/end2end.py index 6624c2b35..3f66ed8c3 100644 --- a/examples/offline_inference/text_to_speech/glm_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/glm_tts/end2end.py @@ -145,7 +145,7 @@ def main(args): omni = Omni( model=args.model, - stage_configs_path=deploy_config_path, + deploy_config=deploy_config_path, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, ) diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py index 7c0ce125e..75f562af6 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py @@ -490,10 +490,10 @@ def parse_args(): help="Path to a .txt file with one prompt per line (preferred).", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, default=None, - help="Path to a stage configs file.", + help="Path to a deploy config YAML.", ) parser.add_argument( "--audio-path", diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index c3c3edbd1..9708c1c50 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -6,7 +6,7 @@ Unit tests for StageConfigFactory and related classes. import importlib import warnings -from dataclasses import dataclass, fields +from dataclasses import dataclass from pathlib import Path from unittest.mock import patch @@ -34,9 +34,8 @@ from vllm_omni.config.stage_config import ( load_deploy_config, merge_pipeline_deploy, pipeline_cfg_resolver, - strip_parent_engine_args, ) -from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys +from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -475,7 +474,7 @@ class TestStageConfigFactory: cli_overrides = { "gpu_memory_utilization": 0.9, "model": "some_model", # Internal - "stage_configs_path": "/path", # Internal + "deploy_config": "/path", # Internal "batch_timeout": 10, # Internal } @@ -483,7 +482,7 @@ class TestStageConfigFactory: assert overrides["gpu_memory_utilization"] == 0.9 assert "model" not in overrides - assert "stage_configs_path" not in overrides + assert "deploy_config" not in overrides assert "batch_timeout" not in overrides def test_per_stage_override_excludes_internal_keys(self): @@ -525,49 +524,6 @@ class TestStageResolutionHelpers: assert "model" not in overrides assert "parallel_config" not in overrides - def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self): - parent_fields = {f.name: f for f in fields(EngineArgs)} - filtered, overridden = strip_parent_engine_args( - { - "model": "some/model", - "stage_configs_path": "/tmp/stages.yaml", - "tensor_parallel_size": 4, - "worker_extension_cls": "some.Extension", - "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"}, - }, - parent_fields=parent_fields, - keep_keys={"worker_extension_cls"}, - strip_keys={"stage_configs_path"}, - no_warn_keys={"model"}, - ) - - assert filtered == { - "worker_extension_cls": "some.Extension", - "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"}, - } - assert overridden == ["tensor_parallel_size"] - - def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self): - parent_fields = {f.name: f for f in fields(EngineArgs)} - filtered, overridden = strip_parent_engine_args( - { - "model": "some/model", - "stage_configs_path": "/tmp/stages.yaml", - "allowed_local_media_path": "/data/qwentts", - "allowed_media_domains": ["example.com"], - }, - parent_fields=parent_fields, - keep_keys={"allowed_local_media_path", "allowed_media_domains"}, - strip_keys={"stage_configs_path"}, - no_warn_keys={"model"}, - ) - - assert filtered == { - "allowed_local_media_path": "/data/qwentts", - "allowed_media_domains": ["example.com"], - } - assert overridden == [] - class TestPipelineDiscovery: """Tests for the central pipeline registry (``OMNI_PIPELINES``).""" diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 617083398..0da5b6dd3 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -199,7 +199,7 @@ def _generate_bagel_image( config_path = get_deploy_config_path("ci/bagel.yaml") omni_kwargs: dict[str, Any] = { "model": "ByteDance-Seed/BAGEL-7B-MoT", - "stage_configs_path": config_path, + "deploy_config": config_path, "stage_init_timeout": 300, } if diffusion_quantization_config: diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py index 57997905c..3b9fcbb64 100644 --- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py @@ -314,7 +314,7 @@ def test_bagel_text2img_mooncake_connector(run_level): temp_config_file = _resolve_deploy_config(temp_config_file, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", - stage_configs_path=temp_config_file, + deploy_config=temp_config_file, stage_init_timeout=300, ) as runner: generated_image = _generate_bagel_image(runner.omni) diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py index b97161335..72b56d3e9 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py @@ -251,7 +251,7 @@ def test_bagel_img2img_shared_memory_connector(run_level): config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", - stage_configs_path=config_path, + deploy_config=config_path, ) as runner: generated_image = _generate_bagel_img2img(runner.omni, input_image) if run_level == "advanced_model": @@ -266,7 +266,7 @@ def test_bagel_text2img_shared_memory_connector(run_level): config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", - stage_configs_path=config_path, + deploy_config=config_path, ) as runner: generated_image = _generate_bagel_text2img(runner.omni) if run_level == "advanced_model": diff --git a/tests/e2e/offline_inference/test_audiox_model_expansion.py b/tests/e2e/offline_inference/test_audiox_model_expansion.py index af944b144..5739cc463 100644 --- a/tests/e2e/offline_inference/test_audiox_model_expansion.py +++ b/tests/e2e/offline_inference/test_audiox_model_expansion.py @@ -17,7 +17,7 @@ from vllm_omni.platforms import current_omni_platform # Override in CI if needed: AUDIOX_TEST_MODEL= models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")] -# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize +# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize _OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models] pytestmark = [ diff --git a/tests/e2e/offline_inference/test_bagel_expansion.py b/tests/e2e/offline_inference/test_bagel_expansion.py index 980fc633f..e8ac6c19f 100644 --- a/tests/e2e/offline_inference/test_bagel_expansion.py +++ b/tests/e2e/offline_inference/test_bagel_expansion.py @@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT" BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml") DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" -# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` +# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` _OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG) diff --git a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py index fa1bdfead..7e57433ef 100644 --- a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py +++ b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py @@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s ] -# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` +# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` _cosy_deployment = get_deploy_config_path("cosyvoice3.yaml") _cosy_model_path = str(_resolve_model_dir()) _OMNI_RUNNER_PARAMS = [ diff --git a/tests/e2e/offline_inference/test_ernie_image_expansion.py b/tests/e2e/offline_inference/test_ernie_image_expansion.py index 19793ff00..a43db7d7e 100644 --- a/tests/e2e/offline_inference/test_ernie_image_expansion.py +++ b/tests/e2e/offline_inference/test_ernie_image_expansion.py @@ -17,7 +17,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "baidu/ERNIE-Image" -# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize +# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize _OMNI_RUNNER_PARAM = ( MODEL, None, diff --git a/tests/e2e/offline_inference/test_glm_tts_expansion.py b/tests/e2e/offline_inference/test_glm_tts_expansion.py index e54dd5480..2bde38e6b 100644 --- a/tests/e2e/offline_inference/test_glm_tts_expansion.py +++ b/tests/e2e/offline_inference/test_glm_tts_expansion.py @@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None: with OmniRunner( MODEL, - stage_configs_path=_get_deploy_config(async_chunk=async_chunk), + deploy_config=_get_deploy_config(async_chunk=async_chunk), stage_init_timeout=600, trust_remote_code=True, ) as omni_runner: diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py index 4bdb42c34..84a9291b0 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py @@ -52,7 +52,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] def omni() -> Generator[Omni, None, None]: with OmniRunner( MODEL_NAME, - stage_configs_path=str(STAGE_CONFIG_PATH), + deploy_config=str(STAGE_CONFIG_PATH), trust_remote_code=True, ) as runner: yield runner.omni diff --git a/tests/e2e/offline_inference/test_longcat_video_avatar.py b/tests/e2e/offline_inference/test_longcat_video_avatar.py index f0fa81eb3..404ad9485 100644 --- a/tests/e2e/offline_inference/test_longcat_video_avatar.py +++ b/tests/e2e/offline_inference/test_longcat_video_avatar.py @@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]: with OmniRunner( model, seed=42, - stage_configs_path=None, model_class_name="LongCatVideoAvatarPipeline", additional_config={ "model_type": "avatar-v1.5", diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py index ec1545f6a..f570a6ea7 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py @@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16 MODEL_PATH = "bytedance-research/MammothModa2-Preview" T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml") -_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG}) +_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG) # Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate. _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json" diff --git a/tests/e2e/offline_inference/test_minicpmo_4_5.py b/tests/e2e/offline_inference/test_minicpmo_4_5.py index 6725ce2aa..5a518a7f7 100644 --- a/tests/e2e/offline_inference/test_minicpmo_4_5.py +++ b/tests/e2e/offline_inference/test_minicpmo_4_5.py @@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"] _CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml") -test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models] +test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models] def get_question(prompt_type: str = "text") -> str: diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py index b894fe4ea..41e2229d3 100644 --- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py @@ -25,7 +25,7 @@ from vllm_omni import Omni MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano" STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml") -# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize +# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize _OMNI_RUNNER_PARAM = ( MODEL_NAME, STAGE_CONFIG, diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py index 5b4c7a935..dd65cbdcc 100644 --- a/tests/e2e/offline_inference/test_omnivoice_expansion.py +++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py @@ -20,13 +20,11 @@ from tests.helpers.stage_config import get_deploy_config_path MODEL = "k2-fsa/OmniVoice" DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml") -# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs. -# The migrated test passes deploy_config through extra Omni kwargs. +# OmniRunner tuple: model, deploy config path, extra Omni kwargs. _OMNI_RUNNER_PARAM = ( MODEL, - None, + DEPLOY_CONFIG, { - "deploy_config": DEPLOY_CONFIG, "trust_remote_code": True, "log_stats": True, }, diff --git a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py index 4f5c97d11..42510d7c6 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py @@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None: """Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation.""" lora_request = _make_file_lora_request(tmp_path / "sensenova_lora") - with OmniRunner(MODEL, stage_configs_path=None) as runner: + with OmniRunner(MODEL) as runner: omni = runner.omni baseline = _generate(omni) img_1x = _generate(omni, lora_request, lora_scale=1.0) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py index 1ac76da28..c2739844c 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py @@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img( def test_sensenova_u1_img2img(run_level): """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML).""" input_image = _load_input_image() - with OmniRunner( - "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, - ) as runner: + with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner: generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image) if run_level == "full_model" or run_level == "advanced_model": _validate_pixels(generated_image) @@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit(): input_image = _load_input_image() with OmniRunner( "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, cache_backend="cache_dit", ) as runner: _generate_sensenova_u1_img2img(runner.omni, input_image) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py index 092da16b9..2e07e2076 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py @@ -151,10 +151,7 @@ def _generate_sensenova_u1_image( @hardware_test(res={"cuda": "H100"}) def test_sensenova_u1_text2img(run_level): """Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML).""" - with OmniRunner( - "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, - ) as runner: + with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner: generated_image = _generate_sensenova_u1_image(runner.omni) if run_level == "full_model" or run_level == "advanced_model": _validate_pixels(generated_image) @@ -168,7 +165,6 @@ def test_sensenova_u1_text2img_cache_dit(): """Test SenseNova-U1 text2img with Cache-DiT enabled.""" with OmniRunner( "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, cache_backend="cache_dit", ) as runner: _generate_sensenova_u1_image(runner.omni) diff --git a/tests/e2e/offline_inference/test_soulxsinger_expansion.py b/tests/e2e/offline_inference/test_soulxsinger_expansion.py index c3a6091d6..a008be526 100644 --- a/tests/e2e/offline_inference/test_soulxsinger_expansion.py +++ b/tests/e2e/offline_inference/test_soulxsinger_expansion.py @@ -159,7 +159,7 @@ def test_soulxsinger_multistage_from_audio( with OmniRunner( model, - stage_configs_path=get_deploy_config_path(deploy_yaml), + deploy_config=get_deploy_config_path(deploy_yaml), async_chunk=False, ) as runner: sampling = OmniDiffusionSamplingParams( @@ -212,7 +212,7 @@ def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> with OmniRunner( str(base_dir), - stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"), + deploy_config=get_deploy_config_path("soulxsinger_svs.yaml"), async_chunk=False, ) as runner: sampling = OmniDiffusionSamplingParams( diff --git a/tests/e2e/offline_inference/test_step_audio2_expansion.py b/tests/e2e/offline_inference/test_step_audio2_expansion.py index 483396b23..3a1b60c94 100644 --- a/tests/e2e/offline_inference/test_step_audio2_expansion.py +++ b/tests/e2e/offline_inference/test_step_audio2_expansion.py @@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput MODEL = "stepfun-ai/Step-Audio-2-mini" DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml") -TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})] +TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})] pytestmark = [pytest.mark.slow, pytest.mark.tts] diff --git a/tests/e2e/offline_inference/test_voxtral_tts_expansion.py b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py index 2441c46dc..dcff41cda 100644 --- a/tests/e2e/offline_inference/test_voxtral_tts_expansion.py +++ b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py @@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming(): async def _run(): async_omni = AsyncOmni( model=MODEL, - stage_configs_path=STAGE_CONFIG, + deploy_config=STAGE_CONFIG, stage_init_timeout=300, enforce_eager=True, ) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 055866bd5..ed1031b64 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -4,7 +4,6 @@ invariant to the specific attributes of vLLM config except in cases where we explicitly patch values that differ from vLLM. """ -import argparse import inspect import json import os @@ -20,7 +19,6 @@ from vllm.engine.arg_utils import EngineArgs from vllm_omni.config.model import OmniModelConfig from vllm_omni.engine.arg_utils import OmniEngineArgs -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import build_engine_args_dict pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -149,19 +147,6 @@ def test_qwen3_tts_startup_task_type_is_validated(): ) -def test_from_cli_args_picks_up_stage_configs_path(): - """from_cli_args should pick up stage_configs_path from namespace.""" - ns = argparse.Namespace( - model="facebook/opt-125m", - stage_configs_path="/some/path.yaml", - custom_pipeline_args=None, - ) - - args = OmniEngineArgs.from_cli_args(ns) - assert args.stage_configs_path == "/some/path.yaml" - assert args.custom_pipeline_args is None - - def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch): """Ensure Code2Wav mirrors stage max_model_len into nested HF overrides. @@ -418,130 +403,6 @@ def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size(): assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size -def test_stage_configs_path_field(): - """OmniEngineArgs with stage_configs_path should construct without error.""" - args = OmniEngineArgs(stage_configs_path="/some/path.yaml") - assert args.stage_configs_path == "/some/path.yaml" - - -def test_strip_single_engine_args(): - """_strip_single_engine_args should remove EngineArgs fields but keep omni fields.""" - kwargs = { - # Parent EngineArgs fields — stripped unless explicitly allowlisted - "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}', - "tensor_parallel_size": 4, - "gpu_memory_utilization": 0.9, - "model": "some/model", - # Parent field that should be kept (allowlisted) - "worker_extension_cls": "some.Extension", - # OmniEngineArgs-only / non-engine fields — should pass through - "stage_configs_path": "/path/to/yaml", - "custom_pipeline_args": {"pipeline_class": "my.Pipeline"}, - "mode": "text-to-image", - "lora_path": "/some/lora", - } - - filtered = AsyncOmniEngine._strip_single_engine_args(kwargs) - - # Stripped — parent EngineArgs fields - assert "compilation_config" not in filtered - assert filtered["tensor_parallel_size"] == 4 - assert "gpu_memory_utilization" not in filtered - assert "model" not in filtered - - # Stripped — orchestrator-level OmniEngineArgs field - assert "stage_configs_path" not in filtered - - # Kept - assert filtered["worker_extension_cls"] == "some.Extension" - assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"} - assert filtered["mode"] == "text-to-image" - assert filtered["lora_path"] == "/some/lora" - - -def test_strip_single_engine_args_model_does_not_trigger_warning(mocker): - """model is always in kwargs (callers set it via from_cli_args/asdict), - so it should not cause the override warning by itself or appear in it.""" - mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning") - - # Typical caller kwargs: model is always present, no other parent - # EngineArgs fields are explicitly overridden. - AsyncOmniEngine._strip_single_engine_args( - { - "model": "some/model", - "custom_pipeline_args": {"pipeline_class": "my.Pipeline"}, - } - ) - mock_warn.assert_not_called() - - # When there *are* genuinely surprising overrides alongside model, - # the warning should mention them but not model. Keep-listed fields such as - # tensor_parallel_size are intentionally passed through and should not warn. - AsyncOmniEngine._strip_single_engine_args( - { - "model": "some/model", - "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}', - "tensor_parallel_size": 4, - "custom_pipeline_args": {"pipeline_class": "my.Pipeline"}, - } - ) - mock_warn.assert_called_once() - warned_args = mock_warn.call_args[0][-1] # the formatted arg list - assert "compilation_config" in warned_args - assert "tensor_parallel_size" not in warned_args - assert "model" not in warned_args - - -def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path): - """New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs. - - Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping - those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni - accuracy from ~78% to ~70%. - """ - from vllm_omni.entrypoints.utils import is_new_format_deploy_config - - yaml_path = tmp_path / "minicpm_deploy.yaml" - yaml_path.write_text( - """ -pipeline: minicpmo_4_5 -async_chunk: true -stages: - - stage_id: 0 - devices: "0" -""".strip(), - encoding="utf-8", - ) - path = str(yaml_path) - assert is_new_format_deploy_config(path) - - kwargs = { - "stage_configs_path": path, - "trust_remote_code": True, - "interleave_mm_strings": True, - "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}, - "gpu_memory_utilization": 0.55, - } - # Mirror AsyncOmniEngine._resolve_stage_configs strip decision. - if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]): - base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs)) - else: - base = dict(kwargs) - assert base.get("interleave_mm_strings") is True - assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}} - # Parent EngineArgs fields also remain available for deploy merge. - assert base.get("gpu_memory_utilization") == 0.55 - - legacy = tmp_path / "legacy.yaml" - legacy.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8") - assert not is_new_format_deploy_config(str(legacy)) - legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy)) - stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs) - assert "interleave_mm_strings" not in stripped - assert "media_io_kwargs" not in stripped - assert "gpu_memory_utilization" not in stripped - - # For https://github.com/vllm-project/vllm-omni/issues/3293 def test_tensor_parallel_size_none_is_handled(): """Ensure the tensor parallel size of None isn't forwarded.""" diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py index 40f89df84..56b25704d 100644 --- a/tests/engine/test_async_omni_engine_abort.py +++ b/tests/engine/test_async_omni_engine_abort.py @@ -68,7 +68,7 @@ async def test_abort(): with ExitStack() as after: engine = AsyncOmni( model=model, - stage_configs_path=stage_config, + deploy_config=stage_config, ) after.callback(engine.shutdown) diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index bde8ba490..65ec024d5 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -428,7 +428,7 @@ class TestSingleStageModeDetection: assert engine.single_stage_mode is False assert engine._single_stage_id_filter is None - def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture): + def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture): duplex_session = DuplexSessionRuntimeConfig(max_sessions=2) get_pipeline_config = mocker.patch( "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", @@ -441,7 +441,7 @@ class TestSingleStageModeDetection: engine = self._make_engine_no_thread( mocker, - stage_configs_path="/fake/duplex.yaml", + deploy_config="/fake/duplex.yaml", ) get_pipeline_config.assert_called_once_with( diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 6120d5ea4..a672c32e9 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -364,7 +364,7 @@ async def test_omni_generate_request_id(): Text modality only; asserts caller-visible ids are preserved across the multi-stage orchestrator path on H100. """ - engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG) + engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG) try: for req_id in _OMNI_REQ_IDS: async for output in engine.generate( diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 2fa8c92bd..75a9de2a7 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -478,12 +478,11 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) - engine._strip_single_engine_args = lambda kwargs: kwargs _, stage_configs = engine._resolve_stage_configs( "dummy-model", { - "stage_configs_path": "dummy.yaml", + "deploy_config": "dummy.yaml", "additional_config": {"torchair_graph_config": {"enabled": True}}, }, trust_remote_code=False, @@ -516,12 +515,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage( ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) - engine._strip_single_engine_args = lambda kwargs: kwargs _, stage_configs = engine._resolve_stage_configs( "dummy-model", { - "stage_configs_path": "dummy.yaml", + "deploy_config": "dummy.yaml", "quantization_config": {"method": "bitsandbytes"}, }, trust_remote_code=False, diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 04a29edd1..6aa5d6855 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -193,7 +193,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) with pytest.raises(ValueError, match="No stage config found for stage_id=0"): run_headless(args) - assert captured["args"][1] is None + assert len(captured["args"]) == 2 assert captured["deploy_config_path"] == "/tmp/deploy.yaml" assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}} assert captured["strategy_config_path"] == "/tmp/strategy.yaml" diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index c6d8e0354..a20c97a58 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -381,7 +381,6 @@ class TestLoadAndResolveStageConfigs: config_path, stage_configs, _ = load_and_resolve_stage_configs( model="black-forest-labs/FLUX.2-klein-4B", - stage_configs_path=None, kwargs=kwargs, trust_remote_code=False, default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs), @@ -390,9 +389,7 @@ class TestLoadAndResolveStageConfigs: assert len(stage_configs) == 1 assert "dtype" in stage_configs[0]["engine_args"] - def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas( - self, tmp_path, mocker: MockerFixture - ): + def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture): deploy_path = tmp_path / "qwen3_multi.yaml" deploy_path.write_text( 'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n', @@ -413,18 +410,23 @@ class TestLoadAndResolveStageConfigs: "vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=(returned_stage_configs, None), ) + cli_overrides = { + "interleave_mm_strings": True, + "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}, + "gpu_memory_utilization": 0.55, + } config_path, stage_configs, _ = load_and_resolve_stage_configs( model="dummy-model", - stage_configs_path=str(deploy_path), - kwargs={}, + kwargs=cli_overrides, trust_remote_code=True, + deploy_config_path=str(deploy_path), ) load_stage_configs.assert_called_once_with( "dummy-model", trust_remote_code=True, - base_engine_args={}, + base_engine_args=cli_overrides, deploy_config_path=str(deploy_path), stage_overrides=None, strategy_config_path=None, diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 82fcc308d..cfc8cc75d 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -2635,7 +2635,7 @@ class OmniRunner: # affects the test runner wrapper. init_timeout: int = 1800, log_stats: bool = False, - stage_configs_path: str | None = None, + deploy_config: str | None = None, **kwargs, ) -> None: startup_t0 = time.perf_counter() @@ -2653,7 +2653,7 @@ class OmniRunner: stage_init_timeout=stage_init_timeout, batch_timeout=batch_timeout, init_timeout=init_timeout, - stage_configs_path=stage_configs_path, + deploy_config=deploy_config, **kwargs, ) startup_s = time.perf_counter() - startup_t0 @@ -3179,12 +3179,7 @@ def iter_omni_server( print("OmniServer stopping...") else: if stage_config_path is not None: - from vllm_omni.entrypoints.utils import is_new_format_deploy_config - - if is_new_format_deploy_config(stage_config_path): - server_args += ["--deploy-config", stage_config_path] - else: - server_args += ["--stage-configs-path", stage_config_path] + server_args += ["--deploy-config", stage_config_path] with ( OmniServer( @@ -3237,7 +3232,7 @@ def iter_omni_runner( model = model_prefix + model if run_level == "core_model" and request.node.get_closest_marker("diffusion"): model = resolve_tiny_model_path(model) - with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner: + with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner: print("OmniRunner started successfully") yield runner print("OmniRunner stopping...") diff --git a/tools/baselines/omni_hash_baseline.py b/tools/baselines/omni_hash_baseline.py index 1bb93099b..8d8f5ff13 100644 --- a/tools/baselines/omni_hash_baseline.py +++ b/tools/baselines/omni_hash_baseline.py @@ -147,7 +147,7 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]: # (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time. omni = Omni( model, - stage_configs_path=yaml_path, + deploy_config=yaml_path, stage_init_timeout=1800, init_timeout=3600, async_chunk=async_mode, diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 4b5c73a39..32ee2c826 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -53,7 +53,6 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset( "model", "omni", "output_modalities", - "stage_configs_path", "stage_id", "tokenizer", } diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 496aea208..86a6fed1b 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -4,7 +4,6 @@ from __future__ import annotations -import dataclasses import functools import re import warnings @@ -56,7 +55,7 @@ def build_stage_runtime_overrides( ``internal_keys`` defaults to the union of ``arg_utils.internal_blacklist_keys()`` and ``arg_utils.SHARED_FIELDS`` so that neither orchestrator-only fields nor shared-pipeline fields - (``model`` / ``stage_configs_path`` / ``log_stats`` / ``stage_id``) leak + (``model`` / ``log_stats`` / ``stage_id``) leak into a stage's per-stage runtime overrides — the orchestrator sets those uniformly for every stage, they are not per-stage knobs. Callers can pass an explicit set for tests or specialized flows. @@ -90,46 +89,6 @@ def build_stage_runtime_overrides( return result -def strip_parent_engine_args( - kwargs: dict[str, Any], - *, - parent_fields: dict[str, dataclasses.Field], - keep_keys: set[str] | frozenset[str] = frozenset(), - strip_keys: set[str] | frozenset[str] = frozenset(), - no_warn_keys: set[str] | frozenset[str] = frozenset(), -) -> tuple[dict[str, Any], list[str]]: - """Strip parent ``EngineArgs`` fields before merging into stage YAML.""" - overridden: list[str] = [] - result: dict[str, Any] = {} - - for key, value in kwargs.items(): - if key in strip_keys: - continue - - if key not in parent_fields or key in keep_keys: - result[key] = value - continue - - field_def = parent_fields[key] - if field_def.default is not dataclasses.MISSING: - default = field_def.default - elif field_def.default_factory is not dataclasses.MISSING: - default = field_def.default_factory() - else: - default = dataclasses.MISSING - - if default is dataclasses.MISSING or value is None: - continue - - if dataclasses.is_dataclass(default) and not isinstance(default, type): - default = asdict(default) - - if value != default and key not in no_warn_keys: - overridden.append(key) - - return result, sorted(overridden) - - def _apply_diffusion_parallel_runtime_overrides( engine_args: dict[str, Any], runtime_overrides: dict[str, Any], diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a3aa96f60..29c84a1c9 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -167,9 +167,6 @@ class OmniEngineArgs(EngineArgs): Required when single-stage mode is active. omni_master_port: TCP port for the OmniMasterServer registration socket. Required when single-stage mode is active. - stage_configs_path: Optional path to a JSON/YAML file containing - stage configurations for the multi-stage pipeline. If None, - stage configs are resolved from the model's default configuration. output_modalities: Optional list of output modality names to enable (e.g. ["text", "audio"]). If None, all modalities supported by the model are used. @@ -231,7 +228,6 @@ class OmniEngineArgs(EngineArgs): omni_dp_size_local: int = 1 omni_lb_policy: str = "random" omni_heartbeat_timeout: float = 30.0 - stage_configs_path: str | None = None output_modalities: list[str] | None = None log_stats: bool = False custom_pipeline_args: dict[str, Any] | None = None @@ -490,7 +486,6 @@ class OrchestratorArgs: ray_address: str | None = None # === Config Files === - stage_configs_path: str | None = None deploy_config: str | None = None stage_overrides: str | None = None # raw JSON string; parsed downstream # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays @@ -582,7 +577,6 @@ SHARED_FIELDS: frozenset[str] = frozenset( "model", # orch: detect model_type; engine: load weights "stage_id", # orch: route (headless); engine: identity "log_stats", # both want the flag - "stage_configs_path", # orch: load legacy YAML; engine: may reference for validation "async_chunk", # orch: read from CLI, redistribute; engine: per-stage flag "tokenizer", # orch: detect model type; engine: tokenization } diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index aa284067b..a613ebc80 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -9,7 +9,6 @@ from __future__ import annotations import asyncio import concurrent.futures -import dataclasses import json import queue import threading @@ -24,7 +23,6 @@ import janus import torch from omegaconf import OmegaConf from vllm import envs as vllm_envs -from vllm.engine.arg_utils import EngineArgs from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest @@ -34,7 +32,6 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot from vllm_omni.config.stage_config import ( DuplexSessionRuntimeConfig, load_deploy_config, - strip_parent_engine_args, ) from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config from vllm_omni.diffusion.diffusion_engine import supports_audio_output @@ -69,7 +66,6 @@ from vllm_omni.engine.stage_runtime import ( ) from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import ( - is_new_format_deploy_config, load_and_resolve_stage_configs, parse_stage_overrides, ) @@ -85,39 +81,6 @@ if TYPE_CHECKING: _STARTUP_POLL_INTERVAL_S = 1.0 _REQUEST_QUEUE_MAXSIZE = 256 -# ============================================================================ -# Parent-EngineArgs field-routing contracts (consumed by -# AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set). -# ============================================================================ - -# Fields that must survive the "equal to default → strip" filter because -# diffusion stages need them even when equal to vllm's default value -# (e.g. colocate worker setup relies on worker_extension_cls being forwarded). -_PARENT_ARGS_KEEP: frozenset[str] = frozenset( - { - "worker_extension_cls", - "allowed_local_media_path", - "allowed_media_domains", - # Legacy stage-config YAMLs may intentionally leave parallel or - # distributed knobs unspecified at the stage level and rely on - # top-level CLI values to fill them in during the per-stage merge. - # Keep these fields so stages that omit them can inherit CLI values, - # while stages with explicit YAML values still win because the legacy - # stage-config loader prefers stage-local engine args. - "tensor_parallel_size", - } -) - -# Omni orchestrator-level fields consumed by ``_resolve_stage_configs`` that -# must never leak into per-stage EngineArgs (``stage_configs_path`` would -# trigger the ``create_model_config`` guard). -_PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"}) - - -# Fields always populated by callers (via ``from_cli_args`` / ``asdict``) so -# their presence as an override is never a surprise — suppress the -# "override ignored" warning for these. -_PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"}) class AsyncOmniEngine: @@ -210,10 +173,8 @@ class AsyncOmniEngine: self._omni_master_port, ) - # Stage resolution pops deploy_config, so get pipeline-wide settings - # beforehand. The stage CLI exposes the same deploy YAML through - # stage_configs_path. - deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path") + # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand. + deploy_config_path = kwargs.get("deploy_config") # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not # specified" so stage-config resolution can defer to the deploy yaml's # per-stage value (see ``with_trust_remote_code_override``). The @@ -1116,40 +1077,6 @@ class AsyncOmniEngine: default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" return default_stage_cfg - @staticmethod - def _strip_single_engine_args(kwargs: dict[str, Any]) -> dict[str, Any]: - """Remove parent ``EngineArgs`` fields from *kwargs*. - - When ``stage_configs_path`` is set, per-stage engine args are defined - in the YAML. Top-level single-engine fields (``compilation_config``, - ``tensor_parallel_size``, …) must not leak into per-stage configs via - the ``base_engine_args`` merge in ``load_stage_configs_from_yaml`` — - they can cause type errors (e.g. ``compilation_config`` as a JSON - string rejected by ``VllmConfig``) or silently override YAML values. - - Logs a warning for any parent field whose value differs from the - dataclass default, so users know their explicit overrides are ignored. - See the module-level ``_PARENT_ARGS_*`` constants for the routing - contracts this method enforces. - """ - parent_fields: dict[str, dataclasses.Field] = {f.name: f for f in dataclasses.fields(EngineArgs)} - result, overridden = strip_parent_engine_args( - kwargs, - parent_fields=parent_fields, - keep_keys=_PARENT_ARGS_KEEP, - strip_keys=_PARENT_ARGS_STRIP, - no_warn_keys=_PARENT_ARGS_NO_WARN, - ) - - if overridden: - logger.warning( - "stage_configs_path is set — the following top-level engine " - "args are ignored (per-stage YAML takes precedence): %s", - ", ".join(sorted(overridden)), - ) - - return result - def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: """Apply a strategy-derived ``omni_lb_policy`` to the engine. @@ -1188,7 +1115,6 @@ class AsyncOmniEngine: ) -> tuple[str, list[Any]]: """Resolve stage configs and inject defaults shared by orchestrator/headless.""" - stage_configs_path = kwargs.get("stage_configs_path", None) deploy_config_path = kwargs.pop("deploy_config", None) strategy_config_path = kwargs.pop("strategy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) @@ -1196,26 +1122,15 @@ class AsyncOmniEngine: if explicit_stage_configs is not None: logger.warning( "`stage_configs` is not part of the public API. " - "Ignoring it and resolving stages from stage_configs_path/model factory." + "Ignoring it and resolving stages from deploy_config/model factory." ) - # Legacy ``stage_args`` YAMLs own per-stage EngineArgs completely, so - # strip parent CLI fields. New-format deploy YAMLs (also accepted via - # ``--stage-configs-path``) merge CLI overrides the same way as - # ``--deploy-config`` — do not strip, or flags like - # ``interleave_mm_strings`` / ``media_io_kwargs`` are silently dropped. - if stage_configs_path is not None and not is_new_format_deploy_config(stage_configs_path): - base_kwargs = self._strip_single_engine_args(kwargs) - else: - base_kwargs = kwargs - # Parse --stage-overrides JSON string if provided stage_overrides = parse_stage_overrides(stage_overrides_json) config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs( model, - stage_configs_path, - base_kwargs, + kwargs, trust_remote_code=trust_remote_code, default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs), deploy_config_path=deploy_config_path, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index f6899deaa..33905b254 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -123,9 +123,8 @@ class AsyncOmni(EngineClient, OmniBase): Args: model: Model name or path to load. **kwargs: Additional keyword arguments. - - stage_configs_path: Optional path to YAML file containing stage - configurations. If None, configurations are resolved from model - pipeline factory. + - deploy_config: Optional path to a deploy YAML. If None, + configurations are resolved from the model pipeline factory. - log_stats: Whether to enable statistics logging. - stage_init_timeout: Timeout for per-stage initialization. - init_timeout: Total timeout for orchestrator startup. diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index dfa222cd1..62ee10d94 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -945,8 +945,6 @@ def run_headless(args: TrackingNamespace) -> None: config_path, stage_configs, _ = load_and_resolve_stage_configs( model, - # The serve CLI no longer accepts legacy stage_args YAMLs. - None, args_dict, # store_true cannot express an explicit False: absent maps to None # ("not specified") so the deploy yaml's per-stage value applies. diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 0a8daa055..16114bcb6 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -679,7 +679,7 @@ async def build_async_omni_from_stage_config( EngineClient instance (AsyncOmni) ready for use Note: - Stage configurations are loaded from args.stage_configs_path if provided, + Stage configurations are loaded from ``args.deploy_config`` when provided, otherwise from the model's default configuration. """ @@ -1156,7 +1156,7 @@ async def omni_init_app_state( state.openai_serving_duplex = None if state.openai_serving_chat is not None and should_enable_duplex_endpoint( state.stage_configs, - config_path=getattr(args, "stage_configs_path", None) or getattr(args, "deploy_config", None), + config_path=getattr(args, "deploy_config", None), ): from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index ff0cd4f9e..07b45f374 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -6,7 +6,6 @@ from dataclasses import fields, is_dataclass from pathlib import Path from typing import Any, get_args, get_origin -import yaml from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.transformers_utils.config import get_config, get_hf_file_to_dict @@ -30,15 +29,6 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent logger = init_logger(__name__) -def is_new_format_deploy_config(path: str | None) -> bool: - """Return True when *path* is a deploy YAML (``stages:`` without legacy ``stage_args:``).""" - if not path or not os.path.exists(path): - return False - with open(path, encoding="utf-8") as f: - peek = yaml.safe_load(f) or {} - return isinstance(peek, dict) and "stages" in peek and "stage_args" not in peek - - _DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = { "GlmImagePipeline": "glm_image", } @@ -606,7 +596,6 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None: def load_and_resolve_stage_configs( model: str, - stage_configs_path: str | None, kwargs: dict | None, *, trust_remote_code: bool | None, @@ -615,17 +604,15 @@ def load_and_resolve_stage_configs( stage_overrides: dict[str, dict[str, Any]] | None = None, strategy_config_path: str | None = None, ) -> tuple[str, list, str | None]: - """Load stage configurations from model or YAML file with fallback to defaults. + """Load stage configurations from a deploy YAML or model defaults. Args: model: Model name or path - stage_configs_path: Optional path to legacy YAML (stage_args format) kwargs: Engine arguments to merge with stage configs trust_remote_code: Whether to trust remote code while resolving the model config. default_stage_cfg_factory: Optional callable that takes no args and returns default stage config list when no configs are found - deploy_config_path: Optional path to deploy YAML (new format). - Mutually exclusive with ``stage_configs_path``. + deploy_config_path: Optional path to a deploy YAML. stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON. Keys are stage_id strings, values are dicts of overrides. strategy_config_path: Optional path to a composable-parallel @@ -636,65 +623,21 @@ def load_and_resolve_stage_configs( the strategy-derived pipeline-wide load-balancer policy (``None`` when no strategy set one), returned for the engine to apply. """ - if stage_configs_path is not None and deploy_config_path is not None: - raise ValueError( - "--stage-configs-path and --deploy-config are mutually exclusive: " - "they use different path resolution rules and loading paths. " - "Use --deploy-config for new-format YAMLs (preferred); " - "--stage-configs-path is kept only for the legacy `stage_args` format " - "and will be removed in a future release." - ) - if stage_configs_path is not None and deploy_config_path is None: - if not os.path.exists(stage_configs_path): - raise FileNotFoundError( - f"--stage-configs-path {stage_configs_path!r} does not exist. " - "Legacy `stage_configs/` yamls were replaced by `vllm_omni/deploy/.yaml`; " - "use --deploy-config. See docs/configuration/stage_configs.md." - ) - if is_new_format_deploy_config(stage_configs_path): - deploy_config_path = stage_configs_path - stage_configs_path = None + config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model) + stage_configs, omni_lb_policy = load_stage_configs_from_model( + model, + trust_remote_code=trust_remote_code, + base_engine_args=kwargs, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ) + if not stage_configs: + if default_stage_cfg_factory is not None: + default_stage_cfg = default_stage_cfg_factory() + stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg)) else: - logger.warning( - "--stage-configs-path is deprecated; migrate %r and use --deploy-config.", - stage_configs_path, - ) - - omni_lb_policy: str | None = None - if deploy_config_path is not None: - config_path = deploy_config_path - stage_configs, omni_lb_policy = load_stage_configs_from_model( - model, - trust_remote_code=trust_remote_code, - base_engine_args=kwargs, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ) - if not stage_configs: - if default_stage_cfg_factory is not None: - default_stage_cfg = default_stage_cfg_factory() - stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg)) - else: - stage_configs = [] - elif stage_configs_path is None: - config_path = resolve_model_config_path(model) - stage_configs, omni_lb_policy = load_stage_configs_from_model( - model, - trust_remote_code=trust_remote_code, - base_engine_args=kwargs, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ) - if not stage_configs: - if default_stage_cfg_factory is not None: - default_stage_cfg = default_stage_cfg_factory() - stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg)) - else: - stage_configs = [] - else: - config_path = stage_configs_path - stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=kwargs) + stage_configs = [] stage_configs = filter_stages(config_path, stage_configs, kwargs) logger.debug(f"stage_configs: {stage_configs}") From 83d1f11449a0eaf2910a0a8e34c055df20e4f169 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Fri, 14 Aug 2026 00:13:14 +0800 Subject: [PATCH 27/41] [Feat] LTX Standard Two-Stage Pipeline (#5500) Signed-off-by: mglyn <1203789601@qq.com> --- .buildkite/cuda/test-nightly.yml | 10 +- docs/models/supported_models.md | 3 +- docs/user_guide/quantization/fp8.md | 4 +- recipes/LTX/LTX-2.md | 118 +++-- .../models/ltx2/test_ltx2_latents.py | 43 +- .../models/ltx2/test_ltx2_parallel.py | 14 + .../models/ltx2/test_ltx2_phase_adapter.py | 483 ++++++++++++++++++ .../models/ltx2/test_ltx2_pipeline.py | 372 +++++++++++++- .../models/ltx2/test_ltx2_transformer.py | 93 +++- tests/e2e/accuracy/ltx/run_ltx_reference.py | 155 ++++-- .../ltx/test_ltx_official_similarity.py | 444 ++++++++++++---- tests/e2e/online_serving/test_ltx.py | 109 ++++ .../e2e/online_serving/test_ltx2_expansion.py | 91 ---- tests/model_extras/test_model_extras.py | 11 +- tests/model_tests/diffusion/test_alignment.py | 1 + vllm_omni/diffusion/models/ltx2/__init__.py | 6 +- .../models/ltx2/ltx2_adapter_parser.py | 156 ++++++ .../diffusion/models/ltx2/ltx2_components.py | 133 ++++- .../diffusion/models/ltx2/ltx2_denoise.py | 23 +- .../diffusion/models/ltx2/ltx2_guidance.py | 32 +- .../diffusion/models/ltx2/ltx2_latents.py | 50 +- .../models/ltx2/ltx2_phase_adapter.py | 430 ++++++++++++++++ .../diffusion/models/ltx2/ltx2_recipes.py | 50 ++ .../diffusion/models/ltx2/ltx2_request.py | 31 +- .../diffusion/models/ltx2/ltx2_runtime.py | 38 +- .../diffusion/models/ltx2/ltx2_transformer.py | 63 ++- .../models/ltx2/pipeline_ltx2_two_stage.py | 45 +- vllm_omni/diffusion/registry.py | 6 + vllm_omni/model_extras/registry.py | 1 + 29 files changed, 2644 insertions(+), 371 deletions(-) create mode 100644 tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py create mode 100644 tests/e2e/online_serving/test_ltx.py delete mode 100644 tests/e2e/online_serving/test_ltx2_expansion.py create mode 100644 vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py create mode 100644 vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index d3246a98f..ea74847b9 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -456,7 +456,7 @@ steps: timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" - - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model" + - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" @@ -477,10 +477,16 @@ steps: commands: - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model - - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model - pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model mirror_hardwares: h100_2 + - label: ":full_moon: Diffusion X2V · LTX Accuracy Test" + timeout_in_minutes: 180 + commands: + - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model + - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model + mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2V · Perf Test" key: nightly-diffusion-x2v-performance timeout_in_minutes: 180 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b9686047a..73fdc55e4 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -42,7 +42,8 @@ th { | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — | | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — | -| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — | +| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — | +| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — | | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — | | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index ddca3ea19..05473b15e 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -80,12 +80,14 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) ## Model Type Support -### Diffusion Model (Qwen-Image, Wan2.2) +### Diffusion Models | Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization | |-------|-----------|:-------:|:------:|----------------|------------------|------------------| | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | | | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | | +| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | +| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ | | FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | | | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ | diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index 3142f9fd1..123880d09 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -4,11 +4,20 @@ ## Pipelines -| Model | Checkpoint | Task | `--model-class-name` | Batching | -|---|---|---|---|---| -| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes | -| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No | -| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes | +| `--model-class-name` | Task | Required checkpoint repositories | +|---|---|---| +| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` | +| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` | +| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` | +| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` | +| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`
`Lightricks/LTX-2.3` | +| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`
`Lightricks/LTX-2.3` | + +Repositories in the table are download units. A full pipeline repository +contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler, +and tokenizer; an additional repository supplies LoRA or upsampler sidecars. +The runtime first searches the model root for sidecars under their official +filenames, then downloads them from the matching Lightricks Hub repository. `LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial @@ -16,10 +25,10 @@ image selects I2V. Both one-stage repositories declare this class, so `--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint; the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable. -`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V -or provide one initial image for I2V. As in the official LTX pipelines, -distilled mode is selected explicitly rather than inferred from the checkpoint -name or metadata; always pass `--model-class-name LTX2DistilledPipeline`. +`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples, +then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully +merged distilled Transformer in both stages. Both entries support T2V and I2V; +select their class explicitly. ## API Migration @@ -68,24 +77,28 @@ The model recipe also supplies the default negative prompt. Top-level API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts default to `121`. -## Distilled Two-Stage Defaults +## Two-Stage Defaults -| Parameter | Value | -|---|---:| -| Final width × height | 1536 × 1024 | -| Stage 1 width × height | 768 × 512 | -| Frames / frame rate | 121 / 24 | -| Stage 1 / Stage 2 denoise steps | 8 / 3 | -| Guidance | Fixed positive-only | +| Parameter | Ordinary | Full-distilled | +|---|---:|---:| +| Final width × height | 1536 × 1024 | 1536 × 1024 | +| Stage 1 width × height | 768 × 512 | 768 × 512 | +| Frames / frame rate | 121 / 24 | 121 / 24 | +| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 | +| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only | -The dimensions passed to the API are the final output dimensions and must be -divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`. -The distilled recipe fixes both sigma schedules; `num_inference_steps` must be -`8` when supplied. +API dimensions are final dimensions and must be divisible by 64. All LTX +requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or +LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so +`num_inference_steps`, when supplied, must be `8`. Both entries reject custom +sigmas and input latents. + +Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer +and automatically switches to dynamic LoRA when quantization is enabled. ## Serving -Start either checkpoint: +Start either one-stage checkpoint: ```bash vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600 @@ -95,11 +108,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600 vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600 ``` -Start the distilled checkpoint with its two-stage entry: +Start a two-stage checkpoint with its explicit class: ```bash vllm serve rootonchair/LTX-2-19b-distilled --omni \ --model-class-name LTX2DistilledPipeline --stage-init-timeout 600 +# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent +vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \ + --model-class-name LTX2DistilledPipeline --stage-init-timeout 600 +# Ordinary LTX-2.3 two-stage +vllm serve diffusers/LTX-2.3-Diffusers --omni \ + --model-class-name LTX2TwoStagePipeline \ + --enable-layerwise-offload \ + --stage-init-timeout 600 ``` The same server handles T2V and I2V. A T2V request using the selected recipe's @@ -135,8 +156,9 @@ together with `input_reference`. ## Guidance -LTX supports independent video/audio CFG, spatio-temporal guidance (STG), -cross-modality guidance, and rescaling: +One-stage and ordinary Stage 1 support independent video/audio CFG, +spatio-temporal guidance (STG), cross-modality guidance, and rescaling. +Distilled stages and ordinary Stage 2 are fixed positive-only. | Parameter | Default | Effect | Alias | |---|---:|---|---| @@ -172,9 +194,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers text CFG, STG, and cross-modality guidance passes; guidance rescaling is applied after all pass predictions are gathered. -The default one-stage recipe has four Transformer passes per denoise step: -`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced -configurations are therefore: +The default one-stage recipe and ordinary Stage 1 have four Transformer passes +per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod` +(cross-modality). The useful balanced configurations are therefore: | `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes | |---:|---:|---:|---| @@ -195,9 +217,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \ ``` The total device count is the product of `cfg_parallel_size` and the other -configured parallel dimensions. Distilled pipelines use positive-only -guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for -them. +configured parallel dimensions. Positive-only distilled phases do not benefit +from `cfg_parallel_size > 1`. ### Python API @@ -242,22 +263,22 @@ noted below. | `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. | | `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. | | `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. | -| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. | +| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. | | `width` | `int`, `None` | Same precedence and alignment as `height`. | | `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. | | `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. | -| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. | +| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. | | `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. | | `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. | -| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. | +| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. | | `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. | | `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. | | `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. | | `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. | -| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. | -| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. | +| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. | +| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. | | `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. | -| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. | +| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. | | `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. | | `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. | | `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. | @@ -273,13 +294,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`. ### Recipe-Specific Request Capabilities -| Override | One-stage | Distilled two-stage | -|---|---|---| -| Guidance | Supported | Fixed positive-only | -| Negative prompt/embeddings | Supported | Rejected | -| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 | -| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules | -| Video/audio latents | Supported | Rejected | +| Override | One-stage | Ordinary two-stage | Distilled two-stage | +|---|---|---|---| +| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only | +| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected | +| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 | +| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules | +| Video/audio latents | Supported | Rejected | Rejected | These capability checks apply equally to direct `forward` keywords and values in `OmniDiffusionSamplingParams`; unsupported values fail instead of @@ -309,6 +330,10 @@ bundled offline CLI do not currently expose `sigmas`. `--max-num-seqs 1`. - `--cfg-parallel-size` shards the complete LTX guidance plan, including STG, modality guidance, and rescale-compatible prediction gathering. +- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding; + Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a + divisible request shape. +- Cache-DiT is one-stage only; multi-stage configurations are rejected. ## Operational Notes @@ -323,14 +348,15 @@ bundled offline CLI do not currently expose `sigmas`. request parameter. - For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with `tests/dfx/perf/scripts/run_diffusion_benchmark.py`. -- Current two-stage support is limited to the distilled checkpoint; official - LTX two-stage/HQ execution remains out of scope. +- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution + remains out of scope. ## References - - - +- - [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md) - [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md) - [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md) diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index 36f7e9687..cad5c300d 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -23,7 +23,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1): return pipeline -def test_prepare_video_latents_samples_directly_in_packed_token_space(): +def test_prepare_video_latents_matches_official_values_and_token_major_layout(): pipeline = _make_pipeline(LTX2Pipeline) pipeline.vae_spatial_compression_ratio = 8 pipeline.vae_temporal_compression_ratio = 8 @@ -44,6 +44,7 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space(): ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert actual.stride()[1:] == (1, actual.shape[1]) def test_prepare_audio_latents_samples_directly_in_packed_token_space(): @@ -81,6 +82,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(): assert original_num_frames == 1 assert padded_num_frames == 2 assert latents.shape == (1, 2, 128) + torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:])) + + +def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding(): + pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1) + pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4) + generator_sp1 = torch.Generator().manual_seed(42) + generator_sp4 = torch.Generator().manual_seed(42) + + latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=3, + dtype=torch.float32, + device=torch.device("cpu"), + generator=generator_sp1, + ) + latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=3, + dtype=torch.float32, + device=torch.device("cpu"), + generator=generator_sp4, + ) + + torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0) + torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:])) + torch.testing.assert_close( + torch.randn(8, generator=generator_sp4), + torch.randn(8, generator=generator_sp1), + rtol=0, + atol=0, + ) def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents(): @@ -137,7 +174,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(): assert original_num_frames == 10 assert padded_num_frames == 12 assert padded.shape == (1, 12, 8) - torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents)) + packed = latent_ops.pack_audio_latents(latents) + torch.testing.assert_close(padded[:, :10], packed[:, :10]) + torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:])) def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(): diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 8d8c22e79..95ca6e868 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -12,6 +12,19 @@ import torch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker): + from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0 + + sigma = mocker.MagicMock() + sigma.to.return_value.item.return_value = 0.5 + + actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma) + + sigma.to.assert_called_once_with(torch.float32) + sigma.to.return_value.item.assert_called_once_with() + torch.testing.assert_close(actual, torch.tensor([2.0])) + + class TestCFGParallelHelpers: """Test LTX-2.3 CFG helper math without loading model weights.""" @@ -199,6 +212,7 @@ class TestCFGParallelHelpers: ), attention_kwargs=None, audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])), + original_audio_num_frames=state.audio.shape[1], ) actual_video, actual_audio = executor.predict_parallel_guidance( diff --git a/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py new file mode 100644 index 000000000..4b1d18cb5 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py @@ -0,0 +1,483 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the fixed LTX refinement-phase adapter.""" + +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter +from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import ( + AdapterTarget, + parse_ltx_adapter, +) +from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE +from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime +from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_TWO_STAGE_RECIPE, + LTX_DISTILLED_ADAPTER_SLOT, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _target( + lora_a, + lora_b, + *, + source_module="proj", + module="proj", + target_slice=None, +): + return AdapterTarget( + source_module=source_module, + module=module, + slice=target_slice, + a_key="a", + b_key="b", + a_shape=tuple(lora_a.shape), + b_shape=tuple(lora_b.shape), + ) + + +def _wrap( + layer, + lora_a, + lora_b, + *, + source_module="proj", + module="proj", + target_slice=None, + layer_fused=False, +): + target = _target( + lora_a, + lora_b, + source_module=source_module, + module=module, + target_slice=target_slice, + ) + wrapper = ltx2_phase_adapter._PhaseAdapterLinear( + layer, + [target], + adapter_dtype=lora_a.dtype, + layer_fused=layer_fused, + ) + wrapper.load_target(target, lora_a, lora_b) + wrapper.set_enabled(True) + return wrapper + + +class _LoRAAttention(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.to_qkv = torch.nn.Identity() + + +class _LoRATransformer(torch.nn.Module): + stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),) + + def __init__(self) -> None: + super().__init__() + self.proj_in = torch.nn.Linear(4, 6, bias=False) + block = torch.nn.Module() + block.attn1 = _LoRAAttention() + self.transformer_blocks = torch.nn.ModuleList([block]) + + +def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path): + adapter_path = tmp_path / "distilled.safetensors" + save_file( + { + "diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4), + "diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4), + "diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2), + "diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4), + "diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2), + }, + adapter_path, + ) + + manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path)) + assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == { + ("proj_in", "proj_in", None), + ("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"), + } + + +def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path): + adapter_path = tmp_path / "distilled.safetensors" + save_file( + {"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)}, + adapter_path, + ) + + with pytest.raises(ValueError, match="Missing paired LoRA A tensors"): + parse_ltx_adapter(_LoRATransformer(), str(adapter_path)) + + +def _pipeline(*, dtype, quantization_config): + return SimpleNamespace( + pipeline_recipe=LTX2_TWO_STAGE_RECIPE, + component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE, + transformer=object(), + od_config=SimpleNamespace( + model="unused", + lora_path=None, + dtype=dtype, + quantization_config=quantization_config, + ), + ) + + +@pytest.mark.parametrize( + ("dtype", "quantization_config", "layer_fused"), + [ + (torch.float32, object(), False), + (torch.bfloat16, None, True), + ], +) +def test_ltx_phase_adapter_selects_execution_from_quantization( + monkeypatch, + dtype, + quantization_config, + layer_fused, +): + pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config) + manifest = object() + captured = {} + + def resolve_artifact(*_args, **kwargs): + assert not kwargs + return "adapter" + + def parse(transformer, path): + assert transformer is pipeline.transformer + assert path == "adapter" + return manifest + + class Runtime: + def __init__(self, transformer, parsed_manifest, **kwargs): + captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs) + + monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact) + monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse) + monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime) + + runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline) + + assert isinstance(runtime, Runtime) + assert captured == { + "transformer": pipeline.transformer, + "manifest": manifest, + "dtype": dtype, + "layer_fused": layer_fused, + } + + +def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch): + pipeline = _pipeline(dtype=torch.float16, quantization_config=None) + monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter") + + with pytest.raises(ValueError, match="bfloat16"): + ltx2_phase_adapter.build_ltx_phase_adapter(pipeline) + + +def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path): + class TinyTransformer(torch.nn.Module): + stacked_params_mapping = () + + def __init__(self) -> None: + super().__init__() + self.proj = torch.nn.Linear(2, 2, bias=False) + + adapter_path = tmp_path / "distilled.safetensors" + save_file( + { + "diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]), + "diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]), + }, + adapter_path, + ) + transformer = TinyTransformer() + runtime = LTXPhaseAdapterRuntime( + transformer, + parse_ltx_adapter(transformer, str(adapter_path)), + dtype=torch.float32, + ) + + assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight" + assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"} + with torch.no_grad(): + transformer.proj.base_layer.weight.copy_(torch.eye(2)) + runtime.finalize() + + x = torch.tensor([[1.0, 1.0]]) + runtime.activate(None) + torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]])) + runtime.activate(LTX_DISTILLED_ADAPTER_SLOT) + torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]])) + + +def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion(): + layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16) + base_weight = torch.tensor( + [ + [0.5, -0.25, 0.125, 0.75], + [-0.375, 0.625, 0.25, -0.5], + [0.875, -0.125, -0.75, 0.375], + ], + dtype=torch.bfloat16, + ) + lora_a = torch.tensor( + [ + [0.1015625, -0.203125, 0.3046875, -0.40625], + [0.5078125, -0.609375, 0.7109375, -0.8125], + ], + dtype=torch.bfloat16, + ) + lora_b = torch.tensor( + [[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]], + dtype=torch.bfloat16, + ) + with torch.no_grad(): + layer.weight.copy_(base_weight) + layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16)) + wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True) + + expected_weight = torch.matmul(lora_b, lora_a) + expected_weight.add_(base_weight) + x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16) + expected = torch.nn.functional.linear(x, expected_weight, layer.bias) + + assert torch.equal(wrapper(x), expected) + assert torch.equal(layer.weight, base_weight) + + +def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + class FakeRow(torch.nn.Module): + input_size = 4 + output_size = 2 + input_size_per_partition = 2 + input_is_parallel = True + tp_rank = 0 + tp_size = 2 + reduce_results = True + skip_bias_add = False + return_bias = False + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16)) + self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16)) + self.quant_method = UnquantizedLinearMethod() + + monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow) + reductions = [] + + def all_reduce(value): + reductions.append(value.clone()) + return value + 1 + + monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce) + layer = FakeRow() + lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16) + lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16) + wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True) + + expected_weight = torch.matmul(lora_b, lora_a[:, :2]) + expected_weight.add_(layer.weight) + x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16) + expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias) + + actual = wrapper(x) + assert len(reductions) == 1 + assert torch.equal(reductions[0], expected_local) + assert torch.equal(actual, expected_local + 1) + + +def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + class FakeColumn(torch.nn.Module): + input_size = 3 + output_size = 4 + output_size_per_partition = 2 + tp_rank = 1 + tp_size = 2 + gather_output = False + skip_bias_add = False + return_bias = False + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16) + ) + self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16)) + self.quant_method = UnquantizedLinearMethod() + + monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn) + layer = FakeColumn() + base_weight = layer.weight.detach().clone() + lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16) + lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16) + wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True) + + expected_weight = torch.matmul(lora_b[2:], lora_a) + expected_weight.add_(base_weight) + x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16) + expected = torch.nn.functional.linear(x, expected_weight, layer.bias) + + assert torch.equal(wrapper(x), expected) + assert torch.equal(layer.weight, base_weight) + + +def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + class FakeColumn(torch.nn.Module): + pass + + class FakeQKV(FakeColumn): + input_size = 4 + head_size = 2 + v_head_size = 2 + total_num_heads = 2 + total_num_kv_heads = 2 + num_heads = 1 + num_kv_heads = 1 + num_kv_head_replicas = 1 + tp_rank = 1 + tp_size = 2 + gather_output = False + skip_bias_add = True + return_bias = True + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16 + ) + self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8) + self.quant_method = UnquantizedLinearMethod() + + monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn) + monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV) + layer = FakeQKV() + base_weight = layer.weight.detach().clone() + lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16) + lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16) + wrapper = _wrap( + layer, + lora_a, + lora_b, + source_module="to_k", + module="to_qkv", + target_slice="k", + layer_fused=True, + ) + + expected_weight = base_weight.clone() + expected_k = torch.matmul(lora_b[2:], lora_a) + expected_k.add_(base_weight[2:4]) + expected_weight[2:4].copy_(expected_k) + x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16) + expected_output = torch.nn.functional.linear(x, expected_weight) + + actual_output, returned_bias = wrapper(x) + assert torch.equal(actual_output, expected_output) + assert returned_bias is layer.bias + assert torch.equal(layer.weight, base_weight) + + +@pytest.mark.parametrize( + ("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"), + [ + ( + lambda: type( + "FakeColumn", + (torch.nn.Module,), + { + "input_size": 4, + "output_size": 4, + "output_size_per_partition": 2, + "tp_rank": 1, + "gather_output": False, + "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)), + }, + )(), + None, + torch.tensor([[1.0, 10.0, 100.0, 1000.0]]), + torch.tensor([[1.0], [2.0], [3.0], [4.0]]), + torch.tensor([[12963.0, 17284.0]]), + ), + ( + lambda: type( + "FakeRow", + (torch.nn.Module,), + { + "input_size": 4, + "output_size": 2, + "input_size_per_partition": 2, + "tp_rank": 1, + "reduce_results": False, + "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)), + }, + )(), + None, + torch.tensor([[1.0, 10.0, 100.0, 1000.0]]), + torch.tensor([[2.0], [3.0]]), + torch.tensor([[8600.0, 12900.0]]), + ), + ( + lambda: type( + "FakeQKV", + (torch.nn.Module,), + { + "input_size": 4, + "head_size": 2, + "v_head_size": 2, + "total_num_heads": 2, + "total_num_kv_heads": 2, + "num_heads": 1, + "num_kv_heads": 1, + "num_kv_head_replicas": 1, + "tp_rank": 1, + "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)), + }, + )(), + "k", + torch.tensor([[1.0, 10.0, 100.0, 1000.0]]), + torch.tensor([[1.0], [2.0], [3.0], [4.0]]), + torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]), + ), + ], +) +def test_ltx_phase_adapter_uses_local_tp_lora_slices( + monkeypatch, + layer_factory, + target_slice, + expected_a, + expected_b, + expected_output, +): + layer = layer_factory() + layer.register_buffer("device_anchor", torch.empty(0)) + layer_type = type(layer) + if target_slice == "k": + monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type) + elif "Column" in layer_type.__name__: + monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type) + else: + monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type) + + wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice) + torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output) + expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size) + assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index d70820f41..2f3ae522e 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -13,13 +13,21 @@ import pytest import torch from vllm_omni.diffusion.data import DiffusionOutput -from vllm_omni.diffusion.models.ltx2 import ltx2_components +from vllm_omni.diffusion.models.ltx2 import ( + ltx2_components, + ltx2_transformer, +) from vllm_omni.diffusion.models.ltx2.ltx2_components import ( LTX2_COMPONENT_PROFILE, LTX2_DISTILLED_COMPONENT_PROFILE, + LTX2_TWO_STAGE_COMPONENT_PROFILE, LTX23_COMPONENT_PROFILE, + LTX23_DISTILLED_COMPONENT_PROFILE, + LTX23_TWO_STAGE_COMPONENT_PROFILE, _install_connector_attention, detect_ltx_model_version, + resolve_ltx_artifact, + resolve_ltx_component_profile, ) from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin from vllm_omni.diffusion.models.ltx2.ltx2_denoise import ( @@ -33,6 +41,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( LTXGuidancePlan, LTXGuidanceSpec, LTXModalityGuidance, + _repeat_batch, build_perturbation_kwargs, combine_guided_x0, ) @@ -40,18 +49,29 @@ from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( LTX2_DISTILLED_TWO_STAGE_RECIPE, LTX2_ONE_STAGE_RECIPE, + LTX2_TWO_STAGE_RECIPE, + LTX23_DISTILLED_TWO_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE, + LTX23_TWO_STAGE_RECIPE, LTX_DEFAULT_NEGATIVE_PROMPT, LTX_POSITIVE_ONLY_RECIPE, + resolve_ltx_pipeline_recipe, +) +from vllm_omni.diffusion.models.ltx2.ltx2_request import ( + LTXRequestInputs, + validate_ltx_checkpoint, + validate_pipeline_request, ) -from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( LTX2I2VDMD2Pipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( + LTX2DistilledPipeline, + LTX2TwoStagePipeline, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -104,10 +124,22 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries(): ) assert issubclass(LTX2Pipeline, LTXRuntime) + assert issubclass(LTX2TwoStagePipeline, LTXRuntime) assert issubclass(LTX2DistilledPipeline, LTXRuntime) assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE + assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE + assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE + for pipeline_cls, profile in ( + (LTX2Pipeline, LTX2_COMPONENT_PROFILE), + (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE), + (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE), + ): + assert pipeline_cls._dit_modules == list(profile.dit_modules) + assert pipeline_cls._encoder_modules == list(profile.encoder_modules) + assert pipeline_cls._vae_modules == list(profile.vae_modules) + assert pipeline_cls._resident_modules == list(profile.resident_modules) assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video @@ -123,6 +155,72 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path): assert detect_ltx_model_version(str(tmp_path)) == "2" +def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch): + filename = "ltx-sidecar.safetensors" + calls = [] + + def fake_download(**kwargs): + calls.append(kwargs) + return "/cache/ltx-sidecar.safetensors" + + monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) + + assert ( + resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors" + ) + assert calls == [{"repo_id": "Lightricks/test", "filename": filename}] + + +def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): + filename = "ltx-sidecar.safetensors" + expected = tmp_path / filename + expected.write_bytes(b"sidecar") + monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup")) + + assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected) + + +@pytest.mark.parametrize( + ("expected_kind", "scheduler_config"), + [ + ("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}), + ("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}), + ], +) +def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config): + validate_ltx_checkpoint( + scheduler_config, + expected_kind=expected_kind, + pipeline_name="TestLTXPipeline", + ) + + +@pytest.mark.parametrize( + ("expected_kind", "scheduler_config", "error"), + [ + ("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"), + ("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"), + ], +) +def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error): + with pytest.raises(ValueError, match=error): + validate_ltx_checkpoint( + scheduler_config, + expected_kind=expected_kind, + pipeline_name="TestLTXPipeline", + ) + + +def test_ltx_rejects_advanced_uaa_before_component_initialization(): + od_config = SimpleNamespace( + model="unused", + parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"), + ) + + with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"): + LTX2Pipeline(od_config=od_config) + + def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch): from vllm_omni.diffusion.models.ltx2 import ltx2_runtime @@ -144,16 +242,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m assert pipe.reports_stage_durations -def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts(): +def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]})) + initialized = False + + def stub_components(*_args, **_kwargs): + nonlocal initialized + initialized = True + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + + with pytest.raises(ValueError, match="only one-stage LTX recipes"): + LTX2TwoStagePipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + cache_backend="cache_dit", + ) + ) + + assert not initialized + + +def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]})) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace()) + pipe = LTX2TwoStagePipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + enable_diffusion_pipeline_profiler=False, + lora_path=None, + model_config={}, + ) + ) + + assert pipe.model_version == "2.3" + assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE + + +def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]})) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = LTX2DistilledPipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + enable_diffusion_pipeline_profiler=False, + ) + ) + + assert pipe.model_version == "2.3" + assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE + assert pipe._phase_adapter is None + assert pipe.preserve_sp_padded_audio_duration + assert pipe.reports_stage_durations + + +def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts(): assert LTX2Pipeline.supports_request_batch assert LTX2Pipeline.support_image_input assert LTX2Pipeline.unified_text_image_entry assert LTX2Pipeline.distributed_video_decode + assert not LTX2TwoStagePipeline.supports_request_batch + assert LTX2TwoStagePipeline.support_image_input + assert LTX2TwoStagePipeline.unified_text_image_entry assert not LTX2DistilledPipeline.supports_request_batch assert LTX2DistilledPipeline.support_image_input assert LTX2DistilledPipeline.unified_text_image_entry +def test_ltx_two_stage_resolves_request_embedded_images(): + pipe = object.__new__(LTX2TwoStagePipeline) + request_inputs = SimpleNamespace(latents=None) + image = torch.zeros(3, 8, 8) + req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}]) + + assert pipe._resolve_request_image(req, None, request_inputs) is image + + def test_ltx_dmd2_entries_retain_their_task_boundaries(): assert not LTX2T2VDMD2Pipeline.support_image_input assert LTX2I2VDMD2Pipeline.support_image_input @@ -330,6 +516,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics(): assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",) +def test_ltx_positive_only_guidance_preserves_token_major_layout(): + tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2) + + repeated = _repeat_batch(tensor, 1) + + assert repeated is tensor + assert repeated.stride() == tensor.stride() + + def test_ltx_guidance_rescale_is_invariant_to_request_batching(): cond = torch.tensor( [ @@ -380,6 +575,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction(): torch.testing.assert_close(actual, torch.zeros_like(cond)) +def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens(): + logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]]) + logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]]) + guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0) + expected = combine_guided_x0( + cond=logical_cond, + uncond_text=logical_uncond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1) + padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1) + + actual = combine_guided_x0( + cond=padded_cond, + uncond_text=padded_uncond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + rescale_token_count=2, + ) + + torch.testing.assert_close(actual[:, :2], expected) + + def test_ltx_guidance_builds_per_sample_stg_and_modality_masks(): plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1)) @@ -547,11 +768,19 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi def test_ltx2_distilled_pipeline_shares_recipe_runtime(): assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin) - assert issubclass(LTX2DistilledPipeline, LTX2Pipeline) + assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline) assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe assert "_forward_request" not in LTX2DistilledPipeline.__dict__ +def test_ltx_two_stage_entry_is_a_thin_pipeline_variant(): + assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__ + assert "_forward_request" not in LTX2TwoStagePipeline.__dict__ + assert "load_weights" not in LTX2TwoStagePipeline.__dict__ + assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin) + assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline) + + def test_ltx_variants_share_denoise_loop_and_i2v_conditioning(): assert "_denoise_loop" not in LTX2Pipeline.__dict__ assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin) @@ -594,7 +823,8 @@ def test_denoise_executor_owns_progress_and_interrupt(): torch.testing.assert_close(state.audio, torch.tensor(11.0)) -def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): +@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline]) +def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls): request_inputs = LTXRequestInputs( prompt="prompt", negative_prompt="", @@ -671,12 +901,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) prompt_context_sentinel = prompt_context - pipeline = object.__new__(LTX2DistilledPipeline) + pipeline = object.__new__(pipeline_cls) torch.nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") pipeline.vae_spatial_compression_ratio = 32 pipeline.vae_temporal_compression_ratio = 8 pipeline.latent_upsampler = FakeUpsampler() + activated_slots = [] + if pipeline_cls is LTX2TwoStagePipeline: + pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append) object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) object.__setattr__(pipeline, "run_phase", run_phase) object.__setattr__(pipeline, "decode_phase", decode_phase) @@ -686,14 +919,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): assert len(phase_calls) == 2 assert phase_calls[1][2] is prompt_context_sentinel + if pipeline_cls is LTX2TwoStagePipeline: + assert activated_slots == [None, "ltx_distilled"] torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0)) - torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0)) + expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0 + torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio)) -def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning(): +@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline]) +def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls): from vllm_omni.diffusion.models.ltx2 import ltx2_latents - pipeline = object.__new__(LTX2DistilledPipeline) + pipeline = object.__new__(pipeline_cls) torch.nn.Module.__init__(pipeline) pipeline.vae_spatial_compression_ratio = 32 pipeline.vae_temporal_compression_ratio = 8 @@ -750,6 +987,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only(): assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536) +def test_ltx23_full_distilled_route_uses_merged_weights_without_lora(): + assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE + assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE + assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases) + assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None + assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls + assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == ( + "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" + ) + + @pytest.mark.parametrize( ("request_inputs", "error"), [ @@ -958,6 +1207,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps) +@pytest.mark.parametrize( + ("recipe", "steps", "stg_block"), + [ + (LTX2_TWO_STAGE_RECIPE, 40, 29), + (LTX23_TWO_STAGE_RECIPE, 30, 28), + ], +) +def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block): + stage1, stage2 = recipe.phases + assert (recipe.width, recipe.height) == (1536, 1024) + assert recipe.num_inference_steps == steps + assert stage1.spatial_downscale == 2 + assert stage1.adapter_slot is None + assert stage1.guidance.video.stg_blocks == (stg_block,) + assert stage1.guidance.audio.stg_blocks == (stg_block,) + assert stage2.input_transform == "spatial_upsample" + assert stage2.adapter_slot == "ltx_distilled" + assert stage2.guidance == LTXGuidanceSpec.positive_only() + assert stage2.num_inference_steps == 3 + assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0) + # Official Stage 2 refines video only; its audio result is discarded. + assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0) + + +def test_ltx_two_stage_entries_select_the_official_adapter_slots(): + phase_switches = [] + ordinary_pipeline = object.__new__(LTX2TwoStagePipeline) + ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append) + for phase in LTX2_TWO_STAGE_RECIPE.phases: + ordinary_pipeline._enter_phase(phase) + assert phase_switches == [None, "ltx_distilled"] + + distilled_pipeline = object.__new__(LTX2DistilledPipeline) + for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases: + distilled_pipeline._enter_phase(phase) + + +def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch): + class FakeColumn(torch.nn.Module): + def __init__( + self, + input_size, + output_size, + *, + bias, + gather_output, + return_bias, + quant_config, + prefix, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.bias_enabled = bias + self.gather_output = gather_output + self.return_bias = return_bias + self.quant_config = quant_config + self.prefix = prefix + + class FakeRow(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + class FakeAttention(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + + monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn) + monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow) + monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention) + monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2) + + attention = ltx2_transformer.LTX2Attention( + query_dim=16, + heads=4, + dim_head=4, + qk_norm="rms_norm", + apply_gated_attention=True, + pack_qkv=False, + prefix="transformer_blocks.0.attn1", + ) + + gate = attention.to_gate_logits + assert isinstance(gate, FakeColumn) + assert gate.input_size == 16 + assert gate.output_size == 4 + assert gate.gather_output is False + assert gate.return_bias is False + assert gate.quant_config is None + assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits" + + class TestLTXRequestParsing: def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self): from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -1432,6 +1773,7 @@ class TestPipelineComponents: def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self): assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler") + assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler") class TestLTX23DecodeConditioning: @@ -1477,6 +1819,11 @@ class TestRegistryIntegration: from vllm_omni.diffusion.registry import _DIFFUSION_MODELS assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline") + assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2TwoStagePipeline", + ) assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == ( "ltx2", "pipeline_ltx2_two_stage", @@ -1488,6 +1835,9 @@ class TestRegistryIntegration: "LTX23ImageToVideoPipeline", "LTX2TwoStagesPipeline", "LTX2ImageToVideoTwoStagesPipeline", + "LTX2ImageToVideoTwoStagePipeline", + "LTX23TwoStagePipeline", + "LTX23ImageToVideoTwoStagePipeline", } assert removed_entries.isdisjoint(_DIFFUSION_MODELS) @@ -1497,6 +1847,7 @@ class TestRegistryIntegration: expected = [ "LTX2Pipeline", + "LTX2TwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", @@ -1509,6 +1860,7 @@ class TestRegistryIntegration: "model_class_name", [ "LTX2Pipeline", + "LTX2TwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index 73e10463e..f74b2991d 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -1,13 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest +import torch from torch import nn +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +from vllm_omni.diffusion.models.ltx2.ltx2_transformer import ( + LTX2AudioVideoAttnProcessor, + LTX2VideoTransformer3DModel, + _LTX2ParallelAttention, + _make_rms_norm, +) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -66,3 +73,83 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type): assert video_timestep.expected_dims == audio_timestep.expected_dims == 2 assert not video_timestep.split_output assert not audio_timestep.split_output + + +def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch): + monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True)) + processor = LTX2AudioVideoAttnProcessor() + hidden_states = torch.zeros(1, 3, 4) + encoder_hidden_states = torch.zeros(1, 4, 4) + key_padding_mask = torch.tensor([[True, True, True, False]]) + + actual = processor._prepare_attention_mask( + SimpleNamespace(), + hidden_states, + encoder_hidden_states, + key_padding_mask, + batch_size=1, + sequence_length=4, + ) + + torch.testing.assert_close(actual, key_padding_mask[:, None, None, :]) + + +def _fake_ltx_parallel_attention(): + calls = [] + + class Backend: + @staticmethod + def get_name(): + return "FLASH_ATTN" + + @staticmethod + def supports_attention_mask(): + return True + + def run(name): + def forward(query, _key, _value, _metadata): + calls.append(name) + return torch.zeros_like(query) + + return forward + + layer = object.__new__(_LTX2ParallelAttention) + nn.Module.__init__(layer) + layer.attn_backend = Backend + layer.attention = SimpleNamespace(forward=run("native")) + layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa")) + layer.backend_pref = "FLASH_ATTN" + return layer, calls + + +@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)]) +def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length): + layer, calls = _fake_ltx_parallel_attention() + query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16) + key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16) + key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]]) + metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :]) + + layer._run_local_attention(query, key, value, metadata) + + assert calls == ["sdpa"] + + +def test_ltx_masked_self_attention_uses_sdpa(): + layer, calls = _fake_ltx_parallel_attention() + query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16) + key_padding_mask = torch.tensor([[True, True, True, False]]) + metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :]) + + layer._run_local_attention(query, key, value, metadata) + + assert calls == ["sdpa"] + + +def test_ltx_unmasked_cross_attention_keeps_native_backend(): + layer, calls = _fake_ltx_parallel_attention() + query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16) + + layer._run_local_attention(query, key, value, None) + + assert calls == ["native"] diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py index a4a180047..3ab794662 100644 --- a/tests/e2e/accuracy/ltx/run_ltx_reference.py +++ b/tests/e2e/accuracy/ltx/run_ltx_reference.py @@ -20,6 +20,11 @@ from PIL import Image def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--backend", choices=("official", "omni"), required=True) + parser.add_argument( + "--pipeline-kind", + choices=("one_stage", "distilled", "two_stage"), + default="one_stage", + ) parser.add_argument("--request", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--model") @@ -27,6 +32,8 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--official-root", type=Path) parser.add_argument("--checkpoint", type=Path) parser.add_argument("--gemma-root", type=Path) + parser.add_argument("--spatial-upsampler", type=Path) + parser.add_argument("--distilled-lora", type=Path) parser.add_argument("--enable-layerwise-offload", action="store_true") return parser.parse_args() @@ -91,9 +98,13 @@ def _configure_official_sdpa(pipeline: Any) -> None: attention=attention, masked_attention=attention, ) - owners_and_attributes = ( - (pipeline.stage, "_transformer_builder"), + owners_and_attributes = [ (pipeline.prompt_encoder, "_embeddings_processor_builder"), + ] + owners_and_attributes.extend( + (stage, "_transformer_builder") + for name in ("stage", "stage_1", "stage_2") + if (stage := getattr(pipeline, name, None)) is not None ) for owner, attribute in owners_and_attributes: builder = getattr(owner, attribute) @@ -104,6 +115,12 @@ def _configure_official_sdpa(pipeline: Any) -> None: ) +def _require_path(path: Path | None, description: str) -> str: + if path is None or not path.is_file(): + raise ValueError(f"Official {description} is required and must exist: {path}") + return str(path) + + @torch.inference_mode() def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: if args.official_root is None or args.checkpoint is None or args.gemma_root is None: @@ -113,16 +130,54 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: torch.backends.cuda.enable_cudnn_sdp(False) from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps + from ltx_pipelines.distilled import DistilledPipeline from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline from ltx_pipelines.utils.args import ImageConditioningInput from ltx_pipelines.utils.types import OffloadMode - pipeline = TI2VidOneStagePipeline( - checkpoint_path=str(args.checkpoint), - gemma_root=str(args.gemma_root), - loras=(), - offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE, - ) + if request.get("pipeline_kind", "one_stage") != args.pipeline_kind: + raise ValueError( + f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}" + ) + checkpoint = _require_path(args.checkpoint, "checkpoint") + gemma_root = str(args.gemma_root) + if not Path(gemma_root).is_dir(): + raise ValueError(f"Official Gemma root is required and must be a directory: {gemma_root}") + offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE + pipeline: Any + if args.pipeline_kind == "one_stage": + pipeline = TI2VidOneStagePipeline( + checkpoint_path=checkpoint, + gemma_root=gemma_root, + loras=(), + offload_mode=offload_mode, + ) + elif args.pipeline_kind == "distilled": + pipeline = DistilledPipeline( + distilled_checkpoint_path=checkpoint, + spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"), + gemma_root=gemma_root, + loras=(), + offload_mode=offload_mode, + ) + else: + distilled_lora = [ + LoraPathStrengthAndSDOps( + path=_require_path(args.distilled_lora, "distilled LoRA"), + strength=1.0, + sd_ops=LTXV_LORA_COMFY_RENAMING_MAP, + ) + ] + pipeline = TI2VidTwoStagesPipeline( + checkpoint_path=checkpoint, + distilled_lora=distilled_lora, + spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"), + gemma_root=gemma_root, + loras=(), + offload_mode=offload_mode, + ) _configure_official_sdpa(pipeline) image_path = request.get("image") images = ( @@ -140,34 +195,45 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: ) ] ) - video, audio = pipeline( - prompt=request["prompt"], - negative_prompt=request["negative_prompt"], - seed=request["seed"], - height=request["height"], - width=request["width"], - num_frames=request["num_frames"], - frame_rate=request["fps"], - num_inference_steps=request["num_inference_steps"], - video_guider_params=MultiModalGuiderParams( - cfg_scale=request["video_cfg_scale"], - stg_scale=request["video_stg_scale"], - rescale_scale=request["video_rescale_scale"], - modality_scale=request["video_modality_scale"], - skip_step=0, - stg_blocks=request["video_stg_blocks"], - ), - audio_guider_params=MultiModalGuiderParams( - cfg_scale=request["audio_cfg_scale"], - stg_scale=request["audio_stg_scale"], - rescale_scale=request["audio_rescale_scale"], - modality_scale=request["audio_modality_scale"], - skip_step=0, - stg_blocks=request["audio_stg_blocks"], - ), - images=images, - max_batch_size=4, - ) + if args.pipeline_kind == "distilled": + video, audio = pipeline( + prompt=request["prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + images=images, + ) + else: + video, audio = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + ) video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0) _save_outputs( args.output_dir, @@ -179,6 +245,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: "attention_backend": "torch_sdpa", "official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"), "checkpoint": str(args.checkpoint), + "pipeline_kind": args.pipeline_kind, }, ) @@ -253,6 +320,10 @@ def _canonical_video(video: Any) -> torch.Tensor: def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: if args.model is None or args.model_class_name is None: raise ValueError("Omni backend requires --model and --model-class-name") + if request.get("pipeline_kind", "one_stage") != args.pipeline_kind: + raise ValueError( + f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}" + ) attention_config = {"default": {"backend": "TORCH_SDPA"}} @@ -273,7 +344,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: parallel_config=DiffusionParallelConfig(), ) try: - model_class_name = get_model_class_name(omni) + detected_model_class_name = get_model_class_name(omni) + model_class_name = args.model_class_name sampling_params = OmniDiffusionSamplingParams( height=request["height"], width=request["width"], @@ -289,10 +361,9 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_") } apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance) - prompt: dict[str, Any] = { - "prompt": request["prompt"], - "negative_prompt": request["negative_prompt"], - } + prompt: dict[str, Any] = {"prompt": request["prompt"]} + if request.get("negative_prompt"): + prompt["negative_prompt"] = request["negative_prompt"] image_path = request.get("image") if image_path is not None: with Image.open(str(image_path)) as source_image: @@ -312,6 +383,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: "attention_backend": "torch_sdpa", "model": args.model, "model_class_name": model_class_name, + "model_config_class_name": detected_model_class_name, + "pipeline_kind": args.pipeline_kind, }, ) finally: diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py index 647ece7a7..8a3be4a99 100644 --- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py +++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py @@ -3,10 +3,9 @@ """E2E accuracy guard against a pinned Lightricks LTX pipeline revision. -The comparison runs both runtimes through PyTorch SDPA and uses -``max_batch_size=4`` in the official reference to match Omni's fused guidance -batch. Video and audio guidance use the official non-HQ one-stage defaults; -only the generation shape and step count are reduced for CI runtime. +The original reduced one-stage guards remain unchanged. Four complementary +default-shape cases cover both model versions, both two-stage weight families, +and T2V/I2V without expanding every combination into a duplicate golden. """ from __future__ import annotations @@ -17,8 +16,9 @@ import shutil import subprocess import sys import time -from dataclasses import dataclass +from dataclasses import asdict, dataclass, replace from pathlib import Path +from typing import Literal import numpy as np import pytest @@ -53,74 +53,260 @@ NEGATIVE_PROMPT = ( # Both runtimes use PyTorch SDPA with the current Torch dispatch defaults. ATTENTION_BACKEND = "torch_sdpa" -VIDEO_SSIM_MEAN_THRESHOLD = 0.95 -VIDEO_SSIM_MIN_THRESHOLD = 0.90 -VIDEO_PSNR_MEAN_THRESHOLD = 30.0 -AUDIO_RELATIVE_L2_THRESHOLD = 0.2 -AUDIO_COSINE_THRESHOLD = 0.95 + + +@dataclass(frozen=True) +class LTXArtifact: + repo_id: str + filename: str + revision: str + env: str + repo_type: Literal["model", "dataset"] = "model" + + +@dataclass(frozen=True) +class LTXAccuracyThresholds: + video_ssim_mean: float + video_ssim_min: float + video_psnr_mean_db: float + # Zero disables gating for waveform-sensitive audio that is not expected + # to align with the official output. + audio_relative_l2: float + audio_cosine_similarity: float + + +STRICT_THRESHOLDS = LTXAccuracyThresholds( + video_ssim_mean=0.95, + video_ssim_min=0.90, + video_psnr_mean_db=30.0, + audio_relative_l2=0.20, + audio_cosine_similarity=0.95, +) @dataclass(frozen=True) class LTXAccuracyCase: name: str + pipeline_kind: Literal["one_stage", "distilled", "two_stage"] model_id: str model_revision: str model_env: str model_class_name: str - checkpoint_repo: str - checkpoint_filename: str - checkpoint_revision: str - checkpoint_env: str - stg_block: int + checkpoint: LTXArtifact + width: int + height: int + num_frames: int + num_inference_steps: int + seed: int + stg_block: int | None + thresholds: LTXAccuracyThresholds prompt: str = PROMPT - image_repo: str | None = None - image_filename: str | None = None - image_revision: str | None = None + gemma_model_id: str | None = None + gemma_model_revision: str | None = None + gemma_model_env: str | None = None + image: LTXArtifact | None = None + spatial_upsampler: LTXArtifact | None = None + distilled_lora: LTXArtifact | None = None + # Official uses block streaming; Omni uses layerwise offload. Both preserve + # bf16 arithmetic while keeping full-resolution two-stage cases on one H100. + enable_layerwise_offload: bool = False -CASES = ( +LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a" +LTX23_REVISION = "4229404625088d21c4f112eb640fb04a0900ee25" +LTX2_CHECKPOINT = LTXArtifact( + repo_id="Lightricks/LTX-2", + filename="ltx-2-19b-dev.safetensors", + revision=LTX2_REVISION, + env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT", +) +LTX2_DISTILLED_CHECKPOINT = replace( + LTX2_CHECKPOINT, + filename="ltx-2-19b-distilled.safetensors", + env="VLLM_TEST_LTX2_DISTILLED_OFFICIAL_CHECKPOINT", +) +LTX2_UPSAMPLER = replace( + LTX2_CHECKPOINT, + filename="ltx-2-spatial-upscaler-x2-1.0.safetensors", + env="VLLM_TEST_LTX2_UPSAMPLER", +) +LTX2_DISTILLED_LORA = replace( + LTX2_CHECKPOINT, + filename="ltx-2-19b-distilled-lora-384.safetensors", + env="VLLM_TEST_LTX2_DISTILLED_LORA", +) +LTX23_CHECKPOINT = LTXArtifact( + repo_id="Lightricks/LTX-2.3", + filename="ltx-2.3-22b-dev.safetensors", + revision=LTX23_REVISION, + env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT", +) +LTX23_DISTILLED_CHECKPOINT = replace( + LTX23_CHECKPOINT, + # The pinned Diffusers checkpoint contains the original merged distilled + # Transformer. Keep the reference Transformer on the same version; the + # spatial upsampler is versioned independently and remains on 1.1 below. + filename="ltx-2.3-22b-distilled.safetensors", + env="VLLM_TEST_LTX23_DISTILLED_OFFICIAL_CHECKPOINT", +) +LTX23_UPSAMPLER = replace( + LTX23_CHECKPOINT, + filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", + env="VLLM_TEST_LTX23_UPSAMPLER", +) +LTX23_DISTILLED_LORA = replace( + LTX23_CHECKPOINT, + filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors", + env="VLLM_TEST_LTX23_DISTILLED_LORA", +) +I2V_IMAGE = LTXArtifact( + repo_id="huggingface/documentation-images", + filename="diffusers/svd/rocket.png", + revision="645d8364f0c7a101180b364811b5a11a362e4010", + env="VLLM_TEST_LTX_I2V_IMAGE", + repo_type="dataset", +) + + +LEGACY_CASES = ( LTXAccuracyCase( name="ltx2", + pipeline_kind="one_stage", model_id="Lightricks/LTX-2", - model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a", + model_revision=LTX2_REVISION, model_env="VLLM_TEST_LTX2_MODEL", model_class_name="LTX2Pipeline", - checkpoint_repo="Lightricks/LTX-2", - checkpoint_filename="ltx-2-19b-dev.safetensors", - checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a", - checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT", + checkpoint=LTX2_CHECKPOINT, + width=512, + height=384, + num_frames=25, + num_inference_steps=20, + seed=42, stg_block=29, + thresholds=STRICT_THRESHOLDS, ), LTXAccuracyCase( name="ltx2_3", + pipeline_kind="one_stage", model_id="diffusers/LTX-2.3-Diffusers", model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf", model_env="VLLM_TEST_LTX23_MODEL", model_class_name="LTX2Pipeline", - checkpoint_repo="Lightricks/LTX-2.3", - checkpoint_filename="ltx-2.3-22b-dev.safetensors", - checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25", - checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT", + checkpoint=LTX23_CHECKPOINT, + width=512, + height=384, + num_frames=25, + num_inference_steps=20, + seed=42, stg_block=28, + thresholds=STRICT_THRESHOLDS, ), LTXAccuracyCase( name="ltx2_3_i2v", + pipeline_kind="one_stage", model_id="diffusers/LTX-2.3-Diffusers", model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf", model_env="VLLM_TEST_LTX23_MODEL", model_class_name="LTX2Pipeline", - checkpoint_repo="Lightricks/LTX-2.3", - checkpoint_filename="ltx-2.3-22b-dev.safetensors", - checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25", - checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT", + checkpoint=LTX23_CHECKPOINT, + width=512, + height=384, + num_frames=25, + num_inference_steps=20, + seed=42, stg_block=28, - image_repo="huggingface/documentation-images", - image_filename="diffusers/svd/rocket.png", - image_revision="645d8364f0c7a101180b364811b5a11a362e4010", + thresholds=STRICT_THRESHOLDS, + image=I2V_IMAGE, ), ) +TWO_STAGE_CASES = ( + LTXAccuracyCase( + name="ltx2_distilled_t2v", + pipeline_kind="distilled", + model_id="rootonchair/LTX-2-19b-distilled", + model_revision="388e2846f54aae51687498ffb6b27c7c2c9ce9e5", + model_env="VLLM_TEST_LTX2_DISTILLED_MODEL", + model_class_name="LTX2DistilledPipeline", + checkpoint=LTX2_DISTILLED_CHECKPOINT, + spatial_upsampler=LTX2_UPSAMPLER, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=8, + seed=10, + stg_block=None, + gemma_model_id="Lightricks/LTX-2", + gemma_model_revision=LTX2_REVISION, + gemma_model_env="VLLM_TEST_LTX2_MODEL", + thresholds=STRICT_THRESHOLDS, + ), + LTXAccuracyCase( + name="ltx23_distilled_i2v", + pipeline_kind="distilled", + model_id="diffusers/LTX-2.3-Distilled-Diffusers", + model_revision="432e0d3c2d1769aaa4d295f9243f7062bf6b47ee", + model_env="VLLM_TEST_LTX23_DISTILLED_MODEL", + model_class_name="LTX2DistilledPipeline", + checkpoint=LTX23_DISTILLED_CHECKPOINT, + spatial_upsampler=LTX23_UPSAMPLER, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=8, + seed=10, + stg_block=None, + gemma_model_id="diffusers/LTX-2.3-Diffusers", + gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf", + gemma_model_env="VLLM_TEST_LTX23_MODEL", + thresholds=STRICT_THRESHOLDS, + image=I2V_IMAGE, + ), + LTXAccuracyCase( + name="ltx2_two_stage_layer_fused_t2v", + pipeline_kind="two_stage", + model_id="Lightricks/LTX-2", + model_revision=LTX2_REVISION, + model_env="VLLM_TEST_LTX2_MODEL", + model_class_name="LTX2TwoStagePipeline", + checkpoint=LTX2_CHECKPOINT, + spatial_upsampler=LTX2_UPSAMPLER, + distilled_lora=LTX2_DISTILLED_LORA, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=40, + seed=10, + stg_block=29, + enable_layerwise_offload=True, + thresholds=STRICT_THRESHOLDS, + ), + LTXAccuracyCase( + name="ltx23_two_stage_layer_fused_i2v", + pipeline_kind="two_stage", + model_id="diffusers/LTX-2.3-Diffusers", + model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf", + model_env="VLLM_TEST_LTX23_MODEL", + model_class_name="LTX2TwoStagePipeline", + checkpoint=LTX23_CHECKPOINT, + spatial_upsampler=LTX23_UPSAMPLER, + distilled_lora=LTX23_DISTILLED_LORA, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=30, + seed=10, + stg_block=28, + enable_layerwise_offload=True, + thresholds=STRICT_THRESHOLDS, + image=I2V_IMAGE, + ), +) + +CASES = (*LEGACY_CASES, *TWO_STAGE_CASES) + + def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None: start = time.perf_counter() subprocess.run(command, env=env, timeout=timeout, check=True) @@ -209,6 +395,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path: "model_index.json", "audio_vae/*", "connectors/*", + "latent_upsampler/*", "processor/*", "scheduler/*", "text_encoder/config.json", @@ -223,74 +410,82 @@ def _resolve_model(case: LTXAccuracyCase) -> Path: ) -def _resolve_gemma_root(model: Path) -> Path: +def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path: configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT") if configured_root: root = Path(configured_root) assert root.is_dir(), f"Gemma root not found: {root}" return root + if case.gemma_model_id is not None: + configured_model = os.environ.get(case.gemma_model_env or "") + if configured_model and Path(configured_model).is_dir(): + return Path(configured_model) + return Path( + snapshot_download( + repo_id=case.gemma_model_id, + revision=case.gemma_model_revision, + allow_patterns=[ + "processor/*", + "text_encoder/*", + "tokenizer/*", + ], + ) + ) return model -def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path: - configured_checkpoint = os.environ.get(case.checkpoint_env) - if configured_checkpoint: - checkpoint = Path(configured_checkpoint) - assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}" - return checkpoint - model_checkpoint = model / case.checkpoint_filename - if model_checkpoint.is_file(): - return model_checkpoint +def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path: + configured_path = os.environ.get(artifact.env) + if configured_path: + path = Path(configured_path) + assert path.is_file(), f"Configured LTX artifact not found: {path}" + return path + if model is not None: + model_path = model / artifact.filename + if model_path.is_file(): + return model_path return Path( hf_hub_download( - repo_id=case.checkpoint_repo, - filename=case.checkpoint_filename, - revision=case.checkpoint_revision, + repo_id=artifact.repo_id, + repo_type=None if artifact.repo_type == "model" else artifact.repo_type, + filename=artifact.filename, + revision=artifact.revision, ) ) def _resolve_image(case: LTXAccuracyCase) -> Path | None: - if case.image_filename is None: + if case.image is None: return None - if case.image_repo is None or case.image_revision is None: - raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.") - configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE") - if configured_image: - image = Path(configured_image) - assert image.is_file(), f"LTX I2V conditioning image not found: {image}" - return image - return Path( - hf_hub_download( - repo_id=case.image_repo, - repo_type="dataset", - filename=case.image_filename, - revision=case.image_revision, - ) - ) + return _resolve_artifact(case.image) def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]: request: dict[str, object] = { + "pipeline_kind": case.pipeline_kind, "prompt": case.prompt, - "negative_prompt": NEGATIVE_PROMPT, - "width": 512, - "height": 384, - "num_frames": 25, + "negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT, + "width": case.width, + "height": case.height, + "num_frames": case.num_frames, "fps": 24, - "num_inference_steps": 20, - "seed": 42, - "video_cfg_scale": 3.0, - "audio_cfg_scale": 7.0, - "video_stg_scale": 1.0, - "audio_stg_scale": 1.0, - "video_modality_scale": 3.0, - "audio_modality_scale": 3.0, - "video_rescale_scale": 0.7, - "audio_rescale_scale": 0.7, - "video_stg_blocks": [case.stg_block], - "audio_stg_blocks": [case.stg_block], + "num_inference_steps": case.num_inference_steps, + "seed": case.seed, } + if case.pipeline_kind != "distilled": + assert case.stg_block is not None + request.update( + video_cfg_scale=3.0, + audio_cfg_scale=7.0, + video_stg_scale=1.0, + audio_stg_scale=1.0, + video_modality_scale=3.0, + audio_modality_scale=3.0, + video_rescale_scale=0.7, + audio_rescale_scale=0.7, + video_stg_blocks=[case.stg_block], + audio_stg_blocks=[case.stg_block], + ) if image is not None: request["image"] = str(image.resolve()) return request @@ -299,20 +494,30 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]: def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]: assert reference.shape == prediction.shape assert reference.ndim == 4 and reference.shape[-1] == 3 + device = torch.device("cpu") + ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device) + psnr = PeakSignalNoiseRatio(data_range=1.0).to(device) ssim_scores: list[float] = [] psnr_scores: list[float] = [] - for reference_frame, prediction_frame in zip(reference, prediction, strict=True): - reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0) - prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0) - ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor))) - psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor))) - difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64)) + max_abs = 0.0 + absolute_error_sum = 0.0 + with torch.inference_mode(): + for reference_frame, prediction_frame in zip(reference, prediction, strict=True): + reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device) + prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device) + ssim_scores.append(float(ssim(prediction_tensor, reference_tensor))) + psnr_scores.append(float(psnr(prediction_tensor, reference_tensor))) + ssim.reset() + psnr.reset() + difference = np.abs(reference_frame.astype(np.float64) - prediction_frame.astype(np.float64)) + max_abs = max(max_abs, float(difference.max())) + absolute_error_sum += float(difference.sum()) return { "ssim_mean": float(np.mean(ssim_scores)), "ssim_min": float(np.min(ssim_scores)), "psnr_mean_db": float(np.mean(psnr_scores)), - "max_abs": float(difference.max()), - "mean_abs": float(difference.mean()), + "max_abs": max_abs, + "mean_abs": absolute_error_sum / reference.size, } @@ -343,13 +548,19 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f @pytest.mark.benchmark @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None: - """Compare official and Omni raw AV outputs from the same E2E request.""" +def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None: + """Compare official and Omni raw AV outputs from the same single-H100 request.""" + configured_artifact_root = os.environ.get("VLLM_TEST_LTX_ARTIFACT_ROOT") + if configured_artifact_root: + accuracy_artifact_root = Path(configured_artifact_root) + accuracy_artifact_root.mkdir(parents=True, exist_ok=True) output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name) official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official") model = _resolve_model(case) - gemma_root = _resolve_gemma_root(model) - checkpoint = _resolve_checkpoint(case, model) + gemma_root = _resolve_gemma_root(case, model) + checkpoint = _resolve_artifact(case.checkpoint, model) + spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None + distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None image = _resolve_image(case) request_path = output_root / "request.json" request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n") @@ -359,8 +570,13 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact str(runner), "--request", str(request_path), + "--pipeline-kind", + case.pipeline_kind, ] - if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}: + enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get( + "VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "" + ).lower() in {"1", "true", "yes", "on"} + if enable_layerwise_offload: runner_args.append("--enable-layerwise-offload") env = os.environ.copy() env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision @@ -372,10 +588,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact ) official_output = output_root / "official" - _run( - _official_runner_prefix() - + runner_args - + [ + official_args = _official_runner_prefix() + runner_args + official_args.extend( + [ "--backend", "official", "--output-dir", @@ -386,12 +601,16 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact str(checkpoint), "--gemma-root", str(gemma_root), - ], - env=env, + ] ) + if spatial_upsampler is not None: + official_args.extend(["--spatial-upsampler", str(spatial_upsampler)]) + if distilled_lora is not None: + official_args.extend(["--distilled-lora", str(distilled_lora)]) + _run(official_args, env=env) omni_output = output_root / "omni" - _run( + omni_args = ( [sys.executable] + runner_args + [ @@ -403,9 +622,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact str(model), "--model-class-name", case.model_class_name, - ], - env=env, + ] ) + _run(omni_args, env=env) official_metadata = json.loads((official_output / "metadata.json").read_text()) omni_metadata = json.loads((omni_output / "metadata.json").read_text()) @@ -423,18 +642,27 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact result = { "case": case.name, "task": "i2v" if image is not None else "t2v", + "pipeline_kind": case.pipeline_kind, + "model_class_name": case.model_class_name, "attention_backend": ATTENTION_BACKEND, "official_revision": official_revision, "model_revision": case.model_revision, - "checkpoint_revision": case.checkpoint_revision, + "checkpoint_revision": case.checkpoint.revision, + "spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None), + "distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None, + "enable_layerwise_offload": enable_layerwise_offload, + "thresholds": asdict(case.thresholds), + "request": _request(case, image), "video": video_metrics, "audio": audio_metrics, } (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n") print(json.dumps(result, indent=2)) - assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD - assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD - assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD - assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD - assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD + assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean + assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min + assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db + if case.thresholds.audio_relative_l2 > 0: + assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2 + if case.thresholds.audio_cosine_similarity > 0: + assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity diff --git a/tests/e2e/online_serving/test_ltx.py b/tests/e2e/online_serving/test_ltx.py new file mode 100644 index 000000000..43df5df2d --- /dev/null +++ b/tests/e2e/online_serving/test_ltx.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Online-serving smoke tests for the public LTX-2 pipeline entries.""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers") +BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers") +PROMPT = "A serene lake at sunset with mountains in the background." +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) +HSDP_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) + + +def _server( + model: str, + model_class_name: str, + *extra_args: str, +) -> OmniServerParams: + return OmniServerParams( + model=model, + server_args=[ + "--enforce-eager", + *extra_args, + "--model-class-name", + model_class_name, + ], + ) + + +def _cases(): + return [ + pytest.param( + _server(BASE_MODEL, "LTX2Pipeline"), + 3, + id="one_stage", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + _server(DISTILLED_MODEL, "LTX2DistilledPipeline"), + 8, + id="distilled", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + _server( + DISTILLED_MODEL, + "LTX2DistilledPipeline", + "--use-hsdp", + "--hsdp-shard-size", + "2", + ), + 8, + id="distilled_hsdp2", + marks=HSDP_MARKS, + ), + pytest.param( + _server( + BASE_MODEL, + "LTX2TwoStagePipeline", + "--enable-layerwise-offload", + ), + 3, + id="two_stage", + marks=SINGLE_CARD_MARKS, + ), + ] + + +@pytest.mark.parametrize( + ("omni_server", "num_inference_steps"), + _cases(), + indirect=["omni_server"], +) +def test_ltx_pipeline_entries( + omni_server: OmniServer, + num_inference_steps: int, + openai_client: OpenAIClientHandler, + subtests, +): + for task in ("t2v", "i2v"): + with subtests.test(task=task): + form_data = { + "prompt": PROMPT, + "model": omni_server.model, + "height": 128, + "width": 128, + "num_frames": 9, + "fps": 8, + "num_inference_steps": num_inference_steps, + "seed": 42, + } + request_config = { + "model": omni_server.model, + "form_data": form_data, + } + if task == "i2v": + request_config["image_reference"] = ( + f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + ) + + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py deleted file mode 100644 index 4d0eb6c90..000000000 --- a/tests/e2e/online_serving/test_ltx2_expansion.py +++ /dev/null @@ -1,91 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Expansion tests for LTX-2 two-stage pipelines. - -Coverage: -- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video) -""" - -import os - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler - -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] - -MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled") -PROMPT = "A serene lake at sunset with mountains in the background." -PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) -HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"] - - -def _cases(): - cases = [] - - cases.append( - pytest.param( - OmniServerParams( - model=MODEL, - server_args=[ - *HSDP_ARGS, - "--model-class-name", - "LTX2DistilledPipeline", - ], - ), - False, - id="t2v_hsdp", - marks=PARALLEL_MARKS, - ) - ) - - cases.append( - pytest.param( - OmniServerParams( - model=MODEL, - server_args=[ - *HSDP_ARGS, - "--model-class-name", - "LTX2DistilledPipeline", - ], - ), - True, - id="i2v_hsdp", - marks=PARALLEL_MARKS, - ) - ) - - return cases - - -@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"]) -def test_ltx2_two_stage_hsdp( - omni_server: OmniServer, - is_i2v: bool, - openai_client: OpenAIClientHandler, -): - # Keep CI small while exercising the fixed 8+3-step distilled recipe. - # Height and width describe the final output; Stage 1 runs at half size. - form_data = { - "prompt": PROMPT, - "model": omni_server.model, - "height": 128, - "width": 128, - "num_frames": 9, - "fps": 8, - "num_inference_steps": 8, - "seed": 42, - } - - request_config = { - "model": omni_server.model, - "form_data": form_data, - } - - if is_i2v: - request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" - - openai_client.send_video_diffusion_request(request_config) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index d57671adc..557fb5f6c 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -233,10 +233,13 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None: } ) - assert get_extra_body_params("LTX2Pipeline") == expected - assert get_extra_output_params("LTX2Pipeline") == frozenset() - assert get_extra_body_params("LTX2DistilledPipeline") == expected - assert get_extra_output_params("LTX2DistilledPipeline") == frozenset() + for pipeline_name in ( + "LTX2Pipeline", + "LTX2TwoStagePipeline", + "LTX2DistilledPipeline", + ): + assert get_extra_body_params(pipeline_name) == expected + assert get_extra_output_params(pipeline_name) == frozenset() @pytest.mark.core_model diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index be064f603..f5d6c18c4 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -27,6 +27,7 @@ EXCLUDED_MODELS = [ "OvisImagePipeline", "WanPipeline", "WanVACEPipeline", + "LTX2TwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index 300295bd8..20ada6adc 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -12,12 +12,16 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( LTX2Pipeline, LTX2T2VDMD2Pipeline, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( + LTX2DistilledPipeline, + LTX2TwoStagePipeline, +) __all__ = [ "LTX2Pipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", + "LTX2TwoStagePipeline", "LTX2DistilledPipeline", "get_ltx2_post_process_func", "load_transformer_config", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py b/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py new file mode 100644 index 000000000..d8228370e --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Parse official LTX distilled adapters without materializing their tensors. + +The official LTX adapters are raw safetensors files rather than PEFT +checkpoints. This module is deliberately LTX-specific: it translates the +official names into the locally constructed Diffusers transformer names and +describes the resulting adapter targets. Applying those targets is the job +of :mod:`ltx2_phase_adapter`. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.nn as nn +from safetensors import safe_open + +_IGNORED_OFFICIAL_LORA_MODULES = {"text_embedding_projection.aggregate_embed"} +_OFFICIAL_TO_DIFFUSERS_PREFIXES = ( + ("audio_prompt_adaln_single.", "audio_prompt_adaln."), + ("prompt_adaln_single.", "prompt_adaln."), + ("audio_adaln_single.", "audio_time_embed."), + ("adaln_single.", "time_embed."), + ("audio_patchify_proj", "audio_proj_in"), + ("patchify_proj", "proj_in"), + ("av_ca_video_scale_shift_adaln_single.", "av_cross_attn_video_scale_shift."), + ("av_ca_audio_scale_shift_adaln_single.", "av_cross_attn_audio_scale_shift."), + ("av_ca_a2v_gate_adaln_single.", "av_cross_attn_video_a2v_gate."), + ("av_ca_v2a_gate_adaln_single.", "av_cross_attn_audio_v2a_gate."), +) + + +@dataclass(frozen=True) +class AdapterTarget: + """One A/B pair mapped onto a local transformer linear layer.""" + + # Keep the original Diffusers module name for diagnostics after LTX packs + # Q/K/V into the runtime module path. + source_module: str + module: str + slice: str | None + a_key: str + b_key: str + a_shape: tuple[int, ...] + b_shape: tuple[int, ...] + + +@dataclass(frozen=True) +class AdapterManifest: + """Tensor-free description of the fixed refinement adapter.""" + + path: str + targets: tuple[AdapterTarget, ...] + + +def _to_diffusers_module_name(name: str) -> str: + for official_prefix, diffusers_prefix in _OFFICIAL_TO_DIFFUSERS_PREFIXES: + if name.startswith(official_prefix): + return diffusers_prefix + name[len(official_prefix) :] + return name + + +def _resolve_lora_target(transformer: nn.Module, module_name: str) -> tuple[str, str | int | None] | None: + modules = dict(transformer.named_modules()) + if module_name in modules: + return module_name, None + for packed_suffix, sub_suffix, shard_id in getattr(transformer, "stacked_params_mapping", ()): + if sub_suffix in module_name: + packed_name = module_name.replace(sub_suffix, packed_suffix) + if packed_name in modules: + return packed_name, shard_id + return None + + +def parse_ltx_adapter(transformer: nn.Module, path: str) -> AdapterManifest: + """Translate one official raw LTX adapter into local Transformer targets.""" + path = str(Path(path)) + targets: list[AdapterTarget] = [] + unresolved: list[str] = [] + seen: set[tuple[str, str | int | None]] = set() + + with safe_open(path, framework="pt", device="cpu") as handle: + tensor_names = set(handle.keys()) + official_name_set = { + tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0] + for tensor_name in tensor_names + if tensor_name.endswith(".lora_A.weight") + } + official_b_names = { + tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0] + for tensor_name in tensor_names + if tensor_name.endswith(".lora_B.weight") + } + missing_a = sorted(official_b_names - official_name_set - _IGNORED_OFFICIAL_LORA_MODULES) + if missing_a: + raise ValueError(f"Missing paired LoRA A tensors for {missing_a} in {path}.") + + official_names = sorted(official_name_set) + for official_name in official_names: + if official_name in _IGNORED_OFFICIAL_LORA_MODULES: + continue + key_prefix = f"diffusion_model.{official_name}" + key_a = f"{key_prefix}.lora_A.weight" + key_b = f"{key_prefix}.lora_B.weight" + if key_b not in tensor_names: + raise ValueError(f"Missing paired LoRA tensor {key_b!r} in {path}.") + + a_shape = tuple(handle.get_slice(key_a).get_shape()) + b_shape = tuple(handle.get_slice(key_b).get_shape()) + if len(a_shape) != 2 or len(b_shape) != 2 or b_shape[1] != a_shape[0]: + raise ValueError(f"Invalid LoRA pair for {official_name}: A={a_shape}, B={b_shape}.") + + source_module = _to_diffusers_module_name(official_name) + resolved = _resolve_lora_target(transformer, source_module) + if resolved is None: + unresolved.append(official_name) + continue + module, shard_id = resolved + target_id = (module, shard_id) + if target_id in seen: + raise ValueError(f"Official LTX LoRA contains duplicate target {module!r}, slice {shard_id!r}.") + seen.add(target_id) + targets.append( + AdapterTarget( + source_module=source_module, + module=module, + slice=shard_id, + a_key=key_a, + b_key=key_b, + a_shape=a_shape, + b_shape=b_shape, + ) + ) + + if unresolved: + raise ValueError(f"Official LTX LoRA contains unmapped modules: {unresolved}.") + if not targets: + raise ValueError(f"No Transformer LoRA tensors were found in {path}.") + return AdapterManifest(path=path, targets=tuple(targets)) + + +def iter_adapter_tensors(manifest: AdapterManifest) -> Iterator[tuple[AdapterTarget, torch.Tensor, torch.Tensor]]: + """Stream adapter pairs in manifest order without a full checkpoint copy.""" + + with safe_open(manifest.path, framework="pt", device="cpu") as handle: + for target in manifest.targets: + yield ( + target, + handle.get_tensor(target.a_key), + handle.get_tensor(target.b_key), + ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 54c866eb7..7f2f31a60 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -9,7 +9,8 @@ import inspect import json import logging import os -from dataclasses import dataclass +from dataclasses import dataclass, replace +from pathlib import Path from typing import TYPE_CHECKING, Any import torch @@ -19,6 +20,8 @@ from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from diffusers.video_processor import VideoProcessor from huggingface_hub import hf_hub_download +from safetensors import safe_open +from safetensors.torch import load_file from transformers import AutoTokenizer, Gemma3ForConditionalGeneration from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata @@ -32,6 +35,7 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint from .ltx2_transformer import ( LTX2VideoTransformer3DModel, apply_interleaved_rotary_emb, @@ -70,6 +74,9 @@ class LTXComponentProfile: video_vae_cls: type = AutoencoderKLLTX2Video vocoder_cls: type = LTX2Vocoder vocoder_fallback_cls: type | None = None + artifact_repo_id: str | None = None + latent_upsampler_filename: str | None = None + distilled_lora_filename: str | None = None LTX2_COMPONENT_PROFILE = LTXComponentProfile( @@ -101,16 +108,96 @@ LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( video_vae_cls=DistributedAutoencoderKLLTX2Video, ) +LTX23_DISTILLED_COMPONENT_PROFILE = replace( + LTX23_COMPONENT_PROFILE, + name="ltx2_3_distilled", + resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2.3", + latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", +) + +LTX2_TWO_STAGE_COMPONENT_PROFILE = replace( + LTX2_COMPONENT_PROFILE, + name="ltx2_two_stage", + resident_modules=(*LTX2_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2", + latent_upsampler_filename="ltx-2-spatial-upscaler-x2-1.0.safetensors", + distilled_lora_filename="ltx-2-19b-distilled-lora-384.safetensors", +) + +LTX23_TWO_STAGE_COMPONENT_PROFILE = replace( + LTX23_COMPONENT_PROFILE, + name="ltx2_3_two_stage", + resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2.3", + latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", + distilled_lora_filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors", +) + _COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = { ("one_stage", "2"): LTX2_COMPONENT_PROFILE, ("one_stage", "2.3"): LTX23_COMPONENT_PROFILE, + ("two_stage", "2"): LTX2_TWO_STAGE_COMPONENT_PROFILE, + ("two_stage", "2.3"): LTX23_TWO_STAGE_COMPONENT_PROFILE, ("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE, + ("distilled_two_stage", "2.3"): LTX23_DISTILLED_COMPONENT_PROFILE, ("dmd2", "2"): LTX2_COMPONENT_PROFILE, ("dmd2", "2.3"): LTX23_COMPONENT_PROFILE, } +def resolve_ltx_artifact( + model: str, + repo_id: str, + filename: str, +) -> str: + """Resolve an official LTX sidecar from the model root or its Hub repository.""" + candidate = Path(model) / filename + if candidate.is_file(): + return str(candidate) + + try: + return hf_hub_download(repo_id=repo_id, filename=filename) + except Exception as exc: + raise FileNotFoundError( + f"Unable to resolve LTX artifact {filename!r}. Searched {candidate}; " + f"place the file in the model root or make {repo_id} available." + ) from exc + + +def _load_ltx_latent_upsampler_single_file(path: str, dtype: torch.dtype) -> LTX2LatentUpsamplerModel: + """Load an official single-file upsampler into the Diffusers module.""" + with safe_open(path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + raw_config = json.loads(metadata.get("config", "{}")) + config = { + "in_channels": raw_config.get("in_channels", 128), + "mid_channels": raw_config.get("mid_channels", 1024), + "num_blocks_per_stage": raw_config.get("num_blocks_per_stage", 4), + "dims": raw_config.get("dims", 3), + "spatial_upsample": raw_config.get("spatial_upsample", True), + "temporal_upsample": raw_config.get("temporal_upsample", False), + "rational_spatial_scale": raw_config.get("spatial_scale", 2.0), + "use_rational_resampler": raw_config.get("rational_resampler", False), + } + with torch.device("cpu"): + upsampler = LTX2LatentUpsamplerModel(**config).to(dtype=dtype) + + state_dict = load_file(path, device="cpu") + if "upsampler.0.weight" in state_dict and hasattr(upsampler.upsampler, "conv"): + state_dict["upsampler.conv.weight"] = state_dict.pop("upsampler.0.weight") + state_dict["upsampler.conv.bias"] = state_dict.pop("upsampler.0.bias") + missing, unexpected = upsampler.load_state_dict(state_dict, strict=False) + unresolved_missing = set(missing) - {"upsampler.blur_down.kernel"} + if unresolved_missing or unexpected: + raise ValueError( + f"Invalid LTX latent upsampler {path}: missing={sorted(unresolved_missing)}, " + f"unexpected={sorted(unexpected)}." + ) + return upsampler + + def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile: """Resolve component construction independently from execution recipes.""" try: @@ -392,16 +479,34 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: ) if "latent_upsampler" in profile.resident_modules: - # BlurDownsample constructs an integer kernel that must be initialized - # on CPU; component placement is handled uniformly after construction. - with torch.device("cpu"): - pipeline.latent_upsampler = _load_component( - LTX2LatentUpsamplerModel, + upsampler_config = os.path.join(model, "latent_upsampler", "config.json") + if os.path.isfile(upsampler_config) or not local_files_only: + try: + pipeline.latent_upsampler = _load_component( + LTX2LatentUpsamplerModel, + model, + "latent_upsampler", + local_files_only=local_files_only, + dtype=dtype, + ) + except (OSError, ValueError): + if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None: + raise + upsampler_path = resolve_ltx_artifact( + model, + profile.artifact_repo_id, + profile.latent_upsampler_filename, + ) + pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype) + else: + if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None: + raise FileNotFoundError(f"LTX latent upsampler component not found under {model}.") + upsampler_path = resolve_ltx_artifact( model, - "latent_upsampler", - local_files_only=local_files_only, - dtype=dtype, + profile.artifact_repo_id, + profile.latent_upsampler_filename, ) + pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype) transformer_config = load_transformer_config(model, "transformer", local_files_only) quant_config = getattr(od_config, "quantization_config", None) @@ -412,6 +517,16 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: subfolder="scheduler", local_files_only=local_files_only, ) + expected_checkpoint_kind: LTXCheckpointKind | None = None + if pipeline.pipeline_kind == "two_stage": + expected_checkpoint_kind = "regular" + elif pipeline.pipeline_kind == "distilled_two_stage": + expected_checkpoint_kind = "distilled" + validate_ltx_checkpoint( + pipeline.scheduler.config, + expected_kind=expected_checkpoint_kind, + pipeline_name=type(pipeline).__name__, + ) pipeline.vae_spatial_compression_ratio = pipeline.vae.spatial_compression_ratio pipeline.vae_temporal_compression_ratio = pipeline.vae.temporal_compression_ratio diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index 4a149be85..dd55f3953 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -62,13 +62,14 @@ class LTXDenoiseContext: audio_latents: torch.Tensor video_coords: torch.Tensor audio_coords: torch.Tensor + audio_attention_mask: torch.Tensor | None = None conditioning_mask: torch.Tensor | None = None conditioning_mask_for_model: torch.Tensor | None = None @dataclass class LTXPhaseResult: - """Denoised, unpacked AV latents and the context used to produce them.""" + """Denoised AV latents and the context used to produce them.""" forward_context: LTXForwardContext video: torch.Tensor @@ -284,7 +285,8 @@ def prepare_rope_coords_stage( latents: torch.Tensor, audio_latents: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - video_coords = pipeline.transformer.rope.prepare_video_coords( + transformer = pipeline.transformer + video_coords = transformer.rope.prepare_video_coords( latents.shape[0], forward_ctx.latent_num_frames, forward_ctx.latent_height, @@ -292,7 +294,7 @@ def prepare_rope_coords_stage( latents.device, fps=forward_ctx.request_inputs.frame_rate, ) - audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords( + audio_coords = transformer.audio_rope.prepare_audio_coords( audio_latents.shape[0], forward_ctx.padded_audio_num_frames, audio_latents.device, @@ -331,6 +333,7 @@ def build_transformer_kwargs( # with a learned register, making all output context tokens valid. "encoder_attention_mask": None, "audio_encoder_attention_mask": None, + "audio_attention_mask": denoise_ctx.audio_attention_mask, "num_frames": forward_ctx.latent_num_frames, "height": forward_ctx.latent_height, "width": forward_ctx.latent_width, @@ -442,11 +445,25 @@ class LTXPhaseExecutor: video_audio_step_adapter=video_audio_step_adapter, ) video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents) + ring_degree = getattr(pipeline.od_config.parallel_config, "ring_degree", 1) or 1 + if padded_audio_num_frames > original_audio_num_frames and ring_degree > 1: + raise ValueError( + "LTX audio padding requires an attention mask, which Ring sequence parallelism does not support. " + "Use Ulysses-only SP or choose a request whose audio latent length is divisible by the SP size." + ) denoise_ctx = LTXDenoiseContext( latents=latents, audio_latents=audio_latents, video_coords=video_coords, audio_coords=audio_coords, + audio_attention_mask=( + torch.arange(padded_audio_num_frames, device=audio_latents.device) + .lt(original_audio_num_frames) + .unsqueeze(0) + .expand(audio_latents.shape[0], -1) + if padded_audio_num_frames > original_audio_num_frames + else None + ), conditioning_mask=conditioning_mask, ) denoise_ctx = pipeline._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py index f35f5b998..bc84e7c32 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -126,7 +126,10 @@ def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch. def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: """Convert x0 back to velocity using the official fp32 arithmetic order.""" - sigma = sigma.to(torch.float32) + # Official LTX converts the scalar scheduler sigma to a host value before + # division. A CUDA scalar tensor selects different kernel arithmetic and + # can cross bf16 rounding boundaries late in the denoise trajectory. + sigma = sigma.to(torch.float32).item() return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype) @@ -148,6 +151,7 @@ def combine_guided_x0( uncond_perturbed: torch.Tensor | float, uncond_modality: torch.Tensor | float, guidance: LTXModalityGuidance, + rescale_token_count: int | None = None, ) -> torch.Tensor: dtype = cond.dtype cond = cond.float() @@ -163,9 +167,18 @@ def combine_guided_x0( # uses its full-tensor standard deviation. In Omni, dimension 0 may # instead contain independently scheduled requests. Reduce each item # separately so its result cannot depend on unrelated co-batched work. - reduce_dims = tuple(range(1, pred.ndim)) - cond_std = cond.std(dim=reduce_dims, keepdim=True) - pred_std = pred.std(dim=reduce_dims, keepdim=True) + stats_cond = cond + stats_pred = pred + if rescale_token_count is not None: + if cond.ndim < 2 or not 0 < rescale_token_count <= cond.shape[1]: + raise ValueError( + f"Guidance rescale token count must be in [1, {cond.shape[1]}], got {rescale_token_count}." + ) + stats_cond = cond[:, :rescale_token_count] + stats_pred = pred[:, :rescale_token_count] + reduce_dims = tuple(range(1, stats_pred.ndim)) + cond_std = stats_cond.std(dim=reduce_dims, keepdim=True) + pred_std = stats_pred.std(dim=reduce_dims, keepdim=True) eps = torch.finfo(pred.dtype).eps factor = cond_std / pred_std.clamp_min(eps) factor = torch.where(pred_std > eps, factor, torch.ones_like(factor)) @@ -175,6 +188,8 @@ def combine_guided_x0( def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor: + if repeats == 1: + return tensor return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1)) @@ -278,6 +293,11 @@ class LTXGuidanceExecutor: if model_pass_count > 1: denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, model_pass_count) denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, model_pass_count) + if denoise_ctx.audio_attention_mask is not None: + denoise_ctx.audio_attention_mask = _repeat_batch( + denoise_ctx.audio_attention_mask, + model_pass_count, + ) return denoise_ctx @staticmethod @@ -322,6 +342,7 @@ class LTXGuidanceExecutor: guidance: LTXModalityGuidance, *, model_sigma: torch.Tensor | None = None, + rescale_token_count: int | None = None, ) -> torch.Tensor: model_sigma = sigma if model_sigma is None else model_sigma # Official guidance reduces contiguous BSC tensors. The fused @@ -334,6 +355,7 @@ class LTXGuidanceExecutor: uncond_perturbed=x0.get("ptb", 0.0), uncond_modality=x0.get("mod", 0.0), guidance=guidance, + rescale_token_count=rescale_token_count, ) return velocity_from_x0(sample, guided, sigma) @@ -434,6 +456,7 @@ class LTXGuidanceExecutor: audio_splits, forward_ctx.audio_scheduler.sigmas[index], plan.spec.audio, + rescale_token_count=forward_ctx.original_audio_num_frames, ), ) @@ -509,6 +532,7 @@ class LTXGuidanceExecutor: audio_splits, forward_ctx.audio_scheduler.sigmas[index], plan.spec.audio, + rescale_token_count=forward_ctx.original_audio_num_frames, ), ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py index bb0155e79..11929d331 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -117,7 +117,9 @@ def create_noised_state( device=latents.device, dtype=latents.dtype, ) - return noise_scale * noise + (1 - noise_scale) * latents + if isinstance(noise_scale, torch.Tensor): + noise_scale = noise_scale.float() + return torch.lerp(latents.float(), noise.float(), noise_scale).to(latents.dtype) def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor: @@ -135,6 +137,16 @@ def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: return latents[:, :num_frames] +def clear_audio_padding(latents: torch.Tensor, num_frames: int) -> torch.Tensor: + """Keep Omni's SP-only audio padding outside the sampler state.""" + if not 0 < num_frames <= latents.shape[1]: + raise ValueError(f"Audio frame count must be in [1, {latents.shape[1]}], got {num_frames}.") + if num_frames == latents.shape[1]: + return latents + padding = latents.new_zeros(latents.shape[0], latents.shape[1] - num_frames, latents.shape[2]) + return torch.cat([latents[:, :num_frames], padding], dim=1) + + def get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int: if sp_size > 1: audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size @@ -205,7 +217,12 @@ def prepare_video_latents( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) - return randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + # Official LTX samples the 5D video latent first and then patchifies it, + # leaving token-major storage behind the [B, tokens, channels] view. The + # logical values are unchanged, but preserving that layout is important: + # CUDA selects a different bf16 GEMM path for a contiguous token tensor. + return latents.transpose(1, 2).contiguous().transpose(1, 2) def prepare_audio_latents( @@ -226,6 +243,16 @@ def prepare_audio_latents( sp_size = getattr(pipeline.od_config.parallel_config, "sequence_parallel_size", 1) or 1 padded_latent_length = get_sp_padded_audio_latent_length(original_latent_length, int(sp_size)) + def pad_logical_latents(logical_latents: torch.Tensor) -> torch.Tensor: + if padded_latent_length == original_latent_length: + return logical_latents + padding = logical_latents.new_zeros( + logical_latents.shape[0], + padded_latent_length - original_latent_length, + logical_latents.shape[2], + ) + return torch.cat([logical_latents, padding], dim=1) + if latents is not None: if latents.ndim == 4: latents = pack_audio_latents(latents) @@ -237,29 +264,24 @@ def prepare_audio_latents( pipeline.audio_vae.latents_mean, pipeline.audio_vae.latents_std, ) - latents = create_noised_state(latents, noise_scale, generator) if latents.shape[1] not in {original_latent_length, padded_latent_length}: raise ValueError( "Provided `audio_latents` has incompatible audio frame count " f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}." ) - if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length: - padding = torch.zeros( - latents.shape[0], - padded_latent_length - original_latent_length, - latents.shape[2], - dtype=latents.dtype, - device=latents.device, - ) - latents = torch.cat([latents, padding], dim=1) + # Padding is an Omni implementation detail, not part of the official + # audio state. Draw request RNG only for logical tokens so later phases + # remain seed-invariant across SP degrees, then append deterministic 0s. + latents = create_noised_state(latents[:, :original_latent_length], noise_scale, generator) + latents = pad_logical_latents(latents) return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length - shape = (batch_size, padded_latent_length, num_channels_latents * latent_mel_bins) + shape = (batch_size, original_latent_length, num_channels_latents * latent_mel_bins) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = pad_logical_latents(randn_tensor(shape, generator=generator, device=device, dtype=dtype)) return latents, original_latent_length, padded_latent_length diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py new file mode 100644 index 000000000..a205fce6b --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py @@ -0,0 +1,430 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Construction and execution of a fixed LTX refinement-phase adapter. + +Checkpoint key parsing remains in :mod:`ltx2_adapter_parser`; this module +resolves the declared artifact and makes the existing Transformer linears +adapter-capable while retaining their quantization and tensor-parallel paths. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +import vllm.envs as envs +from vllm.distributed import ( + split_tensor_along_last_dim, + tensor_model_parallel_all_reduce, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.batch_invariant import linear_batch_invariant +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.utils import dispatch_unquantized_gemm +from vllm.platforms import current_platform + +from .ltx2_adapter_parser import AdapterManifest, AdapterTarget, iter_adapter_tensors, parse_ltx_adapter +from .ltx2_components import resolve_ltx_artifact +from .ltx2_recipes import LTX_DISTILLED_ADAPTER_SLOT + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class _AdapterLayout: + """Rank-local storage and execution layout for one LoRA pair.""" + + kind: str + input_size: int + output_size: int + a_input_start: int + a_input_size: int + b_output_start: int + b_output_size: int + output_start: int + reduce_results: bool = False + + +def _module_device(module: nn.Module) -> torch.device: + for tensor in (*module.parameters(), *module.buffers()): + if tensor.device.type != "meta": + return tensor.device + raise RuntimeError(f"Cannot materialize an LTX phase adapter for meta module {module._get_name()}.") + + +def _build_layout(layer: nn.Module, target: AdapterTarget) -> _AdapterLayout: + if isinstance(layer, QKVParallelLinear): + if target.slice not in {"q", "k", "v"}: + raise ValueError( + f"LTX packed QKV target {target.module!r} requires a q/k/v adapter slice, got {target.slice!r}." + ) + q_local = layer.num_heads * layer.head_size + k_local = layer.num_kv_heads * layer.head_size + v_local = layer.num_kv_heads * layer.v_head_size + if target.slice == "q": + output_size = layer.total_num_heads * layer.head_size + b_output_start = layer.tp_rank * q_local + b_output_size = q_local + output_start = 0 + elif target.slice == "k": + output_size = layer.total_num_kv_heads * layer.head_size + shard_rank = layer.tp_rank // layer.num_kv_head_replicas + b_output_start = shard_rank * k_local + b_output_size = k_local + output_start = q_local + else: + output_size = layer.total_num_kv_heads * layer.v_head_size + shard_rank = layer.tp_rank // layer.num_kv_head_replicas + b_output_start = shard_rank * v_local + b_output_size = v_local + output_start = q_local + k_local + return _AdapterLayout( + kind="qkv", + input_size=layer.input_size, + output_size=output_size, + a_input_start=0, + a_input_size=layer.input_size, + b_output_start=b_output_start, + b_output_size=b_output_size, + output_start=output_start, + ) + + if isinstance(layer, ColumnParallelLinear): + if layer.gather_output: + raise ValueError( + f"LTX phase adapter does not support gathered ColumnParallelLinear target {target.module!r}." + ) + return _AdapterLayout( + kind="column", + input_size=layer.input_size, + output_size=layer.output_size, + a_input_start=0, + a_input_size=layer.input_size, + b_output_start=layer.tp_rank * layer.output_size_per_partition, + b_output_size=layer.output_size_per_partition, + output_start=0, + ) + + if isinstance(layer, RowParallelLinear): + return _AdapterLayout( + kind="row", + input_size=layer.input_size, + output_size=layer.output_size, + a_input_start=layer.tp_rank * layer.input_size_per_partition, + a_input_size=layer.input_size_per_partition, + b_output_start=0, + b_output_size=layer.output_size, + output_start=0, + reduce_results=layer.reduce_results, + ) + + if isinstance(layer, nn.Linear): + return _AdapterLayout( + kind="replicated", + input_size=layer.in_features, + output_size=layer.out_features, + a_input_start=0, + a_input_size=layer.in_features, + b_output_start=0, + b_output_size=layer.out_features, + output_start=0, + ) + + raise TypeError(f"LTX phase adapter target {target.module!r} has unsupported layer type {type(layer).__name__}.") + + +class _AdapterPiece(nn.Module): + """One rank-local A/B pair stored as non-persistent buffers.""" + + def __init__(self, target: AdapterTarget, layout: _AdapterLayout) -> None: + super().__init__() + self.target = target + self.layout = layout + # The fixed graph is installed before base checkpoint loading. The + # empty buffers are replaced only after the base layer's quant/TP layout + # has settled and before offload/compile snapshots are taken. + self.register_buffer("lora_a", torch.empty(0), persistent=False) + self.register_buffer("lora_b", torch.empty(0), persistent=False) + + def load(self, lora_a: torch.Tensor, lora_b: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> None: + if tuple(lora_a.shape) != self.target.a_shape or tuple(lora_b.shape) != self.target.b_shape: + raise ValueError( + f"LTX adapter source shape changed for {self.target.source_module!r}: " + f"A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}." + ) + if lora_a.ndim != 2 or lora_b.ndim != 2: + raise ValueError(f"LTX adapter target {self.target.source_module!r} must use rank-2 A/B tensors.") + rank = self.target.a_shape[0] + if lora_a.shape[1] != self.layout.input_size or lora_b.shape[0] != self.layout.output_size: + raise ValueError( + f"LTX adapter shape does not match {self.target.module!r}: " + f"expected A=({rank}, {self.layout.input_size}), " + f"B=({self.layout.output_size}, {rank}); " + f"got A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}." + ) + + local_a = lora_a.narrow(1, self.layout.a_input_start, self.layout.a_input_size) + local_b = lora_b.narrow(0, self.layout.b_output_start, self.layout.b_output_size) + self.lora_a = local_a.to(device=device, dtype=dtype) + self.lora_b = local_b.to(device=device, dtype=dtype) + + def delta(self, input_: torch.Tensor) -> torch.Tensor: + if self.lora_a.numel() == 0 or self.lora_b.numel() == 0: + raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.") + if self.layout.kind == "row" and input_.shape[-1] != self.layout.a_input_size: + input_ = input_.narrow(-1, self.layout.a_input_start, self.layout.a_input_size) + input_ = input_.to(dtype=self.lora_a.dtype) + delta = F.linear(F.linear(input_, self.lora_a), self.lora_b) + if self.layout.kind == "row" and self.layout.reduce_results: + delta = tensor_model_parallel_all_reduce(delta) + return delta + + def weight_delta(self) -> torch.Tensor: + """Materialize the rank-local BF16 weight delta used by official fusion.""" + if self.lora_a.numel() == 0 or self.lora_b.numel() == 0: + raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.") + return torch.matmul(self.lora_b, self.lora_a) + + +def _apply_unquantized_gemm_with_weight( + layer: nn.Module, + input_: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike(): + return linear_batch_invariant(input_, weight, bias) + return dispatch_unquantized_gemm()(layer, input_, weight, bias) + + +def _forward_with_weight( + layer: nn.Module, + input_: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: + """Run an unquantized base layer with an ephemeral fused weight.""" + if isinstance(layer, nn.Linear): + return F.linear(input_, weight, layer.bias) + + quant_method = getattr(layer, "quant_method", None) + if not isinstance(quant_method, UnquantizedLinearMethod): + raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.") + + if isinstance(layer, RowParallelLinear): + if layer.input_is_parallel: + input_parallel = input_ + else: + input_parallel = split_tensor_along_last_dim(input_, num_partitions=layer.tp_size)[ + layer.tp_rank + ].contiguous() + bias = None if layer.tp_rank > 0 or layer.skip_bias_add else layer.bias + output_parallel = _apply_unquantized_gemm_with_weight(layer, input_parallel, weight, bias) + if layer.reduce_results and layer.tp_size > 1: + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + elif isinstance(layer, ColumnParallelLinear): + bias = layer.bias if not layer.skip_bias_add else None + output = _apply_unquantized_gemm_with_weight(layer, input_, weight, bias) + else: + raise TypeError(f"Unsupported LTX layer-fused linear type {type(layer).__name__}.") + + if not layer.return_bias: + return output + return output, layer.bias if layer.skip_bias_add else None + + +class _PhaseAdapterLinear(nn.Module): + """Delegate the base linear unchanged and add an enabled adapter slot.""" + + def __init__( + self, + base_layer: nn.Module, + targets: Iterable[AdapterTarget], + *, + adapter_dtype: torch.dtype, + layer_fused: bool = False, + ) -> None: + super().__init__() + self.base_layer = base_layer + self.adapter_dtype = adapter_dtype + if layer_fused and not ( + isinstance(base_layer, nn.Linear) + or isinstance(getattr(base_layer, "quant_method", None), UnquantizedLinearMethod) + ): + raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.") + self.layer_fused = layer_fused + pieces = [_AdapterPiece(target, _build_layout(base_layer, target)) for target in targets] + self.pieces = nn.ModuleList(pieces) + self._pieces_by_target = {piece.target: piece for piece in pieces} + self._enabled = False + + def load_target(self, target: AdapterTarget, lora_a: torch.Tensor, lora_b: torch.Tensor) -> None: + try: + piece = self._pieces_by_target[target] + except KeyError as exc: + raise ValueError(f"Adapter target {target.module!r} is not installed on this layer.") from exc + piece.load(lora_a, lora_b, device=_module_device(self.base_layer), dtype=self.adapter_dtype) + + def set_enabled(self, enabled: bool) -> None: + self._enabled = enabled + + def _add_adapter_delta(self, input_: torch.Tensor, output: torch.Tensor) -> torch.Tensor: + if not self._enabled: + return output + if output.requires_grad: + output = output.clone() + for piece in self.pieces: + delta = piece.delta(input_).to(dtype=output.dtype) + start = piece.layout.output_start + output[..., start : start + piece.layout.b_output_size].add_(delta) + return output + + @torch.no_grad() + def _fused_weight(self) -> torch.Tensor: + base_weight = self.base_layer.weight + fused_weight = base_weight.clone() + for piece in self.pieces: + start = piece.layout.output_start + base_slice = base_weight.narrow(0, start, piece.layout.b_output_size) + delta = piece.weight_delta() + if delta.shape != base_slice.shape: + raise ValueError( + f"LTX layer-fused delta shape {tuple(delta.shape)} does not match base weight slice " + f"{tuple(base_slice.shape)} for {piece.target.module!r}." + ) + # Match official fusion exactly: materialize B@A in the adapter + # dtype, then add the base tensor in-place before the main GEMM. + delta.add_(base_slice) + fused_weight.narrow(0, start, piece.layout.b_output_size).copy_(delta) + return fused_weight + + def forward(self, input_: torch.Tensor): + if not self._enabled: + return self.base_layer(input_) + if self.layer_fused: + return _forward_with_weight(self.base_layer, input_, self._fused_weight()) + output = self.base_layer(input_) + if isinstance(output, tuple): + return (self._add_adapter_delta(input_, output[0]), *output[1:]) + return self._add_adapter_delta(input_, output) + + +class LTXPhaseAdapterRuntime: + """Install and operate one fixed phase adapter on an existing transformer.""" + + def __init__( + self, + transformer: nn.Module, + manifest: AdapterManifest, + *, + dtype: torch.dtype, + layer_fused: bool = False, + ) -> None: + self.transformer = transformer + self.manifest = manifest + self._wrappers: dict[str, _PhaseAdapterLinear] = {} + self._target_names: tuple[str, ...] = () + self._materialized = False + grouped: dict[str, list[AdapterTarget]] = defaultdict(list) + for target in manifest.targets: + grouped[target.module].append(target) + + original_layers = {name: self.transformer.get_submodule(name) for name in grouped} + for name, targets in grouped.items(): + wrapper = _PhaseAdapterLinear( + original_layers[name], + targets, + adapter_dtype=dtype, + layer_fused=layer_fused, + ) + parent_name, _, child_name = name.rpartition(".") + parent = self.transformer.get_submodule(parent_name) if parent_name else self.transformer + if child_name not in parent._modules: + raise ValueError(f"LTX phase adapter target {name!r} is no longer a direct module child.") + parent._modules[child_name] = wrapper + self._wrappers[name] = wrapper + + self._target_names = tuple(sorted(self._wrappers, key=len, reverse=True)) + # LTX's custom loader receives the original checkpoint names. Keep the + # adaptation local to that loader rather than changing AutoWeightsLoader. + self.transformer._phase_adapter_parameter_name = self.base_parameter_name + logger.info( + "Installed %d LTX %s phase-adapter linears", + len(self._wrappers), + "layer-fused" if layer_fused else "dynamic", + ) + + def base_parameter_name(self, name: str) -> str: + for target_name in self._target_names: + prefix = f"{target_name}." + if name.startswith(prefix): + return f"{target_name}.base_layer.{name[len(prefix) :]}" + return name + + def finalize(self) -> None: + """Stream rank-local A/B buffers after base weight processing completes.""" + if self._materialized: + return + for target, lora_a, lora_b in iter_adapter_tensors(self.manifest): + self._wrappers[target.module].load_target(target, lora_a, lora_b) + self._materialized = True + logger.info("Materialized %d LTX phase-adapter tensor pairs", len(self.manifest.targets)) + + def activate(self, adapter_slot: str | None) -> None: + if adapter_slot is not None: + if adapter_slot != LTX_DISTILLED_ADAPTER_SLOT: + raise ValueError(f"Unknown LTX phase adapter slot {adapter_slot!r}.") + if not self._materialized: + raise RuntimeError("LTX phase adapter data must be materialized before activation.") + for wrapper in self._wrappers.values(): + wrapper.set_enabled(adapter_slot is not None) + + +def build_ltx_phase_adapter(pipeline: Any) -> LTXPhaseAdapterRuntime | None: + """Build the fixed adapter required by a multi-phase LTX recipe.""" + adapter_slots = {phase.adapter_slot for phase in pipeline.pipeline_recipe.phases if phase.adapter_slot is not None} + if not adapter_slots: + return None + if adapter_slots != {LTX_DISTILLED_ADAPTER_SLOT}: + raise ValueError(f"Unsupported LTX phase adapter slots: {sorted(adapter_slots)}.") + + profile = pipeline.component_profile + if profile.distilled_lora_filename is None or profile.artifact_repo_id is None: + raise ValueError(f"{profile.name} does not declare the required distilled adapter artifact.") + if getattr(pipeline.od_config, "lora_path", None) is not None: + raise ValueError( + f"{pipeline.__class__.__name__} reserves LoRA execution for its phase adapter; " + "request or static LoRA composition is not supported yet." + ) + + quantization_config = getattr(pipeline.od_config, "quantization_config", None) + layer_fused = quantization_config is None + adapter_path = resolve_ltx_artifact( + pipeline.od_config.model, + profile.artifact_repo_id, + profile.distilled_lora_filename, + ) + + if layer_fused and pipeline.od_config.dtype is not torch.bfloat16: + raise ValueError("LTX layer-fused phase LoRA requires a bfloat16 pipeline dtype.") + + manifest = parse_ltx_adapter(pipeline.transformer, adapter_path) + runtime = LTXPhaseAdapterRuntime( + pipeline.transformer, + manifest, + dtype=pipeline.od_config.dtype, + layer_fused=layer_fused, + ) + return runtime diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py index a70f349ac..bb42bc0e9 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py @@ -24,6 +24,7 @@ LTX_DEFAULT_NEGATIVE_PROMPT = ( LTX_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0) LTX_STAGE_2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0) +LTX_DISTILLED_ADAPTER_SLOT = "ltx_distilled" @dataclass(frozen=True) @@ -36,6 +37,7 @@ class LTXPhaseRecipe: sigmas: tuple[float, ...] | None = None noise_scale: float = 0.0 input_transform: Literal["initial", "spatial_upsample"] = "initial" + adapter_slot: str | None = None allow_guidance_override: bool = True use_official_sigma_schedule: bool = True @@ -157,11 +159,59 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe( fixed_num_inference_steps=True, ) +# LTX-2.3 full-distilled uses the same official fixed 8 + 3 sigma schedules +# and request contract. Its distinct component profile selects the 22B merged +# distilled Transformer, BWE vocoder, and matching spatial upsampler; no +# distilled LoRA is loaded for either phase. +LTX23_DISTILLED_TWO_STAGE_RECIPE = LTX2_DISTILLED_TWO_STAGE_RECIPE + + +def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipelineRecipe: + return LTXPipelineRecipe( + height=one_stage_recipe.height * 2, + width=one_stage_recipe.width * 2, + num_frames=one_stage_recipe.num_frames, + frame_rate=one_stage_recipe.frame_rate, + num_inference_steps=one_stage_recipe.num_inference_steps, + negative_prompt=one_stage_recipe.negative_prompt, + phases=( + LTXPhaseRecipe( + name="generate_lowres", + guidance=one_stage_recipe.request_guidance, + spatial_downscale=2, + noise_scale=1.0, + ), + LTXPhaseRecipe( + name="refine", + guidance=LTXGuidanceSpec.positive_only(), + sigmas=LTX_STAGE_2_DISTILLED_SIGMAS, + noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0], + input_transform="spatial_upsample", + adapter_slot=LTX_DISTILLED_ADAPTER_SLOT, + allow_guidance_override=False, + use_official_sigma_schedule=False, + ), + ), + video_output_phase=1, + # The official second stage refines video only and deliberately + # discards its audio result. Decode the full-context Stage-1 audio. + audio_output_phase=0, + allow_request_sigmas=False, + allow_request_latents=False, + ) + + +LTX2_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX2_ONE_STAGE_RECIPE) +LTX23_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX23_ONE_STAGE_RECIPE) + _PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = { ("one_stage", "2"): LTX2_ONE_STAGE_RECIPE, ("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE, + ("two_stage", "2"): LTX2_TWO_STAGE_RECIPE, + ("two_stage", "2.3"): LTX23_TWO_STAGE_RECIPE, ("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE, + ("distilled_two_stage", "2.3"): LTX23_DISTILLED_TWO_STAGE_RECIPE, ("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE, ("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE, } diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index e2eef19df..91d62943a 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -5,8 +5,9 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, replace -from typing import Any +from typing import Any, Literal import torch @@ -49,6 +50,34 @@ class LTXRequestInputs: return self.guidance.video.cfg_scale +LTXCheckpointKind = Literal["distilled", "regular"] + + +def validate_ltx_checkpoint( + scheduler_config: Mapping[str, Any], + *, + expected_kind: LTXCheckpointKind | None, + pipeline_name: str, +) -> None: + """Ensure scheduler metadata matches the checkpoint required by a profile.""" + if expected_kind is None: + return + + if expected_kind == "distilled": + if scheduler_config.get("use_dynamic_shifting", True) or scheduler_config.get("shift_terminal") is not None: + raise ValueError( + f"{pipeline_name} requires a merged distilled LTX2 checkpoint; " + "the loaded scheduler metadata describes a non-distilled checkpoint." + ) + return + + if not scheduler_config.get("use_dynamic_shifting", False) or scheduler_config.get("shift_terminal") is None: + raise ValueError( + f"{pipeline_name} requires a regular non-distilled LTX checkpoint; " + "the loaded scheduler metadata describes a distilled checkpoint." + ) + + def validate_pipeline_request( request_inputs: LTXRequestInputs, *, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 4b74ebcb8..cb4130259 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -47,6 +47,7 @@ from .ltx2_guidance import ( LTXGuidanceExecutor, LTXGuidancePlan, ) +from .ltx2_phase_adapter import LTXPhaseAdapterRuntime, build_ltx_phase_adapter from .ltx2_recipes import ( LTXPhaseRecipe, LTXPipelineRecipe, @@ -137,9 +138,21 @@ class LTXRuntime( def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: del prefix + parallel_config = getattr(od_config, "parallel_config", None) + if getattr(parallel_config, "ulysses_mode", "strict") == "advanced_uaa": + raise ValueError( + f"{self.__class__.__name__} does not support ulysses_mode='advanced_uaa'. " + "Use the default ulysses_mode='strict' for LTX sequence parallelism." + ) self.model_version = detect_ltx_model_version(od_config.model) self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version) self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version) + if getattr(od_config, "cache_backend", "none") == "cache_dit" and len(self.pipeline_recipe.phases) > 1: + raise ValueError( + f"{self.__class__.__name__} does not support cache_backend='cache_dit'. " + "Cache-DiT currently supports only one-stage LTX recipes; multi-stage recipes require " + "phase-aware cache enable and refresh." + ) self._dit_modules = list(self.component_profile.dit_modules) self._encoder_modules = list(self.component_profile.encoder_modules) self._vae_modules = list(self.component_profile.vae_modules) @@ -150,6 +163,7 @@ class LTXRuntime( super().__init__() self._guidance_plan = LTXGuidancePlan.build(self.pipeline_recipe.request_guidance) initialize_pipeline_components(self, od_config) + self._phase_adapter: LTXPhaseAdapterRuntime | None = build_ltx_phase_adapter(self) self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) @@ -205,6 +219,8 @@ class LTXRuntime( max_sequence_length=max_sequence_length, ) image = self._resolve_request_image(req, image, request_inputs) + if image is not None and not self.support_image_input: + raise ValueError(f"{self.__class__.__name__} does not support `image` input.") request_sigmas = self._resolve_request_sigmas(req, sigmas) validate_pipeline_request( request_inputs, @@ -214,6 +230,13 @@ class LTXRuntime( pipeline_name=self.__class__.__name__, request_sigmas=request_sigmas, ) + phase_adapter = getattr(self, "_phase_adapter", None) + if phase_adapter is not None and any( + getattr(sampling, "lora_request", None) is not None for sampling in req.sampling_params_list + ): + raise ValueError( + f"{self.__class__.__name__} cannot compose a request LoRA with its internal phase adapter." + ) return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image) def _run_recipe( @@ -258,8 +281,20 @@ class LTXRuntime( return self.decode_phase(output_phase) def _enter_phase(self, phase: LTXPhaseRecipe) -> None: - """Hook for a future phase-weight strategy.""" self._active_phase_name = phase.name + phase_adapter = getattr(self, "_phase_adapter", None) + if phase_adapter is None: + if phase.adapter_slot is not None: + raise RuntimeError(f"LTX phase {phase.name!r} requires adapter slot {phase.adapter_slot!r}.") + return + phase_adapter.activate(phase.adapter_slot) + + def eval(self): + result = super().eval() + phase_adapter = getattr(self, "_phase_adapter", None) + if phase_adapter is not None: + phase_adapter.finalize() + return result def prepare_latents( self, @@ -692,6 +727,7 @@ class LTXRuntime( noise_pred_audio, timestep, ) + audio = latent_ops.clear_audio_padding(audio, forward_ctx.original_audio_num_frames) return latent_ops.LTXAVState(video=video, audio=audio) def _unpack_and_denormalize_stage( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 95724bd19..a59655cd3 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -355,6 +355,19 @@ def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor: return attention_mask +class _LTX2ParallelAttention(Attention): + """Preserve LTX SP collectives while applying model-specific padding masks.""" + + def _run_local_attention(self, query, key, value, attn_metadata): + attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None + if attention_mask is not None: + # LTX masks describe K/V padding. Keep this semantic local to LTX: + # Flash varlen applies a 2D mask to both Q and K, while SP uses a + # broadcast 4D key mask that only SDPA consumes correctly. + return self.sdpa_fallback.forward(query, key, value, attn_metadata) + return super()._run_local_attention(query, key, value, attn_metadata) + + class LTX2AudioVideoAttnProcessor: r""" Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model. @@ -396,11 +409,10 @@ class LTX2AudioVideoAttnProcessor: return None if self._is_sp_enabled(): - # In SP, Ulysses expects a 2D padding mask that matches query length. - # For cross-attention, encoder sequence length != query length, so drop the mask. - if encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]: - return None - return to_ltx_padding_mask(attention_mask) + # Keep the full key mask replicated while strict Ulysses reshards + # Q/K/V. The broadcast Q dimension also works when cross-attention + # has different Q and K lengths, without changing generic Ulysses. + return to_ltx_padding_mask(attention_mask)[:, None, None, :] attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) @@ -696,7 +708,7 @@ class LTX2Attention(torch.nn.Module): torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(), ] ) - self.attn = Attention( + self.attn = _LTX2ParallelAttention( num_heads=self.query_num_heads, head_size=dim_head, num_kv_heads=self.kv_num_heads, @@ -709,7 +721,15 @@ class LTX2Attention(torch.nn.Module): # LTX-2.3: per-head gated attention # leave unquantized for this linear if apply_gated_attention: - self.to_gate_logits = nn.Linear(query_dim, self.query_num_heads, bias=True) + self.to_gate_logits = ColumnParallelLinear( + query_dim, + heads, + bias=True, + gather_output=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.to_gate_logits" if prefix else "to_gate_logits", + ) else: self.to_gate_logits = None @@ -1498,6 +1518,14 @@ class LTX2VideoTransformer3DModel(nn.Module): _layerwise_offload_blocks_attrs = ["transformer_blocks"] _hsdp_shard_conditions = [is_transformer_block_module] _sp_plan: dict[str, Any] | None = None + stacked_params_mapping = ( + (".attn1.to_qkv", ".attn1.to_q", "q"), + (".attn1.to_qkv", ".attn1.to_k", "k"), + (".attn1.to_qkv", ".attn1.to_v", "v"), + (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"), + (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"), + (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"), + ) packed_modules_mapping = { "to_qkv": ["to_q", "to_k", "to_v"], } @@ -1828,6 +1856,7 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_sigma: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, + audio_attention_mask: torch.Tensor | None = None, num_frames: int | None = None, height: int | None = None, width: int | None = None, @@ -1861,6 +1890,8 @@ class LTX2VideoTransformer3DModel(nn.Module): Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`. audio_encoder_attention_mask (`torch.Tensor`, *optional*): Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling. + audio_attention_mask (`torch.Tensor`, *optional*): + Optional audio-token key padding mask shared by audio self-attention and audio-to-video attention. num_frames (`int`, *optional*): The number of latent video frames. Used if calculating the video coordinates for RoPE. height (`int`, *optional*): @@ -2044,6 +2075,8 @@ class LTX2VideoTransformer3DModel(nn.Module): "ca_audio_rotary_emb": audio_cross_attn_rotary_emb, "encoder_attention_mask": encoder_attention_mask, "audio_encoder_attention_mask": audio_encoder_attention_mask, + "audio_self_attention_mask": audio_attention_mask, + "a2v_cross_attention_mask": audio_attention_mask, "video_self_attention_perturbation_mask": perturbation_mask_for("video_self_attention", block_idx), "audio_self_attention_perturbation_mask": perturbation_mask_for("audio_self_attention", block_idx), "a2v_cross_attention_perturbation_mask": perturbation_mask_for("a2v_cross_attention", block_idx), @@ -2097,16 +2130,8 @@ class LTX2VideoTransformer3DModel(nn.Module): Returns: Set of parameter names that were successfully loaded. """ - stacked_params_mapping = [ - (".attn1.to_qkv", ".attn1.to_q", "q"), - (".attn1.to_qkv", ".attn1.to_k", "k"), - (".attn1.to_qkv", ".attn1.to_v", "v"), - (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"), - (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"), - (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"), - ] - params_dict = dict(self.named_parameters()) + phase_adapter_parameter_name = getattr(self, "_phase_adapter_parameter_name", None) tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0 loaded_params: set[str] = set() @@ -2127,10 +2152,12 @@ class LTX2VideoTransformer3DModel(nn.Module): return weight for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, shard_id in self.stacked_params_mapping: if weight_name not in name: continue packed_name = name.replace(weight_name, param_name) + if callable(phase_adapter_parameter_name): + packed_name = phase_adapter_parameter_name(packed_name) if packed_name not in params_dict: continue name = packed_name @@ -2139,6 +2166,8 @@ class LTX2VideoTransformer3DModel(nn.Module): weight_loader(param, loaded_weight, shard_id) break else: + if callable(phase_adapter_parameter_name): + name = phase_adapter_parameter_name(name) if name not in params_dict: logger.warning( "Skipping transformer weight %s -- not found in model " diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py index 0e5bbc558..3cb6b84f1 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py @@ -1,27 +1,48 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Distilled two-stage entry points for the LTX model family.""" +"""Two-stage entry points for the LTX model family.""" -from vllm_omni.diffusion.data import OmniDiffusionConfig +from typing import ClassVar -from .ltx2_components import LTX2_DISTILLED_COMPONENT_PROFILE +from .ltx2_components import ( + LTX2_DISTILLED_COMPONENT_PROFILE, + LTX2_TWO_STAGE_COMPONENT_PROFILE, +) from .ltx2_components import ( get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401 ) -from .ltx2_recipes import LTX2_DISTILLED_TWO_STAGE_RECIPE -from .pipeline_ltx2 import LTX2Pipeline +from .ltx2_conditioning import LTXI2VConditioningMixin +from .ltx2_recipes import ( + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX2_TWO_STAGE_RECIPE, +) +from .ltx2_runtime import LTXRuntime -class LTX2DistilledPipeline(LTX2Pipeline): - """Unified LTX-2 distilled two-stage T2V/I2V entry.""" +class _LTX2TwoStageBase(LTXI2VConditioningMixin, LTXRuntime): + """Static discovery contract shared by the two-stage public entries.""" - pipeline_kind = "distilled_two_stage" - component_profile = LTX2_DISTILLED_COMPONENT_PROFILE - pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE + _dit_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.dit_modules) + _encoder_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.encoder_modules) + _vae_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.vae_modules) + _resident_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.resident_modules) supports_request_batch = False support_image_input = True unified_text_image_entry = True - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__(od_config=od_config, prefix=prefix) + +class LTX2TwoStagePipeline(_LTX2TwoStageBase): + """Regular checkpoint with low-resolution generation and LoRA refinement.""" + + pipeline_kind = "two_stage" + component_profile = LTX2_TWO_STAGE_COMPONENT_PROFILE + pipeline_recipe = LTX2_TWO_STAGE_RECIPE + + +class LTX2DistilledPipeline(_LTX2TwoStageBase): + """Unified LTX-2/LTX-2.3 full-distilled two-stage T2V/I2V entry.""" + + pipeline_kind = "distilled_two_stage" + component_profile = LTX2_DISTILLED_COMPONENT_PROFILE + pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 75ef7391a..0e814886b 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -71,6 +71,11 @@ _DIFFUSION_MODELS = { "pipeline_ltx2", "LTX2Pipeline", ), + "LTX2TwoStagePipeline": ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2TwoStagePipeline", + ), "LTX2DistilledPipeline": ( "ltx2", "pipeline_ltx2_two_stage", @@ -516,6 +521,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "WanPipeline": "get_wan22_post_process_func", "WanVACEPipeline": "get_wan22_vace_post_process_func", "LTX2Pipeline": "get_ltx2_post_process_func", + "LTX2TwoStagePipeline": "get_ltx2_post_process_func", "LTX2DistilledPipeline": "get_ltx2_post_process_func", "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func", "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 8d7ee1e29..bc923ea86 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -205,6 +205,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { } for model_class_name in ( "LTX2Pipeline", + "LTX2TwoStagePipeline", "LTX2DistilledPipeline", ) }, From dbc0dd6d85272aab507c78b35f837d8651807087 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Thu, 13 Aug 2026 10:00:50 -0700 Subject: [PATCH 28/41] [Bugfix][Higgs-Audio-V3] Disable XQA decode (#6068) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Sysy <1370724210@qq.com> --- tests/config/test_omni_config.py | 1 + tests/engine/test_stage_engine_args.py | 2 + .../higgs_audio_v3/test_higgs_audio_v3.py | 8 ++- vllm_omni/config/omni_config.py | 2 + vllm_omni/deploy/README_higgs_audio_v3.md | 59 +++++++++++++++++++ vllm_omni/deploy/higgs_multimodal_qwen3.yaml | 5 ++ ...iggs_multimodal_qwen3_high_throughput.yaml | 5 ++ .../higgs_multimodal_qwen3_low_latency.yaml | 5 ++ 8 files changed, 86 insertions(+), 1 deletion(-) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1fb32239d..cc08ed397 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -471,6 +471,7 @@ def test_sub_config_fields_match_structured_scopes(): "trust_remote_code", "dtype", "attention_backend", + "attention_config", "moe_backend", "hf_overrides", "limit_mm_per_prompt", diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 75660cbc7..55ae43603 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -139,6 +139,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st engine_extras={ "model": str(stage_model), "attention_backend": "FLASHINFER", + "attention_config": {"use_trtllm_attention": False}, "hf_overrides": {"rope_scaling": {"factor": 2}}, "limit_mm_per_prompt": {"audio": 1}, "logits_processors": ["test.custom.LogitsProcessor"], @@ -229,6 +230,7 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path): assert thinker_args["omni_kv_config"] == {"need_send_cache": True} assert thinker_args["worker_cls"] == "test.custom.Worker" assert thinker_args["attention_backend"] == "FLASHINFER" + assert thinker_args["attention_config"] == {"use_trtllm_attention": False} assert thinker_args["trust_remote_code"] is True assert thinker_args["dtype"] == "float16" assert thinker_args["distributed_executor_backend"] == "mp" diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 6d81f300a..77a32ad7b 100755 --- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py @@ -1326,7 +1326,9 @@ class TestRegistry: assert "higgs_multimodal_qwen3" in OMNI_PIPELINES - def test_deploy_yaml_exists(self): + def test_deploy_yaml_uses_native_flashinfer(self): + from vllm_omni.config.stage_config import load_deploy_config + for name in ( "higgs_multimodal_qwen3.yaml", "higgs_multimodal_qwen3_high_throughput.yaml", @@ -1334,6 +1336,10 @@ class TestRegistry: ): yaml_path = get_deploy_config_path(name) assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}" + deploy_config = load_deploy_config(yaml_path) + stage_0 = next(stage for stage in deploy_config.stages if stage.stage_id == 0) + assert stage_0.engine_extras["attention_backend"] == "FLASHINFER" + assert stage_0.engine_extras["attention_config"] == {"use_trtllm_attention": False} # ---- AC-3: Prompt Builder ---- diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 32ee2c826..8c4a933e2 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -90,6 +90,7 @@ class _ModelEngineOverrides(TypedDict, total=False): trust_remote_code: bool dtype: Any attention_backend: Any + attention_config: Any moe_backend: str hf_overrides: Any limit_mm_per_prompt: dict[str, Any] @@ -321,6 +322,7 @@ class OmniStageModelConfig: trust_remote_code: bool = False dtype: Any = "auto" attention_backend: Any = None + attention_config: Any = None moe_backend: str = "auto" hf_overrides: Any = None limit_mm_per_prompt: dict[str, Any] | None = None diff --git a/vllm_omni/deploy/README_higgs_audio_v3.md b/vllm_omni/deploy/README_higgs_audio_v3.md index ea1363069..44f7f75c7 100644 --- a/vllm_omni/deploy/README_higgs_audio_v3.md +++ b/vllm_omni/deploy/README_higgs_audio_v3.md @@ -44,6 +44,65 @@ vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml ``` +## Why Stage 0 Pins Native FlashInfer + +The three Stage 0 profiles set `attention_backend: FLASHINFER` and also set: + +```yaml +attention_config: + use_trtllm_attention: false +``` + +`FLASHINFER` names the overall vLLM attention backend. Starting with vLLM +0.25, that backend can internally route Hopper decode to the TRT-LLM XQA +kernel. The XQA eligibility check answers whether the kernel can run, not +whether it is the fastest kernel for the current workload. In particular, the +decode builder can select XQA statically when all of the following hold: + +- the explicit `use_trtllm_attention` override is not `false`; +- the device is Hopper/SM90 and the XQA cubin is available; and +- the query-head count is divisible by the KV-head count. + +Higgs has 32 query heads and 8 KV heads, so it passes the head-layout check. +The static gate does not include a model dtype, KV-cache dtype, context-length, +batch-shape, or measured-performance comparison. This matters because kernel +support is not a performance guarantee. + +Higgs also has an unusual output shape that must not be confused with its +attention shape. One autoregressive audio step advances one token per request. +The resulting hidden row is then projected into eight codebook distributions +with shape `[1, 8, codebook_size]`. The eight codebooks therefore do not turn +the attention query into `q_len=8`; at concurrency one, XQA still receives a +single BF16 query row. At this small per-step shape, any XQA advantage for +other KV dtypes, batch sizes, or context lengths does not necessarily apply. +The difference is paid in every transformer layer and every generated audio +step, so even a small per-call disadvantage accumulates in TTFP and RTF. + +The upstream XQA integration also showed that the result is dtype-sensitive: +its FP8 benchmark favored XQA, while its BF16 benchmark reported 21.16 ms mean +TPOT for XQA versus 19.47 ms for FA3. That is not a direct native-FlashInfer +comparison, but it demonstrates why SM90 compatibility alone is insufficient +to choose the fastest BF16 kernel. See the +[vLLM XQA benchmark](https://github.com/vllm-project/vllm/pull/43232#issuecomment-4501176335). + +The end-to-end H800 A/B in +[issue #5584](https://github.com/vllm-project/vllm-omni/issues/5584#issuecomment-5188700218) +isolated the routing decision on the same vLLM 0.25 stack: + +| Decode path | TTFP (ms) | RTF | +|---|---:|---:| +| XQA (automatic default) | 175.76 | 0.3405 | +| Native FlashInfer (`use_trtllm_attention: false`) | 166.39 | 0.3233 | +| vLLM 0.24 reference | 166.68 | 0.3209 | + +This proves that XQA selection causes the observed local regression and that +native FlashInfer recovers the pre-v0.25 range. It does not by itself identify +one defective XQA instruction or memory transaction. Establishing that deeper +kernel-level cause requires an attention-only microbenchmark and GPU profiler +trace for the exact Higgs batch/context distribution. Until vLLM has a +workload-aware selector or XQA is faster for this BF16 shape, the model-level +pin is the deterministic choice. + ## Notes - Stage 1 remains `enforce_eager: true` in both profiles. diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml index da299415b..ae4e373cb 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml @@ -38,6 +38,11 @@ stages: gpu_memory_utilization: 0.6 enforce_eager: true attention_backend: FLASHINFER + # Eight codebooks share one attention row, so decode remains BF16 q_len=1 + # per request. vLLM's SM90 XQA capability gate does not benchmark this + # shape; keep the faster native path (README_higgs_audio_v3.md, #5584). + attention_config: + use_trtllm_attention: false trust_remote_code: true # Stage-0 prefix caching is on by default. The talker's HS opt-out and # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml index 8a751483c..bb37bdb4c 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml @@ -37,6 +37,11 @@ stages: gpu_memory_utilization: 0.6 enforce_eager: true attention_backend: FLASHINFER + # Eight codebooks share one attention row, so decode remains BF16 q_len=1 + # per request. vLLM's SM90 XQA capability gate does not benchmark this + # shape; keep the faster native path (README_higgs_audio_v3.md, #5584). + attention_config: + use_trtllm_attention: false trust_remote_code: true # Stage-0 prefix caching is on by default. The talker's HS opt-out and # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml index 0d6ae8218..13c6961e6 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml @@ -43,6 +43,11 @@ stages: gpu_memory_utilization: 0.6 enforce_eager: false attention_backend: FLASHINFER + # Eight codebooks share one attention row, so decode remains BF16 q_len=1 + # per request. vLLM's SM90 XQA capability gate does not benchmark this + # shape; keep the faster native path (README_higgs_audio_v3.md, #5584). + attention_config: + use_trtllm_attention: false trust_remote_code: true enable_prefix_caching: true async_scheduling: false From 90a08c0edaa15eda9f0b7bc740fd226703b77470 Mon Sep 17 00:00:00 2001 From: "Lucas(Ruijie)" <39816436+Moore-Z@users.noreply.github.com> Date: Thu, 13 Aug 2026 11:58:33 -0600 Subject: [PATCH 29/41] =?UTF-8?q?[Model][TTS][WIP]=20Add=20dots.tts=20(red?= =?UTF-8?q?note-hilab)=20=E2=80=94=20continuous-AR=2048kHz=20TTS=20(#4765)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Moore-Z --- docs/models/supported_models.md | 1 + .../text_to_speech/README.md | 26 + .../text_to_speech/dots_tts/end2end.py | 125 ++ recipes/rednote-hilab/dots.tts.md | 152 ++ .../models/dots_tts/__init__.py | 2 + .../models/dots_tts/test_talker.py | 137 ++ vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/dots_tts.yaml | 28 + .../models/dots_tts/__init__.py | 8 + .../models/dots_tts/dots_tts_dit.py | 457 ++++++ .../models/dots_tts/dots_tts_patch_encoder.py | 355 ++++ .../models/dots_tts/dots_tts_prompt.py | 70 + .../dots_tts/dots_tts_speaker_encoder.py | 732 +++++++++ .../models/dots_tts/dots_tts_talker.py | 1426 ++++++++++++++++ .../models/dots_tts/dots_tts_vocoder.py | 1434 +++++++++++++++++ .../models/dots_tts/pipeline.py | 46 + vllm_omni/model_executor/models/registry.py | 6 + .../transformers_utils/configs/__init__.py | 3 + .../transformers_utils/configs/dots_tts.py | 102 ++ 19 files changed, 5112 insertions(+) create mode 100644 examples/offline_inference/text_to_speech/dots_tts/end2end.py create mode 100644 recipes/rednote-hilab/dots.tts.md create mode 100644 tests/model_executor/models/dots_tts/__init__.py create mode 100644 tests/model_executor/models/dots_tts/test_talker.py create mode 100644 vllm_omni/deploy/dots_tts.yaml create mode 100644 vllm_omni/model_executor/models/dots_tts/__init__.py create mode 100644 vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py create mode 100644 vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py create mode 100644 vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py create mode 100644 vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py create mode 100644 vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py create mode 100644 vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py create mode 100644 vllm_omni/model_executor/models/dots_tts/pipeline.py create mode 100644 vllm_omni/transformers_utils/configs/dots_tts.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 73fdc55e4..995afc28d 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -59,6 +59,7 @@ th { | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — | +| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — | | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index d246f9e45..64e9b4d42 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -23,6 +23,7 @@ list of supported architectures across all modalities, see | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz | | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz | +| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz | | IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz | | IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz | @@ -430,6 +431,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [ --- +## dots.tts + +Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ + --model rednote-hilab/dots.tts-soar \ + --text "Hello, this is a test of dots TTS running on vLLM Omni." +``` + +### Voice cloning +Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet. + +### Streaming +The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example. + +### Notes +- Output: 48 kHz mono WAV. +- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`). +- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet. +- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap. + +--- + ## IndexTTS-2 and IndexTTS-2.5 Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning. diff --git a/examples/offline_inference/text_to_speech/dots_tts/end2end.py b/examples/offline_inference/text_to_speech/dots_tts/end2end.py new file mode 100644 index 000000000..8705f03f3 --- /dev/null +++ b/examples/offline_inference/text_to_speech/dots_tts/end2end.py @@ -0,0 +1,125 @@ +"""Offline dots.tts inference example (single-stage native AR pipeline). + +Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml). +Reference checkpoint: rednote-hilab/dots.tts-soar. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import soundfile as sf +import torch + +from vllm_omni import Omni +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +REPO_ROOT = Path(__file__).resolve().parents[4] +SAMPLE_RATE = 48_000 + + +def parse_args(): + parser = TrackingArgumentParser(description="Offline dots.tts native AR inference") + parser.add_argument( + "--model", + type=str, + default="rednote-hilab/dots.tts-soar", + help="dots.tts model path or HuggingFace repo ID.", + ) + parser.add_argument( + "--text", + type=str, + default="Hello, this is a test of dots TTS running on vLLM Omni.", + help="Text to synthesize.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="output_audio", + help="Directory for output WAV files.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path. If unset, auto-loads " + "vllm_omni/deploy/dots_tts.yaml based on the HF model_type.", + ) + return parser.parse_args() + + +def extract_audio(multimodal_output: dict) -> torch.Tensor: + """Extract the final complete audio tensor from multimodal output. + + The output processor concatenates per-step delta tensors under + ``model_outputs``. Falls back to ``audio`` for backwards compat. + """ + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}") + + if isinstance(audio, list): + valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] + if not valid: + raise ValueError("Audio list is empty or all elements are None.") + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] + + return torch.as_tensor(audio).float().cpu().reshape(-1) + + +def main(): + args = parse_args() + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + engine = Omni( + model=args.model, + deploy_config=args.deploy_config, + ) + + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import ( + build_dots_tts_prompt, + ) + + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + + prompt = build_dots_tts_prompt( + tokenizer=tokenizer, + text=args.text, + ) + + print(f"Model : {args.model}") + print(f"Text : {args.text}") + print(f"Output dir : {output_dir}") + print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens") + + t_start = time.perf_counter() + outputs = engine.generate([prompt]) + elapsed = time.perf_counter() - t_start + + request_output = outputs[0] + mm = request_output.outputs[0].multimodal_output + audio = extract_audio(mm) + + duration = audio.numel() / SAMPLE_RATE + rtf = elapsed / duration if duration > 0 else float("inf") + + output_path = output_dir / "output.wav" + sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV") + + print(f"Saved : {output_path}") + print(f"Duration : {duration:.2f}s") + print(f"Inference : {elapsed:.2f}s") + print(f"RTF : {rtf:.3f}") + + +if __name__ == "__main__": + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + main() diff --git a/recipes/rednote-hilab/dots.tts.md b/recipes/rednote-hilab/dots.tts.md new file mode 100644 index 000000000..d5c7868cb --- /dev/null +++ b/recipes/rednote-hilab/dots.tts.md @@ -0,0 +1,152 @@ +# dots.tts + +> Offline continuous-AR TTS at 48 kHz (rednote-hilab) + +## Summary + +- Vendor: rednote-hilab +- Model: `rednote-hilab/dots.tts-soar` +- Task: Text-to-speech, zero-shot synthesis only +- Mode: Offline end-to-end example only (not yet wired for online serving — + see [Known limitations](#known-limitations)) +- Maintainer: Community + +## When to use this recipe + +Use this recipe as a known-good starting point for running +`rednote-hilab/dots.tts-soar` offline on vLLM-Omni on consumer-class GPUs. +dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM ++ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio. +It follows the same "vLLM-native base LM + side-path computation" pattern as +VoxCPM2 — single-stage pipeline +`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR +loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone +instead of MiniCPM4, and no FSQ / residual-LM stage. + +This is an early integration (see [Known limitations](#known-limitations) +below before depending on it for anything beyond offline experimentation). + +## References + +- Offline end-to-end script: + [`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py) +- Example guide: + [`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts) +- Default deploy config: + [`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) +- Talker / pipeline source: + [`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/) +- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts) + +## Hardware Support + +This recipe documents one tested 16 GB consumer-GPU configuration. Other +vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as +community validation lands. + +## GPU + +### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended) + +dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++ +speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The +bundled default config at +[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) +(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`, +`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model +weights and ~0.3 GiB peak activation; the remainder of the configured +budget is available for KV cache. Total resident footprint at idle is +roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph +roadmap would be step 8's graph capture (not implemented yet; this +release runs `enforce_eager: true`, so it doesn't apply today). + +#### Environment + +- OS: Linux (WSL2) +- Python: 3.12 +- Driver / runtime: NVIDIA driver 595.95 +- torch: 2.11.0+cu130 +- vLLM: 0.26.0 +- vLLM-Omni: 0.22.1.dev (current `main`) + +#### Command + +```bash +python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ + --model rednote-hilab/dots.tts-soar \ + --text "Hello, this is a test of dots TTS running on vLLM Omni." +``` + +The deploy config at +[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) +is loaded automatically by the model registry (HF `model_type=dots_tts`). +Pass `--deploy-config ` to override. + +#### Verification + +**T1 — offline zero-shot synthesis**: + +```bash +python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ + --model rednote-hilab/dots.tts-soar \ + --text "Hello, this is a test of dots TTS running on vLLM Omni." +``` + +Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single +one-shot process (init → one `generate()` → exit), so the reported +numbers include engine init and first-request warmup, not just steady-state +per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`, +`RTF: 1.569`. + +Weight-loading breakdown from the same run (all tensors matched, no +missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270 +patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder. + +Whisper transcription of the output matched the input text with no +dropped leading word (confirms the streaming-vocoder patch-boundary fix +described in [Known limitations](#known-limitations)). + +#### Notes + +- Output: 48 kHz mono WAV. +- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default + used throughout this recipe. `rednote-hilab/dots.tts-base` shares the + same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf` + (MeanFlow, 2-4 step) is not supported — see below. +- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy + config are load-bearing, not just conservative defaults: with prefix + caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path + does not preserve the `sparse_audio` marker this model relies on to + route audio output correctly, and generation silently truncates to a + single ~160 ms patch. Do not override `enable_prefix_caching` for this + model until that framework-level gap is fixed. + +## Known limitations + +- **No online serving yet.** dots.tts is not registered in + `vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS + models in this repo use for `/v1/audio/speech`), so `vllm serve + rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible + speech endpoint does not know how to build request params for it yet. + Offline `Omni()` (as used in `end2end.py`) is the only supported path + today. +- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder + weights load and are exercised by `load_weights()`, but nothing in the + prompt builder consumes a reference audio yet — generation is zero-shot + only. `end2end.py` has no `--ref-audio`/`--ref-text` flags. +- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only + the fixed 10-step Euler DiT sampler used by `dots.tts-soar` / + `dots.tts-base` is implemented. +- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three + captured graphs (base LM decode, CFM solver, VAE decode) have no + dots.tts equivalent yet. +- **Concurrent requests do not scale.** Each request's 10-step DiT Euler + integration runs serially in the side path (no cross-request batching, + unlike voxcpm2's `enable_batched_cfm`). A community review of this + integration measured no throughput gain at `c=4` concurrent requests + versus `c=1`. +- **`SamplingParams.seed` does not control audio-generation randomness.** + The DiT's flow-matching noise is deterministically derived per-request + from a fixed internal seed (reproducible run-to-run, matching + voxcpm2's `deterministic_cfm_seed` convention) rather than from the + caller-supplied `seed` field — the same limitation voxcpm2 has today. diff --git a/tests/model_executor/models/dots_tts/__init__.py b/tests/model_executor/models/dots_tts/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/tests/model_executor/models/dots_tts/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/model_executor/models/dots_tts/test_talker.py b/tests/model_executor/models/dots_tts/test_talker.py new file mode 100644 index 000000000..ed3c6bcaa --- /dev/null +++ b/tests/model_executor/models/dots_tts/test_talker.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the dots.tts talker's stop-signal pairing and +per-request state lifecycle. + +Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes +the methods under test read) — no weights, no GPU, no engine. Mirrors +tests/model_executor/models/voxcpm2/test_talker_state_eviction.py. +""" + +from __future__ import annotations + +import functools +from types import SimpleNamespace + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +torch = pytest.importorskip("torch") + +_VOCAB_SIZE = 151672 + + +@functools.lru_cache(maxsize=1) +def _dots_tts_talker_mod(): + """Defer talker import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import ( + DotsTTSForConditionalGeneration, + _RequestState, + ) + + return DotsTTSForConditionalGeneration, _RequestState + + +def _make_bare_talker(): + DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod() + talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration) + talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE) + talker._active_states = {} + talker._pending_requests = [] + talker._results_queue = [] + talker._audio_queue = [] + talker._deferred_cleanup_ids = set() + talker._beta_trace = False + return talker + + +def _add_state(talker, req_id: str, *, is_stopping: bool = False): + _, _RequestState = _dots_tts_talker_mod() + state = _RequestState(request_id=req_id) + state.prefill_completed = True + state.is_stopping = is_stopping + talker._active_states[req_id] = state + return state + + +def _hidden(bsz: int) -> torch.Tensor: + return torch.zeros(bsz, 1536, dtype=torch.float32) + + +class TestComputeLogitsPairing: + """compute_logits pairs _results_queue entries with batch rows + positionally, so every request must push exactly one entry per step — + including prefills, which push a (req_id, None) placeholder.""" + + def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None: + talker = _make_bare_talker() + _add_state(talker, "prefill-req") + decode_state = _add_state(talker, "decode-req", is_stopping=True) + decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]]) + + # Row 0 is the prefill (None placeholder), row 1 the stopping decode. + talker._results_queue = [ + ("prefill-req", None), + ("decode-req", torch.tensor([[0.1, 0.9]])), + ] + logits = talker.compute_logits(_hidden(2)) + + assert logits.shape == (2, _VOCAB_SIZE) + # Prefill row: forced continue. + assert logits[0, 0].item() == 1.0 + assert logits[0, 1].item() == float("-inf") + # Stopping decode row: stop slot wins. + assert logits[1, 1].item() == 1.0 + assert logits[1, 0].item() == 0.0 + # Everything else stays masked. + assert logits[0, 2].item() == float("-inf") + assert logits[1, 2].item() == float("-inf") + # Queue drained. + assert talker._results_queue == [] + + def test_below_threshold_stop_logits_force_continue(self) -> None: + """Raw (continue, stop) softmax must not reach the sampler: a + prob_stop in (0.5, 0.8] would flip greedy argmax even though + is_stopping stayed False (the 0.8-threshold contract).""" + talker = _make_bare_talker() + _add_state(talker, "req", is_stopping=False) + talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))] + + logits = talker.compute_logits(_hidden(1)) + + assert logits[0, 0].item() == 1.0 + assert logits[0, 1].item() == float("-inf") + + def test_empty_queue_defaults_to_all_continue(self) -> None: + talker = _make_bare_talker() + logits = talker.compute_logits(_hidden(3)) + + assert torch.all(logits[:, 0] == 1.0) + assert torch.all(logits[:, 1] == float("-inf")) + + +class TestStateEviction: + """on_requests_finished fires BEFORE forward(), so eviction must be + deferred to _flush_deferred_cleanup at the end of the step — the + finishing request's audio still drains through the current forward.""" + + def test_finished_request_survives_until_flush(self) -> None: + talker = _make_bare_talker() + _add_state(talker, "req-a") + _add_state(talker, "req-b") + + talker.on_requests_finished(["req-a"]) + assert "req-a" in talker._active_states # not evicted yet + + talker._flush_deferred_cleanup() + assert "req-a" not in talker._active_states + assert "req-b" in talker._active_states + assert talker._deferred_cleanup_ids == set() + + def test_unknown_request_id_is_ignored(self) -> None: + talker = _make_bare_talker() + talker.on_requests_finished(["never-seen"]) + talker._flush_deferred_cleanup() + assert talker._active_states == {} diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index f02a3578e..f06259eaa 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -47,6 +47,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import ( ) from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE +from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE @@ -144,6 +145,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE, "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, + "dots_tts": DOTS_TTS_PIPELINE, "cosyvoice3": COSYVOICE3_PIPELINE, "audex_tts": AUDEX_TTS_PIPELINE, "audex_tta": AUDEX_TTA_PIPELINE, diff --git a/vllm_omni/deploy/dots_tts.yaml b/vllm_omni/deploy/dots_tts.yaml new file mode 100644 index 000000000..acd28de52 --- /dev/null +++ b/vllm_omni/deploy/dots_tts.yaml @@ -0,0 +1,28 @@ +# dots.tts deployment config (single-stage AR talker). +# +# Conservative first-integration settings: eager execution, no CUDA +# Graphs, no batched side-path optimizations yet — see voxcpm2.yaml for +# the runtime knobs a mature deployment of this model family grows. +# +# Verified locally on RTX 5080 16GB. + +async_chunk: false +dtype: bfloat16 + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.8 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 4096 + repetition_penalty: 1.0 diff --git a/vllm_omni/model_executor/models/dots_tts/__init__.py b/vllm_omni/model_executor/models/dots_tts/__init__.py new file mode 100644 index 000000000..fb6738a29 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py new file mode 100644 index 000000000..086ac109a --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone +# +# Merged from two upstream files (preserved as-is, inference behaviour +# unchanged; only internal imports were folded inline): +# - modules/backbone/layers.py -> Dropout, Mlp, rotate_half, +# apply_rotary_pos_emb, RotaryEmbedding, +# MultiHeadAttention +# (Conv1d / ConvTranspose1d are NOT +# re-vendored here — they already live in +# dots_tts_vocoder.py) +# - modules/backbone/dit.py -> modulate, TimestepEmbedder, FinalLayer, +# DiTBlock, DiT +# +# Per-section upstream attributions are preserved at each section header. +# Nothing in this file diverges from upstream behaviour. + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +# ============================================================================ +# Building blocks (adapted from modules/backbone/layers.py) +# ---------------------------------------------------------------------------- +# Conv1d / ConvTranspose1d from layers.py are already vendored in +# dots_tts_vocoder.py (they are used by the BigVGAN-style AudioVAE). The +# pieces below are the ones consumed by the DiT flow-matching head. +# ============================================================================ + + +class Dropout(nn.Module): + def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **_kwargs): + super().__init__() + if p < 0.0 or p > 1.0: + raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}") + self.p = p + self.inplace = inplace + self.force_drop = force_drop + + def forward(self, x, **_kwargs): + return F.dropout( + x, + p=self.p, + training=True if self.force_drop else self.training, + inplace=self.inplace, + ) + + +class Mlp(nn.Module): + def __init__( + self, + hidden_size, + ffn_hidden_size=4096, + act_layer=nn.GELU, + dropout=0.0, + **_kwargs, + ): + super().__init__() + self.fc1 = nn.Linear(hidden_size, ffn_hidden_size) + self.act = act_layer() + self.fc2 = nn.Linear(ffn_hidden_size, hidden_size) + self.drop = Dropout(dropout) + + def forward(self, x, _mask=None): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + return self.drop(x) + + +def rotate_half(x): + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +@torch.autocast(enabled=False, device_type="cuda") +def apply_rotary_pos_emb(pos, t): + if pos.dim() == 3: + pos = pos.unsqueeze(1) + return t * pos.cos() + rotate_half(t) * pos.sin() + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, theta=50000): + super().__init__() + self.register_buffer( + "inv_freq", + 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)), + persistent=False, + ) + self._theta = float(theta) + + def _apply(self, fn): + inv_freq = self.inv_freq + super()._apply(fn) + self.inv_freq = inv_freq.to(device=self.inv_freq.device, dtype=torch.float32) + return self + + @torch.autocast(enabled=False, device_type="cuda") + def forward(self, t): + inv_freq = self.inv_freq + if inv_freq.device != t.device: + raise RuntimeError(f"RotaryEmbedding buffer device mismatch: inv_freq={inv_freq.device} input={t.device}.") + t = t.to(dtype=inv_freq.dtype) + if t.dim() == 1: + freqs = torch.einsum("i , j -> i j", t, inv_freq) + else: + freqs = torch.einsum("bi, j -> bij", t, inv_freq) + return torch.cat((freqs, freqs), dim=-1) + + +class MultiHeadAttention(nn.Module): + """Multi-head attention""" + + def __init__( + self, + hidden_size: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_norm: bool = False, + attn_drop: float = 0.0, + dropout: float = 0.0, + norm_layer: str = "LayerNorm", + rotary_bias: bool = False, + rotary_theta: float | None = 50000, + **_kwargs, + ): + super().__init__() + assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads" + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scale = self.head_dim**-0.5 + self.rotary_bias = rotary_bias + + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + + norm_layer = getattr(nn, norm_layer) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + + self.attn_drop = Dropout(attn_drop) + self.o_proj = nn.Linear(hidden_size, hidden_size) + self.o_dropout = Dropout(dropout) + + if self.rotary_bias: + self.rotary = RotaryEmbedding(self.head_dim, theta=rotary_theta) + + def forward(self, q, k=None, v=None, mask=None, pos_ids=None, **_kwargs): + k = k or q + v = v or q + B, L, _ = q.shape + _, S, _ = v.shape + if mask is not None: + if mask.ndim == 2: # [B, L] + assert L == S + mask = rearrange(mask, "b j -> b 1 1 j") + mask = mask.expand(-1, self.num_heads, L, -1) + elif mask.ndim == 3: # [B, L, S] + assert mask.size(1) == L and mask.size(2) == S + mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1) + + q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v) + q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) + k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads) + v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads) + q, k = self.q_norm(q), self.k_norm(k) + + # Apply rotary + if self.rotary_bias: + if L == S: + if pos_ids is None: + rotary_emb = self.rotary(torch.arange(L, device=q.device)) + else: + rotary_emb = self.rotary(pos_ids) + q, k = (apply_rotary_pos_emb(rotary_emb, tensor) for tensor in (q, k)) + else: + q_rotary_emb = self.rotary(torch.arange(L, device=q.device)) + k_rotary_emb = self.rotary(torch.arange(S, device=k.device)) + q = apply_rotary_pos_emb(q_rotary_emb, q) + k = apply_rotary_pos_emb(k_rotary_emb, k) + + attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device) + + if mask is not None: + attn_bias.masked_fill_(mask.logical_not(), float("-inf")) + + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_bias, + dropout_p=self.attn_drop.p if self.training else 0.0, + ) + + out = rearrange(out, "b h n d -> b n (h d)") + return self.o_dropout(self.o_proj(out)) + + def decode_step(self, x, *, cache, positions: torch.Tensor): + if x.size(1) <= 0: + raise ValueError("MultiHeadAttention.decode_step expects a non-empty input.") + if positions.ndim != 1 or positions.size(0) != x.size(1): + raise ValueError("MultiHeadAttention.decode_step positions must match the decode block length.") + + q = self.q_proj(x) + k = self.k_proj(x) + v = self.v_proj(x) + + q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) + k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads) + v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads) + q, k = self.q_norm(q), self.k_norm(k) + block_len = q.size(2) + + if self.rotary_bias: + rotary_emb = self.rotary(positions) + q = apply_rotary_pos_emb(rotary_emb, q) + k = apply_rotary_pos_emb(rotary_emb, k) + + cached_k, cached_v = cache + cached_k.index_copy_(2, positions, k) + cached_v.index_copy_(2, positions, v) + + cache_capacity = cached_k.size(2) + key_positions = torch.arange( + cache_capacity, + device=x.device, + dtype=torch.long, + ).unsqueeze(0) + query_positions = positions.unsqueeze(1) + causal_mask = key_positions <= query_positions + valid_mask = key_positions <= positions[-1] + attn_bias = torch.zeros( + q.size(0), + self.num_heads, + block_len, + cache_capacity, + dtype=q.dtype, + device=q.device, + ) + attn_bias.masked_fill_( + (causal_mask & valid_mask).unsqueeze(0).unsqueeze(0).logical_not(), + float("-inf"), + ) + + out = F.scaled_dot_product_attention( + q, + cached_k, + cached_v, + attn_mask=attn_bias, + dropout_p=self.attn_drop.p if self.training else 0.0, + ) + out = rearrange(out, "b h n d -> b n (h d)") + return self.o_dropout(self.o_proj(out)), cache + + +# ============================================================================ +# DiT flow-matching head (adapted from modules/backbone/dit.py) +# ---------------------------------------------------------------------------- +# DiT.__init__ takes a ``transformer_config`` object exposing ``.to_dict()``, +# ``.hidden_size`` and ``.num_layers``. Upstream constructs it from a pydantic +# config class higher up the stack; vLLM-Omni callers can pass anything that +# satisfies the same minimal interface (e.g. a stdlib @dataclass with a +# ``to_dict`` method or SimpleNamespace). See dots_tts_talker.py for the +# concrete soar-checkpoint config factory. +# ============================================================================ + + +def modulate(x, shift, scale, **_kwargs): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + return self.mlp(t_freq) + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, output_size): + super().__init__() + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True), + ) + self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-5) + self.linear = nn.Linear(hidden_size, output_size, bias=True) + + def forward(self, x, c, **_kwargs): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm(x), shift, scale) + return self.linear(x) + + +class DiTBlock(nn.Module): + def __init__( + self, + attention: nn.Module, + ffn: nn.Module, + hidden_size: int = 1024, + modulation: bool = False, + eps: float = 1e-5, + **_kwargs, + ): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps) + self.attn = attention + self.ffn = ffn + self.modulation = modulation + if modulation: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 6 * hidden_size, bias=True), + ) + + def forward(self, x, condition=None, mask=None, **kwargs): + if condition is None: + assert not self.modulation, "Without global condition, must set modulation to False" + else: + assert self.modulation, "With global condition, must set modulation to True" + shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(condition).chunk( + 6, dim=1 + ) + + if condition is not None: + pack_indices = kwargs.get("pack_indices") + if pack_indices is not None: + gate_attn = gate_attn[pack_indices] + gate_ffn = gate_ffn[pack_indices] + else: + gate_attn = gate_attn.unsqueeze(1) + gate_ffn = gate_ffn.unsqueeze(1) + + if condition is not None: + x = x + gate_attn * self.attn( + modulate(self.norm1(x), shift_attn, scale_attn, **kwargs), + mask=mask, + **kwargs, + ) + else: + x = x + self.attn(self.norm1(x), mask=mask, **kwargs) + + if condition is not None: + x = x + gate_ffn * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn, **kwargs)) + else: + x = x + self.ffn(self.norm2(x), mask=mask) + return x + + +class DiT(nn.Module): + def __init__( + self, + in_dim, + out_dim, + transformer_config, + *, + mode: str = "flow_matching", + ): + super().__init__() + if mode not in {"flow_matching", "meanflow"}: + raise ValueError(f"DiT mode must be 'flow_matching' or 'meanflow', got {mode!r}.") + + transformer_kwargs = transformer_config.to_dict() + model_dim = transformer_config.hidden_size + self.mode = mode + self.num_layers = transformer_config.num_layers + + self.input_layer = nn.Linear(in_dim, model_dim) + self.time_embedder = TimestepEmbedder(model_dim) + if mode == "meanflow": + self.duration_embedder = TimestepEmbedder(model_dim) + + self.blocks = nn.ModuleList() + for i in range(self.num_layers): + attn_block = MultiHeadAttention(**transformer_kwargs, name=f"layer_{i}") + ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **transformer_kwargs) + self.blocks.append(DiTBlock(attention=attn_block, ffn=ffn_block, **transformer_kwargs)) + + self.output_layer = FinalLayer(model_dim, out_dim) + self.initialize_weights() + + def initialize_weights(self): + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + self.apply(_basic_init) + + nn.init.normal_(self.time_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.time_embedder.mlp[2].weight, std=0.02) + + for block in self.blocks: + if hasattr(block, "adaLN_modulation"): + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + nn.init.constant_(self.output_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.output_layer.adaLN_modulation[-1].bias, 0) + nn.init.constant_(self.output_layer.linear.weight, 0) + nn.init.constant_(self.output_layer.linear.bias, 0) + + def forward( + self, + x, + timesteps, + duration: torch.Tensor | None = None, + mask=None, + attn_mask=None, + g_cond: torch.Tensor | None = None, + **kwargs, + ): + t = self.time_embedder(timesteps) + c = t + duration_embedder = getattr(self, "duration_embedder", None) + if duration_embedder is not None and duration is not None: + c = c + duration_embedder(duration) + if g_cond is not None: + c = c + g_cond + + x = self.input_layer(x) + for block in self.blocks: + x = block(x, c, mask=attn_mask, **kwargs) + return self.output_layer(x, c, **kwargs) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py new file mode 100644 index 000000000..39d1a4973 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone +# +# Single upstream source (preserved as-is, inference behaviour unchanged; +# only the internal `from dots_tts.modules.backbone.layers import ...` line +# was folded to point at our sibling vendor files): +# - modules/backbone/semantic_encoder.py -> SemanticEncoderDecodeState, +# TransformerEncoderLayer, +# SuperviseEncoder, +# VAESemanticEncoder +# +# Building blocks (Conv1d / Mlp / MultiHeadAttention) are imported from the +# sibling vendor files where they already live: +# - Conv1d <- dots_tts_vocoder.py +# - Mlp <- dots_tts_dit.py +# - MultiHeadAttention <- dots_tts_dit.py +# +# Nothing else in this file diverges from upstream behaviour. + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import ( + Mlp, + MultiHeadAttention, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import Conv1d + + +@dataclass +class SemanticEncoderDecodeState: + conv_tail: torch.Tensor + layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...] + seq_len: int + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + hidden_size, + num_heads=16, + ffn_hidden_size=4096, + attn_dropout=0.0, + ffn_dropout=0.0, + norm_layer="LayerNorm", + **kwargs, + ): + super().__init__() + self.attn = MultiHeadAttention( + hidden_size, + num_heads, + attn_drop=attn_dropout, + norm_layer=norm_layer, + **kwargs, + ) + norm_cls = getattr(nn, norm_layer) + self.attn_norm = norm_cls(hidden_size) + self.ffn = Mlp(hidden_size, ffn_hidden_size, dropout=ffn_dropout, act_layer=nn.SiLU) + self.ffn_norm = norm_cls(hidden_size) + self.hidden_size = hidden_size + + def _build_causal_mask(self, T: int, device): # noqa: N803 + return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device)) + + def _build_padding_mask(self, x_lens, max_len: int, device): + B = x_lens.size(0) + positions = torch.arange(max_len, device=device).unsqueeze(0).expand(B, -1) + return positions < x_lens.unsqueeze(1) + + def _fuse_attn_mask(self, causal_mask, padding_mask): + if causal_mask is None and padding_mask is None: + return None + if causal_mask is None: + row = padding_mask.unsqueeze(2) + col = padding_mask.unsqueeze(1) + return row & col + if padding_mask is None: + return causal_mask.unsqueeze(0) + + _B, _T = padding_mask.shape + causal = causal_mask.unsqueeze(0) + row = padding_mask.unsqueeze(2) + col = padding_mask.unsqueeze(1) + pad_2d = row & col + return causal & pad_2d + + def forward( + self, + x, + x_lens=None, + causal=True, + ): + _B, T, C = x.shape + assert self.hidden_size == C + device = x.device + + causal_mask = self._build_causal_mask(T, device) if causal else None + if x_lens is not None: + padding_mask = self._build_padding_mask(x_lens, T, device) + else: + padding_mask = None + fused_mask = self._fuse_attn_mask(causal_mask, padding_mask) + + h = self.attn_norm(x) + h = self.attn( + q=h, + mask=fused_mask, + ) + x = x + h + + h = self.ffn_norm(x) + h = self.ffn(h) + return x + h + + def decode_step( + self, + x, + *, + cache: tuple[torch.Tensor, torch.Tensor], + positions: torch.Tensor, + ): + if x.size(1) <= 0: + raise ValueError("TransformerEncoderLayer.decode_step expects a non-empty input.") + + h = self.attn_norm(x) + h, cache = self.attn.decode_step(h, cache=cache, positions=positions) + x = x + h + + h = self.ffn_norm(x) + h = self.ffn(h) + return x + h, cache + + +class SuperviseEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.get("hidden_size", 1024) + self.layers = nn.ModuleList( + [ + TransformerEncoderLayer( + hidden_size=self.hidden_size, + num_heads=config.get("num_heads", 16), + ffn_hidden_size=config.get("ffn_hidden_size", 4096), + norm_layer=config.get("norm_layer", "LayerNorm"), + ) + for _ in range(config.get("num_layers", 6)) + ] + ) + self.causal = config.get("causal", False) + + def forward(self, x, x_lens=None): + batch_size, seq_len, _ = x.shape + if x_lens is None: + x_lens = torch.full((batch_size,), seq_len, device=x.device, dtype=torch.long) + for layer in self.layers: + x = layer(x, x_lens=x_lens, causal=self.causal) + return x + + def init_decode_state( + self, + *, + batch_size: int, + max_seq_len: int, + device: torch.device, + dtype: torch.dtype, + ): + layer_caches = [] + for layer in self.layers: + cache_shape = ( + batch_size, + layer.attn.num_heads, + max_seq_len, + layer.attn.head_dim, + ) + layer_caches.append( + ( + torch.zeros(cache_shape, dtype=dtype, device=device), + torch.zeros(cache_shape, dtype=dtype, device=device), + ) + ) + return tuple(layer_caches) + + def reset_decode_state( + self, + layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...], + ) -> None: + if len(layer_caches) != len(self.layers): + raise ValueError("Layer cache count does not match encoder depth.") + for key_cache, value_cache in layer_caches: + key_cache.zero_() + value_cache.zero_() + + def decode_step(self, x, *, layer_caches, positions: torch.Tensor): + if len(layer_caches) != len(self.layers): + raise ValueError("Layer cache count does not match encoder depth.") + + for layer, cache in zip(self.layers, layer_caches, strict=True): + x, _ = layer.decode_step(x, cache=cache, positions=positions) + return x + + +class VAESemanticEncoder(nn.Module): + def __init__(self, in_dim, out_dim, config): + super().__init__() + in_ds_rate = 2 + self.patch_size = int(config.patch_size) + self.in_ds_rate = in_ds_rate + self.ds_proj = Conv1d(in_dim, in_dim, kernel_size=in_ds_rate, stride=in_ds_rate, causal=True) + self.in_proj = nn.Linear(in_dim, config.PatchEncoder.hidden_size) + self.encoder = SuperviseEncoder(config.PatchEncoder) + self.out_ds_rate = self.patch_size // in_ds_rate + self.out_proj = nn.Linear(config.PatchEncoder.hidden_size * self.out_ds_rate, out_dim) + + def forward(self, x, x_lens=None): + x = self._downsample(x) + x = self.in_proj(x) + z = self.encoder(x, x_lens=x_lens) + return self._project_embeddings(z) + + def init_decode_state( + self, + *, + max_audio_patch_count: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> SemanticEncoderDecodeState: + return SemanticEncoderDecodeState( + conv_tail=torch.zeros( + (batch_size, self.ds_proj.in_channels, self.ds_proj.left_padding), + dtype=dtype, + device=device, + ), + layer_caches=self.encoder.init_decode_state( + batch_size=batch_size, + max_seq_len=max_audio_patch_count * self.out_ds_rate, + device=device, + dtype=dtype, + ), + seq_len=0, + ) + + def reset_decode_state(self, state: SemanticEncoderDecodeState) -> None: + state.conv_tail.zero_() + self.encoder.reset_decode_state(state.layer_caches) + state.seq_len = 0 + + def prefill( + self, + x, + state: SemanticEncoderDecodeState, + ) -> tuple[torch.Tensor, SemanticEncoderDecodeState]: + if x.ndim != 3: + raise ValueError(f"VAESemanticEncoder.prefill expects rank-3 input, got {tuple(x.shape)}.") + if x.size(1) % self.patch_size != 0: + raise ValueError(f"Prompt latent length {x.size(1)} must be divisible by patch_size={self.patch_size}.") + + if x.size(1) == 0: + return ( + x.new_zeros((x.size(0), 0, self.out_proj.out_features)), + state, + ) + if state.conv_tail.size(0) != x.size(0): + raise ValueError("VAESemanticEncoder.prefill batch size does not match decode state.") + + step_inputs = self.in_proj(self._downsample(x)) + expected_token_count = (x.size(1) // self.patch_size) * self.out_ds_rate + if step_inputs.size(1) != expected_token_count: + raise RuntimeError( + "Patch encoder prefill produced an unexpected token count: " + f"expected={expected_token_count} actual={step_inputs.size(1)}." + ) + + current_seq_len = state.seq_len + next_seq_len = current_seq_len + step_inputs.size(1) + cache_capacity = state.layer_caches[0][0].size(2) + if next_seq_len > cache_capacity: + raise ValueError( + "Patch encoder prefill exceeds decode-state capacity: " + f"required={next_seq_len} capacity={cache_capacity}." + ) + + positions = torch.arange(step_inputs.size(1), device=x.device, dtype=torch.long) + current_seq_len + encoded = self.encoder.decode_step( + step_inputs, + layer_caches=state.layer_caches, + positions=positions, + ) + embedding = self._project_embeddings(encoded) + raw = x.transpose(1, 2) + state.conv_tail.copy_(raw[..., -self.ds_proj.left_padding :]) + state.seq_len = next_seq_len + return embedding, state + + def decode_patch( + self, + latent_patch, + conv_tail: torch.Tensor, + layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...], + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if latent_patch.ndim != 3: + raise ValueError(f"VAESemanticEncoder.decode_patch expects rank-3 input, got {tuple(latent_patch.shape)}.") + if latent_patch.size(1) != self.patch_size: + raise ValueError(f"decode_patch expects patch length {self.patch_size}, got {latent_patch.size(1)}.") + if positions.ndim != 1 or positions.size(0) != self.out_ds_rate: + raise ValueError("decode_patch positions must be a rank-1 tensor matching out_ds_rate.") + + step_inputs, conv_tail = self._downsample_step( + latent_patch, + conv_tail=conv_tail, + ) + if step_inputs.size(1) != self.out_ds_rate: + raise RuntimeError(f"Downsample step produced {step_inputs.size(1)} tokens, expected {self.out_ds_rate}.") + + encoded = self.encoder.decode_step( + step_inputs, + layer_caches=layer_caches, + positions=positions, + ) + embedding = self._project_embeddings(encoded) + return embedding, conv_tail + + def _downsample(self, x): + return self.ds_proj(x.transpose(1, 2)).transpose(1, 2) + + def _project_embeddings(self, z): + if self.out_ds_rate > 1: + z = rearrange(z, "b (s d) h -> b s (d h)", d=self.out_ds_rate) + return self.out_proj(z) + + def _downsample_step(self, latent_patch, *, conv_tail): + raw = latent_patch.transpose(1, 2) + conv_input = torch.cat([conv_tail, raw], dim=-1) + + projected = F.conv1d( + conv_input, + self.ds_proj.weight, + self.ds_proj.bias, + stride=self.ds_proj.stride[0], + padding=0, + dilation=self.ds_proj.dilation[0], + groups=self.ds_proj.groups, + ).transpose(1, 2) + new_conv_tail = raw[..., -self.ds_proj.left_padding :] + return self.in_proj(projected), new_conv_tail diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py new file mode 100644 index 000000000..a2c4ac9f9 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dots.tts prompt builder for vllm-omni's Omni engine + serving layer. + +Mirrors upstream ``build_generation_schedule`` (rednote-hilab/dots.tts @ +a393d2e data/pipelines/tokenizing.py:169) for the +no-prompt-audio + audio-placeholder path — text-only zero-shot TTS. + +V1 returns real Qwen2 token IDs (not placeholders ``[1] * N`` like +voxcpm2's builder). Reason: our talker's preprocess() runs +``self.model.embed_tokens(input_ids)`` directly, so real IDs flow +through naturally. voxcpm2 uses placeholders because its prefill mixes +reference-audio embeddings (no token IDs); dots.tts V1 zero-shot is +purely text + audio_gen_start, all real IDs. + +Voice clone (reference audio + speaker x-vector) is intentionally not +wired here — V1 is zero-shot only. Add reference_audio / +voice_profile / prompt_text branches in a follow-up (mirror voxcpm2 +``build_voxcpm2_prompt``). +""" + +from __future__ import annotations + +from typing import Any + +from vllm.inputs import tokens_input + +# Upstream rednote-hilab/dots.tts @ a393d2e +# data/pipelines/tts_pipeline.py:17-18 +_TTS_TEXT_PREFIX = "[文本]" +_TTS_AUDIO_PREFIX = "[文本对应语音]" +_AUDIO_GEN_START_TOKEN = "<|audio_gen_start|>" + + +def build_dots_tts_prompt( + tokenizer: Any, + text: str, +) -> dict[str, Any]: + """Build a dots.tts prefill prompt dict for ``Omni.generate()``. + + Returns ``{"prompt_token_ids": [...], "additional_information": {...}}`` + where ``prompt_token_ids`` is the real Qwen2 token sequence: + + [text_tokens of "[文本][文本对应语音]", audio_gen_start_id] + + talker's preprocess() embeds these into ``hidden_states`` for the AR + base LM. Decode then drives via ``state.curr_embed_for_next`` (DiT + → patch_encoder AR loopback). + """ + if not isinstance(text, str) or not text.strip(): + raise ValueError(f"text must be a non-empty string, got {text!r}") + + prompt_text = f"{_TTS_TEXT_PREFIX}{text}{_TTS_AUDIO_PREFIX}" + text_ids = tokenizer.encode(prompt_text, add_special_tokens=False) + + audio_gen_start_id = tokenizer.convert_tokens_to_ids(_AUDIO_GEN_START_TOKEN) + if not isinstance(audio_gen_start_id, int) or audio_gen_start_id == tokenizer.unk_token_id: + raise ValueError( + f"Tokenizer does not know {_AUDIO_GEN_START_TOKEN!r} " + f"(got id={audio_gen_start_id!r}). Did you load the " + f"dots.tts tokenizer (with added_tokens.json)?" + ) + prompt_token_ids = list(text_ids) + [audio_gen_start_id] + + prompt = tokens_input(prompt_token_ids=prompt_token_ids) + prompt["additional_information"] = { + # Reserved for future voice clone (ref_audio / speaker x-vector). + # V1 zero-shot leaves this empty; preprocess() ignores it. + } + return prompt diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py new file mode 100644 index 000000000..09961d0ab --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py @@ -0,0 +1,732 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Portions Copyright 3D-Speaker +# (https://github.com/alibaba-damo-academy/3D-Speaker), Apache-2.0. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/speaker +# (whose CAM++ implementation is itself adapted from 3D-Speaker). +# +# Merged from four upstream files (preserved as-is, inference behaviour +# unchanged; only internal cross-module imports were folded inline): +# - modules/speaker/encoder.py -> SpeakerXVectorFeatures +# - modules/speaker/fbank.py -> fbank constants + extract_speaker_fbank +# - modules/speaker/campplus.py -> FCM, CAMPPlus +# - modules/speaker/campplus_layers.py -> get_nonlinear, statistics_pooling, +# StatsPool, TDNNLayer, CAMLayer, +# CAMDenseTDNNLayer, CAMDenseTDNNBlock, +# TransitLayer, DenseLayer, BasicResBlock +# +# Plus two tiny helpers folded from utils/audio.py (used by +# extract_speaker_fbank): +# - high_quality_resample, extract_fbank +# +# Nothing in this file diverges from upstream behaviour. + +from __future__ import annotations + +import math +import random +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +import torchaudio +import torchaudio.compliance.kaldi as Kaldi +import torchaudio.functional as AF +from torch.nn.utils.rnn import pad_sequence + +# ============================================================================ +# Constants (from modules/speaker/fbank.py) +# ============================================================================ + +_SPEAKER_FBANK_SAMPLE_RATE = 16000 +_SPEAKER_FBANK_N_MELS = 80 +_SPEAKER_FBANK_MEAN_NORM = True +_SPEAKER_FBANK_DITHER = 0.0 + + +# ============================================================================ +# Audio helpers (folded from utils/audio.py — used by extract_speaker_fbank) +# ============================================================================ + + +def high_quality_resample(x, orig_sr, target_sr): + return AF.resample( + x, + orig_freq=orig_sr, + new_freq=target_sr, + lowpass_filter_width=64, + rolloff=0.95, + resampling_method="sinc_interp_kaiser", + ) + + +def extract_fbank( + waveform: torch.Tensor, + *, + sample_rate: int, + n_mels: int, + dither: float = 0.0, + mean_norm: bool = False, +) -> torch.Tensor: + if waveform.ndim == 1: + feature_input = waveform.unsqueeze(0) + elif waveform.ndim == 2: + feature_input = waveform if waveform.size(0) == 1 else waveform[0:1, :] + else: + raise ValueError(f"FBank expects a 1D or 2D waveform, got shape {tuple(waveform.shape)}.") + features = Kaldi.fbank( + feature_input, + num_mel_bins=n_mels, + sample_frequency=sample_rate, + dither=dither, + ) + if mean_norm: + features = features - features.mean(dim=0, keepdim=True) + return features + + +# ============================================================================ +# Speaker fbank entry (from modules/speaker/fbank.py) +# ============================================================================ + + +def extract_speaker_fbank( + waveform: torch.Tensor, + *, + sample_rate: int, +) -> torch.Tensor: + feature_input = waveform + if sample_rate != _SPEAKER_FBANK_SAMPLE_RATE: + feature_input = high_quality_resample( + waveform, + orig_sr=sample_rate, + target_sr=_SPEAKER_FBANK_SAMPLE_RATE, + ) + return extract_fbank( + feature_input, + sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, + n_mels=_SPEAKER_FBANK_N_MELS, + dither=_SPEAKER_FBANK_DITHER, + mean_norm=_SPEAKER_FBANK_MEAN_NORM, + ) + + +# ============================================================================ +# CAM++ building blocks (from modules/speaker/campplus_layers.py) +# ---------------------------------------------------------------------------- +# 3D-Speaker, Apache-2.0. +# ============================================================================ + + +def get_nonlinear(config_str, channels): + nonlinear = nn.Sequential() + for name in config_str.split("-"): + if name == "relu": + nonlinear.add_module("relu", nn.ReLU(inplace=True)) + elif name == "prelu": + nonlinear.add_module("prelu", nn.PReLU(channels)) + elif name == "batchnorm": + nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels)) + elif name == "batchnorm_": + nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False)) + else: + raise ValueError(f"Unexpected module ({name}).") + return nonlinear + + +def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, _eps=1e-2): + mean = x.mean(dim=dim) + std = x.std(dim=dim, unbiased=unbiased) + stats = torch.cat([mean, std], dim=-1) + if keepdim: + stats = stats.unsqueeze(dim=dim) + return stats + + +class StatsPool(nn.Module): + def forward(self, x): + return statistics_pooling(x) + + +class TDNNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + bias=False, + config_str="batchnorm-relu", + ): + super().__init__() + if padding < 0: + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.linear = nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + x = self.linear(x) + return self.nonlinear(x) + + +class CAMLayer(nn.Module): + def __init__( + self, + bn_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + bias, + reduction=2, + ): + super().__init__() + self.linear_local = nn.Conv1d( + bn_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1) + self.relu = nn.ReLU(inplace=True) + self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + y = self.linear_local(x) + context = x.mean(-1, keepdim=True) + self.seg_pooling(x) + context = self.relu(self.linear1(context)) + m = self.sigmoid(self.linear2(context)) + return y * m + + def seg_pooling(self, x, seg_len=100, stype="avg"): + if stype == "avg": + seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + elif stype == "max": + seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + else: + raise ValueError("Wrong segment pooling type.") + shape = seg.shape + seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1) + return seg[..., : x.shape[-1]] + + +class CAMDenseTDNNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.memory_efficient = memory_efficient + self.nonlinear1 = get_nonlinear(config_str, in_channels) + self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False) + self.nonlinear2 = get_nonlinear(config_str, bn_channels) + self.cam_layer = CAMLayer( + bn_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + + def bn_function(self, x): + return self.linear1(self.nonlinear1(x)) + + def forward(self, x): + if self.training and self.memory_efficient: + x = cp.checkpoint(self.bn_function, x) + else: + x = self.bn_function(x) + return self.cam_layer(self.nonlinear2(x)) + + +class CAMDenseTDNNBlock(nn.ModuleList): + def __init__( + self, + num_layers, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + for i in range(num_layers): + layer = CAMDenseTDNNLayer( + in_channels=in_channels + i * out_channels, + out_channels=out_channels, + bn_channels=bn_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + bias=bias, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.add_module(f"tdnnd{i + 1}", layer) + + def forward(self, x): + for layer in self: + x = torch.cat([x, layer(x)], dim=1) + return x + + +class TransitLayer(nn.Module): + def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"): + super().__init__() + self.nonlinear = get_nonlinear(config_str, in_channels) + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + + def forward(self, x): + x = self.nonlinear(x) + return self.linear(x) + + +class DenseLayer(nn.Module): + def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"): + super().__init__() + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + if len(x.shape) == 2: + x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1) + else: + x = self.linear(x) + return self.nonlinear(x) + + +class BasicResBlock(nn.Module): + expansion = 1 + + def __init__(self, in_planes, planes, stride=1): + super().__init__() + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=(stride, 1), + bias=False, + ), + nn.BatchNorm2d(self.expansion * planes), + ) + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out += self.shortcut(x) + return F.relu(out) + + +# ============================================================================ +# CAM++ model (from modules/speaker/campplus.py) +# ---------------------------------------------------------------------------- +# 3D-Speaker, Apache-2.0. +# ============================================================================ + + +class FCM(nn.Module): + def __init__( + self, + block=BasicResBlock, + num_blocks=(2, 2), + m_channels=32, + feat_dim=_SPEAKER_FBANK_N_MELS, + ): + super().__init__() + self.in_planes = m_channels + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(m_channels) + + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2) + self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2) + + self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(m_channels) + self.out_channels = m_channels * (feat_dim // 8) + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + x = x.unsqueeze(1) + out = F.relu(self.bn1(self.conv1(x))) + out = self.layer1(out) + out = self.layer2(out) + out = F.relu(self.bn2(self.conv2(out))) + + shape = out.shape + return out.reshape(shape[0], shape[1] * shape[2], shape[3]) + + +class CAMPPlus(nn.Module): + _TDNN_KERNEL_SIZE = 5 + _TDNN_STRIDE = 2 + _TDNN_PADDING = 2 + + def __init__( + self, + feat_dim=_SPEAKER_FBANK_N_MELS, + embedding_size=512, + growth_rate=32, + bn_size=4, + init_channels=128, + config_str="batchnorm-relu", + memory_efficient=True, + ): + super().__init__() + + self.head = FCM(feat_dim=feat_dim) + channels = self.head.out_channels + + self.xvector = nn.Sequential( + OrderedDict( + [ + ( + "tdnn", + TDNNLayer( + channels, + init_channels, + self._TDNN_KERNEL_SIZE, + stride=self._TDNN_STRIDE, + dilation=1, + padding=-1, + config_str=config_str, + ), + ), + ] + ) + ) + channels = init_channels + for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2), strict=True)): + block = CAMDenseTDNNBlock( + num_layers=num_layers, + in_channels=channels, + out_channels=growth_rate, + bn_channels=bn_size * growth_rate, + kernel_size=kernel_size, + dilation=dilation, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.xvector.add_module(f"block{i + 1}", block) + channels = channels + num_layers * growth_rate + self.xvector.add_module( + f"transit{i + 1}", + TransitLayer(channels, channels // 2, bias=False, config_str=config_str), + ) + channels //= 2 + + self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels)) + + self.xvector.add_module("stats", StatsPool()) + self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")) + + for m in self.modules(): + if isinstance(m, (nn.Conv1d, nn.Linear)): + nn.init.kaiming_normal_(m.weight.data) + if m.bias is not None: + nn.init.zeros_(m.bias) + + @staticmethod + def _conv_output_lengths(lengths, kernel_size, stride=1, padding=0, dilation=1): + return ( + torch.div( + lengths + 2 * padding - dilation * (kernel_size - 1) - 1, + stride, + rounding_mode="floor", + ) + + 1 + ) + + @staticmethod + def _make_length_mask(lengths, max_len, device): + lengths = lengths.to(device=device, dtype=torch.long).clamp(min=0, max=max_len) + return torch.arange(max_len, device=device).unsqueeze(0) < lengths.unsqueeze(1) + + def _masked_stats_pooling(self, x, lengths, unbiased=True, eps=1e-2): + lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1, max=x.size(-1)) + mask = self._make_length_mask(lengths, x.size(-1), x.device).unsqueeze(1) + mask = mask.to(dtype=x.dtype) + + denom = lengths.to(dtype=x.dtype).view(-1, 1).clamp_min(1.0) + mean = (x * mask).sum(dim=-1) / denom + + centered = (x - mean.unsqueeze(-1)) * mask + var_denom = (lengths - 1).clamp_min(1).to(dtype=x.dtype).view(-1, 1) if unbiased else denom + var = centered.pow(2).sum(dim=-1) / var_denom + std = torch.sqrt(var.clamp_min(eps)) + return torch.cat([mean, std], dim=1) + + def forward(self, x, lengths=None): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + x = self.head(x) + if lengths is not None: + lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1) + + for name, module in self.xvector.named_children(): + if name == "stats": + x = self._masked_stats_pooling(x, lengths) if lengths is not None else module(x) + continue + + x = module(x) + if name == "tdnn" and lengths is not None: + lengths = self._conv_output_lengths( + lengths, + kernel_size=self._TDNN_KERNEL_SIZE, + stride=self._TDNN_STRIDE, + padding=self._TDNN_PADDING, + ) + + return x + + +# ============================================================================ +# SpeakerXVectorFeatures wrapper (from modules/speaker/encoder.py) +# ============================================================================ + + +class SpeakerXVectorFeatures(nn.Module): + """Speaker embedding extractor based on 3D-Speaker CAM++.""" + + def __init__( + self, + sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, + campplus_embedding_size=512, + max_audio_seconds=10.0, + ): + super().__init__() + + self.sample_rate = sample_rate + self.max_audio_seconds = float(max_audio_seconds) + self.model = CAMPPlus( + feat_dim=_SPEAKER_FBANK_N_MELS, + embedding_size=campplus_embedding_size, + ) + self.resample = None + if self.sample_rate != _SPEAKER_FBANK_SAMPLE_RATE: + self.resample = torchaudio.transforms.Resample( + orig_freq=sample_rate, + new_freq=_SPEAKER_FBANK_SAMPLE_RATE, + ) + + for param in self.model.parameters(): + param.requires_grad = False + + @staticmethod + def _normalize_lengths(lengths, batch_size, max_length, device, *, min_length): + if lengths is None: + return torch.full( + (batch_size,), + max_length, + device=device, + dtype=torch.long, + ) + return lengths.to(device=device, dtype=torch.long).clamp( + min=min_length, + max=max_length, + ) + + def _crop_audio(self, audio, audio_lengths=None): + original_lengths = self._normalize_lengths( + audio_lengths, + audio.size(0), + audio.size(-1), + audio.device, + min_length=0, + ) + if self.max_audio_seconds <= 0: + return audio, original_lengths, original_lengths, torch.zeros_like(original_lengths) + + max_input_length = round(self.sample_rate * self.max_audio_seconds) + cropped_audio = [] + cropped_lengths = [] + starts = [] + + for index, total_length_tensor in enumerate(original_lengths): + total_length = int(total_length_tensor.item()) + cropped_length = min(total_length, max_input_length) + start = random.randint(0, total_length - cropped_length) if total_length > cropped_length else 0 + cropped_audio.append(audio[index, start : start + cropped_length]) + cropped_lengths.append(cropped_length) + starts.append(start) + + return ( + pad_sequence( + cropped_audio, + batch_first=True, + padding_value=0.0, + ), + original_lengths, + torch.tensor( + cropped_lengths, + device=audio.device, + dtype=torch.long, + ), + torch.tensor(starts, device=audio.device, dtype=torch.long), + ) + + def _crop_fbank( + self, + fbank, + fbank_lengths, + original_audio_lengths, + cropped_audio_lengths, + starts, + ): + original_fbank_lengths = self._normalize_lengths( + fbank_lengths, + fbank.size(0), + fbank.size(1), + fbank.device, + min_length=1, + ) + cropped_fbank = [] + cropped_fbank_lengths = [] + + for index, total_feat_length_tensor in enumerate(original_fbank_lengths): + total_audio_length = int(original_audio_lengths[index].item()) + total_feat_length = int(total_feat_length_tensor.item()) + start_audio = int(starts[index].item()) + end_audio = start_audio + int(cropped_audio_lengths[index].item()) + + if total_audio_length > 0: + start_feat = math.floor(start_audio * total_feat_length / total_audio_length) + end_feat = math.ceil(end_audio * total_feat_length / total_audio_length) + else: + start_feat = 0 + end_feat = 1 + + start_feat = min(start_feat, total_feat_length - 1) + end_feat = min(max(end_feat, start_feat + 1), total_feat_length) + cropped_fbank.append(fbank[index, start_feat:end_feat]) + cropped_fbank_lengths.append(end_feat - start_feat) + + return pad_sequence( + cropped_fbank, + batch_first=True, + padding_value=0.0, + ), torch.tensor( + cropped_fbank_lengths, + device=fbank.device, + dtype=torch.long, + ) + + def _extract_fbank_batch(self, audio, audio_lengths): + if self.resample is not None: + audio = self.resample(audio) + audio_lengths = torch.ceil(audio_lengths.float() * (_SPEAKER_FBANK_SAMPLE_RATE / self.sample_rate)).long() + + audio_cpu = audio.detach().cpu() + features = [] + + for index, valid_length_tensor in enumerate(audio_lengths): + valid_length = int(valid_length_tensor.item()) + waveform = audio_cpu[index, :valid_length] + if waveform.numel() == 0: + waveform = audio_cpu.new_zeros(1) + features.append( + extract_speaker_fbank( + waveform, + sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, + ) + ) + + fbank_lengths = torch.tensor( + [feature.size(0) for feature in features], + device=audio.device, + dtype=torch.long, + ) + fbank = pad_sequence( + features, + batch_first=True, + padding_value=0.0, + ).to(device=audio.device, dtype=audio.dtype) + return fbank, fbank_lengths + + @torch.no_grad() + @torch.autocast(enabled=False, device_type="cuda") + def forward(self, audio, audio_lengths=None, fbank=None, fbank_lengths=None, **_kwargs): + self.model.eval() + audio = audio.float() + if audio.dim() == 3: + if audio.size(1) != 1: + raise ValueError(f"Speaker encoder expects mono audio, got shape {tuple(audio.shape)}.") + audio = audio[:, 0] + elif audio.dim() != 2: + raise ValueError(f"Speaker encoder expects a 2D or 3D audio tensor, got shape {tuple(audio.shape)}.") + + audio, original_audio_lengths, cropped_audio_lengths, starts = self._crop_audio( + audio, + audio_lengths=audio_lengths, + ) + + if fbank is None: + fbank, fbank_lengths = self._extract_fbank_batch( + audio, + cropped_audio_lengths, + ) + else: + if not isinstance(fbank, torch.Tensor): + raise TypeError("Speaker encoder expects `fbank` to be a torch.Tensor.") + if fbank.dim() != 3 or fbank.size(0) != audio.size(0): + raise ValueError( + f"Speaker encoder expects `fbank` with shape (B, T, F) and matching batch size, got {tuple(fbank.shape)}." # noqa: E501 + ) + fbank, fbank_lengths = self._crop_fbank( + fbank.to(device=audio.device, dtype=torch.float32), + fbank_lengths, + original_audio_lengths, + cropped_audio_lengths, + starts, + ) + + return self.model(fbank, lengths=fbank_lengths) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py new file mode 100644 index 000000000..368a3fd1e --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py @@ -0,0 +1,1426 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dots.tts talker — vLLM-native AR base LM + audio side path. + +Mirrors upstream rednote-hilab/dots.tts (pinned @ a393d2e): + + * Qwen2.5-1.5B base LM runs under vLLM PagedAttention (stock + ``Qwen2Model``); weights come from the checkpoint's ``llm.model.*`` + namespace. ``DotsTTSConfig`` hoists the LM fields to top-level + attributes. + * Everything else is side computation, invisible to the vLLM + scheduler and driven once per step by ``_finish_decode``: + DiT flow-matching head (N-step Euler, fp32 integration state, + deterministic per-request noise) → io_helper denormalize → + patch_encoder AR loopback (produces the next step's LLM input + embedding) + streaming AudioVAE decode (per-request + ``BigVGANStreamState`` sliding window, so every latent patch is + decoded with real left conv context) → 48 kHz wav chunks pushed to + ``_audio_queue`` → eos_proj stop signal via ``_results_queue`` + (threshold 0.8, upstream default; the stop step also drains the + vocoder's 2-frame lookahead tail via ``stream_flush``). + * Cross-step state is isolated per request in ``_RequestState`` + (keyed by request_id); eviction is deferred to the end of + ``forward()`` so a finishing request's audio still drains. + * The CAM++ speaker encoder is constructed and its weights load, but + voice cloning is not exposed yet — generation is zero-shot and the + DiT falls back to its null conditioning (``fm_null_g_cond``). + +Debugging: set ``DOTS_TTS_BETA_TRACE=1`` for an env-gated per-patch +rms/max trace across the full side path (zero overhead when unset). + +Reference implementations: ``vllm_omni/model_executor/models/voxcpm2/`` +(architectural template: vLLM-native base LM + side path + per-request +state) and ``ming_flash_omni`` (AudioVAE weight-loading pattern). +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import os +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.qwen2 import Qwen2Model +from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import DiT +from vllm_omni.model_executor.models.dots_tts.dots_tts_patch_encoder import ( + VAESemanticEncoder, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_speaker_encoder import ( + SpeakerXVectorFeatures, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import ( + AudioVAE, + AudioVAEConfig, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +def _require_config_block(hf_config: Any, name: str) -> dict[str, Any]: + """Fetch a sub-config dict (``vocoder`` / ``DiT`` / ``PatchEncoder``) + from the checkpoint's config.json, failing loudly when absent — + silently falling back to another checkpoint's values would produce + garbage audio, not an error.""" + block = getattr(hf_config, name, None) + if not isinstance(block, dict) or not block: + raise ValueError( + f"dots.tts checkpoint config.json has no '{name}' block — cannot " + "build the audio side path. Expected the upstream dots.tts config " + "layout (e.g. rednote-hilab/dots.tts-soar)." + ) + return block + + +def _build_audio_vae_config(hf_config: Any) -> AudioVAEConfig: + # AudioVAEConfig declares exactly the checkpoint's 17 ``vocoder`` keys, + # so an unknown / missing key raises TypeError instead of guessing. + return AudioVAEConfig(**_require_config_block(hf_config, "vocoder")) + + +# Architecture constants used by the per-request FM workspace math. All +# released dots.tts checkpoints (soar / base / mf) share these values; +# _validate_architecture_constants refuses a checkpoint that disagrees +# instead of silently generating garbage. +_FM_HIDDEN = 1024 # DiT.hidden_size (also LLM↔DiT projection target) +_LATENT_DIM = 128 # AudioVAE.latent_dim +_LATENT_PATCH_SIZE = 4 # config.patch_size — DiT samples 4 latent frames / audio patch +_HIDDEN_PATCH_SIZE = 1 # upstream core.py hardcodes self.hidden_patch_size = 1 +_MAX_AUDIO_PATCHES = 1024 # ~164 s @ 4×1920/48k; bounds per-request FM static buffer +_DIT_NUM_STEPS = int( + os.environ.get("DOTS_TTS_DIT_NUM_STEPS", "10") +) # env-gated; upstream default 10 for fixed-step Euler +_DIT_GUIDANCE_SCALE = 1.2 # upstream default guidance_scale for soar +_DIT_NOISE_SEED = 20260601 # base seed for per-request FM noise (voxcpm2 parity) +_PATCH_ENCODER_OUT_DS_RATE = 2 # patch_size / in_ds_rate = 4 / 2 (VAESemanticEncoder hardcodes in_ds_rate=2) + + +class _IOHelper: + """Latent stats wrapper. Holds (mean, var) over the 128-dim VAE latent + space and normalizes / denormalizes between DiT-internal space (which + operates on normalized latents) and AudioVAE / patch_encoder input space + (which expects raw latents). + + Mirrors upstream IOHelper (rednote-hilab/dots.tts @ a393d2e + models/dots_tts/core.py:684). Plain class (not nn.Module) so it stays + out of state_dict() and load_weights() doesn't have to mute its buffers. + Stats are loaded from latent_stats.pt — independent of safetensors. + """ + + def __init__(self, latent_stats_path: str | None = None) -> None: + if latent_stats_path is None: + self.global_mean: torch.Tensor | None = None + self.global_var: torch.Tensor | None = None + return + stats = torch.load(latent_stats_path, weights_only=False) + self.global_mean = torch.as_tensor(stats["mean"]) + self.global_var = torch.as_tensor(stats["var"]) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + if self.global_mean is None or self.global_var is None: + return x + mean = self.global_mean.to(device=x.device, dtype=x.dtype) + var = self.global_var.to(device=x.device, dtype=x.dtype) + return (x - mean) / torch.sqrt(var) + + def denormalize(self, x: torch.Tensor) -> torch.Tensor: + if self.global_mean is None or self.global_var is None: + return x + mean = self.global_mean.to(device=x.device, dtype=x.dtype) + var = self.global_var.to(device=x.device, dtype=x.dtype) + return x * torch.sqrt(var) + mean + + +# Per-request state container for AR-loop continuity. All cross-step +# side-path state lives here, keyed by request_id, so concurrent requests +# stay fully isolated (the voxcpm2 _RequestState pattern). +@dataclass +class _RequestState: + request_id: str + # AR loopback: patch_encoder's previous decode output, fed to the LM + # as next-step inputs_embeds. Shape: [1, llm_hidden] = [1, 1536]. + # (patch_encoder collapses its internal _PATCH_ENCODER_OUT_DS_RATE + # positions into one LLM token via out_proj; see + # VAESemanticEncoder._project_embeddings.) _finish_decode writes a + # fresh tensor every call. + curr_embed_for_next: torch.Tensor | None = None + # Stop-signal cache: _finish_decode populates precomputed_stop_logits; + # compute_logits drains it. + precomputed_stop_logits: torch.Tensor | None = None + is_stopping: bool = False + prefill_completed: bool = False + # Per-request FM noise counter: draw #n of this request hashes to a + # deterministic Generator seed (see _run_dit_n_step_euler), so outputs + # are reproducible run-to-run and concurrent requests cannot perturb + # each other's noise streams (voxcpm2 _fill_deterministic_cfm_noise). + noise_step: int = 0 + # Per-request FM static workspace (lazy-allocated by + # _initialize_request_fm_state on first _finish_decode call). Sized for + # _MAX_AUDIO_PATCHES × (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) = 1024 × 5 + # = 5120 positions. + fm_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] + fm_cfg_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] + fm_null_g_cond: torch.Tensor | None = None # [1, _FM_HIDDEN] + fm_seq_len: int = 0 + fm_capacity: int = 0 + # Speaker x-vector after _xvec_proj. Zero-shot generation leaves this + # None → DiT falls back to fm_null_g_cond; voice-clone wiring is a + # follow-up. + g_cond: torch.Tensor | None = None + # AR-loop bookkeeping. + # patch_encoder_state holds conv_tail + per-layer KV caches for the + # patch_encoder's streaming decode (mirrors upstream state.patch_encoder_state). + # Lazy-allocated by _run_patch_encoder_loopback on first audio step. + patch_encoder_state: Any | None = None + # Review M2: streaming vocoder state (BigVGANStreamState — LSTM hidden + # + decoder sliding latent window). Gives each patch real left conv + # context; output lags input by decoder.stream_lookahead (2 frames = + # 3840 samples). Lazy-allocated by _run_vocoder_stream_step. + vocoder_stream_state: Any | None = None + + +# DiT transformer hyperparameters from rednote-hilab/dots.tts-soar config.json +# (the ``DiT`` block). Mirrors the shape of an upstream pydantic config: DiT's +# constructor calls ``transformer_config.to_dict()`` and reads ``.hidden_size`` +# / ``.num_layers``, so this @dataclass plus the ``to_dict`` helper satisfies +# that minimal interface without dragging in upstream's config system. +# +# Note: ``attn_dropout`` in the JSON spells differently from MultiHeadAttention's +# ``attn_drop`` kwarg — the value gets absorbed by ``**_kwargs`` and the MHA +# default 0.0 applies. Same numeric outcome for the soar checkpoint (both 0.0), +# so this is benign; flagging it because it would mask non-default values if +# upstream ever raised this field. +@dataclass +class _DiTConfig: + num_layers: int = 18 + num_heads: int = 16 + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + modulation: bool = True + qkv_bias: bool = False + qk_norm: bool = True + attn_dropout: float = 0.0 + dropout: float = 0.0 + norm_layer: str = "RMSNorm" + alibi_bias: bool = False + rotary_bias: bool = True + rotary_theta: float = 10000.0 + + def to_dict(self) -> dict[str, Any]: + return dataclasses.asdict(self) + + +def _build_dit_config(hf_config: Any) -> _DiTConfig: + # _DiTConfig declares exactly the checkpoint's 13 ``DiT`` keys, so an + # unknown key raises TypeError instead of being silently dropped. + return _DiTConfig(**_require_config_block(hf_config, "DiT")) + + +# Patch encoder hyperparameters from rednote-hilab/dots.tts-soar config.json +# (the top-level ``patch_size`` and the ``PatchEncoder`` sub-block). +# +# Two layers of config: +# * Outer (``_PatchEncoderConfig``): VAESemanticEncoder reads ``.patch_size`` +# and forwards ``.PatchEncoder`` to its inner SuperviseEncoder. +# * Inner (``_PatchEncoderInner``): the constructor uses BOTH attribute +# access (``config.PatchEncoder.hidden_size``) and dict-style +# ``config.get("hidden_size", 1024)``. Upstream's pydantic ConfigBase +# supports both natively; we add a tiny ``.get()`` helper to a dataclass. +# +# Upstream's ``_EncoderConfig`` also defines qk_norm / rotary_bias / +# rotary_theta / input_dim, but ``SuperviseEncoder.__init__`` only reads six +# keys (num_layers / num_heads / hidden_size / ffn_hidden_size / norm_layer / +# causal). The rest are swallowed and MultiHeadAttention falls back to its +# own defaults — matching upstream behaviour means *not* plumbing them through. +@dataclass +class _PatchEncoderInner: + num_layers: int = 24 + num_heads: int = 16 + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + norm_layer: str = "RMSNorm" + causal: bool = True + + def get(self, key: str, default: Any = None) -> Any: + return getattr(self, key, default) + + +@dataclass +class _PatchEncoderConfig: + patch_size: int = 4 + PatchEncoder: _PatchEncoderInner = dataclasses.field(default_factory=_PatchEncoderInner) + + +def _build_patch_encoder_config(hf_config: Any) -> _PatchEncoderConfig: + block = _require_config_block(hf_config, "PatchEncoder") + # SuperviseEncoder reads only the six declared fields; upstream's + # pydantic config swallows the other checkpoint keys, so filtering + # here matches upstream behaviour exactly. + inner_fields = {f.name for f in dataclasses.fields(_PatchEncoderInner)} + inner = _PatchEncoderInner(**{k: v for k, v in block.items() if k in inner_fields}) + return _PatchEncoderConfig( + patch_size=int(hf_config.patch_size), + PatchEncoder=inner, + ) + + +def _validate_architecture_constants(hf_config: Any, dit_config: _DiTConfig) -> None: + """The FM workspace helpers size buffers off module-level constants; + a checkpoint that disagrees would corrupt the DiT conditioning, so + refuse it up front (making these fully config-driven is deferred + until a released checkpoint actually differs).""" + checks = [ + ("latent_dim", getattr(hf_config, "latent_dim", None), _LATENT_DIM), + ("patch_size", getattr(hf_config, "patch_size", None), _LATENT_PATCH_SIZE), + ("DiT.hidden_size", dit_config.hidden_size, _FM_HIDDEN), + ] + mismatches = [f"{name}={got!r} (expected {want})" for name, got, want in checks if got != want] + if mismatches: + raise ValueError( + "dots.tts checkpoint config disagrees with this integration's " + "architecture constants: " + ", ".join(mismatches) + ) + + +class DotsTTSForConditionalGeneration(nn.Module): + """dots.tts AR talker: Qwen2.5 base LM + DiT / AudioVAE / CAM++ side path.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + + # vllm-omni capability flags — gpu_model_runner checks these to decide + # whether to call preprocess() / postprocess() / accept multimodal + # outputs. Without these, the runner bypasses our preprocess hook + # and _pending_requests stays empty (forward never dispatches the + # side-path). Voxcpm2 sets the same three. + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + + self.model = Qwen2Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Side-path module hyperparameters come off the checkpoint's + # config.json blocks (fail loudly if a block is absent); the FM + # workspace constants are cross-checked against the same config. + dit_config = _build_dit_config(self.config) + _validate_architecture_constants(self.config, dit_config) + llm_hidden = self.config.hidden_size + fm_hidden = dit_config.hidden_size + latent_dim = int(self.config.latent_dim) + xvec_dim = int(getattr(self.config, "campplus_embedding_size", 512)) + + self._audio_vae = AudioVAE(_build_audio_vae_config(self.config)) + # Upstream serializes vocoder.safetensors with weight_norm folded into + # plain `weight` on the decoder (encoder kept weight_norm — it's not in + # the synthesis hot path). Match that layout before load_weights runs, + # so checkpoint keys align with our state_dict 1:1. + self._audio_vae.remove_weight_norm() + # AudioVAE / speaker encoder stay fp32 — upstream only casts + # `model.core` to bf16 (runtime.py:81), leaving vocoder + xvector + # extractor in fp32 because their kernels (conv1d / fbank) need + # higher precision. Without this _audio_vae would inherit talker's + # bf16 dtype via vLLM's auto-cast and dtype-mismatch under + # stream_step (input .float() vs bf16 weights). + + # DiT flow-matching head. Dimensions per upstream core.py:101-104: + # in_dim = config.DiT.hidden_size (DiT internal space) + # out_dim = config.latent_dim (AudioVAE input space) + # mode is "flow_matching" for soar/base; "meanflow" only for the mf + # checkpoint (handled by a separate factory in a later step). + self._head = DiT( + in_dim=fm_hidden, + out_dim=latent_dim, + transformer_config=dit_config, + mode="flow_matching", + ) + + # Patch encoder: closes the AR loop by mapping each DiT-emitted audio + # latent patch back into the LLM hidden space, so Qwen2 "hears" what it + # has already generated. Upstream core.py:75 calls this with + # in_dim=latent_dim and out_dim=llm.config.hidden_size. + self._patch_encoder = VAESemanticEncoder( + in_dim=latent_dim, + out_dim=llm_hidden, + config=_build_patch_encoder_config(self.config), + ) + + # Five thin projectors (upstream core.py:82-113). Soar dimensions: + # llm_hidden=1536, fm_hidden=1024, latent_dim=128, xvec_dim=512. + # These are kept as direct attributes (no vendor file) — they are flat + # nn.Linear / nn.Sequential, and live at the top of the checkpoint + # namespace under their own short prefixes (see load_weights). + self._hidden_proj = nn.Linear(llm_hidden, fm_hidden) + self._latent_proj = nn.Linear(latent_dim, fm_hidden) + self._coordinate_proj = nn.Linear(latent_dim, fm_hidden) + self._xvec_proj = nn.Sequential( + nn.Linear(xvec_dim, fm_hidden), + nn.LayerNorm(fm_hidden), + ) + # eos_proj is the stop predictor — feeds LLM hidden through a 2-layer + # MLP and produces a 2-way logit (continue / stop). + self._eos_proj = nn.Sequential( + nn.Linear(llm_hidden, llm_hidden), + nn.SiLU(), + nn.Linear(llm_hidden, 2), + ) + + # CAM++ speaker encoder (3D-Speaker), produces the x-vector from + # reference audio for voice cloning. sample_rate matches AudioVAE's + # input space (48 kHz); the wrapper resamples to 16k internally for + # CAM++. Weights are frozen (upstream does `requires_grad = False`). + self._speaker_encoder = SpeakerXVectorFeatures( + sample_rate=self._audio_vae.sample_rate, + campplus_embedding_size=xvec_dim, + max_audio_seconds=float(getattr(self.config, "xvec_max_audio_seconds", 10.0)), + ) + + # Pin AudioVAE + speaker encoder to fp32 (mirror upstream + # runtime.py:81 where only `model.core` is cast to bf16). vLLM's + # auto-cast would otherwise lower these to bf16, breaking + # stream_step (input .float() vs bf16 weights) and the fbank + # pipeline. + self._audio_vae.float() + self._speaker_encoder.float() + + # Latent stats — independent of safetensors (torch.load on + # latent_stats.pt sitting next to the checkpoint files). Used by + # _finish_decode to bridge between DiT's normalized latent space and + # the raw latent space expected by patch_encoder + AudioVAE. + latent_stats_path = self._resolve_latent_stats_path(vllm_config.model_config.model) + if latent_stats_path is not None: + self._io_helper = _IOHelper(latent_stats_path) + else: + logger.warning( + "latent_stats.pt not resolvable for %s; running with " + "identity normalize/denormalize. Wav will be ~10x quiet " + "and distorted (no sqrt(var) denormalize).", + vllm_config.model_config.model, + ) + self._io_helper = _IOHelper() + + # (β) diagnostic: env-gated per-patch rms/max trace across the full + # side path (LLM hidden → DiT → denormalize → patch_encoder → wav → + # eos). Zero overhead when DOTS_TTS_BETA_TRACE is unset (single + # bool check per step); kept as the model's debugging instrument + # (cf. voxcpm2's enable_profiling runtime knob). + self._beta_trace = bool(os.environ.get("DOTS_TTS_BETA_TRACE")) + + # Per-request state plumbing. + # _active_states: per-request state keyed by request_id (created in + # preprocess prefill, evicted by _flush_deferred_cleanup). + # _pending_requests: per-step list of (req_id, is_prefill, embeds, + # span_len) populated by preprocess(), consumed by forward(). + # _deferred_cleanup_ids: requests vLLM has marked finished but whose + # side-path audio still needs draining in the current forward step; + # evicted from _active_states at the end of forward(). + self._active_states: dict[str, _RequestState] = {} + self._pending_requests: list[tuple[str, bool, torch.Tensor, int]] = [] + self._deferred_cleanup_ids: set[str] = set() + + # Audio side-channel: + # _audio_queue: _finish_decode pushes (req_id, audio_tensor) tuples + # here; make_omni_output drains them into multimodal_outputs each + # step. + # _sample_rate: cached from AudioVAE config (= 48000 for soar). + self._audio_queue: list[tuple[str, torch.Tensor]] = [] + self._sample_rate = self._audio_vae.sample_rate + + # Stop-signal side-channel: + # _results_queue: _finish_decode pushes (req_id, stop_logits) tuples + # here; compute_logits drains them into the [bsz, vocab_size] + # logits tensor each step. + self._results_queue: list[tuple[str, torch.Tensor | None]] = [] + + logger.info( + "DotsTTS talker built (base_lm=Qwen2[28L,12H,1536d], " + "audio_vae=AudioVAE, dit=DiT[18L,16H,1024d], " + "patch_encoder=VAESemanticEncoder[24L,16H,1024d], " + "5 projectors, speaker_encoder=CAM++[7.2M], " + "io_helper=%s, model=%s, dit_steps=%d); " + "_finish_decode: DiT → denormalize → patch_encoder AR + " + "streaming AudioVAE decode → wav + eos_proj → stop signal.", + "loaded" if self._io_helper.global_mean is not None else "identity", + vllm_config.model_config.model, + _DIT_NUM_STEPS, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: + """Drive the base LM, then the per-request audio side path. + + One vLLM step: base LM forward under PagedAttention, then slice + the scaffold hidden states per request (in _pending_requests + order) and run _finish_decode on each — DiT sampling, streaming + AudioVAE decode, AR loopback, and stop signal all happen there. + vLLM only ever sees the returned hidden states; audio leaves + through the _audio_queue side channel. + """ + output = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + if isinstance(output, IntermediateTensors): + return output + if isinstance(output, tuple): + output = output[0] + + # Per-request side-path: slice scaffold_hidden by token span_len + # in the order preprocess() appended to _pending_requests. + token_offset = 0 + for req_id, is_prefill, _embeds, span_len in self._pending_requests: + req_hidden = output[token_offset : token_offset + span_len] + token_offset += span_len + self._finish_decode(req_id, req_hidden, is_prefill) + + # End-of-step cleanup + self._pending_requests.clear() + self._flush_deferred_cleanup() + + return output + + # ── vllm-omni protocol methods ── + + @staticmethod + def _resolve_latent_stats_path(model_arg: str) -> str | None: + """Resolve latent_stats.pt to a local filesystem path. + + ``vllm_config.model_config.model`` is typically an HF repo ID like + ``"rednote-hilab/dots.tts-soar"`` — joining it with ``"latent_ + stats.pt"`` gives a non-existent path. Try the local-dir form + first (covers ``--model /local/path``), then fall back to HF + cache lookup with ``local_files_only=True`` so we never trigger + a download from this side (the file was fetched alongside the + safetensors at load time). + """ + local = os.path.join(model_arg, "latent_stats.pt") + if os.path.exists(local): + return local + try: + from huggingface_hub import hf_hub_download + + return hf_hub_download( + repo_id=model_arg, + filename="latent_stats.pt", + local_files_only=True, + ) + except Exception: + return None + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Map token ids to embeddings via the base LM's embed_tokens. + + Required by vLLM's ``is_text_generation_model`` Protocol check + (interfaces_base._check_vllm_model_embed_input_ids). Without this + the model fails ``--runner generate`` validation at engine startup. + """ + return self.model.embed_tokens(input_ids) + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None = None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Per-step embedding + per-request state registration. + + Prefill: embed the wrapped prompt + ``[文本][文本对应语音]`` via the LM's + embed_tokens. Caller (serving layer or test harness) is + responsible for the wrap — talker just embeds whatever it gets. + + Decode (span_len == 1): use ``state.curr_embed_for_next`` from + the previous step's AR loopback (DiT → patch_encoder output); + zero fallback only if the loopback hasn't produced one yet. + + Voice clone is intentionally not wired here — zero-shot only. + + Returns ``(input_ids unchanged, embeds [span_len, hidden_size], + {})``. ``forward()`` consumes ``self._pending_requests`` to + slice scaffold_hidden per request. + """ + # voxcpm2 protocol compatibility: flatten additional_information. + additional = info_dict.get("additional_information") + if isinstance(additional, dict): + merged = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional.items(): + merged.setdefault(k, v) + info_dict = merged + + span_len = int(input_ids.shape[0]) + dev = input_ids.device + req_id = info_dict.get("request_id", "default") + state = self._active_states.get(req_id) + # Prefill is detected by state lifecycle (not span_len), so decode + # can use span_len=1 with a multi-position curr_embed_for_next. + is_prefill = state is None or not state.prefill_completed + + if is_prefill: + if state is None: + state = _RequestState(request_id=req_id) + self._active_states[req_id] = state + # Reset per-step state on every prefill — previous run is stale. + state.curr_embed_for_next = None + state.precomputed_stop_logits = None + state.is_stopping = False + state.prefill_completed = False + state.noise_step = 0 + # Reset AR-loop state on every prefill. + state.patch_encoder_state = None + state.vocoder_stream_state = None + embeds = self.model.embed_tokens(input_ids) + else: + curr = state.curr_embed_for_next + if curr is not None: + embeds = curr.to(dev) # [1, llm_hidden] + else: + # First decode step before AR loop closed — zero fallback. + embed_dtype = self.model.embed_tokens.weight.dtype + embeds = torch.zeros( + 1, + self.config.hidden_size, + device=dev, + dtype=embed_dtype, + ) + + self._pending_requests.append((req_id, is_prefill, embeds, span_len)) + return input_ids, embeds, {} + + def postprocess(self, *args: Any, **kwargs: Any) -> dict: + return {} + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Mark finished requests for deferred eviction. + + vLLM scheduler calls this BEFORE forward() to notify that certain + requests have completed. We can't drop their _RequestState yet — + the current forward step still needs to drain side-path audio + chunks. Real eviction happens at the end of forward() via + _flush_deferred_cleanup. + """ + for req_id in finished_req_ids: + if req_id in self._active_states: + self._deferred_cleanup_ids.add(req_id) + + def _flush_deferred_cleanup(self) -> None: + """Evict states marked by on_requests_finished. Called at the + tail of forward().""" + for req_id in self._deferred_cleanup_ids: + self._active_states.pop(req_id, None) + self._deferred_cleanup_ids.clear() + + def _finish_decode( + self, + req_id: str, + req_hidden: torch.Tensor, + is_prefill: bool, + ) -> None: + """One request's per-step side path: DiT → AR loopback → wav → stop. + + One vLLM step = one audio patch. patch_encoder.decode_patch + collapses its internal _PATCH_ENCODER_OUT_DS_RATE positions into + a single LLM input embedding (see VAESemanticEncoder._project_ + embeddings), so vLLM's span_len=1 decode matches naturally — no + slot bookkeeping needed. + + Per-call flow (prefill or decode): + 1. _hidden_proj(last_hidden) → fm_sequence (+1 position) + 2. DiT N-step Euler → audio_patch [1, 4, 128] (normalized) + 3. _latent_proj(audio_patch) → fm_sequence (+4 positions, normalized) + 4. io_helper.denormalize(audio_patch) → audio_patch_raw + 5. patch_encoder(audio_patch_raw) → next-step embed [1, 1536] + 6. AudioVAE stream_step(audio_patch_raw) → wav chunk pushed to + _audio_queue (prefill's patch #0 included; output lags 2 + frames behind input) + 7. (decode only) eos_proj(last_hidden) → stop_logits pushed to + _results_queue; threshold > 0.8 flips state.is_stopping and + stream_flush drains the decoder's lookahead tail. + """ + state = self._active_states.get(req_id) + if state is None: + return + + if state.fm_sequence is None: + self._initialize_request_fm_state( + state, + device=req_hidden.device, + dtype=req_hidden.dtype, + ) + + # This step's append would overflow the per-request FM buffer. + # Stop gracefully (same path as a model-decided stop below) instead + # of raising past _append_hidden_chunk/_append_history_chunk — + # those raises are engine-fatal, not request-fatal, and would take + # every other in-flight and future request down with this one. + if state.fm_seq_len + _HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE > state.fm_capacity: + state.is_stopping = True + stop_logits = torch.tensor([[0.0, 1.0]], device=req_hidden.device, dtype=req_hidden.dtype) + state.precomputed_stop_logits = stop_logits + tail = self._audio_vae.stream_flush(state.vocoder_stream_state) + if tail.size(-1) > 0: + self._audio_queue.append((req_id, tail.reshape(-1))) + self._results_queue.append((req_id, stop_logits)) + return + + # 1. Append last LLM hidden to fm_sequence (+1 position). + last_hidden = req_hidden[-_HIDDEN_PATCH_SIZE:].unsqueeze(0) + + self._append_hidden_chunk(state, last_hidden) + + # 2. DiT N-step Euler → audio latent patch [1, 4, 128] (normalized). + audio_patch = self._run_dit_n_step_euler(state) + + # 3. Append latent to fm_sequence (+4 positions). Stays in + # normalized space — DiT's KV cache lives in that space. + self._append_history_chunk(state, audio_patch) + + # 4. Denormalize for downstream consumers (patch_encoder, AudioVAE). + audio_patch_raw = self._io_helper.denormalize(audio_patch) + + # 5. patch_encoder AR loopback → next-step inputs_embeds [1, 1536]. + next_embeds = self._run_patch_encoder_loopback(state, audio_patch_raw) + state.curr_embed_for_next = next_embeds.squeeze(0).detach() + + # 6. Streaming AudioVAE decode → wav chunk. Runs on prefill too: + # upstream zero-shot emits every DiT patch including the first + # (_generate_latents_stream drops one patch only under prompt + # prefill / voice clone), and our prefill patch #0 corresponds + # to upstream's first decode-loop patch. + wav = self._run_vocoder_stream_step(state, audio_patch_raw) + if wav.size(-1) > 0: + self._audio_queue.append((req_id, wav.reshape(-1))) + + if is_prefill: + if self._beta_trace: + self._beta_trace_log( + state, + is_prefill=True, + last_hidden=last_hidden, + audio_patch=audio_patch, + audio_patch_raw=audio_patch_raw, + next_embeds=next_embeds, + wav=wav if wav.size(-1) > 0 else None, + prob_stop=None, + ) + state.prefill_completed = True + # Keep len(_results_queue) == bsz: compute_logits pairs queue + # entries with batch rows positionally, so every request in the + # batch must push exactly one entry per step (voxcpm2_talker.py + # prefill placeholder pattern). None → forced continue via the + # stop_logits-is-None branch in compute_logits. + self._results_queue.append((req_id, None)) + return + + # 7. eos_proj → stop probability. Mirrors upstream + # model.py:1681 — softmax([continue, stop]) on detached LLM + # hidden, threshold the stop slot at 0.8 (upstream default + # eos_threshold). When stop fires for this step, compute_logits + # drains the high-stop logit on the next dispatch and vLLM + # terminates the request. + eos_logits = self._eos_proj(last_hidden.detach()).softmax(dim=-1) + stop_logits = eos_logits.squeeze(0) + state.precomputed_stop_logits = stop_logits + if eos_logits[0, -1, 1].item() > 0.8: + state.is_stopping = True + # Upstream flushes after its decode loop (model.py:1898); our + # request ends when vLLM sees this step's stop token, so drain + # the decoder's 2-frame lookahead tail in the same step. + # Requests cut by max_tokens/abort never reach this path and + # lose the 40 ms tail — no engine hook fires late enough to + # flush after the final step. + tail = self._audio_vae.stream_flush(state.vocoder_stream_state) + if tail.size(-1) > 0: + self._audio_queue.append((req_id, tail.reshape(-1))) + self._results_queue.append((req_id, stop_logits)) + + if self._beta_trace: + self._beta_trace_log( + state, + is_prefill=False, + last_hidden=last_hidden, + audio_patch=audio_patch, + audio_patch_raw=audio_patch_raw, + next_embeds=next_embeds, + wav=wav if wav.size(-1) > 0 else None, + prob_stop=eos_logits[0, -1, 1].item(), + ) + + def _beta_trace_log( + self, + state: _RequestState, + *, + is_prefill: bool, + last_hidden: torch.Tensor, + audio_patch: torch.Tensor, + audio_patch_raw: torch.Tensor, + next_embeds: torch.Tensor, + wav: torch.Tensor | None, + prob_stop: float | None, + ) -> None: + """(β) diagnostic — patch label matches output.wav indexing in + decode rows; prefill row labels itself ``pre``.""" + + def rmsmax(t: torch.Tensor) -> tuple[float, float]: + f = t.detach().float() + return f.pow(2).mean().sqrt().item(), f.abs().max().item() + + patch_label = "pre" if is_prefill else f"{state.fm_seq_len // 5 - 2:2d}" + req_tail = state.request_id[-8:] + hid_rms, hid_max = rmsmax(last_hidden) + dit_rms, dit_max = rmsmax(audio_patch) + raw_rms, raw_max = rmsmax(audio_patch_raw) + nxt_rms, nxt_max = rmsmax(next_embeds) + if wav is not None: + w_rms, w_max = rmsmax(wav) + wav_part = f"wav:rms={w_rms:.3f} max={w_max:.3f}" + else: + wav_part = "wav:rms= - max= - " + stop_part = f"stop={prob_stop:.3f}" if prob_stop is not None else "stop= - " + logger.info( + f"[β req={req_tail} patch={patch_label}] " + f"hid:rms={hid_rms:.3f} max={hid_max:.3f} | " + f"dit:rms={dit_rms:.3f} max={dit_max:.3f} | " + f"raw:rms={raw_rms:.3f} max={raw_max:.3f} | " + f"next:rms={nxt_rms:.3f} max={nxt_max:.3f} | " + f"{wav_part} | {stop_part}" + ) + + # ── FM helpers (per-request workspace + DiT N-step Euler) ── + + def _initialize_request_fm_state( + self, + state: _RequestState, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + """Lazy-allocate per-request FM workspace. + + Mirrors upstream _allocate_generate_state (model.py:413) but per- + request (no workspace sharing) so concurrent requests stay + isolated. + """ + fm_capacity = _MAX_AUDIO_PATCHES * (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) + state.fm_capacity = fm_capacity + state.fm_sequence = torch.zeros( + (1, fm_capacity, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_cfg_sequence = torch.zeros( + (1, fm_capacity, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_null_g_cond = torch.zeros( + (1, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_seq_len = 0 + + def _append_hidden_chunk( + self, + state: _RequestState, + hidden_chunk: torch.Tensor, + ) -> None: + """Project last _HIDDEN_PATCH_SIZE LLM hiddens via _hidden_proj into fm_sequence. + + Mirrors upstream _append_hidden_chunk (model.py:1308): cfg buffer + gets the null-projected counterpart (``_hidden_proj`` of zero) so + CFG dropout sees no past hidden, only past audio history. + """ + last_hidden = hidden_chunk[:, -_HIDDEN_PATCH_SIZE:, :] + projected = self._hidden_proj(last_hidden) + null_projected = self._hidden_proj(torch.zeros_like(last_hidden)) + end = state.fm_seq_len + projected.size(1) + if end > state.fm_capacity: + raise RuntimeError(f"FM buffer overflow on hidden append: end={end} cap={state.fm_capacity}") + state.fm_sequence[:, state.fm_seq_len : end].copy_(projected.to(state.fm_sequence.dtype)) + state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(null_projected.to(state.fm_cfg_sequence.dtype)) + state.fm_seq_len = end + + def _append_history_chunk( + self, + state: _RequestState, + latent_chunk: torch.Tensor, + ) -> None: + """Project audio latent via _latent_proj into fm_sequence (history append). + + Mirrors upstream _append_history_chunk (model.py:1325): cfg + buffer gets the same value (latent history is shared between + conditional and unconditional CFG paths — only the hidden chunk + gets nulled). + """ + history_latent = self._latent_proj(latent_chunk) + end = state.fm_seq_len + history_latent.size(1) + if end > state.fm_capacity: + raise RuntimeError(f"FM buffer overflow on latent append: end={end} cap={state.fm_capacity}") + state.fm_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_sequence.dtype)) + state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_cfg_sequence.dtype)) + state.fm_seq_len = end + + def _build_fm_attn_mask( + self, + state: _RequestState, + total_len: int, + device: torch.device, + ) -> torch.Tensor: + """Build DiT attention mask for one decode step. + + Mirrors upstream _build_fm_attn_mask (model.py:1203). Layout: + * [0, fm_seq_len) — committed history. Causal over the past + up to a non-causal tail block of size _HIDDEN_PATCH_SIZE that + sees itself and the new noise patch (the just-appended hidden). + * [latent_start, total_len) — new noise patch, bidirectional + with the committed history and itself. + """ + attn_mask = torch.zeros((1, total_len, total_len), device=device, dtype=torch.bool) + latent_start = total_len - _LATENT_PATCH_SIZE + block_start = state.fm_seq_len - _HIDDEN_PATCH_SIZE + if block_start > 0: + causal_mask = torch.ones((block_start, block_start), device=device, dtype=torch.bool).triu(1).logical_not() + attn_mask[:, :block_start, :block_start] = causal_mask + attn_mask[:, block_start : state.fm_seq_len, : state.fm_seq_len] = True + attn_mask[:, block_start : state.fm_seq_len, latent_start:] = True + attn_mask[:, latent_start:, : state.fm_seq_len] = True + attn_mask[:, latent_start:, latent_start:] = True + if latent_start > state.fm_seq_len: + padding_indices = torch.arange(state.fm_seq_len, latent_start, device=device) + attn_mask[:, padding_indices, padding_indices] = True + return attn_mask + + def _build_fm_pos_ids( + self, + state: _RequestState, + total_len: int, + device: torch.device, + ) -> torch.Tensor: + """Build DiT position ids for one decode step. + + Mirrors upstream _build_fm_pos_ids (model.py:1236). Position ids + stay fp32 — DiT's RotaryEmbedding does its own dtype handling. + """ + pos_ids = torch.zeros((1, total_len), device=device, dtype=torch.float32) + latent_start = total_len - _LATENT_PATCH_SIZE + if state.fm_seq_len > 0: + pos_ids[:, : state.fm_seq_len] = torch.arange( + state.fm_seq_len, + device=device, + dtype=pos_ids.dtype, + ) + pos_ids[:, latent_start:] = torch.arange( + state.fm_seq_len, + state.fm_seq_len + _LATENT_PATCH_SIZE, + device=device, + dtype=pos_ids.dtype, + ) + return pos_ids + + def _run_dit_n_step_euler( + self, + state: _RequestState, + *, + num_steps: int = _DIT_NUM_STEPS, + guidance_scale: float = _DIT_GUIDANCE_SCALE, + ) -> torch.Tensor: + """DiT N-step Euler integration → audio latent patch [1, 4, 128]. + + Manual Euler loop (upstream uses torchdyn.odeint; we manualize to + avoid the dep and keep control flow explicit). Each step: + 1. z_proj = _coordinate_proj(z), placed into the last + _LATENT_PATCH_SIZE positions of input_sequence / cfg_sequence. + 2. Batch [conditional, unconditional] → DiT velocity field. + 3. CFG blend: v_c + guidance_scale * (v_c - v_u). + 4. Euler: z += v * dt. + + Mirrors upstream fm_solver_step (core.py:295) + _flow_matching_step_fm + (core.py:463), reduced to a fixed-step Euler integrator. + """ + assert state.fm_sequence is not None + assert state.fm_cfg_sequence is not None + assert state.fm_null_g_cond is not None + device = state.fm_sequence.device + dtype = state.fm_sequence.dtype + total_len = state.fm_seq_len + _LATENT_PATCH_SIZE + + # Workspace: committed history + 4 noise slots (overwritten per step). + input_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) + input_sequence[:, : state.fm_seq_len] = state.fm_sequence[:, : state.fm_seq_len] + cfg_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) + cfg_sequence[:, : state.fm_seq_len] = state.fm_cfg_sequence[:, : state.fm_seq_len] + + attn_mask = self._build_fm_attn_mask(state, total_len, device) + pos_ids = self._build_fm_pos_ids(state, total_len, device) + + g_cond = state.g_cond if state.g_cond is not None else state.fm_null_g_cond + g_cond = g_cond.to(device=device, dtype=dtype) + g_cond_batched = torch.cat([g_cond, torch.zeros_like(g_cond)], dim=0) + + latent_start = total_len - _LATENT_PATCH_SIZE + # fp32 ODE integration (same fix as Ming's CFM sampler, PR #4341): + # bf16's 7-bit mantissa loses the small per-step increments in the + # z += v*dt accumulation (hidden-state cos drift vs upstream past + # step ~7) and quantizes the linspace timesteps. Integration + # state stays fp32; DiT matmuls still run bf16 under the autocast + # block below, and the result is cast back to the sequence dtype + # on return. + # Deterministic per-request noise (voxcpm2 _fill_deterministic_cfm_ + # noise pattern): hash seed:request_key:draw# into a private + # Generator instead of the global CUDA RNG, so outputs reproduce + # run-to-run and concurrent requests cannot perturb each other's + # noise streams. request_key strips the engine's per-run "_" + # uuid suffix so replay across runs keys on the stable batch index. + request_key = state.request_id.split("_", 1)[0] + if not request_key.isdigit(): + request_key = state.request_id + noise_key = f"{_DIT_NOISE_SEED}:{request_key}:{state.noise_step}".encode() + digest = hashlib.blake2b(noise_key, digest_size=8).digest() + gen = torch.Generator(device=device) + gen.manual_seed(int.from_bytes(digest, "little") & 0x7FFF_FFFF_FFFF_FFFF) + z = torch.empty((1, _LATENT_PATCH_SIZE, _LATENT_DIM), device=device, dtype=torch.float32) + z.normal_(generator=gen) + state.noise_step += 1 + dt = 1.0 / num_steps + times = torch.linspace(0.0, 1.0, num_steps + 1, device=device, dtype=torch.float32) + + # Match upstream's autocast wrapper (model.py:280). DiT's + # TimestepEmbedder internally forces fp32 (`.float()` on freqs); + # without autocast the fp32 output hits a bf16 Linear → dtype + # mismatch. autocast bridges this for matmuls inside DiT. + use_amp = device.type == "cuda" and dtype in (torch.float16, torch.bfloat16) + with torch.autocast( + device_type=device.type if device.type == "cuda" else "cuda", + dtype=dtype if use_amp else torch.float32, + enabled=use_amp, + ): + for step in range(num_steps): + t = times[step].reshape(1) + z_proj = self._coordinate_proj(z) + z_c = input_sequence.clone() + z_c[:, latent_start:] = z_proj + z_u = cfg_sequence.clone() + z_u[:, latent_start:] = z_proj + z_batched = torch.cat([z_c, z_u], dim=0) + t_batched = t.repeat(2) + vt = self._head( + x=z_batched, + timesteps=t_batched, + attn_mask=attn_mask, + pos_ids=pos_ids, + g_cond=g_cond_batched, + ) + # Upcast DiT output before CFG blend + Euler update so the + # accumulation arithmetic stays fp32 (autocast only affects + # matmuls; these pointwise ops keep their input dtype). + vt = vt[:, latent_start:].float() + vt_c, vt_u = vt[0:1], vt[1:2] + velocity = vt_c + guidance_scale * (vt_c - vt_u) + z = z + velocity * dt + return z.to(dtype) + + def _run_patch_encoder_loopback( + self, + state: _RequestState, + audio_patch_raw: torch.Tensor, + ) -> torch.Tensor: + """Decode one audio latent patch → 1 LLM input embed. + + Mirrors upstream _consume_audio_patch (model.py:1557) — patch_encoder + side only. Lazy-allocates state.patch_encoder_state on first call. + Returns ``llm_embedding`` of shape ``[1, 1, llm_hidden]`` — the + _PATCH_ENCODER_OUT_DS_RATE encoder positions fold into the feature + axis before ``out_proj``, not into separate output positions. + + ``audio_patch_raw`` must already be denormalized (raw VAE latent + space). Caller (_finish_decode) runs io_helper.denormalize on + DiT's normalized output before passing in. + """ + if state.patch_encoder_state is None: + state.patch_encoder_state = self._patch_encoder.init_decode_state( + max_audio_patch_count=_MAX_AUDIO_PATCHES, + batch_size=1, + device=audio_patch_raw.device, + dtype=audio_patch_raw.dtype, + ) + pe_state = state.patch_encoder_state + positions = ( + torch.arange( + _PATCH_ENCODER_OUT_DS_RATE, + device=audio_patch_raw.device, + dtype=torch.long, + ) + + pe_state.seq_len + ) + llm_embedding, conv_tail = self._patch_encoder.decode_patch( + audio_patch_raw, + pe_state.conv_tail, + pe_state.layer_caches, + positions, + ) + pe_state.conv_tail.copy_(conv_tail) + pe_state.seq_len += _PATCH_ENCODER_OUT_DS_RATE + return llm_embedding + + def _run_vocoder_stream_step( + self, + state: _RequestState, + audio_patch_raw: torch.Tensor, + ) -> torch.Tensor: + """Streaming AudioVAE decode of one latent patch (review M2). + + Mirrors upstream generate_audio_stream (model.py:1881-1900): a + per-request BigVGANStreamState carries the decoder's LSTM hidden + and sliding latent window across steps, so each patch is decoded + with real left context instead of inference_from_latents' + zero-padded isolated boundary. Output lags input by + decoder.stream_lookahead (2 frames): the first call emits 2 frames + (3840 samples), steady state 4 (7680); the stop path in + _finish_decode drains the tail via stream_flush. + + audio_patch_raw: [1, T=4, D=128] denormalized latents. + Returns wav chunk [1, 1, n_samples] (n_samples may be 0). + """ + if state.vocoder_stream_state is None: + state.vocoder_stream_state = self._audio_vae.init_stream_state( + batch_size=1, + chunk_size=_LATENT_PATCH_SIZE, + ) + # stream_step expects [B, latent_dim, T]; the vocoder is pinned + # fp32 (__init__), so cast the (possibly bf16) latents up. + return self._audio_vae.stream_step( + audio_patch_raw.transpose(1, 2).float(), + state.vocoder_stream_state, + ) + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + """Drain self._audio_queue into the OmniOutput's multimodal_outputs. + + Output contract is **delta** (SKILL I1): each call emits only + audio produced since the previous step; the engine's + _consolidate_multimodal_tensors concatenates across steps for + offline consumers. + + Same-request multi-chunk: chunks for the same req_id within one + step are cat'd together. Multi-request: one entry per req_id + in multimodal_outputs["model_outputs"]. + + Note: ``sr`` is not a declared OmniPayload key but the engine / + serving consumer reads it (voxcpm2 convention). + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + + mm: dict[str, Any] = {} + if self._audio_queue: + audio_by_req: dict[str, torch.Tensor] = {} + for req_id, audio in self._audio_queue: + if audio is None: + continue + if req_id in audio_by_req: + audio_by_req[req_id] = torch.cat( + [audio_by_req[req_id].reshape(-1), audio.reshape(-1)], + dim=0, + ) + else: + audio_by_req[req_id] = audio + if audio_by_req: + sr = torch.tensor(self._sample_rate, dtype=torch.int32) + ready_req_ids = list(audio_by_req) + chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] + mm["model_outputs"] = chunks + mm["sr"] = [sr for _ in ready_req_ids] + # sparse_audio: ["1"] flags GPUARModelRunner to skip the default + # payload["hidden"] = scaffold_hidden injection (gpu_ar_model_ + # runner.py:1114). Otherwise prefill scaffold (24 hidden × 1536 + # = 29184 numbers observed) leaks into mm["audio"] via the + # output_processor "hidden" → target_key="audio" rename + # (output_processor.py:84), prepending ~0.61 s of noise. + mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} + self._audio_queue.clear() + else: + # Empty-audio step (prefill) still needs the sparse marker so the + # engine doesn't bleed scaffold_hidden into mm["audio"]. + mm["model_outputs"] = [] + mm["sr"] = [] + mm["meta"] = {"req_id": [], "sparse_audio": ["1"]} + + if self._beta_trace: + payload = mm.get("model_outputs") + if payload: + for i, t in enumerate(payload): + logger.info( + f"[β-emit] mm[model_outputs][{i}]: shape={tuple(t.shape)} " + f"dtype={t.dtype} samples={t.numel()} " + f"rms={t.float().pow(2).mean().sqrt().item():.4f} " + f"max={t.abs().max().item():.4f}" + ) + else: + logger.info("[β-emit] mm empty / no model_outputs") + mo_kind = type(model_outputs).__name__ + mo_info = f"shape={tuple(model_outputs.shape)}" if hasattr(model_outputs, "shape") else "no_shape" + logger.info(f"[β-emit] text_hidden_states arg: {mo_kind} {mo_info}") + + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + """Per-step stop signal logits (vLLM sampling protocol). + + Encodes continue/stop into ``logits[i, 0]`` / ``logits[i, 1]`` + (voxcpm2 convention — rest of vocab stays at -inf so the sampler + picks between the two slots). Source of stop logits is + ``self._results_queue``, populated by forward()'s _finish_decode. + When the queue is empty (preprocess-only steps), defaults to + all-continue. + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + + bsz = hidden_states.shape[0] + logits = torch.full( + (bsz, self.config.vocab_size), + float("-inf"), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + if self._results_queue: + for i, (req_id, stop_logits) in enumerate(self._results_queue): + if i >= bsz: + break + state = self._active_states.get(req_id) + if stop_logits is not None: + if state is not None and state.is_stopping: + logits[i, 0] = 0.0 + logits[i, 1] = 1.0 + else: + # is_stopping=False means prob_stop <= 0.8 — force + # continue. Feeding the raw softmax (continue, stop) + # pair as logits would let the greedy sampler compare + # them directly, silently lowering the effective stop + # threshold to 0.5 (early stop, swallowed endings). + logits[i, 0] = 1.0 + if state is not None: + state.precomputed_stop_logits = None + else: + # No stop signal pushed this step (prefill placeholder). + # Default to continue — stop is driven by eos_proj. + logits[i, 0] = 1.0 # continue + self._results_queue.clear() + else: + logits[:, 0] = 1.0 # all continue + + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load all dots.tts checkpoint weights — step 6b. + + Routing branches in a single pass over the input iterator (it can only + be consumed once): + + * **AudioVAE**: keys from ``vocoder.safetensors`` (no parent prefix). + Matched via set membership against ``self._audio_vae.state_dict()``; + acts as the final catch-all so it must run last. + * **DiT**: keys from ``model.safetensors`` under ``velocity_field_predictor.``. + * **patch_encoder**: keys from ``model.safetensors`` under ``patch_encoder.``. + * **Five projectors** (``hidden_proj`` / ``latent_proj`` / + ``coordinate_proj`` / ``xvec_proj`` / ``eos_proj``): each lives at + its own top-level prefix; dispatched through a shared table. + * **Qwen2 base LM**: keys under ``llm.model.``, handed off to vLLM's + Qwen2Model fuser (q/k/v_proj -> qkv_proj, gate/up_proj -> + gate_up_proj). ``llm.lm_head.*`` is silently dropped (soar uses + ``tie_word_embeddings=True`` and our ``self.model`` is a Qwen2Model + with no lm_head module). + * **CAM++ speaker encoder**: keys from ``speaker_encoder.safetensors`` + (no parent prefix, same convention as AudioVAE). Routed by set + membership against ``self._speaker_encoder.state_dict()`` (keys are + ``model.*`` / ``resample.kernel``, no distinguishing top-level + prefix). **Placed before the AudioVAE fallback** so its keys don't + fall through to the catch-all. + """ + vae_state_keys = set(self._audio_vae.state_dict().keys()) + dit_state_keys = set(self._head.state_dict().keys()) + patch_state_keys = set(self._patch_encoder.state_dict().keys()) + speaker_state_keys = set(self._speaker_encoder.state_dict().keys()) + + # (top-level prefix in checkpoint, local attr name, local nn.Module) + projector_specs: list[tuple[str, str, nn.Module]] = [ + ("hidden_proj.", "_hidden_proj", self._hidden_proj), + ("latent_proj.", "_latent_proj", self._latent_proj), + ("coordinate_proj.", "_coordinate_proj", self._coordinate_proj), + ("xvec_proj.", "_xvec_proj", self._xvec_proj), + ("eos_proj.", "_eos_proj", self._eos_proj), + ] + projector_state_keys = {prefix: set(mod.state_dict().keys()) for prefix, _, mod in projector_specs} + projector_matched: dict[str, list[tuple[str, torch.Tensor]]] = {prefix: [] for prefix, _, _ in projector_specs} + + matched_vae: list[tuple[str, torch.Tensor]] = [] + matched_dit: list[tuple[str, torch.Tensor]] = [] + matched_patch: list[tuple[str, torch.Tensor]] = [] + matched_llm: list[tuple[str, torch.Tensor]] = [] + matched_speaker: list[tuple[str, torch.Tensor]] = [] + skipped_lm_head = 0 + + DIT_PREFIX = "velocity_field_predictor." + PATCH_PREFIX = "patch_encoder." + LLM_MODEL_PREFIX = "llm.model." + LLM_LM_HEAD_PREFIX = "llm.lm_head." + for name, tensor in weights: + if name.startswith(DIT_PREFIX): + candidate = name[len(DIT_PREFIX) :] + if candidate in dit_state_keys: + matched_dit.append((candidate, tensor)) + continue + if name.startswith(PATCH_PREFIX): + candidate = name[len(PATCH_PREFIX) :] + if candidate in patch_state_keys: + matched_patch.append((candidate, tensor)) + continue + if name.startswith(LLM_LM_HEAD_PREFIX): + # Qwen2.5-1.5B-Base has tie_word_embeddings=True, and our + # ``self.model`` is a Qwen2Model (no lm_head submodule). Drop. + skipped_lm_head += 1 + continue + if name.startswith(LLM_MODEL_PREFIX): + candidate = name[len(LLM_MODEL_PREFIX) :] + # Hand off untouched — Qwen2Model.load_weights does its own + # fusion (q/k/v_proj -> qkv_proj, gate/up_proj -> gate_up_proj) + # and tolerates unknown keys, so no local key set to gate. + matched_llm.append((candidate, tensor)) + continue + # Projectors: 5 candidate top-level prefixes, first hit wins. + hit_projector = False + for prefix, _attr, _mod in projector_specs: + if name.startswith(prefix): + candidate = name[len(prefix) :] + if candidate in projector_state_keys[prefix]: + projector_matched[prefix].append((candidate, tensor)) + hit_projector = True + break + if hit_projector: + continue + # Speaker encoder: set-membership match (no distinguishing prefix + # in checkpoint; placed before the AudioVAE fallback so its keys + # don't get silently dropped by the catch-all). + if name in speaker_state_keys: + matched_speaker.append((name, tensor)) + continue + candidate = name[len("vocoder.") :] if name.startswith("vocoder.") else name + if candidate in vae_state_keys: + matched_vae.append((candidate, tensor)) + + loaded: set[str] = set() + + if matched_vae: + vae_loader = AutoWeightsLoader(self._audio_vae) + loaded_vae = vae_loader.load_weights(iter(matched_vae)) + loaded.update(f"_audio_vae.{name}" for name in loaded_vae) + logger.info( + "DotsTTS load_weights: loaded %d/%d AudioVAE tensors.", + len(loaded_vae), + len(vae_state_keys), + ) + + if matched_dit: + dit_loader = AutoWeightsLoader(self._head) + loaded_dit = dit_loader.load_weights(iter(matched_dit)) + loaded.update(f"_head.{name}" for name in loaded_dit) + logger.info( + "DotsTTS load_weights: loaded %d/%d DiT tensors.", + len(loaded_dit), + len(dit_state_keys), + ) + + if matched_patch: + patch_loader = AutoWeightsLoader(self._patch_encoder) + loaded_patch = patch_loader.load_weights(iter(matched_patch)) + loaded.update(f"_patch_encoder.{name}" for name in loaded_patch) + logger.info( + "DotsTTS load_weights: loaded %d/%d patch_encoder tensors.", + len(loaded_patch), + len(patch_state_keys), + ) + + any_projector_matched = False + for prefix, attr_name, mod in projector_specs: + matched = projector_matched[prefix] + if not matched: + continue + any_projector_matched = True + proj_loader = AutoWeightsLoader(mod) + loaded_proj = proj_loader.load_weights(iter(matched)) + loaded.update(f"{attr_name}.{name}" for name in loaded_proj) + logger.info( + "DotsTTS load_weights: loaded %d/%d %s tensors.", + len(loaded_proj), + len(projector_state_keys[prefix]), + prefix[:-1], + ) + + if matched_llm: + # vLLM's Qwen2Model.load_weights does its own fusion + tolerates + # unknown keys. It returns the set of param names it actually + # accepted; re-prefix with ``model.`` for our return value (our + # local attr is ``self.model``). + loaded_llm = self.model.load_weights(iter(matched_llm)) + loaded.update(f"model.{name}" for name in loaded_llm) + logger.info( + "DotsTTS load_weights: loaded %d Qwen2 tensors (fed %d; %d llm.lm_head.* skipped — tied embeddings).", + len(loaded_llm), + len(matched_llm), + skipped_lm_head, + ) + + if matched_speaker: + speaker_loader = AutoWeightsLoader(self._speaker_encoder) + loaded_speaker = speaker_loader.load_weights(iter(matched_speaker)) + loaded.update(f"_speaker_encoder.{name}" for name in loaded_speaker) + logger.info( + "DotsTTS load_weights: loaded %d/%d CAM++ speaker_encoder tensors.", + len(loaded_speaker), + len(speaker_state_keys), + ) + + if ( + not matched_vae + and not matched_dit + and not matched_patch + and not any_projector_matched + and not matched_llm + and not matched_speaker + ): + logger.warning( + "DotsTTS load_weights: no AudioVAE / DiT / " + "patch_encoder / projector / Qwen2 / speaker_encoder keys matched." + ) + + return loaded diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py new file mode 100644 index 000000000..fe17b5c0b --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py @@ -0,0 +1,1434 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules +# +# Merged from six upstream files (preserved as-is, inference behaviour +# unchanged; only internal imports were folded inline): +# - modules/vocoder/config.py -> AudioVAEConfig +# - modules/backbone/layers.py -> Conv1d, ConvTranspose1d +# (Mlp/MultiHeadAttention/RotaryEmbedding +# are not vendored here; they belong +# to the DiT side-path module added later) +# - modules/vocoder/alias_free_filter.py -> sinc, kaiser_sinc_filter1d, +# LowPassFilter1d +# - modules/vocoder/alias_free_resample.py-> UpSample1d, DownSample1d +# - modules/vocoder/alias_free_act.py -> Snake, SnakeBeta, Activation1d +# - modules/vocoder/bigvgan.py -> AudioVAE and all its sub-modules +# +# Per-section upstream attributions are preserved at each section header. +# AudioVAEConfig is the only piece that diverges from upstream: pydantic +# ConfigBase replaced with a stdlib @dataclass (extra fields not exercised by +# any shipped checkpoint). + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, field +from fractions import Fraction + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import pow, sin +from torch.nn import Parameter +from torch.nn.utils import remove_weight_norm, weight_norm + +# ============================================================================ +# Configuration (adapted from modules/vocoder/config.py) +# ---------------------------------------------------------------------------- +# Upstream subclasses a pydantic ``ConfigBase`` with ``extra="allow"``; both the +# dots.tts-soar and dots.tts-mf checkpoint ``config.json`` pass only the 17 +# declared fields below, so the escape hatch is never exercised. ``.get()`` +# helper preserves bigvgan.py's ``h.get(...)`` call sites without modification. +# ``num_decoder_lookahead`` / ``num_encoder_lookahead`` are looked up via +# ``h.get(...)`` but not present in any shipped config; they fall back to 2. +# ============================================================================ + + +@dataclass +class AudioVAEConfig: + sample_rate: int = 24000 + upsample_rates: list[int] = field(default_factory=list) + upsample_kernel_sizes: list[int] = field(default_factory=list) + upsample_initial_channel: int = 1536 + resblock: str = "1" + resblock_kernel_sizes: list[int] = field(default_factory=list) + resblock_dilation_sizes: list[list[int]] = field(default_factory=list) + downsample_rates: list[int] = field(default_factory=list) + downsample_channels: list[int] = field(default_factory=list) + activation: str = "snakebeta" + snake_logscale: bool = True + latent_dim: int = 128 + causal: bool = False + mi_num_layers: int = 4 + causal_encoder: bool = False + use_bias_at_final: bool = True + use_tanh_at_final: bool = True + + num_decoder_lookahead: int = 2 + num_encoder_lookahead: int = 2 + + def get(self, key: str, default=None): + return getattr(self, key, default) + + +# ============================================================================ +# Conv layers (vendored from modules/backbone/layers.py) +# ---------------------------------------------------------------------------- +# Only the Conv1d / ConvTranspose1d helpers are needed by the AudioVAE. +# Mlp / MultiHeadAttention / RotaryEmbedding from the same upstream file are +# reserved for the DiT flow-matching head, which gets its own vendored module +# in a later step. +# ============================================================================ + + +class Conv1d(nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 1, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + padding_mode: str = "zeros", + bias: bool = True, + padding=None, + causal: bool = False, + **_kwargs, + ): + self.causal = causal + if padding is None: + if causal: + padding = 0 + self.left_padding = dilation * (kernel_size - 1) + else: + padding = int((kernel_size * dilation - dilation) / 2) + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + padding_mode=padding_mode, + bias=bias, + ) + + self.in_channels = in_channels + + def forward(self, x): + if self.causal: + x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2) + return super().forward(x) + + +class ConvTranspose1d(nn.ConvTranspose1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + output_padding: int = 0, + groups: int = 1, + bias: bool = True, + dilation: int = 1, + padding=None, + padding_mode: str = "zeros", + causal: bool = False, + **_kwargs, + ): + if padding is None: + padding = 0 if causal else (kernel_size - stride) // 2 + if causal: + assert padding == 0, "padding is not allowed in causal ConvTranspose1d." + assert kernel_size == 2 * stride, "kernel_size must be equal to 2*stride in Causal ConvTranspose1d." + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + output_padding=output_padding, + groups=groups, + bias=bias, + dilation=dilation, + padding_mode=padding_mode, + ) + + self.causal = causal + self.stride = stride + + def forward(self, x): + x = super().forward(x) + if self.causal: + x = x[:, :, : -self.stride] + return x + + +# ============================================================================ +# Alias-free low-pass filter (vendored from modules/vocoder/alias_free_filter.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch under +# the Apache License 2.0. The sinc fallback is adopted from adefossez's +# julius.core.sinc (MIT). +# ============================================================================ + + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + # Normalize filter to have sum = 1, otherwise we will have a small leakage + # of the constant component in the input signal. + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + channels: int = 1, + causal: bool = True, + fixed_filter: bool = False, + ): + # kernel_size should be even number for stylegan3 setup, + # in this implementation, odd number is also possible. + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + if causal: + self.pad_left = kernel_size - 1 + self.pad_right = 0 + else: + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + self.fixed_filter = fixed_filter + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + if fixed_filter: + self.register_buffer("filter", filter) + else: + self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) + + # input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + if self.fixed_filter: + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + else: + out = F.conv1d(x, self.filter, stride=self.stride, groups=C) + return out + + +# ============================================================================ +# Alias-free up/down sampling (vendored from modules/vocoder/alias_free_resample.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). +# ============================================================================ + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.channels = channels + self.causal = causal + self.fixed_filter = fixed_filter + if causal: + self.pad = 0 + else: + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) + if self.fixed_filter: + self.register_buffer("filter", filter) + else: + self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) + + # x: [B, C, T] + def forward(self, x): + _, C, _ = x.shape + x = F.pad(x, (self.pad, self.pad), mode="replicate") + if self.fixed_filter: + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + else: + x = self.ratio * F.conv_transpose1d(x, self.filter, stride=self.stride, groups=C) + if self.causal: + x = x[..., : -(self.kernel_size - self.stride)] + else: + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + channels=channels, + causal=causal, + fixed_filter=fixed_filter, + ) + + def forward(self, x): + return self.lowpass(x) + + +# ============================================================================ +# Snake activations + alias-free Activation1d (vendored from +# modules/vocoder/alias_free_act.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). +# Snake / SnakeBeta paper: https://arxiv.org/abs/2006.08195 +# ============================================================================ + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + causal=True, + fixed_filter=False, + ): + super().__init__() + # causal=False + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d( + up_ratio, + up_kernel_size, + activation.in_features, + causal=causal, + fixed_filter=fixed_filter, + ) + self.downsample = DownSample1d( + down_ratio, + down_kernel_size, + activation.in_features, + causal=causal, + fixed_filter=fixed_filter, + ) + + # x: [B,C,T] + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + return self.downsample(x) + + +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + # Snake := x + 1/a * sin^2 (xa) + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + return x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude + of the periodic components. + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + # SnakeBeta := x + 1/b * sin^2 (xa) + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + +# ============================================================================ +# BigVGAN AudioVAE (vendored from modules/vocoder/bigvgan.py) +# ---------------------------------------------------------------------------- +# Copyright (c) 2022 NVIDIA CORPORATION. +# Licensed under the MIT license. +# Upstream layout flattened: the only modification is that the three +# original ``from dots_tts.modules.{backbone,vocoder}.* import ...`` lines +# are dropped, since those names are already defined above in this file. +# ============================================================================ + + +@dataclass(slots=True) +class BigVGANStreamState: + lstm_hidden: tuple[torch.Tensor, torch.Tensor] + decoder: DecoderStreamState + + +@dataclass(slots=True) +class DecoderStreamState: + window: torch.Tensor + chunk_size: int + total_frames: int = 0 + emitted_frames: int = 0 + + +def _empty_chunk( + ref: torch.Tensor, + *, + channels: int | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + return ref.new_zeros( + (ref.size(0), channels or ref.size(1), 0), + dtype=dtype or ref.dtype, + ) + + +def _module_state_device_dtype(module: nn.Module) -> tuple[torch.device, torch.dtype]: + for name in ("weight", "bias", "filter"): + tensor = getattr(module, name, None) + if tensor is not None: + return tensor.device, tensor.dtype + for tensor in itertools.chain(module.parameters(), module.buffers()): + return tensor.device, tensor.dtype + raise RuntimeError(f"Unable to infer state dtype/device for {type(module).__name__}.") + + +def _stream_state_zeros( + batch_size: int, + channels: int, + length: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + return torch.zeros( + (batch_size, channels, max(0, int(length))), + device=device, + dtype=dtype, + ) + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +class Conv1d_S(nn.Module): + "Conv1d for spectral normalisation and orthogonal initialisation" + + def __init__( + self, + in_channels, + out_channels, + kernel_size=1, + stride=1, + dilation=1, + groups=1, + causal=False, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.stride = stride + self.dilation = dilation + self.groups = groups + self.causal = causal + pad = 0 if causal else dilation * (kernel_size - 1) // 2 + self.causal_pad = dilation * (kernel_size - 1) if causal else 0 + + self.layer = weight_norm( + nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=pad, + dilation=dilation, + groups=groups, + ) + ) + + def forward(self, inputs): + if self.causal and self.causal_pad > 0: + inputs = F.pad(inputs, (self.causal_pad, 0)) + return self.layer(inputs) + + +class SLSTM(nn.Module): + """ + LSTM without worrying about the hidden state, nor the layout of the data. + Expects input as convolutional layout. + """ + + def __init__( + self, + dimension: int, + num_layers: int = 2, + skip: bool = True, + bidirectional: bool = False, + ): + super().__init__() + self.skip = skip + self.bidirectional = bidirectional + self.lstm = nn.LSTM( + input_size=dimension, + hidden_size=dimension, + num_layers=num_layers, + bidirectional=bidirectional, + batch_first=True, + ) + self._stream_num_layers = num_layers + self._stream_weight_ih = tuple(getattr(self.lstm, f"weight_ih_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_weight_hh = tuple(getattr(self.lstm, f"weight_hh_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_bias_ih = tuple(getattr(self.lstm, f"bias_ih_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_bias_hh = tuple(getattr(self.lstm, f"bias_hh_l{layer_idx}") for layer_idx in range(num_layers)) + if self.bidirectional: + self.proj_out = nn.Linear(dimension * 2, dimension) + + def forward(self, x): + y, _ = self.lstm(x) + if self.bidirectional: + y = self.proj_out(y) + if self.skip: + y = y + x + return y + + def stream_step( + self, + x: torch.Tensor, + hidden: tuple[torch.Tensor, torch.Tensor] | None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + if self.bidirectional: + raise RuntimeError("Streaming only supports unidirectional SLSTM.") + + residual = x + if hidden is None: + hidden = self.init_stream_state(x.size(0)) + + hidden_h, hidden_c = hidden + next_hidden_h = [] + next_hidden_c = [] + for layer_idx in range(self._stream_num_layers): + layer_input = x + hx = hidden_h[layer_idx] + cx = hidden_c[layer_idx] + outputs = [] + weight_ih = self._stream_weight_ih[layer_idx] + weight_hh = self._stream_weight_hh[layer_idx] + bias_ih = self._stream_bias_ih[layer_idx] + bias_hh = self._stream_bias_hh[layer_idx] + + for frame_idx in range(x.size(1)): + gates = F.linear(layer_input[:, frame_idx, :], weight_ih, bias_ih) + gates = gates + F.linear(hx, weight_hh, bias_hh) + input_gate, forget_gate, cell_gate, output_gate = gates.chunk(4, dim=-1) + input_gate = torch.sigmoid(input_gate) + forget_gate = torch.sigmoid(forget_gate) + cell_gate = torch.tanh(cell_gate) + output_gate = torch.sigmoid(output_gate) + cx = forget_gate * cx + input_gate * cell_gate + hx = output_gate * torch.tanh(cx) + outputs.append(hx) + + x = torch.stack(outputs, dim=1) + next_hidden_h.append(hx) + next_hidden_c.append(cx) + + y = x + if self.skip: + y = y + residual + return y, (torch.stack(next_hidden_h, dim=0), torch.stack(next_hidden_c, dim=0)) + + def init_stream_state( + self, + batch_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_directions = 2 if self.bidirectional else 1 + state_shape = ( + self.lstm.num_layers * num_directions, + batch_size, + self.lstm.hidden_size, + ) + weight = self.lstm.weight_hh_l0 + return ( + weight.new_zeros(state_shape), + weight.new_zeros(state_shape), + ) + + +class ResStack(nn.Module): + def __init__(self, channel, kernel_size=3, base=3, nums=4, causal=False): + super().__init__() + + self.layers = nn.ModuleList([]) + for i in range(nums): + dil = base**i + pad1 = dil * (kernel_size - 1) if causal else dil + pad2 = (kernel_size - 1) if causal else 1 + block = [ + nn.LeakyReLU(), + ] + if causal and pad1 > 0: + block.append(nn.ConstantPad1d((pad1, 0), 0.0)) + block.append( + nn.utils.weight_norm( + nn.Conv1d( + channel, + channel, + kernel_size=kernel_size, + dilation=dil, + padding=0 if causal else pad1, + ) + ) + ) + block.append(nn.LeakyReLU()) + if causal and pad2 > 0: + block.append(nn.ConstantPad1d((pad2, 0), 0.0)) + block.append( + nn.utils.weight_norm( + nn.Conv1d( + channel, + channel, + kernel_size=kernel_size, + dilation=1, + padding=0 if causal else pad2, + ) + ) + ) + self.layers.append(nn.Sequential(*block)) + + def forward(self, x): + for layer in self.layers: + x = x + layer(x) + return x + + +class Encoder(nn.Module): + def __init__( + self, + in_channels=1, + out_channels=100, + base_channels=12, + proj_kernel_size=3, + stack_kernel_size=3, + stack_dilation_base=2, + stacks=6, + channels=(12, 24, 48, 96, 192, 384, 768), + down_sample_factors=(2, 2, 2, 2, 4, 4), + causal=False, + lookahead=0, + ): + super().__init__() + + act_slope = 0.2 + layers = [] + # pre proj_layer + layers += [ + Conv1d_S( + in_channels, + base_channels, + kernel_size=proj_kernel_size, + stride=1, + causal=causal, + ), + nn.LeakyReLU(act_slope, True), + ] + + # channels: [512, 256, 128, 64], upsample_factors: [5, 2, 2] + for (in_c, out_c), down_f in zip(itertools.pairwise(channels), down_sample_factors, strict=True): + layers += [ + Conv1d_S(in_c, out_c, kernel_size=down_f * 2, stride=down_f, causal=causal), + ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks, causal=causal), + nn.LeakyReLU(act_slope, True), + ] + + # post layers + if lookahead > 0: + layers += [ + Conv1d_S( + channels[-1], + out_channels, + kernel_size=lookahead * 2 + 1, + stride=1, + causal=False, + ), + ] + else: + layers += [ + Conv1d_S( + channels[-1], + out_channels, + kernel_size=proj_kernel_size, + stride=1, + causal=causal, + ), + ] + self.generator = nn.Sequential(*layers) + + def forward(self, conditions, _z_inputs=None): + return self.generator(conditions) + + +class AMPBlock1(torch.nn.Module): + def __init__( + self, + h, + channels, + kernel_size=3, + dilation=(1, 3, 5), + activation=None, + causal=True, + ): + super().__init__() + self.h = h + + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + causal=causal, + ) + ), + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + ] + ) + self.convs2.apply(init_weights) + + self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers + + if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=Snake(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs1: + remove_weight_norm(layer) + for layer in self.convs2: + remove_weight_norm(layer) + + +class AMPBlock2(torch.nn.Module): + def __init__(self, h, channels, kernel_size=3, dilation=(1, 3), activation=None, causal=True): + super().__init__() + self.h = h + + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + causal=causal, + ) + ), + ] + ) + self.convs.apply(init_weights) + + self.num_layers = len(self.convs) # total number of conv layers + + if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=Snake(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + for c, a in zip(self.convs, self.activations, strict=True): + xt = a(x) + xt = c(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs: + remove_weight_norm(layer) + + +class Decoder(nn.Module): + def __init__(self, h): + super().__init__() + self.h = h + causal = h.causal + self._stream_window_sizes: dict[int, int] = {} + + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + + num_decoder_lookahead = h.get("num_decoder_lookahead", 2) + # pre conv + self.conv_pre = weight_norm( + Conv1d( + h.latent_dim, + h.upsample_initial_channel, + kernel_size=2 * num_decoder_lookahead + 1, + stride=1, + causal=False, + ) + ) + + # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default + resblock = AMPBlock1 if h.resblock == "1" else AMPBlock2 + + # transposed conv-based upsamplers. does not apply anti-aliasing + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes, strict=True)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + causal=causal, + ) + ) + ] + ) + ) + + # residual blocks using anti-aliased multi-periodicity composition modules (AMP) + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for k, d in zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes, strict=True): + self.resblocks.append(resblock(h, ch, k, d, activation=h.activation, causal=causal)) + + # post conv + if h.activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + activation_post = Snake(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) + elif h.activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal, bias=h.get("use_bias_at_final", True))) + + # weight initialization + for i in range(len(self.ups)): + self.ups[i].apply(init_weights) + self.conv_post.apply(init_weights) + + def forward(self, z): + # pre conv + x = self.conv_pre(z) + + for i in range(self.num_upsamples): + # upsampling + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + # AMP blocks + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + # post conv + x = self.activation_post(x) + x = self.conv_post(x) + if self.h.get("use_tanh_at_final", True): + x = torch.tanh(x) + else: + x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1] + + return x + + @property + def stream_lookahead(self) -> int: + return (self.conv_pre.kernel_size[0] - 1) // 2 + + @staticmethod + def _conv1d_left_context(layer) -> int: + dilation = layer.dilation[0] if isinstance(layer.dilation, tuple) else layer.dilation + kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size + if getattr(layer, "causal", False): + return dilation * (kernel_size - 1) + return layer.padding[0] if isinstance(layer.padding, tuple) else layer.padding + + @staticmethod + def _convtranspose1d_left_context(layer) -> int: + stride = layer.stride[0] if isinstance(layer.stride, tuple) else layer.stride + kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size + if not getattr(layer, "causal", False): + raise NotImplementedError("Streaming only supports causal ConvTranspose1d.") + if kernel_size != 2 * stride: + raise ValueError( + "Streaming ConvTranspose1d expects kernel_size == 2 * stride, got " + f"kernel_size={kernel_size} stride={stride}." + ) + return 1 + + @classmethod + def _activation_left_context(cls, activation: Activation1d) -> int: + upsample = activation.upsample + downsample = activation.downsample.lowpass + if not upsample.causal or not downsample.padding or downsample.pad_right != 0: + raise NotImplementedError("Streaming only supports causal alias-free activations.") + ratio = int(upsample.ratio) + if ratio != int(downsample.stride): + raise ValueError( + "Alias-free activation expects matched up/down ratios, got " + f"up_ratio={ratio} down_ratio={downsample.stride}." + ) + total_left = (upsample.kernel_size - 1) + (downsample.kernel_size - 1) + return (total_left + ratio - 1) // ratio + + @classmethod + def _ampblock_left_context(cls, block) -> int: + if isinstance(block, AMPBlock1): + left_context = 0 + acts1 = block.activations[::2] + acts2 = block.activations[1::2] + for conv1, conv2, act1, act2 in zip( + block.convs1, + block.convs2, + acts1, + acts2, + strict=True, + ): + left_context += ( + cls._activation_left_context(act1) + + cls._conv1d_left_context(conv1) + + cls._activation_left_context(act2) + + cls._conv1d_left_context(conv2) + ) + return left_context + if isinstance(block, AMPBlock2): + left_context = 0 + for conv, activation in zip(block.convs, block.activations, strict=True): + left_context += cls._activation_left_context(activation) + cls._conv1d_left_context(conv) + return left_context + raise TypeError(f"Unsupported resblock type: {type(block).__name__}.") + + def _stream_left_context(self) -> int: + left_context = Fraction(self._conv1d_left_context(self.conv_pre), 1) + current_scale = Fraction(1, 1) + for stage_idx, upsample_layers in enumerate(self.ups): + for upsample in upsample_layers: + left_context += current_scale * self._convtranspose1d_left_context(upsample) + stride = upsample.stride[0] if isinstance(upsample.stride, tuple) else upsample.stride + current_scale /= int(stride) + + stage_start = stage_idx * self.num_kernels + stage_end = stage_start + self.num_kernels + stage_context = max(self._ampblock_left_context(block) for block in self.resblocks[stage_start:stage_end]) + left_context += current_scale * stage_context + + left_context += current_scale * self._activation_left_context(self.activation_post) + left_context += current_scale * self._conv1d_left_context(self.conv_post) + return int(left_context.__ceil__()) + + def stream_window_size(self, chunk_size: int) -> int: + if chunk_size < 1: + raise ValueError(f"chunk_size must be >= 1, got {chunk_size}.") + cached = self._stream_window_sizes.get(chunk_size) + if cached is not None: + return cached + + window_size = chunk_size + self.stream_lookahead + self._stream_left_context() + self._stream_window_sizes[chunk_size] = window_size + return window_size + + def remove_weight_norm(self): + for upsample_layers in self.ups: + for upsample_layer in upsample_layers: + remove_weight_norm(upsample_layer) + for resblock in self.resblocks: + resblock.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + + +class AudioVAE(nn.Module): + def __init__(self, h: AudioVAEConfig): + super().__init__() + self.config = h + + self.h = h + self.hop_size = int(np.prod(h.downsample_rates)) + self.sample_rate = h.sample_rate + self.decoder_lookahead = int(h.get("num_decoder_lookahead", 2)) + + self.audio_encoder = Encoder( + out_channels=h.latent_dim, + down_sample_factors=h.downsample_rates, + channels=h.downsample_channels, + causal=h.causal_encoder, + lookahead=h.get("num_encoder_lookahead", 2), + ) + + intermediate_size = h.latent_dim * 4 + self.enc_mi_layer = nn.Sequential( + nn.Linear(h.latent_dim, intermediate_size), + SLSTM(intermediate_size, num_layers=h.mi_num_layers), + nn.Linear(intermediate_size, h.latent_dim), + ) + self.dec_mi_layer = nn.Sequential( + nn.Linear(h.latent_dim, intermediate_size), + SLSTM(intermediate_size, num_layers=h.mi_num_layers), + nn.Linear(intermediate_size, h.latent_dim), + ) + self.pre_proj = Conv1d( + in_channels=h.latent_dim, + out_channels=h.latent_dim * 2, + kernel_size=1, + stride=1, + ) + self.post_proj = Conv1d(in_channels=h.latent_dim, out_channels=h.latent_dim, kernel_size=1, stride=1) + + self.decoder = Decoder(h) + + def inference(self, data): + latents = self.extract_latents(data["sample"]) + return {"sample": self.inference_from_latents(latents)} + + @torch.autocast(enabled=False, device_type="cuda") + def extract_latents(self, x, do_sample=False): + x = x.float() + x = self.audio_encoder(x) + x = x.permute(0, 2, 1) + x = self.enc_mi_layer(x) + x = x.permute(0, 2, 1) + x = self.pre_proj(x) + if do_sample: + m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) + x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) + return x + + def inference_from_latents(self, x, do_sample=True, noise_scale=1.0): + if do_sample: + assert x.size(1) == self.h.latent_dim * 2, f"Input must be like [B, D, H], got {x.shape}" + m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) + x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) * noise_scale + else: + assert x.size(1) == self.h.latent_dim, f"Input must be like [B, D, H], got {x.shape}" + x = self.post_proj(x) + x = x.permute(0, 2, 1) + x = self.dec_mi_layer(x) + x = x.permute(0, 2, 1) + return self.decoder(x) + + def _validate_stream_latents(self, latents: torch.Tensor) -> None: + if latents.ndim != 3: + raise ValueError( + f"Streaming latents must have shape [batch, latent_dim, frames], got {tuple(latents.shape)}." + ) + if latents.size(1) != self.h.latent_dim: + raise ValueError(f"Streaming latent_dim must be {self.h.latent_dim}, got {latents.size(1)}.") + + def init_stream_state( + self, + batch_size: int = 1, + chunk_size: int = 8, + ) -> BigVGANStreamState: + if not self.h.causal: + raise RuntimeError("Strict streaming requires a causal vocoder.") + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}.") + window_size = self.decoder.stream_window_size(chunk_size) + device, dtype = _module_state_device_dtype(self.decoder.conv_pre) + return BigVGANStreamState( + lstm_hidden=self.dec_mi_layer[1].init_stream_state(batch_size), + decoder=DecoderStreamState( + window=_stream_state_zeros( + batch_size, + self.h.latent_dim, + window_size, + device=device, + dtype=dtype, + ), + chunk_size=int(chunk_size), + ), + ) + + def _decode_stream_latents( + self, + latents: torch.Tensor, + lstm_hidden: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + self._validate_stream_latents(latents) + latents = latents.float() + x = self.post_proj(latents) + x = x.permute(0, 2, 1) + x = self.dec_mi_layer[0](x) + x, lstm_hidden = self.dec_mi_layer[1].stream_step(x, lstm_hidden) + x = self.dec_mi_layer[2](x) + decoder_dtype = self.decoder.conv_pre.weight.dtype + return x.permute(0, 2, 1).to(dtype=decoder_dtype), lstm_hidden + + def _prepare_stream_decoder_input( + self, + latents: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_input, state.lstm_hidden = self._decode_stream_latents( + latents, + state.lstm_hidden, + ) + return decoder_input + + def _append_stream_decoder_input_tensor( + self, + decoder_input: torch.Tensor, + window: torch.Tensor, + valid_frames: torch.Tensor, + ) -> torch.Tensor: + if window.dtype != decoder_input.dtype: + window = window.to(dtype=decoder_input.dtype) + chunk_size = int(decoder_input.size(-1)) + if chunk_size >= window.size(-1): + raise ValueError(f"decoder window size {window.size(-1)} must be larger than chunk_size {chunk_size}.") + positions = torch.arange( + window.size(-1), + device=window.device, + dtype=valid_frames.dtype, + ) + clipped_valid = valid_frames.clamp(min=0, max=window.size(-1)) + combined = torch.cat( + [window, decoder_input.new_zeros(window.size(0), window.size(1), chunk_size)], + dim=-1, + ) + insert_index = clipped_valid + torch.arange( + chunk_size, + device=window.device, + dtype=valid_frames.dtype, + ) + combined.scatter_( + -1, + insert_index.view(1, 1, -1).expand_as(decoder_input), + decoder_input, + ) + new_valid = (clipped_valid + chunk_size).clamp(max=window.size(-1)) + start = (clipped_valid + chunk_size - window.size(-1)).clamp(min=0) + gather_index = (start + positions).clamp(max=combined.size(-1) - 1) + gathered = combined.gather( + -1, + gather_index.view(1, 1, -1).expand_as(window), + ) + mask = (positions < new_valid).to(dtype=window.dtype).view(1, 1, -1) + return gathered * mask + + def _append_stream_decoder_input( + self, + decoder_input: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_state = state.decoder + chunk_size = int(decoder_input.size(-1)) + if chunk_size != decoder_state.chunk_size: + raise ValueError(f"Streaming chunk_size must stay fixed at {decoder_state.chunk_size}, got {chunk_size}.") + window = decoder_state.window + valid_frames = min(decoder_state.total_frames, window.size(-1)) + valid_frames_tensor = window.new_tensor(valid_frames, dtype=torch.int64) + new_window = self._append_stream_decoder_input_tensor( + decoder_input, + window, + valid_frames_tensor, + ) + decoder_state.window = new_window + decoder_state.total_frames += chunk_size + return new_window + + def compiled_stream_step( + self, + latents: torch.Tensor, + hidden_h: torch.Tensor, + hidden_c: torch.Tensor, + window: torch.Tensor, + valid_frames: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + decoder_input, (hidden_h, hidden_c) = self._decode_stream_latents( + latents, + (hidden_h, hidden_c), + ) + new_window = self._append_stream_decoder_input_tensor( + decoder_input, + window, + valid_frames, + ) + audio_window = self.decode_stream_window(new_window) + return audio_window, hidden_h, hidden_c, new_window + + def decode_stream_window(self, window: torch.Tensor) -> torch.Tensor: + decoder_dtype = self.decoder.conv_pre.weight.dtype + if window.dtype != decoder_dtype: + window = window.to(dtype=decoder_dtype) + return self.decoder(window) + + def _slice_stream_audio_window( + self, + audio_window: torch.Tensor, + state: BigVGANStreamState, + *, + final: bool, + ) -> torch.Tensor: + decoder_state = state.decoder + stable_end = ( + decoder_state.total_frames if final else max(0, decoder_state.total_frames - self.decoder.stream_lookahead) + ) + if stable_end <= decoder_state.emitted_frames: + return _empty_chunk(audio_window, channels=1) + + window_size = decoder_state.window.size(-1) + valid_frames = min(decoder_state.total_frames, window_size) + window_start = decoder_state.total_frames - valid_frames + if decoder_state.emitted_frames < window_start: + raise RuntimeError("Decoder stream window is too short for fixed-graph decoding.") + + local_start = decoder_state.emitted_frames - window_start + local_end = stable_end - window_start + sample_start = local_start * self.hop_size + sample_end = local_end * self.hop_size + decoder_state.emitted_frames = stable_end + return audio_window[..., sample_start:sample_end] + + def stream_step( + self, + latents: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_input = self._prepare_stream_decoder_input(latents, state) + window = self._append_stream_decoder_input(decoder_input, state) + audio_window = self.decode_stream_window(window) + return self._slice_stream_audio_window(audio_window, state, final=False) + + def stream_flush(self, state: BigVGANStreamState) -> torch.Tensor: + audio_window = self.decode_stream_window(state.decoder.window) + return self._slice_stream_audio_window(audio_window, state, final=True) + + def remove_weight_norm(self): + self.decoder.remove_weight_norm() + + +__all__ = ["AudioVAEConfig", "AudioVAE"] diff --git a/vllm_omni/model_executor/models/dots_tts/pipeline.py b/vllm_omni/model_executor/models/dots_tts/pipeline.py new file mode 100644 index 000000000..22b50c629 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/pipeline.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dots.tts pipeline topology (frozen). + +Single-stage AR TTS: text → 48 kHz waveform in one pass. +Architecture: Qwen2.5-1.5B base LM (vLLM-native) → DiTAR / MeanFlow head +→ AudioVAE.decode → BigVGAN → 48 kHz mono. + +Speaker conditioning via CAM++ x-vector injected in preprocess(). + +Modeled after voxcpm2's single-stage pipeline (same "vLLM-native base LM ++ side-path computation" pattern), but simplified: no residual LM, no FSQ. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +DOTS_TTS_PIPELINE = PipelineConfig( + model_type="dots_tts", + model_arch="DotsTTSForConditionalGeneration", + default_deploy_config_name="dots_tts.yaml", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="latent_generator", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="audio", + owns_tokenizer=True, + engine_output_type="audio", + sampling_constraints={ + "detokenize": False, + # Stop signal slot in compute_logits (voxcpm2 convention — + # talker emits logits[i, 1] = 1.0 when prob_stop > 0.8, vLLM + # matches against this list to terminate the request). Not + # Qwen2.5 EOS (151643) — that ID would never come out of the + # 2-slot continue/stop softmax our compute_logits builds. + "stop_token_ids": [1], + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 498f9eb85..530fb1712 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -301,6 +301,12 @@ _OMNI_MODELS = { "voxcpm2_talker", "VoxCPM2TalkerForConditionalGeneration", ), + ## dots.tts + "DotsTTSForConditionalGeneration": ( + "dots_tts", + "dots_tts_talker", + "DotsTTSForConditionalGeneration", + ), ## Voxtral TTS "VoxtralTTSForConditionalGeneration": ( "voxtral_tts", diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 88a23dc0c..cdf105d69 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -20,6 +20,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts", "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2", + "DotsTTSConfig": "vllm_omni.transformers_utils.configs.dots_tts", "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3", "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice", @@ -42,6 +43,7 @@ __all__ = [ "FishSpeechFastARConfig", "GLMTTSConfig", "VoxCPM2Config", + "DotsTTSConfig", "VoxtralTTSConfig", "CosyVoice3Config", "OmniVoiceConfig", @@ -70,6 +72,7 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3 # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/dots_tts.py b/vllm_omni/transformers_utils/configs/dots_tts.py new file mode 100644 index 000000000..1950a2b3a --- /dev/null +++ b/vllm_omni/transformers_utils/configs/dots_tts.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HF config for dots.tts. + +dots.tts splits its config across two files in the HF checkpoint: + * ``config.json`` — top-level dots.tts fields (latent_dim / patch_size + / PatchEncoder / DiT / vocoder / etc.). + * ``llm_config.json`` — full Qwen2 LM config (vocab_size / hidden_size / + num_hidden_layers / rope_theta / etc.). + +``AutoConfig.from_pretrained`` only consumes ``config.json``: its dots.tts +sub-blocks (``vocoder`` / ``DiT`` / ``PatchEncoder``) arrive through +``**kwargs`` and land as attributes, which the talker's config factories +consume directly. The Qwen2 LM fields live in ``llm_config.json``, which +AutoConfig never reads — this class hoists them to top-level attributes +with defaults matching the released checkpoints (soar / base / mf all ship +the same Qwen2.5-1.5B backbone). vLLM's ``Qwen2Model.__init__`` then reads +``config.hidden_size`` / ``config.vocab_size`` / etc. through +``get_text_config()``, the same pattern voxcpm2 follows. + +Future work: read ``llm_config.json`` during ``from_pretrained`` so a +checkpoint with a different LM shape works without editing this file. + +Upstream reference: + https://huggingface.co/rednote-hilab/dots.tts-soar/blob/main/llm_config.json +""" + +from __future__ import annotations + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig + + +class DotsTTSConfig(PretrainedConfig): + """Configuration for dots.tts integration (Qwen2.5-1.5B-Base AR backbone).""" + + model_type = "dots_tts" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + # -- top-level dots.tts params -- + lm_config: dict | None = None, + # -- LM defaults (overridden by lm_config dict if present) -- + # All values match dots.tts-soar's llm_config.json. + vocab_size: int = 151672, + hidden_size: int = 1536, + intermediate_size: int = 8960, + num_hidden_layers: int = 28, + num_attention_heads: int = 12, + num_key_value_heads: int = 2, + max_position_embeddings: int = 131072, + rms_norm_eps: float = 1e-6, + rope_theta: float = 1_000_000.0, + tie_word_embeddings: bool = True, + hidden_act: str = "silu", + bos_token_id: int = 151643, + eos_token_id: int = 151643, + **kwargs, + ): + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + # The dots.tts sub-blocks (vocoder / DiT / PatchEncoder / latent_dim + # / patch_size / ...) arrive via **kwargs and are stored as + # attributes by PretrainedConfig; the talker's config factories + # read them from there. + self.lm_config = lm_config or {} + + # -- Hoist LM parameters to top-level for vLLM's Qwen2Model -- + lm = self.lm_config + self.vocab_size = lm.get("vocab_size", vocab_size) + self.hidden_size = lm.get("hidden_size", hidden_size) + self.intermediate_size = lm.get("intermediate_size", intermediate_size) + self.num_hidden_layers = lm.get("num_hidden_layers", num_hidden_layers) + self.num_attention_heads = lm.get("num_attention_heads", num_attention_heads) + self.num_key_value_heads = lm.get("num_key_value_heads", num_key_value_heads) + self.max_position_embeddings = lm.get("max_position_embeddings", max_position_embeddings) + self.rms_norm_eps = lm.get("rms_norm_eps", rms_norm_eps) + self.rope_theta = lm.get("rope_theta", rope_theta) + self.hidden_act = lm.get("hidden_act", hidden_act) + self.head_dim = self.hidden_size // self.num_attention_heads + + # vLLM's Qwen2Model reads ``max_window_layers`` and checks against + # ``num_hidden_layers`` only when ``use_sliding_window`` is true. + # soar disables sliding window — set both to the same value so the + # assertion (if executed) stays vacuous. + self.use_sliding_window = lm.get("use_sliding_window", False) + self.sliding_window = lm.get("sliding_window", None) + self.max_window_layers = lm.get("max_window_layers", self.num_hidden_layers) + + def get_text_config(self, **kwargs): + """Return self — LM fields are hoisted to top-level attributes.""" + return self + + +AutoConfig.register("dots_tts", DotsTTSConfig) + +__all__ = ["DotsTTSConfig"] From 596c16a550aa134faf7f3dcfa0f8adf513ccd9ce Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Fri, 14 Aug 2026 08:46:59 +0800 Subject: [PATCH 30/41] [Kernel] Fuse Q/K RMSNorm and RoPE (#5990) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Signed-off-by: lishunyang12 Signed-off-by: Mengjie Zhao Co-authored-by: lishunyang12 Co-authored-by: Claude Opus 5 --- .../layers/test_fused_qk_norm_rope.py | 82 +++++ tests/e2e/accuracy/ltx/run_ltx_reference.py | 12 +- .../diffusion/cache/teacache/extractors.py | 7 +- .../diffusion/layers/fused_qk_norm_rope.py | 299 ++++++++++++++++++ .../minimax_h3/minimax_h3_transformer.py | 47 ++- 5 files changed, 420 insertions(+), 27 deletions(-) create mode 100644 tests/diffusion/layers/test_fused_qk_norm_rope.py create mode 100644 vllm_omni/diffusion/layers/fused_qk_norm_rope.py diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope.py b/tests/diffusion/layers/test_fused_qk_norm_rope.py new file mode 100644 index 000000000..40888556a --- /dev/null +++ b/tests/diffusion/layers/test_fused_qk_norm_rope.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] + +_HEAD_DIM = 128 +_ROTARY_DIM = 96 +_EPS = 1e-5 + + +def _reference(q, k, q_weight, k_weight, rope_table): + q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS) + k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS) + half = _ROTARY_DIM // 2 + cos = rope_table[..., :half].unsqueeze(1) + sin = rope_table[..., half:].unsqueeze(1) + + def apply(x): + first = x[..., :half] + second = x[..., half:_ROTARY_DIM] + return torch.cat( + ( + first * cos - second * sin, + second * cos + first * sin, + x[..., _ROTARY_DIM:], + ), + dim=-1, + ) + + return apply(q), apply(k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +@pytest.mark.parametrize("seq_len", [1, 257, 1024]) +def test_fused_qk_norm_rope_matches_bf16_reference(seq_len): + from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + fused_qk_norm_rope, + ) + + torch.manual_seed(17) + heads = 14 + qkv = torch.randn( + seq_len, + heads * _HEAD_DIM * 3, + device="cuda", + dtype=torch.bfloat16, + ) + q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM) + k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view( + seq_len, + heads, + _HEAD_DIM, + ) + q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda") + rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16) + + expected_q, expected_k = _reference( + q, + k, + q_weight, + k_weight, + rope_table, + ) + actual_q, actual_k = fused_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + ) + + torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py index 3ab794662..ccabd4596 100644 --- a/tests/e2e/accuracy/ltx/run_ltx_reference.py +++ b/tests/e2e/accuracy/ltx/run_ltx_reference.py @@ -260,15 +260,9 @@ def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]: multimodal_output = frames.multimodal_output or {} audio = multimodal_output.get("audio") audio_sample_rate = multimodal_output.get("audio_sample_rate") - if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput): - frames = frames.request_output - multimodal_output = frames.multimodal_output or {} - audio = multimodal_output.get("audio", audio) - audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate) - if isinstance(frames, OmniRequestOutput): - if not frames.images: - raise ValueError("No video frames found in OmniRequestOutput") - frames = frames.images + if not frames.images: + raise ValueError("No video frames found in OmniRequestOutput") + frames = frames.images if isinstance(frames, list) and len(frames) == 1: frames = frames[0] diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 5f1ca1d22..8689c2c54 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1263,6 +1263,7 @@ def extract_minimax_h3_context( _BF16_DTYPE, _FORWARD_SUPPORTED_KWARGS, MINIMAX_H3_ADALN_MODALITY_NUM, + _build_rope_table, _modulate_scale_shift, _required_kwarg, ) @@ -1318,7 +1319,7 @@ def extract_minimax_h3_context( raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}") device = x.device - rope_freqs = module.rope(img_position_ids).to(device) + rope_table = _build_rope_table(module.rope(img_position_ids).to(device)) decoder_input, t_emb = module._embed( x=x, @@ -1351,7 +1352,7 @@ def extract_minimax_h3_context( def run_transformer_blocks() -> tuple[torch.Tensor, ...]: hidden, block_rope, block_combined = module.sp_prepare( decoder_input, - rope_freqs, + rope_table, combined_indices, ) for block in module.blocks: @@ -1359,7 +1360,7 @@ def extract_minimax_h3_context( hidden, t_emb=t_emb, combined_indices=block_combined, - rope_freqs=block_rope, + rope_table=block_rope, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=seq_len, diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py new file mode 100644 index 000000000..64462d6c2 --- /dev/null +++ b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Q/K RMSNorm followed by packed non-interleaved RoPE. + +The public contract is shared by diffusion attention implementations: + +* ``q`` and ``k`` are ``[tokens, heads, head_dim]``; +* norm weights are one-dimensional ``[head_dim]`` tensors; +* ``rope_table`` is ``[tokens, rotary_dim]`` and stores + ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. + +The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K +or rotary-product intermediates. Unsupported inputs use the eager reference. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch.library import Library +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + +_FUSED_HEAD_DIM = 128 +_FUSED_ROTARY_DIM = 96 +_HEADS_PER_PROGRAM = 8 + + +def _apply_rope_table( + x: torch.Tensor, + rope_table: torch.Tensor, + rotary_dim: int, +) -> torch.Tensor: + half = rotary_dim // 2 + cos = rope_table[..., :half].to(x.dtype).unsqueeze(1) + sin = rope_table[..., half:].to(x.dtype).unsqueeze(1) + first = x[..., :half] + second = x[..., half:rotary_dim] + return torch.cat( + ( + first * cos - second * sin, + second * cos + first * sin, + x[..., rotary_dim:], + ), + dim=-1, + ) + + +if HAS_TRITON: + + @triton.jit + def _rms_norm_rope_kernel( + x_ptr, + weight_ptr, + rope_table_ptr, + out_ptr, + x_stride_t, + x_stride_h, + x_stride_d, + rope_stride_t, + out_stride_t, + out_stride_h, + out_stride_d, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + rotary_half: tl.constexpr, + eps: tl.constexpr, + heads_per_program: tl.constexpr, + ): + token = tl.program_id(0) + head_group = tl.program_id(1) + heads = head_group * heads_per_program + tl.arange(0, heads_per_program) + dims = tl.arange(0, head_dim) + mask = heads[:, None] < num_heads + offsets = token * x_stride_t + heads[:, None] * x_stride_h + dims[None, :] * x_stride_d + + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + weight = tl.load(weight_ptr + dims).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) + normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16) + + rotary_dim = rotary_half * 2 + pair_dims = tl.where( + dims < rotary_half, + dims + rotary_half, + tl.where(dims < rotary_dim, dims - rotary_half, dims), + ) + pair_offsets = token * x_stride_t + heads[:, None] * x_stride_h + pair_dims[None, :] * x_stride_d + pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + pair_weight = tl.load(weight_ptr + pair_dims).to(tl.float32) + pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16) + + freq_dims = tl.where( + dims < rotary_half, + dims, + tl.where(dims < rotary_dim, dims - rotary_half, 0), + ) + table_offsets = token * rope_stride_t + freq_dims + cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32) + sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32) + first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin + second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin + output = tl.where( + dims < rotary_dim, + tl.where(dims < rotary_half, first, second), + normalized.to(tl.float32), + ) + + out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d + tl.store(out_ptr + out_offsets, output, mask=mask) + + +def _eager_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + q_norm = F.rms_norm(q, (head_dim,), q_weight, eps) + k_norm = F.rms_norm(k, (head_dim,), k_weight, eps) + return ( + _apply_rope_table(q_norm, rope_table, rotary_dim), + _apply_rope_table(k_norm, rope_table, rotary_dim), + ) + + +def _fused_cuda_supported( + q: torch.Tensor, + k: torch.Tensor, + head_dim: int, + rotary_dim: int, +) -> bool: + return ( + HAS_TRITON + and current_platform.is_cuda() + and q.is_cuda + and k.is_cuda + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and head_dim == _FUSED_HEAD_DIM + and rotary_dim == _FUSED_ROTARY_DIM + ) + + +def _launch_fused_rms_norm_rope( + x: torch.Tensor, + weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, +) -> torch.Tensor: + tokens, heads, head_dim = x.shape + rotary_half = rope_table.shape[-1] // 2 + out = torch.empty(x.shape, dtype=x.dtype, device=x.device) + if tokens == 0: + return out + grid = (tokens, triton.cdiv(heads, _HEADS_PER_PROGRAM)) + _rms_norm_rope_kernel[grid]( + x, + weight, + rope_table, + out, + x.stride(0), + x.stride(1), + x.stride(2), + rope_table.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + num_heads=heads, + head_dim=head_dim, + rotary_half=rotary_half, + eps=eps, + heads_per_program=_HEADS_PER_PROGRAM, + num_warps=8, + ) + return out + + +def _fused_qk_norm_rope_impl( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + return _eager_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + head_dim, + rotary_dim, + ) + return ( + _launch_fused_rms_norm_rope(q, q_weight, rope_table, eps), + _launch_fused_rms_norm_rope(k, k_weight, rope_table, eps), + ) + + +def _fused_qk_norm_rope_fake( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim + return torch.empty_like(q), torch.empty_like(k) + + +_OMNI_OP_LIB = Library("vllm_omni", "FRAGMENT") +if not hasattr(torch.ops.vllm_omni, "fused_qk_norm_rope"): + direct_register_custom_op( + op_name="fused_qk_norm_rope", + op_func=_fused_qk_norm_rope_impl, + fake_impl=_fused_qk_norm_rope_fake, + mutates_args=[], + target_lib=_OMNI_OP_LIB, + ) + + +def fused_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + *, + head_dim: int | None = None, + rotary_dim: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply Q/K RMSNorm and packed non-interleaved RoPE.""" + if q.ndim != 3 or k.ndim != 3: + raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}") + if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]: + raise ValueError(f"q and k shapes are incompatible: {q.shape} and {k.shape}") + if q.dtype != k.dtype or q.device != k.device: + raise ValueError("q and k must have the same dtype and device") + if q.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise TypeError(f"Fused QK RMSNorm/RoPE requires floating inputs, got {q.dtype}") + + head_dim = q.shape[-1] if head_dim is None else head_dim + rotary_dim = rope_table.shape[-1] if rotary_dim is None else rotary_dim + if q.shape[-1] != head_dim: + raise ValueError(f"Expected q/k head_dim={head_dim}, got {q.shape[-1]}") + if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2: + raise ValueError(f"rotary_dim must be even and in [2, {head_dim}], got {rotary_dim}") + if q_weight.shape != (head_dim,) or k_weight.shape != (head_dim,): + raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}") + if q_weight.device != q.device or k_weight.device != q.device: + raise ValueError("Q/K norm weights must be on the activation device") + if rope_table.device != q.device or rope_table.dtype != q.dtype: + raise ValueError("rope_table must have the same dtype and device as q/k") + if rope_table.shape != (q.shape[0], rotary_dim): + raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}") + + q_weight = q_weight.contiguous() + k_weight = k_weight.contiguous() + rope_table = rope_table.contiguous() + if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + return _fused_qk_norm_rope_impl( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + head_dim, + rotary_dim, + ) + return torch.ops.vllm_omni.fused_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + head_dim, + rotary_dim, + ) + + +__all__ = ["fused_qk_norm_rope"] diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 5e180aa23..659b64bc4 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -33,6 +33,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) +from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding @@ -237,6 +238,15 @@ class MiniMaxH3Rope(nn.Module): return torch.cat((half, half), dim=-1) # [S, 96] +def _build_rope_table(freqs: torch.Tensor) -> torch.Tensor: + """Materialize H3's packed ``[cos(freqs[:48]), sin(freqs[:48])]`` table.""" + half = freqs.shape[-1] // 2 + return torch.cat( + (torch.cos(freqs[..., :half]), torch.sin(freqs[..., :half])), + dim=-1, + ).to(_BF16_DTYPE) + + class MiniMaxH3TimeEmbedder(nn.Module): def __init__( self, @@ -453,7 +463,7 @@ class MiniMaxH3Attention(nn.Module): self, x: torch.Tensor, *, - rope_freqs: torch.Tensor | None, + rope_table: torch.Tensor | None, cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int | None = None, @@ -479,11 +489,18 @@ class MiniMaxH3Attention(nn.Module): q = q.view(total, self.num_heads, self.head_dim) k = k.view(total, self.num_kv_heads, self.head_dim) v = v.view(total, self.num_kv_heads, self.head_dim) - q = self.q_norm(q) - k = self.k_norm(k) - if rope_freqs is not None: - q = self._apply_rope(q, rope_freqs) - k = self._apply_rope(k, rope_freqs) + if rope_table is None: + q = self.q_norm(q) + k = self.k_norm(k) + else: + q, k = fused_qk_norm_rope( + q, + k, + self.q_norm.weight, + self.k_norm.weight, + rope_table, + self.q_norm.variance_epsilon, + ) # The packed layout uses a second document for alignment padding. # Local/Ulysses backends unpad it, while Ring keeps aligned rows for @@ -628,7 +645,7 @@ class MiniMaxH3TokenRefinerBlock(nn.Module): ) -> torch.Tensor: x = x + self.attn( self.norm1(x), - rope_freqs=None, + rope_table=None, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, ) @@ -707,7 +724,7 @@ class MiniMaxH3DiTBlock(nn.Module): *, t_emb: torch.Tensor, combined_indices: torch.Tensor, - rope_freqs: torch.Tensor, + rope_table: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int, @@ -735,7 +752,7 @@ class MiniMaxH3DiTBlock(nn.Module): h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE) h = self.attn( h, - rope_freqs=rope_freqs, + rope_table=rope_table, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=packed_total, @@ -817,10 +834,10 @@ class MiniMaxH3SPPrepare(nn.Module): def forward( self, hidden_states: torch.Tensor, - rope_freqs: torch.Tensor, + rope_table: torch.Tensor, combined_indices: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return hidden_states, rope_freqs, combined_indices + return hidden_states, rope_table, combined_indices class MiniMaxH3SPGather(nn.Module): @@ -1153,8 +1170,8 @@ class MiniMaxH3DiTModel(nn.Module): if inverse_indices.shape[0] != seq_len: raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}") device = x.device - # Compute RoPE frequencies over the full packed sequence. - rope_freqs = self.rope(img_position_ids).to(device) + # Compute the packed RoPE table once for all transformer blocks. + rope_table = _build_rope_table(self.rope(img_position_ids).to(device)) decoder_input, t_emb = self._embed( x=x, @@ -1175,7 +1192,7 @@ class MiniMaxH3DiTModel(nn.Module): hidden = decoder_input cu_seqlens = cu_seqlens.to(device) - block_rope = rope_freqs + block_rope = rope_table block_combined = combined_indices hidden, block_rope, block_combined = self.sp_prepare( @@ -1188,7 +1205,7 @@ class MiniMaxH3DiTModel(nn.Module): hidden, t_emb=t_emb, combined_indices=block_combined, - rope_freqs=block_rope, + rope_table=block_rope, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=seq_len, From 8a5bf2fb3f1ab0f7242c72d25f694652fe957256 Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Fri, 14 Aug 2026 12:04:15 +0800 Subject: [PATCH 31/41] [doc]: Update latest news section for verl-omni release (#6187) Signed-off-by: Samit <285365963@qq.com> --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 418d0b38d..e6abc70d7 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone --- *Latest News* 🔥 +- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0). - [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support. - [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support. - [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage. From 5d50bdc2ab7691b98fa0371ab2d9c534ca801991 Mon Sep 17 00:00:00 2001 From: zhengjia Date: Fri, 14 Aug 2026 12:14:55 +0800 Subject: [PATCH 32/41] [Perf] Cap diffusion worker thread count before spawn (#6165) Signed-off-by: zjli2013 Co-authored-by: Cursor --- vllm_omni/diffusion/executor/multiproc_executor.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 7de46b5cf..7cc6488c1 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -18,6 +18,7 @@ import zmq from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.executor.multiproc_executor import set_multiprocessing_worker_envs from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput from vllm_omni.diffusion.executor.abstract import DiffusionExecutor @@ -294,6 +295,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): logger.info("Starting server...") num_gpus = cast(int, od_config.num_gpus) + # Without this, every worker inherits one Torch thread per core, so an + # N-GPU run oversubscribes the host by N x core_count. Honours a + # user-provided OMP_NUM_THREADS. + set_multiprocessing_worker_envs() mp.set_start_method("spawn", force=True) processes = [] From 2b08419d8be88f3f78852e42bb70a0bd46b289ae Mon Sep 17 00:00:00 2001 From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Date: Fri, 14 Aug 2026 00:36:56 -0700 Subject: [PATCH 33/41] [CI] Per-model, per-entry-mode code coverage collection (#5593) Signed-off-by: Shenglei Fu Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Co-authored-by: Claude Sonnet 5 Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com> Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .buildkite/common/ci_mirror_hardwares.yml | 7 ++ .buildkite/common/scripts/run_cov_split.sh | 82 ++++++++++++++++++++ .buildkite/cuda/test-merge.yml | 28 +++++-- .gitignore | 2 + docs/contributing/ci/ci_settings.md | 90 ++++++++++++++++++++++ pyproject.toml | 15 ++++ tests/buildkite/test_upload_pipeline.py | 46 +++++++++++ 7 files changed, 263 insertions(+), 7 deletions(-) create mode 100755 .buildkite/common/scripts/run_cov_split.sh diff --git a/.buildkite/common/ci_mirror_hardwares.yml b/.buildkite/common/ci_mirror_hardwares.yml index 00c6e4f82..e68cbd382 100644 --- a/.buildkite/common/ci_mirror_hardwares.yml +++ b/.buildkite/common/ci_mirror_hardwares.yml @@ -9,6 +9,11 @@ # a2b3_npu_1 — ascend-a2b3, resource_class npu-1 (A2/B3 CI image) # a2b3_npu_4 — ascend-a2b3, resource_class npu-4 # a3_npu_2 — ascend-a3, resource_class npu-2 +# +# The ``l4_*`` presets run under the ``docker`` plugin, which does not put +# ``buildkite-agent`` on PATH inside the container the way the ``kubernetes`` +# presets do, so they set ``mount-buildkite-agent: true``; without it any step +# calling ``buildkite-agent`` (artifact upload, annotate, meta-data) fails there. x-npu-k8s-plugins: &npu_k8s_plugins - kubernetes: @@ -26,6 +31,7 @@ mirror_hardwares: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT always-pull: true propagate-environment: true + mount-buildkite-agent: true shm-size: "8gb" environment: - HF_HOME=/fsx/hf_cache @@ -41,6 +47,7 @@ mirror_hardwares: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT always-pull: true propagate-environment: true + mount-buildkite-agent: true shm-size: "8gb" environment: - HF_HOME=/fsx/hf_cache diff --git a/.buildkite/common/scripts/run_cov_split.sh b/.buildkite/common/scripts/run_cov_split.sh new file mode 100755 index 000000000..fe068519c --- /dev/null +++ b/.buildkite/common/scripts/run_cov_split.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Run a model's E2E tests once per entry mode as separate coverage runs and upload +# the reports, so per-mode attribution survives (see vllm-project/vllm-omni#5332). +# +# Usage: run_cov_split.sh --model-id [--offline ] [--online ] \ +# -- +# +# --model-id Names the artifacts: coverage--{offline,online}-.xml.gz. +# Use the model's directory name under vllm_omni/*/models/. +# --offline Pytest path(s) for the offline half. Repeatable. +# --online Pytest path(s) for the online half. Repeatable. +# -- Everything after this is forwarded verbatim to every pytest run, +# quoting preserved: the job's own -m expression, --run-level, and +# anything else it needs (e.g. --test-config-file). +# +# At least one of --offline/--online is required, so a job with only one mode runs +# just that half. Total runtime is bounded by the step's timeout_in_minutes. +# +# Exits non-zero if any half or the upload failed. +set -uo pipefail + +MODEL_ID="" +OFFLINE=() +ONLINE=() +PYTEST_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --model-id) MODEL_ID="$2"; shift 2 ;; + --offline) OFFLINE+=("$2"); shift 2 ;; + --online) ONLINE+=("$2"); shift 2 ;; + --) shift; PYTEST_ARGS=("$@"); break ;; + *) echo "run_cov_split.sh: unknown argument: $1" >&2; exit 2 ;; + esac +done + +if [[ -z "${MODEL_ID}" ]]; then + echo "run_cov_split.sh: missing --model-id" >&2 + exit 2 +fi + +if ((${#OFFLINE[@]} == 0 && ${#ONLINE[@]} == 0)); then + echo "run_cov_split.sh: need at least one of --offline / --online" >&2 + exit 2 +fi + +# Two steps can test the same model, and artifacts are scoped to the build (a +# nightly one carries the ready, merge and nightly tiers at once), so unsuffixed +# reports would share a path -- a glob download then races to rename onto one file. +# Not BUILDKITE_JOB_ID: the step id survives retries, so attempts collapse under +# Buildkite's latest-attempt lookup instead of accumulating. A parallel step would +# need BUILDKITE_PARALLEL_JOB as well; no pilot is one. +REPORT_SUFFIX="${BUILDKITE_STEP_ID:-local}" +OFFLINE_REPORT="coverage-${MODEL_ID}-offline-${REPORT_SUFFIX}.xml" +ONLINE_REPORT="coverage-${MODEL_ID}-online-${REPORT_SUFFIX}.xml" + +run_half() { + local mode="$1" + shift + local report="coverage-${MODEL_ID}-${mode}-${REPORT_SUFFIX}.xml" + echo "--- coverage: ${MODEL_ID} ${mode}" + pytest -s -v "$@" "${PYTEST_ARGS[@]}" --cov=vllm_omni --cov-report="xml:${report}" + local status=$? + # Cobertura lists every statement in vllm_omni whether or not it ran, so the + # report is ~7MB no matter how little the job covers. The repetition is what + # dominates it, and gzip takes that to ~400KB. + gzip -9 -f "${report}" || status=1 + return "${status}" +} + +# Clear both modes, not just the ones being run: the upload glob matches both, so +# a report left by an earlier run would otherwise be published as if the mode it +# belongs to had run this time. +rm -f "${OFFLINE_REPORT}" "${ONLINE_REPORT}" "${OFFLINE_REPORT}.gz" "${ONLINE_REPORT}.gz" + +EXIT=0 +((${#OFFLINE[@]})) && { run_half offline "${OFFLINE[@]}" || EXIT=1; } +((${#ONLINE[@]})) && { run_half online "${ONLINE[@]}" || EXIT=1; } + +buildkite-agent artifact upload "coverage-${MODEL_ID}-*-${REPORT_SUFFIX}.xml.gz" || EXIT=1 + +exit "${EXIT}" diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 363e8fbdc..970ab5a5a 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -94,13 +94,19 @@ steps: - vllm_omni/model_executor/models/qwen3_tts/ - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - vllm_omni/deploy/qwen3_tts.yaml + - .buildkite/common/scripts/run_cov_split.sh + - pyproject.toml + # 40 -> 60: the coverage split runs the model twice. + timeout_in_minutes: 60 commands: - | - timeout 40m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model' - " + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 + .buildkite/common/scripts/run_cov_split.sh \ + --model-id qwen3_tts \ + --offline tests/e2e/offline_inference/test_qwen3_tts_base.py \ + --online tests/e2e/online_serving/test_qwen3_tts_base.py \ + -- -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: l4_1 - label: "TTS · Higgs-Audio-V3 Test" @@ -203,9 +209,17 @@ steps: - vllm_omni/diffusion/models/bagel/ - vllm_omni/model_executor/stage_input_processors/bagel.py - vllm_omni/deploy/bagel.yaml - timeout_in_minutes: 20 + - .buildkite/common/scripts/run_cov_split.sh + - pyproject.toml + # 20 -> 40: the coverage split runs the model twice. + timeout_in_minutes: 40 commands: - - pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model' + - | + .buildkite/common/scripts/run_cov_split.sh \ + --model-id bagel \ + --offline tests/e2e/offline_inference/test_bagel.py \ + --online tests/e2e/online_serving/test_bagel.py \ + -- -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_1 - label: "Diffusion · Qwen Image Test" diff --git a/.gitignore b/.gitignore index 3c24e40ae..f11b57388 100644 --- a/.gitignore +++ b/.gitignore @@ -46,6 +46,8 @@ htmlcov/ .cache nosetests.xml coverage.xml +coverage-*.xml +coverage-*.xml.gz *.cover *.py,cover .hypothesis/ diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index ff4b8e107..62717ddc9 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -327,6 +327,96 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for - Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` nightly runs on main. - Prefix labels by model domain: **Omni ·**, **TTS ·**, **Diffusion ·**, **Simple ·**, etc., matching existing steps. +### Per-model coverage + +Pilot (v1): two Merge-tier (L3) jobs in `.buildkite/cuda/test-merge.yml` upload +per-model, per-entry-mode coverage as Buildkite artifacts — "Diffusion · Bagel +Test" and "TTS · Qwen3-TTS Base Test". Both run on a single GPU so the pilot is +cheap to reproduce. This is a pilot, not a rollout: no dashboard/visualization +lives in this repo, and nothing is gated on coverage. + +**Naming convention**: `coverage---.xml.gz`, where +`` is the model's directory name under +`vllm_omni/diffusion/models//` or +`vllm_omni/model_executor/models//` (e.g. `z_image`, `qwen3_omni`), +`` is `online` or `offline`, and `` is `BUILDKITE_STEP_ID` +(`local` outside Buildkite). `gunzip` before feeding a report to a Cobertura +consumer. + +Reports are gzipped because `--cov=vllm_omni` sets coverage's `source`, which +makes it walk the package and emit one `` element per statement for every +file — including the ones a single model's tests never import. That inventory, +not the coverage, is the size: a Bagel report is 7.1 MB, of which 96.5% is +`` elements and 77% belongs to the 840 of 1169 files that run never +imported. A run covering zero lines produces the same 7 MB. gzip takes it to +roughly 410 KB and changes nothing about the data. + +`` is what keeps two steps that cover the same model apart. Artifacts are +scoped to the build, and a nightly build carries the ready, merge and nightly tiers +at once, so several same-model steps land in one namespace — `qwen3_tts` already has +both a Base and a CustomVoice job, and `minicpmo_4_5` a base and a duplex one. Two +uploads on one path do not fail loudly: an exact-name download can report the path +as ambiguous, and a glob download fetches every match concurrently and renames each +onto the same destination, so whichever finishes last wins. The step id remains the +logical step identifier across retries, while Buildkite's default artifact lookup +selects only the latest attempt — `BUILDKITE_JOB_ID` would instead give every +attempt its own filename, leaving a failed attempt's partial report to be picked up +alongside the real one. A step using `parallelism`/matrix would need +`BUILDKITE_PARALLEL_JOB` added, since its jobs share a step id; neither pilot is +one. + +**Opting in a new model**: replace the job's combined `pytest` command with +[`run_cov_split.sh`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/run_cov_split.sh), +which runs each mode, names the reports, uploads them, and fails the step if +either half or the upload failed: + +```yaml +commands: + - | + .buildkite/common/scripts/run_cov_split.sh \ + --model-id \ + --offline \ + --online \ + -- +``` + +Everything after `--` is forwarded verbatim to both runs with quoting preserved, +so an expression like `-m 'advanced_model and cuda'` survives intact and a job can +pass anything else it needs. `--offline`/`--online` are repeatable and only one is +required, so a job with tests in a single mode runs just that half. Runtime is +bounded solely by the step's `timeout_in_minutes` — raise it, since the split +loads the model once per mode. + +Before splitting, confirm each half actually collects at least one test under the job's forwarded filters — pytest exits with code 5 ("no tests collected") if a half is empty, which the script reports as a failure, red-ing a job that used to pass as one combined invocation. + +Also check the tests' `num_cards` against the job's `mirror_hardwares` preset. +`cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a +test asking for more GPUs than the preset provides is silently skipped and its +coverage artifact comes back empty. Splitting such a job produces a file that +looks fine and measures nothing — "Diffusion · Z Image Test" is the current +example (online tests want 4 cards, the job runs on `l4_1`), which is why it is +not a coverage pilot. + +Online-mode coverage depends on the `[tool.coverage.run]` `patch`/`sigterm` +settings in `pyproject.toml`: online tests launch the server with +`subprocess.Popen` and stop it with SIGTERM, and without those settings the XML +reflects only the pytest parent process, not the server's code paths. + +The upload depends on `buildkite-agent` being callable inside the container. The +`kubernetes` presets (`h100_*`, `*_npu_*`) provide it; the `docker` ones (`l4_*`) +only do because they set `mount-buildkite-agent: true`. A new docker preset that +runs a coverage job needs the same. + +List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's +`source_file_dependencies` — both change what the job measures, so without them a +change there is filtered out of normal PR builds and only surfaces in a later +nightly. `tests/buildkite/test_upload_pipeline.py` asserts this for the pilots. +Editing only the surrounding CI YAML still does not schedule them, so a PR that +touches just the wiring needs a full E2E run (or the commands run on a GPU host) +to produce artifacts. When checking a new model's +artifacts, compare `lines-covered` between the online and offline XML rather than +just confirming both files exist. + ### Validation checklist | Check | Command / location | diff --git a/pyproject.toml b/pyproject.toml index 0e91bfae6..48273fd20 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -301,6 +301,21 @@ filterwarnings = [ "ignore:.*does not have '__bases__' attribute.*:UserWarning", ] +[tool.coverage.run] +core = "sysmon" +source = ["vllm_omni"] +# Online-serving tests launch the server via subprocess.Popen and stop it with +# SIGTERM. pytest-cov no longer instruments subprocesses implicitly, and a +# SIGTERM-ed process discards its data, so without patch/sigterm a coverage run +# over an online test measures only the pytest parent. +patch = ["subprocess"] +concurrency = ["multiprocessing", "thread"] +sigterm = true + +[tool.coverage.report] +show_missing = true +skip_covered = true + [tool.typos.default] extend-ignore-identifiers-re = [ ".*_thw", diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 24a2a731e..ee45a64d3 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -5,6 +5,7 @@ import sys from pathlib import Path import pytest +import yaml sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "common" / "scripts")) @@ -139,3 +140,48 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None: assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [ {"name": "swr-secret"}, ] + + +COVERAGE_PILOT_LABELS = ("Diffusion · Bagel Test", "TTS · Qwen3-TTS Base Test") +MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml") + + +def _surviving_labels(changed_files: list[str]) -> set[str]: + doc = yaml.safe_load(MERGE_PIPELINE.read_text(encoding="utf-8")) + rendered = _render_test_pipeline(doc, changed_files=changed_files) + labels: set[str] = set() + + def walk(steps: list | None) -> None: + for step in steps or []: + if not isinstance(step, dict): + continue + if "label" in step: + labels.add(step["label"]) + walk(step.get("steps")) + + walk(rendered.get("steps")) + return labels + + +@pytest.mark.parametrize( + "changed_file", + [ + ".buildkite/common/scripts/run_cov_split.sh", + "pyproject.toml", + ], +) +def test_coverage_pilots_selected_by_their_shared_inputs(changed_file: str) -> None: + """The split helper and the coverage config both change what these jobs measure. + + Without them in source_file_dependencies the jobs are filtered out of normal + PR builds, so a regression would only surface in a later nightly. + """ + labels = _surviving_labels([changed_file]) + for pilot in COVERAGE_PILOT_LABELS: + assert pilot in labels, f"{pilot} dropped when {changed_file} changed" + + +def test_coverage_pilots_not_selected_by_unrelated_change() -> None: + labels = _surviving_labels(["vllm_omni/entrypoints/openai/serving_chat.py"]) + for pilot in COVERAGE_PILOT_LABELS: + assert pilot not in labels From bc73546c065844136b16e3a8e799cc0bd36bc2ca Mon Sep 17 00:00:00 2001 From: TJian Date: Fri, 14 Aug 2026 15:38:47 +0800 Subject: [PATCH 34/41] [ROCm] [CI] Migrate CI to mi300x for v0.27.x (#5886) Signed-off-by: tjtanaa Signed-off-by: TJian Co-authored-by: Hongsheng Liu --- .buildkite/amd/scripts/run-amd-test.sh | 232 +++++++----------- .buildkite/amd/test-amd-merge.yml | 79 +++--- .buildkite/amd/test-amd-ready.yml | 81 +++--- .buildkite/amd/test-template-amd-omni.j2 | 112 ++++++++- docker/Dockerfile.rocm | 11 +- pyproject.toml | 1 + .../ar_diffusion/test_paged_attention.py | 14 +- tests/diffusion/attention/test_flash_attn.py | 13 +- .../distributed/test_sequence_parallel.py | 12 +- .../lora/test_images_generations_lora.py | 16 ++ tests/engine/test_duplex_import_boundary.py | 18 ++ .../openai_api/test_decode_audio_url.py | 3 +- .../ar_diffusion/kv_cache/paged_attention.py | 44 ++++ 13 files changed, 403 insertions(+), 233 deletions(-) diff --git a/.buildkite/amd/scripts/run-amd-test.sh b/.buildkite/amd/scripts/run-amd-test.sh index c49614e70..c1c80c495 100755 --- a/.buildkite/amd/scripts/run-amd-test.sh +++ b/.buildkite/amd/scripts/run-amd-test.sh @@ -1,153 +1,103 @@ #!/bin/bash -# vllm-omni customized version -# Based on: vllm/.buildkite/scripts/hardware_ci/run-amd-test.sh -# Last synced: 2025-12-15 -# Modifications: docker image name for vllm-omni -# This script runs test inside the corresponding ROCm docker container. -set -o pipefail +# Run vLLM-Omni ROCm tests directly in the MI300 Kubernetes pod. The pod's +# container image is selected by test-template-amd-omni.j2; MI300 has no DinD. +set -euo pipefail -# Export Python path -export PYTHONPATH=".." +: "${PYTHONFAULTHANDLER:=1}" +: "${HF_HOME:=/home/buildkite-agent/huggingface}" +: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" +: "${HF_HUB_ETAG_TIMEOUT:=60}" +: "${MIOPEN_DEBUG_CONV_DIRECT:=0}" +: "${MIOPEN_DEBUG_CONV_GEMM:=0}" +: "${VLLM_ROCM_USE_AITER:=0}" +export PYTHONFAULTHANDLER HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT +export MIOPEN_DEBUG_CONV_DIRECT MIOPEN_DEBUG_CONV_GEMM VLLM_ROCM_USE_AITER +export PYTORCH_ROCM_ARCH="" +export PYTHONPATH="${PYTHONPATH:-..}" + +if [[ "${VLLM_CI_DOCKER_DISABLED:-0}" != "1" ]]; then + echo "Error: MI300 CI must run natively with Docker disabled." >&2 + exit 1 +fi + +if [[ -n "${TEST_COMMANDS:-}" ]]; then + commands="${TEST_COMMANDS}" +else + commands="$*" +fi +if [[ -z "${commands}" ]]; then + echo "Error: No test commands provided." >&2 + exit 1 +fi + +echo "--- Native in-pod ROCm CI" + +job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}" +job_id="${job_id//[^A-Za-z0-9_.-]/_}" +native_cache_root="/tmp/vllm-omni-native-${job_id}" +TMPDIR="${native_cache_root}/tmp" +TORCHINDUCTOR_CACHE_DIR="${native_cache_root}/torchinductor" +TRITON_CACHE_DIR="${native_cache_root}/triton" +VLLM_CACHE_ROOT="${native_cache_root}/vllm" +XDG_CACHE_HOME="${native_cache_root}/xdg" +HF_DATASETS_CACHE="${native_cache_root}/huggingface/datasets" +export TMPDIR TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT +export XDG_CACHE_HOME HF_DATASETS_CACHE +mkdir -p "${TMPDIR}" "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" \ + "${VLLM_CACHE_ROOT}" "${XDG_CACHE_HOME}" "${HF_HOME}" \ + "${HF_DATASETS_CACHE}" + +if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then + if ! command -v findmnt >/dev/null 2>&1; then + echo "Error: findmnt is required to verify the Hugging Face cache mount." >&2 + exit 1 + fi + hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true) + if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then + echo "Error: MI300 CI requires a persistent volume at or above ${HF_HOME}." >&2 + exit 1 + fi +fi -# Print ROCm version -echo "--- ROCm info" rocminfo -# cleanup older docker images -cleanup_docker() { - # Get Docker's root directory - docker_root=$(docker info -f '{{.DockerRootDir}}') - if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 - fi - echo "Docker root directory: $docker_root" - # Check disk usage of the filesystem where Docker's root directory is located - disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') - # Define the threshold - threshold=70 - if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=72h" --all - echo "Docker images and volumes cleanup completed." - else - echo "Disk usage is below $threshold%. No cleanup needed." - fi -} +expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}" +python3 - <<'PY' +import os -# Call the cleanup docker function -cleanup_docker +import torch -echo "--- Pulling container" -## Temporary change to use AMD Docker Hub to store the vllm-omni image -# to bypass the rate limit issue with ECR Public Gallery. -# Images are now stored in a separate repository for vllm-omni, instead of vllm-ci. -# TODO: @tjtanaa point back to ECR Public Gallery -# once the amd agents are configured to use ECR Public Gallery. -# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni" -image_name="rocm/vllm-omni:${BUILDKITE_COMMIT}" -container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)" +expected = int(os.environ.get("VLLM_CI_EXPECTED_GPU_COUNT", "1")) +assert torch.version.hip, "PyTorch is not a ROCm build" +assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch" +actual = torch.cuda.device_count() +assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}" +PY -# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery. -# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images -# sudo apt-get update && sudo apt-get install -y awscli -## Use safe docker login helper to prevent race conditions -# source "$(dirname "${BASH_SOURCE[0]}")/../../common/scripts/docker_login_ecr_public.sh" -# safe_docker_login_ecr_public - -## Pull the container from AMD Docker Hub - -docker pull "${image_name}" - -remove_docker_container() { - docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true - docker image rm -f "${image_name}" || true -} -trap remove_docker_container EXIT - -echo "--- Running container" - -HF_CACHE="$(realpath ~)/huggingface" -mkdir -p "${HF_CACHE}" -HF_MOUNT="/root/.cache/huggingface" - -if [[ -n "${TEST_COMMAND:-}" ]]; then - commands="$TEST_COMMAND" -else - commands="$@" -fi -echo "Commands:$commands" - -PARALLEL_JOB_COUNT=8 -MYPYTHONPATH=".." - -# Test that we're launching on the machine that has -# proper access to GPUs -render_gid=$(getent group render | cut -d: -f3) -if [[ -z "$render_gid" ]]; then - echo "Error: 'render' group not found. This is required for GPU access." >&2 - exit 1 -fi - -# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs. -# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels. -if [[ $commands == *"--shard-id="* ]]; then - # assign job count as the number of shards used - commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g') - for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do - # assign shard-id for each shard - commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g') - echo "Shard ${GPU} commands:$commands_gpu" - echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES" - docker run \ - --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ - --network=host \ - --shm-size=16gb \ - --group-add "$render_gid" \ - -e MIOPEN_DEBUG_CONV_DIRECT=0 \ - -e MIOPEN_DEBUG_CONV_GEMM=0 \ - -e VLLM_ROCM_USE_AITER=0 \ - -e HIP_VISIBLE_DEVICES="${GPU}" \ - -e HF_TOKEN \ - -v "${HF_CACHE}:${HF_MOUNT}" \ - -e "HF_HOME=${HF_MOUNT}" \ - -e "PYTHONPATH=${MYPYTHONPATH}" \ - --name "${container_name}_${GPU}" \ - "${image_name}" \ - /bin/bash -c "${commands_gpu}" \ - |& while read -r line; do echo ">>Shard $GPU: $line"; done & - PIDS+=($!) - done - #wait for all processes to finish and collect exit codes - for pid in "${PIDS[@]}"; do - wait "${pid}" - STATUS+=($?) - done - for st in "${STATUS[@]}"; do - if [[ ${st} -ne 0 ]]; then - echo "One of the processes failed with $st" - exit "${st}" +echo "Commands:${commands}" +if /bin/bash -o pipefail -c ' +set -E +test_status=0 +trap '\'' + command_status=$? + # Keep running subsequent commands, but retain a failure for the final exit. + # Prefer a test failure over pytest "no tests collected" when both occur. + if (( test_status == 0 || (test_status == 5 && command_status != 5) )); then + test_status=${command_status} fi - done +'\'' ERR +eval "$1" +exit "${test_status}" +' _ "${commands}"; then + exit 0 else - echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES" - docker run \ - --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ - --network=host \ - --shm-size=16gb \ - --group-add "$render_gid" \ - -e MIOPEN_DEBUG_CONV_DIRECT=0 \ - -e MIOPEN_DEBUG_CONV_GEMM=0 \ - -e VLLM_ROCM_USE_AITER=0 \ - -e HF_TOKEN \ - -v "${HF_CACHE}:${HF_MOUNT}" \ - -e "HF_HOME=${HF_MOUNT}" \ - -e "PYTHONPATH=${MYPYTHONPATH}" \ - --name "${container_name}" \ - "${image_name}" \ - /bin/bash -c "${commands}" + exit_code=$? fi + +if [[ ${exit_code} -eq 5 && "${VLLM_CI_ALLOW_NO_TESTS:-0}" == "1" ]]; then + echo "Pytest collected no tests; VLLM_CI_ALLOW_NO_TESTS=1, treating exit code 5 as success." + exit 0 +fi + +exit "${exit_code}" diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml index 4e80d1968..fbe3b0c0f 100644 --- a/.buildkite/amd/test-amd-merge.yml +++ b/.buildkite/amd/test-amd-merge.yml @@ -7,54 +7,64 @@ steps: - group: ":card_index_dividers: Simple Test" steps: - - label: "Simple · Diffusion & Model Executor Test (1/2)" - agent_pool: mi325_1 + - label: "Simple · Model Executor Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 2 depends_on: amd-build - timeout_in_minutes: 70 mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - # CPU tests on ROCm gpu machines are very slow, so we give them more time. - - "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml" - - label: "Simple · Diffusion & Model Executor Test (2/2)" - agent_pool: mi325_1 + - label: "Simple · Diffusion Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 2 depends_on: amd-build - timeout_in_minutes: 70 mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_ROCM_USE_AITER=0 - # CPU tests on ROCm gpu machines are very slow, so we give them more time. - - "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-diffusion-$$BUILDKITE_PARALLEL_JOB.xml" + + - label: "Simple · Engine&Entrypoints Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" - label: "Simple · Other Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build - timeout_in_minutes: 50 mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" - group: ":card_index_dividers: Diffusion Test" steps: - label: "Diffusion · LoRA Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking + artifact_paths: + - lora-server.log commands: - - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py + - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model" 2>&1 | tee lora-server.log - label: "Diffusion Model CPU offloading Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -62,7 +72,7 @@ steps: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - | - timeout 20m bash -c ' + bash -c ' set +e pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py EXIT1=\$? @@ -72,7 +82,7 @@ steps: ' - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)" - agent_pool: mi325_4 + agent_pool: mi300_4 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -84,7 +94,7 @@ steps: # merge-only tests - label: "Diffusion Tensor Parallelism Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -92,7 +102,7 @@ steps: - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py - label: "Diffusion Model Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -100,7 +110,7 @@ steps: - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py - label: "Diffusion GPU Worker Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -108,7 +118,7 @@ steps: - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py - label: "Diffusion · Wan22 Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -121,7 +131,7 @@ steps: ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 # - label: "Audio Generation Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -132,7 +142,7 @@ steps: # - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model - label: "Engine Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -147,7 +157,7 @@ steps: - group: ":card_index_dividers: E2E Test" steps: - label: "Omni Model Test Qwen2-5-Omni" - agent_pool: mi325_4 + agent_pool: mi300_4 depends_on: amd-build mirror_hardwares: [amdproduction] commands: @@ -155,7 +165,7 @@ steps: - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model" - label: "Omni Model Test Qwen3-Omni" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -163,22 +173,23 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' - - label: "Qwen3-TTS CustomVoice E2E Test" - agent_pool: mi325_1 + - label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - | - timeout 50m bash -c ' + bash -c ' export VLLM_LOGGING_LEVEL=DEBUG export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" + pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB ' - label: "Qwen3-TTS Base E2E Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -196,7 +207,7 @@ steps: # Need to debug before reneable numerical changes across large PRs # # split Bagel Model Test with H100 (Real Weights) into three tests # - label: "Bagel Text2Img Model Test (1/3)" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -208,7 +219,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm" # - label: "Bagel Img2Img Model Test (2/3)" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -220,7 +231,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm" # - label: "Bagel Online Serving Test (3/3)" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 8f99ce909..b8fca2f88 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -7,8 +7,9 @@ steps: - group: ":card_index_dividers: Simple Test" steps: - - label: "Simple · Diffusion & Model Executor Test" - agent_pool: mi325_1 + - label: "Simple · Model Executor Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -16,10 +17,30 @@ steps: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml" + + - label: "Simple · Diffusion Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + + - label: "Simple · Engine&Entrypoints Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" - label: "Simple · Other Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -27,19 +48,19 @@ steps: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" - label: "Custom Pipeline Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" + - pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" - label: "Voxtral TTS CUDA Unit Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -49,33 +70,25 @@ steps: - group: ":card_index_dividers: Diffusion Test" steps: - label: "Diffusion · Batch Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model" + - pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model" - label: "Diffusion · Offloader Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - | - timeout 20m bash -c ' - set +e - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py - EXIT1=\$? - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py - EXIT2=\$? - exit \$((EXIT1 | EXIT2)) - ' + - pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model" - label: "Diffusion · Cache Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -83,10 +96,10 @@ steps: - export GPU_ARCHS=gfx942 - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4" + - pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model" - label: "Diffusion Sequence Parallelism Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -96,7 +109,7 @@ steps: - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model - label: "Diffusion GPU Worker Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -104,7 +117,7 @@ steps: - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py - label: "Diffusion · Model Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -112,7 +125,7 @@ steps: - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model" - label: "Diffusion · Wan22 Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -126,7 +139,7 @@ steps: ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 # - label: "Audio Generation Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -137,7 +150,7 @@ steps: # - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model - label: "Engine Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -149,7 +162,7 @@ steps: ' # - label: "Omni Model Test Qwen3-Omni" -# agent_pool: mi325_2 +# agent_pool: mi300_2 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -162,7 +175,7 @@ steps: - group: ":card_index_dividers: E2E Test" steps: - label: "Qwen3-TTS E2E Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -173,7 +186,7 @@ steps: - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model" - label: "CosyVoice3-TTS E2E Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -186,7 +199,7 @@ steps: # TODO: Bagel test on ROCm is very unstable. @tjtanaa # Need to debug before reneable numerical changes across large PRs # - label: "Bagel Text2Img Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -198,7 +211,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm" # - label: "Bagel Img2Img Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -210,7 +223,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm" # - label: "Bagel Online Serving Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking diff --git a/.buildkite/amd/test-template-amd-omni.j2 b/.buildkite/amd/test-template-amd-omni.j2 index 2f1a9aed3..e2042b96a 100644 --- a/.buildkite/amd/test-template-amd-omni.j2 +++ b/.buildkite/amd/test-template-amd-omni.j2 @@ -1,11 +1,65 @@ {# vllm-omni customized version Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/test-template-amd.j2 - Last synced: 2025-12-15 - Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests + Last synced: 2026-08-04 + Modifications: Use the shared MI300 queues and Omni command-passing contract #} {% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %} {% set default_working_dir = "/app/vllm-omni" %} +{# MI300 runs tests natively in Kubernetes; Docker-in-Docker is unavailable. #} +{% macro amd_native_k8s_plugin(indent, step) -%} +{{ indent }}plugins: +{{ indent }} - kubernetes: +{{ indent }} podSpecPatch: +{{ indent }} automountServiceAccountToken: false +{{ indent }} securityContext: +{{ indent }} seccompProfile: +{{ indent }} type: RuntimeDefault +{{ indent }} imagePullSecrets: +{{ indent }} - name: docker-config +{{ indent }} containers: +{{ indent }} - name: container-0 +{{ indent }} image: {{ docker_image_amd }} +{{ indent }} imagePullPolicy: Always +{{ indent }} securityContext: +{{ indent }} allowPrivilegeEscalation: false +{{ indent }} capabilities: +{{ indent }} add: +{{ indent }} - IPC_LOCK +{{ indent }} resources: +{{ indent }} limits: +{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}" +{{ indent }} requests: +{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}" +{{ indent }} volumeMounts: +{{ indent }} - name: devshm +{{ indent }} mountPath: /dev/shm +{{ indent }} env: +{{ indent }} - name: AMD_CI_RUNTIME +{{ indent }} value: native +{{ indent }} - name: NATIVE_CI +{{ indent }} value: "true" +{{ indent }} - name: VLLM_CI_DOCKER_DISABLED +{{ indent }} value: "1" +{{ indent }} - name: VLLM_CI_EXPECTED_GPU_COUNT +{{ indent }} value: "{{ step.agent_pool[-1:] }}" +{{ indent }} - name: HF_HOME +{{ indent }} value: /home/buildkite-agent/huggingface +{{ indent }} - name: MIOPEN_DEBUG_CONV_DIRECT +{{ indent }} value: "0" +{{ indent }} - name: MIOPEN_DEBUG_CONV_GEMM +{{ indent }} value: "0" +{{ indent }} - name: VLLM_ROCM_USE_AITER +{{ indent }} value: "0" +{{ indent }} - name: PYTORCH_ROCM_ARCH +{{ indent }} value: "" +{{ indent }} volumes: +{{ indent }} - name: devshm +{{ indent }} emptyDir: +{{ indent }} medium: Memory +{{ indent }} sizeLimit: 16Gi +{%- endmacro %} + {% if env %} env: {% for name in env %} @@ -53,19 +107,31 @@ steps: {% if step.agent_pool %} queue: amd_{{ step.agent_pool }} {% else %} - queue: amd_mi325_1 + queue: amd_mi300_1 {% endif %} +{{ amd_native_k8s_plugin(" ", step) }} {% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %} {% set indented_cmd = cmd_body | replace("\n", "\n ") %} command: bash .buildkite/amd/scripts/run-amd-test.sh env: - DOCKER_BUILDKIT: "1" + AMD_CI_RUNTIME: "native" + NATIVE_CI: "true" + VLLM_CI_DOCKER_DISABLED: "1" + VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}" + VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1" + HF_HOME: "/home/buildkite-agent/huggingface" + HF_HUB_DOWNLOAD_TIMEOUT: "300" + HF_HUB_ETAG_TIMEOUT: "60" + MIOPEN_DEBUG_CONV_DIRECT: "0" + MIOPEN_DEBUG_CONV_GEMM: "0" + VLLM_ROCM_USE_AITER: "0" + PYTORCH_ROCM_ARCH: "" {% if step.env %} {% for name in step.env %} {{ name }}: "{{ step.env[name] }}" {% endfor %} {% endif %} - TEST_COMMAND: |- + TEST_COMMANDS: |- (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} {% if env %} {% for name in env %} @@ -82,6 +148,15 @@ steps: {% endif %} {{ indented_cmd | safe }} priority: 100 +{% if step.parallelism %} + parallelism: {{ step.parallelism }} +{% endif %} +{% if step.artifact_paths %} + artifact_paths: +{% for path in step.artifact_paths %} + - "{{ path }}" +{% endfor %} +{% endif %} {% if step.timeout_in_minutes %} timeout_in_minutes: {{ step.timeout_in_minutes }} {% endif %} @@ -100,19 +175,31 @@ steps: {% if step.agent_pool %} queue: amd_{{ step.agent_pool }} {% else %} - queue: amd_mi325_1 + queue: amd_mi300_1 {% endif %} +{{ amd_native_k8s_plugin(" ", step) }} {% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %} {% set indented_cmd = cmd_body | replace("\n", "\n ") %} command: bash .buildkite/amd/scripts/run-amd-test.sh env: - DOCKER_BUILDKIT: "1" + AMD_CI_RUNTIME: "native" + NATIVE_CI: "true" + VLLM_CI_DOCKER_DISABLED: "1" + VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}" + VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1" + HF_HOME: "/home/buildkite-agent/huggingface" + HF_HUB_DOWNLOAD_TIMEOUT: "300" + HF_HUB_ETAG_TIMEOUT: "60" + MIOPEN_DEBUG_CONV_DIRECT: "0" + MIOPEN_DEBUG_CONV_GEMM: "0" + VLLM_ROCM_USE_AITER: "0" + PYTORCH_ROCM_ARCH: "" {% if step.env %} {% for name in step.env %} {{ name }}: "{{ step.env[name] }}" {% endfor %} {% endif %} - TEST_COMMAND: |- + TEST_COMMANDS: |- (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} {% if env %} {% for name in env %} @@ -129,6 +216,15 @@ steps: {% endif %} {{ indented_cmd | safe }} priority: 100 +{% if step.parallelism %} + parallelism: {{ step.parallelism }} +{% endif %} +{% if step.artifact_paths %} + artifact_paths: +{% for path in step.artifact_paths %} + - "{{ path }}" +{% endfor %} +{% endif %} {% if step.timeout_in_minutes %} timeout_in_minutes: {{ step.timeout_in_minutes }} {% endif %} diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 6b5a12309..4283efc77 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0 FROM ${BASE_IMAGE} AS base # Declare a variable to know if we want to use the nightly build or the stable build. @@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base # we should swap over to use stable release ASAP. # We should avoid relying on custom commits. ARG USE_NIGHTLY_BUILD=0 -ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0 +ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0 ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} @@ -79,9 +79,16 @@ RUN printf '%s\n' \ ' os._exit(_exit_code)' \ > ${COMMON_WORKDIR}/vllm-omni/conftest.py +# Some optional dependencies pull in the CPU ONNX Runtime distribution. The +# CPU and ROCm wheels own the same import path and cannot coexist, so remove +# the CPU wheel and reinstall the ROCm wheel after dependency resolution. RUN cd ${COMMON_WORKDIR}/vllm-omni && \ uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \ --no-cache-dir ".[dev]" --no-build-isolation && \ + python3 -m pip uninstall -y onnxruntime && \ + python3 -m pip install --no-cache-dir --force-reinstall --no-deps \ + "onnxruntime-rocm>=1.22.2" && \ + python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \ python3 -c "from torchcodec.decoders import VideoDecoder" RUN ln -sf /usr/bin/python3 /usr/bin/python diff --git a/pyproject.toml b/pyproject.toml index 48273fd20..160c6bc2f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ dev = [ "pytest-asyncio==1.4.0", "pytest-cov==7.1.0", "pytest-mock==3.15.1", + "pytest-shard==0.1.2", "websockets==16.1", "pytest-subtests==0.15.0", "datasets==3.6.0", diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py index 0b66c01a0..27135874f 100644 --- a/tests/diffusion/ar_diffusion/test_paged_attention.py +++ b/tests/diffusion/ar_diffusion/test_paged_attention.py @@ -60,9 +60,18 @@ def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor return torch.einsum("bhqk,bkhd->bqhd", probs, value) -def _cuda_flash_attn_usable() -> bool: +def _gpu_flash_attn_usable() -> bool: if not torch.cuda.is_available(): return False + if torch.version.hip is not None: + for module in ("aiter", "flash_attn"): + try: + imported = __import__(module, fromlist=["flash_attn_varlen_func"]) + if getattr(imported, "flash_attn_varlen_func", None) is not None: + return True + except ImportError: + pass + return False try: spec = find_spec("vllm.vllm_flash_attn") if spec is None or spec.origin is None: @@ -227,12 +236,11 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0) -@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required") +@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required") @pytest.mark.parametrize("history_chunks", [1, 3]) @pytest.mark.parametrize("action_len", [0, 3]) @pytest.mark.parametrize("commit_current", [False, True]) def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current): - pytest.importorskip("vllm.vllm_flash_attn") torch.manual_seed(0) device = torch.device("cuda") dtype = torch.float16 diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index e3f5355bf..5924beabf 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -265,10 +265,15 @@ def test_fa_vs_sdpa(): print(f"Max absolute difference (valid region): {max_diff:.6f}") print(f"Mean absolute difference (valid region): {mean_diff:.6f}") - # Assert that outputs are close - # Using higher tolerance for bfloat16 and different implementations - assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01" - assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001" + # AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by + # about 1.3% relative error. An absolute-only bound misclassifies that + # expected rounding difference for larger output values (for example, + # -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere. + if current_omni_platform.is_rocm(): + torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2) + else: + assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01" + assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001" print("✓ Case 2 PASSED: FA and SDPA outputs are very close!") diff --git a/tests/diffusion/distributed/test_sequence_parallel.py b/tests/diffusion/distributed/test_sequence_parallel.py index d43d18c49..f083e1d26 100644 --- a/tests/diffusion/distributed/test_sequence_parallel.py +++ b/tests/diffusion/distributed/test_sequence_parallel.py @@ -416,8 +416,8 @@ def test_sp_correctness_advanced(model_name: str): @pytest.mark.skipif( - not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), - reason="Only tested on CUDA and XPU", + not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA, ROCM and XPU", ) @pytest.mark.diffusion @pytest.mark.parallel @@ -447,8 +447,8 @@ def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.skipif( - not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), - reason="Only tested on CUDA and XPU", + not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA, ROCM and XPU", ) @pytest.mark.diffusion @pytest.mark.parallel @@ -468,8 +468,8 @@ def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> No @pytest.mark.skipif( - not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), - reason="Only tested on CUDA and XPU", + not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA, ROCM and XPU", ) @pytest.mark.diffusion @pytest.mark.parallel diff --git a/tests/diffusion/lora/test_images_generations_lora.py b/tests/diffusion/lora/test_images_generations_lora.py index dbdf59647..a01bcdae3 100644 --- a/tests/diffusion/lora/test_images_generations_lora.py +++ b/tests/diffusion/lora/test_images_generations_lora.py @@ -24,11 +24,20 @@ from safetensors.torch import save_file from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer, OmniServerParams +from vllm_omni.platforms import current_omni_platform os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" MODEL = "Tongyi-MAI/Z-Image-Turbo" +# ROCm's default FLASH_ATTN diffusion backend resolves to AITER. On a cold CI +# worker, JIT-building its MHA kernel can take longer than the server readiness +# timeout. This test validates LoRA switching rather than attention-backend or +# compiled-execution performance, so use the eager SDPA path on ROCm only. +ROCM_SERVER_ARGS = ( + ["--enforce-eager", "--diffusion-attention-backend", "TORCH_SDPA"] if current_omni_platform.is_rocm() else [] +) + PROMPT = "a photo of a cat sitting on a laptop keyboard" SIZE = "256x256" @@ -42,7 +51,14 @@ server_params = [ server_args=[ "--num-gpus", "1", + *ROCM_SERVER_ARGS, ], + # A cold-cache Z-Image load can take slightly more than the + # orchestrator's 900s default on CI (weights alone have been + # observed at ~690s and full orchestrator readiness at ~906s). + # Match the large-model offline runner's startup allowance without + # weakening timeout detection for every online-serving test. + init_timeout=1800, ), id="zimage_turbo_lora", ), diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py index 6eae2e698..d04ed2e70 100644 --- a/tests/engine/test_duplex_import_boundary.py +++ b/tests/engine/test_duplex_import_boundary.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os import subprocess import sys from pathlib import Path @@ -14,9 +15,18 @@ REPO_ROOT = Path(__file__).resolve().parents[2] def _assert_isolated_import_succeeds(script: str) -> None: + # These tests validate Python import/package boundaries and do not use a + # GPU. Other tests in the combined entrypoints/engine suite may temporarily + # configure only one of the CUDA/ROCm visibility variables. Do not let an + # unrelated, inconsistent device-selection environment prevent the child + # interpreter from reaching the boundary assertion. + env = os.environ.copy() + env.pop("CUDA_VISIBLE_DEVICES", None) + env.pop("HIP_VISIBLE_DEVICES", None) result = subprocess.run( [sys.executable, "-c", script], cwd=REPO_ROOT, + env=env, capture_output=True, text=True, timeout=60, @@ -25,6 +35,14 @@ def _assert_isolated_import_succeeds(script: str) -> None: assert result.returncode == 0, result.stderr +def test_isolated_import_ignores_conflicting_gpu_visibility(monkeypatch: pytest.MonkeyPatch) -> None: + """Import-boundary checks must not depend on GPU state leaked by prior tests.""" + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + + _assert_isolated_import_succeeds("import vllm_omni.outputs") + + def test_stable_engine_imports_do_not_load_experimental_duplex() -> None: _assert_isolated_import_succeeds(""" import sys diff --git a/tests/entrypoints/openai_api/test_decode_audio_url.py b/tests/entrypoints/openai_api/test_decode_audio_url.py index 2a5c5be47..5dff3b052 100644 --- a/tests/entrypoints/openai_api/test_decode_audio_url.py +++ b/tests/entrypoints/openai_api/test_decode_audio_url.py @@ -5,6 +5,7 @@ import asyncio import base64 import os +import tempfile import httpx import pytest @@ -101,7 +102,7 @@ class TestDecodeAudioUrlSuffixSanitization: try: # ".." contains non-alnum dots, so falls back to .wav assert path.endswith(".wav") - assert os.path.dirname(path) == "/tmp" + assert os.path.dirname(path) == tempfile.gettempdir() finally: _cleanup(path) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py index bcd06b63a..f16ecc9e9 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py @@ -365,6 +365,22 @@ def _resolve_fa_version(head_size: int) -> int: return version +def _rocm_flash_attn_varlen_func(): + """Resolve a ROCm varlen kernel, preferring AITER when available. + + The caller gathers paged KV into packed tensors before invoking this + function, avoiding AITER releases whose ``block_table`` kernel is broken. + """ + try: + from aiter import flash_attn_varlen_func + + return flash_attn_varlen_func + except ImportError: + from flash_attn import flash_attn_varlen_func + + return flash_attn_varlen_func + + def ar_diffusion_paged_attention( query: torch.Tensor, key_cache: torch.Tensor, @@ -401,6 +417,34 @@ def ar_diffusion_paged_attention( softmax_scale, causal=causal, ) + elif torch.version.hip is not None: + # vllm.vllm_flash_attn contains CUDA-only extensions. ROCm's AITER and + # upstream flash-attn expose the standard cu_seqlens_k API instead of + # vLLM's seqused_k/fa_version API. The ROCm flash-attn paged kernel also + # requires 128-token blocks, while AR-Diffusion uses frame-aligned + # 16-token blocks, so gather the visible blocks on-device first. + flash_attn_varlen_func = _rocm_flash_attn_varlen_func() + cu_seqlens_k = torch.cat([seq_lens.new_zeros(1), torch.cumsum(seq_lens, dim=0, dtype=torch.int32)]) + positions = torch.arange(int(max_seq_len), device=query_flat.device) + logical_blocks = torch.div(positions, key_cache.shape[1], rounding_mode="floor") + offsets = positions % key_cache.shape[1] + physical_blocks = block_table[:, logical_blocks].long() + gathered_k = key_cache[physical_blocks, offsets] + gathered_v = value_cache[physical_blocks, offsets] + valid = positions.unsqueeze(0) < seq_lens.unsqueeze(1) + packed_k = gathered_k[valid] + packed_v = gathered_v[valid] + out = flash_attn_varlen_func( + q=query_flat, + k=packed_k, + v=packed_v, + cu_seqlens_q=query_start_loc, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=int(max_query_len), + max_seqlen_k=int(max_seq_len), + softmax_scale=float(softmax_scale), + causal=causal, + ) else: from vllm.vllm_flash_attn import flash_attn_varlen_func From b300f666568e6b83a1e8893a3e13fa71afa88aa5 Mon Sep 17 00:00:00 2001 From: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com> Date: Fri, 14 Aug 2026 03:45:55 -0400 Subject: [PATCH 35/41] [Bugfix] Fix HSDP compatibility with the new online FP8 linear method (#5677) Signed-off-by: roG0d --- docs/user_guide/diffusion/parallelism/hsdp.md | 36 +++++++++++++++++++ tests/diffusion/quantization/test_hsdp_fp8.py | 34 ++++++++++++++++-- vllm_omni/diffusion/quantization/hsdp_fp8.py | 15 +++++--- 3 files changed, 79 insertions(+), 6 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism/hsdp.md b/docs/user_guide/diffusion/parallelism/hsdp.md index 96a357c86..8a8aa5ef9 100644 --- a/docs/user_guide/diffusion/parallelism/hsdp.md +++ b/docs/user_guide/diffusion/parallelism/hsdp.md @@ -64,6 +64,41 @@ omni = Omni( ) ``` +### Combined with Online FP8 Quantization + +HSDP can be combined with online FP8 quantization for diffusion models that +support both features. HSDP shards the quantized weights across GPUs to reduce +the per-GPU weight memory footprint. + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="", + quantization="fp8", + parallel_config=DiffusionParallelConfig( + use_hsdp=True, + hsdp_shard_size=2, + ), +) +``` + +For online serving: + +```bash +vllm serve \ + --omni \ + --quantization fp8 \ + --use-hsdp \ + --hsdp-shard-size 2 +``` + +The model must support both HSDP and online FP8 quantization. HSDP remains +incompatible with Tensor Parallelism. See the +[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and +hardware. + --- ## Example Script @@ -147,3 +182,4 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co 1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models 2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism 3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled +4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py index 482235ff2..ce939dbf3 100644 --- a/tests/diffusion/quantization/test_hsdp_fp8.py +++ b/tests/diffusion/quantization/test_hsdp_fp8.py @@ -36,15 +36,23 @@ class _ToyKernel: # --- helper to build a toy module simulating online-FP8 post-load state --- -def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32): +def _make_fp8_toy_module( + out_features: int = 16, + in_features: int = 32, + *, + quant_method_cls=None, +): """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose ``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``. """ from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod + if quant_method_cls is None: + quant_method_cls = Fp8LinearMethod + # Create the quant-method instance without calling __init__ (avoids # heavyweight upstream config dependencies). isinstance() still works. - qm = object.__new__(Fp8LinearMethod) + qm = object.__new__(quant_method_cls) qm.fp8_linear = _ToyKernel() # randn does not support float8 on CPU; zeros does and the test only @@ -104,6 +112,28 @@ def test_rewrites_non_contiguous_weight(): assert module.weight.shape == (16, 32) +def test_rewrites_per_tensor_online_fp8_weight(): + """Regression test for vLLM's new online FP8 frontend (#45463).""" + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + module = _make_fp8_toy_module( + 16, + 32, + quant_method_cls=Fp8PerTensorOnlineLinearMethod, + ) + + n_rewritten_layers = prepare_fp8_layers_for_fsdp(module) + + assert n_rewritten_layers == 1 + assert module.weight.is_contiguous() + assert module.weight.shape == (16, 32) + w, *_ = module.quant_method.fp8_linear._get_layer_params(module) + assert w.shape == (32, 16) + assert w.stride() == (1, 32) + + def test_patched_get_layer_params_returns_transposed_view(): module = _make_fp8_toy_module(16, 32) prepare_fp8_layers_for_fsdp(module) diff --git a/vllm_omni/diffusion/quantization/hsdp_fp8.py b/vllm_omni/diffusion/quantization/hsdp_fp8.py index 7c4666e29..6ffa8f2b4 100644 --- a/vllm_omni/diffusion/quantization/hsdp_fp8.py +++ b/vllm_omni/diffusion/quantization/hsdp_fp8.py @@ -33,9 +33,17 @@ import types from torch import nn from vllm.logger import init_logger from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod +from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, +) logger = init_logger(__name__) +_FP8_TRANSPOSED_WEIGHT_METHODS = ( + Fp8LinearMethod, + Fp8PerTensorOnlineLinearMethod, +) + def _build_transposed_get_layer_params(original_bound_method): """Return a ``_get_layer_params`` replacement that transposes ``w``. @@ -55,9 +63,8 @@ def _build_transposed_get_layer_params(original_bound_method): def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int: """Make online-FP8 linear layers in ``model`` FSDP2-compatible. - For every layer whose ``quant_method`` is an :class:`Fp8LinearMethod` - and whose weight is currently a non-contiguous transpose view, this - function: + For every layer whose quantization method stores FP8 weights as a + non-contiguous transpose view, this function: 1. Replaces ``layer.weight`` with the underlying ``(out, in)`` row-major contiguous storage so FSDP2 ``fully_shard`` accepts it. @@ -76,7 +83,7 @@ def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int: patched_kernel_ids: set[int] = set() for module in model.modules(): qm = getattr(module, "quant_method", None) - if not isinstance(qm, Fp8LinearMethod): + if not isinstance(qm, _FP8_TRANSPOSED_WEIGHT_METHODS): continue weight = getattr(module, "weight", None) From 022d7a7695ffc3ce7329c38656e388cc11aecb8d Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Fri, 14 Aug 2026 15:49:26 +0800 Subject: [PATCH 36/41] [Perf][Diffusion] Avoid redundant MiniMax-H3 reference video scans (#6064) Signed-off-by: Xiaodong Ye --- .../minimax_h3/test_minimax_h3_contract.py | 210 +++++++++++++++++- .../models/minimax_h3/reference_video.py | 133 ++++++++--- 2 files changed, 300 insertions(+), 43 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 60ed7090f..f872e381f 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +import sys from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace from unittest.mock import Mock @@ -1428,34 +1429,219 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti assert transcode_calls[0][1]["target_frame_count"] == 209 -def test_ref2va_qwen_sampling_decodes_prepared_video_once(monkeypatch): +def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module - decoded = np.arange(25 * 2 * 2 * 3, dtype=np.uint8).reshape(25, 2, 2, 3) - load_calls = [] + decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3) + decode_calls = [] monkeypatch.setattr( reference_video_module, "_probe_video", - lambda _path: {"frame_count": len(decoded)}, + lambda _path: { + "frame_count": 25, + "width": 2, + "height": 2, + }, ) monkeypatch.setattr( reference_video_module, - "load_video_frames", - lambda path: load_calls.append(path) or decoded, + "_decode_video_frames_ffmpeg", + lambda path, **kwargs: decode_calls.append((path, kwargs)) or decoded, ) sampled = reference_video_module.sample_reference_video_frames("prepared.mp4") - assert load_calls == ["prepared.mp4"] - for frame, expected in zip( - sampled["frames"], - (decoded[0], decoded[12], decoded[24]), - strict=True, - ): + assert decode_calls == [ + ( + "prepared.mp4", + { + "frame_count": 25, + "indices": [0, 12, 24], + "width": 2, + "height": 2, + }, + ) + ] + for frame, expected in zip(sampled["frames"], decoded, strict=True): np.testing.assert_array_equal(frame, expected) assert sampled["block_timestamps"] == [0.25, 1.0] +@pytest.mark.parametrize( + ("indices", "frame_count"), + [([0, 12], 25), (None, 2)], +) +def test_ref2va_ffmpeg_decode_reads_exact_rgb_frames(monkeypatch, indices, frame_count): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + expected = np.arange(2 * 2 * 3 * 3, dtype=np.uint8).reshape(2, 2, 3, 3) + observed = [] + + def fake_run(command, **kwargs): + observed.append((command, kwargs)) + return SimpleNamespace(stdout=expected.tobytes()) + + monkeypatch.setattr(reference_video_module.subprocess, "run", fake_run) + + actual = reference_video_module._decode_video_frames_ffmpeg( + "prepared.mp4", + frame_count=frame_count, + indices=indices, + width=3, + height=2, + ) + + np.testing.assert_array_equal(actual, expected) + command, kwargs = observed[0] + assert command[0] == "ffmpeg" + assert command[command.index("-frames:v") + 1] == "2" + if indices is None: + assert "-vf" not in command + else: + assert command[command.index("-vf") + 1] == "select=eq(n\\,0)+eq(n\\,12)" + assert command[-1] == "pipe:1" + assert kwargs == {"check": True, "capture_output": True} + + +def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + expected = np.zeros((2, 2, 3, 3), dtype=np.uint8) + decode_calls = [] + monkeypatch.setitem(sys.modules, "decord", None) + monkeypatch.setattr( + reference_video_module, + "_probe_video", + lambda _path: { + "frame_count": 2, + "width": 3, + "height": 2, + }, + ) + monkeypatch.setattr( + reference_video_module, + "_decode_video_frames_ffmpeg", + lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected, + ) + + actual = reference_video_module.load_video_frames("prepared.mp4") + + assert actual is expected + assert decode_calls == [ + ( + "prepared.mp4", + { + "frame_count": 2, + "width": 3, + "height": 2, + }, + ) + ] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"frame_count": 0, "width": 3, "height": 2}, "video has no frames"), + ({"frame_count": 1, "width": 0, "height": 2}, "invalid dimensions"), + ( + {"frame_count": 1, "indices": [], "width": 3, "height": 2}, + "invalid frame indices", + ), + ], +) +def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError, match=message): + reference_video_module._decode_video_frames_ffmpeg("prepared.mp4", **kwargs) + + +def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + calls = [] + probe = { + "streams": [ + { + "codec_type": "video", + "width": 4, + "height": 2, + "r_frame_rate": "24/1", + "nb_frames": "3", + "duration": "0.125", + "sample_aspect_ratio": "1:1", + } + ], + "format": {"duration": "0.125", "size": "123"}, + } + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)), + ) + + metadata = reference_video_module._probe_video("prepared.mp4") + + assert metadata["frame_count"] == 3 + assert len(calls) == 1 + assert "-count_frames" not in calls[0][0] + + +def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + calls = [] + stream = { + "codec_type": "video", + "width": 4, + "height": 2, + "r_frame_rate": "24/1", + "duration": "0.125", + "sample_aspect_ratio": "1:1", + } + first_probe = { + "streams": [stream], + "format": {"duration": "0.125", "size": "123"}, + } + second_probe = { + "streams": [{**stream, "nb_read_frames": "3"}], + "format": {"duration": "0.125", "size": "123"}, + } + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda command, **kwargs: calls.append((command, kwargs)) + or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)), + ) + + metadata = reference_video_module._probe_video("prepared.mp4") + + assert metadata["frame_count"] == 3 + assert len(calls) == 2 + assert "-count_frames" in calls[1][0] + + +def test_ref2va_ffmpeg_decode_rejects_incomplete_output(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + from vllm_omni.errors import OmniClientError + + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace(stdout=b"short"), + ) + + with pytest.raises(OmniClientError, match="decoded 5 bytes.*expected 18"): + reference_video_module._decode_video_frames_ffmpeg( + "prepared.mp4", + frame_count=1, + width=3, + height=2, + ) + + def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module diff --git a/vllm_omni/diffusion/models/minimax_h3/reference_video.py b/vllm_omni/diffusion/models/minimax_h3/reference_video.py index 0ecd69f31..18c8de5b4 100644 --- a/vllm_omni/diffusion/models/minimax_h3/reference_video.py +++ b/vllm_omni/diffusion/models/minimax_h3/reference_video.py @@ -44,31 +44,33 @@ def _nearest_multiple(value: float, multiple: int) -> int: def _probe_video(path: str) -> dict[str, Any]: - result = subprocess.run( - [ - "ffprobe", - "-v", - "error", - "-count_frames", - "-show_entries", - ( - "stream=codec_type,codec_name,width,height,r_frame_rate,nb_read_frames,nb_frames,duration," - "sample_aspect_ratio:format=format_name,size,duration" - ), - "-of", - "json", - path, - ], - check=True, - capture_output=True, - text=True, + show_entries = ( + "stream=codec_type,codec_name,width,height,r_frame_rate,nb_read_frames,nb_frames,duration," + "sample_aspect_ratio:format=format_name,size,duration" ) + command = ["ffprobe", "-v", "error", "-show_entries", show_entries, "-of", "json", path] + result = subprocess.run(command, check=True, capture_output=True, text=True) probe = json.loads(result.stdout) streams = probe.get("streams") or [] video_streams = [stream for stream in streams if stream.get("codec_type") == "video"] if not video_streams: raise OmniClientError(f"media has no video stream: {path}") stream = video_streams[0] + raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") + if raw_count in (None, "", "N/A"): + result = subprocess.run( + ["ffprobe", "-v", "error", "-count_frames", *command[3:]], + check=True, + capture_output=True, + text=True, + ) + probe = json.loads(result.stdout) + streams = probe.get("streams") or [] + video_streams = [stream for stream in streams if stream.get("codec_type") == "video"] + if not video_streams: + raise OmniClientError(f"media has no video stream: {path}") + stream = video_streams[0] + raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") try: numerator, denominator = str(stream["r_frame_rate"]).split("/", 1) fps = float(numerator) / float(denominator) @@ -76,7 +78,6 @@ def _probe_video(path: str) -> dict[str, Any]: raise OmniClientError(f"cannot determine video FPS: {path}") from exc if not math.isfinite(fps) or fps <= 0: raise OmniClientError(f"cannot determine video FPS: {path}") - raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") if raw_count in (None, "", "N/A"): raise OmniClientError(f"cannot determine video frame count: {path}") raw_duration = stream.get("duration") @@ -404,13 +405,13 @@ def load_video_frames(path: str) -> np.ndarray: try: import decord except ImportError: - import av - - with av.open(path) as container: - frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)] - if not frames: - raise OmniClientError(f"video has no frames: {path}") - return np.stack(frames) + meta = _probe_video(path) + return _decode_video_frames_ffmpeg( + path, + frame_count=int(meta["frame_count"]), + width=int(meta["width"]), + height=int(meta["height"]), + ) reader = decord.VideoReader(path) if len(reader) <= 0: @@ -419,6 +420,70 @@ def load_video_frames(path: str) -> np.ndarray: return frames.asnumpy() if hasattr(frames, "asnumpy") else np.asarray(frames) +def _decode_video_frames_ffmpeg( + path: str, + *, + frame_count: int, + width: int, + height: int, + indices: list[int] | None = None, +) -> np.ndarray: + frame_count = int(frame_count) + width = int(width) + height = int(height) + if frame_count <= 0: + raise OmniClientError(f"video has no frames: {path}") + if width <= 0 or height <= 0: + raise OmniClientError(f"video has invalid dimensions {width}x{height}: {path}") + if indices is not None and (not indices or any(index < 0 or index >= frame_count for index in indices)): + raise OmniClientError(f"invalid frame indices for {path}: {indices}") + + output_frame_count = frame_count if indices is None else len(indices) + command = [ + "ffmpeg", + "-loglevel", + "error", + "-threads", + "0", + "-i", + path, + "-map", + "0:v:0", + "-an", + ] + if indices is not None: + select = "+".join(f"eq(n\\,{index})" for index in indices) + command.extend(["-vf", f"select={select}"]) + command.extend( + [ + "-vsync", + "0", + "-frames:v", + str(output_frame_count), + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "pipe:1", + ] + ) + result = subprocess.run( + command, + check=True, + capture_output=True, + ) + frame_size = width * height * 3 + expected_size = output_frame_count * frame_size + if len(result.stdout) != expected_size: + raise OmniClientError(f"decoded {len(result.stdout)} bytes from {path}, expected {expected_size}") + return np.frombuffer(result.stdout, dtype=np.uint8).reshape( + output_frame_count, + height, + width, + 3, + ) + + def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: meta = _probe_video(prepared_path) ratio = MINIMAX_H3_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS @@ -434,11 +499,17 @@ def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: if not indices: raise OmniClientError(f"no frames sampled from {prepared_path}") - # The preparation step emits a lossless RGB stream. Decode it once and - # sample from that array so Qwen3VL sees exactly the same pixels as the - # video VAE, without starting one full-video ffmpeg decode per sample. - decoded_frames = load_video_frames(prepared_path) - frames = [np.asarray(decoded_frames[source_index]) for source_index in indices] + # The preparation step emits a lossless RGB stream. Decode only the sampled + # frames in one ffmpeg process so Qwen3VL sees the exact prepared pixels + # without decoding the entire high-bitrate stream into host memory. + decoded_frames = _decode_video_frames_ffmpeg( + prepared_path, + frame_count=int(meta["frame_count"]), + indices=indices, + width=meta["width"], + height=meta["height"], + ) + frames = [np.asarray(frame) for frame in decoded_frames] timestamps = [index / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))] timestamps += [timestamps[-1]] * ((-len(timestamps)) % MINIMAX_H3_QWEN_TEMPORAL_PATCH) From 728640f4d9bb7c1b68d646e2a4c59ce1ce45de9c Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Fri, 14 Aug 2026 15:51:23 +0800 Subject: [PATCH 37/41] [Bugfix][MUSA][Qwen Image] Avoid complex RoPE alias guards (#6110) Signed-off-by: Xiaodong Ye --- .../diffusion/models/qwen_image/qwen_image_transformer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 67e2d0db7..c82ace9c4 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -642,10 +642,10 @@ class QwenImageCrossAttention(nn.Module): txt_query = self.norm_added_q(txt_query) txt_key = self.norm_added_k(txt_key) - img_cos = vid_freqs.real.to(img_query.dtype) - img_sin = vid_freqs.imag.to(img_query.dtype) - txt_cos = txt_freqs.real.to(txt_query.dtype) - txt_sin = txt_freqs.imag.to(txt_query.dtype) + img_cos = torch.real(vid_freqs).to(img_query.dtype) + img_sin = torch.imag(vid_freqs).to(img_query.dtype) + txt_cos = torch.real(txt_freqs).to(txt_query.dtype) + txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) img_query = self.rope(img_query, img_cos, img_sin) img_key = self.rope(img_key, img_cos, img_sin) From ffd11c181e5fa4e558ba2570f326fb41371c0421 Mon Sep 17 00:00:00 2001 From: Wenjia Li Date: Fri, 14 Aug 2026 17:49:26 +0800 Subject: [PATCH 38/41] [Perf][NPU] Fuse MiniMax H3 Qwen3-VL RoPE (#6061) Signed-off-by: Wenjia Li --- .../test_minimax_h3_qwen3vl_rope.py | 104 ++++++++++++++++++ vllm_omni/platforms/npu/models/minimax_h3.py | 43 ++++++++ vllm_omni/platforms/npu/platform.py | 2 + 3 files changed, 149 insertions(+) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py create mode 100644 vllm_omni/platforms/npu/models/minimax_h3.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py new file mode 100644 index 000000000..9ac3243bb --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize("num_q_heads,num_kv_heads,seq_len", [(8, 2, 16), (64, 8, 1)]) +def test_qwen3_vl_npu_rope_matches_bnsd_reference( + num_q_heads: int, + num_kv_heads: int, + seq_len: int, +) -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu + + device = torch.device("npu") + dtype = torch.bfloat16 + q = torch.randn(1, num_q_heads, seq_len, 128, device=device, dtype=dtype) + k = torch.randn(1, num_kv_heads, seq_len, 128, device=device, dtype=dtype) + cos = torch.randn(1, seq_len, 128, device=device, dtype=dtype) + sin = torch.randn(1, seq_len, 128, device=device, dtype=dtype) + + q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin) + cos_bnsd = cos.unsqueeze(1) + sin_bnsd = sin.unsqueeze(1) + + def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + torch.testing.assert_close( + q_out, + q * cos_bnsd + rotate_half(q) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + torch.testing.assert_close( + k_out, + k * cos_bnsd + rotate_half(k) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_rope_preserves_batched_mrope_positions() -> None: + """Keep per-sample Qwen3-VL M-RoPE frequencies on the BNSD NPU path.""" + from types import SimpleNamespace + + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLTextRotaryEmbedding, + ) + from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu + + device = torch.device("npu") + dtype = torch.bfloat16 + batch_size, seq_len, head_dim = 2, 16, 128 + config = SimpleNamespace( + hidden_size=head_dim, + num_attention_heads=1, + rope_parameters={ + "rope_theta": 1_000_000.0, + "mrope_section": [24, 20, 20], + }, + ) + rope = MiniMaxH3Qwen3VLTextRotaryEmbedding(config).to(device) + positions = torch.arange(seq_len, device=device) + position_ids = torch.stack( + ( + torch.stack((positions, positions + 7)), + torch.stack((positions * 2, positions * 2 + 3)), + torch.stack((positions * 3, positions * 3 + 5)), + ) + ) + q = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn(batch_size, 2, seq_len, head_dim, device=device, dtype=dtype) + cos, sin = rope(q, position_ids) + + q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin) + cos_bnsd = cos.unsqueeze(1) + sin_bnsd = sin.unsqueeze(1) + + def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + assert not torch.equal(cos[0], cos[1]) + torch.testing.assert_close( + q_out, + q * cos_bnsd + rotate_half(q) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + torch.testing.assert_close( + k_out, + k * cos_bnsd + rotate_half(k) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) diff --git a/vllm_omni/platforms/npu/models/minimax_h3.py b/vllm_omni/platforms/npu/models/minimax_h3.py new file mode 100644 index 000000000..a05abd654 --- /dev/null +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""NPU patches for the MiniMax H3 Qwen3-VL text encoder.""" + +from __future__ import annotations + +import torch +from vllm.logger import init_logger + +from vllm_omni.platforms.npu.layers.rotary_embedding import ( + npu_rotary_mul_with_bsnd_fallback, +) + +logger = init_logger(__name__) + +_PATCHED = False + + +def _apply_rotary_pos_emb_npu( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply Qwen3-VL text RoPE with fused BNSD rotary multiplication.""" + return ( + npu_rotary_mul_with_bsnd_fallback(q, cos, sin, unsqueeze_dim=1), + npu_rotary_mul_with_bsnd_fallback(k, cos, sin, unsqueeze_dim=1), + ) + + +def apply_minimax_h3_qwen3vl_patch() -> None: + """Route MiniMax H3 Qwen3-VL text RoPE to the Ascend fused operator.""" + global _PATCHED + if _PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder._apply_rotary_pos_emb = _apply_rotary_pos_emb_npu + _PATCHED = True + logger.debug("Applied NPU fused RoPE patch for MiniMax H3 Qwen3-VL text encoder") diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 525725fcf..16f87b21d 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -36,6 +36,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): from vllm_ascend.utils import adapt_patch from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches + from vllm_omni.platforms.npu.models.minimax_h3 import apply_minimax_h3_qwen3vl_patch from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import ( apply_qwen3_tts_code2wav_patch, ) @@ -44,6 +45,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): ) adapt_patch(is_global_patch=True) + apply_minimax_h3_qwen3vl_patch() apply_qwen3_tts_code2wav_patch() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() From 8ecd1f6d5cc91aab8a475a861213720b336e2f65 Mon Sep 17 00:00:00 2001 From: Wenjia Li Date: Fri, 14 Aug 2026 19:47:17 +0800 Subject: [PATCH 39/41] [Perf][NPU] Fuse MiniMax H3 Qwen3-VL SwiGLU (#6167) Signed-off-by: Wenjia Li --- .../test_minimax_h3_qwen3vl_swiglu.py | 46 +++++++++++++++++++ vllm_omni/platforms/npu/models/minimax_h3.py | 36 +++++++++++++-- vllm_omni/platforms/npu/platform.py | 8 ++++ 3 files changed, 86 insertions(+), 4 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py new file mode 100644 index 000000000..2664990b4 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize("shape", [(1, 16, 256), (2, 1, 256)]) +def test_qwen3_vl_npu_swiglu_matches_split_reference( + shape: tuple[int, int, int], +) -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed + + packed = torch.randn(*shape, device="npu", dtype=torch.bfloat16) + gate, up = packed.chunk(2, dim=-1) + + torch.testing.assert_close( + npu_swiglu_from_packed(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed + + x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16) + weight = torch.randn(256, 128, device="npu", dtype=torch.bfloat16) + packed = F.linear(x, weight) + gate = F.linear(x, weight[:128]) + up = F.linear(x, weight[128:]) + + torch.testing.assert_close( + npu_swiglu_from_packed(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) diff --git a/vllm_omni/platforms/npu/models/minimax_h3.py b/vllm_omni/platforms/npu/models/minimax_h3.py index a05abd654..78e0b3f14 100644 --- a/vllm_omni/platforms/npu/models/minimax_h3.py +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -5,7 +5,11 @@ from __future__ import annotations +from typing import Any + import torch +import torch.nn.functional as F +import torch_npu from vllm.logger import init_logger from vllm_omni.platforms.npu.layers.rotary_embedding import ( @@ -14,7 +18,8 @@ from vllm_omni.platforms.npu.layers.rotary_embedding import ( logger = init_logger(__name__) -_PATCHED = False +_ROPE_PATCHED = False +_SWIGLU_PATCHED = False def _apply_rotary_pos_emb_npu( @@ -32,12 +37,35 @@ def _apply_rotary_pos_emb_npu( def apply_minimax_h3_qwen3vl_patch() -> None: """Route MiniMax H3 Qwen3-VL text RoPE to the Ascend fused operator.""" - global _PATCHED - if _PATCHED: + global _ROPE_PATCHED + if _ROPE_PATCHED: return from vllm_omni.diffusion.models.minimax_h3 import encoder encoder._apply_rotary_pos_emb = _apply_rotary_pos_emb_npu - _PATCHED = True + _ROPE_PATCHED = True logger.debug("Applied NPU fused RoPE patch for MiniMax H3 Qwen3-VL text encoder") + + +def npu_swiglu_from_packed(gate_up: torch.Tensor) -> torch.Tensor: + """Apply fused SwiGLU to a packed gate/up projection tensor.""" + return torch_npu.npu_swiglu(gate_up, dim=-1) + + +def _forward_minimax_h3_qwen3vl_text_mlp_npu(self: Any, x: torch.Tensor) -> torch.Tensor: + """Run the Qwen3-VL MLP with one packed GEMM and fused SwiGLU.""" + gate_up = F.linear(x, self.gate_up_proj.weight) + return self.down_proj(npu_swiglu_from_packed(gate_up)) + + +def apply_minimax_h3_qwen3vl_swiglu_patch() -> None: + """Route MiniMax H3 Qwen3-VL text MLP to the Ascend fused SwiGLU path.""" + global _SWIGLU_PATCHED + if _SWIGLU_PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder.MiniMaxH3Qwen3VLTextMLP.forward = _forward_minimax_h3_qwen3vl_text_mlp_npu + _SWIGLU_PATCHED = True diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 16f87b21d..b17f447c4 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -82,6 +82,14 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): def init_diffusion_model_runner_runtime(cls, vllm_config: Any, od_config: Any, device: torch.device) -> None: from vllm_ascend.ascend_forward_context import set_mc2_mask, set_mc2_tokens_capacity + from vllm_omni.platforms.npu.models.minimax_h3 import ( + apply_minimax_h3_qwen3vl_swiglu_patch, + ) + + # The patch imports the MiniMax encoder, which depends on + # current_omni_platform. Run it only after platform construction has + # completed, but before the diffusion pipeline is loaded. + apply_minimax_h3_qwen3vl_swiglu_patch() set_mc2_tokens_capacity(vllm_config, od_config.max_num_seqs, 1) set_mc2_mask(vllm_config, device) From f7a5146093bc4d16b5607901a8a2509d785777db Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 14 Aug 2026 13:13:00 -0600 Subject: [PATCH 40/41] [Feat/Bugfix] Fix API Server model tag <-> model sync & flag normalization (#3805) Signed-off-by: Alex Brooks --- vllm_omni/entrypoints/openai/api_server.py | 29 ++++++++-------------- 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 16114bcb6..1ba26096f 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -43,6 +43,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, ) +from vllm.entrypoints.openai.cli_args import make_arg_parser # yapf conflicts with isort for this block # yapf: disable @@ -3826,26 +3827,16 @@ async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request): if __name__ == "__main__": - import asyncio - - from vllm.entrypoints.openai.cli_args import make_arg_parser - parser = TrackingArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server") parser = make_arg_parser(parser) - registered_flags = set() - for action in parser._actions: - registered_flags.update(action.option_strings) - - # FIXME - this is broken on `main`. `make_arg_parser` does not handle omni engine args properly. - if "--omni" not in registered_flags: - parser.add_argument("--omni", action="store_true", default=False, help="Enable vLLM-Omni mode.") - if "--enable-sleep-mode" not in registered_flags: - parser.add_argument( - "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode." - ) + # Ensure that passing --omni won't crash the server. + # NOTE: the value here does not matter since we are always running the Omni server + # when __main__ is called, i.e., --omni is only used when called through the entrypoints. + parser.add_argument("--omni", action="store_true", default=False) args = parser.parse_args() - if not hasattr(args, "model_tag"): - setattr(args, "model_tag", args.model) - if hasattr(args, "model_tag") and args.model_tag is None: - args.model_tag = args.model + # sync args.model to model_tag, because if we pass the model positionally, + # args.model will be the default from vLLM's ModelConfig (currently + # Qwen/Qwen3-0.6B) and crash cryptically. + if args.model_tag is not None: + args.model = args.model_tag asyncio.run(omni_run_server(args)) From 4b6d6d5cf79506d2a6ea697da11f1b2bdceccda8 Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Fri, 14 Aug 2026 13:23:41 -0700 Subject: [PATCH 41/41] [Bugfix][XPU][Tests]Scope XPU pytest to explicit paths to fix collection crash (#6175) Signed-off-by: Joshna Medisetty Co-authored-by: Cursor Co-authored-by: Chendi.Xue --- .buildkite/intel/scripts/run-xpu-test.sh | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.buildkite/intel/scripts/run-xpu-test.sh b/.buildkite/intel/scripts/run-xpu-test.sh index 742e39b78..9fb7267c7 100755 --- a/.buildkite/intel/scripts/run-xpu-test.sh +++ b/.buildkite/intel/scripts/run-xpu-test.sh @@ -87,8 +87,9 @@ time timeout -k 30 30m docker run \ echo $ZE_AFFINITY_MASK pip install tblib==3.1.0 cd /workspace/vllm-omni - pytest -v -s -m "core_model and xpu and B60" + XPU_TEST_PATHS="tests/diffusion tests/dfx tests/e2e" + pytest -v -s $XPU_TEST_PATHS -m "core_model and xpu and B60" pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py - pytest -v -s -m "advanced_model and xpu and B60" - pytest -v -s -m "omni and xpu and B60" + pytest -v -s $XPU_TEST_PATHS -m "advanced_model and xpu and B60" + pytest -v -s $XPU_TEST_PATHS -m "omni and xpu and B60" '