[BugFix][Higgs-Audio-V3] Fix transcript and concurrent sampling tests (#6422)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/config/test_config_factory.py

@@ -2399,6 +2399,20 @@ class TestPlatformOverrides:
         assert rocm.stages[0].enforce_eager is None
         assert rocm.stages[1].enforce_eager is True
 
+    def test_higgs_audio_v3_rocm_uses_triton_attention(self):
+        deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
+
+        base = load_deploy_config(deploy_path)
+        assert base.stages[0].engine_extras["attention_backend"] == "FLASHINFER"
+
+        rocm = _apply_platform_overrides(base, platform="rocm")
+        assert rocm.stages[0].engine_extras["attention_backend"] == "TRITON_ATTN"
+
+        pipeline = resolve_pipeline_config("higgs_multimodal_qwen3")
+        assert isinstance(pipeline, PipelineConfig)
+        stages = merge_pipeline_deploy(pipeline, rocm)
+        assert stages[0].yaml_engine_args["attention_backend"] == "TRITON_ATTN"
+
     def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         pipeline = resolve_pipeline_config(

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -48,10 +48,6 @@ DEFAULT_SPEECH_TIMEOUT_S = 180.0
 # silence-only outputs without flagging short legitimate clips.
 _MIN_AUDIO_BYTES = 20_000
 
-_SKIP_ISSUE_6418 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/6418",
-)
-
 # Reuse the shared TTS reference clip (clean ~5 s 24 kHz mono human speech)
 # vendored under tests/assets/qwen3_tts/. Keeps a single WAV across TTS
 # tests rather than duplicating asset bytes.
@@ -65,7 +61,7 @@ class TestHiggsAudioV3OnlineHappyPath:
 
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_plain_text_wav(self, omni_server, online_client) -> None:
         """Single non-streaming WAV request - canonical TTS happy path."""
         online_client.send_audio_speech_request(
@@ -82,10 +78,9 @@ class TestHiggsAudioV3OnlineHappyPath:
             }
         )
 
-    @_SKIP_ISSUE_6418
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_plain_text_with_max_new_tokens(self, omni_server, online_client) -> None:
         """``max_new_tokens`` is one of the few extra fields the v3 validator accepts."""
         online_client.send_audio_speech_request(
@@ -95,15 +90,14 @@ class TestHiggsAudioV3OnlineHappyPath:
                 "stream": False,
                 "response_format": "wav",
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
-                "max_new_tokens": 32,
+                "max_new_tokens": 1024,
                 "min_audio_bytes": _MIN_AUDIO_BYTES,
             }
         )
 
