[Model] Enable Higgs Audio v3 MRV2 streaming and fix mixed-prefill capture bounds (#8226)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/model_executor/models/higgs_audio_v3/test_async_audio_mode.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_model_metadata_reset(enabled):
+    from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+        HiggsAudioV3TalkerForConditionalGeneration as C,
+    )
+
+    t = C.__new__(C)
+    torch.nn.Module.__init__(t)
+    t.config = SimpleNamespace(audio_async_prompt_mode=enabled)
+    t._set_last_step_query_start_loc = lambda x: None
+    t._sync_decode_state_with_batch = lambda x: None
+    t.update_decode_step_metadata(audio_prompt_mode_rows=2)
+    assert t._step_audio_mode_rows == (2 if enabled else 0)
+    t.update_decode_step_metadata()
+    assert t._step_audio_mode_rows == 0
+
+
+@pytest.mark.parametrize("bad_state", [False, True])
+def test_actual_state_guard_and_eos(bad_state):
+    import importlib.util
+    from pathlib import Path
+
+    import vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker as mod
+
+    root = Path(mod.__file__).parents[4]
+    spec = importlib.util.spec_from_file_location(
+        "helper_fixture", root / "tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py"
+    )
+    assert spec is not None and spec.loader is not None
+    helpers = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(helpers)
+    t = helpers.TestSamplerMethods()._make_batched_sampler_talker(2)
+    t._restore_terminal_audio_rows = (
+        mod.HiggsAudioV3TalkerForConditionalGeneration._restore_terminal_audio_rows.__get__(t)
+    )
+    t.config = SimpleNamespace(audio_async_prompt_mode=True, audio_full_sample_graph=False)
+    t._resolve_token_ids = lambda: None
+    t._audio_continuation_id = 99999
+    t._eos_token_id = 151671
+    t._last_logits_hidden = torch.zeros(2, 16)
+    t._last_step_input_ids = torch.tensor([99999, 12345 if bad_state else 151671])
+    t._last_step_query_start_loc = None
+    t._decode_has_codes = torch.tensor([True, False])
+    t._decode_generation_done = torch.tensor([False, False])
+    t._decode_delay_count = torch.zeros(2, dtype=torch.long)
+    t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long)
+    t._fast_audio_direct_rows = 0
+    t._step_audio_tail_rows = 0
+    t._step_audio_mode_rows = 2
+    t._fast_audio_sampler_gpu_fallback_reason = lambda **kw: None
+    t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026)
+    t._apply_delay_pattern_masking_batched = lambda *a, **kw: None
+    t._sample_audio_codes = lambda logits, *a, **kw: torch.zeros(logits.shape[0], dtype=torch.long)
+    seen_codes = {}
+    t._update_delay_state_batched = lambda *a, **kw: seen_codes.update(kw)
+
+    def call():
+        return mod.HiggsAudioV3TalkerForConditionalGeneration.sample(
+            t, torch.zeros(2, 200000), SimpleNamespace(no_penalties=True)
+        )
+
+    if bad_state:
+        with pytest.raises(RuntimeError, match="tail mismatch"):
+            call()
+    else:
+        assert call().sampled_token_ids.tolist() == [[99999], [151671]]
+        assert seen_codes["code_row_mask"].tolist() == [True, False]

tests/model_executor/models/higgs_audio_v3/test_async_payload.py

@@ -0,0 +1,90 @@
+# 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.higgs_audio_v3.higgs_audio_v3_talker import (
+    HiggsAudioV3TalkerForConditionalGeneration as C,
+)
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_two_inflight_steps_keep_owned_codes_and_masks():
+    t = C.__new__(C)
+    torch.nn.Module.__init__(t)
+    t.num_codebooks = 2
+    t.use_async_omni_output = True
+    first = torch.tensor([[10, 11, 1, 0], [20, 21, 0, 0], [30, 31, 1, 1], [40, 41, 1, 0]])
+    t._finish_audio_staging(first, 4, torch.device("cpu"), 4)
+    first.zero_()  # Next graph replay overwrites original GPU storage.
+    a = t.post_sample_multimodal_outputs(
+        req_ids=["b", "a", "c", "invalid"], invalid_req_indices=[3], multimodal_outputs=None
+    )
+    second = torch.tensor([[50, 51, 1, 0]])
+    t._finish_audio_staging(second, 1, torch.device("cpu"), 1)
+    b = t.post_sample_multimodal_outputs(req_ids=["new"], invalid_req_indices=[], multimodal_outputs=None)
+    second.zero_()
+    assert t.post_sample_multimodal_outputs(req_ids=[], invalid_req_indices=[], multimodal_outputs=None) is None
+    # Finish callbacks out of order. They cannot read current model state.
+    rb = t.finalize_multimodal_outputs_from_cpu_snapshot(b)["codes"]["audio"]
+    ra = t.finalize_multimodal_outputs_from_cpu_snapshot(a)["codes"]["audio"]
+    assert [x.tolist() for x in ra] == [[[10, 11]], [], [[30, 31]], []]
+    assert rb[0].tolist() == [[50, 51]]
+    a["_higgs_audio_snapshot"].zero_()
+    assert ra[0].tolist() == [[10, 11]]
+    assert t._fast_audio_direct_rows == 0
+    assert t.finalize_multimodal_outputs_from_cpu_snapshot(None) is None
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_whole_output_explicit_opt_in(enabled):
+    r = object.__new__(GPUARModelRunner)
+    r.use_async_scheduling = True
+    r.speculative_config = None
+    r.model_config = SimpleNamespace(async_chunk=False, enable_return_routed_experts=False)
+    r.model = SimpleNamespace(use_async_omni_output=True, supports_async_whole_payload=enabled, has_postprocess=False)
+    assert r._should_use_async_omni_output() == enabled
+    r.use_async_scheduling = False
+    assert not r._should_use_async_omni_output()
+
+
+@pytest.mark.parametrize("async_output", [False, True])
+def test_terminal_and_invalid_rows_preserve_full_response_codes(async_output):
+    from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import (
+        _OmniConnectorPayloadTransportMixin,
+    )
+    from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import talker2code2wav_full_payload
+
+    model = C.__new__(C)
+    torch.nn.Module.__init__(model)
+    model.num_codebooks = 8
+    model.use_async_omni_output = async_output
+    model._last_audio_host_staging = None
+    model._last_audio_staging_event = None
+    transport = _OmniConnectorPayloadTransportMixin()
+    transport._pending_full_payload_send = {}
+    transport._full_payload_replace_keys_cached = frozenset()
+    expected = torch.arange(80, dtype=torch.int32).reshape(10, 8)
+    # Terminal and invalid rows must append zero rows, never replace earlier codes.
+    for codes, valid, done, invalid in [
+        *[(row, 1, 0, []) for row in expected],
+        (torch.zeros(8), 0, 1, []),
+        (torch.ones(8), 1, 0, [0]),
+    ]:
+        staging = torch.cat((codes.to(torch.int32), torch.tensor([valid, done], dtype=torch.int32))).reshape(1, 10)
+        model._finish_audio_staging(staging, 1, torch.device("cpu"), 1)
+        payload = model.post_sample_multimodal_outputs(
+            req_ids=["request"], invalid_req_indices=invalid, multimodal_outputs=None
+        )
+        payload = model.finalize_multimodal_outputs_from_cpu_snapshot(payload)
+        audio = payload["codes"]["audio"][0]
+        assert audio.shape == (int(bool(valid) and not invalid), 8)
+        transport.accumulate_full_payload_output("request", {"codes.audio": audio}, None)
+    full, _ = transport._materialize_full_payload_entry(transport._pending_full_payload_send["request"])
+    torch.testing.assert_close(full["codes.audio"], expected)
+    assert talker2code2wav_full_payload(None, full, None)["codes"]["audio"].numel() == 16

