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

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

.buildkite/cuda/test-ready.yml

@@ -205,6 +205,16 @@ steps:
           - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
         mirror_hardwares: h100_1
 
+      - label: "TTS · Higgs v3 Full Response Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v3.py
+          - vllm_omni/model_executor/models/higgs_audio_v3/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
+          - vllm_omni/deploy/higgs_multimodal_qwen3.yaml
+        commands:
+          - timeout 20m pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py::test_full_response_v1_preserves_audio -m 'core_model and cuda' --run-level 'core_model'
+        mirror_hardwares: h100_1
+
       - label: "TTS · Qwen3-TTS CustomVoice Test"
         source_file_dependencies: tts_qwen3_tts_function
         commands:

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -315,3 +315,34 @@ class TestHiggsAudioV3OnlineVoiceClone:
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
             }
         )
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize(
+    "omni_server",
+    [
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=STAGE_CONFIG,
+            server_args=[*SERVER_ARGS, "--no-async-chunk"],
+            env_dict=SERVER_ENV,
+        )
+    ],
+    indirect=True,
+)
+def test_full_response_v1_preserves_audio(omni_server, online_client) -> None:
+    """Non-chunk V1 must retain codec rows through the terminal talker step."""
+    online_client.send_audio_speech_request(
+        {
+            "model": omni_server.model,
+            "input": "The quick brown fox jumps over the lazy dog.",
+            "stream": False,
+            "response_format": "wav",
+            "min_audio_bytes": _MIN_AUDIO_BYTES,
+            "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+        },
+        request_num=3,
+    )

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()
 

tests/model_executor/stage_input_processors/test_higgs_audio_v3_async_chunk_batch.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The batched Higgs async-chunk builder must emit exactly what the per-request one does."""
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+    talker2code2wav_async_chunk,
+    talker2code2wav_async_chunk_batch,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _manager():
+    extra = {
+        "codec_chunk_frames": 5,
+        "codec_left_context_frames": 3,
+        "codec_right_holdback_frames": 1,
+        "initial_codec_chunk_frames": 2,
+    }
+    return SimpleNamespace(connector=SimpleNamespace(config={"extra": extra}), code_prompt_token_ids=defaultdict(list))
+
+
+class _Request:
+    def __init__(self, rid, finish_step):
+        self.external_req_id = rid
+        self.finish_step = finish_step
+        self.step = 0
+
+    def is_finished(self):
+        return self.step >= self.finish_step
+
+
+def _payload_items(payload):
+    if payload is None:
+        return None
+    meta = payload.meta
+    return (
+        payload.codes.audio.tolist(),
+        None if meta is None else (meta.left_context_size, meta.right_holdback_size, bool(meta.finished)),
+    )
+
+
+def test_batch_async_chunk_matches_per_request_builder():
+    torch.manual_seed(0)
+    single, batched = _manager(), _manager()
+    finish = [23, 9, 31]
+    requests_single = [_Request(f"r{i}", f) for i, f in enumerate(finish)]
+    requests_batch = [_Request(f"r{i}", f) for i, f in enumerate(finish)]
+    for step in range(32):
+        outputs: list[dict | None] = []
+        flags: list[bool] = []
+        for i, request in enumerate(requests_single):
+            if step > request.finish_step:
+                outputs.append(None)
+            elif step == 7 and i == 2:
+                outputs.append(None)  # a step without an emitted row
+            else:
+                outputs.append({"codes": {"audio": torch.randint(0, 1026, (1, 8))}})
+            flags.append(step == request.finish_step)
+        for request in requests_single + requests_batch:
+            request.step = step
+        active = [i for i, request in enumerate(requests_single) if step <= request.finish_step]
+        want = [
+            _payload_items(talker2code2wav_async_chunk(single, outputs[i], requests_single[i], flags[i]))
+            for i in active
+        ]
+        got = talker2code2wav_async_chunk_batch(
+            batched,
+            [outputs[i] for i in active],
+            [requests_batch[i] for i in active],
+            [flags[i] for i in active],
+        )
+        assert [_payload_items(p) for p in got] == want
+    assert single.higgs_v3_emitted_frames == batched.higgs_v3_emitted_frames == {}

tests/worker/test_higgs_forward_metadata_scope.py

@@ -0,0 +1,96 @@
+# 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.v1.worker.gpu_model_runner import GPUModelRunner
+
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("fail", [False, True])
+def test_runner_closes_metadata_scope_on_return_and_error(monkeypatch, fail):
+    events = []
+
+    def begin(**kw):
+        events.append("begin")
+
+    def end():
+        events.append("end")
+
+    def forward(*args, **kw):
+        assert events == ["begin"]
+        events.extend(["warmup", "capture"])
+        if fail:
+            raise RuntimeError("capture failed")
+        return torch.ones(2)
+
+    r = object.__new__(OmniGPUModelRunner)
+    r.model = SimpleNamespace(
+        supports_omni_decode_step_metadata=True, update_decode_step_metadata=begin, finish_decode_step_forward=end
+    )
+    r.input_batch = SimpleNamespace(req_ids=["a", "b"])
+    r._build_model_kwargs_extra = lambda: {}
+    monkeypatch.setattr(GPUModelRunner, "_model_forward", forward)
+    if fail:
+        with pytest.raises(RuntimeError, match="capture failed"):
+            r._model_forward()
+    else:
+        torch.testing.assert_close(r._model_forward(), torch.ones(2))
+    assert events == ["begin", "warmup", "capture", "end"]
+
+
+@pytest.mark.parametrize("requires_tails", [False, True])
+@pytest.mark.parametrize("async_scheduling", [False, True])
+def test_cpu_tail_metadata_requires_model_capability(monkeypatch, requires_tails, async_scheduling):
+    observed = {}
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(
+        supports_omni_decode_step_metadata=True,
+        requires_cpu_input_tail_ids=requires_tails,
+        update_decode_step_metadata=lambda **kwargs: observed.update(kwargs),
+    )
+    ids = torch.tensor([10, 11, 12])
+    runner.use_async_scheduling = async_scheduling
+    runner.input_ids = SimpleNamespace(gpu=ids, cpu=ids)
+    runner.query_start_loc = SimpleNamespace(cpu=torch.tensor([0, 2, 3]))
+    runner.input_batch = SimpleNamespace(req_ids=["a", "b"])
+    runner._build_model_kwargs_extra = lambda: {}
+    monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: torch.ones(2))
+    runner._model_forward(input_ids=ids)
+    expected = [11, 12] if requires_tails and not async_scheduling else None
+    assert observed["cpu_input_tail_ids"] == expected
+
+
+@pytest.mark.parametrize("generation", [False, True])
+@pytest.mark.parametrize("enabled,fail", [(False, False), (True, False), (True, True)])
+def test_auxiliary_capture_finishes_before_worker_readiness(monkeypatch, generation, enabled, fail):
+    from vllm.v1.worker.gpu_worker import Worker
+
+    from vllm_omni.worker.base import OmniGPUWorkerBase
+    from vllm_omni.worker.gpu_generation_worker import GPUGenerationWorker
+
+    events: list[str] = []
+
+    def capture():
+        assert events == ["warmup"]
+        events.append("capture")
+        if fail:
+            raise RuntimeError("auxiliary capture failed")
+
+    monkeypatch.setattr(Worker, "compile_or_warm_up_model", lambda _: events.append("warmup"))
+    worker = object.__new__(GPUGenerationWorker if generation else OmniGPUWorkerBase)
+    worker.use_v2_model_runner = generation
+    model = SimpleNamespace(capture_auxiliary_graphs=capture) if enabled else SimpleNamespace()
+    worker.model_runner = SimpleNamespace(model=model, profile_run=lambda: events.append("warmup"))
+    if fail:
+        with pytest.raises(RuntimeError, match="auxiliary capture failed"):
+            worker.compile_or_warm_up_model()
+    else:
+        worker.compile_or_warm_up_model()
+        events.append("ready")
+    assert events == ["warmup"] + (["capture"] if enabled else []) + ([] if fail else ["ready"])

tests/worker_v2/test_capture_query_bound_gpu.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.utils import AttentionGroup
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_full_graph_prefill_uses_real_flash_attention_backend():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("requires FA3 on Hopper or newer")
+    _check_prefill()
+
+
+def _check_prefill():
+    torch.manual_seed(17)
+    device = torch.device("cuda:0")
+    c = VllmConfig()
+    c.model_config = SimpleNamespace(
+        get_num_attention_heads=lambda p: 32,
+        rswa_window=None,
+        is_diffusion=False,
+        is_mm_prefix_lm=False,
+        hf_config=SimpleNamespace(model_type="qwen3"),
+    )
+    c.scheduler_config.max_num_seqs = 64
+    c.compilation_config.cudagraph_mode = CUDAGraphMode.FULL
+    c.compilation_config.max_cudagraph_capture_size = 256
+    c.attention_config.flash_attn_version = 3
+    spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16)
+    kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
+    state = OmniModelState.__new__(OmniModelState)
+    state.max_model_len = 256
+    state.supports_mm_inputs = False
+    with set_current_vllm_config(c):
+        group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0)
+        group.create_metadata_builders(c, device)
+        impl = FlashAttentionImpl(32, 128, 128**-0.5, 8, None, None, "auto")
+        assert impl.vllm_flash_attn_version == 3
+        buffers = InputBuffers(64, 256, device)
+        batch = InputBatch.make_dummy(64, 256, buffers)
+        tables = torch.arange(1024, dtype=torch.int32, device=device).reshape(64, 16)
+        slots = torch.zeros((1, 256), dtype=torch.int64, device=device)
+        cache = torch.randn(1024, 8, 16, 256, device=device, dtype=torch.bfloat16)
+        q = torch.randn(256, 32, 128, device=device, dtype=torch.bfloat16)
+        k = v = torch.empty(256, 8, 128, device=device, dtype=torch.bfloat16)
+        layer = SimpleNamespace(
+            _q_scale=torch.ones((), device=device),
+            _k_scale=torch.ones((), device=device),
+            _v_scale=torch.ones((), device=device),
+        )
+        batch.query_start_loc.copy_(torch.arange(65, device=device, dtype=torch.int32) * 4)
+        batch.seq_lens.fill_(4)
+        prepare = OmniModelState.prepare_attn
+        md = prepare(state, batch, CUDAGraphMode.NONE, (tables,), slots, [[group]], kvconfig, for_capture=True)["attn"]
+        out = torch.full_like(q, float("nan"))
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                impl.forward(layer, q, k, v, cache, md, out)
+        stream.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            impl.forward(layer, q, k, v, cache, md, out)
+        # Same captured 64-request shape: one 212-token prompt, 44 single-token
+        # requests, then zero-length padded rows. No backend is mocked.
+        lengths = np.array([212] + [1] * 44 + [0] * 19, dtype=np.int32)
+        batch.query_start_loc.copy_(torch.tensor(np.r_[0, np.cumsum(lengths)], device=device, dtype=torch.int32))
+        batch.seq_lens.copy_(torch.tensor(lengths, device=device))
+        from dataclasses import replace
+
+        runtime = replace(
+            batch,
+            max_query_len=212,
+            num_scheduled_tokens=lengths,
+            query_start_loc_np=np.r_[0, np.cumsum(lengths)],
+            seq_lens_cpu_upper_bound=torch.tensor(lengths),
+        )
+        eager_md = state.prepare_attn(runtime, CUDAGraphMode.NONE, (tables,), slots, [[group]], kvconfig)["attn"]
+        expected = torch.empty_like(q)
+        impl.forward(layer, q, k, v, cache, eager_md, expected)
+        state.prepare_attn(runtime, CUDAGraphMode.NONE, (tables,), slots, [[group]], kvconfig)
+        out.fill_(float("nan"))
+        graph.replay()
+        torch.accelerator.synchronize()
+        assert md.max_query_len == 256
+        torch.testing.assert_close(out[:212], expected[:212], atol=0.002, rtol=0.002)

vllm_omni/deploy/higgs_multimodal_qwen3_mrv2_h200.yaml

@@ -0,0 +1,169 @@
+# Higgs v3 Model Runner V2 streaming profile, validated on one H200.
+# Select an available physical GPU with CUDA_VISIBLE_DEVICES before launch.
+# First audio: 8 frames (320 ms); following chunks: 25 frames (1 s).
+# Codec graph windows are exact: 8 first frames, 25 + 8 left context.
+# AR and codec capacities are 64; tested client concurrencies are 32 and 64.
+# Keep model weights and sampling parameters identical for comparisons.
+# Server flags used in measurements: --disable-log-stats --disable-uvicorn-access-log.
+# Auxiliary graphs are captured after KV sizing and are outside the utilization
+# budget. Leave at least 10 GiB free for codec graph pools on this H200 profile;
+# inspect both allocated and reserved capture logs before raising stage budgets.
+# Sampler graphs cover tensor-valued temperature/top_k/top_p. Disabling top-k
+# or top-p for an entire batch (None metadata) uses eager sampling.
+async_chunk: true
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      codec_streaming: true
+      codec_chunk_frames: 25
+      codec_left_context_frames: 8
+      codec_right_holdback_frames: 0
+      initial_codec_chunk_frames: 8
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 5000
+      connector_get_max_wait: 500
+stages:
+- stage_id: 0
+  max_num_seqs: 64
+  gpu_memory_utilization: 0.6
+  enforce_eager: false
+  attention_backend: FLASH_ATTN
+  trust_remote_code: true
+  enable_prefix_caching: false
+  async_scheduling: true
+  compilation_config:
+    cudagraph_capture_sizes:
+    - 1
+    - 2
+    - 4
+    - 8
+    - 16
+    - 24
+    - 32
+    - 40
+    - 48
+    - 56
+    - 64
+    - 96
+    - 128
+    - 192
+    - 256
+    - 384
+    - 512
+    cudagraph_mode: FULL
+    cudagraph_num_of_warmups: 1
+    max_cudagraph_capture_size: 512
+    cudagraph_copy_inputs: true
+  max_num_batched_tokens: 4096
+  max_model_len: 8192
+  devices: '0'
+  output_connectors:
+    to_stage_1: connector_of_shared_memory
+  default_sampling_params:
+    temperature: 1.0
+    top_p: 0.95
+    top_k: 50
+    max_tokens: 2048
+    seed: 42
+    repetition_penalty: 1.0
+  hf_overrides:
+    audio_async_prompt_mode: true
+    audio_async_payload: true
+    audio_mrv2_static_inputs: true
+    audio_mrv2_slot_parameters: true
+    audio_mrv2_direct_payload: true
+    audio_mrv2_sample_graph_buckets: true
+    audio_mrv2_async_state_indices: true
+    audio_mrv2_fused_state: true
+- stage_id: 1
+  max_num_seqs: 64
+  gpu_memory_utilization: 0.25
+  enforce_eager: true
+  trust_remote_code: true
+  enable_prefix_caching: false
+  async_scheduling: false
+  max_num_batched_tokens: 65536
+  max_model_len: 65536
+  devices: '0'
+  input_connectors:
+    from_stage_0: connector_of_shared_memory
+  default_sampling_params:
+    temperature: 0.0
+    top_p: 1.0
+    top_k: -1
+    max_tokens: 65536
+    seed: 42
+    repetition_penalty: 1.0
+  hf_overrides:
+    codec_graph_single_max_frames: 33
+    codec_graph_batch_sizes:
+    - 1
+    - 2
+    - 3
+    - 4
+    - 5
+    - 6
+    - 7
+    - 8
+    - 9
+    - 10
+    - 11
+    - 12
+    - 13
+    - 14
+    - 15
+    - 16
+    - 17
+    - 18
+    - 19
+    - 20
+    - 21
+    - 22
+    - 23
+    - 24
+    - 25
+    - 26
+    - 27
+    - 28
+    - 29
+    - 30
+    - 31
+    - 32
+    - 33
+    - 34
+    - 35
+    - 36
+    - 37
+    - 38
+    - 39
+    - 40
+    - 41
+    - 42
+    - 43
+    - 44
+    - 45
+    - 46
+    - 47
+    - 48
+    - 49
+    - 50
+    - 51
+    - 52
+    - 53
+    - 54
+    - 55
+    - 56
+    - 57
+    - 58
+    - 59
+    - 60
+    - 61
+    - 62
+    - 63
+    - 64
+    codec_graph_frame_sizes:
+    - 8
+    - 33
+cuda_mps: true
+model_runner: v2

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py

@@ -1280,11 +1280,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             payload_kwarg: pooling_output,
             "request": request,
         }
-        supports_is_finished = getattr(
-            self,
-            "_custom_process_supports_is_finished",
-            self._custom_process_supports_is_finished_kwarg(),
-        )
+        # The signature check is cached at initialization; never re-inspect
+        # the hook for every request of every model step.
+        supports_is_finished = getattr(self, "_custom_process_supports_is_finished", None)
+        if supports_is_finished is None:
+            supports_is_finished = self._custom_process_supports_is_finished_kwarg()
+            self._custom_process_supports_is_finished = supports_is_finished
         is_finished_fn = getattr(request, "is_finished", None)
         if callable(is_finished_fn):
             try:

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.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
 """HiggsAudio codec decoder kernel for higgs-audio v2.
 
 This module hosts the parameter-side building blocks for the higgs-audio-v2
@@ -16,11 +16,14 @@ from __future__ import annotations
 
 import json
 import os
+import types
 from typing import Any
 
 import torch
 import torch.nn as nn
 
+from vllm_omni.platforms import current_omni_platform
+
 __all__ = [
     "HiggsAudioVQLayer",
     "HiggsAudioRVQ",
@@ -99,6 +102,138 @@ class _Snake1d(nn.Module):
         return _snake(x, self.alpha)
 
 
+class _FusedDacSnake1d(nn.Module):
+    """Drop-in for ``transformers`` ``DacModel`` Snake1d with a fused CUDA kernel.
+
+    Same ``alpha`` parameter, so checkpoints load unchanged. Inference on CUDA
+    runs one bit-identical kernel instead of about seven elementwise launches;
+    everything else keeps the transformers expression.
+    """
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+        self._inverse: tuple[tuple[int, int], torch.Tensor] | None = None
+
+    def can_fuse(self, x: torch.Tensor) -> bool:
+        # The kernel uses NVIDIA PTX and float32 intermediates. ROCm tensors
+        # also report is_cuda; retain the eager path there and for float64.
+        return (
+            current_omni_platform.is_cuda()
+            and x.is_cuda
+            and x.dim() == 3
+            and not torch.is_grad_enabled()
+            and x.dtype in (torch.float16, torch.bfloat16, torch.float32)
+            and x.dtype == self.alpha.dtype
+        )
+
+    def fused(
+        self,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        residual: torch.Tensor | None = None,
+        write_sum: bool = False,
+    ) -> tuple[torch.Tensor | None, torch.Tensor]:
+        """Snake of ``residual + (x + bias)`` in one kernel; see ``fused_snake1d``."""
+        from .snake_kernel import fused_snake1d
+
+        alpha = self.alpha.detach()
+        try:
+            key = (alpha.data_ptr(), alpha._version)
+        except RuntimeError:  # inference tensors carry no version counter
+            key = None
+        if key is not None and self._inverse is not None and self._inverse[0] == key:
+            inverse = self._inverse[1]
+        else:
+            # The eager (alpha + 1e-9).reciprocal(), computed once per weight
+            # update in the same dtype instead of on every call.
+            inverse = (alpha + 1e-9).reciprocal().reshape(-1).contiguous()
+            self._inverse = None if key is None else (key, inverse)
+        return fused_snake1d(x, alpha.reshape(-1), inverse, bias=bias, residual=residual, write_sum=write_sum)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.can_fuse(x):
+            return self.fused(x)[1]
+        shape = x.shape
+        x = x.reshape(shape[0], shape[1], -1)
+        x = x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2)
+        return x.reshape(shape)
+
+
+def _conv_without_bias(conv: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    # cuDNN convolutions add the bias as a separate kernel anyway, so leaving
+    # it to the next fused Snake changes no bits.
+    if isinstance(conv, nn.ConvTranspose1d):
+        return nn.functional.conv_transpose1d(
+            x, conv.weight, None, conv.stride, conv.padding, conv.output_padding, conv.groups, conv.dilation
+        )
+    return conv._conv_forward(x, conv.weight, None)
+
+
+def _fused_dac_decoder_forward(decoder: nn.Module, hidden_state: torch.Tensor) -> torch.Tensor:
+    """transformers ``DacDecoder.forward`` with conv biases and residual adds
+    folded into the Snake kernels (5 elementwise passes per residual unit -> 2).
+
+    A value is carried as ``residual + (raw + bias)`` until the next Snake
+    consumes it, and only materialized where a residual connection needs it.
+    """
+    if not decoder.snake1.can_fuse(hidden_state):
+        return type(decoder).forward(decoder, hidden_state)
+    raw, bias, residual = _conv_without_bias(decoder.conv1, hidden_state), decoder.conv1.bias, None
+    for block in decoder.block:
+        _, act = block.snake1.fused(raw, bias, residual)
+        raw, bias, residual = _conv_without_bias(block.conv_t1, act), block.conv_t1.bias, None
+        for unit in (block.res_unit1, block.res_unit2, block.res_unit3):
+            residual, act = unit.snake1.fused(raw, bias, residual, write_sum=True)
+            _, act = unit.snake2.fused(_conv_without_bias(unit.conv1, act), unit.conv1.bias)
+            raw, bias = _conv_without_bias(unit.conv2, act), unit.conv2.bias
+    _, act = decoder.snake1.fused(raw, bias, residual)
+    return decoder.tanh(decoder.conv2(act))
+
+
+def _fuse_dac_decoder(decoder: nn.Module) -> int:
+    """Swap transformers DAC Snake1d modules for fused ones (before weights
+    load) and, when the decoder has the stock layout, fuse its forward."""
+    try:
+        from transformers.models.dac.modeling_dac import DacDecoder
+        from transformers.models.dac.modeling_dac import Snake1d as DacSnake1d
+    except ImportError:
+        return 0
+    replaced = 0
+    for module in list(decoder.modules()):
+        for name, child in list(module.named_children()):
+            if isinstance(child, DacSnake1d):
+                fused = _FusedDacSnake1d(int(child.alpha.shape[1]))
+                fused.to(device=child.alpha.device, dtype=child.alpha.dtype)
+                setattr(module, name, fused)
+                replaced += 1
+    if type(decoder) is DacDecoder and _has_fusable_dac_layout(decoder):
+        decoder.forward = types.MethodType(_fused_dac_decoder_forward, decoder)
+    return replaced
+
+
+def _has_fusable_dac_layout(decoder: nn.Module) -> bool:
+    convs = [decoder.conv1]
+    for block in decoder.block:
+        units = (block.res_unit1, block.res_unit2, block.res_unit3)
+        convs.append(block.conv_t1)
+        for unit in units:
+            conv1, conv2 = unit.conv1, unit.conv2
+            # Length-preserving convs, so the residual add needs no crop.
+            if conv1.stride[0] != 1 or 2 * conv1.padding[0] != conv1.dilation[0] * (conv1.kernel_size[0] - 1):
+                return False
+            if conv2.stride[0] != 1 or conv2.kernel_size[0] != 1 or conv2.padding[0] != 0:
+                return False
+            convs += [conv1, conv2]
+        if not all(isinstance(m, _FusedDacSnake1d) for m in (block.snake1, *(u.snake1 for u in units))):
+            return False
+        if not all(isinstance(u.snake2, _FusedDacSnake1d) for u in units):
+            return False
+    return isinstance(decoder.snake1, _FusedDacSnake1d) and all(
+        conv.bias is not None and conv.padding_mode == "zeros" for conv in convs
+    )
+
+
 def _wn_conv1d(*args: Any, **kwargs: Any) -> nn.Module:
     from torch.nn.utils import weight_norm
 
@@ -248,6 +383,7 @@ def build_higgs_audio_acoustic_decoder(
     dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"])
     dac_model = DacModel(dac_cfg)
     decoder = dac_model.decoder.to(device)
+    _fuse_dac_decoder(decoder)
     adjust_conv_transpose_output_padding(decoder)
     if hasattr(decoder, "tanh"):
         decoder.tanh = nn.Identity()

vllm_omni/model_executor/models/higgs_audio_v2/snake_kernel.py

@@ -0,0 +1,113 @@
+# ruff: noqa: N803
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused DAC Snake activation: ``x + 1 / (alpha + 1e-9) * sin(alpha * x) ** 2``.
+
+The eager expression launches about seven elementwise kernels per call, each
+reading and writing the whole activation; in the codec's CUDA graphs they are
+most of its GPU time. This kernel reads ``x`` once and rounds every
+intermediate to the tensor dtype exactly where the eager ops do, so outputs are
+bit-identical to the unfused expression. It can also absorb the preceding conv
+bias add and residual add, which are otherwise two more full passes each.
+"""
+
+import torch
+from vllm.triton_utils import tl, tldevice, triton
+
+
+@triton.jit
+def _mul_rn(a, b):
+    # Explicit round-to-nearest ops are never contracted into FMA, so each
+    # product/sum rounds exactly where the separate eager kernels round.
+    return tl.inline_asm_elementwise("mul.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1)
+
+
+@triton.jit
+def _add_rn(a, b):
+    return tl.inline_asm_elementwise("add.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1)
+
+
+@triton.jit
+def _dac_snake_kernel(
+    x_ptr,
+    bias_ptr,
+    residual_ptr,
+    alpha_ptr,
+    inverse_ptr,
+    sum_ptr,
+    out_ptr,
+    stride_b,
+    stride_c,
+    t_len,
+    HAS_BIAS: tl.constexpr,
+    HAS_RESIDUAL: tl.constexpr,
+    WRITE_SUM: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+):
+    batch = tl.program_id(0).to(tl.int64)
+    channel = tl.program_id(1).to(tl.int64)
+    t = tl.program_id(2) * BLOCK_T + tl.arange(0, BLOCK_T)
+    mask = t < t_len
+    offset = batch * stride_b + channel * stride_c + t
+    x = tl.load(x_ptr + offset, mask=mask, other=0.0)
+    dtype = x.dtype
+    xf = x.to(tl.float32)
+    if HAS_BIAS:
+        # The conv bias add that cuDNN convolutions run as a separate kernel.
+        xf = _add_rn(xf, tl.load(bias_ptr + channel).to(tl.float32)).to(dtype).to(tl.float32)
+    if HAS_RESIDUAL:
+        residual = tl.load(residual_ptr + offset, mask=mask, other=0.0).to(tl.float32)
+        xf = _add_rn(residual, xf).to(dtype).to(tl.float32)
+    if WRITE_SUM:
+        tl.store(sum_ptr + offset, xf.to(dtype), mask=mask)
+    alpha = tl.load(alpha_ptr + channel).to(tl.float32)
+    inverse = tl.load(inverse_ptr + channel).to(tl.float32)
+    scaled = _mul_rn(alpha, xf).to(dtype).to(tl.float32)
+    # libdevice sinf, as the eager CUDA kernel (tl.sin is the approximation).
+    sine = tldevice.sin(scaled).to(dtype).to(tl.float32)
+    square = _mul_rn(sine, sine).to(dtype).to(tl.float32)
+    term = _mul_rn(inverse, square).to(dtype).to(tl.float32)
+    tl.store(out_ptr + offset, _add_rn(xf, term).to(dtype), mask=mask)
+
+
+def fused_snake1d(
+    x: torch.Tensor,
+    alpha: torch.Tensor,
+    inverse: torch.Tensor,
+    *,
+    bias: torch.Tensor | None = None,
+    residual: torch.Tensor | None = None,
+    write_sum: bool = False,
+) -> tuple[torch.Tensor | None, torch.Tensor]:
+    """Snake of ``residual + (x + bias)``, rounded as the separate eager ops.
+
+    x, residual: (B, C, T); alpha, inverse, bias: (C,), all in x's dtype.
+    Returns ``(sum, snake(sum))``; ``sum`` is only materialized when
+    ``write_sum`` is set (for the next residual connection), else None.
+    """
+    x = x.contiguous()
+    if residual is not None:
+        residual = residual.contiguous()
+    batch, channels, t_len = x.shape
+    out = torch.empty_like(x)
+    total = torch.empty_like(x) if write_sum else None
+    if x.numel() == 0:
+        return total, out
+    block = min(triton.next_power_of_2(t_len), 2048)
+    _dac_snake_kernel[(batch, channels, triton.cdiv(t_len, block))](
+        x,
+        x if bias is None else bias,
+        x if residual is None else residual,
+        alpha,
+        inverse,
+        out if total is None else total,
+        out,
+        x.stride(0),
+        x.stride(1),
+        t_len,
+        HAS_BIAS=bias is not None,
+        HAS_RESIDUAL=residual is not None,
+        WRITE_SUM=write_sum,
+        BLOCK_T=block,
+    )
+    return total, out

vllm_omni/model_executor/models/higgs_audio_v3/full_sample_graph.py

@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Dense decode sampler graphs captured and parity-checked during startup.
+
+Capture compares eager and captured state transitions with identical noise.
+Serving only replays warmed shapes and preserves request torch RNG draws. CPU payload staging
+remains outside capture. This module does not handle mixed prefill batches.
+"""
+
+import copy
+import time
+
+import torch
+from vllm.logger import init_logger
+from vllm.v1.outputs import SamplerOutput
+
+logger = init_logger(__name__)
+
+_STATE = (
+    "_decode_last_codes",
+    "_decode_has_codes",
+    "_decode_delay_count",
+    "_decode_eoc_countdown",
+    "_decode_generation_done",
+)
+
+
+_BUCKETS = (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 96, 128)
+_FIELDS = ("temperature", "top_k", "top_p")
+
+
+def _graph_key(model, batch, hidden, force_audio_inputs, metadata):
+    return (
+        batch,
+        hidden.dtype,
+        hidden.device,
+        bool(force_audio_inputs),
+        bool(metadata.all_greedy),
+        tuple(getattr(metadata, f) is None for f in _FIELDS),
+        tuple(getattr(model, name).data_ptr() for name in _STATE),
+    )
+
+
+def _replay_captured(model, hidden, metadata, force_audio_inputs):
+    """Replay an already captured (bucket) graph writing inputs in place.
+
+    The capture path below pads through temporaries on every call. Once a
+    graph exists, only the live rows are written into its static inputs:
+    padded rows keep finite values from earlier steps and own no request
+    state, forced-audio inputs are constants, sampling parameters are copied
+    only when the cached metadata changes, and request noise is drawn directly
+    into the graph's noise input.
+    """
+    cache = getattr(model, "_dense_sample_graphs", None)
+    if not cache:
+        return None
+    batch = hidden.shape[0]
+    bucket = batch
+    if getattr(model.config, "audio_mrv2_sample_graph_buckets", False):
+        bucket = next((n for n in _BUCKETS if n >= batch), batch)
+    model._ensure_decode_state_capacity(bucket, hidden.device)
+    key = _graph_key(model, bucket, hidden, force_audio_inputs, metadata)
+    entry = cache.get(key)
+    if entry is None:
+        return None
+    graph, static_hidden, static_input_ids, static_noise, static_metadata, output, staging, keepalive = entry
+    books = model.num_codebooks
+    static_hidden[:batch].copy_(hidden)
+    if not force_audio_inputs:
+        static_input_ids[:batch].copy_(model._last_step_input_ids)
+        if batch < bucket:
+            static_input_ids[batch:].fill_(model._audio_continuation_id)
+    if metadata.generators:
+        from vllm_omni.utils.seeded_exponential import fill_exponential_rows
+
+        fill_exponential_rows(static_noise[: batch * books], [metadata.generators.get(i) for i in range(batch)])
+    else:
+        static_noise[: batch * books].exponential_()
+    sources = getattr(model, "_dense_sample_metadata_sources", None)
+    if sources is None:
+        sources = model._dense_sample_metadata_sources = {}
+    if sources.get(key) is not metadata:
+        for name in _FIELDS:
+            value = getattr(metadata, name)
+            if value is not None:
+                getattr(static_metadata, name)[:batch].copy_(value)
+        sources[key] = metadata
+    graph.replay()
+    model._finish_audio_staging(staging[:batch], batch, hidden.device, batch)
+    return SamplerOutput(sampled_token_ids=output.sampled_token_ids[:batch].clone(), logprobs_tensors=None)
+
+
+def run_dense_sample(model, hidden, logits, metadata, force_audio_inputs=False, active_batch=None):
+    batch = hidden.shape[0]
+    if active_batch is None:
+        replayed = _replay_captured(model, hidden, metadata, force_audio_inputs)
+        if replayed is not None:
+            return replayed
+    # Serving never captures a new shape. Unsupported shapes/metadata use the
+    # regular sampler; startup owns capture and its parity checks.
+    if not getattr(model, "_warming_sample_graphs", False):
+        return None
+    if active_batch is None and getattr(model.config, "audio_mrv2_sample_graph_buckets", False):
+        bucket = next((n for n in (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 96, 128) if n >= batch), batch)
+        if bucket != batch:
+            # Dummy rows own no request state. The batch-state gather replaces
+            # these rows before any later request can use them.
+            model._ensure_decode_state_capacity(bucket, hidden.device)
+            padded_hidden = hidden.new_zeros((bucket, hidden.shape[1]))
+            padded_hidden[:batch].copy_(hidden)
+            padded_metadata = copy.copy(metadata)
+            for name in ("temperature", "top_k", "top_p"):
+                value = getattr(metadata, name)
+                if value is not None:
+                    padded = value.new_ones(bucket)
+                    padded[:batch].copy_(value)
+                    setattr(padded_metadata, name, padded)
+            original_ids = model._last_step_input_ids
+            padded_ids = original_ids.new_full((bucket,), model._audio_continuation_id)
+            if not force_audio_inputs:
+                padded_ids[:batch].copy_(original_ids)
+            model._last_step_input_ids = padded_ids
+            try:
+                output = run_dense_sample(
+                    model,
+                    padded_hidden,
+                    logits.new_empty((bucket, logits.shape[1])),
+                    padded_metadata,
+                    force_audio_inputs=force_audio_inputs,
+                    active_batch=batch,
+                )
+            finally:
+                model._last_step_input_ids = original_ids
+            return SamplerOutput(sampled_token_ids=output.sampled_token_ids[:batch], logprobs_tensors=None)
+    actual = batch if active_batch is None else active_batch
+    effective_ids = (
+        torch.full((batch,), model._audio_continuation_id, dtype=model._last_step_input_ids.dtype, device=hidden.device)
+        if force_audio_inputs
+        else model._last_step_input_ids
+    )
+    books = model.num_codebooks
+    vocab = model.modality_head.vocab_size
+    noise = torch.empty((batch * books, vocab), dtype=torch.float32, device=hidden.device)
+    if actual < batch:
+        noise[actual * books :].fill_(1)
+    if metadata.generators:
+        from vllm_omni.utils.seeded_exponential import fill_exponential_rows
+
+        fill_exponential_rows(noise[: actual * books], [metadata.generators.get(i) for i in range(actual)])
+    else:
+        noise[: actual * books].exponential_()
+    fields = _FIELDS
+    # force_audio_inputs controls capture-time Python branches (including
+    # the mixed-tail assertion). Never reuse that graph for real EOS inputs.
+    key = _graph_key(model, batch, hidden, force_audio_inputs, metadata)
+    cache = getattr(model, "_dense_sample_graphs", None)
+    if cache is None:
+        cache = model._dense_sample_graphs = {}
+    entry = cache.get(key)
+    if entry is None:
+        static_hidden = hidden.clone()
+        static_input_ids = effective_ids.clone()
+        static_noise = noise.clone()
+        static_metadata = copy.copy(metadata)
+        for name in fields:
+            value = getattr(metadata, name)
+            setattr(static_metadata, name, None if value is None else value.clone())
+        # Ensure the reusable staging exists before capture. Each graph keeps
+        # its own reference if a later larger batch grows the model buffer.
+        staging = model._get_audio_gpu_staging_buffer(batch, hidden.device)
+        state = {name: getattr(model, name).clone() for name in _STATE}
+        original_ids = model._last_step_input_ids
+        original_qsl = model._last_step_query_start_loc
+        original_mode_rows = getattr(model, "_step_audio_mode_rows", 0)
+        original_tail_rows = getattr(model, "_step_audio_tail_rows", 0)
+        original_direct_rows = getattr(model, "_fast_audio_direct_rows", 0)
+        # Dense decode has one token per row. Own this input instead of
+        # capturing a view of the runner's resizable staging buffer.
+        static_qsl = torch.arange(batch + 1, dtype=torch.int32, device=hidden.device)
+
+        def restore():
+            for name, value in state.items():
+                getattr(model, name).copy_(value)
+
+        try:
+            model._in_audio_sample_graph = True
+            # Outer eligibility and real-input validation already ran. The
+            # owned graph inputs represent exactly one audio-state row each.
+            # Normalize capture-time host dispatch; never capture an outer
+            # mixed-batch assertion or a stock-sampler branch into this graph.
+            model._step_audio_mode_rows = 0
+            model._step_audio_tail_rows = 0
+            model._fast_audio_direct_rows = batch
+            model._audio_graph_noise = static_noise
+            model._last_step_input_ids = static_input_ids
+            model._last_step_query_start_loc = static_qsl
+            model._last_logits_hidden = static_hidden
+            reference = model.sample(logits, static_metadata)
+            expected_ids = reference.sampled_token_ids.clone()
+            expected_staging = staging.clone()
+            expected_state = {name: getattr(model, name).clone() for name in _STATE}
+            restore()
+            model._last_logits_hidden = static_hidden
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                output = model.sample(logits, static_metadata)
+            # CUDA capture records work; it does not execute the transition.
+            # Validate only after replay from the preserved initial state.
+            restore()
+            graph.replay()
+            torch.testing.assert_close(output.sampled_token_ids, expected_ids, atol=0, rtol=0)
+            torch.testing.assert_close(staging, expected_staging, atol=0, rtol=0)
+            for name, expected in expected_state.items():
+                torch.testing.assert_close(getattr(model, name), expected, atol=0, rtol=0)
+            restore()
+            logger.debug("Higgs sample graph parity checked for batch=%d", batch)
+        finally:
+            model._in_audio_sample_graph = False
+            model._step_audio_mode_rows = original_mode_rows
+            model._step_audio_tail_rows = original_tail_rows
+            model._fast_audio_direct_rows = original_direct_rows
+            model._audio_graph_noise = None
+            model._last_step_input_ids = original_ids
+            model._last_step_query_start_loc = original_qsl
+            model._last_logits_hidden = None
+        # Graphs reference storage, not Python owners. Model caches and state
+        # can grow after capture; keep their original allocations alive.
+        keepalive = [getattr(model, name) for name in _STATE]
+        for name in ("_row_index_cache", "_codebook_index_cache", "_boc_frame_cache"):
+            keepalive.extend(getattr(model, name).values())
+        keepalive.append(static_qsl)
+        entry = (graph, static_hidden, static_input_ids, static_noise, static_metadata, output, staging, keepalive)
+        cache[key] = entry
+    graph, static_hidden, static_input_ids, static_noise, static_metadata, output, staging, keepalive = entry
+    static_hidden.copy_(hidden)
+    static_input_ids.copy_(effective_ids)
+    static_noise.copy_(noise)
+    for name in fields:
+        value = getattr(metadata, name)
+        if value is not None:
+            getattr(static_metadata, name).copy_(value)
+    graph.replay()
+    model._finish_audio_staging(staging[:actual], actual, hidden.device, actual)
+    return SamplerOutput(sampled_token_ids=output.sampled_token_ids.clone(), logprobs_tensors=None)
+
+
+@torch.inference_mode()
+def capture_sample_graphs(model, max_num_reqs):
+    """Warm the MRv2 sampler without consuming request RNG or decode state."""
+    if not model._sample_graph_enabled or getattr(model, "_sample_graphs_ready", False):
+        return
+    from vllm.v1.sample.logits_processor import LogitsProcessors
+    from vllm.v1.sample.metadata import SamplingMetadata
+
+    device = model.model.embed_tokens.weight.device
+    if device.type != "cuda":
+        return
+    start = time.perf_counter()
+    memory_before = torch.accelerator.memory_allocated(device)
+    reserved_before = torch.accelerator.memory_reserved(device)
+    model._resolve_token_ids()
+    buckets = (
+        sorted({next((b for b in _BUCKETS if b >= n), n) for n in range(1, max_num_reqs + 1)})
+        if getattr(model.config, "audio_mrv2_sample_graph_buckets", False)
+        else range(1, max_num_reqs + 1)
+    )
+    model._ensure_decode_state_capacity(max(buckets), device)
+    model._get_audio_gpu_staging_buffer(max(buckets), device)
+    saved_state = {name: getattr(model, name).clone() for name in _STATE}
+    fields = (
+        "_last_step_input_ids",
+        "_last_step_query_start_loc",
+        "_last_logits_hidden",
+        "_step_audio_mode_rows",
+        "_step_audio_tail_rows",
+        "_fast_audio_direct_rows",
+    )
+    saved = {name: getattr(model, name, None) for name in fields}
+    try:
+        model._warming_sample_graphs = True
+        with torch.random.fork_rng(devices=[device.index]):
+            for batch in buckets:
+                hidden = model.model.embed_tokens.weight.new_zeros((batch, model.model.embed_tokens.weight.shape[1]))
+                logits = hidden.new_empty((batch, model.config.text_config.vocab_size))
+                model._last_step_input_ids = torch.full(
+                    (batch,), model._audio_continuation_id, device=device, dtype=torch.long
+                )
+                model._last_step_query_start_loc = torch.arange(batch + 1, device=device, dtype=torch.int32)
+                ones = torch.ones(batch, device=device)
+                zeros = torch.zeros_like(ones)
+                for greedy in (False, True):
+                    metadata = SamplingMetadata(
+                        temperature=zeros if greedy else ones,
+                        all_greedy=greedy,
+                        all_random=not greedy,
+                        top_p=ones,
+                        top_k=torch.full((batch,), -1, device=device, dtype=torch.long),
+                        generators={},
+                        max_num_logprobs=None,
+                        no_penalties=True,
+                        prompt_token_ids=None,
+                        frequency_penalties=zeros,
+                        presence_penalties=zeros,
+                        repetition_penalties=ones,
+                        output_token_ids=[[] for _ in range(batch)],
+                        allowed_token_ids_mask=None,
+                        bad_words_token_ids={},
+                        logitsprocs=LogitsProcessors(),
+                    )
+                    for forced in (False, True):
+                        run_dense_sample(model, hidden, logits, metadata, force_audio_inputs=forced)
+    finally:
+        for name, value in saved_state.items():
+            getattr(model, name).copy_(value)
+        for name, value in saved.items():
+            setattr(model, name, value)
+        model._warming_sample_graphs = False
+        model._clear_last_audio_outputs()
+        model._async_audio_gpu_staging = None
+        model._dense_sample_metadata_sources = {}
+    torch.accelerator.synchronize()
+    model._sample_graphs_ready = True
+    logger.info(
+        "Higgs sampler captured %d CUDA graphs in %.2f s "
+        "(allocated delta %.1f MiB, reserved delta %.1f MiB, reserved total %.1f MiB)",
+        len(model._dense_sample_graphs),
+        time.perf_counter() - start,
+        (torch.accelerator.memory_allocated(device) - memory_before) / 2**20,
+        (torch.accelerator.memory_reserved(device) - reserved_before) / 2**20,
+        torch.accelerator.memory_reserved(device) / 2**20,
+    )

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.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
 """Stage 1 (codec decoder) for higgs-audio v3.
 
 Reuses higgs-audio-v2's RVQ + DAC codec decoder but loads weights from the
@@ -10,6 +10,7 @@ prefix) rather than from a standalone tokenizer repo.
 from __future__ import annotations
 
 import os
+import time
 from typing import Any
 
 import torch
@@ -57,6 +58,9 @@ class HiggsAudioV3Code2Wav(nn.Module):
 
     input_modalities = "audio"
 
+    supports_native_payload_input = True
+    batched_gpu_staging_keys = {("codes", "audio")}
+
     def __init__(
         self,
         config: HiggsAudioV3Config | None = None,
@@ -81,6 +85,11 @@ class HiggsAudioV3Code2Wav(nn.Module):
             self._model_path = None
             self.vllm_config = None
 
+        _, self._codec_graph_enabled = self.config.resolve_graph_defaults(
+            use_v2_model_runner=bool(
+                getattr(getattr(self.vllm_config, "model_config", None), "use_v2_model_runner", False)
+            )
+        )
         self.sample_rate: int = int(self.config.sample_rate)
         self.num_codebooks: int = int(self.config.num_codebooks)
         self.num_real_codes: int = int(self.config.num_real_codes)
@@ -97,6 +106,8 @@ class HiggsAudioV3Code2Wav(nn.Module):
         self.quantizer: HiggsAudioRVQ | None = None
         self.fc2: nn.Linear | None = None
         self.acoustic_decoder: nn.Module | None = None
+        self._decode_graphs: dict[tuple[int, int], tuple[Any, torch.Tensor, torch.Tensor]] = {}
+        self._decode_graphs_initialized = False
         self._loaded: bool = False
         # Do NOT eagerly load from standalone tokenizer here.
         # load_weights() will try bundled V3 codec first, then fall back.
@@ -452,6 +463,30 @@ class HiggsAudioV3Code2Wav(nn.Module):
             self._ensure_codec_loaded()
 
         codes = self._validate_codes(audio_codes)
+        if codes.is_cuda and self._codec_graph_enabled:
+            batch, _, frames = codes.shape
+            entry = self._decode_graphs.get((batch, frames))
+            if entry is not None:
+                graph, static_codes, static_audio = entry
+                static_codes[:batch].copy_(codes)
+                graph.replay()
+                # The next group may replay the same graph. Own the output
+                # before reusing its storage or shared graph memory pool.
+                return static_audio[:batch].clone()
+        return self._decode_codes_impl(codes)
+
+    @torch.inference_mode()
+    def capture_auxiliary_graphs(self) -> None:
+        """Capture exact codec shapes before the generation worker is ready."""
+        if not self._codec_graph_enabled or self._decode_graphs_initialized:
+            return
+        if not self._loaded:
+            self._ensure_codec_loaded()
+        device = self.fc2.weight.device
+        if device.type == "cuda":
+            self._capture_decode_graphs(device)
+
+    def _decode_codes_impl(self, codes: torch.Tensor) -> torch.Tensor:
         rvq_codes = codes.transpose(0, 1).long()
         quantized = self.quantizer.decode(rvq_codes)
         quantized = quantized.to(dtype=self.fc2.weight.dtype)
@@ -464,6 +499,55 @@ class HiggsAudioV3Code2Wav(nn.Module):
             audio = audio.unsqueeze(1)
         return audio.to(codes.device)
 
+    def _decode_graph_shapes(self) -> list[tuple[int, int]]:
+        single_max = int(getattr(self.config, "codec_graph_single_max_frames", 150))
+        batches = getattr(self.config, "codec_graph_batch_sizes", (2, 4, 8, 16))
+        frames = getattr(self.config, "codec_graph_frame_sizes", (5, 29, 54))
+        if single_max < 1 or any(int(n) < 1 for n in (*batches, *frames)):
+            raise ValueError("Higgs codec graph dimensions must be positive")
+        return sorted({(1, n) for n in range(1, single_max + 1)} | {(int(b), int(f)) for b in batches for f in frames})
+
+    @torch.inference_mode()
+    def _capture_decode_graphs(self, device: torch.device) -> None:
+        # Never pad time: convolution boundary conditions must stay exact.
+        # Batch size also stays exact: padding can change convolution
+        # algorithms and rounding. Uncovered shapes retain the eager path.
+        start = time.perf_counter()
+        memory_before = torch.accelerator.memory_allocated(device)
+        reserved_before = torch.accelerator.memory_reserved(device)
+        shapes = self._decode_graph_shapes()
+        previous_device = current_omni_platform.current_device()
+        current_omni_platform.set_device(device)
+        try:
+            stream = torch.cuda.Stream(device=device)
+            stream.wait_stream(torch.cuda.current_stream(device))
+            pool = torch.cuda.graph_pool_handle()
+            for batch, frames in shapes:
+                codes = torch.zeros((batch, self.num_codebooks, frames), device=device, dtype=torch.long)
+                stream.wait_stream(torch.cuda.current_stream(device))
+                with torch.cuda.stream(stream):
+                    for _ in range(3):
+                        self._decode_codes_impl(codes)
+                stream.synchronize()
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph, pool=pool, stream=stream):
+                    audio = self._decode_codes_impl(codes)
+                self._decode_graphs[batch, frames] = (graph, codes, audio)
+            torch.cuda.current_stream(device).wait_stream(stream)
+        finally:
+            current_omni_platform.set_device(torch.device(device.type, previous_device))
+        self._decode_graphs_initialized = True
+        torch.accelerator.synchronize()
+        logger.info(
+            "Higgs codec captured %d exact-frame CUDA graphs in %.2f s "
+            "(allocated delta %.1f MiB, reserved delta %.1f MiB, reserved total %.1f MiB)",
+            len(self._decode_graphs),
+            time.perf_counter() - start,
+            (torch.accelerator.memory_allocated(device) - memory_before) / 2**20,
+            (torch.accelerator.memory_reserved(device) - reserved_before) / 2**20,
+            torch.accelerator.memory_reserved(device) / 2**20,
+        )
+
     @torch.inference_mode()
     def forward_chunk(
         self,
@@ -518,6 +602,18 @@ class HiggsAudioV3Code2Wav(nn.Module):
         ids = input_ids.reshape(-1).to(dtype=torch.long)
         request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
 
+        # Native MRV2 uses token placeholders for scheduling. The connector
+        # payload owns the actual codebook-major codec sequence, including an
+        # empty terminal payload which must not decode a placeholder token.
+        if runtime_additional_information is not None:
+            for i, info in enumerate(runtime_additional_information):
+                if i >= len(request_ids_list):
+                    break
+                codes = info.get("codes", {}) if isinstance(info, dict) else {}
+                audio = codes.get("audio") if isinstance(codes, dict) else None
+                if isinstance(audio, torch.Tensor):
+                    request_ids_list[i] = audio.reshape(-1).to(device=ids.device, dtype=torch.long)
+
         left_context_size = [0] * len(request_ids_list)
         right_holdback_size = [0] * len(request_ids_list)
         if runtime_additional_information is not None:
@@ -530,11 +626,14 @@ class HiggsAudioV3Code2Wav(nn.Module):
                 if "right_holdback_size" in meta:
                     right_holdback_size[i] = int(meta["right_holdback_size"])
 
-        wavs: list[torch.Tensor] = []
+        wavs: list[torch.Tensor] = [empty] * len(request_ids_list)
+        # Requests with identical codec windows are independent along B. Keep
+        # time lengths and overlap trims exact rather than padding the temporal
+        # axis, which could change the convolution boundary conditions.
+        groups: dict[tuple[int, int, int], list[tuple[int, torch.Tensor]]] = {}
         for i, req_ids in enumerate(request_ids_list):
             n = int(req_ids.numel())
             if n == 0:
-                wavs.append(empty)
                 continue
             if n % self.num_codebooks != 0:
                 logger.warning(
@@ -542,23 +641,43 @@ class HiggsAudioV3Code2Wav(nn.Module):
                     n,
                     self.num_codebooks,
                 )
-                wavs.append(empty)
                 continue
             frames = n // self.num_codebooks
-            codes_qf = req_ids.reshape(self.num_codebooks, frames)
-            codes_bqf = codes_qf.unsqueeze(0)
+            key = (frames, left_context_size[i], right_holdback_size[i])
+            groups.setdefault(key, []).append((i, req_ids.reshape(self.num_codebooks, frames)))
+
+        for (_, left, right), group in groups.items():
+            codes_bqf = group[0][1].unsqueeze(0) if len(group) == 1 else torch.stack([codes for _, codes in group])
             try:
                 pcm = self.forward_chunk(
                     codes_bqf,
-                    left_context_size=left_context_size[i],
-                    right_holdback_size=right_holdback_size[i],
+                    left_context_size=left,
+                    right_holdback_size=right,
                     hop_length=self.hop_length,
                 )
             except ValueError as exc:
-                logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", exc)
-                wavs.append(empty)
+                if len(group) == 1:
+                    logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", exc)
+                    continue
+                # An invalid row must not discard valid peers. This slow path
+                # retains the previous per-request validation/error behavior.
+                for index, codes in group:
+                    try:
+                        row_pcm = self.forward_chunk(
+                            codes.unsqueeze(0),
+                            left_context_size=left,
+                            right_holdback_size=right,
+                            hop_length=self.hop_length,
+                        )
+                        wavs[index] = row_pcm[0, 0].to(torch.float32)
+                    except ValueError as row_exc:
+                        logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", row_exc)
                 continue
-            wavs.append(pcm.squeeze(0).squeeze(0).to(torch.float32).cpu())
+            pcm = pcm.to(torch.float32)
+            for row, (index, _) in enumerate(group):
+                # The generation runner owns the batched pinned D2H copy and
+                # its completion event; avoid blocking once per request here.
+                wavs[index] = pcm[row, 0]
 
         return OmniOutput(
             text_hidden_states=None,

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py

@@ -108,12 +108,26 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
     prefer_model_sampler: bool = True
     # Tell the runner to call postprocess() to emit per-step audio codes.
     have_multimodal_outputs: bool = True
-    has_postprocess: bool = True
+    has_postprocess: bool = False
     skips_model_sampler_output_token_history: bool = True
     postprocess_uses_hidden_states: bool = False
     postprocess_uses_multimodal_outputs: bool = False
     postprocess_uses_req_infos: bool = False
     supports_omni_query_start_loc: bool = True
+    requires_cpu_input_tail_ids: bool = True
+    # The codec stage consumes audio codes, never backbone hidden states.
+    omni_pooler_payload_include_hidden: bool = False
+
+    def capture_auxiliary_graphs(self) -> None:
+        from .full_sample_graph import capture_sample_graphs
+
+        scheduler = self.vllm_config.scheduler_config
+        capture_sample_graphs(self, scheduler.max_num_seqs)
+
+    def create_omni_model_state(self, vllm_config, encoder_cache, device):
+        from .model_state import HiggsModelState
+
+        return HiggsModelState(vllm_config, self, encoder_cache, device)
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -122,6 +136,17 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
             self.config = hf_config
         else:
             self.config = HiggsAudioV3Config(**hf_config.to_dict())
+        self._sample_graph_enabled, _ = self.config.resolve_graph_defaults(
+            use_v2_model_runner=bool(getattr(vllm_config.model_config, "use_v2_model_runner", False))
+        )
+        self.use_async_omni_output = getattr(self.config, "audio_async_payload", False) is True
+        self.supports_async_whole_payload = self.use_async_omni_output
+        if self.use_async_omni_output:
+            if not vllm_config.scheduler_config.async_scheduling:
+                raise ValueError("audio_async_payload requires async scheduling")
+            if getattr(self.config, "audio_async_prompt_mode", False) is not True:
+                raise ValueError("audio_async_payload requires audio prompt mode classification")
+        self._async_audio_gpu_staging = None
         model_config = getattr(vllm_config, "model_config", None)
         compilation_config = getattr(vllm_config, "compilation_config", None)
         cudagraph_mode = getattr(compilation_config, "cudagraph_mode", None)
@@ -345,10 +370,26 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         omni_query_start_loc: torch.Tensor | None = None,
         req_ids: list[str] | None = None,
+        cpu_input_tail_ids: list[int] | None = None,
+        audio_prompt_mode_rows: int = 0,
         **_: Any,
     ) -> None:
         """Update per-step metadata before runner forward or CUDA graph replay."""
         _ = positions, inputs_embeds
+        self._step_audio_mode_rows = (
+            audio_prompt_mode_rows if getattr(self.config, "audio_async_prompt_mode", False) is True else 0
+        )
+        self._step_audio_tail_rows = 0
+        if cpu_input_tail_ids:
+            self._resolve_token_ids()
+            if (
+                cpu_input_tail_ids
+                and self._audio_continuation_id is not None
+                and req_ids is not None
+                and len(cpu_input_tail_ids) == len(req_ids)
+                and all(token == self._audio_continuation_id for token in cpu_input_tail_ids)
+            ):
+                self._step_audio_tail_rows = len(cpu_input_tail_ids)
         if input_ids is not None:
             self._last_step_input_ids = input_ids
             if self._use_external_decode_cudagraph and input_ids.is_cuda and input_ids.ndim >= 1:
@@ -359,6 +400,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         if req_ids is not None:
             self._sync_decode_state_with_batch(req_ids)
 
+    def finish_decode_step_forward(self) -> None:
+        self._decode_step_metadata_from_runner = False
+
     def _ensure_state_pool_capacity(self, num_rows: int, device: torch.device) -> None:
         if not hasattr(self, "_state_pool_indices"):
             self._state_pool_indices: dict[str, int] = {}
@@ -387,17 +431,53 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         self._state_pool_eoc_countdown = resize(self._state_pool_eoc_countdown, -1)
         self._state_pool_generation_done = resize(self._state_pool_generation_done, False)
 
+    def _state_indices(self, rows: list[int], device: torch.device) -> torch.Tensor:
+        # A fresh pinned allocation is owned by PyTorch's host allocator until
+        # its asynchronous copy completes; never overwrite a shared CPU staging
+        # buffer while the previous decode step may still be reading it.
+        if device.type == "cuda" and getattr(self.config, "audio_mrv2_async_state_indices", False):
+            return torch.tensor(rows, dtype=torch.long, pin_memory=True).to(device, non_blocking=True)
+        return torch.tensor(rows, dtype=torch.long, device=device)
+
     def _sync_decode_state_with_batch(self, req_ids: list[str]) -> None:
         device = self._decode_has_codes.device
         self._ensure_state_pool_capacity(len(req_ids), device)
         previous_req_ids = getattr(self, "_last_batch_req_ids", [])
 
+        if (
+            device.type == "cuda"
+            and getattr(self.config, "audio_mrv2_fused_state", False)
+            and (previous_req_ids != req_ids or any(r not in self._state_pool_indices for r in req_ids))
+        ):
+            from .mrv2_state_kernels import sync_state
+
+            previous_rows = [row for row, req_id in enumerate(previous_req_ids) if req_id in self._state_pool_indices]
+            pool_rows = [self._state_pool_indices[previous_req_ids[row]] for row in previous_rows]
+            current_rows = []
+            for req_id in req_ids:
+                fresh = req_id not in self._state_pool_indices
+                if fresh:
+                    slot = (
+                        self._free_state_pool_indices.pop()
+                        if self._free_state_pool_indices
+                        else len(self._state_pool_indices)
+                    )
+                    self._ensure_state_pool_capacity(slot + 1, device)
+                    self._state_pool_indices[req_id] = slot
+                slot = self._state_pool_indices[req_id]
+                current_rows.append(-slot - 1 if fresh else slot)
+            self._ensure_decode_state_capacity(len(req_ids), device)
+            if req_ids or previous_rows:
+                sync_state(self, previous_rows, pool_rows, current_rows)
+            self._last_batch_req_ids = list(req_ids)
+            return
+
         if previous_req_ids != req_ids or any(req_id not in self._state_pool_indices for req_id in req_ids):
             previous_rows = [row for row, req_id in enumerate(previous_req_ids) if req_id in self._state_pool_indices]
             if previous_rows:
                 pool_rows = [self._state_pool_indices[previous_req_ids[row]] for row in previous_rows]
-                src = torch.tensor(previous_rows, dtype=torch.long, device=device)
-                dst = torch.tensor(pool_rows, dtype=torch.long, device=device)
+                src = self._state_indices(previous_rows, device)
+                dst = self._state_indices(pool_rows, device)
                 self._state_pool_has_codes[dst] = self._decode_has_codes[src]
                 self._state_pool_last_codes[dst] = self._decode_last_codes[src]
                 self._state_pool_delay_count[dst] = self._decode_delay_count[src]
@@ -420,10 +500,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
                 self._state_pool_generation_done[pool_row] = False
 
             self._ensure_decode_state_capacity(len(req_ids), device)
-            pool_rows = torch.tensor(
+            pool_rows = self._state_indices(
                 [self._state_pool_indices[req_id] for req_id in req_ids],
-                dtype=torch.long,
-                device=device,
+                device,
             )
             self._decode_has_codes[: len(req_ids)] = self._state_pool_has_codes[pool_rows]
             self._decode_last_codes[: len(req_ids)] = self._state_pool_last_codes[pool_rows]
@@ -450,8 +529,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         **kwargs: Any,
     ) -> torch.Tensor:
+        # A runner invocation may call forward repeatedly while warming and
+        # capturing a graph. Keep its metadata authoritative for that scope.
         metadata_from_runner = self._decode_step_metadata_from_runner
-        self._decode_step_metadata_from_runner = False
 
         info_dicts = kwargs.get("model_intermediate_buffer")
         if info_dicts is None:
@@ -501,7 +581,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
                 and inputs_embeds is None
                 and not self._is_single_token_decode_step(int(hidden_states.shape[0]))
             )
-        if is_prefill and info_dicts:
+        if is_prefill and info_dicts and inputs_embeds is None:
             # Voice clone: replace marked prompt positions with ref audio embeddings.
             hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, positions, info_dicts)
 
@@ -515,6 +595,11 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         return self._run_qwen3_layers_eager(positions, hidden_states)
 
     def _run_qwen3_layers_eager(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self._use_external_decode_cudagraph:
+            # Enter Qwen3Model's compilation boundary before the runner
+            # captures the decode graph. Audio feedback stays outside it.
+            return self.model(input_ids=None, positions=positions, inputs_embeds=hidden_states)
+
         residual: torch.Tensor | None = None
         for layer in self.model.layers:
             hidden_states, residual = layer(positions, hidden_states, residual)
@@ -619,6 +704,25 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
 
     def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor:
         self._last_logits_hidden = hidden_states
+        ids = getattr(self, "_last_step_input_ids", None)
+        rows = int(hidden_states.shape[0])
+        if (
+            self._audio_continuation_id is not None
+            and sampling_metadata is not None
+            and getattr(sampling_metadata, "no_penalties", False)
+            and (
+                (isinstance(ids, torch.Tensor) and ids.numel() == rows and self._fast_audio_direct_rows == rows)
+                or getattr(self, "_step_audio_tail_rows", 0) == rows
+                or getattr(self, "_step_audio_mode_rows", 0) == rows
+            )
+            and getattr(sampling_metadata, "max_num_logprobs", None) is None
+            and getattr(sampling_metadata, "allowed_token_ids_mask", None) is None
+            and not getattr(sampling_metadata, "bad_words_token_ids", None)
+        ):
+            # sample() already selects continuation/EOS directly in this
+            # state, using the audio head rather than the text head. Keep
+            # the normal vocabulary shape for runner-side processors.
+            return hidden_states.new_empty((rows, self._text_vocab_size), dtype=torch.float32)
         return self.logits_processor(self.lm_head, hidden_states)
 
     def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
@@ -656,7 +760,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         ``audio_placeholder_positions`` ([T] absolute prompt positions). Legacy
         callers that still submit -100 token IDs remain supported.
         """
-        if not info_dicts:
+        if not info_dicts or not any(
+            isinstance(info, dict) and info.get("audio_input_ids") is not None for info in info_dicts
+        ):
             return hidden_states
 
         flat_ids = input_ids.reshape(-1)
@@ -788,10 +894,6 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         if self._decode_active_audio_count == 0 and not external_decode_graph:
             return hidden_states
 
-        req_rows, token_positions = self._decode_request_token_positions(num_tokens, hidden_states.device)
-        if int(req_rows.numel()) == 0:
-            return hidden_states
-
         num_reqs = self._step_request_count(num_tokens)
         self._ensure_decode_state_capacity(num_reqs, hidden_states.device)
 
@@ -800,6 +902,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
             has_codes = self._decode_has_codes[:num_tokens].unsqueeze(-1)
             current_hidden = hidden_states
         else:
+            req_rows, token_positions = self._decode_request_token_positions(num_tokens, hidden_states.device)
+            if int(req_rows.numel()) == 0:
+                return hidden_states
             codes_slice = self._decode_last_codes.index_select(0, req_rows)
             has_codes = self._decode_has_codes.index_select(0, req_rows).unsqueeze(-1)
             current_hidden = hidden_states.index_select(0, token_positions)
@@ -932,6 +1037,44 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
             return None
         return tail_ids == int(audio_id)
 
+    def _step_tail_guard(self, num_rows: int, *, prompt_mode: bool, decode_only: bool) -> None:
+        from .step_kernels import step_tail_guard
+
+        error = getattr(self, "_decode_tail_guard_error", None)
+        device = self._decode_has_codes.device
+        if error is None or error.device != device:
+            error = self._decode_tail_guard_error = torch.zeros(1, dtype=torch.int32, device=device)
+        step_tail_guard(
+            self._last_step_input_ids.reshape(-1),
+            None if decode_only else self._last_step_query_start_loc,
+            self._decode_has_codes[:num_rows],
+            self._decode_generation_done[:num_rows],
+            self._decode_delay_count[:num_rows],
+            self._decode_eoc_countdown[:num_rows],
+            error,
+            audio_id=int(self._audio_continuation_id),
+            eos_id=int(self._eos_token_id),
+            prompt_mode=prompt_mode,
+        )
+        torch._assert_async(error[0] == 0, "CPU/GPU Higgs graph audio tail mismatch")
+
+    def _restore_terminal_audio_rows(self, num_rows):
+        # Async scheduling can execute one final in-flight EOS row after
+        # request retirement has released its state slot. EOS is terminal:
+        # never resurrect that row as a new audio stream with a forced input.
+        eos_id = self._eos_token_id
+        if eos_id is None:
+            return
+        ids = self._last_step_input_ids.reshape(-1)
+        qsl = self._last_step_query_start_loc
+        if isinstance(qsl, torch.Tensor) and qsl.numel() == num_rows + 1:
+            tail = ids.index_select(0, (qsl[1:].long() - 1).clamp_min(0))
+        elif ids.numel() == num_rows:
+            tail = ids
+        else:
+            raise RuntimeError("Cannot validate terminal audio rows")
+        self._decode_generation_done[:num_rows].logical_or_(tail == eos_id)
+
     def _fast_audio_sampler_gpu_fallback_reason(
         self,
         *,
@@ -997,6 +1140,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         and accumulate per-request state.
         """
         self._resolve_token_ids()
+        self._async_audio_gpu_staging = None
 
         audio_id = self._audio_continuation_id
 
@@ -1020,6 +1164,56 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         self._ensure_decode_state_capacity(num_rows, hidden.device)
         ids = getattr(self, "_last_step_input_ids", None)
         decode_only = isinstance(ids, torch.Tensor) and int(ids.numel()) == num_rows
+        prompt_mode_rows = getattr(self, "_step_audio_mode_rows", 0) == num_rows
+        graph_forced_audio = getattr(self, "_step_audio_tail_rows", 0) == num_rows or prompt_mode_rows
+        use_sample_graph = bool(
+            getattr(self, "_sample_graph_enabled", False)
+            and not getattr(self, "_in_audio_sample_graph", False)
+            and hidden.is_cuda
+            and (graph_forced_audio or (decode_only and self._fast_audio_direct_rows == num_rows))
+            and getattr(sampling_metadata, "no_penalties", False)
+            and self._fast_audio_sampler_gpu_fallback_reason(
+                logits=logits, sampling_metadata=sampling_metadata, num_rows=num_rows
+            )
+            is None
+        )
+        qsl = self._last_step_query_start_loc
+        fused_guard = (
+            use_sample_graph
+            and graph_forced_audio
+            and self._eos_token_id is not None
+            and (decode_only or (isinstance(qsl, torch.Tensor) and int(qsl.numel()) == num_rows + 1))
+        )
+        if fused_guard:
+            # One launch restores terminal EOS rows, resets new prefill rows
+            # and validates the audio tail (~20 small eager kernels before).
+            self._step_tail_guard(num_rows, prompt_mode=prompt_mode_rows, decode_only=decode_only)
+        elif prompt_mode_rows:
+            self._restore_terminal_audio_rows(num_rows)
+        if use_sample_graph:
+            from .full_sample_graph import run_dense_sample
+
+            if not fused_guard and graph_forced_audio:
+                actual_audio = self._audio_seed_mask_from_step_input(num_rows, hidden.device)
+                if getattr(self, "_step_audio_mode_rows", 0) == num_rows:
+                    # Reset new prefill rows before accepting old active state.
+                    if not decode_only:
+                        self._reset_decode_state_rows(
+                            self._prefill_row_mask(num_rows, hidden.device), num_rows, hidden.device
+                        )
+                    actual_audio = (
+                        actual_audio | self._decode_has_codes[:num_rows] | self._decode_generation_done[:num_rows]
+                    )
+                torch._assert_async(actual_audio.all(), "CPU/GPU Higgs graph audio tail mismatch")
+            if not decode_only:
+                if not fused_guard:
+                    pfmask = self._prefill_row_mask(num_rows, hidden.device)
+                    self._reset_decode_state_rows(pfmask, num_rows, hidden.device)
+                self._fast_audio_direct_rows = num_rows
+            sampled = run_dense_sample(self, hidden, logits, sampling_metadata, force_audio_inputs=graph_forced_audio)
+            if sampled is not None:
+                return sampled
+            self._last_logits_hidden = None
         if not decode_only:
             self._fast_audio_direct_rows = 0
             pfmask = self._prefill_row_mask(num_rows, hidden.device)
@@ -1042,6 +1236,17 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
 
         if fast_fallback_reason is None:
             direct_audio_batch = decode_only and self._fast_audio_direct_rows == num_rows
+            mixed_direct = (
+                getattr(self, "_step_audio_tail_rows", 0) == num_rows
+                or getattr(self, "_step_audio_mode_rows", 0) == num_rows
+            ) and getattr(sampling_metadata, "no_penalties", False)
+            if mixed_direct:
+                # Guard the CPU classification against runner/device drift.
+                torch._assert_async(
+                    (audio_mask if getattr(self, "_step_audio_mode_rows", 0) == num_rows else prev_audio_mask).all(),
+                    "CPU/GPU Higgs audio tail mismatch",
+                )
+                direct_audio_batch = True
 
             if direct_audio_batch:
                 self._fast_audio_direct_rows = num_rows
@@ -1194,6 +1399,62 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
             self._boc_frame_cache[key] = buf
         return buf[:num_rows]
 
+    def post_sample_multimodal_outputs(self, *, req_ids, invalid_req_indices, multimodal_outputs):
+        """Export current-step codes by current batch row, with owned CPU storage.
+
+        Forward-time payloads refer to previous request-buffer rows, which may
+        be reordered or overwritten when the sampler reuses its host staging.
+        Export only after sampling, before the runner takes its output snapshot.
+        """
+        if getattr(self, "use_async_omni_output", False):
+            staging = self._async_audio_gpu_staging
+            self._async_audio_gpu_staging = None
+            self._fast_audio_direct_rows = 0
+            if staging is None:
+                return None
+            if staging.shape[0] != len(req_ids):
+                raise RuntimeError("Higgs async snapshot row mismatch")
+            return {"_higgs_audio_snapshot": staging, "_higgs_invalid_rows": tuple(invalid_req_indices)}
+        staging = self._last_audio_host_staging
+        if staging is None:
+            self._fast_audio_direct_rows = 0
+            return None
+        if staging.shape[0] != len(req_ids):
+            raise RuntimeError("Higgs audio staging rows do not match the sampled request batch")
+        event = self._last_audio_staging_event
+        if event is not None:
+            event.synchronize()
+            self._last_audio_staging_event = None
+        codes = staging[:, : self.num_codebooks].clone()
+        valid = staging[:, self.num_codebooks].tolist()
+        done = staging[:, self.num_codebooks + 1].tolist()
+        invalid = set(invalid_req_indices)
+        self._fast_audio_direct_rows = len(req_ids) if not invalid and all(v or d for v, d in zip(valid, done)) else 0
+        # Keep the codec-row rank on terminal/invalid steps: a 1-D empty
+        # value would replace previously accumulated full-response rows.
+        empty = codes.new_empty((0, self.num_codebooks))
+        rows = [codes[i : i + 1] if valid[i] and i not in invalid else empty for i in range(len(req_ids))]
+        return {"codes": {"audio": rows}}
+
+    def finalize_multimodal_outputs_from_cpu_snapshot(self, payload):
+        """Pure per-step conversion; never reads mutable decode/request state."""
+        if not payload or "_higgs_audio_snapshot" not in payload:
+            return payload
+        staging = payload["_higgs_audio_snapshot"]
+        if staging.device.type != "cpu":
+            raise RuntimeError("Higgs snapshot finalized before D2H")
+        codes = staging[:, : self.num_codebooks].clone()
+        valid = staging[:, self.num_codebooks].tolist()
+        invalid = set(payload["_higgs_invalid_rows"])
+        # Keep the codec-row rank on terminal/invalid steps: a 1-D empty
+        # value would replace previously accumulated full-response rows.
+        empty = codes.new_empty((0, self.num_codebooks))
+        return {
+            "codes": {
+                "audio": [codes[i : i + 1] if valid[i] and i not in invalid else empty for i in range(staging.shape[0])]
+            }
+        }
+
     # ------------------------------------------------------------------ postprocess
     def postprocess(
         self,
@@ -1450,6 +1711,20 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
             staging[:, num_codebooks].copy_(valid_full.to(torch.int32))
             staging[:, num_codebooks + 1].copy_(done_full.to(torch.int32))
 
+        self._finish_audio_staging(staging, num_rows, device, num_rows if all_rows else int(rows.numel()))
+
+    def _finish_audio_staging(self, staging, num_rows, device, active_rows):
+        if getattr(self, "_in_audio_sample_graph", False):
+            return
+        if getattr(self, "use_async_omni_output", False):
+            # Own this step's GPU values before the next graph replay. The
+            # runner copies the snapshot on its output stream and waits there.
+            self._async_audio_gpu_staging = staging[:num_rows].clone()
+            self._last_audio_host_staging = None
+            self._last_audio_staging_event = None
+            self._last_audio_codes = None
+            self._decode_active_audio_count = active_rows
+            return
         host = self._get_audio_host_staging_buffer(num_rows)
         host.copy_(staging, non_blocking=True)
         if device.type == "cuda":
@@ -1459,7 +1734,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         else:
             self._last_audio_staging_event = None
 
-        self._decode_active_audio_count = num_rows if all_rows else int(rows.numel())
+        self._decode_active_audio_count = active_rows
         # Host staging is the authoritative postprocess path. Keeping another
         # full GPU codes buffer alive adds fill/copy work on the hot path.
         self._last_audio_codes = None
@@ -1514,7 +1789,24 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         fallback = x.argmax(dim=-1)
         safe_x = torch.where(all_masked.unsqueeze(-1), torch.zeros_like(x), x)
         probs = safe_x.softmax(dim=-1)
-        if sampling_metadata.generators:
+        if getattr(self, "_audio_graph_noise", None) is not None:
+            torch._assert_async(torch.isfinite(probs).all(), "Audio sampling probabilities must be finite")
+            sampled = probs.div(self._audio_graph_noise).argmax(dim=-1)
+        elif probs.is_cuda:
+            # Match torch.multinomial's single-sample exponential race, but
+            # validate the whole batch asynchronously. The original per-request
+            # calls each copy validation scalars to the host before drawing.
+            torch._assert_async(torch.isfinite(probs).all(), "Audio sampling probabilities must be finite")
+            q = torch.empty_like(probs)
+            if sampling_metadata.generators:
+                from vllm_omni.utils.seeded_exponential import fill_exponential_rows
+
+                rows = int(probs.shape[0]) // num_codebooks
+                fill_exponential_rows(q, [sampling_metadata.generators.get(i) for i in range(rows)])
+            else:
+                q.exponential_()
+            sampled = probs.div(q).argmax(dim=-1)
+        elif sampling_metadata.generators:
             sampled = torch.cat(
                 [
                     torch.multinomial(

vllm_omni/model_executor/models/higgs_audio_v3/model_state.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Higgs delay-pattern sampling and owned audio output on Model Runner V2."""
+
+from typing import Any
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+from vllm.v1.sample.logits_processor import LogitsProcessors
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+from vllm_omni.worker_v2.omni_sampler import OmniSampler, OmniSamplingOutput
+
+logger = init_logger(__name__)
+
+
+class HiggsModelState(OmniModelState):
+    def __init__(self, vllm_config, model, encoder_cache, device):
+        super().__init__(vllm_config, model, encoder_cache, device)
+        model._use_external_decode_cudagraph = True
+        # Fix decode-state addresses before any sampling graph can capture them.
+        model._ensure_decode_state_capacity(max(self.max_num_tokens, 128), device)
+        logger.info("Higgs MRV2 custom sampler and owned audio snapshots enabled")
+        if getattr(model.config, "audio_mrv2_static_inputs", False):
+            self._static_inputs_embeds = torch.zeros(
+                (self.max_num_tokens, model.model.embed_tokens.weight.shape[1]),
+                dtype=self.dtype,
+                device=device,
+            )
+        self.direct_payload = bool(getattr(model.config, "audio_mrv2_direct_payload", False))
+        self.sampling_slots = None
+        self.request_slots = {}
+        if getattr(model.config, "audio_mrv2_slot_parameters", False):
+            capacity = self.scheduler_config.max_num_seqs
+            self.sampling_slots = {
+                "temperature": torch.zeros(capacity, device=device),
+                "top_p": torch.ones(capacity, device=device),
+                "top_k": torch.full((capacity,), -1, dtype=torch.long, device=device),
+                "zeros": torch.zeros(capacity, device=device),
+                "ones": torch.ones(capacity, device=device),
+            }
+        self.requests = {}
+        self.generators = {}
+        self._ensure_slot_capacity(self.scheduler_config.max_num_seqs)
+        self.metadata_key = None
+        self.metadata = None
+        logger.info(
+            "Higgs MRV2 options: static=%s slot_parameters=%s direct_payload=%s",
+            self._static_inputs_embeds is not None,
+            self.sampling_slots is not None,
+            self.direct_payload,
+        )
+        if not model.use_async_omni_output:
+            raise ValueError("Higgs MRV2 requires audio_async_payload; use higgs_multimodal_qwen3_mrv2_h200.yaml")
+
+    def _ensure_slot_capacity(self, slots):
+        """Per-slot prompt facts for vectorized step classification."""
+        current = getattr(self, "prompt_len_np", None)
+        if current is not None and len(current) >= slots:
+            return
+        size = max(slots, 2 * (0 if current is None else len(current)))
+        old_len, old_tail = current, getattr(self, "prompt_audio_tail_np", None)
+        self.prompt_len_np = np.zeros(size, dtype=np.int64)
+        self.prompt_audio_tail_np = np.zeros(size, dtype=bool)
+        if current is not None:
+            self.prompt_len_np[: len(old_len)] = old_len
+            self.prompt_audio_tail_np[: len(old_tail)] = old_tail
+
+    def add_request(self, req_index, new_req_data):
+        params = new_req_data.sampling_params
+        is_warmup = str(new_req_data.req_id).startswith("_warmup_")
+        if not is_warmup and (
+            params.frequency_penalty != 0
+            or params.presence_penalty != 0
+            or params.repetition_penalty != 1
+            or params.logprobs is not None
+            or params.prompt_logprobs is not None
+            or params.min_tokens != 0
+            or params.allowed_token_ids
+            or params.bad_words
+            or params.logit_bias
+            or params.min_p != 0
+        ):
+            raise ValueError("Higgs MRV2 currently supports temperature/top-k/top-p sampling without text penalties")
+        super().add_request(req_index, new_req_data)
+        self.requests[new_req_data.req_id] = new_req_data
+        self.request_slots[new_req_data.req_id] = req_index
+        if self.sampling_slots is not None:
+            for name in ("temperature", "top_p", "top_k"):
+                self.sampling_slots[name][req_index : req_index + 1].fill_(getattr(params, name))
+        self.metadata_key = None
+        prompt = new_req_data.prompt_token_ids or []
+        self._ensure_slot_capacity(req_index + 1)
+        self.prompt_len_np[req_index] = len(prompt)
+        self.model._resolve_token_ids()
+        self.prompt_audio_tail_np[req_index] = bool(prompt) and prompt[-1] == self.model._audio_continuation_id
+        if params.seed is not None:
+            self.generators[new_req_data.req_id] = torch.Generator(device=self.device).manual_seed(params.seed)
+
+    def remove_request(self, req_index_or_id):
+        index = self._resolve_req_index(req_index_or_id)
+        if index is not None:
+            req_id = self.intermediate_buffer.buffers[index].get("req_id")
+            if req_id is not None:
+                self.model.on_requests_finished({req_id})
+                self.requests.pop(req_id, None)
+                self.request_slots.pop(req_id, None)
+                self.generators.pop(req_id, None)
+                self.metadata_key = None
+        super().remove_request(req_index_or_id)
+
+    def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_descriptor_dispatcher=None):
+        model = self.model
+        model._resolve_token_ids()
+        # Every row must have an audio-continuation prompt whose prefill ends
+        # within this step; one vectorized check replaces a per-row loop.
+        n = input_batch.num_reqs
+        slots = input_batch.idx_mapping_np[:n]
+        end = req_states.num_computed_prefill_tokens[slots] + input_batch.num_scheduled_tokens[:n]
+        audio_mode = bool(np.all(self.prompt_audio_tail_np[slots] & (end >= self.prompt_len_np[slots])))
+        model.update_decode_step_metadata(
+            input_ids=model_inputs["input_ids"][: input_batch.num_tokens],
+            positions=input_batch.positions,
+            omni_query_start_loc=input_batch.query_start_loc[: input_batch.num_reqs + 1],
+            req_ids=input_batch.req_ids,
+            audio_prompt_mode_rows=input_batch.num_reqs if audio_mode else 0,
+        )
+
+        if self._static_inputs_embeds is not None:
+            self._prepare_static_embeddings(input_batch, model_inputs)
+
+    def _prepare_static_embeddings(self, input_batch, model_inputs):
+        # Feedback and reference substitution are request-dependent and belong
+        # outside the captured backbone. Use only actual tokens: graph padding
+        # must not change the one-token-per-request decode classification.
+        n = input_batch.num_tokens
+        ids = model_inputs["input_ids"][:n]
+        if self._fused_decode_embeddings(input_batch, ids):
+            return
+        safe_ids = torch.where(ids < 0, torch.zeros_like(ids), ids)
+        embeds = self.model.model.embed_tokens(safe_ids)
+        info = model_inputs.get("model_intermediate_buffer")
+        if input_batch.has_prefill and info:
+            embeds = self.model._apply_ref_audio_substitution(
+                embeds,
+                ids,
+                input_batch.positions[:n],
+                info,
+            )
+        embeds = self._apply_audio_feedback(embeds, ids, input_batch)
+        self._static_inputs_embeds[:n].copy_(embeds)
+        padded = input_batch.num_tokens_after_padding
+        self._static_inputs_embeds[n:padded].zero_()
+
+    def _fused_decode_embeddings(self, input_batch, ids) -> bool:
+        """Pure decode steps: token/feedback embeddings and padding in one launch."""
+        model = self.model
+        if (
+            input_batch.has_prefill
+            or input_batch.num_tokens != input_batch.num_reqs
+            or not ids.is_cuda
+            or not getattr(model, "_use_external_decode_cudagraph", False)
+        ):
+            return False
+        text = getattr(model.model.embed_tokens, "weight", None)
+        audio = getattr(model.multimodal_embedding, "weight", None)
+        if text is None or audio is None or text.dtype != self._static_inputs_embeds.dtype:
+            return False
+        from .step_kernels import decode_embeddings
+
+        n = int(ids.numel())
+        model._ensure_decode_state_capacity(n, ids.device)
+        decode_embeddings(
+            self._static_inputs_embeds,
+            ids,
+            text,
+            audio,
+            model._decode_last_codes,
+            model._decode_has_codes,
+            int(input_batch.num_tokens_after_padding),
+            int(model.multimodal_embedding.vocab_size),
+        )
+        return True
+
+    def _apply_audio_feedback(self, embeds, ids, input_batch):
+        if input_batch.num_tokens == input_batch.num_reqs:
+            return self.model._apply_audio_feedback(embeds, ids)
+        # Every scheduled request has a nonempty span. Gather its first token
+        # for all rows, then mask on-device. Boolean indexing would first ask
+        # the GPU how many rows matched, synchronizing every mixed prefill.
+        n = input_batch.num_reqs
+        qsl = input_batch.query_start_loc[: n + 1]
+        starts = qsl[:-1].long()
+        model = self.model
+        model._ensure_decode_state_capacity(n, embeds.device)
+        audio = model.multimodal_embedding(model._decode_last_codes[:n]).to(embeds.dtype)
+        use_audio = model._decode_has_codes[:n] & ((qsl[1:] - qsl[:-1]) == 1)
+        current = embeds.index_select(0, starts)
+        result = embeds.clone()
+        result.index_copy_(0, starts, torch.where(use_audio[:, None], audio, current))
+        return result
+
+    def run_postprocess(self, hidden_states, input_batch):
+        self.model.finish_decode_step_forward()
+
+    def postprocess_model_output(self, model_output, input_batch, req_states):
+        return model_output, None
+
+    def sampling_metadata(self, req_ids, slots=None):
+        key = tuple(req_ids)
+        if self.metadata_key != key:
+            params = [self.requests[r].sampling_params for r in req_ids]
+            n = len(params)
+            if self.sampling_slots is not None:
+                if slots is None:
+                    slots = torch.tensor([self.request_slots[r] for r in req_ids], device=self.device)
+                temperature, top_p, top_k = (
+                    self.sampling_slots[name].index_select(0, slots) for name in ("temperature", "top_p", "top_k")
+                )
+                zeros = self.sampling_slots["zeros"][:n]
+                ones = self.sampling_slots["ones"][:n]
+            else:
+                temperature = torch.tensor([p.temperature for p in params], device=self.device)
+                top_p = torch.tensor([p.top_p for p in params], device=self.device)
+                top_k = torch.tensor([p.top_k for p in params], device=self.device)
+                zeros = torch.zeros(n, device=self.device)
+                ones = torch.ones(n, device=self.device)
+            self.metadata = SamplingMetadata(
+                temperature=temperature,
+                all_greedy=all(p.temperature == 0 for p in params),
+                all_random=all(p.temperature != 0 for p in params),
+                top_p=top_p,
+                top_k=top_k,
+                generators={i: self.generators[r] for i, r in enumerate(req_ids) if r in self.generators},
+                max_num_logprobs=None,
+                no_penalties=True,
+                prompt_token_ids=None,
+                frequency_penalties=zeros,
+                presence_penalties=zeros,
+                repetition_penalties=ones,
+                output_token_ids=[[] for _ in params],
+                allowed_token_ids_mask=None,
+                bad_words_token_ids={},
+                logitsprocs=LogitsProcessors(),
+            )
+            self.metadata_key = key
+        return self.metadata
+
+    def custom_sampler(self, sampler):
+        return HiggsSampler(sampler, self), None
+
+    def finalize_audio_snapshot(self, payload: dict[str, Any], num_sampled: list[int]):
+        if payload and "_higgs_audio_snapshot" in payload and getattr(self, "direct_payload", False):
+            from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot
+
+            snapshot = payload["_higgs_audio_snapshot"]
+            if snapshot.device.type != "cpu":
+                raise RuntimeError("Higgs snapshot finalized before D2H")
+            codes = snapshot[:, : self.model.num_codebooks].clone()
+            valid = snapshot[:, self.model.num_codebooks].tolist()
+            invalid = set(payload.get("_higgs_invalid_rows", ()))
+            # One split creates every row view; per-row slicing cost one
+            # dispatcher call per request on the output-materialization path.
+            rows = torch.split(codes, 1)
+            return RequestOutputSnapshot(
+                [
+                    {"codes.audio": rows[i]} if valid[i] and n and i not in invalid else None
+                    for i, n in enumerate(num_sampled)
+                ]
+            )
+        if payload and "_higgs_audio_snapshot" in payload:
+            payload = dict(payload)
+            payload["_higgs_invalid_rows"] = tuple(i for i, n in enumerate(num_sampled) if n == 0)
+        return self.model.finalize_multimodal_outputs_from_cpu_snapshot(payload)
+
+
+class HiggsSampler(OmniSampler):
+    """Higgs audio sampling registered through the upstream sampler factory."""
+
+    def __init__(self, base_sampler, state: HiggsModelState):
+        super().__init__(base_sampler)
+        self.state = state
+
+    def sample_step(self, hidden_states, input_batch, req_states, grammar_output, standard_sample):
+        # Upstream warmup exercises its text sampler, including grammar and
+        # penalties, with synthetic non-audio prompts. Keep that path intact.
+        if all(str(r).startswith("_warmup_") for r in input_batch.req_ids):
+            return super().sample_step(hidden_states, input_batch, req_states, grammar_output, standard_sample)
+        if grammar_output is not None or input_batch.num_draft_tokens:
+            raise ValueError("Higgs MRV2 does not support grammar or speculative sampling")
+        metadata = self.state.sampling_metadata(input_batch.req_ids, input_batch.idx_mapping[: input_batch.num_reqs])
+        hidden = hidden_states[input_batch.logits_indices]
+        logits = self.state.model.compute_logits(hidden, metadata)
+        sampled = self.state.model.sample(logits, metadata)
+        count, rejected = get_num_sampled_and_rejected(
+            torch.ones(input_batch.num_reqs, dtype=torch.int32, device=hidden.device),
+            input_batch.seq_lens,
+            input_batch.cu_num_logits,
+            input_batch.idx_mapping,
+            req_states.prefill_len.gpu,
+        )
+        output = SamplerOutput(sampled.sampled_token_ids, sampled.logprobs_tensors, None, count, rejected)
+        payload = self.state.model.post_sample_multimodal_outputs(
+            req_ids=input_batch.req_ids,
+            invalid_req_indices=[],
+            multimodal_outputs=None,
+        )
+        return OmniSamplingOutput(
+            output,
+            count,
+            rejected,
+            payload,
+            include_hidden_states=False,
+            owns_multimodal_outputs=True,
+            finalize_multimodal=self.state.finalize_audio_snapshot,
+        )

vllm_omni/model_executor/models/higgs_audio_v3/mrv2_state_kernels.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused movement of model-owned delay state when a request batch changes."""
+
+import torch
+from vllm.triton_utils import tl
+from vllm.triton_utils import triton as tr
+
+
+@tr.jit(do_not_specialize=["count"])
+def _move(
+    indices,
+    source_codes,
+    source_has,
+    source_delay,
+    source_eoc,
+    source_done,
+    dest_codes,
+    dest_has,
+    dest_delay,
+    dest_eoc,
+    dest_done,
+    count,
+    books: tl.constexpr,
+    gather: tl.constexpr,
+    block: tl.constexpr,
+):
+    row = tl.program_id(0)
+    col = tl.arange(0, block)
+    if gather:
+        encoded = tl.load(indices + row)
+        fresh = encoded < 0
+        src = tl.where(fresh, -encoded - 1, encoded)
+        dst = row
+    else:
+        src = tl.load(indices + row)
+        dst = tl.load(indices + count + row)
+        fresh = False
+    codes = tl.load(source_codes + src * books + col, col < books, 0)
+    has = tl.load(source_has + src)
+    delay = tl.load(source_delay + src)
+    eoc = tl.load(source_eoc + src)
+    done = tl.load(source_done + src)
+    tl.store(dest_codes + dst * books + col, tl.where(fresh, 0, codes), col < books)
+    tl.store(dest_has + dst, tl.where(fresh, 0, has))
+    tl.store(dest_delay + dst, tl.where(fresh, 0, delay))
+    tl.store(dest_eoc + dst, tl.where(fresh, -1, eoc))
+    tl.store(dest_done + dst, tl.where(fresh, 0, done))
+
+
+def sync_state(model, previous_rows, pool_rows, current_rows):
+    # One owned pinned upload, ordered on the decode stream. Separate scatter
+    # and gather kernels provide a global ordering point for moved pool rows.
+    device = model._decode_has_codes.device
+    packed = torch.tensor(previous_rows + pool_rows + current_rows, dtype=torch.int64, pin_memory=True).to(
+        device, non_blocking=True
+    )
+    names = ("last_codes", "has_codes", "delay_count", "eoc_countdown", "generation_done")
+    decode = [getattr(model, "_decode_" + name) for name in names]
+    pool = [getattr(model, "_state_pool_" + name) for name in names]
+    n = len(previous_rows)
+    block = tr.next_power_of_2(model.num_codebooks)
+    if n:
+        _move[(n,)](packed, *decode, *pool, n, model.num_codebooks, False, block)
+    if current_rows:
+        _move[(len(current_rows),)](
+            packed[2 * n :], *pool, *decode, len(current_rows), model.num_codebooks, True, block
+        )

vllm_omni/model_executor/models/higgs_audio_v3/pipeline.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
 """higgs-audio v3 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM).
 
 Two delivery modes are wired here:
@@ -19,6 +19,9 @@ from vllm_omni.config.stage_config import (
     StageExecutionType,
     StagePipelineConfig,
 )
+from vllm_omni.outputs.output_modality import TensorAccumulationStrategy, register_key_accumulation_strategy
+
+register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0)
 
 _PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v3"
 
@@ -34,6 +37,8 @@ HIGGS_AUDIO_V3_PIPELINE = PipelineConfig(
             execution_type=StageExecutionType.LLM_AR,
             input_sources=(),
             owns_tokenizer=True,
+            supports_native_mrv2_data_plane=True,
+            custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
             engine_output_type="latent",
             # stop_token_ids: the model-owned sampler forces eos at ramp-down
             # completion. Safety stops from the actual V3 checkpoint:
@@ -54,7 +59,9 @@ HIGGS_AUDIO_V3_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="audio",
             model_arch="HiggsAudioV3Code2WavForConditionalGeneration",
-            sync_process_input_func=f"{_PROC}.talker2code2wav",
+            sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
+            supports_native_mrv2_data_plane=True,
+            requires_full_payload_input=True,
             sampling_constraints={"detokenize": True},
         ),
     ),