-    @_SKIP_ISSUE_6418
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_concurrent_plain_text(self, omni_server, online_client) -> None:
         """Three concurrent non-streaming requests - guards per-slot audio state and Stage-0 PC under batching."""
         online_client.send_audio_speech_request(
@@ -119,7 +113,7 @@ class TestHiggsAudioV3OnlineHappyPath:
 
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_plain_text_pcm_streaming(self, omni_server, online_client) -> None:
         """Streaming PCM via talker2code2wav_async_chunk + the connector's codec_streaming path.
 
@@ -150,7 +144,7 @@ class TestHiggsAudioV3OnlineHappyPath:
 
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_concurrent_pcm_streaming(self, omni_server, online_client) -> None:
         """Three concurrent streaming requests - guards per-request frame cursors
         in ``talker2code2wav_async_chunk`` and per-slot delay-pattern state under batched AR.
@@ -185,15 +179,13 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
     tokens (pause, sfx) go inline. SFX tokens must pair with their written
     onomatopoeia.
 
-    These tests check the *serving surface* - the validator accepts the
-    payload, the engine produces audio, and the WAV has non-trivial size.
-    They do not assert audio quality (out of scope for CI).
+    These tests check the serving surface and spoken text. They do not assert
+    whether the requested control changes audio quality.
     """
 
-    @_SKIP_ISSUE_6418
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_inline_emotion_and_expressive(self, omni_server, online_client) -> None:
         """Delivery tokens (emotion + expressive_high) at the start of input."""
         online_client.send_audio_speech_request(
@@ -204,13 +196,13 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
                 "response_format": "wav",
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
                 "min_audio_bytes": _MIN_AUDIO_BYTES,
+                "transcript_expected_text": "Wait, that was actually hilarious.",
             }
         )
 
-    @_SKIP_ISSUE_6418
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_inline_style_whispering(self, omni_server, online_client) -> None:
         """Style token at the start - ``whispering``."""
         online_client.send_audio_speech_request(
@@ -221,13 +213,13 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
                 "response_format": "wav",
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
                 "min_audio_bytes": _MIN_AUDIO_BYTES,
+                "transcript_expected_text": "It is just between you and me, alright?",
             }
         )
 
-    @_SKIP_ISSUE_6418
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_inline_prosody_speed_and_pitch(self, omni_server, online_client) -> None:
         """Two prosody tokens at the start - slow speed plus low pitch."""
         online_client.send_audio_speech_request(
@@ -240,13 +232,13 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
                 "response_format": "wav",
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
                 "min_audio_bytes": _MIN_AUDIO_BYTES,
+                "transcript_expected_text": "The radar shows a storm approaching from the east.",
             }
         )
 
-    @_SKIP_ISSUE_6418
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_inline_pause_mid_text(self, omni_server, online_client) -> None:
         """Positional pause token placed inline between two clauses."""
         online_client.send_audio_speech_request(
@@ -257,13 +249,13 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
                 "response_format": "wav",
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
                 "min_audio_bytes": _MIN_AUDIO_BYTES,
+                "transcript_expected_text": "Hold on a moment. Let me think about it.",
             }
         )
 
-    @_SKIP_ISSUE_6418
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_inline_sfx_with_onomatopoeia(self, omni_server, online_client) -> None:
         """SFX token paired with its written onomatopoeia (``<|sfx:laughter|>Hehe``)."""
         online_client.send_audio_speech_request(
@@ -277,6 +269,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
                 "response_format": "wav",
                 "timeout": DEFAULT_SPEECH_TIMEOUT_S,
                 "min_audio_bytes": _MIN_AUDIO_BYTES,
+                "transcript_expected_text": "I cannot believe that just happened. I am still recovering from it.",
             }
         )
 
@@ -287,7 +280,7 @@ class TestHiggsAudioV3OnlineVoiceClone:
 
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_voice_clone_ref_audio_ref_text(self, omni_server, online_client) -> None:
         """Canonical vllm-omni voice clone via ``ref_audio`` + ``ref_text``."""
         online_client.send_audio_speech_request(
@@ -304,7 +297,7 @@ class TestHiggsAudioV3OnlineVoiceClone:
 
     @pytest.mark.slow
     @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
     def test_voice_clone_references_alias(self, omni_server, online_client) -> None:
         """BosonAI cookbook payload: ``references=[{audio_path, text}]``.
 

tests/helpers/assertions.py

@@ -665,7 +665,7 @@ def _speech_assertion_needs_audio_transcript(request_config: dict[str, Any], run
         return False
     if request_config.get("response_format") == "pcm":
         return False
-    return bool(request_config.get("input"))
+    return bool(request_config.get("transcript_expected_text", request_config.get("input")))
 
 
 def _resolve_audio_transcript(
@@ -880,7 +880,7 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
 
         transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True)
         if transcript is not None:
-            expected_text = request_config.get("input")
+            expected_text = request_config.get("transcript_expected_text", request_config.get("input"))
             if expected_text:
                 print(f"audio content is: {transcript}")
                 print(f"input text is: {expected_text}")

tests/helpers/client.py

@@ -1432,6 +1432,7 @@ class OnlineOmniClient:
           - response_format: audio format such as "wav" or "pcm" (optional)
           - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body)
           - min_audio_bytes: optional minimum ``len(audio_bytes)`` checked in ``assert_audio_speech_response``
+          - transcript_expected_text: local expected spoken text; defaults to ``input``
           - timeout: request timeout in seconds (float, optional, default 120.0)
           - stream: whether to use streaming API (bool, optional, default False)
 

tests/helpers/tests/test_assertions.py

@@ -1,12 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
 import pytest
 
 from tests.helpers import assertions
-from tests.helpers.assertions import _assert_transcript_matches, _resolve_audio_transcript
+from tests.helpers.assertions import (
+    _assert_transcript_matches,
+    _resolve_audio_transcript,
+    assert_audio_speech_response,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -71,6 +75,40 @@ def test_resolve_transcript_leaves_language_unset_by_default(monkeypatch):
     assert captured["model_size"] == "small"
 
 
+@pytest.mark.parametrize(
+    ("request_config", "expected_text"),
+    [
+        ({"input": "spoken words"}, "spoken words"),
+        (
+            {
+                "input": "<|style:whispering|>spoken words",
+                "transcript_expected_text": "spoken words",
+            },
+            "spoken words",
+        ),
+    ],
+)
+def test_speech_transcript_expected_text(monkeypatch, request_config, expected_text):
+    captured: dict = {}
+
+    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "spoken words")
+
+    def capture_match(_transcript, _audio_bytes, expected_text, **_kwargs):
+        captured["expected_text"] = expected_text
+
+    monkeypatch.setattr(assertions, "_assert_transcript_matches", capture_match)
+    response = SimpleNamespace(success=True, audio_bytes=b"fake-wav", audio_format="audio/wav")
+    request_config["response_format"] = "wav"
+
+    assert_audio_speech_response(
+        response,
+        request_config,
+        "advanced_model",
+    )
+
+    assert captured["expected_text"] == expected_text
+
+
 def test_escalated_transcript_keeps_declared_language(monkeypatch):
     # The escalated pass must honour the same language, otherwise a request that
     # pins one silently falls back to auto-detection on retry.

tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Deterministic unit tests for higgs-audio v3.
 
 These tests verify AC-1 (config), AC-3 (prompt), AC-4 (fused modules),
@@ -343,7 +344,22 @@ class TestSamplerMethods:
         return t
 
     @staticmethod
-    def _sampling_metadata(temperature=1.0, top_k=50, top_p=0.95):
+    def _make_state_tracking_talker(num_rows):
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        talker_cls = mod.HiggsAudioV3TalkerForConditionalGeneration
+        talker = talker_cls.__new__(talker_cls)
+        torch.nn.Module.__init__(talker)
+        talker.num_codebooks = 8
+        talker._decode_last_codes = torch.zeros(num_rows, 8, dtype=torch.long)
+        talker._decode_has_codes = torch.zeros(num_rows, dtype=torch.bool)
+        talker._decode_delay_count = torch.zeros(num_rows, dtype=torch.long)
+        talker._decode_eoc_countdown = torch.full((num_rows,), -1, dtype=torch.long)
+        talker._decode_generation_done = torch.zeros(num_rows, dtype=torch.bool)
+        return talker
+
+    @staticmethod
+    def _sampling_metadata(temperature=1.0, top_k=50, top_p=0.95, generators=None):
         return type(
             "SamplingMetadata",
             (),
@@ -352,9 +368,16 @@ class TestSamplerMethods:
                 "top_k": top_k,
                 "top_p": top_p,
                 "all_greedy": temperature is None,
+                "generators": generators or {},
             },
         )()
 
