[CPU][Perf] Add vectorized Sampler Kernel (#53913)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Rehan Khan committed
Sep 30, 2026 at 12:46 UTC
583637f126a5f5193952a81ca9e0630ae9fc9986
7 files changed
+529
-12
benchmarks/kernels/bench_cpu_sampling.py
new
+217
@@ -0,0 +1,217 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Micro-benchmark for CPU sampling kernels.
4
+
5
+Compares fused Gumbel-max / greedy argmax against the baseline
6
+(softmax → exponential → div → argmax) across vocab and batch sizes.
7
+
8
+Usage:
9
+ .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py
10
+ .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py --profile
11
+ .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py --vocab 128256 --batch 16
12
+"""
13
+
14
+import argparse
15
+import time
16
+
17
+import torch
18
+import vllm._C # noqa: F401
19
+from torch.profiler import ProfilerActivity, profile, record_function
20
+
21
+
22
+def baseline_random_sample(logits: torch.Tensor) -> torch.Tensor:
23
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
24
+ q = torch.empty_like(probs)
25
+ q.exponential_()
26
+ return probs.div(q).argmax(dim=-1).view(-1)
27
+
28
+
29
+def baseline_greedy_sample(logits: torch.Tensor) -> torch.Tensor:
30
+ return logits.argmax(dim=-1).view(-1)
31
+
32
+
33
+def bench_latency(fn, args, n_warmup=20, n_iters=500):
34
+ for _ in range(n_warmup):
35
+ fn(*args)
36
+
37
+ t0 = time.perf_counter()
38
+ for _ in range(n_iters):
39
+ fn(*args)
40
+ elapsed = time.perf_counter() - t0
41
+ return elapsed / n_iters * 1e6 # µs
42
+
43
+
44
+def run_profile(logits, seeds, n_iters=50):
45
+ """Run torch.profiler and print comparison tables."""
46
+ # Profile baseline random sampling
47
+ with profile(
48
+ activities=[ProfilerActivity.CPU],
49
+ record_shapes=True,
50
+ ) as prof_base_rand:
51
+ for _ in range(n_iters):
52
+ with record_function("baseline_random"):
53
+ baseline_random_sample(logits)
54
+
55
+ # Profile fused Gumbel-max
56
+ with profile(
57
+ activities=[ProfilerActivity.CPU],
58
+ record_shapes=True,
59
+ ) as prof_fused:
60
+ for _ in range(n_iters):
61
+ with record_function("fused_gumbel_argmax"):
62
+ torch.ops._C.fused_gumbel_argmax(logits, seeds)
63
+
64
+ # Profile baseline greedy
65
+ with profile(
66
+ activities=[ProfilerActivity.CPU],
67
+ record_shapes=True,
68
+ ) as prof_base_grdy:
69
+ for _ in range(n_iters):
70
+ with record_function("baseline_greedy"):
71
+ baseline_greedy_sample(logits)
72
+
73
+ # Profile custom greedy
74
+ with profile(
75
+ activities=[ProfilerActivity.CPU],
76
+ record_shapes=True,
77
+ ) as prof_cust_grdy:
78
+ for _ in range(n_iters):
79
+ with record_function("custom_greedy_argmax"):
80
+ torch.ops._C.greedy_argmax(logits)
81
+
82
+ B, V = logits.shape
83
+ print(f"\n{'=' * 80}")
84
+ print(f"torch.profiler breakdown (batch={B}, vocab={V}, iters={n_iters})")
85
+ print(f"{'=' * 80}")
86
+
87
+ print("\n--- Baseline Random (softmax → exp → div → argmax) ---")
88
+ print(prof_base_rand.key_averages().table(sort_by="cpu_time_total", row_limit=15))
89
+
90
+ print("\n--- Fused Gumbel-max (table lookup + add + argmax) ---")
91
+ print(prof_fused.key_averages().table(sort_by="cpu_time_total", row_limit=15))
92
+
93
+ print("\n--- Baseline Greedy (torch.argmax) ---")
94
+ print(prof_base_grdy.key_averages().table(sort_by="cpu_time_total", row_limit=15))
95
+
96
+ print("\n--- Custom Greedy (vec_op argmax) ---")
97
+ print(prof_cust_grdy.key_averages().table(sort_by="cpu_time_total", row_limit=15))
98
+
99
+ # Summary comparison
100
+ def avg_us(prof, label):
101
+ for e in prof.key_averages():
102
+ if e.key == label:
103
+ return e.cpu_time_total / e.count
104
+ return 0.0
105
+
106
+ t_br = avg_us(prof_base_rand, "baseline_random")
107
+ t_fg = avg_us(prof_fused, "fused_gumbel_argmax")
108
+ t_bg = avg_us(prof_base_grdy, "baseline_greedy")
109
+ t_cg = avg_us(prof_cust_grdy, "custom_greedy_argmax")
110
+
111
+ print(f"\n{'=' * 60}")
112
+ print(f" Summary (batch={B}, vocab={V})")
113
+ print(f"{'=' * 60}")
114
+ print(f" {'Kernel':<30} {'Avg (µs)':>10} {'Speedup':>10}")
115
+ print(f" {'-' * 50}")
116
+ print(f" {'baseline random':<30} {t_br:>10.1f} {'—':>10}")
117
+ print(
118
+ f" {'fused gumbel-max':<30} {t_fg:>10.1f} "
119
+ f"{t_br / t_fg if t_fg > 0 else 0:>9.2f}x"
120
+ )
121
+ print(f" {'baseline greedy':<30} {t_bg:>10.1f} {'—':>10}")
122
+ print(
123
+ f" {'custom greedy':<30} {t_cg:>10.1f} {t_bg / t_cg if t_cg > 0 else 0:>9.2f}x"
124
+ )
125
+ print(f"{'=' * 60}")
126
+
127
+
128
+def main():
129
+ parser = argparse.ArgumentParser(description="Benchmark CPU sampling kernels")
130
+ parser.add_argument(
131
+ "--vocab",
132
+ type=int,
133
+ nargs="+",
134
+ default=[32000, 49152, 128256],
135
+ help="Vocab sizes to benchmark",
136
+ )
137
+ parser.add_argument(
138
+ "--batch",
139
+ type=int,
140
+ nargs="+",
141
+ default=[1, 4, 16],
142
+ help="Batch sizes to benchmark",
143
+ )
144
+ parser.add_argument(
145
+ "--profile",
146
+ action="store_true",
147
+ help="Run torch.profiler and export chrome trace",
148
+ )
149
+ parser.add_argument(
150
+ "--iters", type=int, default=500, help="Iterations per measurement"
151
+ )
152
+ args = parser.parse_args()
153
+
154
+ header = (
155
+ f"{'batch':>5} {'vocab':>7} "
156
+ f"{'base_rand':>10} {'fused_rand':>10} {'rand_spdup':>10} "
157
+ f"{'base_grdy':>10} {'cust_grdy':>10} {'grdy_spdup':>10}"
158
+ )
159
+ units = (
160
+ f"{'':>5} {'':>7} "
161
+ f"{'(µs)':>10} {'(µs)':>10} {'':>10} "
162
+ f"{'(µs)':>10} {'(µs)':>10} {'':>10}"
163
+ )
164
+ print("\n" + "=" * len(header))
165
+ print("CPU Sampling Kernel Benchmark")
166
+ print("=" * len(header))
167
+ print(header)
168
+ print(units)
169
+ print("-" * len(header))
170
+
171
+ for V in args.vocab:
172
+ for B in args.batch:
173
+ logits = torch.randn(B, V, dtype=torch.float32)
174
+ seeds = torch.arange(B, dtype=torch.long)
175
+
176
+ t_base_rand = bench_latency(
177
+ baseline_random_sample, (logits,), n_iters=args.iters
178
+ )
179
+ t_fused = bench_latency(
180
+ torch.ops._C.fused_gumbel_argmax, (logits, seeds), n_iters=args.iters
181
+ )
182
+
183
+ t_base_grdy = bench_latency(
184
+ baseline_greedy_sample, (logits,), n_iters=args.iters
185
+ )
186
+ t_cust_grdy = bench_latency(
187
+ torch.ops._C.greedy_argmax, (logits,), n_iters=args.iters
188
+ )
189
+
190
+ rand_speedup = t_base_rand / t_fused if t_fused > 0 else 0
191
+ grdy_speedup = t_base_grdy / t_cust_grdy if t_cust_grdy > 0 else 0
192
+
193
+ print(
194
+ f"{B:>5} {V:>7} "
195
+ f"{t_base_rand:>10.1f} {t_fused:>10.1f} "
196
+ f"{rand_speedup:>9.2f}x "
197
+ f"{t_base_grdy:>10.1f} {t_cust_grdy:>10.1f} "
198
+ f"{grdy_speedup:>9.2f}x"
199
+ )
200
+
201
+ print("-" * len(header))
202
+
203
+ if args.profile:
204
+ print("\nRunning torch.profiler (batch=16, vocab=128256) ...")
205
+ logits = torch.randn(16, 128256, dtype=torch.float32)
206
+ seeds = torch.arange(16, dtype=torch.long)
207
+
208
+ # warmup
209
+ for _ in range(10):
210
+ baseline_random_sample(logits)
211
+ torch.ops._C.fused_gumbel_argmax(logits, seeds)
212
+
213
+ run_profile(logits, seeds)
214
+
215
+
216
+if __name__ == "__main__":
217
+ main()
cmake/cpu_extension.cmake
+3
@@ -501,6 +501,7 @@ set(VLLM_EXT_SRC
501
"csrc/cpu/cpu_fused_moe.cpp"
502
"csrc/cpu/cpu_attn.cpp"
503
"csrc/cpu/cpu_isa.cpp"
504
+ "csrc/cpu/sampling_kernels.cpp"
505
"csrc/cpu/torch_bindings.cpp")
506
507
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
@@ -589,6 +590,7 @@ if (ENABLE_X86_ISA)
590
"csrc/cpu/cpu_isa.cpp"
591
"csrc/cpu/dnnl_kernels.cpp"
592
"csrc/cpu/mamba_cpu.cpp"
593
+ "csrc/cpu/sampling_kernels.cpp"
594
"csrc/cpu/torch_bindings.cpp"
595
# TODO: Remove these files
596
"csrc/cpu/activation.cpp"
@@ -606,6 +608,7 @@ if (ENABLE_X86_ISA)
608
"csrc/cpu/cpu_isa.cpp"
609
"csrc/cpu/mamba_cpu.cpp"
610
"csrc/cpu/dnnl_kernels.cpp"
611
+ "csrc/cpu/sampling_kernels.cpp"
612
"csrc/cpu/torch_bindings.cpp"
613
# TODO: Remove these files
614
"csrc/cpu/activation.cpp"
csrc/cpu/sampling_kernels.cpp
new
+151
@@ -0,0 +1,151 @@
1
+#include "cpu_types.hpp"
2
+
3
+#include <torch/library.h>
4
+
5
+#include <cmath>
6
+#include <cstdint>
7
+#include <limits>
8
+#include <mutex>
9
+#include <random>
10
+
11
+namespace {
12
+
13
+constexpr int GUMBEL_TABLE_SIZE = 1 << 20;
14
+constexpr int GUMBEL_TABLE_MASK = GUMBEL_TABLE_SIZE - 1;
15
+
16
+static float g_gumbel_table[GUMBEL_TABLE_SIZE];
17
+static std::once_flag g_gumbel_init_flag;
18
+
19
+static void init_gumbel_table() {
20
+ std::mt19937 rng(42);
21
+ std::uniform_real_distribution<float> uniform(1.0e-7f, 1.0f - 1.0e-7f);
22
+ for (int i = 0; i < GUMBEL_TABLE_SIZE; ++i) {
23
+ float u = uniform(rng);
24
+ g_gumbel_table[i] = -std::log(-std::log(u));
25
+ }
26
+}
27
+
28
+static inline void ensure_gumbel_table() {
29
+ std::call_once(g_gumbel_init_flag, init_gumbel_table);
30
+}
31
+
32
+// Fused Gumbel-max kernel
33
+
34
+static void fused_gumbel_argmax_kernel(int64_t* __restrict__ output,
35
+ const float* __restrict__ logits,
36
+ const int64_t* __restrict__ seeds,
37
+ const int64_t batch_size,
38
+ const int64_t vocab_size) {
39
+ ensure_gumbel_table();
40
+ constexpr int VEC_ELEM_NUM = vec_op::FP32Vec8::VEC_ELEM_NUM;
41
+ const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM);
42
+
43
+#pragma omp parallel for schedule(static)
44
+ for (int64_t b = 0; b < batch_size; ++b) {
45
+ const float* row = logits + b * vocab_size;
46
+ const uint32_t seed = static_cast<uint32_t>(seeds[b]);
47
+
48
+ float best_score = -std::numeric_limits<float>::infinity();
49
+ int64_t best_idx = 0;
50
+
51
+ for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) {
52
+ vec_op::FP32Vec8 logit_vec(row + i);
53
+
54
+ float gumbel_buf[VEC_ELEM_NUM];
55
+ for (int k = 0; k < VEC_ELEM_NUM; ++k)
56
+ gumbel_buf[k] = g_gumbel_table[(seed + i + k) & GUMBEL_TABLE_MASK];
57
+ vec_op::FP32Vec8 gumbel_vec(gumbel_buf);
58
+
59
+ vec_op::FP32Vec8 score = logit_vec + gumbel_vec;
60
+
61
+ float score_buf[VEC_ELEM_NUM];
62
+ score.save(score_buf);
63
+ for (int k = 0; k < VEC_ELEM_NUM; ++k) {
64
+ if (score_buf[k] > best_score) {
65
+ best_score = score_buf[k];
66
+ best_idx = i + k;
67
+ }
68
+ }
69
+ }
70
+
71
+ for (int64_t i = vec_end; i < vocab_size; ++i) {
72
+ float s = row[i] + g_gumbel_table[(seed + i) & GUMBEL_TABLE_MASK];
73
+ if (s > best_score) {
74
+ best_score = s;
75
+ best_idx = i;
76
+ }
77
+ }
78
+
79
+ output[b] = best_idx;
80
+ }
81
+}
82
+
83
+static void greedy_argmax_kernel(int64_t* __restrict__ output,
84
+ const float* __restrict__ logits,
85
+ const int64_t batch_size,
86
+ const int64_t vocab_size) {
87
+ constexpr int VEC_ELEM_NUM = vec_op::FP32Vec16::VEC_ELEM_NUM;
88
+ const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM);
89
+
90
+#pragma omp parallel for schedule(static)
91
+ for (int64_t b = 0; b < batch_size; ++b) {
92
+ const float* row = logits + b * vocab_size;
93
+
94
+ vec_op::FP32Vec16 vmax(-std::numeric_limits<float>::infinity());
95
+ for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) {
96
+ vmax = vmax.max(vec_op::FP32Vec16(row + i));
97
+ }
98
+ float best_val = vmax.reduce_max();
99
+ for (int64_t i = vec_end; i < vocab_size; ++i) {
100
+ if (row[i] > best_val) {
101
+ best_val = row[i];
102
+ }
103
+ }
104
+
105
+ int64_t best_idx = 0;
106
+ for (int64_t i = 0; i < vocab_size; ++i) {
107
+ if (row[i] == best_val) {
108
+ best_idx = i;
109
+ break;
110
+ }
111
+ }
112
+ output[b] = best_idx;
113
+ }
114
+}
115
+
116
+} // namespace
117
+
118
+torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits,
119
+ const torch::Tensor& seeds) {
120
+ TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]");
121
+ TORCH_CHECK(logits.scalar_type() == torch::kFloat32,
122
+ "logits must be float32");
123
+ TORCH_CHECK(seeds.dim() == 1 && seeds.size(0) == logits.size(0),
124
+ "seeds must be 1-D with batch_size elements");
125
+
126
+ auto logits_contig = logits.contiguous();
127
+ auto output = torch::empty({logits_contig.size(0)}, torch::kInt64);
128
+ fused_gumbel_argmax_kernel(
129
+ output.data_ptr<int64_t>(), logits_contig.data_ptr<float>(),
130
+ seeds.data_ptr<int64_t>(), logits_contig.size(0), logits_contig.size(1));
131
+ return output;
132
+}
133
+
134
+torch::Tensor greedy_argmax(const torch::Tensor& logits) {
135
+ TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]");
136
+ TORCH_CHECK(logits.scalar_type() == torch::kFloat32,
137
+ "logits must be float32");
138
+
139
+ // A single row is faster in ATen: OpenMP fork/join dominates the scan
140
+ // (~1 ms vs ~50 µs on s390x at vocab=32k).
141
+ if (logits.size(0) <= 1) {
142
+ return logits.argmax(/*dim=*/-1);
143
+ }
144
+
145
+ auto logits_contig = logits.contiguous();
146
+ auto output = torch::empty({logits_contig.size(0)}, torch::kInt64);
147
+ greedy_argmax_kernel(output.data_ptr<int64_t>(),
148
+ logits_contig.data_ptr<float>(), logits_contig.size(0),
149
+ logits_contig.size(1));
150
+ return output;
151
+}
csrc/cpu/torch_bindings.cpp
+12
@@ -393,6 +393,11 @@ void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
393
const c10::optional<at::Tensor>& z,
394
const at::Tensor& cu_seqlens);
395
396
+torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits,
397
+ const torch::Tensor& seeds);
398
+
399
+torch::Tensor greedy_argmax(const torch::Tensor& logits);
400
+
401
void init_cpu_memory_env(std::vector<int64_t> node_ids);
402
403
namespace cpu_utils {
@@ -969,6 +974,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
974
975
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
976
977
+ // Fused sampling kernels
978
+ ops.def("fused_gumbel_argmax(Tensor logits, Tensor seeds) -> Tensor");
979
+ ops.impl("fused_gumbel_argmax", torch::kCPU, &fused_gumbel_argmax);
980
+
981
+ ops.def("greedy_argmax(Tensor logits) -> Tensor");
982
+ ops.impl("greedy_argmax", torch::kCPU, &greedy_argmax);
983
+
984
// Speculative decoding kernels
985
ops.def(
986
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
tests/kernels/test_cpu_fused_sampling.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import tracemalloc
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.platforms import current_platform
10
+
11
+if not current_platform.is_cpu():
12
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
13
+
14
+import vllm._C # noqa: F401, E402
15
+
16
+VOCAB_SIZES = [32000, 49152, 128256]
17
+BATCH_SIZES = [1, 4, 16]
18
+
19
+
20
+def reference_greedy_sample(logits: torch.Tensor) -> torch.Tensor:
21
+ return logits.argmax(dim=-1).view(-1)
22
+
23
+
24
+class TestGreedyArgmax:
25
+ @pytest.mark.parametrize("vocab_size", VOCAB_SIZES)
26
+ @pytest.mark.parametrize("batch_size", BATCH_SIZES)
27
+ def test_exact_match(self, vocab_size: int, batch_size: int):
28
+ logits = torch.randn(batch_size, vocab_size, dtype=torch.float32)
29
+ expected = reference_greedy_sample(logits)
30
+ result = torch.ops._C.greedy_argmax(logits)
31
+ torch.testing.assert_close(result, expected)
32
+
33
+ def test_single_dominant(self):
34
+ logits = torch.full((1, 50000), -1e9, dtype=torch.float32)
35
+ logits[0, 42] = 100.0
36
+ assert torch.ops._C.greedy_argmax(logits).item() == 42
37
+
38
+ def test_negative_logits(self):
39
+ logits = torch.randn(8, 32000, dtype=torch.float32) - 10.0
40
+ expected = reference_greedy_sample(logits)
41
+ result = torch.ops._C.greedy_argmax(logits)
42
+ torch.testing.assert_close(result, expected)
43
+
44
+
45
+class TestFusedGumbelArgmax:
46
+ @pytest.mark.parametrize("vocab_size", VOCAB_SIZES)
47
+ def test_distribution_chi_squared(self, vocab_size: int):
48
+ """Verify sampling distribution via chi-squared goodness of fit."""
49
+ small_vocab = min(vocab_size, 100)
50
+ logits = torch.randn(1, small_vocab, dtype=torch.float32)
51
+ probs = logits.softmax(dim=-1).squeeze(0)
52
+
53
+ n_samples = 100_000
54
+ counts = torch.zeros(small_vocab)
55
+ for trial in range(n_samples):
56
+ seed = torch.tensor([trial * 7 + 13], dtype=torch.long)
57
+ tile = logits.expand(1, -1).contiguous()
58
+ idx = torch.ops._C.fused_gumbel_argmax(tile, seed)
59
+ counts[idx.item()] += 1
60
+
61
+ expected = probs * n_samples
62
+ mask = expected > 5
63
+ chi2 = ((counts[mask] - expected[mask]) ** 2 / expected[mask]).sum()
64
+ dof = mask.sum().item() - 1
65
+ from scipy.stats import chi2 as chi2_dist
66
+
67
+ p_value = 1.0 - chi2_dist.cdf(chi2.item(), dof)
68
+ assert p_value > 0.001, (
69
+ f"Chi-squared test failed: chi2={chi2.item():.1f}, "
70
+ f"dof={dof}, p={p_value:.6f}"
71
+ )
72
+
73
+ def test_deterministic_same_seed(self):
74
+ """Same seed produces same result."""
75
+ logits = torch.randn(4, 32000, dtype=torch.float32)
76
+ seeds = torch.tensor([42, 123, 456, 789], dtype=torch.long)
77
+ r1 = torch.ops._C.fused_gumbel_argmax(logits, seeds)
78
+ r2 = torch.ops._C.fused_gumbel_argmax(logits, seeds)
79
+ torch.testing.assert_close(r1, r2)
80
+
81
+ def test_different_seeds_differ(self):
82
+ logits = torch.zeros(16, 50000, dtype=torch.float32)
83
+ seeds_a = torch.arange(16, dtype=torch.long)
84
+ seeds_b = torch.arange(16, dtype=torch.long) + 1_000_000
85
+ r_a = torch.ops._C.fused_gumbel_argmax(logits, seeds_a)
86
+ r_b = torch.ops._C.fused_gumbel_argmax(logits, seeds_b)
87
+ assert not torch.equal(r_a, r_b)
88
+
89
+ @pytest.mark.parametrize("vocab_size", VOCAB_SIZES)
90
+ @pytest.mark.parametrize("batch_size", BATCH_SIZES)
91
+ def test_output_in_range(self, vocab_size: int, batch_size: int):
92
+ logits = torch.randn(batch_size, vocab_size, dtype=torch.float32)
93
+ seeds = torch.arange(batch_size, dtype=torch.long)
94
+ result = torch.ops._C.fused_gumbel_argmax(logits, seeds)
95
+ assert result.min() >= 0
96
+ assert result.max() < vocab_size
97
+
98
+
99
+class TestMemory:
100
+ def test_fused_no_intermediate_allocs(self):
101
+ """Verify that the fused kernel does not allocate large intermediates."""
102
+ logits = torch.randn(16, 128256, dtype=torch.float32)
103
+ seeds = torch.arange(16, dtype=torch.long)
104
+
105
+ torch.ops._C.fused_gumbel_argmax(logits, seeds)
106
+
107
+ tracemalloc.start()
108
+ snap_before = tracemalloc.take_snapshot()
109
+ for _ in range(50):
110
+ torch.ops._C.fused_gumbel_argmax(logits, seeds)
111
+ snap_after = tracemalloc.take_snapshot()
112
+ tracemalloc.stop()
113
+
114
+ diff = snap_after.compare_to(snap_before, "lineno")
115
+ total_new_bytes = sum(s.size_diff for s in diff if s.size_diff > 0)
116
+ vocab_bytes = 16 * 128256 * 4
117
+ assert total_new_bytes < vocab_bytes, (
118
+ f"Fused kernel allocated {total_new_bytes} bytes, "
119
+ f"expected < {vocab_bytes} (one intermediate tensor)"
120
+ )
vllm/v1/sample/ops/topk_topp_sampler.py
+23
-12
@@ -319,9 +319,11 @@ class TopKTopPSampler(nn.Module):
319
k: torch.Tensor | None,
320
p: torch.Tensor | None,
321
) -> tuple[torch.Tensor, torch.Tensor | None]:
322
- """PyTorch-native implementation of top-k and top-p sampling for CPU.
322
+ """Fused Gumbel-max sampling for CPU.
323
324
- The logits tensor may be updated in-place.
324
+ Uses a precomputed Gumbel table + single-pass argmax over logits,
325
+ skipping softmax and intermediate allocations entirely.
326
+ Falls back to the native path when fp64 Gumbel noise is requested.
327
"""
328
logits = apply_top_k_top_p(logits, k, p)
329
logits_to_return = None
@@ -330,16 +332,25 @@ class TopKTopPSampler(nn.Module):
332
elif self.logprobs_mode == "processed_logprobs":
333
logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
334
333
- if not generators and not self.use_fp64_gumbel:
334
- return compiled_random_sample(logits), logits_to_return
335
-
336
- probs = logits.softmax(dim=-1, dtype=torch.float32)
337
- q = empty_exponential_noise_like(probs, self.use_fp64_gumbel)
338
- q.exponential_()
339
- for i, generator in generators.items():
340
- q[i].exponential_(generator=generator)
341
-
342
- return sample_with_exponential_noise(probs, q), logits_to_return
335
+ if self.use_fp64_gumbel:
336
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
337
+ q = empty_exponential_noise_like(probs, self.use_fp64_gumbel)
338
+ q.exponential_()
339
+ for i, generator in generators.items():
340
+ q[i].exponential_(generator=generator)
341
+ return sample_with_exponential_noise(probs, q), logits_to_return
342
+
343
+ batch_size = logits.shape[0]
344
+ seeds = torch.randint(0, 2**31, (batch_size,), dtype=torch.long)
345
+ for i, gen in generators.items():
346
+ seeds[i] = torch.randint(
347
+ 0, 2**31, (1,), generator=gen, dtype=torch.long
348
+ ).item()
349
+ logits_f32 = logits.to(dtype=torch.float32)
350
+ return (
351
+ torch.ops._C.fused_gumbel_argmax(logits_f32, seeds),
352
+ logits_to_return,
353
+ )
354
355
def _init_aiter_ops(self) -> bool:
356
if self._aiter_ops_import_failed:
vllm/v1/sample/sampler.py
+3
@@ -6,6 +6,7 @@ import torch
6
import torch.nn as nn
7
8
from vllm.config.model import LogprobsMode
9
+from vllm.platforms import current_platform
10
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
11
from vllm.utils.torch_utils import PIN_MEMORY
12
from vllm.v1.outputs import LogprobsTensors, SamplerOutput
@@ -239,6 +240,8 @@ class Sampler(nn.Module):
240
241
@staticmethod
242
def greedy_sample(logits: torch.Tensor) -> torch.Tensor:
243
+ if current_platform.is_cpu():
244
+ return torch.ops._C.greedy_argmax(logits)
245
return logits.argmax(dim=-1).view(-1)
246
247
def sample(