[Bugfix][Higgs-Audio-V3] Propagate Sampling Parameters to Sampler (#5406)

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

tests/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -250,6 +250,7 @@ class TestSamplerMethods:
 
         t = FakeTalker()
         t._sample_audio_codes = mod.HiggsAudioV3TalkerForConditionalGeneration._sample_audio_codes.__get__(t)
+        t._expand_audio_sampling_params = mod.HiggsAudioV3TalkerForConditionalGeneration._expand_audio_sampling_params
         return t
 
     def _make_batched_sampler_talker(self, num_rows=4):
@@ -300,16 +301,30 @@ class TestSamplerMethods:
             "_audio_seed_mask_from_step_input",
         ):
             setattr(t, name, getattr(cls, name).__get__(t))
+        t._expand_audio_sampling_params = cls._expand_audio_sampling_params
         t._device_cache_key = cls._device_cache_key
         return t
 
+    @staticmethod
+    def _sampling_metadata(temperature=1.0, top_k=50, top_p=0.95):
+        return type(
+            "SamplingMetadata",
+            (),
+            {
+                "temperature": temperature,
+                "top_k": top_k,
+                "top_p": top_p,
+                "all_greedy": temperature is None,
+            },
+        )()
+
     def test_sample_respects_mask(self):
         """Tokens masked to -inf must never be sampled."""
         t = self._make_batched_sampler_talker(num_rows=1)
         logits = torch.full((10, 1026), float("-inf"))
         # Only token 500 is allowed for each row
         logits[:, 500] = 0.0
-        result = t._sample_audio_codes(logits)
+        result = t._sample_audio_codes(logits, self._sampling_metadata(), num_codebooks=1)
         assert result.shape == (10,)
         assert (result == 500).all(), f"Expected all 500, got {result.tolist()}"
 
@@ -320,10 +335,76 @@ class TestSamplerMethods:
         # Row 0: all masked
         # Row 1: only token 42 allowed
         logits[1, 42] = 0.0
-        result = t._sample_audio_codes(logits)
+        result = t._sample_audio_codes(logits, self._sampling_metadata(), num_codebooks=1)
         assert result.shape == (2,)
         assert result[1].item() == 42
 
+    def test_top_k_minus_one_disables_top_k_filtering(self, monkeypatch):
+        t = self._make_minimal_talker()
+        seen_probs = []
+
+        def fake_multinomial(probs, num_samples):
+            seen_probs.append(probs.clone())
+            return probs.argmax(dim=-1, keepdim=True)
+
+        monkeypatch.setattr(torch, "multinomial", fake_multinomial)
+
+        logits = torch.tensor([[4.0, 3.0, 2.0, 1.0]])
+        metadata = self._sampling_metadata(
+            temperature=1.0,
+            top_k=-1,
+            top_p=1.0,
+        )
+
+        t._sample_audio_codes(logits, metadata, num_codebooks=1)
+
+        assert len(seen_probs) == 1
+        assert torch.count_nonzero(seen_probs[0][0]).item() == 4
+
+    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)
+        )
+        logits = torch.tensor([[0.0, 2.0], [3.0, 0.0]])
+
+        result = t._sample_audio_codes(
+            logits,
+            self._sampling_metadata(torch.tensor([float("nan"), float("inf")])),
+            num_codebooks=1,
+        )
+
+        assert len(seen) == 1
+        assert torch.isfinite(seen[0]).all()
+        assert result.tolist() == [1, 0]
+
+    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)
+        )
+        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(
+            temperature=torch.tensor([0.5, 1.0]),
+            top_k=torch.tensor([1, 4]),
+            top_p=torch.tensor([1.0, 0.8]),
+        )
+
+        temperatures, top_ks, top_ps = t._expand_audio_sampling_params(metadata, logits, num_codebooks=2)
+        t._sample_audio_codes(logits, metadata, num_codebooks=2)
+
+        torch.testing.assert_close(temperatures, torch.tensor([0.5, 0.5, 1.0, 1.0]))
+        assert top_ks.tolist() == [1, 1, 4, 4]
+        torch.testing.assert_close(top_ps, torch.tensor([1.0, 1.0, 0.8, 0.8]))
+        assert len(seen) == 1
+        torch.testing.assert_close(seen[0][0], seen[0][1])
+        torch.testing.assert_close(seen[0][2], seen[0][3])
+        assert torch.count_nonzero(seen[0][0]).item() == 1
+        assert torch.count_nonzero(seen[0][2]).item() > 1
+
     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
