[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>

vllm/model_executor/layers/activation.py

@@ -169,9 +169,9 @@ class SiluAndMulWithClamp(CustomOp):
     def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
         super().__init__(compile_native=compile_native)
         self.swiglu_limit = float(swiglu_limit)
-        if current_platform.is_rocm():
+        if current_platform.is_rocm() or current_platform.is_xpu():
             self._forward_method = self.forward_native
-        elif current_platform.is_cuda_alike() or current_platform.is_xpu():
+        elif current_platform.is_cuda_alike():
             self.op = torch.ops._C.silu_and_mul_with_clamp
         elif current_platform.is_cpu():
             self._forward_method = self.forward_native
@@ -190,7 +190,7 @@ class SiluAndMulWithClamp(CustomOp):
         return out
 
     def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
-        return self.forward_cuda(x)
+        return self.forward_native(x)
 
 
 # --8<-- [start:mul_and_silu]

vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py

@@ -243,7 +243,11 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
         (scale_inner * tma_aligned_T, 1, tma_aligned_T),
     )
     grid = (tma_aligned_T, n_groups * heads_per_group)
-    pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
+    pdl_kwargs = (
+        {}
+        if current_platform.is_rocm() or current_platform.is_xpu()
+        else {"launch_pdl": False}
+    )
     _fused_inv_rope_fp8_quant_per_head[grid](
         o,
         positions,

vllm/models/deepseek_v4/compressor.py

@@ -296,7 +296,11 @@ class DeepseekCompressor(nn.Module):
         state_cache = self.state_cache.kv_cache
         # kv_state stored in first half, score_state stored in second half
         state_width = state_cache.shape[-1] // 2
-        pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
+        pdl_kwargs = (
+            {}
+            if current_platform.is_rocm() or current_platform.is_xpu()
+            else {"launch_pdl": False}
+        )
 
         # Store the KV and score (with fused APE addition) in the state.
         # NOTE: PDL is disabled — both this kernel and _fused_kernel below

vllm/models/deepseek_v4/nvidia/model.py

@@ -1153,7 +1153,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         x = self.ffn(x, input_ids)
         return x, residual, post_mix, res_mix
 
-    def _forward_rocm(
+    def _forward_native(
         self,
         x: torch.Tensor,
         positions: torch.Tensor,
@@ -1193,8 +1193,8 @@ class DeepseekV4DecoderLayer(nn.Module):
     ) -> tuple[
         torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
     ]:
-        if current_platform.is_rocm():
-            return self._forward_rocm(
+        if current_platform.is_rocm() or current_platform.is_xpu():
+            return self._forward_native(
                 x, positions, input_ids, post_mix, res_mix, residual
             )
 
@@ -1228,10 +1228,10 @@ class DeepseekV4Model(nn.Module):
         # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
         # (compressor kv_score, indexer.weights_proj, indexer.compressor
         # kv_score). fused_wqa_wkv stays on the default stream.
-        # Disable them on ROCm because of hang issues.
+        # Disable them on ROCm / XPU because of hang issues / no overlap.
         aux_stream_list = (
             None
-            if current_platform.is_rocm()
+            if current_platform.is_rocm() or current_platform.is_xpu()
             else [torch.cuda.Stream() for _ in range(3)]
         )
 

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -367,7 +367,11 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             _LAYER_TYPE_C4A: None,
             _LAYER_TYPE_C128A: None,
         }
-        if num_decode_tokens == 0 or current_platform.is_rocm():
+        if (
+            num_decode_tokens == 0
+            or current_platform.is_rocm()
+            or current_platform.is_xpu()
+        ):
             return out
         for layer_type in self._layer_types:
             # get_mla_metadata() is the official FlashMLA entry point that

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -468,12 +468,13 @@ def fp8_mqa_logits_torch(
     seq_len_kv = k_fp8.shape[0]
     k = k_fp8.to(torch.bfloat16)
     q = q.to(torch.bfloat16)
+    device = q.device
 
     mask_lo = (
-        torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None]
+        torch.arange(0, seq_len_kv, device=device)[None, :] >= cu_seqlen_ks[:, None]
     )
     mask_hi = (
-        torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
+        torch.arange(0, seq_len_kv, device=device)[None, :] < cu_seqlen_ke[:, None]
     )
     mask = mask_lo & mask_hi
 
@@ -1355,7 +1356,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
     assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}"
     assert indices.ndim == 1, f"expected indices=[nnz], got {indices.shape}"
     assert indptr.ndim == 1, f"expected indptr=[sq+1], got {indptr.shape}"
-    assert q.is_cuda and kv.is_cuda and indices.is_cuda and indptr.is_cuda
+    assert not q.is_cpu and not kv.is_cpu and not indices.is_cpu and not indptr.is_cpu
 
     indices = _as_int32_contiguous_1d(indices)
     indptr = _as_int32_contiguous_1d(indptr)
@@ -1459,10 +1460,10 @@ def _rocm_sparse_attn_decode_ragged_triton(
     )
     assert main_indptr.ndim == 1, f"expected main_indptr=[b+1], got {main_indptr.shape}"
     assert (
-        q.is_cuda
-        and main_cache.is_cuda
-        and main_indices.is_cuda
-        and main_indptr.is_cuda
+        not q.is_cpu
+        and not main_cache.is_cpu
+        and not main_indices.is_cpu
+        and not main_indptr.is_cpu
     )
 
     main_indices = _as_int32_contiguous_1d(main_indices)