+    @staticmethod
+    def _seeded_generator(seed):
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(seed)
+        return generator
+
     def test_sample_respects_mask(self):
         """Tokens masked to -inf must never be sampled."""
         t = self._make_batched_sampler_talker(num_rows=1)
@@ -442,6 +465,51 @@ class TestSamplerMethods:
         assert torch.count_nonzero(seen[0][0]).item() == 1
         assert torch.count_nonzero(seen[0][2]).item() > 1
 
+    def test_seeded_sampling_is_batch_independent(self):
+        t = self._make_minimal_talker()
+        num_codebooks = 8
+        torch.manual_seed(123)
+        first_logits = torch.randn(num_codebooks, 64)
+        second_logits = torch.randn(num_codebooks, 64)
+        filler_logits = torch.randn(num_codebooks, 64)
+        reference_generator = self._seeded_generator(42)
+        batched_generators = {
+            1: self._seeded_generator(42),
+            2: self._seeded_generator(42),
+        }
+
+        reference_first = t._sample_audio_codes(
+            first_logits,
+            self._sampling_metadata(generators={0: reference_generator}),
+            num_codebooks=num_codebooks,
+        )
+        batched_first = t._sample_audio_codes(
+            torch.cat([filler_logits, first_logits, first_logits]),
+            self._sampling_metadata(generators=batched_generators),
+            num_codebooks=num_codebooks,
+        )
+        assert torch.equal(reference_first, batched_first[num_codebooks : 2 * num_codebooks])
+        assert torch.equal(reference_first, batched_first[2 * num_codebooks :])
+
+        reference_second = t._sample_audio_codes(
+            second_logits,
+            self._sampling_metadata(generators={0: reference_generator}),
+            num_codebooks=num_codebooks,
+        )
+        condensed_second = t._sample_audio_codes(
+            torch.cat([second_logits, second_logits]),
+            self._sampling_metadata(
+                generators={
+                    0: batched_generators[1],
+                    1: batched_generators[2],
+                }
+            ),
+            num_codebooks=num_codebooks,
+        )
+
+        assert torch.equal(reference_second, condensed_second[:num_codebooks])
+        assert torch.equal(reference_second, condensed_second[num_codebooks:])
+
     def test_delay_masking_forces_boc_during_delay(self):
         """During delay phase, codebooks beyond delay_count must have only BOC allowed."""
         from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
