From 6654bcbfca86a8c17c8a377f6813c7a86e64cd6d Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Mon, 21 Sep 2026 13:24:41 +0800 Subject: [PATCH] [MRV2][XPU] use xpu sample kernel in mrv2 sampler (#57277) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- vllm/v1/sample/ops/topk_topp_sampler.py | 98 ++++++++++++++++--------- vllm/v1/watermarking/gpu_sampler.py | 4 +- vllm/v1/worker/gpu/sample/sampler.py | 30 ++++++-- 3 files changed, 89 insertions(+), 43 deletions(-) diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 477e99c2fa..dac8554568 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - import torch import torch.nn as nn @@ -127,6 +126,69 @@ def flashinfer_sampler_supported() -> bool: return False +def xpu_sampler_supported() -> bool: + """Decide whether the fused XPU top-k/top-p sampler kernel can be used. + + Returns False (with appropriate logging) when the platform isn't XPU, when + ``VLLM_XPU_USE_SAMPLER_KERNEL`` is 0. + + Note: callers must additionally ensure no request needs a per-request seed + or greedy sampling, since the kernel draws from the device's default + generator and always samples randomly. + """ + if not current_platform.is_xpu(): + return False + if not envs.VLLM_XPU_USE_SAMPLER_KERNEL: + logger.info_once( + "Fused XPU top-p/top-k sampling disabled via VLLM_XPU_USE_SAMPLER_KERNEL=0." + ) + return False + + logger.info_once("Using the fused XPU kernel for top-p & top-k sampling.") + return True + + +def xpu_sample( + logits: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, + logprobs_mode: LogprobsMode = "raw_logprobs", +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Sample from the logits using the fused XPU top-k/top-p kernel. + + Statistically equivalent to `random_sample`, but avoids sorting the vocab + and never materializes the probability tensor. The noise comes from the + device's default generator, so per-request generators aren't supported. + + Returns the sampled token ids and, for processed logprobs modes, the + post-top-k/top-p logits (or logprobs); otherwise None. + """ + logits = logits.to(dtype=torch.float32).contiguous() + sampled = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device) + logits_to_return = ( + torch.empty_like(logits) if logprobs_mode in PROCESSED_LOGPROBS_MODES else None + ) + + generator = torch.xpu.default_generators[logits.device.index] + state = generator.get_state() + seed, offset = state.view(torch.int64).tolist() + seeds = torch.tensor([seed, offset], dtype=torch.int64, device="cpu") + # The XPU kernel expects k as int64 (Long), but the input batch + # stores top_k as int32. Cast here to avoid dtype mismatch. + if k is not None: + k = k.to(torch.int64) + torch.ops.vllm.xpu_topk_topp_sampler( + sampled, logits_to_return, logits, k, p, logprobs_mode, seeds + ) + # The custom XPU sampler kernel consumes RNG values internally, so advance + # the default generator's offset to keep future draws deterministic. + # pytorch: offset must be multiple of 4 + offset = (offset + logits.numel() + 3) // 4 * 4 + state.view(torch.int64)[1] = offset + generator.set_state(state) + return sampled, logits_to_return + + class TopKTopPSampler(nn.Module): """Module that performs optional top-k and top-p filtering followed by weighted random sampling of logits. @@ -162,7 +224,7 @@ class TopKTopPSampler(nn.Module): else: self.forward = self.forward_cpu elif current_platform.is_xpu(): - if envs.VLLM_XPU_USE_SAMPLER_KERNEL: + if xpu_sampler_supported(): self.forward = self.forward_xpu else: self.forward = self.forward_native @@ -356,37 +418,7 @@ class TopKTopPSampler(nn.Module): "PyTorch-native implementation." ) return self.forward_native(logits, generators, k, p) - random_sampled = torch.empty( - logits.shape[0], dtype=torch.int64, device=logits.device - ) - logits_to_return = None - if self.logprobs_mode in PROCESSED_LOGPROBS_MODES: - logits_to_return = torch.empty_like(logits) - - assert len(generators) != logits.shape[0], ( - "xpu kernel topk_topp_sampler does not support batch-wise generators." - ) - generator = torch.xpu.default_generators[logits.device.index] - - state = generator.get_state() - seed, offset = state.view(torch.int64) - seeds = torch.tensor( - [seed, offset], dtype=torch.int64, device=torch.device("cpu") - ) - # The XPU kernel expects k as int64 (Long), but the input batch - # stores top_k as int32. Cast here to avoid dtype mismatch. - if k is not None: - k = k.to(torch.int64) - torch.ops.vllm.xpu_topk_topp_sampler( - random_sampled, logits_to_return, logits, k, p, self.logprobs_mode, seeds - ) - # The custom XPU sampler kernel consumes RNG values internally, so advance - # the default generator's offset to keep future draws deterministic. - # pytorch: offset must be multiple of 4 - offset = (offset + logits.numel() + 3) // 4 * 4 - state.view(torch.int64)[1] = offset - generator.set_state(state) - return random_sampled, logits_to_return + return xpu_sample(logits, k, p, self.logprobs_mode) # Note: this is a workaround for diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py index c9624b2b25..a4e6e2dd5b 100644 --- a/vllm/v1/watermarking/gpu_sampler.py +++ b/vllm/v1/watermarking/gpu_sampler.py @@ -58,7 +58,7 @@ class GPUWatermarkSampler(Sampler): pos: torch.Tensor, top_k: torch.Tensor | None, top_p: torch.Tensor | None, - use_flashinfer: bool, + use_fused_sampler: bool, ) -> tuple[torch.Tensor, torch.Tensor]: enabled = self.watermarking.np[idx_mapping_np] & ( self.sampling_states.temperature.np[idx_mapping_np] != 0 @@ -71,7 +71,7 @@ class GPUWatermarkSampler(Sampler): pos, top_k, top_p, - use_flashinfer, + use_fused_sampler, ) processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 92e9ecee4d..802adb6f64 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -14,6 +14,8 @@ from vllm.v1.sample.ops.topk_topp_sampler import ( apply_top_k_top_p, flashinfer_sample, flashinfer_sampler_supported, + xpu_sample, + xpu_sampler_supported, ) from vllm.v1.worker.gpu.input_batch import InputBatch, get_num_sampled_and_rejected from vllm.v1.worker.gpu.metrics.logits import get_num_nans @@ -86,6 +88,10 @@ class Sampler: self.use_flashinfer = ( not return_sampling_mask and flashinfer_sampler_supported() ) + # The XPU kernel draws fp32 exponential noise, so it can't honor fp64. + self.use_xpu_sampler = ( + not return_sampling_mask and not use_fp64_gumbel and xpu_sampler_supported() + ) def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: needs_processing = self.sampling_states.add_request(req_idx, sampling_params) @@ -292,15 +298,18 @@ class Sampler: top_k, top_p = self.sampling_states.get_top_k_top_p( expanded_idx_mapping, idx_mapping_np ) - use_flashinfer = self.use_flashinfer and not ( - # Don't use FI sampler if no requests use top_k/top_p, if there are - # any greedy requests or per-request seeds, or if post-processed - # logprobs need to be returned for any requests. + # Don't use a fused sampler if no requests use top_k/top_p, if there are + # any greedy requests or per-request seeds, or if post-processed + # logprobs need to be returned for any requests. + fused_sampler_eligible = not ( (top_k is None and top_p is None) or (return_logprobs and self.logprobs_mode in PROCESSED_LOGPROBS_MODES) or self.sampling_states.any_greedy(idx_mapping_np) or self.sampling_states.any_explicit_seed(idx_mapping_np) ) + use_fused_sampler = ( + self.use_flashinfer or self.use_xpu_sampler + ) and fused_sampler_eligible return self._sample_random( processed_logits, @@ -309,7 +318,7 @@ class Sampler: pos, top_k, top_p, - use_flashinfer, + use_fused_sampler, ) def _sample_random( @@ -320,10 +329,15 @@ class Sampler: pos: torch.Tensor, top_k: torch.Tensor | None, top_p: torch.Tensor | None, - use_flashinfer: bool, + use_fused_sampler: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - if use_flashinfer: - sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64) + if use_fused_sampler: + if self.use_flashinfer: + sampled = flashinfer_sample(processed_logits, top_k, top_p).to( + torch.int64 + ) + else: # Use XPU sampler + sampled, _ = xpu_sample(processed_logits, top_k, top_p) else: processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) sampled = gumbel_sample(