vllm_omni/model_executor/models/higgs_audio_v3/step_kernels.py

@@ -0,0 +1,150 @@
+# ruff: noqa: N803
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused per-step Higgs kernels for the MRV2 decode loop.
+
+Each replaces a chain of small eager launches that ran on every model step
+(terminal-row restore, prefill reset and audio-tail validation before dense
+sampling; token and feedback embeddings on pure decode steps).
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _step_tail_guard_kernel(
+    ids_ptr,
+    qsl_ptr,
+    has_codes_ptr,
+    done_ptr,
+    delay_ptr,
+    eoc_ptr,
+    error_ptr,
+    num_rows,
+    audio_id,
+    eos_id,
+    HAS_QSL: tl.constexpr,
+    PROMPT_MODE: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    rows = tl.arange(0, BLOCK)
+    mask = rows < num_rows
+    if HAS_QSL:
+        start = tl.load(qsl_ptr + rows, mask=mask, other=0)
+        end = tl.load(qsl_ptr + rows + 1, mask=mask, other=1)
+        valid = end > 0
+        tail = tl.load(ids_ptr + tl.maximum(end - 1, 0), mask=mask, other=audio_id)
+        prefill = mask & ((end - start) > 1)
+    else:
+        valid = mask
+        tail = tl.load(ids_ptr + rows, mask=mask, other=audio_id)
+        prefill = rows < 0
+    done = tl.load(done_ptr + rows, mask=mask, other=0) != 0
+    has = tl.load(has_codes_ptr + rows, mask=mask, other=0) != 0
+    if PROMPT_MODE:
+        # An in-flight EOS row after retirement is terminal, never a new stream.
+        done = done | (tail == eos_id)
+    # New prefill rows start a fresh delay-pattern stream.
+    done = done & ~prefill
+    has = has & ~prefill
+    tl.store(done_ptr + rows, done, mask=mask)
+    tl.store(has_codes_ptr + rows, has, mask=mask)
+    tl.store(delay_ptr + rows, 0, mask=prefill)
+    tl.store(eoc_ptr + rows, -1, mask=prefill)
+    ok = (tail == audio_id) & valid
+    if PROMPT_MODE:
+        ok = ok | has | done
+    tl.store(error_ptr, tl.max((mask & ~ok).to(tl.int32), axis=0))
+
+
+def step_tail_guard(
+    ids: torch.Tensor,
+    qsl: torch.Tensor | None,
+    has_codes: torch.Tensor,
+    done: torch.Tensor,
+    delay: torch.Tensor,
+    eoc: torch.Tensor,
+    error: torch.Tensor,
+    *,
+    audio_id: int,
+    eos_id: int,
+    prompt_mode: bool,
+) -> None:
+    """Fused terminal-row restore, prefill reset and audio-tail validation.
+
+    One launch replaces the index/compare/reset chains that ran before every
+    dense sample, for pure decode and mixed prefill steps alike; ``error``
+    becomes nonzero if any row is not an audio row. ``qsl`` (query start
+    locations) selects each row's last token; ``None`` means one token per row.
+    """
+    rows = int(has_codes.numel())
+    _step_tail_guard_kernel[(1,)](
+        ids,
+        ids if qsl is None else qsl,
+        has_codes,
+        done,
+        delay,
+        eoc,
+        error,
+        rows,
+        audio_id,
+        eos_id,
+        HAS_QSL=qsl is not None,
+        PROMPT_MODE=prompt_mode,
+        BLOCK=max(16, triton.next_power_of_2(rows)),
+    )
+
+
+@triton.jit
+def _decode_embeddings_kernel(
+    out_ptr,
+    ids_ptr,
+    text_weight_ptr,
+    audio_weight_ptr,
+    codes_ptr,
+    has_codes_ptr,
+    num_rows,
+    hidden,
+    audio_vocab,
+    BOOKS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    mask = cols < hidden
+    out = out_ptr + row.to(tl.int64) * hidden + cols
+    if row >= num_rows:
+        # Graph padding rows own no request and must stay zero.
+        tl.store(out, tl.zeros((BLOCK,), dtype=out_ptr.dtype.element_ty), mask=mask)
+        return
+    if tl.load(has_codes_ptr + row) != 0:
+        # Audio feedback: sum of the previous frame's codebook embeddings,
+        # accumulated in float32 and rounded once like ``embedding(...).sum()``.
+        acc = tl.zeros((BLOCK,), dtype=tl.float32)
+        for book in range(BOOKS):
+            code = tl.load(codes_ptr + row * BOOKS + book).to(tl.int64) + book * audio_vocab
+            acc += tl.load(audio_weight_ptr + code * hidden + cols, mask=mask, other=0.0).to(tl.float32)
+        tl.store(out, acc.to(out_ptr.dtype.element_ty), mask=mask)
+    else:
+        token = tl.maximum(tl.load(ids_ptr + row), 0).to(tl.int64)
+        tl.store(out, tl.load(text_weight_ptr + token * hidden + cols, mask=mask), mask=mask)
+
+
+def decode_embeddings(
+    out: torch.Tensor,
+    ids: torch.Tensor,
+    text_weight: torch.Tensor,
+    audio_weight: torch.Tensor,
+    codes: torch.Tensor,
+    has_codes: torch.Tensor,
+    padded_rows: int,
+    audio_vocab: int,
+) -> None:
+    """One-token-per-row input embeddings with audio feedback, into ``out[:padded_rows]``."""
+    rows, hidden = int(ids.numel()), int(out.shape[1])
+    books = int(codes.shape[1])
+    block = 1024
+    _decode_embeddings_kernel[(padded_rows, triton.cdiv(hidden, block))](
+        out, ids, text_weight, audio_weight, codes, has_codes, rows, hidden, audio_vocab, BOOKS=books, BLOCK=block
+    )

vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.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
 """Stage-input processor for higgs-audio v3: Talker -> Code2Wav.
 
 Two adapters:
@@ -21,6 +21,7 @@ to avoid corrupting valid tail content.
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -95,6 +96,18 @@ def _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor:
     return frames[valid].t().contiguous()
 
 
+def _whole_utterance_codec_codes(audio_rows: torch.Tensor) -> torch.Tensor:
+    """Convert complete delayed [T, Q] rows to flat codebook-major CPU codes."""
+    codes = _revert_delay_pattern(audio_rows.to(device="cpu", dtype=torch.long).t().contiguous())
+    # Substitute special/padding codes after de-delay. Clamping would map EOC
+    # to a valid codec entry and introduce audible artifacts.
+    codes = torch.where((codes < 0) | (codes >= _NUM_REAL_CODES), 0, codes)
+    # Preserve the existing removal of the residual ramp-down frame.
+    if codes.shape[-1] >= 2:
+        codes = codes[:, :-1]
+    return codes.reshape(-1).contiguous()
+
+
 def talker2code2wav(
     source_outputs: list[Any],
     prompt: Any = None,
@@ -133,43 +146,16 @@ def talker2code2wav(
                 f"got {audio_codes.shape[1]}. Audio codes shape: {tuple(audio_codes.shape)}"
             )
 
-        # Transpose to [Q, T] for delay pattern reversal
-        codes_qt = audio_codes.transpose(0, 1).contiguous().cpu()
-
-        # Step 1: Revert delay pattern
         try:
-            codes_qt = _revert_delay_pattern(codes_qt)
+            codec_codes = _whole_utterance_codec_codes(audio_codes)
         except ValueError as exc:
             logger.warning("Skipping invalid Higgs Audio v3 code sequence for Stage 1: %s", exc)
             code2wav_inputs.append(_empty_code2wav_prompt())
             continue
 
-        # Step 2: Replace out-of-range codes (BOC=1024, EOC=1025, -1) with 0.
-        # Must use torch.where, NOT clamp: clamp(max=1023) turns 1025→1023
-        # which is a valid codec code and decodes to audio artifacts.
-        # Matches sglang's: torch.where(codes >= codec_vocab, 0, codes)
-        codes_qt = torch.where(
-            (codes_qt >= _NUM_REAL_CODES) | (codes_qt < 0),
-            torch.zeros_like(codes_qt),
-            codes_qt,
-        )
-
-        # Step 3: Trim the last frame. After de-delay, the final frame
-        # contains residual ramp-down codes (EOC→0 substituted) that
-        # decode to a brief noise artifact at the end of the audio.
-        if codes_qt.shape[-1] >= 2:
-            codes_qt = codes_qt[:, :-1]
-
-        if codes_qt.numel() == 0:
-            code2wav_inputs.append(_empty_code2wav_prompt())
-            continue
-
-        # Code2Wav expects codebook-major flat: [Q * num_frames]
-        codec_codes = codes_qt.reshape(-1).tolist()
-
         code2wav_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
+                prompt_token_ids=codec_codes.tolist(),
                 multi_modal_data=None,
                 mm_processor_kwargs=None,
                 additional_information=None,
@@ -261,6 +247,13 @@ def talker2code2wav_async_chunk(
     elif not finished:
         return None
 
+    return _flush_async_chunk(
+        transfer_manager, request_id, finished, _async_chunk_config(transfer_manager), emitted_frames
+    )
+
+
+def _async_chunk_config(transfer_manager: Any) -> tuple[int, int, int, int]:
+    """Validated (chunk, left context, right holdback, initial chunk) frames."""
     connector = getattr(transfer_manager, "connector", None)
     raw_cfg = getattr(connector, "config", {}) or {}
     cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
@@ -281,7 +274,18 @@ def talker2code2wav_async_chunk(
             f"codec_right_holdback_frames={right_holdback_size_config}, "
             f"initial_codec_chunk_frames={configured_initial_chunk_size}"
         )
+    return chunk_size, left_context_size_config, right_holdback_size_config, configured_initial_chunk_size
 
+
+def _flush_async_chunk(
+    transfer_manager: Any,
+    request_id: str,
+    finished: bool,
+    config: tuple[int, int, int, int],
+    emitted_frames: dict[str, int],
+) -> OmniPayloadStruct | None:
+    """Emit the next sliding-window codec chunk once enough rows accumulated."""
+    chunk_size, left_context_size_config, right_holdback_size_config, configured_initial_chunk_size = config
     n_rows = len(transfer_manager.code_prompt_token_ids[request_id])
     de_delayed_total = max(0, n_rows - (_NUM_CODEBOOKS - 1))
     emitted = int(emitted_frames.get(request_id, 0))
@@ -385,3 +389,83 @@ def talker2code2wav_async_chunk(
         codes=CodesStruct(audio=codec_codes),
         meta=meta,
     )
+
+
+def talker2code2wav_async_chunk_batch(
+    transfer_manager: Any,
+    pooling_outputs: list[OmniPayload | None],
+    requests: list[Any],
+    is_finished: list[bool],
+) -> list[OmniPayloadStruct | None]:
+    """One talker step of async-chunk payloads with one host conversion per row.
+
+    Same accumulation and flush semantics as ``talker2code2wav_async_chunk``;
+    the chunk config is resolved once per step and each emitted ``[1, Q]``
+    CPU row is converted with a single ``tolist`` instead of several small
+    tensor operations per request.
+    """
+    if not (len(pooling_outputs) == len(requests) == len(is_finished)):
+        raise ValueError("batch codec inputs must have identical lengths")
+    emitted_frames = getattr(transfer_manager, "higgs_v3_emitted_frames", None)
+    if emitted_frames is None:
+        emitted_frames = {}
+        transfer_manager.higgs_v3_emitted_frames = emitted_frames
+    config = None
+    rows_by_request = transfer_manager.code_prompt_token_ids
+    payloads: list[OmniPayloadStruct | None] = []
+    for pooling_output, request, finished_flag in zip(pooling_outputs, requests, is_finished):
+        request_id = request.external_req_id
+        finished = bool(finished_flag or request.is_finished())
+        if isinstance(pooling_output, dict):
+            codes = pooling_output.get("codes")
+            audio = codes.get("audio") if isinstance(codes, dict) else None
+            if isinstance(audio, torch.Tensor) and audio.numel():
+                if audio.ndim == 2 and audio.device.type == "cpu" and not audio.is_floating_point():
+                    row = audio.tolist()[-1]
+                    if len(row) != _NUM_CODEBOOKS:
+                        raise ValueError(
+                            f"talker emit row has {len(row)} codebooks; expected {_NUM_CODEBOOKS} "
+                            "for higgs_audio_v3 async_chunk."
+                        )
+                else:
+                    row = _extract_last_step_row(pooling_output)
+                    row = None if row is None else row.cpu().tolist()
+                if row is not None:
+                    rows_by_request[request_id].append(row)
+        elif not finished:
+            payloads.append(None)
+            continue
+        if config is None:
+            config = _async_chunk_config(transfer_manager)
+        payloads.append(_flush_async_chunk(transfer_manager, request_id, finished, config, emitted_frames))
+    return payloads
+
+
+def talker2code2wav_token_only(source_outputs, prompt=None, _requires_multimodal_data=False):
+    """Control slots for the native full-payload consumer."""
+    result = []
+    for output in source_outputs:
+        if not output.finished:
+            continue
+        mm = getattr(output.outputs[0], "multimodal_output", None)
+        if isinstance(mm, Mapping) and isinstance(mm.get("codes", {}).get("audio"), torch.Tensor):
+            result.extend(talker2code2wav([output], prompt, _requires_multimodal_data))
+        else:
+            result.append(OmniTokensPrompt(prompt_token_ids=[0]))
+    return result
+
+
+def talker2code2wav_full_payload(transfer_manager, pooling_output, request):
+    """De-delay the accumulated raw rows once before native codec delivery."""
+    del transfer_manager, request
+    audio = pooling_output.get("codes.audio") if isinstance(pooling_output, dict) else None
+    if audio is None and isinstance(pooling_output, dict):
+        audio = pooling_output.get("codes", {}).get("audio")
+    codes = torch.empty(0, dtype=torch.long)
+    if isinstance(audio, torch.Tensor) and audio.numel():
+        if audio.numel() % _NUM_CODEBOOKS:
+            raise ValueError("Higgs payload must contain complete codebook rows")
+        rows = audio.reshape(-1, _NUM_CODEBOOKS).to(device="cpu", dtype=torch.long)
+        if rows.shape[0] >= _NUM_CODEBOOKS:
+            codes = _whole_utterance_codec_codes(rows)
+    return {"codes": {"audio": codes}, "meta": {"finished": torch.tensor(True)}}

vllm_omni/transformers_utils/configs/higgs_audio_v3.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
 """Configuration class for higgs-audio v3 (HiggsMultimodalQwen3) in vllm-omni.
 
 ``HiggsAudioV3Config.from_pretrained(model_path)`` returns a config with
@@ -95,6 +95,8 @@ class HiggsAudioV3Config(PretrainedConfig):
         audio_continuation_id: int | None = None,
         enable_flashinfer_api_unwrap: bool = True,
         enable_mlp_cudagraph: bool = True,
+        audio_full_sample_graph: bool | None = None,
+        codec_cuda_graph: bool | None = None,
         **kwargs: Any,
     ) -> None:
         # Legacy perf knob removed: Higgs v3 scheduler tokens now come from
@@ -131,9 +133,18 @@ class HiggsAudioV3Config(PretrainedConfig):
         self.audio_continuation_id = audio_continuation_id
         self.enable_flashinfer_api_unwrap = bool(enable_flashinfer_api_unwrap)
         self.enable_mlp_cudagraph = bool(enable_mlp_cudagraph)
+        self.audio_full_sample_graph = audio_full_sample_graph
+        self.codec_cuda_graph = codec_cuda_graph
 
         super().__init__(**kwargs)
 
+    def resolve_graph_defaults(self, *, use_v2_model_runner: bool) -> tuple[bool, bool]:
+        """Resolve sampler/codec graphs without mutating shared checkpoint config."""
+        return (
+            use_v2_model_runner if self.audio_full_sample_graph is None else self.audio_full_sample_graph,
+            use_v2_model_runner if self.codec_cuda_graph is None else self.codec_cuda_graph,
+        )
+
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs: Any) -> HiggsAudioV3Config:
         """Load config and resolve special token IDs from the checkpoint tokenizer.

vllm_omni/worker/base.py

@@ -13,6 +13,7 @@ from contextlib import AbstractContextManager, nullcontext
 import torch
 from vllm.logger import init_logger
 from vllm.utils.mem_utils import format_gib, memory_profiling
+from vllm.v1.worker.gpu_worker import CompilationTimes
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 
 from vllm_omni.diffusion.data import (
@@ -39,6 +40,17 @@ class OmniGPUWorkerBase(GPUWorker):
     for custom trace naming, background gzip, and trace path collection.
     """
 
