Prhub

#53574 [Bugfix][SM120] DSv4: pass contiguous C128A decode topk indices on SM120

原始 PR 作者 lucifer1004 合并时间 2026-08-31 17:44 文件变更 2 提交数 5 评论 16 代码增减 +31 / -7

执行摘要

修复 SM120 上 DSv4 投机解码启动崩溃,eidx 改全宽连续

该 PR 解决 DSv4-Flash 在 SM120 上启用 FlashInfer sparse-MLA 后端时服务器启动即崩溃的生产阻断问题。PR body 引用崩溃断言 Check failed: (eidx.IsContiguous()) is false: eidx must be contiguous,并定位根因:PR #52823 的自适应 topk 宽度使 build_c128a_topk_metadata 返回保持行 stride 的收窄切片,而 FlashInfer 对超过 64-token 的调用走 paged orchestrator 并强制 eidx 连续;DSpark 验证批次 num_decodes * (1 + K) 恰好越过阈值,因此在 CUDA graph 捕获阶段暴露。maci0 在 DGX Spark(sm_121a)上独立复现同一机制,并确认 C4A 分支不受影响,佐证了根因分析。

值得精读。该 PR 是「供应商 kernel 契约与框架内存布局冲突」的教科书案例:理解 FlashInfer 64-token 分派阈值如何决定崩溃暴露时机、CUDA graph 捕获下地址稳定性与 contiguity 的取舍(为何不能在消费端 .contiguous())、以及 reviewer 坚持「生产代码简洁 + 测试技巧」的工程权衡,对后续 FlashInfer/FlashMLA 集成工作有直接参考价值。

讨论亮点

核心讨论围绕「短期修复与长期方案」的取舍展开:

  • yewentao256 在 CHANGES_REQUESTED 中指出:长期正确修复是让 FlashInfer kernel 接受真实 stride(0),但同意先收短期修复;同时明确要求 SM100 行为不变,通过 if 分支区分后端。
  • lucifer1004 询问内联能力判断后单测如何覆盖两条分支,提出 monkeypatch 平台单例或保留 keyword arg seam 两种方案;yewentao256 回复「让生产代码保持简单,测试可以用技巧」,最终采用 monkeypatch current_platform.get_device_capability。
  • 讨论还涉及去除 docstring 中冗余解释性注释,保留 TODO 记录 SM120 自适应宽度的长期方向。

实现拆解

  1. 入口与问题定位:vllm/models/deepseek_v4/sparse_mla.py 的 build_c128a_topk_metadata 是所有 C128A token(decode + prefill)的单一元数据构建入口,写入预分配缓冲区以保证 CUDA graph 地址稳定。PR #52823 引入自适应 topk 宽度后,decode 视图变为 global_decode_buffer[:num_decode_tokens, :max_compressed_tokens],即按活动宽度收窄的行切片,保留了缓冲区原始行 stride,导致 is_contiguous() == False;该视图随后作为 eidx 传入 FlashInfer SM120 sparse-MLA kernel。
  2. 触发机制确认:FlashInfer 对 num_tokens <= 64 的调用分派到独立 decode kernel(不做 eidx 连续性检查),更大的调用走 paged orchestrator 并执行 CHECK_INPUT_AND_TYPE(eidx, ...)。纯 target-only decode 始终低于 64 阈值;DSpark 投机解码的验证批次携带 num_decodes * (1 + K) 行,越过阈值,因此在服务器启动阶段的 CUDA graph 捕获中即崩溃。maci0 在 DGX Spark(sm_121a)上用 DSpark k=5 独立复现同一错误,并验证了 C4A 分支不受影响。
  3. 核心修复:build_c128a_topk_metadata 内部通过 current_platform.get_device_capability() 判断,仅当 capability major == 12(SM120)时 decode 视图改为全宽行切片 global_decode_buffer[:num_decode_tokens],该切片连续且位于持久 buffer 内;其余平台(SM100 FlashMLA/TRTLLM)保持自适应宽度切片。prefill 视图对所有后端保持收窄。设计依据:decode 消费端按 per-token topk lens 限界读取,全宽带来的陈旧列不会被读取;prefill 的 Triton 消费者统计行内非负条目,不能带陈旧列。
  4. 测试配套:tests/kernels/attention/test_flashmla_sparse.py 的 test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride 参数化 sm120 两个取值,monkeypatch sparse_mla.current_platform.get_device_capability 覆盖两条分支,断言 decode 视图的 shape、is_contiguous() 与 stride,并在 CUDA graph 捕获重放后校验数据正确性,确保 capture-stable 语义不回归。
  5. 演进与收尾:初始版本通过 full_width_decode 关键字参数传递标志,review 后内联为平台能力判断(commit bdff48b9),生产代码保持简单,测试用 monkeypatch 技巧;代码中保留 TODO,明确长期方向是让 FlashInfer SM120 kernel 接受 eidx 的真实 row stride。