tests/model_executor/models/higgs_audio_v3/test_codec_batching.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Codec batching must preserve per-request order and convolution boundaries."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import HiggsAudioV3Code2Wav
+from vllm_omni.transformers_utils.configs.higgs_audio_v3 import HiggsAudioV3Config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def codec(monkeypatch):
+    model = HiggsAudioV3Code2Wav(config=HiggsAudioV3Config())
+    calls = []
+
+    def decode(codes):
+        calls.append(tuple(codes.shape))
+        if (codes < 0).any():
+            raise ValueError("invalid codec row")
+        # A time-dependent decoder makes both trim boundaries observable.
+        pcm = codes[:, :1].float().repeat_interleave(model.hop_length, dim=-1)
+        return pcm
+
+    monkeypatch.setattr(model, "decode_codes", decode)
+    return model, calls
+
+
+def test_ragged_codec_batch_preserves_request_order_and_trims(codec):
+    model, calls = codec
+    codes = [
+        torch.arange(6).expand(8, -1),
+        torch.arange(10, 14).expand(8, -1),
+        torch.arange(20, 26).expand(8, -1),
+    ]
+    output = model(
+        input_ids=torch.cat([row.reshape(-1) for row in codes]),
+        seq_token_counts=[row.numel() for row in codes],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 1, "right_holdback_size": 1}},
+            {"meta": {"left_context_size": 0, "right_holdback_size": 1}},
+            {"meta": {"left_context_size": 1, "right_holdback_size": 1}},
+        ],
+    )
+    assert calls == [(2, 8, 6), (1, 8, 4)]
+    for actual, values in zip(
+        output.multimodal_outputs["model_outputs"], ([1, 2, 3, 4], [10, 11, 12], [21, 22, 23, 24])
+    ):
+        expected = torch.tensor(values, dtype=torch.float32).repeat_interleave(model.hop_length)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("bad_length,bad_value", [(0, 0), (7, 1), (24, -1)])
+def test_invalid_codec_request_does_not_discard_valid_peer(codec, bad_length, bad_value):
+    model, _ = codec
+    valid = torch.arange(3).expand(8, -1).reshape(-1)
+    bad = torch.full((bad_length,), bad_value, dtype=torch.long)
+    output = model(input_ids=torch.cat([bad, valid]), seq_token_counts=[bad.numel(), valid.numel()])
+    audio = output.multimodal_outputs["model_outputs"]
+    assert len(audio) == 2
+    assert audio[0].numel() == 0
+    expected = torch.arange(3, dtype=torch.float32).repeat_interleave(model.hop_length)
+    torch.testing.assert_close(audio[1], expected, rtol=0, atol=0)
+
+
+def test_native_payload_overrides_control_slots_and_empty_terminal(codec):
+    model, calls = codec
+    codes = torch.arange(6).expand(8, -1).reshape(-1)
+    output = model(
+        input_ids=torch.zeros(2, dtype=torch.long),
+        seq_token_counts=[1, 1],
+        runtime_additional_information=[
+            {"codes": {"audio": codes}, "meta": {"left_context_size": 1, "right_holdback_size": 1}},
+            {"codes": {"audio": torch.empty(0, dtype=torch.long)}, "meta": {"finished": True}},
+        ],
+    )
+    assert calls == [(1, 8, 6)]
+    wavs = output.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(wavs[0], torch.arange(1, 5).float().repeat_interleave(model.hop_length))
+    assert wavs[1].numel() == 0

