From 020792e6be54cc2d06593bae900453f54cad759d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E9=91=AB?= <77673148+l-wave@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:41:54 +0800 Subject: [PATCH] [Perf][Qwen3-TTS] Sample code predictor outputs via Gumbel-max (#4261) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 林鑫 <77673148+l-wave@users.noreply.github.com> Co-authored-by: yurain --- .../qwen3_tts/test_code_predictor_dtype.py | 92 +++++++++++++++++++ .../models/common/qwen3_code_predictor.py | 51 ++++++++-- 2 files changed, 136 insertions(+), 7 deletions(-) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index b73c15fc7..519940ec3 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -18,8 +18,11 @@ import types import pytest import torch +import torch.nn.functional as F from pytest_mock import MockerFixture +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + # Direct file import to avoid vllm_omni.__init__ patch dependencies. _MODELS = os.path.join( os.path.dirname(__file__), @@ -542,3 +545,92 @@ class TestCodePredictorWrapperConfig: wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True), ) assert graph_wrapper._prefix_graphs_enabled is True + + +class TestGumbelMaxSampling: + """Sanity tests for the shared Gumbel-max sampling helper.""" + + @staticmethod + def _sample(logits: torch.Tensor, generator=None) -> torch.Tensor: + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + return common_mod.CodePredictorWrapper._sample_codes_gumbel(logits, generator=generator) + + def test_gumbel_max_picks_unique_high_logit(self, loaded_target_classes) -> None: + """If one logit is dominant, Gumbel-max must always pick it.""" + torch.manual_seed(0) + bsz, vocab = 4, 64 + logits = torch.full((bsz, vocab), -1e4) + logits[:, 7] = 0.0 # only finite entry + out = self._sample(logits) + assert out.shape == (bsz, 1) + assert torch.equal(out.squeeze(-1), torch.full((bsz,), 7, dtype=torch.long)) + + def test_gumbel_max_handles_single_finite_entry_rows(self, loaded_target_classes) -> None: + """Degenerate rows with one surviving finite entry should still return + a valid index instead of tripping multinomial-style asserts.""" + bsz, vocab = 2, 8 + logits = torch.full((bsz, vocab), float("-inf")) + logits[0, 3] = 0.0 + logits[1, 5] = 0.0 + out = self._sample(logits) + assert out[0].item() == 3 + assert out[1].item() == 5 + + def test_gumbel_max_generator_makes_results_reproducible(self, loaded_target_classes) -> None: + """Same ``Generator.manual_seed`` should produce identical samples.""" + torch.manual_seed(0) + logits = torch.randn(8, 32) + g1 = torch.Generator() + g1.manual_seed(123) + g2 = torch.Generator() + g2.manual_seed(123) + g3 = torch.Generator() + g3.manual_seed(456) + s1 = self._sample(logits, generator=g1) + s2 = self._sample(logits, generator=g2) + s3 = self._sample(logits, generator=g3) + assert torch.equal(s1, s2) + # Different seeds must (almost surely) yield a different sequence. + assert not torch.equal(s1, s3) + + def test_gumbel_max_per_row_generators_match_individual_rows(self, loaded_target_classes) -> None: + """A per-row generator list should preserve row independence.""" + torch.manual_seed(11) + logits = torch.randn(3, 16) + seeds = [101, 202, 303] + batched_generators = [] + expected_rows = [] + for row, seed in enumerate(seeds): + row_generator = torch.Generator() + row_generator.manual_seed(seed) + batched_generators.append(row_generator) + + solo_generator = torch.Generator() + solo_generator.manual_seed(seed) + expected_rows.append(self._sample(logits[row : row + 1], generator=solo_generator)) + + batched = self._sample(logits, generator=batched_generators) + expected = torch.cat(expected_rows, dim=0) + assert torch.equal(batched, expected) + + def test_gumbel_max_rejects_generator_length_mismatch(self, loaded_target_classes) -> None: + logits = torch.randn(2, 8) + generators = [torch.Generator()] + with pytest.raises(ValueError, match="Expected 2 per-row generators"): + self._sample(logits, generator=generators) + + def test_gumbel_max_distribution_matches_softmax(self, loaded_target_classes) -> None: + """Empirical frequency under Gumbel-max should be close to the + softmax probability mass. Loose tolerance because we draw a + finite number of samples.""" + torch.manual_seed(7) + vocab = 16 + logits = torch.randn(vocab) * 2.0 + target = F.softmax(logits, dim=-1, dtype=torch.float32) + + n_draws = 20000 + batched_logits = logits.unsqueeze(0).expand(n_draws, vocab).contiguous() + samples = self._sample(batched_logits).squeeze(-1) + empirical = torch.bincount(samples, minlength=vocab).float() / n_draws + # L1 distance under 0.03 with 20k draws is comfortable for vocab=16. + assert (empirical - target).abs().sum().item() < 0.05 diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index c8591c32e..a8a259504 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -28,6 +28,9 @@ from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +_GeneratorLike = torch.Generator | Sequence[torch.Generator | None] | None +_UNIFORM_EPS = 1e-20 + # =================================================================== # HF-numerics-compatible layers for code predictor @@ -605,6 +608,30 @@ class CodePredictorWrapper(nn.Module): values.add(parsed) return values + @staticmethod + def _normalize_generators( + generator: _GeneratorLike, batch_size: int + ) -> torch.Generator | list[torch.Generator | None] | None: + if generator is None or isinstance(generator, torch.Generator): + return generator + + row_generators = list(generator) + if len(row_generators) != batch_size: + raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.") + return row_generators + + @classmethod + def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor: + """Sample ``logits`` via Gumbel-max with optional per-row generators.""" + row_generators = cls._normalize_generators(generator, int(logits.shape[0])) + u = torch.empty_like(logits, dtype=torch.float32) + if isinstance(row_generators, list): + for row, row_generator in enumerate(row_generators): + u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator) + else: + u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators) + return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True) + def _prefix_seq_lens(self, max_seq: int) -> list[int]: all_seq_lens = list(range(2, max_seq)) if not self._prefix_graph_seq_lens: @@ -780,6 +807,7 @@ class CodePredictorWrapper(nn.Module): bsz = int(layer0_code.shape[0]) if generators is not None and len(generators) != bsz: raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}") + sample_generator: _GeneratorLike = generators if generators is not None else generator num_groups = self._num_groups device = layer0_code.device @@ -856,9 +884,20 @@ class CodePredictorWrapper(nn.Module): logits = lm_heads[step - 1](hidden_out[:bsz, step, :]) - # Sample next code + # Sample next code via Gumbel-max. + # + # ``argmax_i(logits_i + Gumbel_i)`` with + # ``Gumbel_i = -log(-log(u_i)), u_i ~ Uniform(0, 1)`` is + # distributionally identical to sampling from ``softmax(logits)``. + # In this file the motivations are practical rather than graph + # related: it is measurably cheaper than ``softmax + multinomial`` + # on the B x 2048 shapes used here, it stays well-defined for + # degenerate masked rows with a surviving finite entry (and is more + # defensive than ``multinomial`` around fully-masked/NaN inputs), + # and the helper below can honor either one batch generator or one + # generator per seeded row. if stored_mode: - # "stored" mode: top-k -> top-p -> softmax -> multinomial + # "stored" mode: top-k -> top-p -> Gumbel-max if s_top_k > 0: topk_vals, _ = logits.topk(s_top_k, dim=-1) logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf")) @@ -869,17 +908,15 @@ class CodePredictorWrapper(nn.Module): remove_mask = (cumulative_probs - sorted_probs) >= s_top_p sorted_logits[remove_mask] = float("-inf") logits = sorted_logits.scatter(1, sorted_idx, sorted_logits) - probs = F.softmax(logits, dim=-1, dtype=torch.float32) - code = self._multinomial(probs, generator, generators) + code = self._sample_codes_gumbel(logits, generator=sample_generator) else: - # "per_call" mode: temperature-scaled + top-k + # "per_call" mode: temperature-scaled + top-k -> Gumbel-max if use_sampling: scaled = logits * inv_temperature if top_k > 0: topk_vals, _ = scaled.topk(top_k, dim=-1) scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) - probs = F.softmax(scaled, dim=-1, dtype=torch.float32) - code = self._multinomial(probs, generator, generators) + code = self._sample_codes_gumbel(scaled, generator=sample_generator) else: code = logits.argmax(dim=-1, keepdim=True)