Merge branch 'main' into fix/pr4709-review-followup

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -18,8 +18,11 @@ import types
 
 import pytest
 import torch
+import torch.nn.functional as F
 from pytest_mock import MockerFixture
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 # Direct file import to avoid vllm_omni.__init__ patch dependencies.
 _MODELS = os.path.join(
     os.path.dirname(__file__),
@@ -46,6 +49,11 @@ def _load_module(name: str, filename: str):
 
 def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
     """Build the dict of modules to inject into sys.modules."""
+
+    class NativeCustomOp(torch.nn.Module):
+        def forward(self, *args, **kwargs):
+            return self.forward_native(*args, **kwargs)
+
     platforms_mock = mocker.MagicMock()
     platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False
     platforms_mock.current_omni_platform.is_npu.return_value = False
@@ -73,6 +81,8 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
 
     vllm_parallel_mock = mocker.MagicMock()
     vllm_parallel_mock.VocabParallelEmbedding = torch.nn.Embedding
+    custom_op_mock = types.ModuleType("vllm_omni.diffusion.layers.custom_op")
+    custom_op_mock.CustomOp = NativeCustomOp
 
     return {
         "vllm_omni": mocker.MagicMock(),
@@ -82,6 +92,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
         "vllm.config.vllm": vllm_config_mod,
         "vllm.model_executor.model_loader.weight_utils": weight_utils_mock,
         "vllm.model_executor.layers.vocab_parallel_embedding": vllm_parallel_mock,
+        "vllm_omni.diffusion.layers.custom_op": custom_op_mock,
         "vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"),
         "vllm_omni.model_executor.models": models_pkg,
         "vllm_omni.model_executor.models.common": common_pkg,
@@ -165,6 +176,38 @@ def _make_vllm_config(mocker: MockerFixture, max_num_seqs: int = 4):
     return vllm_config
 
 
+def test_npu_custom_ops_use_fused_norm_and_cached_rope(mocker: MockerFixture, loaded_target_classes) -> None:
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    cp_config, _ = _make_tiny_config(loaded_target_classes)
+    rms_calls = []
+
+    def npu_rms_norm(hidden_states, weight, epsilon):
+        rms_calls.append((hidden_states, weight, epsilon))
+        return hidden_states + 1, None
+
+    mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=True)
+    mocker.patch.object(
+        common_mod,
+        "torch_npu",
+        types.SimpleNamespace(npu_rms_norm=npu_rms_norm),
+        create=True,
+    )
+
+    hidden_states = torch.zeros(1, 2, cp_config.hidden_size, dtype=torch.float16)
+    norm = common_mod._RMSNorm(cp_config.hidden_size, eps=cp_config.rms_norm_eps)
+    torch.testing.assert_close(norm.forward_npu(hidden_states), hidden_states + 1)
+    assert len(rms_calls) == 1
+    assert rms_calls[0][2] == cp_config.rms_norm_eps
+
+    rotary = common_mod._RotaryEmbedding(cp_config)
+    assert rotary.cos_cached.shape == (cp_config.num_code_groups + 1, cp_config.head_dim)
+    assert rotary.sin_cached.shape == rotary.cos_cached.shape
+    position_ids = torch.tensor([[0, 2, 4]])
+    cos, sin = rotary.forward_npu(hidden_states, position_ids)
+    torch.testing.assert_close(cos, rotary.cos_cached[position_ids].to(torch.float16))
+    torch.testing.assert_close(sin, rotary.sin_cached[position_ids].to(torch.float16))
+
+
 class TestCodePredictorDtypeAlignment:
     """Test that code predictor buffers match model parameter dtype."""
 
@@ -339,6 +382,91 @@ class TestCodePredictorDtypeAlignment:
         assert not torch.equal(first[:, 1:], different[:, 1:])
 
 
