执行摘要

• 一句话：XPU 接入 flash_mla_sparse_fwd 支撑 DSV4 稀疏 prefill
• 推荐动作：值得快速阅读。虽然改动只有 5 行，但 review 中关于 return_softmax_lse 语义对齐的讨论很有参考价值：平台差异应尽量在 kernel 侧收敛而不是在调用方补分支。同时可作为 XPU 平台接入 sgl-kernel 算子的最小样例，与历史 PR #28040 的 fused_k_norm_rope_flashmla 接入形成连续脉络。

功能与动机

PR body 明确写着 Depend on sgl-kernel-xpu implementation https://github.com/sgl-project/sgl-kernel-xpu/pull/337，即需要依赖 sgl-kernel-xpu 提供的 flash_mla_sparse_fwd 实现，让 XPU 平台的 DSV4 稀疏 MLA prefill 能够直接调用该算子，而不是退回 flash_mla_with_kvcache 的扩展路径。首个 commit 信息 '[xpu] support sparse mla prefill' 也印证这是 XPU 上的能力补齐，属于 DSV4 XPU 内核接入系列的延续。

实现拆解

1. 变更入口：python/sglang/srt/layers/attention/deepseek_v4_backend.py 的 _forward_prefill_sparse 方法，它是 DSV4 稀疏 MLA prefill（compress_ratio 为 0/4/128）的统一入口，docstring 明确其取代 extend 路径上的 flash_mla_with_kvcache 调用，并按请求把 SWA 窗口与压缩缓存（c4/c128）gather 到 flat bf16 workspace 中。
2. 平台分流导入：原实现固定执行 from sgl_kernel.flash_mla import flash_mla_sparse_fwd；改动后当 _is_xpu 为真时改为 from sgl_kernel import flash_mla_sparse_fwd（sgl-kernel-xpu 的导出位置），否则维持原路径。CUDA 侧行为零变化，XPU 侧获得新算子入口。
3. 调用语义收敛：review 阶段曾尝试在调用处为 XPU 额外传 return_softmax_lse=True（因为 XPU 算子默认不返回 lse 三元组，而 CUDA 默认返回），该 extra_kwargs 改动最终被移除，改为在 sgl-kernel-xpu 侧对齐 CUDA 的默认返回行为，从而保持 SGLang 调用点统一。
4. 测试与依赖：本仓库未新增测试文件；正确性依赖外部 sgl-kernel-xpu#337 的算子实现，并由 run-ci / run-ci-extra 的 XPU 相关 CI 覆盖。后续若 XPU 算子返回语义再出现差异，应在 kernel 仓库对齐而非在调用方打补丁。

关键文件：

• python/sglang/srt/layers/attention/deepseek_v4_backend.py（模块 注意力后端；类别 source；类型 dependency-wiring；符号 _forward_prefill_sparse）: DSV4 稀疏 MLA prefill 的统一入口 _forward_prefill_sparse 中按平台分流 flash_mla_sparse_fwd 的导入来源，是本次 PR 唯一的源码变更，决定 XPU 能否调用 sgl-kernel-xpu 的算子实现。

关键符号：_forward_prefill_sparse

关键源码片段

python/sglang/srt/layers/attention/deepseek_v4_backend.py

DSV4 稀疏 MLA prefill 的统一入口 _forward_prefill_sparse 中按平台分流 flash_mla_sparse_fwd 的导入来源，是本次 PR 唯一的源码变更，决定 XPU 能否调用 sgl-kernel-xpu 的算子实现。

def _forward_prefill_sparse(self, q, layer_id, compress_ratio, forward_batch,
                            token_to_kv_pool, core_attn_metadata, attn_sink):
    """统一 prefill 路径：通过 flash_mla_sparse_fwd 处理 DSV4 稀疏 MLA 注意力。"""
    # XPU 的 sgl-kernel-xpu 把 flash_mla_sparse_fwd 导出在 sgl_kernel 顶层 ;
    # CUDA 的 sgl-kernel 则放在 sgl_kernel.flash_mla 子模块。
    # 这里按平台分流导入，避免 XPU 上出现 ImportError。
    if _is_xpu:
        from sgl_kernel import flash_mla_sparse_fwd
    else:
        from sgl_kernel.flash_mla import flash_mla_sparse_fwd