@@ -468,7 +549,9 @@ class TestSamplerMethods:
         t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: 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 cb_logits, audio_rows, all_rows=False: None
-        t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long)
+        t._sample_audio_codes = lambda logits_2d, *args, **kwargs: torch.zeros(
+            int(logits_2d.shape[0]), dtype=torch.long
+        )
         t._update_delay_state_batched = lambda *args, **kwargs: None
         t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t)
 
@@ -499,7 +582,9 @@ class TestSamplerMethods:
         t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: 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 cb_logits, audio_rows, all_rows=False: None
-        t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long)
+        t._sample_audio_codes = lambda logits_2d, *args, **kwargs: torch.zeros(
+            int(logits_2d.shape[0]), dtype=torch.long
+        )
         t._update_delay_state_batched = lambda *args, **kwargs: None
         t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t)
 

tests/worker/test_gpu_ar_model_runner.py

@@ -1,10 +1,17 @@
+import asyncio
 from collections.abc import Sequence
 from types import SimpleNamespace
 
 import numpy as np
 import pytest
 import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker import gpu_input_batch
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
+from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_ar_model_runner import (
     ExecuteModelState,
@@ -25,6 +32,79 @@ def _make_runner(engine_output_type: str | None, downstream_req_ids: set[str]) -
     return runner
 
 
+def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatch):
+    requested = {"temperature": 0.7, "top_p": 0.8, "top_k": 17}
+    request = OpenAICreateSpeechRequest.model_validate(
+        {"input": "Hello", "extra_params": requested},
+    )
+
+    class Adapter:
+        def validate(self, request):
+            return None
+
+        async def build(self, request, sampling_params_list, has_inline_ref_audio):
+            return PreparedRequest(prompt={"prompt": request.input}, tts_params={}, model_type="higgs_audio_v3")
+
+    engine = SimpleNamespace(
+        errored=False,
+        default_sampling_params_list=[SamplingParams(temperature=1.0, top_p=0.95, top_k=50)],
+        generate=lambda **kwargs: kwargs["sampling_params_list"],
+    )
+    serving = object.__new__(OmniOpenAIServingSpeech)
+    serving.engine_client = engine
+    serving.model_config = SimpleNamespace(async_chunk=True)
+    serving._tts_model_type = "higgs_audio_v3"
+    serving._get_tts_adapter = lambda: Adapter()
+    serving._track_ref_audio_artifact_warmup = lambda *args, **kwargs: None
+
+    _, stage_sampling_params, _ = asyncio.run(serving._prepare_speech_generation(request, request_id="speech-test"))
+    stage0_params = stage_sampling_params[0]
+
+    monkeypatch.setattr(gpu_input_batch, "PIN_MEMORY", False)
+    input_batch = InputBatch(
+        max_num_reqs=1,
+        max_model_len=8,
+        max_num_batched_tokens=8,
+        device=torch.device("cpu"),
+        vocab_size=1024,
+        block_sizes=[1],
+        kernel_block_sizes=[1],
+    )
+    input_batch.add_request(
+        CachedRequestState(
+            req_id="speech-test",
+            prompt_token_ids=[1],
+            mm_features=[],
+            sampling_params=stage0_params,
+            generator=None,
+            block_ids=([],),
+            num_computed_tokens=0,
+            output_token_ids=[],
+        )
+    )
+    input_batch.sampling_metadata = input_batch._make_sampling_metadata()
+
+    received = []
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = input_batch
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        skips_model_sampler_output_token_history=True,
+        sample=lambda logits, metadata: received.append(metadata) or "model-sampler",
+    )
+    runner.sampler = SimpleNamespace()
+    logits = torch.zeros((1, 4))
+
+    output = runner._sample(logits, spec_decode_metadata=None)
+
+    assert output == "model-sampler"
+    assert len(received) == 1
+    sampling_metadata = received[0]
+    torch.testing.assert_close(sampling_metadata.temperature, torch.tensor([requested["temperature"]]))
+    torch.testing.assert_close(sampling_metadata.top_p, torch.tensor([requested["top_p"]]))
+    assert sampling_metadata.top_k.tolist() == [requested["top_k"]]
+
+
 def test_resolve_pooler_payload_req_ids_audio_terminal_stage_keeps_payload():
     runner = _make_runner(engine_output_type="audio", downstream_req_ids=set())
 

vllm_omni/entrypoints/openai/serving_speech.py

@@ -3183,6 +3183,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             import copy
 
             sampling_params_list = copy.deepcopy(sampling_params_list)