tests/model_executor/models/higgs_audio_v3/test_codec_graph_shapes.py

@@ -0,0 +1,65 @@
+# 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.higgs_audio_v3.full_sample_graph import run_dense_sample
+from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
+    HiggsAudioV3Code2WavForConditionalGeneration as Codec,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_matched_streaming_graphs_cover_every_exact_batch():
+    model = SimpleNamespace(
+        config=SimpleNamespace(
+            codec_graph_single_max_frames=33,
+            codec_graph_batch_sizes=list(range(1, 65)),
+            codec_graph_frame_sizes=[8, 33],
+        )
+    )
+    shapes = Codec._decode_graph_shapes(model)
+    assert len(shapes) == 159
+    assert all((b, f) in shapes for b in range(1, 65) for f in (8, 33))
+    assert all((1, f) in shapes for f in range(1, 34))
+    assert (2, 32) not in shapes
+
+
+def test_default_graph_shapes_are_preserved():
+    shapes = Codec._decode_graph_shapes(SimpleNamespace(config=SimpleNamespace()))
+    assert len(shapes) == 162
+    assert (1, 150) in shapes and (16, 54) in shapes
+
+
+def test_invalid_graph_shape_rejected():
+    with pytest.raises(ValueError):
+        Codec._decode_graph_shapes(SimpleNamespace(config=SimpleNamespace(codec_graph_batch_sizes=[0])))
+
+
+@pytest.mark.parametrize("use_v2", [False, True])
+@pytest.mark.parametrize("sampler", [None, False, True])
+@pytest.mark.parametrize("codec", [None, False, True])
+def test_graph_defaults_follow_runner_and_preserve_explicit_overrides(use_v2, sampler, codec):
+    from vllm_omni.transformers_utils.configs.higgs_audio_v3 import HiggsAudioV3Config
+
+    config = HiggsAudioV3Config(audio_full_sample_graph=sampler, codec_cuda_graph=codec)
+    assert config.resolve_graph_defaults(use_v2_model_runner=use_v2) == (
+        use_v2 if sampler is None else sampler,
+        use_v2 if codec is None else codec,
+    )
+    # Resolving a stage must not bake its defaults into shared model config.
+    assert config.audio_full_sample_graph is sampler
+    assert config.codec_cuda_graph is codec
+
+
+def test_sampler_cache_miss_does_not_capture_or_advance_rng():
+    model = SimpleNamespace(_dense_sample_graphs={})
+    torch.manual_seed(7)
+    before = torch.random.get_rng_state()
+    assert run_dense_sample(model, torch.zeros(2, 4), None, None) is None
+    assert torch.equal(before, torch.random.get_rng_state())
+    assert not model._dense_sample_graphs