@@ -678,6 +746,58 @@ class TestSamplerMethods:
         assert not t._decode_generation_done[2:5].any()
         assert t._decode_eoc_countdown[2:5].eq(-1).all()
 
+    def test_decode_state_follows_request_after_batch_condensation(self):
+        talker = self._make_state_tracking_talker(3)
+
+        talker._sync_decode_state_with_batch(["A", "B", "C"])
+        c_codes = torch.arange(300, 308)
+        talker._decode_last_codes[:] = torch.stack((torch.arange(100, 108), torch.zeros(8, dtype=torch.long), c_codes))
+        talker._decode_has_codes[:] = torch.tensor([True, False, True])
+        talker._decode_delay_count[:] = torch.tensor([3, 0, 8])
+
+        talker.on_requests_finished({"B"})
+        talker._sync_decode_state_with_batch(["A", "C"])
+
+        assert talker._decode_has_codes[1].item() is True
+        assert talker._decode_delay_count[1].item() == 8
+        assert torch.equal(talker._decode_last_codes[1], c_codes)
+
+    def test_finished_request_state_is_cleared_before_pool_slot_reuse(self):
+        talker = self._make_state_tracking_talker(1)
+
+        talker._sync_decode_state_with_batch(["finished"])
+        talker._decode_has_codes[0] = True
+        talker._decode_last_codes[0] = 7
+        talker._decode_delay_count[0] = 5
+        talker._decode_eoc_countdown[0] = 0
+        talker._decode_generation_done[0] = True
+        talker.on_requests_finished({"finished"})
+        talker._sync_decode_state_with_batch(["new"])
+
+        assert talker._decode_has_codes[0].item() is False
+        assert talker._decode_last_codes[0].eq(0).all()
+        assert talker._decode_delay_count[0].item() == 0
+        assert talker._decode_eoc_countdown[0].item() == -1
+        assert talker._decode_generation_done[0].item() is False
+
+    def test_finished_request_state_is_cleared_when_request_id_is_reused(self):
+        talker = self._make_state_tracking_talker(1)
+
+        talker._sync_decode_state_with_batch(["request"])
+        talker._decode_has_codes[0] = True
+        talker._decode_last_codes[0] = 7
+        talker._decode_delay_count[0] = 5
+        talker._decode_eoc_countdown[0] = 0
+        talker._decode_generation_done[0] = True
+        talker.on_requests_finished({"request"})
+        talker._sync_decode_state_with_batch(["request"])
+
+        assert talker._decode_has_codes[0].item() is False
+        assert talker._decode_last_codes[0].eq(0).all()
+        assert talker._decode_delay_count[0].item() == 0
+        assert talker._decode_eoc_countdown[0].item() == -1
+        assert talker._decode_generation_done[0].item() is False
+
     def test_mixed_batch_prefill_mask_targets_request_rows(self):
         """Mixed prefill/decode must reset only prefill request rows."""
         t = self._make_batched_sampler_talker(num_rows=3)