+            for name in ("temperature", "top_p", "top_k"):
+                if (value := request.extra_params.get(name)) is not None:
+                    setattr(sampling_params_list[0], name, value)
             if sampling_params_list[0].extra_args is None:
                 sampling_params_list[0].extra_args = {}
             sampling_params_list[0].extra_args.update(request.extra_params)

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py

@@ -35,6 +35,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.qwen3 import Qwen3Model
 from vllm.platforms import current_platform
 from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.higgs_audio_v3 import (
@@ -983,7 +985,11 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
 
         # Sample per-codebook
         cb_logits_2d = cb_logits.reshape(-1, cb_logits.shape[-1])
-        codes_2d = self._sample_audio_codes(cb_logits_2d)
+        codes_2d = self._sample_audio_codes(
+            cb_logits_2d,
+            sampling_metadata,
+            int(cb_logits.shape[1]),
+        )
         codes_flat = codes_2d.view(cb_logits.shape[0], cb_logits.shape[1]).to(torch.long)
 
         self._update_delay_state_batched(
@@ -1337,23 +1343,42 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         self._postprocess_audio_rows = num_rows
         self._postprocess_audio_active_rows = 0
 
-    def _sample_audio_codes(self, logits_2d: torch.Tensor) -> torch.Tensor:
+    @staticmethod
+    def _expand_audio_sampling_params(
+        sampling_metadata: SamplingMetadata, logits_2d: torch.Tensor, num_codebooks: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size, remainder = divmod(int(logits_2d.shape[0]), num_codebooks)
+        if remainder:
+            raise ValueError(f"audio codebook logits rows must be divisible by {num_codebooks}")
+
+        def expand(value, default, dtype):
+            if value is None:
+                value = torch.full((batch_size,), default, device=logits_2d.device, dtype=dtype)
+            else:
+                value = torch.as_tensor(value, device=logits_2d.device, dtype=dtype).expand(batch_size)
+            return value.repeat_interleave(num_codebooks)
+
+        return (
+            expand(sampling_metadata.temperature, 0.0 if sampling_metadata.all_greedy else 1.0, logits_2d.dtype),
+            expand(sampling_metadata.top_k, logits_2d.shape[-1], torch.long),
+            expand(sampling_metadata.top_p, 1.0, logits_2d.dtype),
+        )
+
+    def _sample_audio_codes(
+        self, logits_2d: torch.Tensor, sampling_metadata: SamplingMetadata, num_codebooks: int = 8
+    ) -> torch.Tensor:
         """Replicate upstream sampling: temperature → top-k → top-p → multinomial."""
         x = logits_2d.float()
-        top_k = 50
-        top_p = 0.95
-        if 0 < top_k < x.shape[-1]:
-            kth = x.topk(top_k, dim=-1).values[..., -1:]
-            x = x.masked_fill(x < kth, float("-inf"))
-        if 0.0 < top_p < 1.0:
-            sorted_logits, sorted_idx = x.sort(dim=-1, descending=True)
-            cumprobs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
-            sorted_mask = cumprobs > top_p
-            sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
-            sorted_mask[..., 0] = False
-            mask = torch.zeros_like(x, dtype=torch.bool)
-            mask.scatter_(-1, sorted_idx, sorted_mask)
-            x = x.masked_fill(mask, float("-inf"))
+        temperatures, top_ks, top_ps = self._expand_audio_sampling_params(sampling_metadata, x, num_codebooks)
+        top_ks = torch.where(top_ks <= 0, torch.full_like(top_ks, x.shape[-1]), top_ks)
+        top_ks = top_ks.clamp(max=x.shape[-1])
+        top_ps = top_ps.clamp(min=0.0, max=1.0)
+        temperatures = torch.where(torch.isfinite(temperatures), temperatures, torch.ones_like(temperatures))
+        top_ps = torch.where(torch.isfinite(top_ps), top_ps, torch.ones_like(top_ps))
+        greedy = temperatures <= 0
+        safe_temperatures = torch.where(greedy, torch.ones_like(temperatures), temperatures)
+        x.div_(safe_temperatures.unsqueeze(-1))
+        x = apply_top_k_top_p(x, top_ks, top_ps)
         # Detect all-masked rows BEFORE softmax (softmax of all-inf yields NaN).
         has_finite = torch.isfinite(x).any(dim=-1)
         all_masked = ~has_finite
@@ -1362,7 +1387,8 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
         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)
-        return torch.where(all_masked, fallback, sampled)
+        sampled = torch.where(all_masked, fallback, sampled)
+        return torch.where(greedy, logits_2d.argmax(dim=-1), sampled)
 
     # ------------------------------------------------------------------ omni output
     def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: