执行摘要

• 一句话：修复 Intel XPU 上 MLA workspace 计算错误
• 推荐动作：该 PR 是 Intel XPU 平台的关键修正，解决了平台差异导致的正确性问题。值得关注其参数替换思路，为后续多平台适配提供参考。

功能与动机

Intel XPU 设备不支持 CUDA 的 sm_count 概念，原调用中 flash_mla_get_workspace_size 传入了 sm_count=get_device_core_count()，导致 workspace 大小计算错误。PR 将此参数替换为 num_heads 和 page_size，符合 Intel XPU 设备的工作空间分配要求。

实现拆解

1. 移除 get_device_core_count 导入及相关无用代码。
2. 在 __init__ 中添加 num_attention_heads、tp_size、num_local_heads 属性，用于后续参数计算。
3. 修改 flash_mla_get_workspace_size 调用：将位置参数改为关键字参数，用 num_heads=self.num_local_heads 和 page_size=self.page_size 替换 sm_count=get_device_core_count()。

关键文件：

• python/sglang/srt/layers/attention/xpu_backend.py（模块 注意力后端；类别 source；类型 dependency-wiring）: 唯一变更文件，修复了 MLA workspace 计算中的平台相关参数，确保 Intel XPU 设备正确工作。

关键符号：XPUAttentionBackend.init, XPUAttentionBackend.init_forward_metadata

关键源码片段

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

唯一变更文件，修复了 MLA workspace 计算中的平台相关参数，确保 Intel XPU 设备正确工作。

# python/sglang/srt/layers/attention/xpu_backend.py
# 移除了 get_device_core_count 导入，因为它不适用于 Intel XPU
# XPU 设备没有 CUDA 的 SM 概念，需要使用 num_local_heads 和 page_size
​
from sgl_kernel import flash_mla_decode, flash_mla_get_workspace_size, merge_state_v2
from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
​
class XPUAttentionBackend(AttentionBackend):
    def __init__(self, model_runner, ...):
        # ... 其他初始化 ...
        # 新增 : 获取模型参数，用于计算本地注意力头数
        self.num_attention_heads = (
            model_runner.model_config.hf_text_config.num_attention_heads
        )
        self.tp_size = model_runner.tp_size
        assert self.num_attention_heads % self.tp_size == 0
        self.num_local_heads = self.num_attention_heads // self.tp_size
​
    def init_forward_metadata(self, forward_batch):
        # ... 其他处理 ...
        if self.use_mla:
            # 修复 : 使用关键字参数，避免平台相关参数
            # 原代码使用 sm_count=get_device_core_count()，但 Intel XPU 不支持
            workspace_size = flash_mla_get_workspace_size(
                max_seq_len=self.max_context_len,
                num_batches=batch_size,
                num_heads=self.num_local_heads, # 新增 : 使用本地头数
                page_size=self.page_size, # 新增 : 指定 page size
                num_kv_splits=-1, # 保持默认
            )

评论区精华

review 中讨论了 workspace 大小计算的两种策略：一是基于 max_context_len 计算 num_blocks（可能导致长序列块过多），二是基于 num_cores 控制块数量以优化 occupancy。评审者 mingfeima 和 pralay-das 一致认为基准测试中第一种策略表现更好，本 PR 采用前者。

• Workspace 大小计算策略 (performance): 采用 num_blocks = div_up(max_context_len, block_size) 方法，即基于 max_seq_len 和 page_size 计算。

风险与影响

• 风险：仅修改了 Intel XPU 后端的 MLA workspace 计算逻辑，不涉及 CUDA 或其他后端。风险较低，但需确保 XPU 设备上的 page_size 和 num_local_heads 计算正确，否则可能导致 workspace 分配不足或溢出。
• 影响：影响范围局限于 Intel XPU 后端，且仅影响 MLA 注意力机制的 workspace 计算。正确修复后，Intel XPU 上的 DeepSeek 等 MLA 模型应能正常分配 workspace，避免崩溃或性能下降。
• 风险标记：平台特定适配, 缺少测试覆盖

关联脉络

• PR #24005 [AMD] Enable dual-stream MoE on ROCm: 同为厂商特定硬件支持，类似的多平台适配模式。
• PR #24550 [R3] Avoid implicit CUDA sync in routed experts DP slicing: 同为性能优化和平台适配相关。