From 583637f126a5f5193952a81ca9e0630ae9fc9986 Mon Sep 17 00:00:00 2001 From: Rehan Khan Date: Wed, 30 Sep 2026 12:46:28 +0530 Subject: [PATCH] [CPU][Perf] Add vectorized Sampler Kernel (#53913) Signed-off-by: Rehan Khan Co-authored-by: Li, Jiang --- benchmarks/kernels/bench_cpu_sampling.py | 217 +++++++++++++++++++++++ cmake/cpu_extension.cmake | 3 + csrc/cpu/sampling_kernels.cpp | 151 ++++++++++++++++ csrc/cpu/torch_bindings.cpp | 12 ++ tests/kernels/test_cpu_fused_sampling.py | 120 +++++++++++++ vllm/v1/sample/ops/topk_topp_sampler.py | 33 ++-- vllm/v1/sample/sampler.py | 3 + 7 files changed, 528 insertions(+), 11 deletions(-) create mode 100644 benchmarks/kernels/bench_cpu_sampling.py create mode 100644 csrc/cpu/sampling_kernels.cpp create mode 100644 tests/kernels/test_cpu_fused_sampling.py diff --git a/benchmarks/kernels/bench_cpu_sampling.py b/benchmarks/kernels/bench_cpu_sampling.py new file mode 100644 index 0000000000..fb7fdeb3dd --- /dev/null +++ b/benchmarks/kernels/bench_cpu_sampling.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Micro-benchmark for CPU sampling kernels. + +Compares fused Gumbel-max / greedy argmax against the baseline +(softmax → exponential → div → argmax) across vocab and batch sizes. + +Usage: + .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py + .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py --profile + .venv/bin/python benchmarks/kernels/bench_cpu_sampling.py --vocab 128256 --batch 16 +""" + +import argparse +import time + +import torch +import vllm._C # noqa: F401 +from torch.profiler import ProfilerActivity, profile, record_function + + +def baseline_random_sample(logits: torch.Tensor) -> torch.Tensor: + probs = logits.softmax(dim=-1, dtype=torch.float32) + q = torch.empty_like(probs) + q.exponential_() + return probs.div(q).argmax(dim=-1).view(-1) + + +def baseline_greedy_sample(logits: torch.Tensor) -> torch.Tensor: + return logits.argmax(dim=-1).view(-1) + + +def bench_latency(fn, args, n_warmup=20, n_iters=500): + for _ in range(n_warmup): + fn(*args) + + t0 = time.perf_counter() + for _ in range(n_iters): + fn(*args) + elapsed = time.perf_counter() - t0 + return elapsed / n_iters * 1e6 # µs + + +def run_profile(logits, seeds, n_iters=50): + """Run torch.profiler and print comparison tables.""" + # Profile baseline random sampling + with profile( + activities=[ProfilerActivity.CPU], + record_shapes=True, + ) as prof_base_rand: + for _ in range(n_iters): + with record_function("baseline_random"): + baseline_random_sample(logits) + + # Profile fused Gumbel-max + with profile( + activities=[ProfilerActivity.CPU], + record_shapes=True, + ) as prof_fused: + for _ in range(n_iters): + with record_function("fused_gumbel_argmax"): + torch.ops._C.fused_gumbel_argmax(logits, seeds) + + # Profile baseline greedy + with profile( + activities=[ProfilerActivity.CPU], + record_shapes=True, + ) as prof_base_grdy: + for _ in range(n_iters): + with record_function("baseline_greedy"): + baseline_greedy_sample(logits) + + # Profile custom greedy + with profile( + activities=[ProfilerActivity.CPU], + record_shapes=True, + ) as prof_cust_grdy: + for _ in range(n_iters): + with record_function("custom_greedy_argmax"): + torch.ops._C.greedy_argmax(logits) + + B, V = logits.shape + print(f"\n{'=' * 80}") + print(f"torch.profiler breakdown (batch={B}, vocab={V}, iters={n_iters})") + print(f"{'=' * 80}") + + print("\n--- Baseline Random (softmax → exp → div → argmax) ---") + print(prof_base_rand.key_averages().table(sort_by="cpu_time_total", row_limit=15)) + + print("\n--- Fused Gumbel-max (table lookup + add + argmax) ---") + print(prof_fused.key_averages().table(sort_by="cpu_time_total", row_limit=15)) + + print("\n--- Baseline Greedy (torch.argmax) ---") + print(prof_base_grdy.key_averages().table(sort_by="cpu_time_total", row_limit=15)) + + print("\n--- Custom Greedy (vec_op argmax) ---") + print(prof_cust_grdy.key_averages().table(sort_by="cpu_time_total", row_limit=15)) + + # Summary comparison + def avg_us(prof, label): + for e in prof.key_averages(): + if e.key == label: + return e.cpu_time_total / e.count + return 0.0 + + t_br = avg_us(prof_base_rand, "baseline_random") + t_fg = avg_us(prof_fused, "fused_gumbel_argmax") + t_bg = avg_us(prof_base_grdy, "baseline_greedy") + t_cg = avg_us(prof_cust_grdy, "custom_greedy_argmax") + + print(f"\n{'=' * 60}") + print(f" Summary (batch={B}, vocab={V})") + print(f"{'=' * 60}") + print(f" {'Kernel':<30} {'Avg (µs)':>10} {'Speedup':>10}") + print(f" {'-' * 50}") + print(f" {'baseline random':<30} {t_br:>10.1f} {'—':>10}") + print( + f" {'fused gumbel-max':<30} {t_fg:>10.1f} " + f"{t_br / t_fg if t_fg > 0 else 0:>9.2f}x" + ) + print(f" {'baseline greedy':<30} {t_bg:>10.1f} {'—':>10}") + print( + f" {'custom greedy':<30} {t_cg:>10.1f} {t_bg / t_cg if t_cg > 0 else 0:>9.2f}x" + ) + print(f"{'=' * 60}") + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark CPU sampling kernels") + parser.add_argument( + "--vocab", + type=int, + nargs="+", + default=[32000, 49152, 128256], + help="Vocab sizes to benchmark", + ) + parser.add_argument( + "--batch", + type=int, + nargs="+", + default=[1, 4, 16], + help="Batch sizes to benchmark", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Run torch.profiler and export chrome trace", + ) + parser.add_argument( + "--iters", type=int, default=500, help="Iterations per measurement" + ) + args = parser.parse_args() + + header = ( + f"{'batch':>5} {'vocab':>7} " + f"{'base_rand':>10} {'fused_rand':>10} {'rand_spdup':>10} " + f"{'base_grdy':>10} {'cust_grdy':>10} {'grdy_spdup':>10}" + ) + units = ( + f"{'':>5} {'':>7} " + f"{'(µs)':>10} {'(µs)':>10} {'':>10} " + f"{'(µs)':>10} {'(µs)':>10} {'':>10}" + ) + print("\n" + "=" * len(header)) + print("CPU Sampling Kernel Benchmark") + print("=" * len(header)) + print(header) + print(units) + print("-" * len(header)) + + for V in args.vocab: + for B in args.batch: + logits = torch.randn(B, V, dtype=torch.float32) + seeds = torch.arange(B, dtype=torch.long) + + t_base_rand = bench_latency( + baseline_random_sample, (logits,), n_iters=args.iters + ) + t_fused = bench_latency( + torch.ops._C.fused_gumbel_argmax, (logits, seeds), n_iters=args.iters + ) + + t_base_grdy = bench_latency( + baseline_greedy_sample, (logits,), n_iters=args.iters + ) + t_cust_grdy = bench_latency( + torch.ops._C.greedy_argmax, (logits,), n_iters=args.iters + ) + + rand_speedup = t_base_rand / t_fused if t_fused > 0 else 0 + grdy_speedup = t_base_grdy / t_cust_grdy if t_cust_grdy > 0 else 0 + + print( + f"{B:>5} {V:>7} " + f"{t_base_rand:>10.1f} {t_fused:>10.1f} " + f"{rand_speedup:>9.2f}x " + f"{t_base_grdy:>10.1f} {t_cust_grdy:>10.1f} " + f"{grdy_speedup:>9.2f}x" + ) + + print("-" * len(header)) + + if args.profile: + print("\nRunning torch.profiler (batch=16, vocab=128256) ...") + logits = torch.randn(16, 128256, dtype=torch.float32) + seeds = torch.arange(16, dtype=torch.long) + + # warmup + for _ in range(10): + baseline_random_sample(logits) + torch.ops._C.fused_gumbel_argmax(logits, seeds) + + run_profile(logits, seeds) + + +if __name__ == "__main__": + main() diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 414fbc5e23..94dd188822 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -501,6 +501,7 @@ set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/cpu_isa.cpp" + "csrc/cpu/sampling_kernels.cpp" "csrc/cpu/torch_bindings.cpp") if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND @@ -589,6 +590,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/cpu_isa.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/mamba_cpu.cpp" + "csrc/cpu/sampling_kernels.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files "csrc/cpu/activation.cpp" @@ -606,6 +608,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/cpu_isa.cpp" "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/dnnl_kernels.cpp" + "csrc/cpu/sampling_kernels.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files "csrc/cpu/activation.cpp" diff --git a/csrc/cpu/sampling_kernels.cpp b/csrc/cpu/sampling_kernels.cpp new file mode 100644 index 0000000000..3e2f00a0ba --- /dev/null +++ b/csrc/cpu/sampling_kernels.cpp @@ -0,0 +1,151 @@ +#include "cpu_types.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace { + +constexpr int GUMBEL_TABLE_SIZE = 1 << 20; +constexpr int GUMBEL_TABLE_MASK = GUMBEL_TABLE_SIZE - 1; + +static float g_gumbel_table[GUMBEL_TABLE_SIZE]; +static std::once_flag g_gumbel_init_flag; + +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 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); + +#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; + } +} + +static void greedy_argmax_kernel(int64_t* __restrict__ output, + const float* __restrict__ logits, + const int64_t batch_size, + const int64_t vocab_size) { + constexpr int VEC_ELEM_NUM = vec_op::FP32Vec16::VEC_ELEM_NUM; + const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM); + +#pragma omp parallel for schedule(static) + for (int64_t b = 0; b < batch_size; ++b) { + const float* row = logits + b * vocab_size; + + vec_op::FP32Vec16 vmax(-std::numeric_limits::infinity()); + for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) { + vmax = vmax.max(vec_op::FP32Vec16(row + i)); + } + float best_val = vmax.reduce_max(); + for (int64_t i = vec_end; i < vocab_size; ++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; + } + } + output[b] = best_idx; + } +} + +} // namespace + +torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits, + const torch::Tensor& seeds) { + 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(seeds.dim() == 1 && seeds.size(0) == logits.size(0), + "seeds must be 1-D with batch_size elements"); + + auto logits_contig = logits.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)); + return output; +} + +torch::Tensor greedy_argmax(const torch::Tensor& logits) { + TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]"); + TORCH_CHECK(logits.scalar_type() == torch::kFloat32, + "logits must be float32"); + + // A single row is faster in ATen: OpenMP fork/join dominates the scan + // (~1 ms vs ~50 µs on s390x at vocab=32k). + if (logits.size(0) <= 1) { + return logits.argmax(/*dim=*/-1); + } + + auto logits_contig = logits.contiguous(); + auto output = torch::empty({logits_contig.size(0)}, torch::kInt64); + greedy_argmax_kernel(output.data_ptr(), + logits_contig.data_ptr(), logits_contig.size(0), + logits_contig.size(1)); + return output; +} diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index accd5fcec3..918daf5a8b 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -393,6 +393,11 @@ void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states, const c10::optional& z, const at::Tensor& cu_seqlens); +torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits, + const torch::Tensor& seeds); + +torch::Tensor greedy_argmax(const torch::Tensor& logits); + void init_cpu_memory_env(std::vector node_ids); namespace cpu_utils { @@ -969,6 +974,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); + // Fused sampling kernels + ops.def("fused_gumbel_argmax(Tensor logits, Tensor seeds) -> Tensor"); + ops.impl("fused_gumbel_argmax", torch::kCPU, &fused_gumbel_argmax); + + ops.def("greedy_argmax(Tensor logits) -> Tensor"); + ops.impl("greedy_argmax", torch::kCPU, &greedy_argmax); + // Speculative decoding kernels ops.def( "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, " diff --git a/tests/kernels/test_cpu_fused_sampling.py b/tests/kernels/test_cpu_fused_sampling.py new file mode 100644 index 0000000000..8a3857fe9e --- /dev/null +++ b/tests/kernels/test_cpu_fused_sampling.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import tracemalloc + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +import vllm._C # noqa: F401, E402 + +VOCAB_SIZES = [32000, 49152, 128256] +BATCH_SIZES = [1, 4, 16] + + +def reference_greedy_sample(logits: torch.Tensor) -> torch.Tensor: + return logits.argmax(dim=-1).view(-1) + + +class TestGreedyArgmax: + @pytest.mark.parametrize("vocab_size", VOCAB_SIZES) + @pytest.mark.parametrize("batch_size", BATCH_SIZES) + def test_exact_match(self, vocab_size: int, batch_size: int): + logits = torch.randn(batch_size, vocab_size, dtype=torch.float32) + expected = reference_greedy_sample(logits) + result = torch.ops._C.greedy_argmax(logits) + torch.testing.assert_close(result, expected) + + def test_single_dominant(self): + logits = torch.full((1, 50000), -1e9, dtype=torch.float32) + logits[0, 42] = 100.0 + assert torch.ops._C.greedy_argmax(logits).item() == 42 + + def test_negative_logits(self): + logits = torch.randn(8, 32000, dtype=torch.float32) - 10.0 + expected = reference_greedy_sample(logits) + result = torch.ops._C.greedy_argmax(logits) + torch.testing.assert_close(result, expected) + + +class TestFusedGumbelArgmax: + @pytest.mark.parametrize("vocab_size", VOCAB_SIZES) + def test_distribution_chi_squared(self, vocab_size: int): + """Verify sampling distribution via chi-squared goodness of fit.""" + small_vocab = min(vocab_size, 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 + + expected = probs * n_samples + mask = expected > 5 + chi2 = ((counts[mask] - expected[mask]) ** 2 / expected[mask]).sum() + dof = mask.sum().item() - 1 + from scipy.stats import chi2 as chi2_dist + + p_value = 1.0 - chi2_dist.cdf(chi2.item(), dof) + assert p_value > 0.001, ( + f"Chi-squared test failed: chi2={chi2.item():.1f}, " + f"dof={dof}, p={p_value:.6f}" + ) + + def test_deterministic_same_seed(self): + """Same seed produces same result.""" + logits = torch.randn(4, 32000, dtype=torch.float32) + seeds = torch.tensor([42, 123, 456, 789], dtype=torch.long) + r1 = torch.ops._C.fused_gumbel_argmax(logits, seeds) + r2 = torch.ops._C.fused_gumbel_argmax(logits, seeds) + torch.testing.assert_close(r1, r2) + + def test_different_seeds_differ(self): + logits = torch.zeros(16, 50000, dtype=torch.float32) + seeds_a = torch.arange(16, dtype=torch.long) + seeds_b = torch.arange(16, dtype=torch.long) + 1_000_000 + r_a = torch.ops._C.fused_gumbel_argmax(logits, seeds_a) + r_b = torch.ops._C.fused_gumbel_argmax(logits, seeds_b) + assert not torch.equal(r_a, r_b) + + @pytest.mark.parametrize("vocab_size", VOCAB_SIZES) + @pytest.mark.parametrize("batch_size", BATCH_SIZES) + def test_output_in_range(self, vocab_size: int, batch_size: int): + logits = torch.randn(batch_size, vocab_size, dtype=torch.float32) + seeds = torch.arange(batch_size, dtype=torch.long) + result = torch.ops._C.fused_gumbel_argmax(logits, seeds) + assert result.min() >= 0 + assert result.max() < vocab_size + + +class TestMemory: + def test_fused_no_intermediate_allocs(self): + """Verify that the fused kernel does not allocate large intermediates.""" + logits = torch.randn(16, 128256, dtype=torch.float32) + seeds = torch.arange(16, dtype=torch.long) + + torch.ops._C.fused_gumbel_argmax(logits, seeds) + + tracemalloc.start() + snap_before = tracemalloc.take_snapshot() + for _ in range(50): + torch.ops._C.fused_gumbel_argmax(logits, seeds) + snap_after = tracemalloc.take_snapshot() + tracemalloc.stop() + + diff = snap_after.compare_to(snap_before, "lineno") + total_new_bytes = sum(s.size_diff for s in diff if s.size_diff > 0) + vocab_bytes = 16 * 128256 * 4 + assert total_new_bytes < vocab_bytes, ( + f"Fused kernel allocated {total_new_bytes} bytes, " + f"expected < {vocab_bytes} (one intermediate tensor)" + ) diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 69ec368309..a3084a09c2 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -319,9 +319,11 @@ class TopKTopPSampler(nn.Module): k: torch.Tensor | None, p: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """PyTorch-native implementation of top-k and top-p sampling for CPU. + """Fused Gumbel-max sampling for CPU. - The logits tensor may be updated in-place. + 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. """ logits = apply_top_k_top_p(logits, k, p) logits_to_return = None @@ -330,16 +332,25 @@ class TopKTopPSampler(nn.Module): elif self.logprobs_mode == "processed_logprobs": logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32) - if not generators and not self.use_fp64_gumbel: - return compiled_random_sample(logits), logits_to_return + if self.use_fp64_gumbel: + probs = logits.softmax(dim=-1, dtype=torch.float32) + q = empty_exponential_noise_like(probs, self.use_fp64_gumbel) + q.exponential_() + for i, generator in generators.items(): + q[i].exponential_(generator=generator) + return sample_with_exponential_noise(probs, q), logits_to_return - probs = logits.softmax(dim=-1, dtype=torch.float32) - q = empty_exponential_noise_like(probs, self.use_fp64_gumbel) - q.exponential_() - for i, generator in generators.items(): - q[i].exponential_(generator=generator) - - return sample_with_exponential_noise(probs, q), logits_to_return + batch_size = logits.shape[0] + seeds = torch.randint(0, 2**31, (batch_size,), dtype=torch.long) + for i, gen in generators.items(): + seeds[i] = torch.randint( + 0, 2**31, (1,), generator=gen, dtype=torch.long + ).item() + logits_f32 = logits.to(dtype=torch.float32) + return ( + torch.ops._C.fused_gumbel_argmax(logits_f32, seeds), + logits_to_return, + ) def _init_aiter_ops(self) -> bool: if self._aiter_ops_import_failed: diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index 7e8d597d88..d19f65f6d8 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -6,6 +6,7 @@ import torch import torch.nn as nn from vllm.config.model import LogprobsMode +from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.outputs import LogprobsTensors, SamplerOutput @@ -239,6 +240,8 @@ class Sampler(nn.Module): @staticmethod def greedy_sample(logits: torch.Tensor) -> torch.Tensor: + if current_platform.is_cpu(): + return torch.ops._C.greedy_argmax(logits) return logits.argmax(dim=-1).view(-1) def sample(