iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[MRV2][XPU] use xpu sample kernel in mrv2 sampler (#57277)

Browse Source 
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
liuzhenwei
2026-09-21 13:24:41 +08:00
committed by GitHub
co-authored by Kunshang Ji
parent 76ee20c1ec
commit 6654bcbfca
3 changed files with 89 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/sample/ops/topk_topp_sampler.py
+65 -33
View File
@@ -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
vllm/v1/watermarking/gpu_sampler.py
+2 -2
View File
@@ -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)
vllm/v1/worker/gpu/sample/sampler.py
+22 -8
View File
@@ -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(