From 76e4a78afc17d5fac34c8e5be83531b9e5daf432 Mon Sep 17 00:00:00 2001 From: boatman <109857087+sphinxkkkbc@users.noreply.github.com> Date: Sat, 25 Jul 2026 23:25:35 +0800 Subject: [PATCH] [Bugfix][Higgs-Audio-V3] Propagate Sampling Parameters to Sampler (#5406) Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> --- .../higgs_audio_v3/test_higgs_audio_v3.py | 93 ++++++++++++++++++- tests/worker/test_gpu_ar_model_runner.py | 80 ++++++++++++++++ .../entrypoints/openai/serving_speech.py | 3 + .../higgs_audio_v3/higgs_audio_v3_talker.py | 60 ++++++++---- 4 files changed, 215 insertions(+), 21 deletions(-) diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py index b92b0095b..1d4bdb48d 100755 --- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py +++ b/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) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 83d270da4..daab10ba0 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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()) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index ca3b8c21f..2bb402caf 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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) 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 f51adafa9..579d0cd13 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 @@ -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: