[Perf][Qwen3-TTS] Sample code predictor outputs via Gumbel-max (#4261)

Signed-off-by: 林鑫 <77673148+l-wave@users.noreply.github.com>
Co-authored-by: yurain <yurain26@users.noreply.github.com>

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

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)