+class TestCodePredictorPerRowGenerators:
+    """Seeded requests must stay deterministic inside a multi-row batch (#4883)."""
+
+    def _make_predictor(self, mocker: MockerFixture, loaded_target_classes):
+        _, _, code_predictor_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_npu", return_value=False)
+        cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        predictor = code_predictor_wrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor._wrapper_config.use_cuda_graphs = False
+        return predictor, talker_config
+
+    def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        multinomial = common_mod.CodePredictorWrapper._multinomial
+        torch.manual_seed(7)
+        probs = torch.softmax(torch.randn(3, 64), dim=-1)
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=probs.device)
+            generator.manual_seed(seed)
+            return generator
+
+        batched = multinomial(probs, None, [seeded(11), None, seeded(22)])
+        assert batched.shape == (3, 1)
+        # Each seeded row must reproduce a standalone draw from the same seed,
+        # independent of what the other rows in the batch consume.
+        assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11)))
+        assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22)))
+
+    def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        torch.manual_seed(123)
+        row_embed = torch.randn(1, hidden)
+        row_hidden = torch.randn(1, hidden)
+        # Three identical rows: two share a seed, one differs.
+        layer0_code = torch.zeros(3, dtype=torch.long)
+        layer0_embed = row_embed.expand(3, hidden).contiguous()
+        last_talker_hidden = row_hidden.expand(3, hidden).contiguous()
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=layer0_code.device)
+            generator.manual_seed(seed)
+            return generator
+
+        def run(generators):
+            return predictor(
+                layer0_code=layer0_code,
+                layer0_embed=layer0_embed,
+                last_talker_hidden=last_talker_hidden,
+                do_sample=True,
+                temperature=0.9,
+                top_k=50,
+                top_p=1.0,
+                generators=generators,
+            )
+
+        first = run([seeded(1234), seeded(1234), seeded(4321)])
+        second = run([seeded(1234), seeded(1234), seeded(4321)])
+
+        # Same per-row seeds -> the whole batch reproduces across calls.
+        assert torch.equal(first, second)
+        # Identical inputs + identical seeds -> identical rows within a call.
+        assert torch.equal(first[0], first[1])
+        # A different seed on the same inputs must diverge in the residual layers.
+        assert not torch.equal(first[0, 1:], first[2, 1:])
+
+    def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        with pytest.raises(ValueError, match="one entry per row"):
+            predictor(
+                layer0_code=torch.zeros(2, dtype=torch.long),
+                layer0_embed=torch.randn(2, hidden),
+                last_talker_hidden=torch.randn(2, hidden),
+                generators=[torch.Generator()],
+            )
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
@@ -457,3 +585,92 @@ class TestCodePredictorWrapperConfig:
             wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True),
         )
         assert graph_wrapper._prefix_graphs_enabled is True
