diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 3fdce975d3..d54688b574 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.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 \ diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 689ca92d6a..1f598b277b 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.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 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index cfc628d382..d3c8f65790 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.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 diff --git a/csrc/cpu/sampling_kernels.cpp b/csrc/cpu/sampling_kernels.cpp index 3e2f00a0ba..1d30b92a8d 100644 --- a/csrc/cpu/sampling_kernels.cpp +++ b/csrc/cpu/sampling_kernels.cpp @@ -1,82 +1,73 @@ #include "cpu_types.hpp" -#include +#include #include #include #include -#include -#include 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 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::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(seeds[b]); - - float best_score = -std::numeric_limits::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::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(), logits_contig.data_ptr(), - seeds.data_ptr(), logits_contig.size(0), logits_contig.size(1)); + fused_gumbel_argmax_kernel(output.data_ptr(), + logits_contig.data_ptr(), + seeds_contig.data_ptr(), + 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"); diff --git a/tests/kernels/test_cpu_fused_sampling.py b/tests/kernels/test_cpu_fused_sampling.py index 8a3857fe9e..cdb10ab7e3 100644 --- a/tests/kernels/test_cpu_fused_sampling.py +++ b/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): diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index a3084a09c2..daa94ebc3a 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/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