[Bugfix][CPU] Fix fused Gumbel-max wrapping-window bias (#60221)

Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
Co-authored-by: Bob <bob@ibm.com>

.buildkite/hardware_tests/cpu.yaml

@@ -23,6 +23,7 @@ steps:
   - tests/test_zentorch_version_pin.py
   - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
   - tests/kernels/test_onednn.py
+  - tests/kernels/test_cpu_fused_sampling.py
   - tests/kernels/test_awq_int4_to_int8.py
   - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
   - tests/kernels/quantization/test_zen_da8w4.py
@@ -49,6 +50,7 @@ steps:
       tests/test_zentorch_version_pin.py \
       tests/kernels/mamba/test_cpu_short_conv.py \
       tests/kernels/test_onednn.py \
+      tests/kernels/test_cpu_fused_sampling.py \
       tests/kernels/test_awq_int4_to_int8.py \
       tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \
       tests/kernels/quantization/test_zen_da8w4.py \

.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh

@@ -43,6 +43,7 @@ kernel_tests() {
   docker exec "$CONTAINER_NAME" bash -c "
     set -e
     pytest -x -v -s tests/kernels/test_onednn.py
+    pytest -x -v -s tests/kernels/test_cpu_fused_sampling.py
     pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
     pytest -x -v -s tests/kernels/core/test_cpu_activation.py
     pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py

.buildkite/test_areas/kernels.yaml

@@ -256,6 +256,7 @@ steps:
         --ignore=kernels/test_fused_qk_norm_rope_gate.py
         --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
         --ignore=kernels/test_top_k_per_row.py
+        --ignore=kernels/test_cpu_fused_sampling.py
         --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py
         --ignore=kernels/test_bf16x3_router_gemm_rocm.py
         --ignore=kernels/test_ll_bf16_gemm.py

csrc/cpu/sampling_kernels.cpp

@@ -1,82 +1,73 @@
 #include "cpu_types.hpp"
 
-#include <torch/library.h>
+#include <torch/all.h>
 
 #include <cmath>
 #include <cstdint>
 #include <limits>
-#include <mutex>
-#include <random>
 
 namespace {
 
-constexpr int GUMBEL_TABLE_SIZE = 1 << 20;
-constexpr int GUMBEL_TABLE_MASK = GUMBEL_TABLE_SIZE - 1;
+// Gumbel-max: argmax_i(logit_i + g_i) with g_i ~ Gumbel(0, 1) matches
+// softmax. Noise is SplitMix64(seed, i), so tokens are independent and
+// consecutive seeds are not a 1-token table shift (vllm#59786).
+constexpr uint64_t SPLITMIX_GAMMA = 0x9E3779B97F4A7C15ULL;
+constexpr double GUMBEL_MAX = 37.5;  // 53-bit u in (0,1) => -log(-log u) < 37.5
 
-static float g_gumbel_table[GUMBEL_TABLE_SIZE];
-static std::once_flag g_gumbel_init_flag;
+static inline uint64_t splitmix64(uint64_t z) {
+  z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
+  z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
+  return z ^ (z >> 31);
+}
 
-static void init_gumbel_table() {
-  std::mt19937 rng(42);
-  std::uniform_real_distribution<float> uniform(1.0e-7f, 1.0f - 1.0e-7f);
-  for (int i = 0; i < GUMBEL_TABLE_SIZE; ++i) {
-    float u = uniform(rng);
-    g_gumbel_table[i] = -std::log(-std::log(u));
+static void fused_gumbel_argmax_row(int64_t* __restrict__ output,
+                                    const float* __restrict__ row,
+                                    uint64_t seed, int64_t vocab_size) {
+  const uint64_t key = splitmix64(seed);
+
+  // Compare in (logit, noise) space: x + g > best_x + best_g is
+  // (x - best_x) + g > best_g. Adding g into a 1e38 logit loses it in
+  // double ULP, so ties at the float max would always keep the lower index.
+  double best_logit = 0.0;
+  double best_g = 0.0;
+  int64_t best_idx = 0;
+  bool have_best = false;
+  for (int64_t i = 0; i < vocab_size; ++i) {
+    const double x = row[i];
+    if (x == -std::numeric_limits<double>::infinity()) continue;
+    if (have_best && (x - best_logit) + GUMBEL_MAX <= best_g) continue;
+    const uint64_t bits =
+        splitmix64(key + uint64_t(i + 1) * SPLITMIX_GAMMA) >> 11;
+    const double one_minus_u =
+        (double((1ULL << 53) - 1 - bits) + 0.5) * 0x1.0p-53;
+    if (have_best && one_minus_u >= std::exp((x - best_logit) - best_g))
+      continue;
+    const double u = (double(bits) + 0.5) * 0x1.0p-53;
+    const double g = -std::log(-std::log(u));
+    if (!have_best || (x - best_logit) + g > best_g) {
+      have_best = true;
+      best_logit = x;
+      best_g = g;
+      best_idx = i;
+    }
   }
+  *output = best_idx;
 }
 
-static inline void ensure_gumbel_table() {
-  std::call_once(g_gumbel_init_flag, init_gumbel_table);
-}
-
-// Fused Gumbel-max kernel
-
 static void fused_gumbel_argmax_kernel(int64_t* __restrict__ output,
                                        const float* __restrict__ logits,
                                        const int64_t* __restrict__ seeds,
                                        const int64_t batch_size,
                                        const int64_t vocab_size) {
-  ensure_gumbel_table();
-  constexpr int VEC_ELEM_NUM = vec_op::FP32Vec8::VEC_ELEM_NUM;
-  const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM);
-
+  if (batch_size <= 0) return;
+  if (batch_size == 1) {
+    fused_gumbel_argmax_row(output, logits, uint64_t(seeds[0]), vocab_size);
+    return;
+  }
 #pragma omp parallel for schedule(static)
   for (int64_t b = 0; b < batch_size; ++b) {
-    const float* row = logits + b * vocab_size;
-    const uint32_t seed = static_cast<uint32_t>(seeds[b]);
-
-    float best_score = -std::numeric_limits<float>::infinity();
-    int64_t best_idx = 0;
-
-    for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) {
-      vec_op::FP32Vec8 logit_vec(row + i);
-
-      float gumbel_buf[VEC_ELEM_NUM];
-      for (int k = 0; k < VEC_ELEM_NUM; ++k)
-        gumbel_buf[k] = g_gumbel_table[(seed + i + k) & GUMBEL_TABLE_MASK];
-      vec_op::FP32Vec8 gumbel_vec(gumbel_buf);
-
-      vec_op::FP32Vec8 score = logit_vec + gumbel_vec;
-
-      float score_buf[VEC_ELEM_NUM];
-      score.save(score_buf);
-      for (int k = 0; k < VEC_ELEM_NUM; ++k) {
-        if (score_buf[k] > best_score) {
-          best_score = score_buf[k];
-          best_idx = i + k;
-        }
-      }
-    }
-
-    for (int64_t i = vec_end; i < vocab_size; ++i) {
-      float s = row[i] + g_gumbel_table[(seed + i) & GUMBEL_TABLE_MASK];
-      if (s > best_score) {
-        best_score = s;
-        best_idx = i;
-      }
-    }
-
-    output[b] = best_idx;
+    fused_gumbel_argmax_row(output + b, logits + b * vocab_size,
+                            uint64_t(seeds[b]), vocab_size);
   }
 }
 
@@ -91,22 +82,41 @@ static void greedy_argmax_kernel(int64_t* __restrict__ output,
   for (int64_t b = 0; b < batch_size; ++b) {
     const float* row = logits + b * vocab_size;
 
+    // Vector max has no consistent NaN rule, so NaNs are caught via the sum:
+    // any NaN (or inf + -inf, a harmless false positive) makes it NaN.
     vec_op::FP32Vec16 vmax(-std::numeric_limits<float>::infinity());
+    vec_op::FP32Vec16 vsum(0.0f);
     for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) {
-      vmax = vmax.max(vec_op::FP32Vec16(row + i));
+      vec_op::FP32Vec16 v(row + i);
+      vmax = vmax.max(v);
+      vsum = vsum + v;
     }
     float best_val = vmax.reduce_max();
+    bool maybe_nan = std::isnan(vsum.reduce_sum());
     for (int64_t i = vec_end; i < vocab_size; ++i) {
+      maybe_nan |= std::isnan(row[i]);
       if (row[i] > best_val) {
         best_val = row[i];
       }
     }
 
-    int64_t best_idx = 0;
-    for (int64_t i = 0; i < vocab_size; ++i) {
-      if (row[i] == best_val) {
-        best_idx = i;
-        break;
+    int64_t best_idx = -1;
+    if (maybe_nan) {
+      // torch.argmax treats NaN as the maximum: return the first NaN.
+      for (int64_t i = 0; i < vocab_size; ++i) {
+        if (std::isnan(row[i])) {
+          best_idx = i;
+          break;
+        }
+      }
+    }
+    if (best_idx < 0) {
+      best_idx = 0;
+      for (int64_t i = 0; i < vocab_size; ++i) {
+        if (row[i] == best_val) {
+          best_idx = i;
+          break;
+        }
       }
     }
     output[b] = best_idx;
@@ -117,21 +127,28 @@ static void greedy_argmax_kernel(int64_t* __restrict__ output,
 
 torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits,
                                   const torch::Tensor& seeds) {
+  TORCH_CHECK(logits.device().is_cpu(), "logits must be a CPU tensor");
   TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]");
   TORCH_CHECK(logits.scalar_type() == torch::kFloat32,
               "logits must be float32");
+  TORCH_CHECK(logits.size(1) > 0, "vocab_size must be positive");
+  TORCH_CHECK(seeds.device().is_cpu(), "seeds must be a CPU tensor");
+  TORCH_CHECK(seeds.scalar_type() == torch::kInt64, "seeds must be int64");
   TORCH_CHECK(seeds.dim() == 1 && seeds.size(0) == logits.size(0),
               "seeds must be 1-D with batch_size elements");
 
   auto logits_contig = logits.contiguous();
+  auto seeds_contig = seeds.contiguous();
   auto output = torch::empty({logits_contig.size(0)}, torch::kInt64);
-  fused_gumbel_argmax_kernel(
-      output.data_ptr<int64_t>(), logits_contig.data_ptr<float>(),
-      seeds.data_ptr<int64_t>(), logits_contig.size(0), logits_contig.size(1));
+  fused_gumbel_argmax_kernel(output.data_ptr<int64_t>(),
+                             logits_contig.data_ptr<float>(),
+                             seeds_contig.data_ptr<int64_t>(),
+                             logits_contig.size(0), logits_contig.size(1));
   return output;
 }
 
 torch::Tensor greedy_argmax(const torch::Tensor& logits) {
+  TORCH_CHECK(logits.device().is_cpu(), "logits must be a CPU tensor");
   TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]");
   TORCH_CHECK(logits.scalar_type() == torch::kFloat32,
               "logits must be float32");

tests/kernels/test_cpu_fused_sampling.py

@@ -43,20 +43,20 @@ class TestGreedyArgmax:
 
 
 class TestFusedGumbelArgmax:
-    @pytest.mark.parametrize("vocab_size", VOCAB_SIZES)
-    def test_distribution_chi_squared(self, vocab_size: int):
+    def test_distribution_chi_squared(self):
         """Verify sampling distribution via chi-squared goodness of fit."""
-        small_vocab = min(vocab_size, 100)
+        # Small support is enough for chi-squared power; large-vocab bias
+        # is covered by test_noise_not_sliding_window (#59786).
+        small_vocab = 100
         logits = torch.randn(1, small_vocab, dtype=torch.float32)
         probs = logits.softmax(dim=-1).squeeze(0)
 
         n_samples = 100_000
-        counts = torch.zeros(small_vocab)
-        for trial in range(n_samples):
-            seed = torch.tensor([trial * 7 + 13], dtype=torch.long)
-            tile = logits.expand(1, -1).contiguous()
-            idx = torch.ops._C.fused_gumbel_argmax(tile, seed)
-            counts[idx.item()] += 1
+        seeds = torch.arange(n_samples, dtype=torch.long) * 7 + 13
+        idx = torch.ops._C.fused_gumbel_argmax(
+            logits.expand(n_samples, -1).contiguous(), seeds
+        )
+        counts = torch.bincount(idx, minlength=small_vocab).float()
 
         expected = probs * n_samples
         mask = expected > 5
@@ -70,6 +70,39 @@ class TestFusedGumbelArgmax:
             f"dof={dof}, p={p_value:.6f}"
         )
 
+    def test_noise_not_sliding_window(self):
+        """Table[(seed+i)&MASK] made consecutive seeds a 1-token shift.
+
+        On uniform logits that yields P(w(s+1) == w(s)-1) ~ 1-2/V. Independent
+        per-token noise yields ~1/V. See #59786.
+        """
+        vocab, n = 128, 2048
+        logits = torch.zeros(n, vocab, dtype=torch.float32)
+        seeds = torch.arange(n, dtype=torch.long)
+        winners = torch.ops._C.fused_gumbel_argmax(logits, seeds)
+        shifted = (winners[:-1] - 1) % vocab
+        frac = (winners[1:] == shifted).float().mean().item()
+        assert frac < 0.05, (
+            f"consecutive-seed winners follow a table shift (frac={frac:.3f})"
+        )
+
+    def test_zipf_head_matches_softmax(self):
+        """Zipf head probability should match softmax (#59786)."""
+        vocab, n_samples, batch = 8192, 16384, 256
+        logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log()
+        p0 = torch.softmax(logits.double(), dim=-1)[0].item()
+        row = logits.view(1, -1).expand(batch, vocab).contiguous()
+        n_top = 0
+        for start in range(0, n_samples, batch):
+            seeds = torch.arange(start, start + batch, dtype=torch.long)
+            tokens = torch.ops._C.fused_gumbel_argmax(row, seeds)
+            n_top += int((tokens == 0).sum())
+        q0 = n_top / n_samples
+        se = (p0 * (1.0 - p0) / n_samples) ** 0.5
+        assert abs(q0 - p0) < 8.0 * se, (
+            f"P(top-1) kernel={q0:.4f} exact={p0:.4f} se={se:.4f}"
+        )
+
     def test_deterministic_same_seed(self):
         """Same seed produces same result."""
         logits = torch.randn(4, 32000, dtype=torch.float32)
@@ -95,6 +128,108 @@ class TestFusedGumbelArgmax:
         assert result.min() >= 0
         assert result.max() < vocab_size
 
+    def test_samples_only_from_finite_logits(self):
+        """top-k/top-p mask with -inf; those IDs must never win."""
+        allowed = torch.tensor([3, 17, 99, 1024, 8191], dtype=torch.long)
+        logits = torch.full((64, 32000), float("-inf"), dtype=torch.float32)
+        logits[:, allowed] = torch.randn(64, allowed.numel())
+        tokens = torch.ops._C.fused_gumbel_argmax(
+            logits, torch.arange(64, dtype=torch.long)
+        )
+        assert torch.isin(tokens, allowed).all()
+
+    def test_single_unmasked_token(self):
+        logits = torch.full((8, 4096), float("-inf"), dtype=torch.float32)
+        logits[:, 123] = 0.0
+        seeds = torch.tensor([0, -1, -(2**63), 2**63 - 1, 42, 99, 1000, 7])
+        tokens = torch.ops._C.fused_gumbel_argmax(logits, seeds)
+        assert (tokens == 123).all()
+
+    def test_tokens_beyond_gumbel_max_cannot_win(self):
+        """A gap larger than GUMBEL_MAX is unreachable."""
+        logits = torch.zeros(256, 4, dtype=torch.float32)
+        logits[:, 1:] = -80.0
+        tokens = torch.ops._C.fused_gumbel_argmax(
+            logits, torch.arange(256, dtype=torch.long)
+        )
+        assert (tokens == 0).all()
+
+    def test_tied_huge_logits_not_index_biased(self):
+        """Noise must still break ties when logits sit at float32 max."""
+        logits = torch.full(
+            (4096, 2), torch.finfo(torch.float32).max, dtype=torch.float32
+        )
+        tokens = torch.ops._C.fused_gumbel_argmax(
+            logits, torch.arange(4096, dtype=torch.long)
+        )
+        frac0 = (tokens == 0).float().mean().item()
+        assert set(tokens.tolist()) == {0, 1}, tokens.unique().tolist()
+        assert 0.35 < frac0 < 0.65, f"P(index 0)={frac0:.3f} expected ~0.5"
+
+    def test_serial_batch1_matches_batched(self):
+        """OpenMP rows must match the batch=1 serial path."""
+        logits = torch.randn(8, 2048, dtype=torch.float32)
+        seeds = torch.arange(8, dtype=torch.long)
+        batched = torch.ops._C.fused_gumbel_argmax(logits, seeds)
+        serial = torch.stack(
+            [
+                torch.ops._C.fused_gumbel_argmax(logits[i : i + 1], seeds[i : i + 1])[0]
+                for i in range(8)
+            ]
+        )
+        torch.testing.assert_close(batched, serial)
+
+    def test_noncontiguous_matches_contiguous(self):
+        logits = torch.randn(8, 1024, dtype=torch.float32)[:, ::2]
+        seeds = torch.arange(16, dtype=torch.long)[::2]
+        assert not logits.is_contiguous() and not seeds.is_contiguous()
+        expected = torch.ops._C.fused_gumbel_argmax(
+            logits.contiguous(), seeds.contiguous()
+        )
+        actual = torch.ops._C.fused_gumbel_argmax(logits, seeds)
+        torch.testing.assert_close(actual, expected)
+
+    def test_empty_batch(self):
+        tokens = torch.ops._C.fused_gumbel_argmax(
+            torch.empty(0, 7), torch.empty(0, dtype=torch.long)
+        )
+        assert tokens.shape == (0,) and tokens.dtype == torch.int64
+
+    def test_masked_support_matches_softmax(self):
+        """Finite subset after -inf mask should match softmax."""
+        from scipy.stats import chisquare
+
+        logits = torch.tensor(
+            [0.0, -0.2, -0.7, -1.5, -3.0, float("-inf")], dtype=torch.float32
+        )
+        n_samples = 20_000
+        tokens = torch.ops._C.fused_gumbel_argmax(
+            logits.expand(n_samples, -1).contiguous(),
+            torch.arange(n_samples, dtype=torch.long),
+        )
+        counts = torch.bincount(tokens, minlength=logits.numel()).double()
+        assert counts[-1] == 0
+        probs = logits.double().softmax(-1)
+        _, p_value = chisquare(counts[:-1].numpy(), (probs[:-1] * n_samples).numpy())
+        assert p_value > 1e-7, f"masked-support p-value: {p_value}"
+
+    def test_zipf_tail_matches_softmax(self):
+        """Independent 53-bit Gumbel must not starve the Zipf tail (#59786)."""
+        vocab, n_samples, batch = 8192, 16384, 256
+        logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log()
+        p_tail = torch.softmax(logits.double(), dim=-1)[vocab // 2 :].sum().item()
+        row = logits.view(1, -1).expand(batch, vocab).contiguous()
+        n_tail = 0
+        for start in range(0, n_samples, batch):
+            seeds = torch.arange(start, start + batch, dtype=torch.long)
+            tokens = torch.ops._C.fused_gumbel_argmax(row, seeds)
+            n_tail += int((tokens >= vocab // 2).sum())
+        q_tail = n_tail / n_samples
+        se = (p_tail * (1.0 - p_tail) / n_samples) ** 0.5
+        assert abs(q_tail - p_tail) < 8.0 * se, (
+            f"P(tail) kernel={q_tail:.4f} exact={p_tail:.4f} se={se:.4f}"
+        )
+
 
 class TestMemory:
     def test_fused_no_intermediate_allocs(self):

vllm/v1/sample/ops/topk_topp_sampler.py

@@ -319,11 +319,11 @@ class TopKTopPSampler(nn.Module):
         k: torch.Tensor | None,
         p: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Fused Gumbel-max sampling for CPU.
+        """CPU Gumbel-max sampling with per-token SplitMix64 noise.
 
-        Uses a precomputed Gumbel table + single-pass argmax over logits,
-        skipping softmax and intermediate allocations entirely.
-        Falls back to the native path when fp64 Gumbel noise is requested.
+        Independent (seed, token) streams avoid the wrapping-window bias
+        (#59786). Falls back to native sampling when fp64 Gumbel noise
+        is requested.
         """
         logits = apply_top_k_top_p(logits, k, p)
         logits_to_return = None