tests/worker/test_omni_gpu_model_runner.py

@@ -1,9 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
 import torch
 from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 from vllm_omni.worker.gpu_model_runner import (
@@ -107,6 +111,23 @@ class DummyReqState:
     pass
 
 
+def test_model_forward_passes_request_ids_to_decode_metadata(monkeypatch):
+    received = {}
+    model = SimpleNamespace(
+        supports_omni_decode_step_metadata=True,
+        update_decode_step_metadata=lambda **kwargs: received.update(kwargs),
+    )
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = model
+    runner.input_batch = DummyInputBatch(["request-a", "request-b"])
+    runner._build_model_kwargs_extra = lambda: {}
+    monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *_args, **_kwargs: torch.zeros(1))
+
+    OmniGPUModelRunner._model_forward(runner, input_ids=torch.ones(2, dtype=torch.long))
+
+    assert received["req_ids"] == ["request-a", "request-b"]
+
+
 class MiMoAudioForConditionalGeneration(torch.nn.Module):
     """Dummy model whose class name must exactly match the production check."""
 

vllm_omni/deploy/higgs_multimodal_qwen3.yaml

@@ -83,3 +83,10 @@ stages:
       max_tokens: 65536
       seed: 42
       repetition_penalty: 1.0
+
+platforms:
+  rocm:
+    stages:
+      - stage_id: 0
+        # FlashInfer is not supported on ROCm.
+        attention_backend: TRITON_ATTN

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.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-0 talker for higgs-audio v3 (Qwen3 backbone, fused multi-codebook).
 
 Architecture:
@@ -344,6 +344,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         positions: torch.Tensor | None = None,
         inputs_embeds: torch.Tensor | None = None,
         omni_query_start_loc: torch.Tensor | None = None,
+        req_ids: list[str] | None = None,
         **_: Any,
     ) -> None:
         """Update per-step metadata before runner forward or CUDA graph replay."""
@@ -355,6 +356,91 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         self._set_last_step_query_start_loc(omni_query_start_loc)
         self._decode_step_metadata_from_runner = True
 
