From 3b4fac5b99670f3e8bcb5bc9ecf04b821e7b8de5 Mon Sep 17 00:00:00 2001 From: Xuan Liao Date: Wed, 5 Aug 2026 21:05:53 +0800 Subject: [PATCH] [XPU] DeepSeek V4: use sgl-kernel-xpu implemetation of flash_mla_sparse_fwd for prefill (#31865) Co-authored-by: Ma Mingfei --- python/sglang/srt/layers/attention/deepseek_v4_backend.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index f85b3bfc6..402d3baa0 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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)