tests/model_executor/models/higgs_audio_v3/test_codec_snake_kernel.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The fused DAC Snake / decoder forward is bit-identical to transformers' DacDecoder."""
+
+import copy
+
+import pytest
+import torch
+import torch.nn as nn
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model]
+
+requires_cuda = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA")
+
+_DAC_CONFIG = {
+    "codebook_dim": 8,
+    "codebook_size": 1024,
+    "decoder_hidden_size": 256,
+    "downsampling_ratios": [8, 5, 4, 2, 3],
+    "encoder_hidden_size": 16,
+    "hidden_size": 64,
+    "hop_length": 960,
+    "n_codebooks": 9,
+    "sampling_rate": 16000,
+    "upsampling_ratios": [8, 5, 4, 2, 3],
+}
+
+
+def _dac_decoder_pair(device, dtype):
+    from transformers import DacConfig, DacModel
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import (
+        _fuse_dac_decoder,
+        adjust_conv_transpose_output_padding,
+    )
+
+    torch.manual_seed(0)
+    reference = DacModel(DacConfig(**_DAC_CONFIG)).decoder
+    adjust_conv_transpose_output_padding(reference)
+    with torch.no_grad():
+        for name, param in reference.named_parameters():
+            if name.endswith("alpha"):
+                param.copy_(torch.rand_like(param) * 2 + 0.05)
+            elif name.endswith("bias"):
+                param.normal_(0.0, 0.1)
+    reference = reference.to(device, dtype).eval()
+    fused = copy.deepcopy(reference)
+    assert _fuse_dac_decoder(fused) == 1 + 7 * len(reference.block)
+    fused.load_state_dict(reference.state_dict())
+    return reference, fused
+
+
+@requires_cuda
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("shape", [(1, 64, 33), (3, 1024, 8), (2, 96, 5000)])
+@torch.inference_mode()
+def test_fused_snake_matches_transformers_snake(dtype, shape):
+    from transformers.models.dac.modeling_dac import Snake1d
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d
+
+    torch.manual_seed(0)
+    reference = Snake1d(shape[1]).to("cuda", dtype)
+    reference.alpha.copy_(torch.rand_like(reference.alpha) * 4 + 0.05)
+    fused = _FusedDacSnake1d(shape[1]).to("cuda", dtype)
+    fused.load_state_dict(reference.state_dict())
+    x = (torch.randn(shape, device="cuda") * 3).to(dtype)
+    residual = (torch.randn(shape, device="cuda") * 3).to(dtype)
+    bias = (torch.randn(shape[1], device="cuda") * 3).to(dtype)
+    assert torch.equal(fused(x), reference(x))
+
+    total, out = fused.fused(x, bias, residual, write_sum=True)
+    expected_total = residual + (x + bias.reshape(1, -1, 1))
+    assert torch.equal(total, expected_total)
+    assert torch.equal(out, reference(expected_total))
+    assert fused.fused(x, bias)[0] is None
+    assert torch.equal(fused.fused(x, bias)[1], reference(x + bias.reshape(1, -1, 1)))
+
+
+@requires_cuda
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("batch,frames", [(1, 33), (4, 8)])
+def test_fused_dac_decoder_matches_transformers_decoder(dtype, batch, frames):
+    reference, fused = _dac_decoder_pair("cuda", dtype)
+    x = torch.randn(batch, _DAC_CONFIG["hidden_size"], frames, device="cuda").to(dtype)
+    with torch.inference_mode():
+        assert torch.equal(fused(x), reference(x))
+
+    # Reloaded weights must not reuse the cached Snake reciprocals.
+    with torch.no_grad():
+        for module in (reference, fused):
+            module.snake1.alpha.mul_(0.5)
+    with torch.inference_mode():
+        assert torch.equal(fused(x), reference(x))
+
+
+@pytest.mark.cpu
+def test_fuse_dac_decoder_keeps_state_dict_and_falls_back_on_cpu():
+    from transformers.models.dac.modeling_dac import Snake1d
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d
+
+    reference, fused = _dac_decoder_pair("cpu", torch.float32)
+    assert not any(isinstance(module, Snake1d) for module in fused.modules())
+    assert set(fused.state_dict()) == set(reference.state_dict())
+    x = torch.randn(2, _DAC_CONFIG["hidden_size"], 5)
+    with torch.no_grad():
+        assert torch.equal(fused(x), reference(x))
+
+    decoder = nn.Sequential(Snake1d(8), nn.Conv1d(8, 8, 3))
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _fuse_dac_decoder
+
+    assert _fuse_dac_decoder(decoder) == 1
+    assert isinstance(decoder[0], _FusedDacSnake1d)
+    assert "forward" not in vars(decoder)
+
+
+@requires_cuda
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@torch.inference_mode()
+def test_fused_decoder_graph_replays_changed_inputs(dtype):
+    reference, fused = _dac_decoder_pair("cuda", dtype)
+    x = torch.randn(2, _DAC_CONFIG["hidden_size"], 8, device="cuda").to(dtype)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            fused(x)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        output = fused(x)
+    for _ in range(2):
+        x.normal_()
+        graph.replay()
+        assert torch.equal(output, reference(x))
+
+
+@requires_cuda
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_fused_snake_uses_eager_on_non_nvidia_platform(monkeypatch):
+    from transformers.models.dac.modeling_dac import Snake1d
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d
+
+    reference = Snake1d(8).cuda()
+    fused = _FusedDacSnake1d(8).cuda()
+    fused.load_state_dict(reference.state_dict())
+    x = torch.randn(2, 8, 33, device="cuda")
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    assert not fused.can_fuse(x)
+    assert torch.equal(fused(x), reference(x))
+
+
+@requires_cuda
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_fused_snake_preserves_float64_precision():
+    from transformers.models.dac.modeling_dac import Snake1d
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d
+
+    reference = Snake1d(8).to("cuda", torch.float64)
+    fused = _FusedDacSnake1d(8).to("cuda", torch.float64)
+    fused.load_state_dict(reference.state_dict())
+    x = torch.randn(2, 8, 33, device="cuda", dtype=torch.float64)
+    assert not fused.can_fuse(x)
+    assert torch.equal(fused(x), reference(x))

tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py

@@ -24,6 +24,27 @@ from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v3 import HiggsAudioV
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_external_decode_enters_backbone_compile_boundary():
+    from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+        HiggsAudioV3TalkerForConditionalGeneration,
+    )
+
+    class Backbone(torch.nn.Module):
+        def forward(self, input_ids, positions, inputs_embeds):
+            assert input_ids is None
+            return inputs_embeds + positions[:, None]
+
+        @property
+        def layers(self):
+            pytest.fail("External decode bypassed the Qwen3Model compilation boundary")
+
+    talker = SimpleNamespace(_use_external_decode_cudagraph=True, model=Backbone())
+    hidden = torch.randn(2, 4)
+    positions = torch.tensor([11, 27])
+    actual = HiggsAudioV3TalkerForConditionalGeneration._run_qwen3_layers_eager(talker, positions, hidden)
+    torch.testing.assert_close(actual, hidden + positions[:, None])
+
+
 # ---- AC-1: Configuration ----
 
 
@@ -284,7 +305,7 @@ class TestSamplerMethods:
         """Create a minimal talker-like object with sampler/masking methods."""
         from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
 
-        class FakeTalker:
+        class FakeTalker(SimpleNamespace):
             num_codebooks = 8
             codebook_size = 1026
 
@@ -297,9 +318,10 @@ class TestSamplerMethods:
         """Create a fake talker with GPU-resident sampler state helpers."""
         from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
 
-        class FakeTalker:
+        class FakeTalker(SimpleNamespace):
             num_codebooks = 8
             codebook_size = 1026
+            config = type("Config", (), {"audio_full_sample_graph": False})()
 
         t = FakeTalker()
         t._decode_last_codes = torch.arange(num_rows * t.num_codebooks, dtype=torch.long).view(
@@ -337,6 +359,7 @@ class TestSamplerMethods:
             "_get_audio_host_staging_buffer",
             "_apply_delay_pattern_masking_batched",
             "_update_delay_state_batched",
+            "_finish_audio_staging",
             "_prefill_row_mask",
             "_audio_seed_mask_from_step_input",
         ):
@@ -426,9 +449,12 @@ class TestSamplerMethods:
     def test_sampling_metadata_invalid_temperature_is_sanitized(self, monkeypatch):
         t = self._make_minimal_talker()
         seen = []