+    def _capture_auxiliary_graphs(self) -> None:
+        """Let opt-in models warm valid inputs before the worker becomes ready."""
+        capture = getattr(self.model_runner.model, "capture_auxiliary_graphs", None)
+        if callable(capture):
+            capture()
+
+    def compile_or_warm_up_model(self) -> CompilationTimes:
+        result = super().compile_or_warm_up_model()
+        self._capture_auxiliary_graphs()
+        return result
+
     def load_model(self, *args, **kwargs):
         with self._maybe_get_memory_pool_context("weights"):
             res = super().load_model(*args, **kwargs)
@@ -225,7 +237,7 @@ class OmniGPUWorkerBase(GPUWorker):
         logger.info(f"[LLM Worker {self.rank}] Wake-up complete.")
         return True
 
-    def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK:
+    def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK | None:
         "Handle deterministic Sleep command from the main process"
         try:
             if isinstance(task, dict):
@@ -279,7 +291,7 @@ class OmniGPUWorkerBase(GPUWorker):
                     pass
             return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e))
 
-    def handle_wake_task(self, task: OmniWakeTask) -> OmniACK:
+    def handle_wake_task(self, task: OmniWakeTask) -> OmniACK | None:
         "Handle deterministic Wakeup command from the main process"
         try:
             if isinstance(task, dict):

