[XPU] DeepSeek V4: use sgl-kernel-xpu implemetation of flash_mla_sparse_fwd for prefill (#31865)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
2026-08-05 21:05:53 +08:00
parent 99709f734d
commit 3b4fac5b99
1 changed files with 4 additions and 1 deletions
@@ -1780,7 +1780,10 @@ class DeepseekV4AttnBackend(
        indices. Chunk-invariant scaffolding lives in
        ``self.forward_metadata.sparse_prefill_cache``.
        """
-        from sgl_kernel.flash_mla import flash_mla_sparse_fwd
+        if _is_xpu:
+            from sgl_kernel import flash_mla_sparse_fwd
+        else:
+            from sgl_kernel.flash_mla import flash_mla_sparse_fwd

        # q is (b, 1, h_q, d_qk); flash_mla_sparse_fwd takes (s_q, h_q, d_qk).
        q_flat = q.squeeze(1)