-        monkeypatch.setattr(
-            torch, "multinomial", lambda probs, num_samples: seen.append(probs) or probs.argmax(-1, True)
-        )
+
+        def sample(probs, num_samples):
+            seen.append(probs)
+            return probs.argmax(-1, True)
+
+        monkeypatch.setattr(torch, "multinomial", sample)
         logits = torch.tensor([[0.0, 2.0], [3.0, 0.0]])
 
         result = t._sample_audio_codes(
@@ -444,9 +470,12 @@ class TestSamplerMethods:
     def test_sampling_metadata_is_expanded_per_request_and_vq(self, monkeypatch):
         t = self._make_minimal_talker()
         seen = []
-        monkeypatch.setattr(
-            torch, "multinomial", lambda probs, num_samples: seen.append(probs) or probs.argmax(-1, True)
-        )
+
+        def sample(probs, num_samples):
+            seen.append(probs)
+            return probs.argmax(-1, True)
+
+        monkeypatch.setattr(torch, "multinomial", sample)
         cb_logits = torch.tensor([[[4.0, 3.0, 2.0, 1.0]] * 2] * 2)  # [2 requests, 2 VQs, 4 vocab]
         logits = cb_logits.reshape(-1, cb_logits.shape[-1])
         metadata = self._sampling_metadata(
@@ -840,7 +869,7 @@ class TestFeedbackMethods:
         """postprocess() should return codes from _last_audio_codes."""
         from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
 
-        class FakeTalker:
+        class FakeTalker(SimpleNamespace):
             _last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]])
             _last_audio_code_valid = [True]
             _last_audio_host_staging = None
@@ -859,7 +888,7 @@ class TestFeedbackMethods:
         """postprocess() should skip rows with -1 (no audio)."""
         from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
 
-        class FakeTalker:
+        class FakeTalker(SimpleNamespace):
             _last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]])
             _last_audio_code_valid = [False]
             _last_audio_host_staging = None
@@ -876,7 +905,7 @@ class TestFeedbackMethods:
         """postprocess() should advance cursor by 1 per call."""
         from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
 
-        class FakeTalker:
+        class FakeTalker(SimpleNamespace):
             _last_audio_codes = torch.tensor(
                 [[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]]
             )
@@ -901,7 +930,7 @@ class TestReferenceAudioSubstitution:
     def _make_talker(self, query_start_loc):
         from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
 
-        class FakeTalker:
+        class FakeTalker(SimpleNamespace):
             _last_step_query_start_loc = query_start_loc
             _apply_ref_audio_substitution = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_ref_audio_substitution
 
@@ -981,7 +1010,7 @@ class TestAudioFeedback:
         embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
         torch.nn.init.ones_(embed.weight)
 
-        class FakeTalker:
+        class FakeTalker(SimpleNamespace):
             num_codebooks = 8
             codebook_size = 1026
             _audio_continuation_id = 99999  # fake audio token

tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py