+
+
+class TestGumbelMaxSampling:
+    """Sanity tests for the shared Gumbel-max sampling helper."""
+
+    @staticmethod
+    def _sample(logits: torch.Tensor, generator=None) -> torch.Tensor:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        return common_mod.CodePredictorWrapper._sample_codes_gumbel(logits, generator=generator)
+
+    def test_gumbel_max_picks_unique_high_logit(self, loaded_target_classes) -> None:
+        """If one logit is dominant, Gumbel-max must always pick it."""
+        torch.manual_seed(0)
+        bsz, vocab = 4, 64
+        logits = torch.full((bsz, vocab), -1e4)
+        logits[:, 7] = 0.0  # only finite entry
+        out = self._sample(logits)
+        assert out.shape == (bsz, 1)
+        assert torch.equal(out.squeeze(-1), torch.full((bsz,), 7, dtype=torch.long))
+
+    def test_gumbel_max_handles_single_finite_entry_rows(self, loaded_target_classes) -> None:
+        """Degenerate rows with one surviving finite entry should still return
+        a valid index instead of tripping multinomial-style asserts."""
+        bsz, vocab = 2, 8
+        logits = torch.full((bsz, vocab), float("-inf"))
+        logits[0, 3] = 0.0
+        logits[1, 5] = 0.0
+        out = self._sample(logits)
+        assert out[0].item() == 3
+        assert out[1].item() == 5
+
+    def test_gumbel_max_generator_makes_results_reproducible(self, loaded_target_classes) -> None:
+        """Same ``Generator.manual_seed`` should produce identical samples."""
+        torch.manual_seed(0)
+        logits = torch.randn(8, 32)
+        g1 = torch.Generator()
+        g1.manual_seed(123)
+        g2 = torch.Generator()
+        g2.manual_seed(123)
+        g3 = torch.Generator()
+        g3.manual_seed(456)
+        s1 = self._sample(logits, generator=g1)
+        s2 = self._sample(logits, generator=g2)
+        s3 = self._sample(logits, generator=g3)
+        assert torch.equal(s1, s2)
+        # Different seeds must (almost surely) yield a different sequence.
+        assert not torch.equal(s1, s3)
+
+    def test_gumbel_max_per_row_generators_match_individual_rows(self, loaded_target_classes) -> None:
+        """A per-row generator list should preserve row independence."""
+        torch.manual_seed(11)
+        logits = torch.randn(3, 16)
+        seeds = [101, 202, 303]
+        batched_generators = []
+        expected_rows = []
+        for row, seed in enumerate(seeds):
+            row_generator = torch.Generator()
+            row_generator.manual_seed(seed)
+            batched_generators.append(row_generator)
+
+            solo_generator = torch.Generator()
+            solo_generator.manual_seed(seed)
+            expected_rows.append(self._sample(logits[row : row + 1], generator=solo_generator))
+
+        batched = self._sample(logits, generator=batched_generators)
+        expected = torch.cat(expected_rows, dim=0)
+        assert torch.equal(batched, expected)
+
+    def test_gumbel_max_rejects_generator_length_mismatch(self, loaded_target_classes) -> None:
+        logits = torch.randn(2, 8)
+        generators = [torch.Generator()]
+        with pytest.raises(ValueError, match="Expected 2 per-row generators"):
+            self._sample(logits, generator=generators)
+
+    def test_gumbel_max_distribution_matches_softmax(self, loaded_target_classes) -> None:
+        """Empirical frequency under Gumbel-max should be close to the
+        softmax probability mass.  Loose tolerance because we draw a
+        finite number of samples."""
+        torch.manual_seed(7)
+        vocab = 16
+        logits = torch.randn(vocab) * 2.0
+        target = F.softmax(logits, dim=-1, dtype=torch.float32)
+
+        n_draws = 20000
+        batched_logits = logits.unsqueeze(0).expand(n_draws, vocab).contiguous()
+        samples = self._sample(batched_logits).squeeze(-1)
+        empirical = torch.bincount(samples, minlength=vocab).float() / n_draws
+        # L1 distance under 0.03 with 20k draws is comfortable for vocab=16.
+        assert (empirical - target).abs().sum().item() < 0.05

