From c41925689d11e99ff20f0cd98a6140d795686790 Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Sun, 30 Aug 2026 14:22:15 +0530 Subject: [PATCH] [BugFix][Higgs-Audio-V3] Fix transcript and concurrent sampling tests (#6422) Signed-off-by: akshatvishu --- tests/config/test_config_factory.py | 14 ++ .../e2e/online_serving/test_higgs_audio_v3.py | 47 +++---- tests/helpers/assertions.py | 4 +- tests/helpers/client.py | 1 + tests/helpers/tests/test_assertions.py | 42 +++++- .../higgs_audio_v3/test_higgs_audio_v3.py | 122 +++++++++++++++++- tests/worker/test_omni_gpu_model_runner.py | 21 +++ vllm_omni/deploy/higgs_multimodal_qwen3.yaml | 7 + .../higgs_audio_v3/higgs_audio_v3_talker.py | 102 ++++++++++++++- vllm_omni/worker/gpu_model_runner.py | 4 + 10 files changed, 330 insertions(+), 34 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 078fab250..19481d895 100644 --- a/tests/config/test_config_factory.py +++ b/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( diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py index 8550f74a4..e9ba68909 100644 --- a/tests/e2e/online_serving/test_higgs_audio_v3.py +++ b/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}]``. diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 8b43233c7..755cf0a4e 100644 --- a/tests/helpers/assertions.py +++ b/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}") diff --git a/tests/helpers/client.py b/tests/helpers/client.py index edda87790..9ef4a7a0c 100644 --- a/tests/helpers/client.py +++ b/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) diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py index 9cb1560b9..98deaa511 100644 --- a/tests/helpers/tests/test_assertions.py +++ b/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. diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 77a32ad7b..f5a7094ef 100755 --- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py +++ b/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) diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index cd5c559d7..81d2a6b8a 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/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.""" diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml index ae4e373cb..41e1a9677 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml +++ b/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 diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index 579d0cd13..186e48714 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/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) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index a4716d682..eca15c16e 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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(