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

.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 @@ def test_voice_clone_references_alias(self, omni_server, online_client) -> None:
                 "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