@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+    HiggsAudioV3TalkerForConditionalGeneration,
+)
+from vllm_omni.model_executor.models.higgs_audio_v3.model_state import HiggsModelState
+from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def make_state():
+    state = object.__new__(HiggsModelState)
+    state.sampling_slots = None
+    state.request_slots = {}
+    state.device = torch.device("cpu")
+    state.rope_state = state.prompt_embeds_state = state._static_inputs_embeds = None
+    state.requests = {}
+    state.generators = {}
+    state.metadata_key = state.metadata = None
+    state.intermediate_buffer = OmniIntermediateBuffer(4)
+    state.model = SimpleNamespace(
+        on_requests_finished=Mock(),
+        _resolve_token_ids=Mock(),
+        _audio_continuation_id=99,
+        update_decode_step_metadata=Mock(),
+    )
+    return state
+
+
+def request(name, seed, **kwargs):
+    return SimpleNamespace(
+        req_id=name, prompt_token_ids=[1, 2, 99], mm_features=[], sampling_params=SamplingParams(seed=seed, **kwargs)
+    )
+
+
+@pytest.mark.parametrize("slots", [False, True])
+def test_sampling_follows_request_identity_and_reinitializes_reused_slots(slots):
+    s = make_state()
+    if slots:
+        s.sampling_slots = {
+            name: torch.zeros(4, dtype=torch.long if name == "top_k" else torch.float32)
+            for name in ("temperature", "top_p", "top_k", "zeros", "ones")
+        }
+        s.sampling_slots["ones"].fill_(1)
+    s.add_request(0, request("a", 42, temperature=0.5))
+    s.add_request(2, request("b", 99, temperature=0.9))
+    a = s.generators["a"]
+    b = s.generators["b"]
+    first = s.sampling_metadata(["a", "b"])
+    assert first.generators == {0: a, 1: b}
+    torch.testing.assert_close(first.temperature, torch.tensor([0.5, 0.9]))
+    swapped = s.sampling_metadata(["b", "a"])
+    assert swapped.generators == {0: b, 1: a}
+    torch.testing.assert_close(swapped.temperature, torch.tensor([0.9, 0.5]))
+    s.remove_request("a")
+    s.add_request(0, request("new", 7))
+    assert "a" not in s.generators and "a" not in s.requests
+    s.model.on_requests_finished.assert_called_once_with({"a"})
+    assert s.sampling_metadata(["new"]).generators[0].initial_seed() == 7
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"repetition_penalty": 1.2},
+        {"logprobs": 1},
+        {"min_p": 0.1},
+        {"frequency_penalty": 0.2},
+        {"prompt_logprobs": 1},
+        {"min_tokens": 1},
+    ],
+)
+def test_unsupported_sampling_is_rejected_before_admission(kwargs):
+    s = make_state()
+    with pytest.raises(ValueError, match="temperature/top-k/top-p"):
+        s.add_request(0, request("a", 42, **kwargs))
+    assert not s.requests
+
+
+def test_partial_prefill_not_classified_as_audio_and_slot_order_is_authoritative():
+    s = make_state()
+    s.add_request(0, request("a", 42))
+    s.add_request(2, request("b", 99))
+    batch = SimpleNamespace(
+        req_ids=["b", "a"],
+        idx_mapping_np=np.array([2, 0]),
+        num_scheduled_tokens=np.array([1, 1]),
+        positions=torch.tensor([2, 1]),
+        query_start_loc=torch.tensor([0, 1, 2]),
+        num_reqs=2,
+        num_tokens=2,
+    )
+    states = SimpleNamespace(num_computed_prefill_tokens=np.array([1, 0, 2, 0]))
+    s.run_preprocess(batch, {"input_ids": torch.tensor([99, 2])}, states)
+    assert s.model.update_decode_step_metadata.call_args.kwargs["audio_prompt_mode_rows"] == 0
+    states.num_computed_prefill_tokens[0] = 2
+    s.run_preprocess(batch, {"input_ids": torch.tensor([99, 99])}, states)
+    assert s.model.update_decode_step_metadata.call_args.kwargs["audio_prompt_mode_rows"] == 2
+
+
+def test_finalizer_filters_unsampled_prefill_and_keeps_owned_audio():
+    s = make_state()
+    s.model = object.__new__(HiggsAudioV3TalkerForConditionalGeneration)
+    torch.nn.Module.__init__(s.model)
+    s.model.num_codebooks = 2
+    snapshot = torch.tensor([[11, 12, 1, 0], [21, 22, 1, 0], [31, 32, 0, 1]])
+    payload = {"_higgs_audio_snapshot": snapshot, "_higgs_invalid_rows": ()}
+    result = s.finalize_audio_snapshot(payload, [1, 0, 1])
+    snapshot.zero_()
+    assert [x.tolist() for x in result["codes"]["audio"]] == [[[11, 12]], [], []]
+    assert payload["_higgs_invalid_rows"] == ()
+
+
+def test_upstream_warmup_retains_generic_text_sampling():
+    s = make_state()
+    s.add_request(0, request("_warmup_0_", 0, repetition_penalty=1.2, logprobs=1))
+    batch = SimpleNamespace(req_ids=["_warmup_0_"])
+    sampler, rejection = s.custom_sampler(Mock())
+    expected = (object(), torch.ones(1), torch.zeros(1))
+    standard = Mock(return_value=expected)
+    output = sampler.sample_step(None, batch, None, None, standard)
+    assert output.sampler_output is expected[0] and output.multimodal_outputs is None
+    assert rejection is None
+    standard.assert_called_once_with(None, batch, None)
+    s.remove_request("_warmup_0_")
+    assert not s.requests
+
+
+def test_static_embeddings_exclude_padding_and_preserve_storage():
+    s = make_state()
+    embed = torch.nn.Embedding(100, 3)
+    s.model.model = SimpleNamespace(embed_tokens=embed)
+    s.model._apply_audio_feedback = Mock(side_effect=lambda x, ids: x + 2)
+    s.model._apply_ref_audio_substitution = Mock(side_effect=lambda x, *args: x + 5)
+    s._static_inputs_embeds = torch.full((8, 3), 999.0)
+    ptr = s._static_inputs_embeds.data_ptr()
+    ids = torch.tensor([1, 2, 3, 0])
+    batch = SimpleNamespace(num_tokens=3, num_reqs=3, num_tokens_after_padding=4, has_prefill=False)
+    s._prepare_static_embeddings(batch, {"input_ids": ids})
+    assert s._static_inputs_embeds.data_ptr() == ptr
+    torch.testing.assert_close(s._static_inputs_embeds[:3], embed(ids[:3]) + 2)
+    assert s._static_inputs_embeds[3].count_nonzero() == 0
+    assert s.model._apply_audio_feedback.call_args.args[1].numel() == 3
+    s.model._apply_ref_audio_substitution.assert_not_called()
+    batch.has_prefill = True
+    batch.positions = torch.arange(4)
+    s._prepare_static_embeddings(batch, {"input_ids": ids, "model_intermediate_buffer": [{}]})
+    torch.testing.assert_close(s._static_inputs_embeds[:3], embed(ids[:3]) + 7)
+
+
+def test_mixed_feedback_updates_only_single_token_spans_with_codes():
+    s = make_state()
+    s.model._ensure_decode_state_capacity = Mock()
+    s.model._decode_last_codes = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+    s.model._decode_has_codes = torch.tensor([True, True, False, True])
+    s.model.multimodal_embedding = lambda codes: codes.sum(1)[:, None].expand(-1, 3)
+    batch = SimpleNamespace(num_tokens=7, num_reqs=4, query_start_loc=torch.tensor([0, 1, 4, 5, 7]))
+    embeds = torch.arange(21).reshape(7, 3).float()
+    output = s._apply_audio_feedback(embeds, torch.ones(7).long(), batch)
+    expected = embeds.clone()
+    expected[0] = 3
+    torch.testing.assert_close(output, expected)
+    assert embeds[0].tolist() == [0, 1, 2]
+
+
+def test_direct_payload_owns_rows_and_filters_partial_prefill():
+    s = make_state()
+    s.direct_payload = True
+    s.model.num_codebooks = 2
+    snapshot = torch.tensor([[11, 12, 1, 0], [21, 22, 1, 0], [31, 32, 0, 1]])
+    out = s.finalize_audio_snapshot({"_higgs_audio_snapshot": snapshot}, [1, 0, 1])
+    snapshot.zero_()
+    assert out.inter_stage[0]["codes.audio"].tolist() == [[11, 12]]
+    assert out.inter_stage[1:] == [None, None]
+
+
+def test_registered_audio_sampler_preserves_rows_counts_and_owned_output(mocker):
+    state = make_state()
+    metadata, payload = object(), {"_higgs_audio_snapshot": torch.ones(2, 4)}
+    state.sampling_metadata = Mock(return_value=metadata)
+    state.model.compute_logits = Mock(return_value=torch.ones(2, 8))
+    state.model.sample = Mock(
+        return_value=SimpleNamespace(sampled_token_ids=torch.tensor([[9], [8]]), logprobs_tensors=None)
+    )
+    state.model.post_sample_multimodal_outputs = Mock(return_value=payload)
+    count, rejected = torch.tensor([1, 0]), torch.tensor([0, 0])
+    mocker.patch(
+        "vllm_omni.model_executor.models.higgs_audio_v3.model_state.get_num_sampled_and_rejected",
+        return_value=(count, rejected),
+    )
+    batch = SimpleNamespace(
+        req_ids=["b", "a"],
+        num_reqs=2,
+        num_draft_tokens=0,
+        idx_mapping=torch.tensor([3, 1]),
+        logits_indices=torch.tensor([2, 0]),
+        seq_lens=None,
+        cu_num_logits=None,
+    )
+    hidden = torch.arange(12).reshape(3, 4).float()
+    standard = Mock(side_effect=AssertionError("must not sample twice"))
+    sampler, rejection = state.custom_sampler(Mock())
+    result = sampler.sample_step(hidden, batch, SimpleNamespace(prefill_len=SimpleNamespace(gpu=None)), None, standard)
+    state.sampling_metadata.assert_called_once()
+    torch.testing.assert_close(state.model.compute_logits.call_args.args[0], hidden[[2, 0]])
+    assert state.model.compute_logits.call_args.args[1] is metadata
+    assert result.num_sampled is count and result.num_rejected is rejected
+    assert result.multimodal_outputs is payload and result.owns_multimodal_outputs
+    assert not result.include_hidden_states and result.finalize_multimodal == state.finalize_audio_snapshot
+    assert rejection is None
+    standard.assert_not_called()
+
+
+@pytest.mark.parametrize("draft,grammar", [(1, None), (0, object())])
+def test_registered_audio_sampler_rejects_unsupported_modes(draft, grammar):
+    sampler, _ = make_state().custom_sampler(Mock())
+    batch = SimpleNamespace(req_ids=["real-request"], num_draft_tokens=draft)
+    with pytest.raises(ValueError, match="grammar or speculative"):
+        sampler.sample_step(None, batch, None, grammar, Mock())