+        if req_ids is not None:
+            self._sync_decode_state_with_batch(req_ids)
+
+    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] = {}
+            self._free_state_pool_indices: set[int] = set()
+            self._state_pool_has_codes = torch.empty(0, dtype=torch.bool, device=device)
+            self._state_pool_last_codes = torch.empty((0, self.num_codebooks), dtype=torch.long, device=device)
+            self._state_pool_delay_count = torch.empty(0, dtype=torch.long, device=device)
+            self._state_pool_eoc_countdown = torch.empty(0, dtype=torch.long, device=device)
+            self._state_pool_generation_done = torch.empty(0, dtype=torch.bool, device=device)
+
+        current_rows = self._state_pool_has_codes.shape[0]
+        if current_rows >= num_rows:
+            return
+
+        new_rows = max(current_rows * 2, num_rows, 64)
+
+        def resize(tensor: torch.Tensor, fill_value: int | bool) -> torch.Tensor:
+            result = torch.empty((new_rows, *tensor.shape[1:]), dtype=tensor.dtype, device=device)
+            result[:current_rows] = tensor
+            result[current_rows:] = fill_value
+            return result
+
+        self._state_pool_has_codes = resize(self._state_pool_has_codes, False)
+        self._state_pool_last_codes = resize(self._state_pool_last_codes, 0)
+        self._state_pool_delay_count = resize(self._state_pool_delay_count, 0)
+        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 _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 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)
+                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]
+                self._state_pool_eoc_countdown[dst] = self._decode_eoc_countdown[src]
+                self._state_pool_generation_done[dst] = self._decode_generation_done[src]
+
+            for req_id in req_ids:
+                if req_id in self._state_pool_indices:
+                    continue
+                if self._free_state_pool_indices:
+                    pool_row = self._free_state_pool_indices.pop()
+                else:
+                    pool_row = len(self._state_pool_indices)
+                    self._ensure_state_pool_capacity(pool_row + 1, device)
+                self._state_pool_indices[req_id] = pool_row
+                self._state_pool_has_codes[pool_row] = False
+                self._state_pool_last_codes[pool_row] = 0
+                self._state_pool_delay_count[pool_row] = 0
+                self._state_pool_eoc_countdown[pool_row] = -1
+                self._state_pool_generation_done[pool_row] = False
+
+            self._ensure_decode_state_capacity(len(req_ids), device)
+            pool_rows = torch.tensor(
+                [self._state_pool_indices[req_id] for req_id in req_ids],
+                dtype=torch.long,
+                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]
+            self._decode_delay_count[: len(req_ids)] = self._state_pool_delay_count[pool_rows]
+            self._decode_eoc_countdown[: len(req_ids)] = self._state_pool_eoc_countdown[pool_rows]
+            self._decode_generation_done[: len(req_ids)] = self._state_pool_generation_done[pool_rows]
+
+        self._last_batch_req_ids = list(req_ids)
+
+    def on_requests_finished(self, finished_req_ids: set[str]) -> None:
+        if not hasattr(self, "_state_pool_indices"):
+            return
+        for req_id in finished_req_ids:
+            pool_row = self._state_pool_indices.pop(req_id, None)
+            if pool_row is not None:
+                self._free_state_pool_indices.add(pool_row)
+
     # ------------------------------------------------------------------ forward
     def forward(
         self,
@@ -1386,7 +1472,19 @@ 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)
-        sampled = torch.multinomial(probs, num_samples=1).squeeze(-1)
+        if sampling_metadata.generators:
+            sampled = torch.cat(
+                [
+                    torch.multinomial(
+                        probs[req_idx * num_codebooks : (req_idx + 1) * num_codebooks],
+                        num_samples=1,
+                        generator=sampling_metadata.generators.get(req_idx),
+                    )
+                    for req_idx in range(int(probs.shape[0]) // num_codebooks)
+                ]
+            ).squeeze(-1)
+        else:
+            sampled = torch.multinomial(probs, num_samples=1).squeeze(-1)
         sampled = torch.where(all_masked, fallback, sampled)
         return torch.where(greedy, logits_2d.argmax(dim=-1), sampled)
 

vllm_omni/worker/gpu_model_runner.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import contextlib
 import inspect
 from collections.abc import Callable
@@ -1988,6 +1991,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 positions=positions,
                 inputs_embeds=inputs_embeds,
                 omni_query_start_loc=model_kwargs_extra.get("omni_query_start_loc"),
+                req_ids=self.input_batch.req_ids,
             )
 
         model_output = super()._model_forward(