A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 06:15 UTC
[XPU] use xpu topk topp sample kernel (#39285)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
May 5, 2026 at 18:05 UTC
2ceea429583ad6429a3bb12ba2e909c5cb99f059
3 files changed
+88
vllm/_xpu_ops.py
+35
@@ -185,6 +185,35 @@ def _xpu_ops_deepseek_scaling_rope_fake(
185
return query, key
186
187
188
+def _topk_topp_sample_impl(
189
+ random_sampled: torch.Tensor,
190
+ logits_to_return: torch.Tensor | None,
191
+ logits: torch.Tensor,
192
+ k: torch.Tensor | None,
193
+ p: torch.Tensor | None,
194
+ logprobs_mode: str,
195
+ seeds: torch.Tensor | None,
196
+ lambda_: float = 1.0,
197
+) -> None:
198
+ torch.ops._xpu_C.topk_topp_sampler(
199
+ random_sampled, logits_to_return, logits, k, p, logprobs_mode, seeds, lambda_
200
+ )
201
+ return
202
+
203
+
204
+def _topk_topp_sample_fake(
205
+ random_sampled: torch.Tensor,
206
+ logits_to_return: torch.Tensor | None,
207
+ logits: torch.Tensor,
208
+ k: torch.Tensor | None,
209
+ p: torch.Tensor | None,
210
+ logprobs_mode: str,
211
+ seeds: torch.Tensor | None,
212
+ lambda_: float = 1.0,
213
+) -> None:
214
+ return
215
+
216
+
217
def _xpu_mxfp8_quantize_impl(
218
x: torch.Tensor, dtype: torch.dtype | None = None
219
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -691,6 +720,12 @@ class xpu_ops:
720
fake_impl=_gdn_attention_core_xpu_fake,
721
)
722
723
+ direct_register_custom_op(
724
+ op_name="xpu_topk_topp_sampler",
725
+ op_func=_topk_topp_sample_impl,
726
+ fake_impl=_topk_topp_sample_fake,
727
+ )
728
+
729
_OPS_REGISTERED = True
730
731
vllm/envs.py
+5
@@ -266,6 +266,7 @@ if TYPE_CHECKING:
266
VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True
267
VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32
268
VLLM_XPU_ENABLE_XPU_GRAPH: bool = False
269
+ VLLM_XPU_USE_SAMPLER_KERNEL: bool = True
270
VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
271
272
@@ -1775,6 +1776,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1776
"VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool(
1777
int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0"))
1778
),
1779
+ # whether use xpu specific sample kernel
1780
+ "VLLM_XPU_USE_SAMPLER_KERNEL": lambda: bool(
1781
+ int(os.getenv("VLLM_XPU_USE_SAMPLER_KERNEL", "1"))
1782
+ ),
1783
# Enable simple KV offload.
1784
"VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool(
1785
int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))
vllm/v1/sample/ops/topk_topp_sampler.py
+48
@@ -82,6 +82,11 @@ class TopKTopPSampler(nn.Module):
82
self.forward = self.forward_native
83
else:
84
self.forward = self.forward_cpu
85
+ elif current_platform.is_xpu():
86
+ if envs.VLLM_XPU_USE_SAMPLER_KERNEL:
87
+ self.forward = self.forward_xpu
88
+ else:
89
+ self.forward = self.forward_native
90
elif (
91
logprobs_mode not in ("processed_logits", "processed_logprobs")
92
and rocm_aiter_ops.is_enabled()
@@ -243,6 +248,49 @@ class TopKTopPSampler(nn.Module):
248
return torch.multinomial(renorm_probs, num_samples=1).view(-1)
249
raise RuntimeError("aiter_sample was called with no active top-k or top-p.")
250
251
+ def forward_xpu(
252
+ self,
253
+ logits: torch.Tensor,
254
+ generators: dict[int, torch.Generator],
255
+ k: torch.Tensor | None,
256
+ p: torch.Tensor | None,
257
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
258
+ if generators:
259
+ logger.warning_once(
260
+ "xpu kernel topk_topp_sampler does not support "
261
+ "per-request generators. Falling back to "
262
+ "PyTorch-native implementation."
263
+ )
264
+ return self.forward_native(logits, generators, k, p)
265
+ random_sampled = torch.empty(
266
+ logits.shape[0], dtype=torch.int64, device=logits.device
267
+ )
268
+ logits_to_return = None
269
+ if (
270
+ self.logprobs_mode == "processed_logits"
271
+ or self.logprobs_mode == "processed_logprobs"
272
+ ):
273
+ logits_to_return = torch.empty_like(logits)
274
+
275
+ assert len(generators) != logits.shape[0], (
276
+ "xpu kernel topk_topp_sampler does not support batch-wise generators."
277
+ )
278
+ generator = torch.xpu.default_generators[logits.device.index]
279
+
280
+ state = generator.get_state()
281
+ seed, offset = state.view(torch.int64)
282
+ seeds = torch.tensor(
283
+ [seed, offset], dtype=torch.int64, device=torch.device("cpu")
284
+ )
285
+ # The XPU kernel expects k as int64 (Long), but the input batch
286
+ # stores top_k as int32. Cast here to avoid dtype mismatch.
287
+ if k is not None:
288
+ k = k.to(torch.int64)
289
+ torch.ops.vllm.xpu_topk_topp_sampler(
290
+ random_sampled, logits_to_return, logits, k, p, self.logprobs_mode, seeds
291
+ )
292
+ return random_sampled, logits_to_return
293
+
294
295
# Note: this is a workaround for
296
# https://github.com/pytorch/pytorch/pull/151218