vllm_omni/worker/gpu_ar_model_runner.py

@@ -1655,7 +1655,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         model_config = getattr(self, "model_config", None)
         if model_config is None:
             model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
-        if not bool(getattr(model_config, "async_chunk", False)):
+        if not bool(getattr(model_config, "async_chunk", False)) and not self._model_omni_flag(
+            getattr(self, "model", None), "supports_async_whole_payload"
+        ):
             return False
         if bool(getattr(model_config, "enable_return_routed_experts", False)):
             return False
@@ -2110,6 +2112,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 scheduler_output.total_num_scheduled_tokens,
             )
 
+        post_sample_mm = getattr(getattr(self, "model", None), "post_sample_multimodal_outputs", None)
+        if callable(post_sample_mm):
+            multimodal_outputs = post_sample_mm(
+                req_ids=req_ids_output_copy,
+                invalid_req_indices=invalid_req_indices,
+                multimodal_outputs=multimodal_outputs,
+            )
+
         multimodal_outputs = self._run_post_sample_talker_mtp(
             req_ids=req_ids_output_copy,
             valid_sampled_token_ids=valid_sampled_token_ids,
@@ -2184,12 +2194,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             if output_tensor_snapshot.async_payload is not None:
                 with record_function_or_nullcontext("omni_async_output:wait_cpu_payload"):
                     output_tensor_snapshot.async_payload.wait()
+            mm_snapshot = output_tensor_snapshot.multimodal_outputs
+            finalize_snapshot = getattr(
+                getattr(self, "model", None), "finalize_multimodal_outputs_from_cpu_snapshot", None
+            )
+            if callable(finalize_snapshot):
+                mm_snapshot = finalize_snapshot(mm_snapshot)
             with record_function_or_nullcontext("omni_output_builder:total"):
                 output = self._build_omni_model_runner_output_from_snapshot(
                     scheduler_output=scheduler_output_snapshot,
                     hidden_states=output_tensor_snapshot.hidden_states,
                     staged_hidden_states_cpu=output_tensor_snapshot.staged_hidden_states_cpu,
-                    multimodal_outputs=output_tensor_snapshot.multimodal_outputs,
+                    multimodal_outputs=mm_snapshot,
                     req_ids_output_copy=req_ids_output_snapshot,
                     req_id_to_index_output_copy=req_id_to_index_output_snapshot,
                     valid_sampled_token_ids=valid_sampled_token_ids_snapshot,

vllm_omni/worker/gpu_generation_worker.py

@@ -126,4 +126,5 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
         start = time.perf_counter()
         self.model_runner.profile_run()
+        self._capture_auxiliary_graphs()
         return CompilationTimes(language_model=time.perf_counter() - start, encoder=0.0)

vllm_omni/worker/gpu_model_runner.py

@@ -2148,7 +2148,19 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
         model_kwargs_extra = self._build_model_kwargs_extra()
         update_decode_metadata = getattr(self.model, "update_decode_step_metadata", None)
         if getattr(self.model, "supports_omni_decode_step_metadata", False) and callable(update_decode_metadata):
+            cpu_input_tail_ids = None
+            if (
+                getattr(self.model, "requires_cpu_input_tail_ids", False)
+                and not self.use_async_scheduling
+                and input_ids is not None
+                and input_ids.data_ptr() == self.input_ids.gpu.data_ptr()
+            ):
+                count = len(self.input_batch.req_ids)
+                tails = self.query_start_loc.cpu[1 : count + 1].to(torch.long) - 1
+                if count and bool((tails >= 0).all()) and bool((tails < input_ids.numel()).all()):
+                    cpu_input_tail_ids = self.input_ids.cpu.index_select(0, tails).tolist()
             update_decode_metadata(
+                cpu_input_tail_ids=cpu_input_tail_ids,
                 input_ids=input_ids,
                 positions=positions,
                 inputs_embeds=inputs_embeds,
@@ -2156,14 +2168,19 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
                 req_ids=self.input_batch.req_ids,
             )
 
-        model_output = super()._model_forward(
-            input_ids=input_ids,
-            positions=positions,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=inputs_embeds,
-            **model_kwargs,
-            **model_kwargs_extra,
-        )
+        try:
+            model_output = super()._model_forward(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                **model_kwargs,
+                **model_kwargs_extra,
+            )
+        finally:
+            finish_decode_step = getattr(self.model, "finish_decode_step_forward", None)
+            if callable(finish_decode_step):
+                finish_decode_step()
         # CUDAGraphWrapper's weak_ref_tensors preserves the fields but turns
         # NamedTuple outputs into plain tuples. Restore the Omni envelope
         # before a model adapter or extraction discards its multimodal fields.

vllm_omni/worker_v2/model_states/__init__.py

@@ -31,6 +31,9 @@ def init_omni_model_state(
     capabilities; otherwise
     delegates to the upstream v2 factory.
     """
+    factory = getattr(model, "create_omni_model_state", None)
+    if factory is not None:
+        return factory(vllm_config, encoder_cache, device)
     uses_omni_lifecycle = any(
         getattr(model, flag, False) is True for flag in ("has_preprocess", "has_postprocess", "have_multimodal_outputs")
     )

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -17,6 +17,7 @@ import inspect
 import threading
 import types
 from collections.abc import Callable
+from dataclasses import replace
 from typing import Any, cast
 
 import numpy as np
@@ -26,10 +27,12 @@ from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.states import RequestState
+from vllm.v1.worker.utils import AttentionGroup
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
 from vllm_omni.platforms import current_omni_platform
@@ -413,6 +416,36 @@ class OmniModelState(DefaultModelState):
     # Input preparation
     # ------------------------------------------------------------------
 
+    def prepare_attn(
+        self,
+        input_batch: InputBatch,
+        cudagraph_mode: CUDAGraphMode,
+        block_tables: tuple[torch.Tensor, ...],
+        slot_mappings: torch.Tensor,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+        for_capture: bool = False,
+        ubatch_idx: int = 0,
+    ) -> dict[str, Any]:
+        if for_capture and input_batch.max_query_len is None:
+            # vLLM 0.30 distributes dummy tokens evenly across requests. For
+            # an unconstrained FULL graph that split is not a query-length
+            # bound: replay may put the entire token bucket in one request.
+            # Attention launch parameters are fixed at capture, so use the
+            # bucket's worst-case query length, not the dummy per-row length.
+            # Keep explicit bounds (e.g. varlen decode) and runtime metadata.
+            input_batch = replace(input_batch, max_query_len=input_batch.num_tokens)
+        return super().prepare_attn(
+            input_batch,
+            cudagraph_mode,
+            block_tables,
+            slot_mappings,
+            attn_groups,
+            kv_cache_config,
+            for_capture=for_capture,
+            ubatch_idx=ubatch_idx,
+        )
+
     def prepare_inputs_embeds(
         self,
         scheduled_encoder_inputs: dict[str, list[int]],