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 3, 00:30 UTC
[XPU]fix: add XPU platform guards to DeepSeek-V4 ops (#42950)
Signed-off-by: Ma Jian <jian1.ma@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Ma Jian committed
May 23, 2026 at 06:29 UTC
8de5cabeb70df7a9ba835c423b66c2ec51dbeefc
6 files changed
+31
-18
vllm/model_executor/layers/activation.py
+3
-3
@@ -169,9 +169,9 @@ class SiluAndMulWithClamp(CustomOp):
169
def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
170
super().__init__(compile_native=compile_native)
171
self.swiglu_limit = float(swiglu_limit)
172
- if current_platform.is_rocm():
172
+ if current_platform.is_rocm() or current_platform.is_xpu():
173
self._forward_method = self.forward_native
174
- elif current_platform.is_cuda_alike() or current_platform.is_xpu():
174
+ elif current_platform.is_cuda_alike():
175
self.op = torch.ops._C.silu_and_mul_with_clamp
176
elif current_platform.is_cpu():
177
self._forward_method = self.forward_native
@@ -190,7 +190,7 @@ class SiluAndMulWithClamp(CustomOp):
190
return out
191
192
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
193
- return self.forward_cuda(x)
193
+ return self.forward_native(x)
194
195
196
# --8<-- [start:mul_and_silu]
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+5
-1
@@ -243,7 +243,11 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
243
(scale_inner * tma_aligned_T, 1, tma_aligned_T),
244
)
245
grid = (tma_aligned_T, n_groups * heads_per_group)
246
- pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
246
+ pdl_kwargs = (
247
+ {}
248
+ if current_platform.is_rocm() or current_platform.is_xpu()
249
+ else {"launch_pdl": False}
250
+ )
251
_fused_inv_rope_fp8_quant_per_head[grid](
252
o,
253
positions,
vllm/models/deepseek_v4/compressor.py
+5
-1
@@ -296,7 +296,11 @@ class DeepseekCompressor(nn.Module):
296
state_cache = self.state_cache.kv_cache
297
# kv_state stored in first half, score_state stored in second half
298
state_width = state_cache.shape[-1] // 2
299
- pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
299
+ pdl_kwargs = (
300
+ {}
301
+ if current_platform.is_rocm() or current_platform.is_xpu()
302
+ else {"launch_pdl": False}
303
+ )
304
305
# Store the KV and score (with fused APE addition) in the state.
306
# NOTE: PDL is disabled — both this kernel and _fused_kernel below
vllm/models/deepseek_v4/nvidia/model.py
+5
-5
@@ -1153,7 +1153,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1153
x = self.ffn(x, input_ids)
1154
return x, residual, post_mix, res_mix
1155
1156
- def _forward_rocm(
1156
+ def _forward_native(
1157
self,
1158
x: torch.Tensor,
1159
positions: torch.Tensor,
@@ -1193,8 +1193,8 @@ class DeepseekV4DecoderLayer(nn.Module):
1193
) -> tuple[
1194
torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
1195
]:
1196
- if current_platform.is_rocm():
1197
- return self._forward_rocm(
1196
+ if current_platform.is_rocm() or current_platform.is_xpu():
1197
+ return self._forward_native(
1198
x, positions, input_ids, post_mix, res_mix, residual
1199
)
1200
@@ -1228,10 +1228,10 @@ class DeepseekV4Model(nn.Module):
1228
# DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
1229
# (compressor kv_score, indexer.weights_proj, indexer.compressor
1230
# kv_score). fused_wqa_wkv stays on the default stream.
1231
- # Disable them on ROCm because of hang issues.
1231
+ # Disable them on ROCm / XPU because of hang issues / no overlap.
1232
aux_stream_list = (
1233
None
1234
- if current_platform.is_rocm()
1234
+ if current_platform.is_rocm() or current_platform.is_xpu()
1235
else [torch.cuda.Stream() for _ in range(3)]
1236
)
1237
vllm/v1/attention/backends/mla/sparse_swa.py
+5
-1
@@ -367,7 +367,11 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
367
_LAYER_TYPE_C4A: None,
368
_LAYER_TYPE_C128A: None,
369
}
370
- if num_decode_tokens == 0 or current_platform.is_rocm():
370
+ if (
371
+ num_decode_tokens == 0
372
+ or current_platform.is_rocm()
373
+ or current_platform.is_xpu()
374
+ ):
375
return out
376
for layer_type in self._layer_types:
377
# get_mla_metadata() is the official FlashMLA entry point that
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+8
-7
@@ -468,12 +468,13 @@ def fp8_mqa_logits_torch(
468
seq_len_kv = k_fp8.shape[0]
469
k = k_fp8.to(torch.bfloat16)
470
q = q.to(torch.bfloat16)
471
+ device = q.device
472
473
mask_lo = (
473
- torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None]
474
+ torch.arange(0, seq_len_kv, device=device)[None, :] >= cu_seqlen_ks[:, None]
475
)
476
mask_hi = (
476
- torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
477
+ torch.arange(0, seq_len_kv, device=device)[None, :] < cu_seqlen_ke[:, None]
478
)
479
mask = mask_lo & mask_hi
480
@@ -1355,7 +1356,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
1356
assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}"
1357
assert indices.ndim == 1, f"expected indices=[nnz], got {indices.shape}"
1358
assert indptr.ndim == 1, f"expected indptr=[sq+1], got {indptr.shape}"
1358
- assert q.is_cuda and kv.is_cuda and indices.is_cuda and indptr.is_cuda
1359
+ assert not q.is_cpu and not kv.is_cpu and not indices.is_cpu and not indptr.is_cpu
1360
1361
indices = _as_int32_contiguous_1d(indices)
1362
indptr = _as_int32_contiguous_1d(indptr)
@@ -1459,10 +1460,10 @@ def _rocm_sparse_attn_decode_ragged_triton(
1460
)
1461
assert main_indptr.ndim == 1, f"expected main_indptr=[b+1], got {main_indptr.shape}"
1462
assert (
1462
- q.is_cuda
1463
- and main_cache.is_cuda
1464
- and main_indices.is_cuda
1465
- and main_indptr.is_cuda
1463
+ not q.is_cpu
1464
+ and not main_cache.is_cpu
1465
+ and not main_indices.is_cpu
1466
+ and not main_indptr.is_cpu
1467
)
1468
1469
main_indices = _as_int32_contiguous_1d(main_indices)