文件 模块 状态 重要度
vllm/models/deepseek_v4/sparse_mla.py 稀疏 MLA modified 6.61
tests/kernels/attention/test_flashmla_sparse.py 内核测试 modified 5.39

关键符号

build_c128a_topk_metadata test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride

关键源码片段

vllm/models/deepseek_v4/sparse_mla.py data-contract

修复核心:`build_c128a_topk_metadata` 按设备能力决定 C128A decode 视图宽度,SM120 返回全宽连续行切片以满足 FlashInfer eidx 连续性检查并保持 CUDA graph 地址稳定。

# vllm/models/deepseek_v4/sparse_mla.py(PR 后状态)
# 函数入口处已完成基本校验:num_tokens 由 positions 推出,max_compressed_tokens
# 满足对齐与容量约束,缓冲区最后一个维度的 stride 为 1
​
    # TODO: support adaptive-width decode on SM120
    # (需要 FlashInfer SM120 kernel 接受 eidx 的真实 row stride,届时可移除分支)
    capability = current_platform.get_device_capability()
    if capability is not None and capability.major == 12:
        # SM120(FlashInfer sparse-MLA 后端):decode 取全宽行切片。
        # 整行连续,满足 paged orchestrator 对 eidx 的 contiguity 检查;
        # 且切片指向持久 buffer 内部,CUDA graph 捕获期间地址稳定。
        # 不能改为消费端 .contiguous():每步拷贝会分配新张量、破坏捕获。
        global_decode = global_decode_buffer[:num_decode_tokens]
    else:
        # SM100 等平台保持自适应宽度切片:FlashMLA / TRTLLM 接受带 stride 的
        # 视图,收窄可减小 kernel 扫描范围,行为与 PR 前完全一致。
        global_decode = global_decode_buffer[
            :num_decode_tokens, :max_compressed_tokens
        ]
    decode_lens = decode_lens_buffer[:num_decode_tokens]
    # prefill 视图对所有后端保持收窄:Triton 消费者统计行内非负条目,
    # 若带回全宽会误扫尾部陈旧列(padding 的 -1 位)。
    prefill_local = prefill_buffer[:num_prefill_tokens, :max_compressed_tokens]
    assert global_decode.stride(0) == global_decode_buffer.stride(0)
    assert prefill_local.stride(0) == prefill_buffer.stride(0)
​
    if num_tokens == 0:
        return global_decode, decode_lens, prefill_local
​
    # 启动 Triton kernel 写入缓冲区,随后返回三个视图
    _build_c128a_topk_metadata_kernel[(num_tokens,)](
        global_decode_buffer,
        global_decode_buffer.stride(0),
        decode_lens_buffer,
        prefill_buffer,
        prefill_buffer.stride(0),
        positions,
        compress_ratio,
        max_compressed_tokens,
        num_decode_tokens,
        token_to_req_indices,
        block_table,
        block_table.stride(0),
        block_size,
        slot_mapping,
        BLOCK_SIZE=1024,
    )
    return global_decode, decode_lens, prefill_local
tests/kernels/attention/test_flashmla_sparse.py test-coverage

回归测试:参数化 `sm120` 双分支,monkeypatch 平台能力,断言解码视图全宽 / 连续性与 CUDA graph 重放正确性,覆盖本次修复的核心契约。

# tests/kernels/attention/test_flashmla_sparse.py(PR 后状态)
​
@pytest.mark.parametrize('sm120', [False, True])
def test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride(
    monkeypatch: pytest.MonkeyPatch, sm120: bool,
):
    from vllm.models.deepseek_v4 import sparse_mla
    from vllm.platforms.interface import DeviceCapability
​
    # monkeypatch 平台能力,让同一测试在任意 GPU 上可覆盖
    # SM120 全宽视图与 SM100 自适应宽度切片两条分支
    monkeypatch.setattr(
        sparse_mla.current_platform,
        'get_device_capability',
        lambda: DeviceCapability(12, 0) if sm120 else DeviceCapability(10, 0),
    )
​
    # 构造 global_decode_buffer / prefill_buffer 与 kwargs 的样板代码省略
    captured_decode, _, captured_prefill = sparse_mla.build_c128a_topk_metadata(
        max_compressed_tokens=256, **kwargs,
    )
    # SM120 下 decode 应为全宽且连续;其他平台收窄到 max_compressed_tokens;
    # prefill 视图始终收窄。stride(0) 不变说明仍是持久 buffer 的行切片。
    expected_decode_width = capacity_width if sm120 else 256
    assert captured_decode.shape == (2, expected_decode_width)
    assert captured_prefill.shape == (2, 256)
    assert captured_decode.stride(0) == captured_prefill.stride(0) == capacity_width
    assert captured_decode.is_contiguous() == sm120
