From d93e1a8e09d2a0d386c4a58424470d685ee7743c Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Thu, 1 Oct 2026 22:22:54 -0700 Subject: [PATCH] [Model] Optimize Qwen3-Omni MRv2 handoff, code prediction and audio decoding (#8223) Signed-off-by: Sy03 <1370724210@qq.com> --- .../common/ci_source_file_dependencies.yml | 8 + docs/design/index.md | 1 + docs/design/qwen3_omni_mrv2_performance.md | 22 + tests/config/test_environment_variables.py | 2 +- .../models/common/test_fused_sampling.py | 90 ++++ .../test_short_kv_attention.py | 2 +- .../qwen3_omni/test_first_frame_cuda.py | 32 ++ .../test_qwen3_omni_code2wav_cudnn.py | 87 ++++ .../test_qwen3_omni_code2wav_snake.py | 88 ++++ .../qwen3_omni/test_qwen3_omni_first_audio.py | 208 +++++++++ .../models/qwen3_omni/test_qwen3_omni_mrv2.py | 116 +++++ .../qwen3_omni/test_qwen3_omni_pp_capture.py | 1 + .../qwen3_tts/test_code_predictor_dtype.py | 21 + .../models/test_talker_first_audio.py | 70 +++ .../test_qwen3_omni_chunk_ramp.py | 116 +++++ .../test_qwen3_omni_first_token_handoff.py | 146 ++++++ .../test_qwen3_omni_streaming_helpers.py | 77 ++++ tests/worker_v2/test_omni_model_state.py | 6 + .../config/environment_variable_inventory.py | 1 + vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml | 50 +++ .../deploy/qwen3_omni_moe_mrv2_h200.yaml | 27 ++ .../models/common/fused_sampling.py | 97 ++++ .../models/common/qwen3_code_predictor.py | 251 ++++++++++- .../short_kv_attention.py | 0 .../models/common/talker_first_audio.py | 29 ++ .../models/qwen3_omni/first_frame_decoder.py | 99 +++++ .../models/qwen3_omni/pipeline.py | 5 + .../models/qwen3_omni/qwen3_omni.py | 418 +++++++++++++++--- .../models/qwen3_omni/qwen3_omni_code2wav.py | 218 ++++++++- .../models/qwen3_tts/cached_code_predictor.py | 3 +- .../models/qwen3_tts/first_audio.py | 18 +- .../qwen3_tts_code_predictor_vllm.py | 3 + .../chunk_size_utils.py | 17 +- .../stage_input_processors/qwen3_omni.py | 179 ++++++-- .../model_states/omni_model_state.py | 16 +- 35 files changed, 2363 insertions(+), 161 deletions(-) create mode 100644 docs/design/qwen3_omni_mrv2_performance.md create mode 100644 tests/model_executor/models/common/test_fused_sampling.py rename tests/model_executor/models/{qwen3_tts => common}/test_short_kv_attention.py (95%) create mode 100644 tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py create mode 100644 tests/model_executor/models/test_talker_first_audio.py create mode 100644 tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py create mode 100644 tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py create mode 100644 vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml create mode 100644 vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml create mode 100644 vllm_omni/model_executor/models/common/fused_sampling.py rename vllm_omni/model_executor/models/{qwen3_tts => common}/short_kv_attention.py (100%) create mode 100644 vllm_omni/model_executor/models/common/talker_first_audio.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 9594e62e5..b34e9d271 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -16,6 +16,9 @@ x-qwen3-tts: &qwen3_tts - vllm_omni/worker/gpu_model_runner.py - vllm_omni/model_executor/models/qwen3_tts/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/fused_sampling.py + - vllm_omni/model_executor/models/common/short_kv_attention.py + - vllm_omni/model_executor/models/common/talker_first_audio.py - vllm_omni/model_executor/models/common/snake_activation.py - vllm_omni/model_executor/models/common/alias_free_activation.py - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -85,11 +88,16 @@ x-qwen3-omni: &qwen3_omni - vllm_omni/entrypoints/openai/serving_chat.py - vllm_omni/model_executor/models/qwen3_omni/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/fused_sampling.py + - vllm_omni/model_executor/models/common/short_kv_attention.py + - vllm_omni/model_executor/models/common/talker_first_audio.py - vllm_omni/model_executor/models/common/snake_activation.py - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - vllm_omni/deploy/qwen3_omni_moe.yaml + - vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml + - vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml x-wan22: &wan22 diff --git a/docs/design/index.md b/docs/design/index.md index 748e1d9ca..af8dbf438 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -116,4 +116,5 @@ The pre-#5137 pages are preserved in the [legacy module archive](module/archive/README.md) for historical reference and are not active design contracts. +- [Qwen3-Omni MRv2 performance](qwen3_omni_mrv2_performance.md) - [MiniCPM-o 4.5 turn-mode MRv2 performance](minicpm_o45_mrv2_performance.md) diff --git a/docs/design/qwen3_omni_mrv2_performance.md b/docs/design/qwen3_omni_mrv2_performance.md new file mode 100644 index 000000000..c81b742dd --- /dev/null +++ b/docs/design/qwen3_omni_mrv2_performance.md @@ -0,0 +1,22 @@ +# Qwen3-Omni MRv2 performance + +This model adaptation builds on #8184's existing stage runtime and output +contracts. It adds Talker CPU metadata parsing, sampled-token embedding +handoff, an optional first-frame decoder, length-grouped Code2Wav graph +buckets, and opt-in incremental short-KV code prediction. The H200 profile +also enables conservative fused sampling and bounded cuDNN algorithm search. +The shared predictor keeps the upstream re-prefill helper and its default +sampling behavior; short KV is explicitly selected by the deployment profile. + +Use `vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml` for the experimental H200 +profile. Preserve a control with `code_predictor_fused_sampling=false` and +`codec_cudnn_benchmark=false` when C64 P99 is the primary requirement. +cuDNN search restores process settings on exit and has startup cost. Kernel +arithmetic is not bitwise-equivalent; the first-frame decoder uses additional +model memory. Full-duplex performance is not established by turn-mode tests. + +First-frame delivery requires `talker_first_audio: true` in connector extra, +a one-frame initial codec chunk, CUDA, a local single-rank Talker and prefix caching disabled. +Other deployments retain the regular Code2Wav path. +`codec_fused_snake` explicitly opts CUDA profiles into fused decoder blocks; +the base V1 profile retains its original blocks and indexed text cache. diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 6f004b6b6..e1f6aa651 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -169,7 +169,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, EnvironmentVariableCategory.MODEL_SPECIFIC: 75, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, - EnvironmentVariableCategory.INTERNAL: 6, + EnvironmentVariableCategory.INTERNAL: 7, } disposition_counts = Counter( diff --git a/tests/model_executor/models/common/test_fused_sampling.py b/tests/model_executor/models/common/test_fused_sampling.py new file mode 100644 index 000000000..45ac3f81d --- /dev/null +++ b/tests/model_executor/models/common/test_fused_sampling.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"), +] + + +def _reference(logits, uniforms, top_k, top_p): + if top_k > 0: + threshold = logits.topk(top_k, dim=-1).values[:, -1:] + logits = logits.masked_fill(logits < threshold, -torch.inf) + if top_p < 1: + values, indices = logits.sort(dim=-1, descending=True) + probs = values.softmax(dim=-1, dtype=torch.float32) + values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf + logits = values.scatter(1, indices, values) + return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) + + +@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)]) +def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + torch.manual_seed(173) + logits = torch.randn(batch, vocab, device="cuda", dtype=dtype) + # The MTP caller supplies a noncontiguous row slice of [B, 15, vocab]. + uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7) + expected = _reference(logits, uniforms, top_k, top_p) + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p) + logits.add_(torch.randn_like(logits) * 0.1) + uniforms.uniform_(1e-7, 1 - 1e-7) + graph.replay() + torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0) + + +@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)]) +def test_threshold_ties_and_degenerate_rows(top_k, top_p): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + logits = torch.tensor( + [[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]], + device="cuda", + dtype=torch.float32, + ) + uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1) + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p) + torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_native_nucleus_boundaries(dtype): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + torch.manual_seed(1923) + logits = torch.randn(1, 2048, device="cuda", dtype=dtype) + values, indices = logits.sort(descending=True) + probabilities = values.softmax(-1, dtype=torch.float32) + before = probabilities.cumsum(-1) - probabilities + for rank in [1, 19, 49, 1000, 2000]: + boundary = before[0, rank] + uniforms = torch.full_like(logits, 0.001, dtype=torch.float32) + uniforms[0, indices[0, rank]] = 1 - 1e-7 + for direction in [-torch.inf, torch.inf]: + top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item() + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p) + torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0) + + +@pytest.mark.parametrize("top_p", [0.8, 1.0]) +def test_strided_logits_and_non_power_of_two_vocabulary(top_p): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + for seed in range(16): + torch.manual_seed(seed) + logits = torch.randn(128, 2050, device="cuda")[:, ::2] + uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7) + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p) + torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0) diff --git a/tests/model_executor/models/qwen3_tts/test_short_kv_attention.py b/tests/model_executor/models/common/test_short_kv_attention.py similarity index 95% rename from tests/model_executor/models/qwen3_tts/test_short_kv_attention.py rename to tests/model_executor/models/common/test_short_kv_attention.py index 15dcecb14..000872617 100644 --- a/tests/model_executor/models/qwen3_tts/test_short_kv_attention.py +++ b/tests/model_executor/models/common/test_short_kv_attention.py @@ -16,7 +16,7 @@ pytestmark = [ @pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)]) @pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]]) def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions): - from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention + from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention with torch.inference_mode(): torch.manual_seed(42) diff --git a/tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py b/tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py new file mode 100644 index 000000000..0d75d68ee --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("rows", [16, 19]) +@torch.inference_mode() +def test_first_frame_graph_outputs_own_storage(dtype, rows): + class Decoder(torch.nn.Linear): + config = SimpleNamespace(num_quantizers=1) + + def forward(self, codes): + return codes.to(dtype) * self.weight + + model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype) + model.weight.fill_(1) + decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000) + decoder.capture() + codes = torch.arange(rows, device="cuda").reshape(-1, 1) + first = decoder.decode(codes) + second = decoder.decode(codes + 32) + assert first.dtype == second.dtype == torch.float32 + torch.testing.assert_close(first, codes.float(), rtol=0, atol=0) + torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py new file mode 100644 index 000000000..d345473c9 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"), +] + + +def _codec(enabled): + config = Qwen3OmniMoeCode2WavConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=8, + decoder_dim=32, + upsample_rates=[2, 2], + upsampling_ratios=[2], + codebook_size=16, + num_quantizers=2, + ) + config._attn_implementation = "sdpa" + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}} + ) + ) + model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval() + model.precompute_snake_caches() + return model + + +def _flags(): + c = torch.backends.cudnn + return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32 + + +@torch.inference_mode() +def test_autotuned_codec_graph_replay_restores_flags(): + torch.manual_seed(21) + model = _codec(True) + codes = torch.randint(0, 16, (2, 2, 8), device="cuda") + original = _flags() + for _ in range(3): + model(codes) + assert _flags() == original + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = model(codes) + assert _flags() == original + codes.random_(0, 16) + graph.replay() + expected = model(codes) + torch.testing.assert_close(captured, expected) + assert _flags() == original + + +@pytest.mark.parametrize("enabled", [False, True]) +@torch.inference_mode() +def test_codec_failure_restores_cudnn_settings(mocker, enabled): + model = _codec(enabled) + codes = torch.randint(0, 16, (1, 2, 8), device="cuda") + original = _flags() + + def fail(hidden): + if enabled: + assert torch.backends.cudnn.benchmark + assert torch.backends.cudnn.benchmark_limit == 10 + else: + assert _flags() == original + raise RuntimeError("decode failed") + + mocker.patch.object(model, "_decode_waveform", side_effect=fail) + with pytest.raises(RuntimeError, match="decode failed"): + model(codes) + assert _flags() == original diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py new file mode 100644 index 000000000..e18aaa109 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import copy +from types import SimpleNamespace + +import pytest +import torch +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeCode2WavDecoderBlock, + Qwen3OmniMoeSnakeBeta, +) + +from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _blocks(): + config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2]) + torch.manual_seed(0) + blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval() + with torch.no_grad(): + for module in blocks.modules(): + if isinstance(module, Qwen3OmniMoeSnakeBeta): + module.alpha.uniform_(-0.5, 0.5) + module.beta.uniform_(-0.5, 0.5) + return blocks + + +def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names(): + blocks = _blocks() + hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules()) + keys = set(blocks.state_dict()) + + assert use_fused_snake(blocks) == hf_count > 0 + assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules()) + assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count + assert set(blocks.state_dict()) == keys # checkpoint weights load unchanged + + +def test_fused_snake_blocks_match_hf_blocks(): + reference = _blocks() + fused = copy.deepcopy(reference) + use_fused_snake(fused) + fused.load_state_dict(reference.state_dict()) + x = torch.randn(2, 32, 20) + with torch.no_grad(): + expected = x + actual = x + for ref_block, fused_block in zip(reference, fused): + expected = ref_block(expected) + actual = fused_block(actual) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize( + ("option", "device", "cuda", "fused"), + [ + (False, "cuda", True, False), + (True, "cuda", True, True), + (True, "cpu", True, False), + (True, "cuda", False, False), + ], +) +def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda) + config = Qwen3OmniMoeCode2WavConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + decoder_dim=32, + ) + model = Qwen3OmniMoeCode2Wav( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}} + ), + device_config=SimpleNamespace(device=device), + ) + ) + assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py new file mode 100644 index 000000000..5ca8b7856 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig + +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import ( + Qwen3OmniMoeCode2Wav, + plan_decode_groups, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_Q = 4 +_CODEBOOK = 16 + + +def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav: + torch.manual_seed(0) + config = Qwen3OmniMoeCode2WavConfig( + codebook_size=_CODEBOOK, + num_quantizers=quantizers, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + decoder_dim=32, + sliding_window=8, + ) + model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config))) + return model.eval() + + +def test_first_frame_pcm_equals_code2wav_streaming_chunk0(): + code2wav = _code2wav() + decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000) + codes = torch.randint(0, _CODEBOOK, (3, _Q)) + with torch.inference_mode(): + for row in range(codes.shape[0]): + pcm = decoder.decode(codes[row : row + 1]) + chunk0 = code2wav.chunked_decode_streaming( + codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q] + )[0] + assert pcm.dtype == torch.float32 + assert torch.equal(pcm[0], chunk0.reshape(-1).float()) + assert 0 < pcm.shape[-1] < int(code2wav.total_upsample) + batched = decoder.decode(codes) + assert batched.shape == (3, pcm.shape[-1]) + + +def test_first_frame_is_a_prefix_of_a_longer_first_chunk(): + code2wav = _code2wav() + decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000) + codes = torch.randint(0, _CODEBOOK, (_Q, 4)) + with torch.inference_mode(): + first = decoder.decode(codes[:, :1].T)[0] + chunk0 = code2wav.chunked_decode_streaming( + codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q] + )[0] + torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first) + + +def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch): + import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module + + real = wrapper_module.CUDAGraphDecoderWrapper + wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes) + monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper) + extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]} + code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q)) + code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph( + code2wav, device=torch.device("cuda"), **kw + ) + model = object.__new__(Qwen3OmniMoeForConditionalGeneration) + torch.nn.Module.__init__(model) + model.code2wav = code2wav + model.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra}) + ) + for batches in ([], [2, 4, 8]): + extra["codec_graph_batch_sizes"] = batches + Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model) + args = wrapper.call_args.kwargs + sizes = args["capture_sizes"] + assert {1, 3, 7, 15, 31, 50} <= set(sizes) + assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes) + assert set(args["extra_capture_shapes"]) == { + (b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325 + } + assert code2wav._streaming_batch_sizes[50] == [1, 2, 4] + extra.pop("codec_chunk_ramp") + extra.pop("codec_graph_batch_sizes") + Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model) + assert wrapper.call_args.kwargs["capture_sizes"] is None + + +_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64] + + +def _bucket(length): + return next((size for size in _SIZES if size >= length), None) + + +def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph(): + assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)] + calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16]) + assert [len(rows) for rows, _length in calls] == [16, 4] + calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4]) + assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)] + assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == [] + + +@pytest.mark.parametrize("frames", [65, 129]) +def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames): + from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper + + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.enabled = wrapper._warmed_up = True + wrapper._bucket_sizes = [8, 32, 64] + wrapper._compiled_shapes = {(1, 64), (8, 64)} + wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock()) + wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,)) + assert groups == [([0], frames), ([1], frames)] + codes = torch.arange(2 * frames).reshape(2, 1, frames) + for rows, width in groups: + actual = wrapper.decode(codes[rows, :, :width]) + torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0) + assert wrapper.decoder.call_count == 2 + for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()): + graph.replay.assert_not_called() + + +class _PaddingGraphs: + def __init__(self, model, batch_sizes): + self.model = model + self.batch_sizes = sorted(batch_sizes) + self.calls = [] + + def _get_padded_size(self, length): + return _bucket(length) + + def decode(self, codes): + rows, _q, frames = codes.shape + size = _bucket(frames) + batch = next(b for b in self.batch_sizes if b >= rows) + self.calls.append((rows, frames)) + padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype) + padded[:rows, :, :frames] = codes + out = self.model(padded) + return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone() + + +def test_grouped_streaming_decode_matches_one_padded_batch(): + code2wav = _code2wav() + lengths = [1, 3, 50, 7, 31] + lefts = [0, 1, 25, 3, 15] + codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long) + for row, length in enumerate(lengths): + codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length)) + counts = [length * _Q for length in lengths] + with torch.inference_mode(): + reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts) + graphs = _PaddingGraphs(code2wav, [1, 2, 4]) + code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True + code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES} + grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts) + + for want, got in zip(reference, grouped, strict=True): + assert want.shape == got.shape + torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)]) +def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers): + from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + + codec = _code2wav(quantizers) + model = object.__new__(Qwen3OmniMoeForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config + model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming( + codes, left_context_size=left, seq_token_counts=counts + ) + codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames)) + with torch.inference_mode(): + reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2) + first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0] + out = model.forward( + input_ids=codes.flatten(), + positions=None, + seq_token_counts=[quantizers * frames] * 2, + runtime_additional_information=[ + {"meta": {"left_context_size": 0, "first_audio": True}}, + {"meta": {"left_context_size": 0, "first_audio": False}}, + ], + ).multimodal_outputs + torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten()) + torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1)) + assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False] diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py new file mode 100644 index 000000000..6f5d8f6c7 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_H = 4 +_Q = 3 + + +class _Talker(nn.Module): + num_code_groups = _Q + + def __init__(self): + super().__init__() + self.text_projection = nn.Identity() + self.code_predictor = SimpleNamespace( + _top_k=50, + _top_p=0.8, + set_sampling_params=lambda **kwargs: setattr(self, "sampling", kwargs), + ) + + +def _talker(*, async_chunk: bool = True) -> Qwen3OmniMoeForConditionalGeneration: + model = object.__new__(Qwen3OmniMoeForConditionalGeneration) + nn.Module.__init__(model) + model.model_stage = "talker" + model.talker = _Talker() + model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=async_chunk, use_v2_model_runner=True)) + model.tts_eos_embed = torch.full((_H,), -1.0) + model.tts_pad_embed = torch.full((_H,), -2.0) + model._codec_codebook_size = 2048 + return model + + +def _rows(*values: float) -> torch.Tensor: + return torch.tensor([[value] * _H for value in values]) + + +def _step(model, payload: dict) -> torch.Tensor: + update: dict = {} + text_step = model._thinker_decode_to_talker_decode(payload, torch.device("cpu"), update) + for key, value in update.items(): + payload.setdefault(key, {}).update(value) + return text_step.reshape(-1) + + +def test_mtp_frame_valid_rejects_codec_eos_and_special_ids(): + layer0 = torch.tensor([0, 2047, 2048, 2150, -1]) + assert _talker().mtp_frame_valid(layer0).tolist() == [True, True, False, False, False] + + +def test_decode_rows_are_consumed_one_per_step_in_arrival_order(): + model = _talker() + # Two rows were already pending (cached at prefill); two more arrive at once. + payload = {"embed": {"cached_decode": _rows(1, 2), "decode": _rows(3, 4)}, "meta": {}} + steps = [_step(model, payload)[0].item()] + for arrival in (None, _rows(5), None, None, None, None): + payload["embed"]["decode"] = arrival + steps.append(_step(model, payload)[0].item()) + assert steps == [1, 2, 3, 4, 5, -1, -2] + assert payload["embed"]["cached_decode"] is None + + +@pytest.mark.parametrize(("v2", "resumable"), [(True, True), (False, False), (False, True)]) +def test_v1_and_resumable_requests_keep_the_indexed_cache_path(v2, resumable): + model = _talker() + model.vllm_config.model_config.use_v2_model_runner = v2 + payload = {"embed": {"cached_decode": _rows(1, 2, 3), "decode": None}, "meta": {"resumable": resumable}} + payload["meta"]["num_processed_tokens"] = 2 + assert _step(model, payload)[0].item() == 3 + + +def test_make_omni_output_aligns_codes_to_request_spans(): + model = _talker() + hidden = torch.zeros((8, _H)) # 5 real tokens + graph padding + decode_codes = torch.tensor([[5, 6, 7]]) + buffers = [ + {"req_id": "_warmup_0_"}, # no preprocess, no codes + {"req_id": "a", "codes": {"audio": torch.zeros((3, _Q), dtype=torch.long)}}, + {"req_id": "b", "codes": {"audio": decode_codes}}, + ] + out = model.make_omni_output( + hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 1), (1, 4), (4, 5)] + ) + codes = out.multimodal_outputs["codes"]["audio"] + assert codes.shape == (5, _Q) + assert codes[4].tolist() == [5, 6, 7] + assert out.text_hidden_states.shape[0] == 5 + + +def test_make_omni_output_without_spans_keeps_v1_concat(): + model = _talker() + buffers = [{"codes": {"audio": torch.tensor([[1, 2, 3]])}}, {"codes": {"audio": torch.tensor([[4, 5, 6]])}}] + out = model.make_omni_output(torch.zeros((4, _H)), model_intermediate_buffer=buffers) + assert out.multimodal_outputs["codes"]["audio"].tolist() == [[1, 2, 3], [4, 5, 6]] + assert out.text_hidden_states.shape[0] == 2 + + +def test_eager_output_is_token_major_with_invalid_rows(): + model = _talker() + model.eager_frames_active = True + hidden = torch.zeros((6, _H)) + buffers = [{"req_id": "a"}, {"req_id": "b"}] + out = model.make_omni_output(hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 3), (3, 4)]) + mm = out.multimodal_outputs + assert mm["codes"]["audio"].shape == (4, _Q) + assert mm["meta"]["codec_frame_valid"].tolist() == [0, 0, 0, 0] + assert out.text_hidden_states is hidden diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py index de8d02e3e..fffb796ad 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py @@ -104,6 +104,7 @@ def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, e if staged: model_config.model_stage = "thinker" vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None) + vllm_config.speculative_config = None thinker = nn.Module() thinker.weight = nn.Parameter(torch.ones(1)) thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start)) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index fe19b44a9..52712a29c 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -1432,3 +1432,24 @@ def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_cl expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms) torch.testing.assert_close(actual, expected) + + +def test_tts_cached_predictor_keeps_its_own_step_dispatch(mocker: MockerFixture, loaded_target_classes) -> None: + """The TTS cache must not enter Omni's uninitialized incremental loop.""" + _, _, wrapper, _, _ = loaded_target_classes + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object(common_mod.current_omni_platform, "is_cuda", return_value=True) + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False) + mocker.patch.object(wrapper, "_stage_connector_extra_config", return_value={"code_predictor_kv_cache": True}) + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + predictor = wrapper(vllm_config=_make_vllm_config(mocker), config=cp_config, talker_config=talker_config) + assert predictor._kv_requested + predictor._model_dtype = torch.float32 + predictor._codec_embeds_list = list(predictor.model.codec_embedding) + predictor._bucket_sizes = [1] + mocker.patch.object(predictor, "_setup_compile") + step = mocker.patch.object(predictor, "_predict_step_logits", return_value=torch.ones(1, cp_config.vocab_size)) + codes = predictor(torch.zeros(1, dtype=torch.long), torch.zeros(1, 32), torch.zeros(1, 32), do_sample=False) + assert step.call_count == cp_config.num_code_groups - 1 + assert codes.shape == (1, cp_config.num_code_groups) + assert predictor._kv_buf is None diff --git a/tests/model_executor/models/test_talker_first_audio.py b/tests/model_executor/models/test_talker_first_audio.py new file mode 100644 index 000000000..8d054290d --- /dev/null +++ b/tests/model_executor/models/test_talker_first_audio.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("env", "extra", "async_chunk", "v2", "expected"), + [ + (None, {}, True, True, False), + ("1", {}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True), + ("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False), + (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True), + (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False), + ], +) +def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected): + if env is None: + monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False) + else: + monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env) + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + config = SimpleNamespace( + model_config=SimpleNamespace( + async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra} + ), + device_config=SimpleNamespace(device="cuda"), + parallel_config=SimpleNamespace( + tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + ) + assert talker_first_audio_enabled(config) is expected + for section, field, value in [ + ("device_config", "device", "cpu"), + ("parallel_config", "tensor_parallel_size", 2), + ("parallel_config", "pipeline_parallel_size", 2), + ("parallel_config", "distributed_executor_backend", "mp"), + ("cache_config", "enable_prefix_caching", True), + ]: + with monkeypatch.context() as patch: + patch.setattr(getattr(config, section), field, value) + assert not talker_first_audio_enabled(config) + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + assert not talker_first_audio_enabled(config) + + +@pytest.mark.parametrize("option", [False, True]) +def test_tts_requires_connector_opt_in(monkeypatch, option): + from vllm_omni.model_executor.models.qwen3_tts import first_audio + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import ( + Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor, + ) + + monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option}) + monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True) + assert first_audio.talker_first_audio_enabled(None) is option diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py new file mode 100644 index 000000000..22fb3c9f3 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3 +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ramp_decode_windows + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_Q = 2 +_RAMP = [1, 2, 4, 8, 16, 25] + + +def _manager(**extra): + config = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "initial_codec_chunk_frames": 4, **extra} + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int), + ramp_chunk_count=defaultdict(int), + connector=SimpleNamespace(config={"extra": config}), + ) + + +def _request(request_id="req"): + return SimpleNamespace(external_req_id=request_id, sampling_params=SimpleNamespace(stop_token_ids=[])) + + +def _frame(index: int) -> dict: + return { + "codes": {"audio": torch.tensor([[index, 1000 + index]], dtype=torch.long)}, + "meta": {"codec_frame_valid": torch.ones(1, dtype=torch.int8)}, + } + + +def _stream(manager, num_frames, finish_after_last=True): + request = _request() + chunks = [] + + def emit(payload): + if payload is None: + return + codes = payload.codes.audio.reshape(_Q, -1) + chunks.append((codes[0].tolist(), int(payload.meta.left_context_size), bool(payload.meta.finished))) + # The transport counts every enqueued chunk. + manager.put_req_chunk[request.external_req_id] += 1 + manager.ramp_chunk_count[request.external_req_id] += 1 + + for index in range(num_frames): + last = finish_after_last and index == num_frames - 1 + emit(q3.talker2code2wav_async_chunk(manager, _frame(index), request, is_finished=last)) + return chunks + + +def test_ramp_decode_windows_include_left_context(): + assert ramp_decode_windows(_RAMP, 25) == [1, 3, 7, 15, 31, 50] + assert ramp_decode_windows([1, 2], 0) == [1, 2] + + +def test_ramp_chunks_grow_and_resend_left_context(): + chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 56 + 25, finish_after_last=False) + + new_frames = [len(window) - left for window, left, _finished in chunks] + assert new_frames == [1, 2, 4, 8, 16, 25, 25] + assert [left for _window, left, _finished in chunks] == [0, 1, 3, 7, 15, 25, 25] + assert [len(window) for window, _left, _finished in chunks] == [1, 3, 7, 15, 31, 50, 50] + emitted = 0 + for window, left, _finished in chunks: + # Each window ends at the newest frame and re-sends `left` earlier ones. + assert window == list(range(emitted - left, emitted + len(window) - left)) + emitted += len(window) - left + # The first chunk is the stream's first frame alone (no context). + assert chunks[0][0] == [0] + + +def test_ramp_flushes_partial_tail_with_left_context(): + chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 10) + + assert [len(window) - left for window, left, _finished in chunks] == [1, 2, 4, 3] + window, left, finished = chunks[-1] + assert finished and left == 7 and window == list(range(0, 10)) + + +def test_ramp_finish_on_chunk_boundary_leaves_the_marker_to_the_transport(): + manager = _manager(codec_chunk_ramp=_RAMP) + request = _request() + for index in range(3): + payload = q3.talker2code2wav_async_chunk(manager, _frame(index), request) + if payload is not None: + manager.ramp_chunk_count[request.external_req_id] += 1 + # Frames 0 and 1-2 went out as chunks 0 and 1; the finish step has no frame. + finish = {"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}, "meta": {"codec_frame_valid": torch.zeros(1)}} + assert q3.talker2code2wav_async_chunk(manager, finish, request, is_finished=True) is None + + +def test_without_ramp_the_initial_chunk_setting_is_unchanged(): + chunks = _stream(_manager(), 29 + 3) + + assert [(len(window), left) for window, left, _finished in chunks] == [(4, 0), (29, 4), (28, 25)] + + +def test_first_audio_marker_only_reaches_initial_codec_chunk(): + manager = _manager(codec_chunk_ramp=[1, 1]) + for index in range(2): + frame = _frame(index) + frame["meta"]["first_audio"] = torch.tensor([True]) + payload = q3.talker2code2wav_async_chunk(manager, frame, _request()) + assert bool(payload.meta.first_audio) is (index == 0) + manager.put_req_chunk["req"] += 1 + manager.ramp_chunk_count["req"] += 1 diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py new file mode 100644 index 000000000..7e32f2c15 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +ACCEPT = 24 +HIDDEN = 8 +PROMPT = [151644, 872, 5, 6, 151645, 151644, 77091, 198] +EOS = 151645 + + +class _Manager: + def __init__(self) -> None: + self.put_req_chunk: defaultdict[str, int] = defaultdict(int) + self.request_payload: dict[str, dict] = {} + self.config = SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=ACCEPT)) + ) + + def process(self, output, request, is_finished=False): + payload = q3.thinker2talker_async_chunk(self, output, request, is_finished=is_finished) + if payload is not None: + self.put_req_chunk[request.external_req_id] += 1 + return payload + + +def _request(output_token_ids, *, history=True, max_tokens=64): + # The native data plane keeps token history only while chunk 0 is pending. + return SimpleNamespace( + external_req_id="r", + prompt_token_ids=list(PROMPT) if history else [], + output_token_ids=list(output_token_ids) if history else [], + all_token_ids=list(PROMPT) + list(output_token_ids) if history else [], + output_token_count=len(output_token_ids), + last_output_token_id=output_token_ids[-1] if output_token_ids else None, + resumable=False, + additional_information=None, + sampling_params=SimpleNamespace(all_stop_token_ids={EOS}, max_tokens=max_tokens), + ) + + +def _embed(token_id: int) -> torch.Tensor: + return torch.full((1, HIDDEN), float(token_id), dtype=torch.bfloat16) + + +def _output(token_ids, sampled=None): + layers = { + 0: torch.cat([_embed(t) for t in token_ids]), + ACCEPT: torch.randn(len(token_ids), HIDDEN).to(torch.bfloat16), + } + embed = {"tts_bos": _embed(1), "tts_eos": _embed(2), "tts_pad": _embed(3)} + if sampled is not None: + embed["sampled"] = sampled + return {"hidden_states": {"layers": layers}, "embed": embed} + + +def _run_turn(tokens, *, publish, max_tokens=64, prefill_chunks=(len(PROMPT),)): + manager = _Manager() + rows, chunks, start = [], 0, 0 + for i, size in enumerate(prefill_chunks): + last = i == len(prefill_chunks) - 1 + sampled = _embed(tokens[0]) if last else torch.empty(0, dtype=torch.bfloat16) + request = _request(tokens[:1] if last else [], max_tokens=max_tokens) + payload = manager.process(_output(PROMPT[start : start + size], sampled if publish else None), request) + start += size + if payload is not None: + chunks += 1 + rows.extend(payload.embed.prefill[len(PROMPT) :]) + for step in range(1, len(tokens)): + history = manager.put_req_chunk["r"] == 0 + request = _request(tokens[: step + 1], history=history, max_tokens=max_tokens) + output = _output([tokens[step - 1]], _embed(tokens[step]) if publish else None) + payload = manager.process(output, request) + if payload is None: + continue + chunks += 1 + rows.extend(payload.embed.prefill[len(PROMPT) :] if payload.embed.prefill is not None else payload.embed.decode) + return [int(row[0]) for row in rows], chunks, manager + + +@pytest.mark.parametrize( + ("tokens", "max_tokens"), + [ + ([11, 12, 13, EOS], 64), # ends on a stop token + ([11, 12, 13], 3), # ends on the max_tokens-th token + ([11, EOS], 64), + ], +) +def test_sampled_stream_matches_captured_stream(tokens, max_tokens): + late, _chunks, _manager = _run_turn(tokens, publish=False, max_tokens=max_tokens) + early, _chunks, manager = _run_turn(tokens, publish=True, max_tokens=max_tokens) + # The final token is never processed, so neither path feeds it to the Talker. + assert early == late == tokens[:-1] + assert q3._SAMPLED_TEXT_STREAM in manager.request_payload["r"] + + +def test_sampled_stream_sends_each_token_one_step_earlier(): + manager = _Manager() + chunk0 = manager.process(_output(PROMPT, sampled=_embed(11)), _request([11])) + assert chunk0 is not None # at the prefill step, not after the next decode step + assert torch.equal(chunk0.embed.prefill, torch.cat([_embed(t) for t in [*PROMPT, 11]])) + assert chunk0.hidden_states.output.shape[0] == len(PROMPT) + assert chunk0.ids.all == [*PROMPT, 11] + # Decode step 1 processes 11 and samples 12: it carries 12. + chunk1 = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False)) + assert torch.equal(chunk1.embed.decode, _embed(12)) + + +def test_chunked_prefill_uses_sample_of_final_chunk_only(): + early, chunks, _manager = _run_turn([11, 12, EOS], publish=True, prefill_chunks=(5, 3)) + assert early == [11, 12] and chunks == 2 + + +def test_terminal_first_token_keeps_captured_stream(): + manager = _Manager() + assert manager.process(_output(PROMPT, sampled=_embed(EOS)), _request([EOS])) is None + assert q3._SAMPLED_TEXT_STREAM not in manager.request_payload["r"] + + +@pytest.mark.parametrize("prefill_rows", [0, 1, 3]) +def test_missing_sample_on_sampled_stream(prefill_rows): + manager = _Manager() + manager.process(_output(PROMPT, sampled=_embed(11)), _request([11])) + chunk = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False)) + assert torch.equal(chunk.embed.decode, _embed(12)) + if not prefill_rows: + with pytest.raises(RuntimeError, match="sample embedding missing"): + manager.process(_output([12]), _request([11, 12], history=False)) + else: + before = manager.put_req_chunk["r"] + for tokens in (PROMPT[:prefill_rows], [11]): + partial = _output(tokens, sampled=torch.empty(0)) + assert manager.process(partial, _request([11, 12], history=False)) is None + assert manager.put_req_chunk["r"] == before + chunk = manager.process(_output([12], sampled=_embed(13)), _request([11, 12, 13], history=False)) + assert torch.equal(chunk.embed.decode, _embed(13)) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index dfbb9d91d..f67fc7a58 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -1047,3 +1047,80 @@ def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None: payload = thinker2talker_full_payload(None, {"any": "thing"}, None) assert payload is None + + +def _codec_transfer_manager(request_id: str, *, initial_frames: int = 2, chunk_frames: int = 4): + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int, {request_id: 0}), + connector=SimpleNamespace( + config={ + "extra": { + "initial_codec_chunk_frames": initial_frames, + "codec_chunk_frames": chunk_frames, + "codec_left_context_frames": 4, + } + } + ), + ) + + +def _codec_request(request_id: str) -> SimpleNamespace: + return SimpleNamespace( + external_req_id=request_id, + sampling_params=SimpleNamespace(stop_token_ids=[2150], stop_token_id=None), + ) + + +def test_talker2code2wav_async_chunk_reads_last_valid_row_of_multi_row_span() -> None: + # MRv2 eager frames: a prefill span carries its first frame in its last row. + request_id = "eager_prefill" + manager = _codec_transfer_manager(request_id) + codes = torch.zeros((5, 3), dtype=torch.long) + codes[-1] = torch.tensor([11, 12, 13]) + valid = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8) + + payload = q3.talker2code2wav_async_chunk( + manager, {"codes": {"audio": codes}, "meta": {"codec_frame_valid": valid}}, _codec_request(request_id) + ) + + assert payload is None # below the initial chunk size + assert [frame.tolist() for frame in manager.code_prompt_token_ids[request_id]] == [[[11, 12, 13]]] + + +def test_talker2code2wav_async_chunk_skips_invalid_frame_and_flushes_on_finish() -> None: + request_id = "eager_eos" + manager = _codec_transfer_manager(request_id) + manager.code_prompt_token_ids[request_id] = [torch.tensor([[7, 8, 9]])] + # The step that sampled codec EOS carries an invalid (placeholder) frame. + output = { + "codes": {"audio": torch.tensor([[2150, 0, 0]])}, + "meta": {"codec_frame_valid": torch.tensor([0], dtype=torch.int8)}, + } + + payload = q3.talker2code2wav_async_chunk(manager, output, _codec_request(request_id), is_finished=True) + + assert len(manager.code_prompt_token_ids[request_id]) == 1 + assert payload is not None + assert payload.codes.audio.tolist() == [7, 8, 9] + assert payload.meta.finished.item() is True + + +def test_talker2code2wav_async_chunk_without_validity_keeps_one_row_contract() -> None: + # V1 / MRv2 deferred frames: one row per step, zero rows are prefill + # placeholders and a stop-token CB0 marks the post-EOS step. + request_id = "legacy_rows" + manager = _codec_transfer_manager(request_id) + request = _codec_request(request_id) + + assert ( + q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.zeros((4, 3), dtype=torch.long)}}, request) + is None + ) + assert q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[2150, 1, 1]])}}, request) is None + assert manager.code_prompt_token_ids[request_id] == [] + q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[5, 6, 7]])}}, request) + payload = q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[8, 9, 10]])}}, request) + + assert payload is not None # two frames fill the initial chunk + assert payload.codes.audio.tolist() == [5, 8, 6, 9, 7, 10] diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index b0178e9a5..4edfed603 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -620,6 +620,12 @@ def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) - assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)] assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)] + batch.is_prefilling_np[:] = True + batch.num_computed_prefill_tokens_np[:] = 0 + batch.num_scheduled_tokens = [1] * 4 + extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]])) + assert all(row.numel() == 0 for row in extra["embed"]["sampled"]) + # Speculative steps sample several tokens per row: not published. assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 69d9306a7..ee7f599f2 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -260,6 +260,7 @@ _BENCHMARK_TRANSITIONAL = ( ) _INTERNAL = ( + "VLLM_OMNI_TALKER_FIRST_AUDIO", "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS", "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S", "VLLM_OMNI_CONNECTOR_RECV_POLL_MS", diff --git a/vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml b/vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml new file mode 100644 index 000000000..c40a29209 --- /dev/null +++ b/vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml @@ -0,0 +1,50 @@ +# Opt-in CUDA Model Runner V2 profile for Qwen3-Omni-MoE: the Thinker (stage 0), +# Talker (stage 1) and Code2Wav (stage 2) run on MRv2 with its native data +# plane. The base profile keeps every stage on V1. +# +# On MRv2 the Thinker hands each generated token's embedding to the Talker with +# the step that sampled it (one Thinker step earlier than its next forward would +# capture the same embedding) and replays decode in FULL CUDA graphs together +# with the layers it captures for the Talker. +# +# When the stages share one GPU, run them under NVIDIA MPS (see +# docs/configuration/stage_configs.md); startup warns when no MPS server is found. +# +# MRv2 downstream stages support turn-based sessions only: realtime (resumable) +# requests and duplex sessions need the V1 chunk adapter. +base_config: qwen3_omni_moe.yaml + +connectors: + connector_of_shared_memory: + extra: + # First Code2Wav chunk is the stream's first frame alone, then chunks + # double up to the steady codec_chunk_frames (replaces the 4-frame + # initial chunk). Code2Wav re-decodes up to codec_left_context_frames + # before each chunk, so windows are 1, 3, 7, 15, 31, then 50 frames. + # A one-frame first chunk is also what lets the Talker deliver it + # itself. Opt in explicitly; VLLM_OMNI_TALKER_FIRST_AUDIO=0 disables it. + talker_first_audio: true + codec_chunk_ramp: [1, 2, 4, 8, 16, 25] + # Code2Wav decodes a streaming step as a few length groups on multi-row + # CUDA graphs instead of one eager call padded to its longest window + # (with the ramp, over half of those padded frames were waste). + codec_graph_batch_sizes: [2, 3, 4, 6, 8, 12, 16] + codec_fused_snake: true + +stages: + - stage_id: 0 + model_runner: v2 + - stage_id: 1 + model_runner: v2 + - stage_id: 2 + model_runner: v2 + +platforms: + npu: + model_runner: v1 + xpu: + model_runner: v1 + rocm: + model_runner: v1 + musa: + model_runner: v1 diff --git a/vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml b/vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml new file mode 100644 index 000000000..a5835dc02 --- /dev/null +++ b/vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml @@ -0,0 +1,27 @@ +# Opt-in single-H200 MRv2 profile, measured at client concurrency 32/64. +# Run all stages under NVIDIA MPS and warm representative shapes first. +# Retains Triton MoE and native nucleus reductions, fusing only the final +# sampling mask/Gumbel selection. cuDNN searches convolution algorithms +# during warmup; short-KV attention and convolution selection change rounding. +# This is not a claim of bitwise audio or universal P99/throughput improvement. +# Evidence: docs/design/qwen3_omni_mrv2_performance.md +base_config: qwen3_omni_moe_mrv2.yaml + +connectors: + connector_of_shared_memory: + extra: + code_predictor_kv_cache: true + code_predictor_short_kv: true + code_predictor_fused_sampling: true + codec_cudnn_benchmark: true + +stages: + - stage_id: 0 + devices: "0" + gpu_memory_utilization: 0.60 + - stage_id: 1 + devices: "0" + gpu_memory_utilization: 0.22 + - stage_id: 2 + devices: "0" + gpu_memory_utilization: 0.10 diff --git a/vllm_omni/model_executor/models/common/fused_sampling.py b/vllm_omni/model_executor/models/common/fused_sampling.py new file mode 100644 index 000000000..a4bfce13e --- /dev/null +++ b/vllm_omni/model_executor/models/common/fused_sampling.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused Gumbel selection after native nucleus probability reductions. + +Keep ATen's top-k threshold, sort, softmax and cumulative sum, including +boundary ties. Random numbers remain indexed by original token id. Only the +final mask, Gumbel transform and argmax are fused; no vocabulary scatter is +needed. This avoids changing which tokens survive a nucleus boundary through +different reduction orders. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _sorted_nucleus_gumbel( + logits, + indices, + probabilities, + cumulative, + uniforms, + output, + vocab: tl.constexpr, + logits_stride: tl.constexpr, + token_stride: tl.constexpr, + uniform_stride: tl.constexpr, + top_p: tl.constexpr, + block: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, block) + valid = offsets < vocab + values = tl.load(logits + row * logits_stride + offsets * token_stride, valid, other=-float("inf")).to(tl.float32) + if top_p < 1.0: + ids = tl.load(indices + row * vocab + offsets, valid, other=0) + probs = tl.load(probabilities + row * vocab + offsets, valid, other=0) + cdf = tl.load(cumulative + row * vocab + offsets, valid, other=0) + values = tl.where((cdf - probs) >= top_p, -float("inf"), values) + else: + ids = offsets + noise = tl.load(uniforms + row * uniform_stride + ids, valid, other=0.5) + score = values - tl.extra.cuda.libdevice.log(-tl.extra.cuda.libdevice.log(noise)) + # torch.argmax returns the first original token for a tie or NaN. + first_nan = tl.min(tl.where(valid & (score != score), ids, 2147483647), 0) + maximum = tl.max(tl.where(valid, score, -float("inf")), 0) + winner = tl.min(tl.where(valid & (score == maximum), ids, 2147483647), 0) + tl.store(output + row, tl.where(first_nan < 2147483647, first_nan, winner)) + + +def sample_top_k_top_p_gumbel( + logits: torch.Tensor, uniforms: torch.Tensor, *, top_k: int, top_p: float +) -> torch.Tensor: + """Sample CUDA rows without scattering sorted logits back to the vocabulary. + + Uniforms are caller-owned FP32 values with the same shape as logits; their + generation and generator advancement are deliberately outside this kernel. + Softmax/scan use the reference ATen operations. The Gumbel logarithms + still use libdevice, so this is not a universal bitwise-equivalence claim. + """ + if logits.ndim != 2 or not logits.is_cuda or uniforms.device != logits.device or torch.version.hip is not None: + raise ValueError("fused sampling requires two-dimensional CUDA logits and uniforms") + if uniforms.shape != logits.shape or uniforms.dtype != torch.float32 or uniforms.stride(1) != 1: + raise ValueError("uniforms must be FP32 with the logits shape and contiguous rows") + rows, vocab = logits.shape + # The stored-mode caller disables top-k for both 0 and negative values + # (vLLM commonly uses -1). Keep that behavior when fusion is enabled. + top_k = max(top_k, 0) + if vocab == 0 or top_k > vocab or not 0 < top_p <= 1: + raise ValueError("invalid top-k/top-p sampling parameters") + if top_k: + threshold = logits.topk(top_k, dim=-1).values[:, -1:] + logits = logits.masked_fill(logits < threshold, -float("inf")) + indices = probabilities = cumulative = None + if top_p < 1: + logits, indices = logits.sort(dim=-1, descending=True) + probabilities = logits.softmax(dim=-1, dtype=torch.float32) + cumulative = probabilities.cumsum(-1) + result = torch.empty((rows, 1), dtype=torch.long, device=logits.device) + if rows: + _sorted_nucleus_gumbel[(rows,)]( + logits, + indices, + probabilities, + cumulative, + uniforms, + result, + vocab, + logits.stride(0), + logits.stride(1), + uniforms.stride(0), + top_p, + triton.next_power_of_2(vocab), + num_warps=4, + enable_fp_fusion=False, + ) + return result diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index ce2cd1e5e..37a9c2388 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -18,6 +18,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models. from __future__ import annotations import dataclasses +import time from collections.abc import Iterable, Sequence import torch @@ -158,8 +159,8 @@ class CodePredictorAttention(nn.Module): """Multi-head self-attention for code predictor. Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm. - No KV cache -- the code predictor always re-prefills the full (short) - sequence each AR step. + The default path re-prefills the short sequence; the opt-in cached path + appends only new positions. Input : [B, seq_len, hidden_size] Output: [B, seq_len, hidden_size] @@ -316,6 +317,53 @@ class CodePredictorAttention(nn.Module): attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) return self.o_proj(attn_out) + def forward_cached( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + write_positions: torch.Tensor, + key_mask: torch.Tensor, + ) -> torch.Tensor: + """Attend only the new positions, reusing earlier positions' keys and values. + + ``k_cache``/``v_cache`` are ``[B, num_kv_heads, max_seq, head_dim]``; + the new positions' keys and values are written at ``write_positions`` + (``[n]``). ``key_mask`` (``[1, 1, n, max_seq]``, bool) keeps key + positions up to each query's, so slots not yet written in this call are + never read and every single-position step has the same shapes. + """ + bsz, seq_len, _ = hidden_states.shape + qkv = self.qkv_proj(hidden_states) + q_raw, k_raw, v_raw = self._split_qkv(qkv) + q = self.q_norm(q_raw.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2) + k = self.k_norm(k_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2) + v = v_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + q = (q * cos) + (_rotate_half(q) * sin) + k = (k * cos) + (_rotate_half(k) * sin) + k_cache.index_copy_(2, write_positions, k.to(k_cache.dtype)) + v_cache.index_copy_(2, write_positions, v.to(v_cache.dtype)) + if getattr(self, "use_short_kv_attention", False): + from .short_kv_attention import short_kv_attention + + attn_out = short_kv_attention(q, k_cache, v_cache, write_positions, self.scaling) + else: + attn_out = F.scaled_dot_product_attention( + q, + k_cache, + v_cache, + attn_mask=key_mask, + scale=self.scaling, + enable_gqa=self.is_gqa, + ) + attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) + return self.o_proj(attn_out) + # =================================================================== # MLP @@ -351,7 +399,7 @@ class CodePredictorMLP(nn.Module): class CodePredictorDecoderLayer(nn.Module): - """Transformer decoder layer (SDPA, no KV cache).""" + """Transformer decoder layer with default re-prefill and opt-in KV reuse.""" def __init__(self, config, *, prefix: str = "") -> None: super().__init__() @@ -383,9 +431,29 @@ class CodePredictorDecoderLayer(nn.Module): hidden_states = residual + hidden_states return hidden_states + def forward_cached( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + write_positions: torch.Tensor, + key_mask: torch.Tensor, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn.forward_cached( + hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + # =================================================================== -# Base Transformer Model (re-prefill, no KV cache) +# Base Transformer Model (re-prefill, or incremental with a K/V buffer) # =================================================================== @@ -454,6 +522,29 @@ class CodePredictorBaseModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states.to(input_dtype) + def forward_cached( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + kv_caches: Sequence[tuple[torch.Tensor, torch.Tensor]], + write_positions: torch.Tensor, + key_mask: torch.Tensor, + ) -> torch.Tensor: + """``forward`` for the new positions only (``position_ids`` ``[B, n]``), reusing earlier K/V.""" + input_dtype = inputs_embeds.dtype + use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type not in ("cpu", "npu") + if use_fp32: + inputs_embeds = inputs_embeds.float() + hidden_states = inputs_embeds + with torch.amp.autocast(inputs_embeds.device.type, enabled=use_fp32, dtype=torch.float32): + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer, (k_cache, v_cache) in zip(self.layers, kv_caches): + hidden_states = layer.forward_cached( + hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask + ) + hidden_states = self.norm(hidden_states) + return hidden_states.to(input_dtype) + # Mapping from HF checkpoint shard name to the (fused parameter, # slice index) pair used to assemble the fused tensor. See # ``load_weights`` below. @@ -626,8 +717,14 @@ class CodePredictorWrapper(nn.Module): 4. torch.compile on inner transformer. 5. Inline sampling (top-k + top-p) -- no custom op overhead. 6. Optional manual CUDA graph capture per batch-size bucket. + + ``code_predictor_kv_cache`` (CUDA, opt-in) switches to incremental K/V + decoding: at large batches the re-prefill FLOPs are not negligible. """ + _kv_cache_enabled: bool = False + _fused_sampling_enabled: bool = False + def __init__( self, *, @@ -672,6 +769,7 @@ class CodePredictorWrapper(nn.Module): # Sampling defaults for "stored" mode self._top_k: int = 50 self._top_p: float = 0.8 + self._do_sample: bool = True # Lazily initialised state self._proj_buf: torch.Tensor | None = None @@ -686,6 +784,7 @@ class CodePredictorWrapper(nn.Module): # ``None`` keeps the legacy power-of-two bucket derivation. self._execution_batch_buckets: list[int] | None = None prefix_graph_cfg = self._stage_connector_extra_config(vllm_config) + self._fused_sampling_enabled = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_fused_sampling")) prefix_graphs_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_prefix_graphs")) is_npu = current_omni_platform.is_npu() # MRv2 captures the whole Talker MTP call. It can therefore capture @@ -705,15 +804,28 @@ class CodePredictorWrapper(nn.Module): prefix_graph_cfg.get("code_predictor_prefix_graph_seq_lens") ) self._prefix_reprefill_seq_lens = tuple(self._prefix_seq_lens(self._num_groups + 1)) + # Incremental decoding (opt-in, CUDA): each AR step runs only its new + # position against a static K/V buffer instead of re-prefilling the + # whole prefix (at batch 64 the re-prefill runs 17x the positions). + self._kv_cache_enabled = ( + self._parse_bool_config(prefix_graph_cfg.get("code_predictor_kv_cache")) and current_omni_platform.is_cuda() + ) + self._short_kv_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_short_kv")) + self._short_kv_enabled = False + self._kv_buf: torch.Tensor | None = None + self._compiled_kv_fwd = None + # Per AR step: (new positions [n], key mask [1, 1, n, max_seq]). + self._kv_steps: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} def get_input_embeddings(self) -> nn.ModuleList: return self.model.get_input_embeddings() - def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8) -> None: - """Configure sampling parameters to maintain consistency with previous implementation.""" + def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8, do_sample: bool = True) -> None: + """Configure "stored"-mode sampling; ``do_sample=False`` selects greedy argmax.""" self._top_k = top_k self._top_p = top_p - logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f", top_k, top_p) + self._do_sample = do_sample + logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f, do_sample=%s", top_k, top_p, do_sample) # ------------------------------------------------------------------ # Lazy-init helpers @@ -730,6 +842,35 @@ class CodePredictorWrapper(nn.Module): ): return self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device) + if self._kv_cache_enabled: + cfg = self.config + head_dim = int(getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads) + self._kv_buf = torch.zeros( + int(cfg.num_hidden_layers), + 2, + bsz, + int(cfg.num_key_value_heads), + max_seq, + head_dim, + dtype=dtype, + device=device, + ) + + def _kv_views(self, bsz: int) -> list[tuple[torch.Tensor, torch.Tensor]]: + kv_buf = self._kv_buf + assert kv_buf is not None + return [(kv_buf[i, 0, :bsz], kv_buf[i, 1, :bsz]) for i in range(kv_buf.shape[0])] + + def _ensure_kv_steps(self, device: torch.device) -> None: + if self._kv_steps: + return + max_seq = self._num_groups + 1 + keys = torch.arange(max_seq, device=device) + for step in range(1, self._num_groups): + # Step 1 runs positions 0 (talker hidden) and 1 (CB0); step s > 1 runs position s. + positions = torch.arange(0, 2, device=device) if step == 1 else torch.tensor([step], device=device) + key_mask = (keys[None, :] <= positions[:, None]).view(1, 1, positions.numel(), max_seq) + self._kv_steps[step] = (positions, key_mask) def _setup_compile(self) -> None: """Lazily set up torch.compile with optional device graph capture.""" @@ -742,6 +883,9 @@ class CodePredictorWrapper(nn.Module): self._model_dtype = next(self.model.parameters()).dtype self._lm_heads_list = list(self.lm_head) self._codec_embeds_list = list(self.model.codec_embedding) + if self._model_dtype == torch.float16 or not current_omni_platform.supports_torch_inductor(): + # The fp16 body runs in fp32 under autocast; keep re-prefill there. + self._kv_cache_enabled = False # Torch 2.13 XPU Dynamo can double-register built-in handlers when # spawned workers compile this predictor. Keep this narrow path eager @@ -768,10 +912,28 @@ class CodePredictorWrapper(nn.Module): dynamic=False, options={"epilogue_fusion": False}, ) + self._short_kv_enabled = ( + self._kv_cache_enabled + and self._short_kv_requested + and current_omni_platform.is_cuda() + and self._model_dtype == torch.bfloat16 + ) + for layer in self.model.layers: + layer.self_attn.use_short_kv_attention = self._short_kv_enabled + if self._kv_cache_enabled: + self._compiled_kv_fwd = torch.compile( + self.model.forward_cached, + dynamic=self._short_kv_enabled, + options={"epilogue_fusion": False}, + ) with torch._dynamo.config.patch(cache_size_limit=self._compile_cache_size_limit()): self._warmup_buckets() - if self._wrapper_config.use_cuda_graphs: + if self._kv_cache_enabled: + # Fifteen short steps per call: they rely on the runner capturing + # the whole MTP call (MRv2) rather than on per-step inner graphs. + logger.info("code_predictor: torch.compile + incremental K/V decoding") + elif self._wrapper_config.use_cuda_graphs: self._capture_cuda_graphs() logger.info("code_predictor: torch.compile (no epilogue fusion) + CUDA graphs") else: @@ -840,13 +1002,13 @@ class CodePredictorWrapper(nn.Module): return row_generators @classmethod - def _sample_codes_gumbel( + def _sampling_uniforms( cls, logits: torch.Tensor, generator: _GeneratorLike = None, uniforms: torch.Tensor | None = None, ) -> torch.Tensor: - """Sample ``logits`` via Gumbel-max with optional precomputed noise.""" + """Draw one FP32 uniform per original token, preserving generator state.""" if uniforms is None: row_generators = cls._normalize_generators(generator, int(logits.shape[0])) u = torch.empty_like(logits, dtype=torch.float32) @@ -862,6 +1024,17 @@ class CodePredictorWrapper(nn.Module): f"logits={tuple(logits.shape)}" ) u = uniforms + return u + + @classmethod + def _sample_codes_gumbel( + cls, + logits: torch.Tensor, + generator: _GeneratorLike = None, + uniforms: torch.Tensor | None = None, + ) -> torch.Tensor: + """Sample ``logits`` via Gumbel-max with optional precomputed noise.""" + u = cls._sampling_uniforms(logits, generator, uniforms) return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True) def _prefix_seq_lens(self, max_seq: int) -> list[int]: @@ -924,6 +1097,8 @@ class CodePredictorWrapper(nn.Module): if needs_full_graph: full_graph_entries += len(prefix_buckets) required_entries = full_graph_entries + len(prefix_buckets) * len(prefix_seq_lens) + if self._kv_cache_enabled: + required_entries = 4 # {first step, later steps} x {batch 1, dynamic batch} return max(torch._dynamo.config.cache_size_limit, required_entries) def _warmup_buckets(self) -> None: @@ -938,7 +1113,29 @@ class CodePredictorWrapper(nn.Module): self._ensure_buffers(device, self._model_dtype, max(self._bucket_sizes)) proj_buf = self._proj_buf - if self._prefix_reprefill_enabled: + if self._kv_cache_enabled: + self._ensure_kv_steps(device) + started = time.perf_counter() + for bsz in self._bucket_sizes: + caches = self._kv_views(bsz) + for step in (1, 2): + positions, key_mask = self._kv_steps[step] + lo = 0 if step == 1 else step + hidden = proj_buf[:bsz, lo : step + 1, :] + pos_ids = positions.unsqueeze(0).expand(bsz, -1) + if bsz > 1: + # One graph for every batch >= 2 per step shape (4 compiles, not 2 per bucket): + # a cold per-bucket compile outlasted the stage handshake timeout. + for tensor in (hidden, pos_ids, *(t for kv in caches for t in kv)): + torch._dynamo.mark_dynamic(tensor, 0) + for _ in range(2): + self._compiled_kv_fwd(hidden, pos_ids, caches, positions, key_mask) + logger.info( + "code_predictor: incremental K/V decoding warmed for buckets %s in %.1f s", + self._bucket_sizes, + time.perf_counter() - started, + ) + elif self._prefix_reprefill_enabled: prefix_seq_lens = self._prefix_reprefill_seq_lens needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq)) for bsz in self._bucket_sizes: @@ -1206,9 +1403,29 @@ class CodePredictorWrapper(nn.Module): all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device) all_codes[:, 0] = layer0_code.reshape(bsz) + kv_caches = None + if self._kv_cache_enabled: + # Masked slots still enter the attention matmul with zero weight, so + # clear them: a non-finite value left by an earlier call would leak. + if not self._short_kv_enabled: + self._kv_buf[:, :, :padded_bsz].zero_() + kv_caches = self._kv_views(padded_bsz) # Autoregressive loop: predict layers 1..G-1 for step in range(1, num_groups): - logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing) + if kv_caches is not None: + # Step 1 runs positions 0-1; later steps append one position. + positions, key_mask = self._kv_steps[step] + lo = 0 if step == 1 else step + new_hidden = self._compiled_kv_fwd( + proj_buf[:padded_bsz, lo : step + 1, :], + positions.unsqueeze(0).expand(padded_bsz, -1), + kv_caches, + positions, + key_mask, + ) + logits = self._lm_heads_list[step - 1](new_hidden[:bsz, -1, :]) + else: + logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing) # Sample next code via Gumbel-max. # @@ -1222,7 +1439,15 @@ class CodePredictorWrapper(nn.Module): # defensive than ``multinomial`` around fully-masked/NaN inputs), # and the helper below can honor either one batch generator or one # generator per seeded row. - if stored_mode: + if stored_mode and not self._do_sample: + code = logits.argmax(dim=-1, keepdim=True) + elif stored_mode and self._fused_sampling_enabled and logits.is_cuda and torch.version.hip is None: + from .fused_sampling import sample_top_k_top_p_gumbel + + step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None + uniforms = self._sampling_uniforms(logits, sample_generator, step_uniforms) + code = sample_top_k_top_p_gumbel(logits, uniforms, top_k=s_top_k, top_p=s_top_p) + elif stored_mode: # "stored" mode: top-k -> top-p -> Gumbel-max if s_top_k > 0: topk_vals, _ = logits.topk(s_top_k, dim=-1) diff --git a/vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py b/vllm_omni/model_executor/models/common/short_kv_attention.py similarity index 100% rename from vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py rename to vllm_omni/model_executor/models/common/short_kv_attention.py diff --git a/vllm_omni/model_executor/models/common/talker_first_audio.py b/vllm_omni/model_executor/models/common/talker_first_audio.py new file mode 100644 index 000000000..ad67aff50 --- /dev/null +++ b/vllm_omni/model_executor/models/common/talker_first_audio.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared eligibility; each model separately controls its first-audio opt-in.""" + +from typing import Any + +import torch + +from vllm_omni.platforms import current_omni_platform + + +def supports_talker_first_audio(vllm_config: Any) -> bool: + """Require CUDA MRv2 streaming, local single-rank execution and no prefix replay. + + Other configurations retain the regular codec path. These conditions do + not enable the feature: model-specific options must also accept it. + """ + model = vllm_config.model_config + if not (getattr(model, "use_v2_model_runner", False) and getattr(model, "async_chunk", False)): + return False + parallel = vllm_config.parallel_config + return ( + current_omni_platform.is_cuda() + and torch.device(vllm_config.device_config.device).type == "cuda" + and parallel.tensor_parallel_size == 1 + and parallel.pipeline_parallel_size == 1 + and parallel.distributed_executor_backend in (None, "uni") + and not vllm_config.cache_config.enable_prefix_caching + ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py b/vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py new file mode 100644 index 000000000..4720071fe --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""First codec frame -> PCM inside the Qwen3-Omni Talker process. + +Code2Wav decodes each streaming chunk statelessly from its codes and left +context. With a one-frame first chunk (``codec_chunk_ramp`` starting at 1), +chunk 0 is a context-free decode of the stream's first frame, which the Talker +can run as soon as eager MTP completes that frame, removing the connector hop +and the Code2Wav step from time to first audio. The orchestrator then trims +Code2Wav's copy of those samples. + +The decoder is a second ``Qwen3OmniMoeCode2Wav`` built and loaded exactly like +the Code2Wav stage's, and captured with the same one-frame input shape as that +stage's ``(batch=1, frames=1)`` graph. Optional cuDNN autotuning can select +different algorithms in the two processes and change floating-point rounding. +Each bucket gets its own graph pool because it replays on a side stream while +the Talker's own graphs may still run. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Sequence +from typing import Any + +import torch +import torch.nn as nn +from torch.cuda import CUDAGraph +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorWrapper +from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp + +logger = init_logger(__name__) + + +def talker_first_audio_enabled(vllm_config: Any) -> bool: + """Opt in only when chunk 0 carries the first frame in the same step.""" + extra = CodePredictorWrapper._stage_connector_extra_config(vllm_config) + ramp = parse_chunk_ramp(extra) + initial_frames = ramp[0] if ramp is not None else int(extra.get("initial_codec_chunk_frames") or 0) + return ( + CodePredictorWrapper._parse_bool_config(extra.get("talker_first_audio")) + and initial_frames == 1 + and os.environ.get("VLLM_OMNI_TALKER_FIRST_AUDIO", "1") == "1" + and supports_talker_first_audio(vllm_config) + ) + + +class Qwen3OmniFirstFrameDecoder(nn.Module): + def __init__(self, code2wav: nn.Module, sample_rate: int) -> None: + super().__init__() + self.code2wav = code2wav + self.num_quantizers = int(code2wav.config.num_quantizers) + self.sample_rate = int(sample_rate) + self._graphs: dict[int, tuple[CUDAGraph, torch.Tensor, torch.Tensor]] = {} + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load the checkpoint's ``code2wav.*`` weights as the Code2Wav stage does.""" + loaded = self.code2wav.load_weights(weights) + if hasattr(self.code2wav, "precompute_snake_caches"): + self.code2wav.precompute_snake_caches() + return loaded + + @torch.inference_mode() + def capture(self, batch_sizes: Sequence[int] = (1, 2, 4, 8)) -> None: + device = next(self.code2wav.parameters()).device + if device.type != "cuda": + return + for batch_size in sorted(set(batch_sizes)): + static_input = torch.zeros(batch_size, self.num_quantizers, 1, dtype=torch.long, device=device) + self.code2wav(static_input) + torch.accelerator.synchronize(device) + graph = CUDAGraph() + with torch.cuda.graph(graph, pool=torch.cuda.graph_pool_handle()): + static_output = self.code2wav(static_input) + self._graphs[batch_size] = (graph, static_input, static_output) + logger.info("Captured Qwen3-Omni Talker first-frame decoder graphs for batch sizes %s", sorted(self._graphs)) + + @torch.inference_mode() + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """``codes`` [n, num_quantizers] long on device -> float32 PCM [n, samples] (new tensor).""" + if not self._graphs: + return self.code2wav(codes.to(torch.long).unsqueeze(-1))[:, 0, :].float() + outputs = [] + largest = max(self._graphs) + for start in range(0, int(codes.shape[0]), largest): + chunk = codes[start : start + largest] + rows = int(chunk.shape[0]) + batch_size = min(size for size in self._graphs if size >= rows) + graph, static_input, static_output = self._graphs[batch_size] + static_input.zero_() + static_input[:rows, :, 0].copy_(chunk) + graph.replay() + # A replay overwrites graph-owned storage; .float() aliases FP32 outputs. + outputs.append(static_output[:rows, 0, :].to(dtype=torch.float32, copy=True)) + return torch.cat(outputs, dim=0) diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 9caff6d27..7a6c07326 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -45,6 +45,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( engine_output_type="latent", custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_full_payload"), async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"), + # Takes effect only when the deploy selects model_runner v2 for this stage. + supports_native_mrv2_data_plane=True, sampling_constraints={"detokenize": True}, ), StagePipelineConfig( @@ -57,6 +59,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( sync_process_input_func=f"{_PROC}.thinker2talker_token_only", custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"), async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"), + # Each stage can use MRv2 when its deploy selects model_runner v2. + supports_native_mrv2_data_plane=True, sampling_constraints={ "detokenize": False, "stop_token_ids": [2150], @@ -72,6 +76,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( final_output_type="audio", hf_config_name="thinker_config", engine_output_type="audio", + supports_native_mrv2_data_plane=True, sampling_constraints={"detokenize": True}, requires_full_payload_input=True, ), diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index cfba92e98..50601d1fa 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -4,7 +4,7 @@ """Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav).""" import asyncio -from collections.abc import AsyncGenerator, Iterable +from collections.abc import AsyncGenerator, Iterable, Mapping from functools import cached_property from typing import Any @@ -50,6 +50,7 @@ from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload from vllm_omni.metrics import definitions as defs from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled from vllm_omni.model_executor.models.qwen3_omni.quantization import ( apply_outer_quant_config_mapping, ) @@ -61,7 +62,9 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.device_copy import to_device_nonblocking # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -88,6 +91,13 @@ TALKER_CODEC_THINK_EOS_ID = 4205 # Think mode end logger = init_logger(__name__) +def _cpu_long(values: Any) -> torch.Tensor: + """Token ids as a host tensor (payload ids arrive as lists or CPU tensors).""" + if isinstance(values, torch.Tensor): + return values.detach().to("cpu", torch.long) + return torch.as_tensor(values, dtype=torch.long) + + @MULTIMODAL_REGISTRY.register_processor( Qwen3OmniMoeThinkerMultiModalProcessor, info=Qwen3OmniMoeThinkerProcessingInfo, @@ -135,6 +145,11 @@ class Qwen3OmniMoeForConditionalGeneration( realtime_max_tokens = 64 + # Set by the MRv2 model state once eager Talker-MTP frames are enabled. + eager_frames_active = False + # Postprocess only stores the last hidden for the deferred decode MTP. + eager_frames_need_postprocess = False + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.have_multimodal_outputs = True @@ -180,9 +195,18 @@ class Qwen3OmniMoeForConditionalGeneration( # thinker layers and the forward must return what stock vLLM expects. self.is_staged_run = getattr(vllm_config.model_config, "model_stage", None) is not None self._returns_tuple = self.model_stage == "thinker" and self.is_staged_run + # The captured layers are a fixed-structure tensor dict on the token + # axis, so MRv2 may replay decode in FULL graphs and return them too. + self.supports_mrv2_full_graph_aux_outputs = self._returns_tuple if self.model_stage == "thinker": self.use_async_omni_output = True + # MRv2: the Talker's prompt ends with the first generated token's + # embedding; publishing it with the prefill output lets the first + # Thinker->Talker chunk leave one decode step earlier. + # Speculative sampling can change the token width between steps. + # Keep the existing capture-stream handoff for the entire request. + self.publishes_sampled_embeddings = self.is_staged_run and vllm_config.speculative_config is None # Initialize thinker model (multimodal processing + text generation) # Create a new vllm_config with thinker_config as the hf_config thinker_vllm_config = vllm_config.with_hf_config( @@ -233,6 +257,19 @@ class Qwen3OmniMoeForConditionalGeneration( ) self.model = self.talker self.code2wav = None + # Residual codebooks (MTP) output key, shared by the V1 runner + # (talker_mtp_*) and the MRv2 model state (mtp_*). + self.talker_mtp_output_key = ("codes", "audio") + # MRv2: the step that samples CB0 may complete the frame itself, + # since the residual codebooks depend only on CB0 and that step's + # last hidden, and the next input is their embedding sum plus the + # next text step (see ``talker_mtp``). + self.mtp_eager_frames = True + self._codec_codebook_size = int(talker_config.code_predictor_config.vocab_size) + self._apply_subtalker_sampling_params(getattr(vllm_config.model_config, "subtalker_sampling_params", None)) + + # Decode each stream's first frame here (see talker_first_audio_enabled). + self.first_frame_decoder = self._build_first_frame_decoder(vllm_config, code2wav_config, prefix) # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights self._init_special_tokens_embeddings() @@ -247,6 +284,9 @@ class Qwen3OmniMoeForConditionalGeneration( # talker MTP codec codes must stay on GPU to avoid a per-step D2H # sync stall; build_mm_cpu handles the eventual D2H at payload time. ("codes", "audio"), + # The growing thinker-embed cache is read and extended every + # decode step; on CPU it cost a blocking D2H per request per step. + ("embed", "cached_decode"), } elif self.model_stage == "code2wav": @@ -285,6 +325,74 @@ class Qwen3OmniMoeForConditionalGeneration( ["hidden_states", "residual", *incoming_captures], thinker_config.text_config.hidden_size ) + @staticmethod + def _build_first_frame_decoder( + vllm_config: VllmConfig, + code2wav_config: Qwen3OmniMoeCode2WavConfig, + prefix: str, + ) -> nn.Module | None: + """A Code2Wav copy for the Talker's first-frame audio, or None when disabled. + + Only a streaming (async-chunk) Talker hands frames to Code2Wav chunk by + chunk, so only there is a stream's first chunk a context-free decode. + The copy is built like the Code2Wav stage's module (same class, config, + dtype and device context). Optional cuDNN autotuning is also shared, + although per-process algorithm selection can change PCM rounding. + """ + if not talker_first_audio_enabled(vllm_config): + return None + from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder + + code2wav = init_vllm_registered_model( + vllm_config=vllm_config.with_hf_config(code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"]), + prefix=maybe_prefix(prefix, "first_frame_decoder.code2wav"), + hf_config=code2wav_config, + architectures=["Qwen3OmniMoeCode2Wav"], + ) + logger.info("Qwen3-Omni Talker decodes first-frame audio in-stage") + return Qwen3OmniFirstFrameDecoder(code2wav, defs.resolve_audio_sample_rate(code2wav_config)) + + def capture_first_frame_graphs(self) -> None: + """Called by the MRv2 runner after its own graph capture.""" + decoder = getattr(self, "first_frame_decoder", None) + if decoder is not None: + decoder.capture() + + def _apply_subtalker_sampling_params(self, params: Any) -> None: + """Apply the deploy's ``subtalker_sampling_params`` to the residual-codebook predictor. + + The predictor samples in its "stored" mode (top-k 50, top-p 0.8 by + default) and ignores per-call sampling arguments; ``do_sample: false`` + makes it greedy. + """ + if not isinstance(params, Mapping) or not params: + return + predictor = self.talker.code_predictor + predictor.set_sampling_params( + top_k=int(params.get("top_k", predictor._top_k)), + top_p=float(params.get("top_p", predictor._top_p)), + do_sample=bool(params.get("do_sample", True)), + ) + + # MRv2 capability names read the V1 canonical values, so a platform patch + # of ``talker_mtp_graph_safe`` applies to both runners. + @property + def mtp(self): + """MRv2 name for ``talker_mtp``; only the Talker stage has an MTP.""" + return self.talker_mtp if self.model_stage == "talker" else None + + @property + def mtp_output_key(self) -> tuple[str, str]: + return getattr(self, "talker_mtp_output_key", ("codes", "audio")) + + @property + def mtp_graph_safe(self) -> bool: + return bool(getattr(self, "talker_mtp_graph_safe", False)) + + def mtp_frame_valid(self, layer0: torch.Tensor) -> torch.Tensor: + """Rows whose CB0 is a codec id rather than codec EOS (or another special id).""" + return (layer0 >= 0) & (layer0 < self._codec_codebook_size) + def _thinker_capture_layer_indices(self) -> list[int]: accept_layer = getattr(self.talker_config, "accept_hidden_layer", None) return [0, int(accept_layer)] if self.is_staged_run and accept_layer is not None else [] @@ -496,38 +604,47 @@ class Qwen3OmniMoeForConditionalGeneration( # ========== Stage 3: Code2Wav ========== elif self.model_stage == "code2wav": seq_token_counts: list[int] | None = kwargs.get("seq_token_counts") + num_quantizers = int(self.code2wav.config.num_quantizers) # Extract codec codes from input - if input_ids.shape[0] % 16 == 0: + if input_ids.shape[0] % num_quantizers == 0: if seq_token_counts is not None: - max_seq_len = max(seq_token_counts) // 16 + max_seq_len = max(seq_token_counts) // num_quantizers batch_size = len(seq_token_counts) split_codes = torch.split(input_ids, seq_token_counts, dim=0) - codes = torch.zeros((batch_size, 16, max_seq_len), device=input_ids.device, dtype=input_ids.dtype) + codes = torch.zeros( + (batch_size, num_quantizers, max_seq_len), device=input_ids.device, dtype=input_ids.dtype + ) for idx, code in enumerate(split_codes): - seq_len = code.shape[0] // 16 - codes[idx, :, :seq_len] = code.reshape(16, seq_len) + seq_len = code.shape[0] // num_quantizers + codes[idx, :, :seq_len] = code.reshape(num_quantizers, seq_len) else: - codes = input_ids.reshape(1, 16, -1) + codes = input_ids.reshape(1, num_quantizers, -1) else: if seq_token_counts is None: logger.debug( - "Code2Wav warmup input length %s is not divisible by 16; padding with zeros.", + "Code2Wav warmup input length %s is not divisible by %s; padding with zeros.", input_ids.shape[0], + num_quantizers, ) else: logger.warning_once( - "Code2Wav input length is not divisible by 16; padding with zeros. " - "This is expected only during cudagraph warmup." + "Code2Wav input length is not divisible by %s; padding with zeros. " + "This is expected only during cudagraph warmup.", + num_quantizers, ) input_ids_flatten = input_ids.reshape(-1) input_ids_flatten = torch.cat( [ input_ids_flatten, - torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device), + torch.zeros( + num_quantizers - input_ids.shape[0] % num_quantizers, + dtype=torch.long, + device=input_ids.device, + ), ] ) - codes = input_ids_flatten.reshape(1, 16, -1) + codes = input_ids_flatten.reshape(1, num_quantizers, -1) # Generate audio from codec codes # Get every request's left_context_size from runtime_additional_information (passed via kwargs) @@ -541,6 +658,38 @@ class Qwen3OmniMoeForConditionalGeneration( logger.debug("No additional_information provided to code2wav stage.") audio_tensors = self.generate_audio(codes, left_context_size, seq_token_counts) + # Keep the skip decision attached to this forward's output, not + # mutable model state: asynchronous materialization can overlap steps. + from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + + flags = [ + bool(info.get("meta", {}).get("first_audio", False)) for info in (runtime_additional_information or []) + ] + if flags and any(flags): + if len(flags) != len(audio_tensors): + raise ValueError("First-audio flags must align with Code2Wav requests") + sample_rate = defs.resolve_audio_sample_rate(self.code2wav_config) + frame_counts = ( + [count // num_quantizers for count in seq_token_counts] if seq_token_counts else [codes.shape[-1]] + ) + # The causal decoder withholds a right-edge tail. The first + # frame has fewer samples than total_upsample; trim its actual + # prefix length, including when chunk 0 contains several frames. + prefix_lengths = [ + max(0, audio.numel() - (frames - 1) * int(self.code2wav.total_upsample)) if skip else 0 + for audio, frames, skip in zip(audio_tensors, frame_counts, flags, strict=True) + ] + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": [ + audio.reshape(1, -1)[..., prefix:] + for audio, prefix in zip(audio_tensors, prefix_lengths, strict=True) + ], + "sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in audio_tensors], + FIRST_AUDIO_REQUIRED_KEY: [torch.tensor(skip) for skip in flags], + }, + ) return audio_tensors # Fallback (shouldn't reach here) @@ -605,7 +754,13 @@ class Qwen3OmniMoeForConditionalGeneration( if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + # MRv2 passes each request's token span. + spans = kwargs.get("request_token_spans") + if spans is not None and self.eager_frames_active and self.vllm_config.model_config.async_chunk: + return self._make_eager_talker_output(talker_hidden, spans) code_predictor_codes = [info.get("codes", {}).get("audio") for info in info_dicts] + if spans is not None: + code_predictor_codes = self._align_codes_to_spans(code_predictor_codes, spans, talker_hidden.device) audio_codes = torch.cat(code_predictor_codes, dim=0) multimodal_outputs: OmniPayload = {"codes": {"audio": audio_codes}} span_len = audio_codes.shape[0] @@ -627,6 +782,44 @@ class Qwen3OmniMoeForConditionalGeneration( return model_outputs + def _align_codes_to_spans( + self, + codes: list[Any], + spans: list[tuple[int, int]], + device: torch.device, + ) -> list[torch.Tensor]: + """One ``codes.audio`` block per request with exactly its span's rows. + + MRv2 also runs requests that never went through preprocess (vLLM's + ``_warmup_`` requests), whose buffers carry no codes; they get zero + rows, which the downstream processor skips like prefill rows. + """ + q = int(self.talker.num_code_groups) + aligned: list[torch.Tensor] = [] + for block, (start, end) in zip(codes, spans, strict=True): + if isinstance(block, torch.Tensor) and block.ndim == 2 and block.shape[0] == end - start: + aligned.append(block) + else: + aligned.append(torch.zeros((end - start, q), dtype=torch.long, device=device)) + return aligned + + def _make_eager_talker_output(self, hidden: torch.Tensor, spans: list[tuple[int, int]]) -> OmniOutput: + """Async-chunk output when the MRv2 runner completes frames after sampling. + + ``run_eager_mtp`` writes each sampled row's frame and validity into the + last token row of its request span; every other row carries no frame. + """ + num_tokens = spans[-1][1] if spans else 0 + q = int(self.talker.num_code_groups) + mm: OmniPayload = { + "codes": {"audio": torch.zeros((num_tokens, q), dtype=torch.long, device=hidden.device)}, + "meta": { + "codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), + "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), + }, + } + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + # ==================== Audio Generation ==================== def generate_audio( @@ -849,7 +1042,7 @@ class Qwen3OmniMoeForConditionalGeneration( def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor: if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0: - xin = _ensure_1x1(x_opt).to(module_device) + xin = to_device_nonblocking(_ensure_1x1(x_opt), module_device) else: xin = torch.zeros( (1, thinker_embed.shape[-1]), @@ -882,27 +1075,26 @@ class Qwen3OmniMoeForConditionalGeneration( voice_type = str(voice_type).lower().strip() start_index = meta.get("num_processed_tokens", 0) end_index = start_index + input_embeds.shape[0] - # Read thinker outputs for prefill - thinker_sequence_embeds = embed["prefill"].to( - device=self._module_device(self.talker), dtype=torch.bfloat16 + # Read thinker outputs for prefill. Every copy here is H2D from the + # connector payload; pageable copies would stall the whole Talker batch. + talker_device = self._module_device(self.talker) + thinker_sequence_embeds = to_device_nonblocking(embed["prefill"], talker_device).to( + torch.bfloat16 ) # Tensor [P,H] - thinker_hidden_states = hs["output"].to( - device=self._module_device(self.talker), dtype=torch.bfloat16 - ) # Tensor [K,H] + thinker_hidden_states = to_device_nonblocking(hs["output"], talker_device).to(torch.bfloat16) # Tensor [K,H] + # Token ids stay on CPU too: prefill parsing reads them on the host. + thinker_sequences_cpu = None if ids.get("all") is None else _cpu_long(ids["all"]) + thinker_chatml_ids_cpu = None if ids.get("prompt") is None else _cpu_long(ids["prompt"]) thinker_sequences = ( - ids.get("all") - if ids.get("all") is None - else torch.as_tensor(ids["all"], device=self._module_device(self.talker)) + None if thinker_sequences_cpu is None else to_device_nonblocking(thinker_sequences_cpu, talker_device) ) thinker_chatml_ids = ( - ids.get("prompt") - if ids.get("prompt") is None - else torch.as_tensor(ids["prompt"], device=self._module_device(self.talker)) + None if thinker_chatml_ids_cpu is None else to_device_nonblocking(thinker_chatml_ids_cpu, talker_device) ) - tts_bos_thinker = embed["tts_bos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16) - tts_eos_thinker = embed["tts_eos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16) - tts_pad_thinker = embed["tts_pad"].to(device=self._module_device(self.talker), dtype=torch.bfloat16) + tts_bos_thinker = to_device_nonblocking(embed["tts_bos"], talker_device).to(torch.bfloat16) + tts_eos_thinker = to_device_nonblocking(embed["tts_eos"], talker_device).to(torch.bfloat16) + tts_pad_thinker = to_device_nonblocking(embed["tts_pad"], talker_device).to(torch.bfloat16) if thinker_sequence_embeds is None or thinker_hidden_states is None: raise ValueError( @@ -924,6 +1116,8 @@ class Qwen3OmniMoeForConditionalGeneration( ) if ids_chatml.ndim == 1: ids_chatml = ids_chatml.unsqueeze(0) + if thinker_chatml_ids_cpu is not None and thinker_chatml_ids_cpu.ndim == 1: + thinker_chatml_ids_cpu = thinker_chatml_ids_cpu.unsqueeze(0) else: # Fallback: create dummy ids if not provided ids_chatml = torch.zeros( @@ -932,6 +1126,7 @@ class Qwen3OmniMoeForConditionalGeneration( device=self._module_device(self.talker), ) thinker_sequences = ids_chatml + thinker_chatml_ids_cpu = thinker_sequences_cpu = None speaker_id = self._get_text_spk_token_id(voice_type) req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill( @@ -944,6 +1139,8 @@ class Qwen3OmniMoeForConditionalGeneration( tts_bos_thinker=tts_bos_thinker, tts_eos_thinker=tts_eos_thinker, tts_pad_thinker=tts_pad_thinker, + input_ids_cpu=thinker_chatml_ids_cpu, + thinker_result_ids_cpu=thinker_sequences_cpu, ) # Queue trailing_text_hidden for decode (drop first for next steps), @@ -993,12 +1190,11 @@ class Qwen3OmniMoeForConditionalGeneration( if cached_thinker_decode_embeds is None: update_dict.setdefault("embed", {})["cached_decode"] = thinker_decode_embeds else: - cached_thinker_decode_embeds = cached_thinker_decode_embeds.to( - device=self._module_device(self.talker), dtype=torch.bfloat16 - ) - thinker_decode_embeds = thinker_decode_embeds.to( - device=self._module_device(self.talker), dtype=torch.bfloat16 + talker_device = self._module_device(self.talker) + cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, talker_device).to( + torch.bfloat16 ) + thinker_decode_embeds = to_device_nonblocking(thinker_decode_embeds, talker_device).to(torch.bfloat16) update_dict.setdefault("embed", {})["cached_decode"] = torch.cat( [cached_thinker_decode_embeds, thinker_decode_embeds], dim=0 ) @@ -1015,26 +1211,31 @@ class Qwen3OmniMoeForConditionalGeneration( tts_bos_thinker: torch.Tensor | None = None, tts_eos_thinker: torch.Tensor | None = None, tts_pad_thinker: torch.Tensor | None = None, + input_ids_cpu: torch.Tensor | None = None, + thinker_result_ids_cpu: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """ Project thinker outputs to talker inputs during prefill stage. + ``input_ids_cpu`` / ``thinker_result_ids_cpu`` (host copies of the same + ids) move segment parsing off the GPU: each ``.item()`` / ``nonzero`` / + mask test on device ids is a host sync that stalls the Talker batch. + Returns: (input_ids, input_embeds) for talker """ target_len = thinker_result_ids.shape[-1] - im_start_indexes = torch.cat( - ( - torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(-1), - torch.tensor([target_len], device=input_ids.device, dtype=input_ids.dtype), - ), - dim=-1, - ) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here. + # Take batch 0 since batched inference is not supported here. + parse_ids = input_ids_cpu if input_ids_cpu is not None else input_ids + result_ids = thinker_result_ids_cpu if thinker_result_ids_cpu is not None else thinker_result_ids + im_start_indexes = torch.nonzero(parse_ids[0] == self.config.im_start_token_id).flatten().tolist() + im_start_indexes.append(target_len) multimodal_mask = ( - (thinker_result_ids == self.thinker_config.audio_token_id) | - (thinker_result_ids == self.thinker_config.image_token_id) | - (thinker_result_ids == self.thinker_config.video_token_id) - ).to(input_ids.device) # [t] # fmt: skip + (result_ids == self.thinker_config.audio_token_id) | + (result_ids == self.thinker_config.image_token_id) | + (result_ids == self.thinker_config.video_token_id) + ) # [t] # fmt: skip + role_tokens = parse_ids[0].tolist() tts_bos_embed, tts_eos_embed, tts_pad_embed = self._get_tts_embed( thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker @@ -1045,21 +1246,21 @@ class Qwen3OmniMoeForConditionalGeneration( trailing_text_hidden_all: torch.Tensor | None = None # For every chatml parts for i in range(len(im_start_indexes) - 1): - im_start_index = im_start_indexes[i].item() - segment_end_index = im_start_indexes[i + 1].item() - role_token = input_ids[0][im_start_index + 1] + im_start_index = im_start_indexes[i] + segment_end_index = im_start_indexes[i + 1] + role_token = role_tokens[im_start_index + 1] # Talker should ignore thinker system prompt - if (role_token == self.config.system_token_id).item(): + if role_token == self.config.system_token_id: continue # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs - elif (role_token == self.config.user_token_id).item(): + elif role_token == self.config.user_token_id: talker_user_part = self._get_talker_user_parts( im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed ) talker_input_embeds.append(talker_user_part) talker_input_ids.append(thinker_result_ids[im_start_index:segment_end_index]) # Take assistant output (for now) - elif (role_token == self.config.assistant_token_id).item() and i == len(im_start_indexes) - 2: + elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2: talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts( im_start_index, segment_end_index, @@ -1078,7 +1279,7 @@ class Qwen3OmniMoeForConditionalGeneration( except Exception: pass # History assistant output (ignore for now) - elif (role_token == self.config.assistant_token_id).item() and i != len(im_start_indexes) - 2: + elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2: continue else: raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)") @@ -1103,20 +1304,22 @@ class Qwen3OmniMoeForConditionalGeneration( cached_thinker_decode_embeds = embed.get("cached_decode", None) thinker_decode_embed = embed.get("decode", None) + if getattr(self.vllm_config.model_config, "use_v2_model_runner", False) and not meta.get("resumable", False): + return self._next_thinker_decode_text_step( + cached_thinker_decode_embeds, thinker_decode_embed, meta, device, update_dict + ) start_index = meta.get("num_processed_tokens", 0) if cached_thinker_decode_embeds is not None and start_index < cached_thinker_decode_embeds.shape[0]: - cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(device) + cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, device) thinker_embed = cached_thinker_decode_embeds[start_index] if thinker_decode_embed is not None: - thinker_decode_embed = thinker_decode_embed.to(device) + thinker_decode_embed = to_device_nonblocking(thinker_decode_embed, device) cached_thinker_decode_embeds = torch.cat([cached_thinker_decode_embeds, thinker_decode_embed], dim=0) update_dict.setdefault("embed", {})["cached_decode"] = cached_thinker_decode_embeds elif thinker_decode_embed is not None: - thinker_embed = thinker_decode_embed - if thinker_embed.device != device: - thinker_embed = thinker_embed.to(device) + thinker_embed = to_device_nonblocking(thinker_decode_embed, device) else: # When the tokens output by the thinker are exhausted, an EOS token needs to be appended. @@ -1129,6 +1332,41 @@ class Qwen3OmniMoeForConditionalGeneration( update_dict.setdefault("embed", {})["decode"] = None return self.talker.text_projection(thinker_embed).to(device) + def _next_thinker_decode_text_step( + self, + pending: torch.Tensor | None, + incoming: torch.Tensor | None, + meta: OmniPayloadMeta, + device: torch.device, + update_dict: OmniPayload, + ) -> torch.Tensor: + """Take the oldest unconsumed Thinker decode row as this step's text. + + Rows arrive in order but not one per Talker step: the MRv2 receiver + hands over every row that arrived since the previous step (the V1 + adapter delivers one chunk per step). New rows are appended to the + pending queue ``embed.cached_decode`` and each step consumes one. Once + the queue is empty the Thinker's text is exhausted: EOS once, then pad. + """ + embed_update = update_dict.setdefault("embed", {}) + embed_update["decode"] = None + if isinstance(incoming, torch.Tensor) and incoming.numel() > 0: + incoming = to_device_nonblocking(incoming, device).reshape(-1, incoming.shape[-1]) + if isinstance(pending, torch.Tensor) and pending.shape[0] > 0: + pending = to_device_nonblocking(pending, device) + pending = torch.cat([pending, incoming.to(pending.dtype)], dim=0) + else: + pending = incoming + if not isinstance(pending, torch.Tensor) or pending.shape[0] == 0: + embed_update["cached_decode"] = None + if meta.get("eos_emitted", False): + return self.tts_pad_embed.to(device) + update_dict.setdefault("meta", {})["eos_emitted"] = True + return self.tts_eos_embed.to(device) + pending = to_device_nonblocking(pending, device) + embed_update["cached_decode"] = pending[1:] if pending.shape[0] > 1 else None + return self.talker.text_projection(pending[:1]).to(device) + def talker_preprocess_decode( self, input_ids: torch.Tensor, input_embeds: torch.Tensor, update_dict: OmniPayload, payload: OmniPayload ): @@ -1203,6 +1441,21 @@ class Qwen3OmniMoeForConditionalGeneration( ) user_mm_mask = multimodal_mask[im_start_index:segment_end_index] + if user_mm_mask.device.type == "cpu": + # Host mask: no device sync for the test, and a text-only segment + # (the common TTS case) needs no masked gather at all. + if not bool(user_mm_mask.any()): + user_thinker_embed = thinker_embed[im_start_index:segment_end_index] + return self.talker.text_projection(user_thinker_embed).to(thinker_hidden.device, torch.bfloat16) + mm_index = to_device_nonblocking(torch.nonzero(user_mm_mask).flatten(), thinker_hidden.device) + text_index = to_device_nonblocking(torch.nonzero(~user_mm_mask).flatten(), thinker_hidden.device) + segment_hidden = thinker_hidden[im_start_index:segment_end_index] + segment_embed = thinker_embed[im_start_index:segment_end_index] + mm_hidden = self.talker.hidden_projection(segment_hidden.index_select(0, mm_index)) + user_talker_part.index_copy_(0, mm_index, mm_hidden.to(user_talker_part.dtype)) + text_hidden = self.talker.text_projection(segment_embed.index_select(0, text_index)) + user_talker_part.index_copy_(0, text_index, text_hidden.to(user_talker_part.dtype)) + return user_talker_part # Multimodal data exists if user_mm_mask.any(): user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask] @@ -1236,17 +1489,19 @@ class Qwen3OmniMoeForConditionalGeneration( ), dim=0, ) - codec_special_tokens = torch.tensor( - [ - self.config.talker_config.codec_nothink_id, - self.config.talker_config.codec_think_bos_id, - self.config.talker_config.codec_think_eos_id, - speaker_id, - self.config.talker_config.codec_pad_id, - self.config.talker_config.codec_bos_id, - ], - device=tts_pad_embed.device, - dtype=torch.long, + codec_special_tokens = to_device_nonblocking( + torch.tensor( + [ + self.config.talker_config.codec_nothink_id, + self.config.talker_config.codec_think_bos_id, + self.config.talker_config.codec_think_eos_id, + speaker_id, + self.config.talker_config.codec_pad_id, + self.config.talker_config.codec_bos_id, + ], + dtype=torch.long, + ), + tts_pad_embed.device, ) embed_input_ids = self.talker.embed_input_ids(codec_special_tokens).to( device=tts_pad_embed.device, dtype=torch.bfloat16 @@ -1381,18 +1636,31 @@ class Qwen3OmniMoeForConditionalGeneration( # ==================== Weight Loading ==================== - def _get_codec_frame_config(self) -> tuple[int, int]: - """Extract codec_chunk_frames and codec_left_context_frames from stage connector config.""" + def _get_codec_connector_extra(self) -> dict[str, Any]: model_cfg = getattr(self.vllm_config, "model_config", None) connector_cfg = getattr(model_cfg, "stage_connector_config", None) if isinstance(connector_cfg, dict): extra = connector_cfg.get("extra", {}) else: extra = getattr(connector_cfg, "extra", None) or {} + return extra if isinstance(extra, Mapping) else {} + + def _get_codec_frame_config(self) -> tuple[int, int]: + """Extract codec_chunk_frames and codec_left_context_frames from stage connector config.""" + extra = self._get_codec_connector_extra() chunk_frames = int(extra.get("codec_chunk_frames", 0) or 0) left_frames = int(extra.get("codec_left_context_frames", 0) or 0) return chunk_frames, left_frames + def _get_codec_ramp_windows(self) -> list[int]: + """Decode windows of the connector's ``codec_chunk_ramp`` (empty without a ramp).""" + extra = self._get_codec_connector_extra() + chunk_frames = int(extra.get("codec_chunk_frames", 25) or 25) + ramp = parse_chunk_ramp(dict(extra), steady=chunk_frames) + if not ramp: + return [] + return ramp_decode_windows(ramp, int(extra.get("codec_left_context_frames", 25) or 0)) + def _maybe_enable_code2wav_cudagraph(self) -> None: """Enable the inner Code2Wav CUDA graph unless this stage runs in eager mode.""" if not self.code2wav or not hasattr(self.code2wav, "enable_cudagraph"): @@ -1404,9 +1672,15 @@ class Qwen3OmniMoeForConditionalGeneration( return chunk_frames, left_frames = self._get_codec_frame_config() + streaming = bool(self.vllm_config.model_config.async_chunk) + # Opt-in multi-row graphs for streaming windows (connector extra + # ``codec_graph_batch_sizes``, e.g. [2, 3, 4, 6, 8, 12, 16]). + batch_sizes = self._get_codec_connector_extra().get("codec_graph_batch_sizes") or () self.code2wav.enable_cudagraph( codec_chunk_frames=chunk_frames, codec_left_context_frames=left_frames, + extra_capture_sizes=self._get_codec_ramp_windows() if streaming else (), + streaming_batch_sizes=[int(b) for b in batch_sizes] if streaming else (), ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1439,6 +1713,12 @@ class Qwen3OmniMoeForConditionalGeneration( loaded_weights.update(talker_loaded) loaded_weights.update(self._init_special_tokens_embeddings()) + # The Talker's first-frame decoder is a Code2Wav copy (same code2wav.* weights). + first_frame_decoder = getattr(self, "first_frame_decoder", None) + if first_frame_decoder is not None and code2wav_weights: + decoder_loaded = first_frame_decoder.load_weights(code2wav_weights) + loaded_weights.update(add_prefix_to_loaded_weights(decoder_loaded, "first_frame_decoder.code2wav")) + # Load code2wav weights if self.code2wav and code2wav_weights: code2wav_loaded = self.code2wav.load_weights(code2wav_weights) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index bc8f6b3b4..9ae711870 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -5,7 +5,8 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Callable, Iterable, Sequence +from contextlib import nullcontext import numpy as np import torch @@ -19,6 +20,7 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( Qwen3OmniMoeCode2WavDecoderBlock, Qwen3OmniMoeCode2WavTransformerModel, Qwen3OmniMoeConvNeXtBlock, + Qwen3OmniMoeSnakeBeta, ) from vllm.config import VllmConfig # type: ignore from vllm.logger import init_logger # type: ignore @@ -31,9 +33,93 @@ from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta from vllm_omni.model_executor.models.qwen3_omni.quantization import ( Qwen3OmniNestedSupportsQuant, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +# Fixed cost of one streaming decode call (graph replay or eager launch), in +# decoded-frame equivalents: a call on a single frame costs about as much GPU +# time as decoding this many extra frames inside a larger call. +_DECODE_CALL_COST_FRAMES = 32 + + +def use_fused_snake(module: nn.Module) -> int: + """Replace HF ``Qwen3OmniMoeSnakeBeta`` activations under ``module`` with the fused one. + + The HF decoder blocks recompute ``exp(alpha)`` / ``exp(beta)`` and launch + about seven broadcast kernels per call; the shared ``SnakeBeta`` computes the + same ``x + 1/b * sin^2(a*x)`` in one kernel from precomputed caches. + Parameter names and shapes match, so this runs before weight loading. + """ + count = 0 + for parent in list(module.modules()): + for name, child in list(parent.named_children()): + if isinstance(child, Qwen3OmniMoeSnakeBeta): + setattr(parent, name, SnakeBeta(child.in_features)) + count += 1 + return count + + +def plan_decode_groups( + lengths: Sequence[int], + bucket_of: Callable[[int], int | None], + batch_sizes_for: Callable[[int], Sequence[int]], + call_cost_frames: int = _DECODE_CALL_COST_FRAMES, +) -> list[tuple[list[int], int]]: + """Split a streaming batch into decode calls of rows with similar lengths. + + Decoding every row at the batch's longest window wastes most of the work + when short ramp chunks (1-15 frames) share a step with steady 50-frame + windows. Rows are sorted by length and cut into contiguous groups of + graph buckets (``bucket_of``) so that the modelled cost, per call + ``call_cost_frames`` plus padded rows (to the next captured row count of + that bucket, ``batch_sizes_for``) times the bucket length, is smallest. + Returns ``(row indices, decode length)`` per call; the decode length is + the group's longest row. + """ + count = len(lengths) + if count == 0: + return [] + order = sorted(range(count), key=lambda row: lengths[row]) + buckets = [bucket_of(int(lengths[row])) or int(lengths[row]) for row in order] + distinct = sorted(set(buckets)) + members = [[row for row, bucket in zip(order, buckets) if bucket == size] for size in distinct] + + def sizes_for(size: int) -> list[int]: + return sorted({int(b) for b in batch_sizes_for(size) if int(b) > 0} | {1}) + + def cost(rows: int, size: int) -> int: + batch_sizes = sizes_for(size) + calls, rest = divmod(rows, batch_sizes[-1]) + padded = calls * batch_sizes[-1] + if rest: + calls += 1 + padded += next(b for b in batch_sizes if b >= rest) + return calls * call_cost_frames + padded * size + + best = [0] + [None] * len(distinct) + cut = [0] * (len(distinct) + 1) + for end in range(1, len(distinct) + 1): + rows = 0 + for start in range(end - 1, -1, -1): + rows += len(members[start]) + total = best[start] + cost(rows, distinct[end - 1]) + if best[end] is None or total < best[end]: + best[end], cut[end] = total, start + groups: list[tuple[list[int], int]] = [] + end = len(distinct) + while end > 0: + start = cut[end] + groups.append(([row for bucket_rows in members[start:end] for row in bucket_rows], distinct[end - 1])) + end = start + calls: list[tuple[list[int], int]] = [] + for rows, size in reversed(groups): + max_batch = sizes_for(size)[-1] + for offset in range(0, len(rows), max_batch): + part = rows[offset : offset + max_batch] + calls.append((part, max(int(lengths[row]) for row in part))) + return calls + class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): """ @@ -74,6 +160,14 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): super().__init__() self.config: Qwen3OmniMoeCode2WavConfig = vllm_config.model_config.hf_config + connector = getattr(vllm_config.model_config, "stage_connector_config", None) + extra = connector.get("extra", {}) if isinstance(connector, dict) else getattr(connector, "extra", {}) + # Both the separate stage and the Talker's first-frame copy receive + # this connector config. Keep algorithm selection consistent between + # them, and leave other models' cuDNN settings unchanged. + self._cudnn_benchmark = (extra or {}).get("codec_cudnn_benchmark", False) + if not isinstance(self._cudnn_benchmark, bool): + raise ValueError("codec_cudnn_benchmark must be a boolean") # Calculate total upsampling factor self.total_upsample = np.prod(self.config.upsample_rates + self.config.upsampling_ratios) @@ -122,10 +216,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): Qwen3OmniMoeCausalConvNet(output_dim, 1, kernel_size=7), ] self.decoder = nn.ModuleList(decoder) + # Keep existing V1 and non-CUDA decoder blocks unchanged. + if ( + (extra or {}).get("codec_fused_snake", False) + and current_omni_platform.is_cuda() + and torch.device(vllm_config.device_config.device).type == "cuda" + ): + use_fused_snake(self.decoder) # CUDA Graph support — reuses CUDAGraphDecoderWrapper from Qwen3-TTS self._cudagraph_enabled = False self._cudagraph_wrapper = None + # Captured row counts per graph size; empty keeps one padded decode per batch. + self._streaming_batch_sizes: dict[int, list[int]] = {} def precompute_snake_caches(self): """Precompute exp(alpha) and 1/(exp(beta)+eps) for all SnakeBeta modules.""" @@ -142,8 +245,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): device: torch.device | None = None, codec_chunk_frames: int = 0, codec_left_context_frames: int = 0, + extra_capture_sizes: Iterable[int] = (), + streaming_batch_sizes: Iterable[int] = (), ): - """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav).""" + """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav). + + ``extra_capture_sizes`` adds exact-size graphs (e.g. a chunk ramp's + decode windows) to the default buckets, so those chunks replay without + zero padding. ``streaming_batch_sizes`` (e.g. ``[2, 4, 8, 16]``) also + captures multi-row graphs for every bucket up to one streaming window + (chunk + left context); a streaming batch is then decoded as a few + length groups (``plan_decode_groups``) instead of one eager call + padded to its longest row. + """ from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import ( CUDAGraphDecoderWrapper, ) @@ -154,8 +268,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): logger.warning("Cannot enable CUDA Graph: not on CUDA device (got %s)", device) return + extra = {int(size) for size in extra_capture_sizes if int(size) > 0} + batch_sizes = sorted({int(b) for b in streaming_batch_sizes if int(b) > 1}) + capture_sizes = None + extra_shapes: list[tuple[int, int]] = [] + if extra or batch_sizes: + capture_sizes = sorted( + set( + CUDAGraphDecoderWrapper.compute_capture_sizes( + codec_chunk_frames=codec_chunk_frames, + codec_left_context_frames=codec_left_context_frames, + ) + ) + | extra + ) + window = codec_chunk_frames + codec_left_context_frames + streaming_sizes = [size for size in capture_sizes if window <= 0 or size <= window] + # A multi-row graph never holds more frames than the largest + # single-row one, so the shared graph pool does not grow. + largest = max(capture_sizes) + extra_shapes = [(b, size) for b in batch_sizes for size in streaming_sizes if b * size <= largest] wrapper = CUDAGraphDecoderWrapper( decoder=self, + capture_sizes=capture_sizes, + extra_capture_shapes=extra_shapes, num_quantizers=self.config.num_quantizers, enabled=True, ) @@ -172,10 +308,18 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): raise self._cudagraph_wrapper = wrapper self._cudagraph_enabled = True + if extra_shapes and torch.cuda.is_available(): + # Return the eager warm-up activations: this stage shares the GPU. + torch.accelerator.empty_cache() + self._streaming_batch_sizes = { + size: [1, *sorted(b for b, shape_size in extra_shapes if shape_size == size)] + for size in {shape_size for _b, shape_size in extra_shapes} + } logger.info( - "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s", + "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s, multi-row shapes=%s", self.config.num_quantizers, self._cudagraph_wrapper.capture_sizes, + extra_shapes, ) def forward(self, codes: torch.Tensor) -> torch.Tensor: @@ -200,6 +344,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state # Shape: [batch, seq_len, hidden_size] + # Algorithm search happens during eager warmup before graph capture. + # Restore process settings afterward, including on a failed decode. + # Captured replays use the selected kernels without this Python path. + cudnn = torch.backends.cudnn + context = ( + cudnn.flags( + enabled=cudnn.enabled, + benchmark=True, + benchmark_limit=10, + deterministic=cudnn.deterministic, + allow_tf32=cudnn.allow_tf32, + ) + if ( + self._cudnn_benchmark + and hidden.is_cuda + and torch.version.hip is None + and not torch.cuda.is_current_stream_capturing() + ) + else nullcontext() + ) + with context: + return self._decode_waveform(hidden) + + def _decode_waveform(self, hidden: torch.Tensor) -> torch.Tensor: # Stage 3: Upsampling hidden = hidden.permute(0, 2, 1) # [batch, hidden_size, seq_len] for blocks in self.upsample: @@ -305,28 +473,38 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): "defaulting to left_context_size=zeros(len(codes)). This is expected during cudagraph warmup." ) left_context_size = [0] * codes.shape[0] - # Decode chunk - wavs = [] - if self._cudagraph_enabled and self._cudagraph_wrapper is not None: - batch_wav = self._cudagraph_wrapper.decode(codes) - else: - batch_wav = self(codes) if seq_token_counts is not None: code_seq_lens = [n // self.config.num_quantizers for n in seq_token_counts] else: # Fallback: assume all batch elements share the same sequence length. code_seq_lens = [codes.shape[-1]] * codes.shape[0] - # The decoder trims the same right-edge tail from every batch row. - # Infer it from the decoded window, not a shorter request's length: - # padding that request with codec zeros does not provide valid context. - tail = max(0, int(codes.shape[-1] * self.total_upsample) - batch_wav.shape[-1]) - for idx, code_seq_len in enumerate(code_seq_lens): - # Refill the previous chunk's tail using the re-decoded left context, - # and exclude the current row's padded tail just as singleton decode does. - start = max(0, left_context_size[idx] * self.total_upsample - tail) - end = max(0, code_seq_len * self.total_upsample - tail) - wav_chunk = batch_wav[idx, :, start:end] - wavs.append(wav_chunk) + batch, width = int(codes.shape[0]), int(codes.shape[-1]) + graphs = self._cudagraph_wrapper if self._cudagraph_enabled else None + batch_sizes = getattr(self, "_streaming_batch_sizes", None) + if graphs is not None and batch_sizes and batch > 1: + groups = plan_decode_groups( + code_seq_lens, graphs._get_padded_size, lambda size: batch_sizes.get(size, (1,)) + ) + else: + groups = [(list(range(batch)), width)] + wavs: list[torch.Tensor | None] = [None] * batch + for rows, length in groups: + if len(rows) == batch and length == width: + group_codes = codes + else: + # Basic indexing only: a device index tensor would need a host copy. + group_codes = torch.stack([codes[row, :, :length] for row in rows]) + batch_wav = graphs.decode(group_codes) if graphs is not None else self(group_codes) + # The decoder trims the same right-edge tail from every batch row. + # Infer it from the decoded window, not a shorter request's length: + # padding that request with codec zeros does not provide valid context. + tail = max(0, int(length * self.total_upsample) - batch_wav.shape[-1]) + for position, row in enumerate(rows): + # Refill the previous chunk's tail using the re-decoded left context, + # and exclude the current row's padded tail just as singleton decode does. + start = max(0, left_context_size[row] * self.total_upsample - tail) + end = max(0, code_seq_lens[row] * self.total_upsample - tail) + wavs[row] = batch_wav[position, :, start:end] return wavs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py b/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py index 81eb8fdda..4481eb50d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py +++ b/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py @@ -18,8 +18,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( CodePredictorDecoderLayer, _rotate_half, ) - -from .short_kv_attention import short_kv_attention +from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention def _attention_cached( diff --git a/vllm_omni/model_executor/models/qwen3_tts/first_audio.py b/vllm_omni/model_executor/models/qwen3_tts/first_audio.py index 1396c0dc4..4b10b7507 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/first_audio.py +++ b/vllm_omni/model_executor/models/qwen3_tts/first_audio.py @@ -4,9 +4,7 @@ from typing import Any -import torch - -from vllm_omni.platforms import current_omni_platform +from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio def talker_first_audio_enabled(vllm_config: Any) -> bool: @@ -18,16 +16,4 @@ def talker_first_audio_enabled(vllm_config: Any) -> bool: from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor extra = Predictor._stage_connector_extra_config(vllm_config) - model = vllm_config.model_config - parallel = vllm_config.parallel_config - return ( - Predictor._parse_bool_config(extra.get("talker_first_audio")) - and current_omni_platform.is_cuda() - and torch.device(vllm_config.device_config.device).type == "cuda" - and bool(getattr(model, "use_v2_model_runner", False)) - and bool(getattr(model, "async_chunk", False)) - and parallel.tensor_parallel_size == 1 - and parallel.pipeline_parallel_size == 1 - and parallel.distributed_executor_backend in (None, "uni") - and not vllm_config.cache_config.enable_prefix_caching - ) + return Predictor._parse_bool_config(extra.get("talker_first_audio")) and supports_talker_first_audio(vllm_config) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 675c5eecd..9620d0d86 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -55,6 +55,9 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp talker_hidden_size=int(talker_config.hidden_size), prefix=prefix, ) + # TTS owns its frame-local cache and overrides _predict_step_logits. + # Do not also activate the shared wrapper's Omni incremental loop. + self._kv_cache_enabled = False # Store talker_config for backward compat (accessed by some callers) self.talker_config = talker_config self._vllm_config = vllm_config diff --git a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py index 9b9462ad9..774e5940e 100644 --- a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py +++ b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import logging from dataclasses import dataclass, field @@ -99,6 +99,21 @@ def ramp_cumulative(index: int, ramp: list[int], steady: int) -> int: return sum(ramp) + (index + 1 - ramp_len) * steady +def ramp_decode_windows(ramp: list[int], left_context: int) -> list[int]: + """Decode window (left context + new frames) of each ramp chunk for a stateless decoder. + + Chunk ``i`` re-decodes up to ``left_context`` earlier frames with its + ``ramp[i]`` new ones, e.g. ``[1, 2, 4, 8, 16, 25]`` with 25 frames of + context -> ``[1, 3, 7, 15, 31, 50]``. + """ + windows: list[int] = [] + emitted = 0 + for new_frames in ramp: + windows.append(new_frames + min(emitted, max(0, left_context))) + emitted += new_frames + return windows + + def compute_ramp_emit( length: int, chunk_index: int, diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 3cef0eed8..3fee73ac3 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The Qwen team. """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition.""" import logging from collections.abc import Mapping from dataclasses import dataclass, field -from typing import Any +from typing import Any, cast import torch from vllm.inputs import TextPrompt @@ -23,6 +23,7 @@ from vllm_omni.data_entry_keys import ( ) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import compute_ramp_emit, parse_chunk_ramp from vllm_omni.model_executor.stage_input_processors.tts_utils import ( extract_language_from_prompt, extract_language_from_request, @@ -76,6 +77,43 @@ def _get_accept_hidden_layer_index(transfer_manager: Any) -> int: return int(mc.hf_config.talker_config.accept_hidden_layer) +# Set once chunk 0 carried the first generated token taken from the prefill +# sample: from then on each chunk carries the token its step sampled. +_SAMPLED_TEXT_STREAM = "_qwen3_omni_sampled_text_stream" + + +def _is_final_token(request: Any) -> bool | None: + """Whether the token sampled last ends the Thinker's turn (``None``: unknown).""" + output_token_ids = _ensure_list(getattr(request, "output_token_ids", None) or []) + count = getattr(request, "output_token_count", None) + count = len(output_token_ids) if count is None else count + token = getattr(request, "last_output_token_id", None) + if token is None and output_token_ids: + token = output_token_ids[-1] + if token is None or count == 0: + return None + params = getattr(request, "sampling_params", None) + stop_ids = getattr(params, "all_stop_token_ids", None) or getattr(params, "stop_token_ids", None) or () + max_tokens = getattr(params, "max_tokens", None) + return token in stop_ids or (max_tokens is not None and count >= max_tokens) + + +def _sampled_token_embed(thinker_embed: Mapping[str, Any], request: Any) -> torch.Tensor | None: + """Embedding of the token this step sampled, when published and not final. + + The Talker is fed the embeddings of the Thinker's generated tokens except + the final one (a stop token, or the last one ``max_tokens`` allows), which + no forward ever processes. For any other token the published embedding is + exactly what the next step's layer-0 capture holds, one step later. + """ + sampled = thinker_embed.get("sampled") + if not isinstance(sampled, torch.Tensor) or sampled.dim() != 2 or sampled.shape[0] != 1: + return None + if getattr(request, "resumable", False) or _is_final_token(request) is not False: + return None + return sampled.detach().cpu() + + def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int: im_start_token_id = 151644 system_token_id = 8948 @@ -298,6 +336,7 @@ def _get_qwen3_streaming_state( streaming_context: Any | None, ) -> _Qwen3OmniStreamingState: bridge_states = getattr(streaming_context, "bridge_states", None) + assert bridge_states is not None per_model_state = bridge_states.setdefault("qwen3_omni", {}) state = per_model_state.get(request_id) if state is None: @@ -387,7 +426,7 @@ def thinker2talker_async_chunk( return None thinker_hs = multimodal_output.get("hidden_states", {}) - thinker_layers = thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {} + thinker_layers = cast(Mapping[int | str, Any], thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {}) thinker_embed_raw = multimodal_output.get("embed", {}) thinker_embed = thinker_embed_raw if isinstance(thinker_embed_raw, dict) else {} @@ -441,27 +480,48 @@ def thinker2talker_async_chunk( speaker=speaker, language=language, ) - if transfer_manager.request_payload.get(request_id) is None: - if not is_finished: - transfer_manager.request_payload[request_id] = to_dict(payload) - return None - else: - save_payload = transfer_manager.request_payload.pop(request_id) + assert payload.embed is not None and payload.embed.prefill is not None + assert payload.hidden_states is not None and payload.hidden_states.output is not None + save_payload = transfer_manager.request_payload.pop(request_id, None) + if save_payload is not None: payload.embed.prefill = torch.cat( (save_payload.get("embed", {}).get("prefill"), payload.embed.prefill), dim=0 ) payload.hidden_states.output = torch.cat( (save_payload.get("hidden_states", {}).get("output"), payload.hidden_states.output), dim=0 ) - prefill_shape = payload.embed.prefill.shape[0] - if not is_finished and prefill_shape <= len(prompt_token_ids): - transfer_manager.request_payload[request_id] = to_dict(payload) - return None + if not is_finished and payload.embed.prefill.shape[0] == len(prompt_token_ids): + # The Talker prompt ends with the first generated token's embedding, + # which the next decode step would capture; take it from the + # prefill step's sample instead when the producer published it. + first_token_embed = _sampled_token_embed(thinker_embed, request) + if first_token_embed is not None: + payload.embed.prefill = torch.cat((payload.embed.prefill, first_token_embed), dim=0) + transfer_manager.request_payload[request_id] = {_SAMPLED_TEXT_STREAM: True} + if not is_finished and payload.embed.prefill.shape[0] <= len(prompt_token_ids): + transfer_manager.request_payload[request_id] = to_dict(payload) + return None else: if request.resumable: return _construct_thinker2talker_streaming_input_async_chunk( is_finished, request, thinker_emb, thinker_hid, transfer_manager ) + stream = transfer_manager.request_payload.get(request_id) + if stream is not None and stream.get(_SAMPLED_TEXT_STREAM): + # This step's capture was already sent as the previous step's sample. + sampled = _sampled_token_embed(thinker_embed, request) + if sampled is None: + # A partial (including resumed) prefill has no accepted sample. + # Its captures replay text already sent before preemption. + raw_sample = thinker_embed.get("sampled") + if isinstance(raw_sample, torch.Tensor) and raw_sample.numel() == 0: + return None + if _is_final_token(request) is not True and not is_finished: + raise RuntimeError(f"Thinker sample embedding missing for request {request_id}") + # The final token is never part of the text; a finish goes out + # as the bare finish marker. + return None + thinker_emb = sampled if thinker_emb.shape[0] > 1: logger.warning( "Unexpected multiple embeddings in thinker2talker_async_chunk for chunk_id %d: " @@ -486,7 +546,7 @@ def thinker2talker_full_payload( transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, -) -> dict[str, Any] | None: +) -> OmniPayload | None: """Pack complete thinker output for the non-async connector path.""" rid = getattr(request, "request_id", None) if not isinstance(pooling_output, Mapping): @@ -612,7 +672,7 @@ def thinker2talker_token_only( ) thinker_sequences = prompt_token_ids + output_ids thinker_input_ids = prompt_token_ids - info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} + info_for_len: OmniPayload = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu") # Keep this fallback until the connector reliably preserves voice metadata. additional_information = to_dict( @@ -637,6 +697,14 @@ def thinker2talker_token_only( # ========================= +def _last_codec_row(codes: Any) -> torch.Tensor | None: + """The last ``[1, Q]`` row of token-major codec codes, or None if there is none.""" + if not isinstance(codes, torch.Tensor) or codes.ndim != 2 or codes.shape[0] == 0: + return None + # Own the row: it outlives this step's (batch-wide) output tensor. + return codes[-1:].to(torch.long).clone() + + def talker2code2wav_async_chunk( transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any], @@ -648,27 +716,47 @@ def talker2code2wav_async_chunk( """ request_id = request.external_req_id code_predictor_codes = None + frame_valid = None + first_audio = False if isinstance(multimodal_output, Mapping): talker_codes = multimodal_output.get("codes", {}) if isinstance(talker_codes, dict): code_predictor_codes = talker_codes.get("audio") + meta = multimodal_output.get("meta", {}) + if isinstance(meta, Mapping): + frame_valid = meta.get("codec_frame_valid") + flag = meta.get("first_audio", False) + first_audio = ( + bool(flag.numel() and flag.reshape(-1)[-1].item()) if isinstance(flag, torch.Tensor) else bool(flag) + ) - sampling_params = getattr(request, "sampling_params", None) - stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) - stop_token_id = getattr(sampling_params, "stop_token_id", None) - if stop_token_id is not None: - stop_token_ids.add(stop_token_id) + if isinstance(frame_valid, torch.Tensor) and frame_valid.numel() > 0: + # Token-major codes with explicit per-row validity (MRv2 eager + # frames): a step's frame, if any, is the last row of its span, and a + # multi-row prefill span carries its first frame there. + frame = _last_codec_row(code_predictor_codes) + append_codes = frame is not None and bool(frame_valid.reshape(-1)[-1].item()) + code_predictor_codes = frame + else: + # One row per step (V1 and MRv2 deferred frames): all-zero rows are + # prefill placeholders and a stop-token CB0 marks the post-EOS step. + sampling_params = getattr(request, "sampling_params", None) + stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) + stop_token_id = getattr(sampling_params, "stop_token_id", None) + if stop_token_id is not None: + stop_token_ids.add(stop_token_id) - append_codes = ( - isinstance(code_predictor_codes, torch.Tensor) - and code_predictor_codes.numel() > 0 - and bool(code_predictor_codes.any()) - ) - if append_codes: - first_codebook = int(code_predictor_codes[0, 0].item()) - if first_codebook in stop_token_ids: - logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) - append_codes = False + append_codes = ( + isinstance(code_predictor_codes, torch.Tensor) + and code_predictor_codes.numel() > 0 + and bool(code_predictor_codes.any()) + ) + if append_codes: + assert isinstance(code_predictor_codes, torch.Tensor) + first_codebook = int(code_predictor_codes[0, 0].item()) + if first_codebook in stop_token_ids: + logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) + append_codes = False if append_codes: transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes) elif not is_finished: @@ -680,12 +768,42 @@ def talker2code2wav_async_chunk( chunk_size_config = int(cfg.get("codec_chunk_frames", 25)) left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0) + # The connector config is fixed for the adapter's lifetime: parse once. + if not hasattr(transfer_manager, "_qwen3_omni_chunk_ramp"): + transfer_manager._qwen3_omni_chunk_ramp = parse_chunk_ramp(cfg, steady=chunk_size_config) + ramp = transfer_manager._qwen3_omni_chunk_ramp chunk_id = transfer_manager.put_req_chunk[request_id] length = len(transfer_manager.code_prompt_token_ids[request_id]) if length <= 0: return None + if ramp is not None: + # Ramp (replaces initial_codec_chunk_frames): chunk i carries ramp[i] + # new frames, then codec_chunk_frames. Code2Wav decodes statelessly, so + # every chunk re-sends up to codec_left_context_frames earlier frames. + # The counter restarts with each segment, like code_prompt_token_ids. + chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) + emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size_config, bool(is_finished)) + if not emit or context_length <= 0: + # Finished on a chunk boundary: the transport sends the finish marker. + return None + left_context_size = min(length - context_length, left_context_size_config) + end_index = left_context_size + context_length + codes = ( + torch.cat(transfer_manager.code_prompt_token_ids[request_id][-end_index:], dim=0) + .transpose(0, 1) + .reshape(-1) + ) + return OmniPayloadStruct( + codes=CodesStruct(audio=codes), + meta=MetaStruct( + first_audio=torch.tensor(first_audio and chunk_id == 0), + left_context_size=left_context_size, + finished=torch.tensor(is_finished, dtype=torch.bool), + ), + ) + if configured_initial_chunk_size > 0: if chunk_id == 0: chunk_size_config = configured_initial_chunk_size @@ -715,6 +833,7 @@ def talker2code2wav_async_chunk( return OmniPayloadStruct( codes=CodesStruct(audio=codes), meta=MetaStruct( + first_audio=torch.tensor(first_audio and chunk_id == 0), left_context_size=left_context_size, finished=torch.tensor(is_finished, dtype=torch.bool), ), diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 6700e26b5..f20a0f2db 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -890,15 +890,15 @@ class OmniModelState(DefaultModelState): finishing = computed + scheduled >= input_batch.prefill_len_np[:num_reqs] kept = ~input_batch.is_prefilling_np[:num_reqs] | finishing rows = np.flatnonzero(kept).tolist() - if not rows: - return None device = sampled_token_ids.device - first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device)) - embeds = self.model.embed_input_ids(first.long()) - empty = torch.empty(0, dtype=embeds.dtype) - sampled: list[torch.Tensor] = [empty] * num_reqs - for k, row in enumerate(rows): - sampled[row] = embeds[k : k + 1] + # Preserve the empty per-row prefill marker even when every request + # is prefilling; absence would look like a broken decode handoff. + sampled: list[torch.Tensor] = [torch.empty(0)] * num_reqs + if rows: + first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device)) + embeds = self.model.embed_input_ids(first.long()) + for k, row in enumerate(rows): + sampled[row] = embeds[k : k + 1] done = torch.cuda.Event() done.record() extra: dict[str, Any] = {"embed": {"sampled": sampled}}