执行摘要

• 一句话：SM120 flash_mla 分页缓冲改在推断模式外分配
• 推荐动作：值得精读，但仅面向关注 SM120 内核和 CUDA graph 相关机制的工程师。核心设计决策是保持惰性缓冲区分配与使用模式一致性，并修复不对称保护。可借鉴其注释说明写法，但实现简单，无需深度分析。

功能与动机

作者在 PR body 中明确说明 _split_kv_pages_to_64 中存在两个惰性分配的持久缓冲区（buf 和 mbuf），但只有 mbuf 在分配时被 torch.inference_mode(False) 保护。buf 同样会在 CUDA graph 捕获期间被写入，若在 autotune 推理模式下首次分配，则会在捕获时因修改推理张量而触发运行时错误。该修复基于 @moxcat 在 SM120 上运行 DSV4 时的失败报告，当时修复覆盖了两个缓冲区，但 buf 的修复在后续代码中丢失。

实现拆解

1. 定位文件 python/sglang/kernels/ops/attention/flash_mla_sm120.py 中的 _split_kv_pages_to_64 函数。
2. 将 buf = torch.empty(...) 的分配操作包裹在 with torch.inference_mode(False): 上下文中，确保缓冲区在正常模式下分配，可被后续 CUDA graph 捕获修改。
3. 保持缓冲区形状、缓存键、生命周期和填写逻辑不变，仅改变分配时的模式。
4. 无测试、配置或部署配套改动。

关键文件：

• python/sglang/kernels/ops/attention/flash_mla_sm120.py（模块 内核模块；类别 infra；类型 infrastructure；符号 _split_kv_pages_to_64）: 修改了 _split_kv_pages_to_64 中 buf 缓冲区的分配模式，修复潜在运行时错误，是核心变更文件。

关键符号：_split_kv_pages_to_64

关键源码片段

python/sglang/kernels/ops/attention/flash_mla_sm120.py

修改了 _split_kv_pages_to_64 中 buf 缓冲区的分配模式，修复潜在运行时错误，是核心变更文件。

# 关键片段：`_split_kv_pages_to_64` 的惰性缓冲区分配部分
key = f"flash_mla_sm120_split:{dev}"
buf = buffers.get(key)
if buf is None or buf.shape[0] < num_dst_pages:
    # 该缓冲区可能在 autotune（推理模式）下首次分配，
    # 但会在 CUDA graph 捕获（非推理模式）期间被再次写入，
    # 因此必须确保分配在非推理模式下进行，以避免后续修改时出错。
    with torch.inference_mode(False):
        buf = torch.empty(
            num_dst_pages,
            _BYTES_PER_DST_PAGE_PADDED,
            dtype=torch.uint8,
            device=dev,
        )
    buffers[key] = buf
out = buf[:num_dst_pages]

评论区精华

• @moxcat 在评论中提供了详细复现与分析：原始配置仍然存在，已在 v0.5.16 backport 栈和当前 main 上复现。发现缺失的拆分缓冲区崩溃是“写入者类型”问题，而非分配模式问题；arm 排序是真实存在但与分支相关。

• 作者指出相似模式同样存在于 sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py，但本次仅修复 flash_mla 一处。

• 缺失修复的复现与验证 (question): @moxcat 的复现证实了修复的必要性，但无法提供跨分支一致的 before/after 用例。

风险与影响

• 风险：
  • 低风险：仅将分配操作移入 torch.inference_mode(False)，不改变缓冲区逻辑。
  • 潜在回归：若分配确实在非推理模式下发生，则此改动为 no-op；若在推理模式下发生，则修复了潜在崩溃。
  • 风险确认：该修复未能提供可复现的 before/after 用例，但基于代码对称性和 @moxcat 的历史报告，修复方向正确。
• 影响：
  • 影响范围：仅影响 SM120 平台上的 flash_mla 内核路径，具体是 DSV4 等模型在 autotune+CUDA graph 场景下的稳定性。
  • 对用户：消除 SM120 上潜在的运行时崩溃，提升部署稳定性。
  • 对团队：作为 #29927（SM120 支持）的前置依赖，虽不在其范围内，但被其依赖。
  • 风险标记：缺少可复现用例, 无测试覆盖

关联脉络

• PR #29927 [SM120] enablement work (未提供标题): 本 PR 从 #29927 拆分而来，且 #29927 依赖它完成 SM120 支持。