tests/model_executor/models/higgs_audio_v3/test_native_payload.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+    talker2code2wav,
+    talker2code2wav_full_payload,
+    talker2code2wav_token_only,
+)
+from vllm_omni.outputs.mm_outputs import MultimodalPayload
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("wrapped", [False, True])
+@pytest.mark.parametrize("rows", [0, 4, 8, 9, 30])
+@pytest.mark.parametrize("offset", [0, -16, 1016])
+def test_native_full_payload_preserves_legacy_dedelay_and_tail(rows, offset, wrapped):
+    audio = torch.arange(rows * 8).reshape(rows, 8) + offset
+    payload = {"codes": {"audio": audio}}
+    if wrapped:
+        payload = MultimodalPayload(metadata=payload)
+    out = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=payload)])
+    legacy = talker2code2wav([out])[0]["prompt_token_ids"]
+    native = talker2code2wav_full_payload(None, {"codes.audio": audio}, None)
+    assert native["codes"]["audio"].tolist() == legacy
+    # Independent expected layout: book q reads delayed rows q + frame.
+    frames = max(rows - 7, 0) if rows >= 8 else 0
+    kept = frames - 1 if frames >= 2 else frames
+    expected = []
+    for q in range(8):
+        for frame in range(kept):
+            value = (q + frame) * 8 + q + offset
+            expected.append(value if 0 <= value < 1024 else 0)
+    assert legacy == expected
+    assert native["meta"]["finished"]
+    assert talker2code2wav_token_only([out])[0]["prompt_token_ids"] == legacy
+
+
+def test_native_control_slot_without_audio_payload():
+    out = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=None)])
+    assert talker2code2wav_token_only([out])[0]["prompt_token_ids"] == [0]