​
    # 把视图封进 CUDA graph 捕获并重放:重新填充 buffer 后重放仍读到
    # 正确数据,验证视图地址在捕获期间稳定(capture-stable stride 回归保障)
    captured_rows = torch.empty((4, 4), dtype=torch.int32, device=device)
    captured_rows[:2].copy_(captured_decode[:, :4])
    captured_rows[2:].copy_(captured_prefill[:, :4])
    torch.accelerator.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        captured_rows[:2].copy_(captured_decode[:, :4])
        captured_rows[2:].copy_(captured_prefill[:, :4])
​
    global_decode_buffer.fill_(-99)
    prefill_buffer.fill_(-99)
    sparse_mla.build_c128a_topk_metadata(max_compressed_tokens=128, **kwargs)
    graph.replay()
​
    # 重放后应读到与捕获时一致的 slot ids,证明视图地址在捕获期间稳定
    assert captured_rows.cpu().tolist() == [
        [1536, 1537, -1, -1],
        [2560, 2561, 2562, 2563],
        [0, 1, 2, -1],
        [0, 1, 2, 3],
    ]

评论区精华

SM100 与 SM120 后端行为差异需分分支处理 设计

yewentao256 指出该切片方式对 SM100 的 FlashMLA 有效,但可能不适用于 SM120 的 FlashInfer sparse-MLA,要求加 if 分支且保持 SM100 不变。

结论:接受建议:decode 视图按设备能力分支,SM120 全宽、SM100 收窄。 · 已解决

去掉 full_width_decode 参数,内联能力判断 设计

yewentao256 建议移除额外参数,把结构判断内联进 build_c128a_topk_metadata,并加 TODO 记录 SM120 adaptive width 支持。

结论:已按建议内联,lucifer1004 完成于 commit bdff48b9。 · 已解决

内联后单测如何覆盖两条分支 测试

lucifer1004 担心能力判断内联后分支无法从单测直接选择,询问是否接受 monkeypatch 平台单例,或保留 keyword arg seam。yewentao256 回复「让生产代码保持简单,测试用技巧」。

结论:采用 monkeypatch current_platform.get_device_capability 覆盖双分支。 · 已解决

docstring 冗余注释删减 style

yewentao256 两次 suggestion 要求删除 docstring 中过长的解释性段落,认为注释不是必需。

结论:按建议精简注释,保留 TODO。 · 已解决

长期修复方向:让 FlashInfer kernel 接受真实 row stride 设计

yewentao256 在 CHANGES_REQUESTED 中明确指出长期正确修复是让 FlashInfer kernel 接受真实 stride(0),本次是短期解除阻塞;PR body 与代码 TODO 均记录该方向。

结论:接受短期修复,TODO 留存长期方向。 · 待处理

风险与影响

  1. 全宽视图的陈旧列:SM120 decode 视图包含超出活动 topk 宽度的历史数据列。当前所有 decode 消费端均按 per-token topk lens 限界读取,但若未来新增消费端直接对整行做规约或拷贝,会读到陈旧索引。相关逻辑集中在 vllm/models/deepseek_v4/sparse_mla.py 与 DeepseekV4FlashInferSM120Attention._forward_decode。
  2. 依赖 FlashInfer 分派实现:修复依赖「64-token 以下 standalone decode kernel 不检查 eidx 连续性」的当前 FlashInfer 行为。若上游未来对 standalone kernel 增加同款检查,纯 target-only decode 也会触发同一崩溃,届时需要 kernel 支持真实 row stride 或统一走全宽视图。
  3. 平台判断分支:current_platform.get_device_capability() 返回 None 或非 SM120 时走收窄分支,行为与旧版一致,回归面小;但该判断属于运行时平台探测,无法静态保证,测试通过 monkeypatch 覆盖。
  4. CI 干扰项:首次 CI 中 Multimodal Models (Standard) 4 超时被判定为基础设施 flake(MOSS-Audio 测试挂起 71 分钟),与本次改动无关,已由维护者重跑确认。

用户影响:解除 RTX 6000 Pro、DGX Spark 等 SM120 设备上 DeepSeek-V4-Flash + DSpark 投机解码的启动崩溃;PR body 记录 4x RTX 6000 Pro 上 8K/1K 随机 benchmark 零失败请求。系统影响:改动仅限于 C128A 元数据构建的视图选取,无接口或 schema 变化;CUDA graph 捕获路径更稳定。团队影响:测试通过 monkeypatch 提供双分支覆盖,为同类「平台相关内存视图」的测试提供了可复用模式;代码中 TODO 明确了与 FlashInfer 上游的协作方向。

SM120 条件分支 核心注意力路径变更 依赖 FlashInfer 64-token 分派 全宽视图含陈旧列 长期 TODO 待跟进

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论