[MRV2][XPU] use xpu sample kernel in mrv2 sampler (#57277)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
liuzhenwei committed
Sep 21, 2026 at 13:24 UTC
6654bcbfca86a8c17c8a377f6813c7a86e64cd6d
3 files changed
+89
-43
vllm/v1/sample/ops/topk_topp_sampler.py
+65
-33
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-
4
import torch
5
import torch.nn as nn
6
@@ -127,6 +126,69 @@ def flashinfer_sampler_supported() -> bool:
126
return False
127
128
129
+def xpu_sampler_supported() -> bool:
130
+ """Decide whether the fused XPU top-k/top-p sampler kernel can be used.
131
+
132
+ Returns False (with appropriate logging) when the platform isn't XPU, when
133
+ ``VLLM_XPU_USE_SAMPLER_KERNEL`` is 0.
134
+
135
+ Note: callers must additionally ensure no request needs a per-request seed
136
+ or greedy sampling, since the kernel draws from the device's default
137
+ generator and always samples randomly.
138
+ """
139
+ if not current_platform.is_xpu():
140
+ return False
141
+ if not envs.VLLM_XPU_USE_SAMPLER_KERNEL:
142
+ logger.info_once(
143
+ "Fused XPU top-p/top-k sampling disabled via VLLM_XPU_USE_SAMPLER_KERNEL=0."
144
+ )
145
+ return False
146
+
147
+ logger.info_once("Using the fused XPU kernel for top-p & top-k sampling.")
148
+ return True
149
+
150
+
151
+def xpu_sample(
152
+ logits: torch.Tensor,
153
+ k: torch.Tensor | None,
154
+ p: torch.Tensor | None,
155
+ logprobs_mode: LogprobsMode = "raw_logprobs",
156
+) -> tuple[torch.Tensor, torch.Tensor | None]:
157
+ """Sample from the logits using the fused XPU top-k/top-p kernel.
158
+
159
+ Statistically equivalent to `random_sample`, but avoids sorting the vocab
160
+ and never materializes the probability tensor. The noise comes from the
161
+ device's default generator, so per-request generators aren't supported.
162
+
163
+ Returns the sampled token ids and, for processed logprobs modes, the
164
+ post-top-k/top-p logits (or logprobs); otherwise None.
165
+ """
166
+ logits = logits.to(dtype=torch.float32).contiguous()
167
+ sampled = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device)
168
+ logits_to_return = (
169
+ torch.empty_like(logits) if logprobs_mode in PROCESSED_LOGPROBS_MODES else None
170
+ )
171
+
172
+ generator = torch.xpu.default_generators[logits.device.index]
173
+ state = generator.get_state()
174
+ seed, offset = state.view(torch.int64).tolist()
175
+ seeds = torch.tensor([seed, offset], dtype=torch.int64, device="cpu")
176
+ # The XPU kernel expects k as int64 (Long), but the input batch
177
+ # stores top_k as int32. Cast here to avoid dtype mismatch.
178
+ if k is not None:
179
+ k = k.to(torch.int64)
180
+ torch.ops.vllm.xpu_topk_topp_sampler(
181
+ sampled, logits_to_return, logits, k, p, logprobs_mode, seeds
182
+ )
183
+ # The custom XPU sampler kernel consumes RNG values internally, so advance
184
+ # the default generator's offset to keep future draws deterministic.
185
+ # pytorch: offset must be multiple of 4
186
+ offset = (offset + logits.numel() + 3) // 4 * 4
187
+ state.view(torch.int64)[1] = offset
188
+ generator.set_state(state)
189
+ return sampled, logits_to_return
190
+
191
+
192
class TopKTopPSampler(nn.Module):
193
"""Module that performs optional top-k and top-p filtering followed by
194
weighted random sampling of logits.
@@ -162,7 +224,7 @@ class TopKTopPSampler(nn.Module):
224
else:
225
self.forward = self.forward_cpu
226
elif current_platform.is_xpu():
165
- if envs.VLLM_XPU_USE_SAMPLER_KERNEL:
227
+ if xpu_sampler_supported():
228
self.forward = self.forward_xpu
229
else:
230
self.forward = self.forward_native
@@ -356,37 +418,7 @@ class TopKTopPSampler(nn.Module):
418
"PyTorch-native implementation."
419
)
420
return self.forward_native(logits, generators, k, p)
359
- random_sampled = torch.empty(
360
- logits.shape[0], dtype=torch.int64, device=logits.device
361
- )
362
- logits_to_return = None
363
- if self.logprobs_mode in PROCESSED_LOGPROBS_MODES:
364
- logits_to_return = torch.empty_like(logits)
365
-
366
- assert len(generators) != logits.shape[0], (
367
- "xpu kernel topk_topp_sampler does not support batch-wise generators."
368
- )
369
- generator = torch.xpu.default_generators[logits.device.index]
370
-
371
- state = generator.get_state()
372
- seed, offset = state.view(torch.int64)
373
- seeds = torch.tensor(
374
- [seed, offset], dtype=torch.int64, device=torch.device("cpu")
375
- )
376
- # The XPU kernel expects k as int64 (Long), but the input batch
377
- # stores top_k as int32. Cast here to avoid dtype mismatch.
378
- if k is not None:
379
- k = k.to(torch.int64)
380
- torch.ops.vllm.xpu_topk_topp_sampler(
381
- random_sampled, logits_to_return, logits, k, p, self.logprobs_mode, seeds
382
- )
383
- # The custom XPU sampler kernel consumes RNG values internally, so advance
384
- # the default generator's offset to keep future draws deterministic.
385
- # pytorch: offset must be multiple of 4
386
- offset = (offset + logits.numel() + 3) // 4 * 4
387
- state.view(torch.int64)[1] = offset
388
- generator.set_state(state)
389
- return random_sampled, logits_to_return
421
+ return xpu_sample(logits, k, p, self.logprobs_mode)
422
423
424
# Note: this is a workaround for
vllm/v1/watermarking/gpu_sampler.py
+2
-2
@@ -58,7 +58,7 @@ class GPUWatermarkSampler(Sampler):
58
pos: torch.Tensor,
59
top_k: torch.Tensor | None,
60
top_p: torch.Tensor | None,
61
- use_flashinfer: bool,
61
+ use_fused_sampler: bool,
62
) -> tuple[torch.Tensor, torch.Tensor]:
63
enabled = self.watermarking.np[idx_mapping_np] & (
64
self.sampling_states.temperature.np[idx_mapping_np] != 0
@@ -71,7 +71,7 @@ class GPUWatermarkSampler(Sampler):
71
pos,
72
top_k,
73
top_p,
74
- use_flashinfer,
74
+ use_fused_sampler,
75
)
76
77
processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
vllm/v1/worker/gpu/sample/sampler.py
+22
-8
@@ -14,6 +14,8 @@ from vllm.v1.sample.ops.topk_topp_sampler import (
14
apply_top_k_top_p,
15
flashinfer_sample,
16
flashinfer_sampler_supported,
17
+ xpu_sample,
18
+ xpu_sampler_supported,
19
)
20
from vllm.v1.worker.gpu.input_batch import InputBatch, get_num_sampled_and_rejected
21
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
@@ -86,6 +88,10 @@ class Sampler:
88
self.use_flashinfer = (
89
not return_sampling_mask and flashinfer_sampler_supported()
90
)
91
+ # The XPU kernel draws fp32 exponential noise, so it can't honor fp64.
92
+ self.use_xpu_sampler = (
93
+ not return_sampling_mask and not use_fp64_gumbel and xpu_sampler_supported()
94
+ )
95
96
def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
97
needs_processing = self.sampling_states.add_request(req_idx, sampling_params)
@@ -292,15 +298,18 @@ class Sampler:
298
top_k, top_p = self.sampling_states.get_top_k_top_p(
299
expanded_idx_mapping, idx_mapping_np
300
)
295
- use_flashinfer = self.use_flashinfer and not (
296
- # Don't use FI sampler if no requests use top_k/top_p, if there are
297
- # any greedy requests or per-request seeds, or if post-processed
298
- # logprobs need to be returned for any requests.
301
+ # Don't use a fused sampler if no requests use top_k/top_p, if there are
302
+ # any greedy requests or per-request seeds, or if post-processed
303
+ # logprobs need to be returned for any requests.
304
+ fused_sampler_eligible = not (
305
(top_k is None and top_p is None)
306
or (return_logprobs and self.logprobs_mode in PROCESSED_LOGPROBS_MODES)
307
or self.sampling_states.any_greedy(idx_mapping_np)
308
or self.sampling_states.any_explicit_seed(idx_mapping_np)
309
)
310
+ use_fused_sampler = (
311
+ self.use_flashinfer or self.use_xpu_sampler
312
+ ) and fused_sampler_eligible
313
314
return self._sample_random(
315
processed_logits,
@@ -309,7 +318,7 @@ class Sampler:
318
pos,
319
top_k,
320
top_p,
312
- use_flashinfer,
321
+ use_fused_sampler,
322
)
323
324
def _sample_random(
@@ -320,10 +329,15 @@ class Sampler:
329
pos: torch.Tensor,
330
top_k: torch.Tensor | None,
331
top_p: torch.Tensor | None,
323
- use_flashinfer: bool,
332
+ use_fused_sampler: bool,
333
) -> tuple[torch.Tensor, torch.Tensor]:
325
- if use_flashinfer:
326
- sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64)
334
+ if use_fused_sampler:
335
+ if self.use_flashinfer:
336
+ sampled = flashinfer_sample(processed_logits, top_k, top_p).to(
337
+ torch.int64
338
+ )
339
+ else: # Use XPU sampler
340
+ sampled, _ = xpu_sample(processed_logits, top_k, top_p)
341
else:
342
processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
343
sampled = gumbel_sample(