tests/model_executor/models/higgs_audio_v3/test_sampling_kernels.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused Higgs sampling kernels must be bit-identical to the eager code they replace."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA")]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("numel", [8 * 1026, 1000, 300000])
+def test_batched_generator_noise_matches_per_request_exponential(numel):
+    from vllm_omni.utils.seeded_exponential import fill_exponential_rows
+
+    device = torch.device("cuda")
+    seeds = [42, 42, None, 2**63 + 5]
+
+    def generators():
+        gens = [None if s is None else torch.Generator(device=device).manual_seed(s) for s in seeds]
+        resumed = gens[1]
+        assert resumed is not None
+        resumed.set_offset(resumed.get_offset() + 4)  # a request that already sampled
+        return gens
+
+    default = torch.cuda.default_generators[device.index or 0]
+    start = default.get_offset()
+    expected_gens = generators()
+    expected = torch.empty(len(seeds), numel, device=device)
+    for row, generator in enumerate(expected_gens):
+        expected[row].exponential_(generator=generator)
+    expected_default = default.get_offset()
+
+    default.set_offset(start)
+    actual_gens = generators()
+    actual = fill_exponential_rows(torch.empty_like(expected), actual_gens)
+    assert torch.equal(actual, expected)
+    assert default.get_offset() == expected_default
+    for got, want in zip(actual_gens, expected_gens):
+        if want is not None:
+            assert got.get_offset() == want.get_offset()
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_fused_decode_embeddings_match_feedback_and_zero_padding():
+    from vllm_omni.model_executor.models.higgs_audio_v3.step_kernels import decode_embeddings
+
+    torch.manual_seed(0)
+    vocab, hidden, books, audio_vocab = 1000, 96, 8, 17
+    text = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16)
+    audio = torch.randn(books * audio_vocab, hidden, device="cuda", dtype=torch.bfloat16)
+    ids = torch.randint(0, vocab, (5,), device="cuda")
+    ids[0] = -1
+    codes = torch.randint(0, audio_vocab, (8, books), device="cuda")
+    has = torch.tensor([True, False, True, False, True, True, False, False], device="cuda")
+    offsets = torch.arange(books, device="cuda") * audio_vocab
+    expected = torch.where(
+        has[:5, None], F.embedding(codes[:5] + offsets, audio).sum(-2), F.embedding(ids.clamp_min(0), text)
+    )
+    out = torch.full((8, hidden), 3.0, device="cuda", dtype=torch.bfloat16)
+    decode_embeddings(out, ids, text, audio, codes, has, 7, audio_vocab)
+    assert torch.equal(out[:5], expected)
+    assert torch.count_nonzero(out[5:7]) == 0
+    assert torch.all(out[7] == 3)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("prompt_mode", [False, True])
+def test_step_tail_guard_matches_eager_restore_reset_and_validation(prompt_mode):
+    from vllm_omni.model_executor.models.higgs_audio_v3.step_kernels import step_tail_guard
+
+    audio, eos = 7, 9
+    # Rows: decode audio, decode EOS, 3-token prefill ending in audio, decode audio.
+    ids = torch.tensor([audio, eos, 1, 2, audio, audio], device="cuda")
+    qsl = torch.tensor([0, 1, 2, 5, 6], device="cuda", dtype=torch.int32)
+    has = torch.tensor([True, True, True, False], device="cuda")
+    done = torch.tensor([False, False, True, False], device="cuda")
+    delay = torch.tensor([3, 4, 5, 6], device="cuda")
+    eoc = torch.tensor([1, 2, 3, 4], device="cuda")
+    error = torch.zeros(1, device="cuda", dtype=torch.int32)
+    step_tail_guard(ids, qsl, has, done, delay, eoc, error, audio_id=audio, eos_id=eos, prompt_mode=prompt_mode)
+    prefill = torch.tensor([False, False, True, False], device="cuda")
+    tail = ids[(qsl[1:] - 1).long()]
+    want_done = torch.tensor([False, False, True, False], device="cuda")
+    if prompt_mode:
+        want_done |= tail == eos
+    want_done &= ~prefill
+    assert torch.equal(done, want_done)
+    assert torch.equal(has, torch.tensor([True, True, False, False], device="cuda"))
+    assert delay.tolist() == [3, 4, 0, 6] and eoc.tolist() == [1, 2, -1, 4]
+    # The EOS row is valid only as a terminal prompt-mode row.
+    assert int(error.item()) == (0 if prompt_mode else 1)

tests/model_executor/models/higgs_audio_v3/test_text_head_skip.py

@@ -0,0 +1,41 @@
+# 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.higgs_audio_v3.higgs_audio_v3_talker import (
+    HiggsAudioV3TalkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("fallback", [None, "prefill", "logprobs", "allowed", "bad_words", "penalty", "inactive"])
+def test_text_head_only_skipped_for_direct_audio_decode(fallback):
+    hidden = torch.randn(3, 4)
+    calls = []
+
+    def logits_processor(head, h):
+        calls.append(h)
+        return torch.ones(3, 7)
+
+    model = SimpleNamespace(
+        _audio_continuation_id=5,
+        _last_step_input_ids=torch.ones(4 if fallback == "prefill" else 3, dtype=torch.long),
+        _fast_audio_direct_rows=0 if fallback == "inactive" else 3,
+        _text_vocab_size=7,
+        lm_head=None,
+        logits_processor=logits_processor,
+    )
+    metadata = SimpleNamespace(
+        no_penalties=fallback != "penalty",
+        max_num_logprobs=1 if fallback == "logprobs" else None,
+        allowed_token_ids_mask=torch.ones(3, 7) if fallback == "allowed" else None,
+        bad_words_token_ids={0: [1]} if fallback == "bad_words" else {},
+    )
+    logits = HiggsAudioV3TalkerForConditionalGeneration.compute_logits(model, hidden, metadata)
+    assert logits.shape == (3, 7)
+    assert model._last_logits_hidden is hidden
+    assert len(calls) == (fallback is not None)

tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py

@@ -215,6 +215,8 @@ class _MiniDiT(nn.Module):
             cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous()
             dt = x.shape[1]
             att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1)
+            if att_b is not None:
+                att_cache_buffer[b_idx][:, :, dt:, :] = att_b
         x = self.final_layer(x)
         x = x.transpose(1, 2)
         return x
@@ -248,6 +250,9 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
     with torch.inference_mode():
         for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)):
             inputs = _cfm_inputs(2, chunk_size, old_att_len)
+            # A cache output must overwrite every row, including the old tail.
+            inputs[4].fill_(float("nan"))
+            inputs[5].fill_(float("nan"))
 
             eager_inputs = tuple(v.clone() for v in inputs)
             with torch.no_grad():
@@ -268,6 +273,8 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
             torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
             torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
             torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
+            assert torch.isfinite(graph_att).all()
+            torch.testing.assert_close(graph_att[:, :, :, chunk_size:, :], inputs[3])
 
     wrapper._flush()