​
    # q 的形状是 (b, 1, h_q, d_qk), 而算子期望 (s_q, h_q, d_qk),
    # 因此先 squeeze 掉长度固定为 1 的序列维。
    q_flat = q.squeeze(1)
​
    # sparse prefill 的 workspace/ 索引缓存按 batch 构建 ,
    # 同一请求的多个 chunk 复用 , 避免重复分配。
    cache = self.forward_metadata.sparse_prefill_cache
    if cache is None:
        seq_lens_cpu = forward_batch.seq_lens_cpu
        assert seq_lens_cpu is not None
        # 注意 : token_to_kv_pool.swa_window_size 是存储页大小 ,
        # 模型真实 SWA 窗口由 SWA_WINDOW 传入 , 两者必须显式区分。
        cache = SparsePrefillChunkCache.build(
            seq_lens=forward_batch.seq_lens.to(torch.int32),
            extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32),
            req_pool_indices=forward_batch.req_pool_indices.to(torch.int32),
            req_to_token=self.req_to_token,
            full_to_swa=token_to_kv_pool.full_to_swa_index_mapping,
            swa_window_size=SWA_WINDOW,
            swa_page_size=token_to_kv_pool.swa_window_size,
            num_qo_tokens=q_flat.shape[0],
            max_seq_len=int(seq_lens_cpu.max().item()),
        )
        self.forward_metadata.sparse_prefill_cache = cache

评论区精华

polisettyvarma 在 diff 中提出重要语义差异："XPU can also always return 3 with unused as None return_softmax_lse is False"，即 XPU 算子只有在 return_softmax_lse 为 True 时才返回 (out, max_logits, lse) 三元组，而 CUDA 默认返回三元组，因此原改动在调用处为 XPU 加了 extra_kwargs = {"return_softmax_lse": True}。Valentine233 回复 "Thanks. I'll remove the change and update the xpu's kernel."，接受了建议，最终将兼容逻辑下沉到 sgl-kernel-xpu 侧，本仓库只保留纯净的导入分流。这一取舍保证了调用方代码不因平台产生语义分支。

• XPU 算子的 return_softmax_lse 返回语义差异 (correctness): Valentine233 接受建议，移除调用处的 extra_kwargs，改为在 sgl-kernel-xpu 侧更新 kernel，使默认返回行为与 CUDA 对齐，保持 SGLang 调用点统一。

风险与影响

• 风险：
  1. 平台分支隔离：新增的 _is_xpu 条件分支位于 deepseek_v4_backend.py 的 prefill 热路径，若 _is_xpu 判定与实际运行环境不一致，XPU 上会走错导入路径并触发 ImportError。
  2. 外部 kernel 依赖：正确性完全依赖 sgl-kernel-xpu#337 的算子实现，本仓库没有单元测试直接覆盖该导入分支，sgl-kernel-xpu 侧若未完全对齐 CUDA 的返回语义，会导致解包不一致。
  3. CI 覆盖不足：PR 未新增测试文件，回归保障依赖 run-ci-extra 的 XPU CI，如果 CI 未覆盖 DSV4 sparse prefill 场景，问题可能漏检。
     - 影响：影响范围较小且明确：仅影响 XPU 平台上 DeepSeek V4 稀疏 MLA prefill 路径，CUDA 及其他平台行为不变。对用户而言，XPU 上 DSV4 的 prefill 可获得 flash_mla_sparse_fwd 内核带来的性能与功能支持；对团队而言，这是 Intel XPU 接入 sgl-kernel 系内核的又一步，后续同类算子接入可以沿用这套平台分流导入模式。
     - 风险标记：核心路径变更, 缺少测试覆盖, 依赖外部 kernel 仓库

关联脉络

• PR #28040 [Intel GPU] DeepSeek V4 8/N: use sgl-kernel implementation of fused_k_norm_rope_flashmla on XPU: 同一功能线：XPU 平台接入 sgl-kernel/sgl-kernel-xpu 内核以优化 DeepSeek V4 推理，本次 PR 是其延续。
• PR #33140 [DSV4] Add official DSV4 reasoning effort support: 同为 DeepSeek V4 的功能演进，反映 DSV4 在 SGLang 中的持续完善脉络。