tests/model_executor/models/test_kv_cache_scale_mapper.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for vLLM's KV-cache scale mapper load path."""
+
+from collections.abc import Iterable
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_STALE_API_FILES = [
+    "vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py",
+    "vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py",
+    "vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py",
+    "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py",
+]
+
+_SOURCE_SCALE_NAME = "layers.0.self_attn.k_proj.output_scale"
+_MAPPED_SCALE_NAME = "layers.0.self_attn.attn.k_scale"
+
+
+class _RecordingCacheScaleMapper:
+    def __init__(self) -> None:
+        self.called = False
+
+    def apply(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        self.called = True
+        for name, weight in weights:
+            yield name.replace(".k_proj.output_scale", ".attn.k_scale"), weight
+
+
+class _QuantConfigWithoutOldCacheScale:
+    def __init__(self) -> None:
+        self.mapper = _RecordingCacheScaleMapper()
+
+    def get_cache_scale_mapper(self) -> _RecordingCacheScaleMapper:
+        return self.mapper
+
+
+class _FakeModel:
+    def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None:
+        self.quant_config = quant_config
+        self.param = torch.nn.Parameter(torch.zeros(()), requires_grad=False)
+
+    def named_parameters(
+        self,
+        remove_duplicate: bool = True,
+    ) -> Iterable[tuple[str, torch.nn.Parameter]]:
+        del remove_duplicate
+        return [(_MAPPED_SCALE_NAME, self.param)]
+
+
+class _FakeHunyuanModel(_FakeModel):
+    def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None:
+        super().__init__(quant_config)
+        self.config = SimpleNamespace(
+            tie_word_embeddings=False,
+            num_attention_heads=1,
+            num_key_value_heads=1,
+            use_cla=False,
+        )
+
+    def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]:
+        return [], {}
+
+    def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return weight
+
+
+@pytest.mark.parametrize("rel_path", _STALE_API_FILES)
+def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> None:
+    """Affected loaders must not call the removed vLLM get_cache_scale API."""
+    source = (_REPO_ROOT / rel_path).read_text(encoding="utf-8")
+    assert ".get_cache_scale(" not in source
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name", "fake_model_cls"),
+    [
+        (
+            "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3",
+            "HunyuanModel",
+            _FakeHunyuanModel,
+        ),
+        (
+            "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old",
+            "Qwen2Model",
+            _FakeModel,
+        ),
+    ],
+)
+def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale(
+    module_name: str,
+    class_name: str,
+    fake_model_cls: type[_FakeModel],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Delegated loaders must not expect the removed get_cache_scale method."""
+    module = __import__(module_name, fromlist=[class_name])
+    monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False)
+
+    fake_model = fake_model_cls(_QuantConfigWithoutOldCacheScale())
+    load_weights = getattr(module, class_name).load_weights
+
+    loaded = load_weights(fake_model, [(_MAPPED_SCALE_NAME, torch.ones(()))])
+
+    assert loaded == {_MAPPED_SCALE_NAME}
+    torch.testing.assert_close(fake_model.param.data, torch.ones(()))
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name"),
+    [
+        (
+            "vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2",
+            "MammothModa2Qwen2ForCausalLM",
+        ),
+        (
+            "vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm",
+            "MiMoAudioLLMForConditionalGeneration",
+        ),
+    ],
+)
+def test_direct_custom_loaders_apply_cache_scale_mapper(
+    module_name: str,
+    class_name: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Direct custom loaders must apply the new mapper before fall-through load."""
+    module = __import__(module_name, fromlist=[class_name])
+    monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False)
+
+    quant_config = _QuantConfigWithoutOldCacheScale()
+    fake_model = _FakeModel(quant_config)
+    load_weights = getattr(module, class_name).load_weights
+
+    loaded = load_weights(fake_model, [(_SOURCE_SCALE_NAME, torch.ones(()))])
+
+    assert quant_config.mapper.called
+    assert loaded == {_MAPPED_SCALE_NAME}
+    torch.testing.assert_close(fake_model.param.data, torch.ones(()))

tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Static regression checks for VoxCPM2 runner-assisted unified graph hooks.
+
+These tests intentionally avoid importing torch/vLLM so they can run in a
+lightweight local checkout. Runtime audio quality still requires a CUDA test.
+"""
+
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py"
+RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py"
+TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py"
+VOXCPM2_RUNTIME_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "runtime_config.py"
+VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py"
+VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py"
+OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py"
+DEPLOY = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm2.yaml"
+
+
+def test_ar_runner_exposes_runner_assisted_full_metadata_hook():
+    source = RUNNER.read_text()
+
+    assert "Models without this hook keep the normal runner path" in source
+    assert "get_runner_assisted_full_attention_metadata_request" in source
+    assert "set_runner_assisted_full_attention_metadata_context" in source
+    assert "runner_assisted_full_attn" in source
+    assert "runner_assisted_full_attn_capture" in source
+    assert "pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn" in source
+    assert "for_cudagraph_capture=runner_assisted_full_attn_capture" in source
+    assert "cudagraph_runtime_mode=(" in source
+    assert "CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode" in source
+    refresh_source = source[source.index("def _refresh_runner_assisted_full_attention_metadata_buffers") :]
+    refresh_source = refresh_source[: refresh_source.index("def _set_runner_assisted_full_attention_metadata_context")]
+    assert "num_computed_tokens_cpu" in refresh_source
+    assert "num_scheduled_tokens_np" in refresh_source
+    assert "optimistic_seq_lens_cpu[:num_reqs]=1" not in "".join(refresh_source.split())
+    padding_source = source[
+        source.index("runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request") :
+    ]
+    padding_source = padding_source[: padding_source.index("ubatch_slices, ubatch_slices_padded")]
+    assert "runner_assisted_full_attn_request.num_reqs_padded" in padding_source
+    assert "runner_assisted_full_attn_request.for_cudagraph_capture" in padding_source
+    assert "num_reqs_padded * max_num_scheduled_tokens" in padding_source
+
+
+def test_ar_runner_without_model_hook_stays_on_normal_path():
+    source = RUNNER.read_text()
+    contract_source = RUNNER_ASSISTED_METADATA.read_text()
+    request_source = source[source.index("def _get_runner_assisted_full_attention_metadata_request") :]
+    request_source = request_source[
+        : request_source.index("def _refresh_runner_assisted_full_attention_metadata_buffers")
+    ]
+    compact_request_source = "".join(request_source.split())
+
+    assert "Models without this hook keep the normal runner path" in request_source
+    assert "class RunnerAssistedAttentionMetadataProvider(Protocol)" in contract_source
+    assert "class RunnerAssistedFullAttentionMetadataRequest(NamedTuple)" in contract_source
+    assert "tuple[int, bool]" not in contract_source
+    assert 'hook=getattr(self.model,"get_runner_assisted_full_attention_metadata_request",None)' in (
+        compact_request_source
+    )
+    assert "ifnotcallable(hook):returnNone" in compact_request_source
+    assert ".tolist()" not in compact_request_source
+    assert "isinstance(request,RunnerAssistedFullAttentionMetadataRequest)" in compact_request_source
+    assert "request.num_reqs_padded" in request_source
+    assert "request.for_cudagraph_capture" in request_source
+    assert "exceptException" not in compact_request_source
+
+    set_context_source = source[source.index("def _set_runner_assisted_full_attention_metadata_context") :]
+    set_context_source = set_context_source[: set_context_source.index("def _deferred_prefix_cache_mm_keys")]
+    assert "except Exception" not in set_context_source
+
+
+def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
+    source = TALKER.read_text()
+    runtime_config_source = VOXCPM2_RUNTIME_CONFIG.read_text()
+    compact_source = "".join(source.split())
+
+    assert "_voxcpm2_compile_without_inductor_cudagraphs" in source
+    assert "def _compile_without_inductor_cudagraphs" not in source
+    assert "_voxcpm2_unwrap_torch_compile" in source
+    assert "_orig_mod" in source
+    compile_source = source[source.index("def _voxcpm2_compile_without_inductor_cudagraphs") :]
+    compile_source = compile_source[: compile_source.index("def _setup_torch_compile")]
+    assert "torch._inductor.config" not in compile_source
+    assert '"triton.cudagraphs": False' in compile_source
+    assert '"triton.cudagraph_trees": False' in compile_source
+
+    assert "self._enable_unified_decode_graph=self._runtime_config.unified_decode_graph_available(" in compact_source
+    assert "andnotself.deterministic_cfm_noise" in "".join(runtime_config_source.split())
+    assert "decode_tail_graph" not in source
+
+    unified_source = source[source.index("def _forward_unified_decode") :]
+    unified_source = unified_source[: unified_source.index("# -------------------- vllm hooks")]
+    assert "except Exception:" not in unified_source
+    assert "_failed_unified_graph_sizes" not in source
+    assert "self._enable_unified_decode_graph = False" not in unified_source
+    assert "_forward_unified_decode_fallback" not in source
+    assert "capture_failed" not in compact_source
+    assert "record_fallback" not in source
+    assert "fallbacks" not in source
+    assert "**info_dict: Unpack[VoxCPM2PreprocessInput]" in source
+    assert "**info: Unpack[VoxCPM2PostprocessInput]" in source
+    assert "reference_audio: Any" not in source
+    assert "prompt_audio: Any" not in source
+    assert "def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike" in source
+    assert "class _PrefillInputs(NamedTuple)" in source
+    assert "-> _PrefillInputs" in source
+    assert 'inputs["audio_feat"]' not in source
+    assert 'inputs["text_token"]' not in source
+
+    capture_source = source[source.index("def _capture_unified_decode_graph") :]
+    capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")]
+    assert "_voxcpm2_compile_unified_capture_estimator" in capture_source
+    assert "_voxcpm2_compile_unified_capture_feat_encoder" in capture_source
+    graph_body = capture_source[capture_source.index("with torch.cuda.graph") :]
+    assert "g.cfm_noise.normal_()" not in graph_body
+    assert "g.cfm_noise.normal_()" in unified_source
+
+    assert "_pre_capture_unified_graphs" not in source
+    assert "_unified_graph_pre_capture_sizes" not in source
+    assert "unified_decode_graph_pre_capture_sizes" not in source
+
+
+def test_voxcpm2_batch_unified_graph_requires_runner_metadata_marker():
+    source = TALKER.read_text()
+    compact_source = "".join(source.split())
+
+    assert "enable_runner_assisted_unified_decode_graph" not in source
+    assert "allow_unified_decode_graph_batch_attention" not in source
+    assert "get_runner_assisted_full_attention_metadata_request" in source
+    assert "set_runner_assisted_full_attention_metadata_context" in source
+    assert "RunnerAssistedFullAttentionMetadataRequest" in source
+    assert "_runner_assisted_unified_decode_graph_active" in source
+    assert "runner_full_metadata_missing" in source
+    assert "_build_unified_graph_bucket_sizes" in source
+    assert "_select_unified_graph_bucket_size" in source
+
+    capture_source = source[source.index("def _capture_unified_decode_graph") :]
+    capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")]
+    assert "override_forward_context" in capture_source
+    assert "_nullify_volatile_metadata" in capture_source
+    assert "capture_context = override_forward_context(self._nullify_volatile_metadata(ctx))" in capture_source
+    assert "ifsize>1:continue" not in compact_source
+    forward_source = source[source.index("def _forward_unified_decode") :]
+    forward_source = forward_source[: forward_source.index("# -------------------- vllm hooks")]
+    assert "graph_size = self._select_unified_graph_bucket_size(num_reqs)" in forward_source
+    assert "self._unified_graphs[graph_size]" in forward_source
+    assert "g.input_embeds[num_reqs:graph_size].zero_()" not in forward_source
+    assert "padded rows as valid duplicates" in forward_source
+    assert "g.input_embeds[last : last + 1].expand" in forward_source
+    assert "ifnum_reqs>1andnotself._runner_assisted_unified_decode_graph_active" in compact_source
+    assert "andnotcfg.enable_runner_assisted_unified_decode_graph" not in compact_source
+    needs_source = source[source.index("def get_runner_assisted_full_attention_metadata_request") :]
+    needs_source = needs_source[: needs_source.index("def set_runner_assisted_full_attention_metadata_context")]
+    assert "_select_unified_graph_bucket_size(num_reqs)" in needs_source
+    assert "_should_use_decode_graph(num_reqs)" not in needs_source
+    assert "RunnerAssistedFullAttentionMetadataRequest(" in needs_source
+
+
+def test_voxcpm2_unified_skip_preserves_segmented_decode_graphs():
+    source = TALKER.read_text()
+    compact_source = "".join(source.split())
+    forward_source = source[source.index("def forward(") :]
+    forward_source = forward_source[: forward_source.index("# -------------------- prefill / decode helpers")]
+
+    assert "use_segmented_decode_graph" in forward_source
+    assert "unified_skip_reasonisnotNone" in compact_source
+    assert "notself._enable_unified_decode_graphorunified_skip_reasonisnotNone" in compact_source
+    assert "ifuse_segmented_decode_graph:" in compact_source
+    assert "anduse_segmented_decode_graph" in compact_source
+
+
+def test_voxcpm2_scheduler_policy_stays_model_local():
+    common_source = OMNI_AR_SCHEDULER.read_text()
+    voxcpm2_scheduler_source = VOXCPM2_SCHEDULER.read_text()
+    pipeline_source = VOXCPM2_PIPELINE.read_text()
+
+    assert "VoxCPM2TalkerForConditionalGeneration" not in common_source
+    assert "voxcpm2_runtime_config" not in common_source
+    assert "pure_decode_graph" not in common_source
+    assert "_schedule_with_optional_waiting_deferral" not in common_source
+    assert "def _should_defer_waiting_admission(self) -> bool:" in common_source
+
+    assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source
+    assert "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source
+    assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source
+    assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" in voxcpm2_scheduler_source
+    assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source
+    assert "def schedule(" not in voxcpm2_scheduler_source
+    assert "create_request_queue" not in voxcpm2_scheduler_source
+    assert "original_waiting.prepend_requests(deferred_waiting)" not in voxcpm2_scheduler_source
+    assert "return self._should_defer_waiting_for_unified_decode_graph()" in voxcpm2_scheduler_source
+    assert "vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler" in pipeline_source
+
+
+def test_voxcpm2_deploy_defaults_to_full_unified_graph_only():
+    source = DEPLOY.read_text()
+
+    assert "max_num_seqs: 8" in source
+    assert "enable_unified_decode_graph: true" in source
+    assert "unified_decode_graph_max_batch_size: 8" in source
+    assert "unified_decode_graph_pre_capture_sizes" not in source
+    assert "enable_runner_assisted_unified_decode_graph" not in source
+    assert "allow_unified_decode_graph_batch_attention" not in source
+    assert "decode_graph_capture_policy" not in source

tests/model_executor/models/voxcpm2/test_talker_state_eviction.py

@@ -52,6 +52,10 @@ def _make_bare_talker():
     talker._cfm_buffers = None
     talker._last_audio_output_req_ids = []
     talker._batched_fsq_fusion_max_batch = 32
+    # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to
+    # default (False) to avoid AttributeError when test exercises forward
+    # path that checks this flag.
+    talker._enable_unified_decode_graph = False
     return talker
 
 
@@ -329,23 +333,27 @@ class TestDecodeBatchContract:
 
         for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
             torch.testing.assert_close(actual_res, expected_res)
-            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
+            torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden)
         assert [state.decode_step_count for state in batch_states] == [1, 1]
 
     def test_batch_decode_matches_per_request_state_updates(self) -> None:
+        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+            _DecodeResidualMeta,
+        )
+
         seq_talker = self._make_decode_talker()
         batch_talker = self._make_decode_talker()
 
         seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
-            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 
         for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
-            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0))
 
         batch_talker._finish_decode_batch(
             [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
@@ -360,12 +368,16 @@ class TestDecodeBatchContract:
             torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)
 
     def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
+        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+            _DecodeResidualMeta,
+        )
+
         talker = self._make_decode_talker()
         talker._enable_batched_cfm = False
         states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
-            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 

tests/model_executor/models/voxtral_tts/test_tts_preprocess.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import functools
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@functools.lru_cache(maxsize=1)
+def _voxtral_tts_model_cls():
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+    import vllm.model_executor.models.utils  # noqa: F401
+
+    from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
+        VoxtralTTSForConditionalGeneration,
+    )
+
+    return VoxtralTTSForConditionalGeneration
+
+
+class FakeAudioGeneration(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed_multimodal_calls = []
+
+    def embed_multimodal(self, **kwargs):
+        self.embed_multimodal_calls.append(kwargs)
+        return [torch.full((1, 4), 7.0)]
+
+
+def _make_voxtral_tts_model():
+    model_cls = _voxtral_tts_model_cls()
+    model = model_cls.__new__(model_cls)
+    nn.Module.__init__(model)
+    model.model_stage = "audio_generation"
+    model.model = FakeAudioGeneration()
+    model._audio_token_id = 42
+    return model
+
+
+def test_tts_preprocess_consumes_nested_codes_audio_feedback():
+    model = _make_voxtral_tts_model()
+    input_ids = torch.tensor([model._audio_token_id])
+    input_embeds = torch.zeros((1, 4))
+    audio_tokens = torch.tensor([[1, 2, 3, 4]])
+
+    _, output_embeds, _ = model.tts_preprocess(
+        input_ids=input_ids,
+        input_embeds=input_embeds,
+        codes={"audio": audio_tokens},
+    )
+
+    assert len(model.model.embed_multimodal_calls) == 1
+    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
+    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))
+
+
+def test_tts_preprocess_keeps_legacy_top_level_audio_feedback():
+    model = _make_voxtral_tts_model()
+    input_ids = torch.tensor([model._audio_token_id])
+    input_embeds = torch.zeros((1, 4))
+    audio_tokens = torch.tensor([[5, 6, 7, 8]])
+
+    _, output_embeds, _ = model.tts_preprocess(
+        input_ids=input_ids,
+        input_embeds=input_embeds,
+        audio=audio_tokens,
+    )
+
+    assert len(model.model.embed_multimodal_calls) == 1
+    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
+    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))