[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