执行摘要

• 一句话：加固 Kimi-K3 AttnRes 分派：强制 packed 行并校验 op 存在
• 推荐动作：值得花约 10 分钟快读。它体现了两个高质量工程原则：一是防御性校验放在内核入口而非依赖调用方自觉——静默错误（无 fault、无 NaN 却超容差 83 倍）比崩溃更难排查，因此用 STD_TORCH_CHECK 转为显式报错；二是 dispatch 校验应检查符号存在性而非仅设备能力，跨编译选项的构建差异需要用 hasattr 兜底。PR body 与 commit message 中给出数值证据、精确定位与复现条件，是高质量变更描述的范式，适合作为 code review 与 bug 报告的模板。对 Kimi-K3/Blackwell 内核维护者尤其相关。

功能与动机

PR body 明确指出两种失效模式：'The kernel hardcodes row stride H, but the dispatch only checked stride(-1) == 1, not stride(0) == hidden_size. A row-padded input silently corrupts output instead of failing (max|out - ref| = 6.66 vs. the suite' 8e-2 tolerance on a repro case)'；以及 'torch.ops._C.kimi_k3_attn_res only exists under VLLM_ENABLE_KIMI_K3_ATTN_RES (CUDA >= 13.0), but the dispatch guard was a device check, not an op-existence check — a CUDA 12.x Blackwell build (e.g. cu129, #50102) raises AttributeError'。作者还引用了 @gau-nernst 的观点，确认 packed rows 是内核的预期前置条件，因此选择 fail loud 而非放宽内核。

实现拆解

本 PR 的修复分为三步：

1. 内核入口强制 packed rows（csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu）：在既有的设备族检查（properties->major == 10）之后、进入 sm100::fwd_prod_v2 逻辑之前，新增一段 STD_TORCH_CHECK。以 prefix.size(1) 作为 hidden_size，要求 prefix、delta、output 三个张量的 stride(0) 均等于 hidden_size，否则在错误消息中回显三者实际 stride 与 hidden_size。原因是内核内部（约 46/386/443/728 行）以硬编码 H 作为行跨度索引，仅对 blocks 透传真实 stride，行填充输入此前会静默读错位置，把这种最难的静默损坏转化为进入内核前即报错。

2. 分派守卫补充 op 存在性检查（vllm/models/kimi_k3/nvidia/ops/attn_res.py）：在原生内核分派条件（hidden_size == 7168、delta is not None、output_norm_weight is not None、num_blocks > 0、block_write_idx < 0、SM100 设备族）末尾新增 hasattr(torch.ops._C, "kimi_k3_attn_res")。原生 op 只在 CUDA >= 13.0、且开启 VLLM_ENABLE_KIMI_K3_ATTN_RES 时编译注册，旧的设备族检查在 CUDA 12.x 的 Blackwell 构建上会通过但 op 不存在，直接访问抛 AttributeError；新检查让这类构建自然落入下方的 Triton 回退路径。文件注释也同步更新，说明设备族检查不足以推断 op 存在。

3. 验证与测试配套：作者声明无测试改动，tests/models/kimi_k3/test_attn_res.py 在 B200（sm_100）上 19 个用例前后均通过；并手工构造 stride(0) = 7175 的 padded 输入，确认现在抛出 RuntimeError: Kimi K3 AttnRes requires densely packed rows; got strides 7175, 7175, 7168 for hidden_size 7168。注意：断言触发路径与 hasattr 回退路径均未新增自动化测试用例，是本次变更的覆盖缺口。

关键文件：

• vllm/models/kimi_k3/nvidia/ops/attn_res.py（模块 模型分派；类别 source；类型 core-logic；符号 attn_res）: 分派守卫扩展的核心文件：在原生内核分派条件中新增 hasattr(torch.ops._C, 'kimi_k3_attn_res') 检查，修复 CUDA 12.x Blackwell 构建（如 cu129）因 op 未编译而抛 AttributeError 的崩溃，并确保此类构建回退到 Triton 路径。
• csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu（模块 内核入口；类别 source；类型 core-logic；符号 kimi_k3_attn_res）: 内核入口新增 STD_TORCH_CHECK，强制 prefix/delta/output 为紧密行布局（stride(0) == hidden_size），把 row-padded 输入从静默数据损坏（超容差 83 倍且无 fault/NaN）转为显式 RuntimeError，是本次 fail-loud 防御的核心落地。

关键符号：attn_res, kimi_k3_attn_res

关键源码片段

vllm/models/kimi_k3/nvidia/ops/attn_res.py

分派守卫扩展的核心文件：在原生内核分派条件中新增 hasattr(torch.ops._C, 'kimi_k3_attn_res') 检查，修复 CUDA 12.x Blackwell 构建（如 cu129）因 op 未编译而抛 AttributeError 的崩溃，并确保此类构建回退到 Triton 路径。

# vllm/models/kimi_k3/nvidia/ops/attn_res.py（节选）
# 函数签名与前面逻辑省略；这里是原生内核与 Triton 路径的分派决策点。
# 原生 SM100 内核以硬编码 H 作为 prefix / delta / output 的行跨度索引，
# 因此除连续内存（stride(-1) == 1）外，还必须保证行间距为 0 padding，
# 即 stride(0) == hidden_size，否则会静默产出超容差结果（B200 实测
# max|out - ref| = 6.66，远超市容差 8e-2）。
assert qk_weight.stride(-1) == 1
assert output_norm_weight is None or output_norm_weight.stride(-1) == 1
​
# 原生 op 只有在开启 VLLM_ENABLE_KIMI_K3_ATTN_RES（要求 CUDA >= 13.0）
# 时才会编译注册，设备族检查无法覆盖 CUDA 12.x 的 Blackwell 构建
# （如 cu129），此时调用不存在的符号会抛 AttributeError，
# 所以必须用 hasattr 确认 op 存在，否则回退到 Triton 实现。
if (
    hidden_size == 7168
    and delta is not None
    and output_norm_weight is not None
    and num_blocks > 0
    and block_write_idx < 0
    and current_platform.is_device_capability_family(100)
    and hasattr(torch.ops._C, "kimi_k3_attn_res")
):
    # 原生路径：覆盖 fused-add + output-norm 的常见组合
    # 与最终的 pre-norm 输出（其余参数省略）
    return ops.kimi_k3_attn_res(prefix, ...)
​
# Triton 路径：处理 block 边界与最终的 pre-norm 输出（本片段省略）

csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu

内核入口新增 STD_TORCH_CHECK，强制 prefix/delta/output 为紧密行布局（stride(0) == hidden_size），把 row-padded 输入从静默数据损坏（超容差 83 倍且无 fault/NaN）转为显式 RuntimeError，是本次 fail-loud 防御的核心落地。

// csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu（节选）
// 内核入口：在设备族检查之后、进入计算逻辑之前，强制输入张量的行布局。
// 内核内部（约 46 / 386 / 443 / 728 行）以硬编码 H 作为 prefix / delta /
// output 的行跨度，仅对 blocks 透传真实 stride；若调用方传入 row-padded
// 张量，数据会被静默读出错误位置，且不会产生 fault 或 NaN，极难排查。
// （函数完整签名省略，此处仅保留入口校验段）
void kimi_k3_attn_res(/* torch::stable::Tensor& prefix、delta、output，
                         blocks 及其它入参 */) {
    cudaDeviceProp const* properties = get_device_prop();
    // AttnRes 内核只面向 SM100 家族（major == 10）
    STD_TORCH_CHECK(properties->major == 10,
                    "Kimi K3 AttnRes requires the SM100 family");
​
    // packed rows 是内核的预期前置条件（@gau-nernst 确认），
    // 因此这里用显式断言让错误输入 " 大声失败 " 而不是静默损坏。
    int64_t const hidden_size = prefix.size(1);
    STD_TORCH_CHECK(
        prefix.stride(0) == hidden_size &&
            delta.stride(0) == hidden_size &&
            output.stride(0) == hidden_size,
        "Kimi K3 AttnRes requires densely packed rows; got strides ",
        prefix.stride(0), ", ", delta.stride(0), ", ",
        output.stride(0), " for hidden_size ", hidden_size);
​
    using namespace sm100::fwd_prod_v2;
    // 两个源 chunk 与两个常驻 CTA 的配置就绪后，进入内核主逻辑（省略）
}

评论区精华

本 PR 的 review 区没有成形讨论线程：claude[bot] 因 PR 来自 fork 而自动跳过审查，gau-nernst 直接 Approved 且未留书面评论。可提炼的技术共识来自 PR body 与 commit message：

• 作者给出了可复现的数值证据与精确定位：'Feeding a row-padded prefix on a B200 (T=320, H=7168, num_blocks=4, prefix.stride(0) == 7175) gave max|out - ref| = 6.66 against the suite's 8e-2 tolerance, with no fault and no NaN'，并指明内核中硬编码行跨度的位置（46/386/443/728 行），让静默损坏的机理一目了然。
• 设计决策的边界清晰：'Packed rows are an intended precondition (per @gau-nernst), so this adds a STD_TORCH_CHECK to fail loudly'——即明确区分'应当支持的输入形态'与'应当拒绝的输入形态'，而不是为了兼容错误输入去扩展内核。

• 该 PR 是 #50090 与 #50000 的 Follow-up，体现了'先合功能、再在真实负载反馈下补防御'的演进节奏。

• Packed rows 作为内核前置条件的确认 (design): 采纳 fail-loud 策略：内核入口新增 STD_TORCH_CHECK 强制 stride(0) == hidden_size，gau-nernst 直接 Approved。

风险与影响

• 风险：
  • 正确性回归（低）：新增的 stride 断言只会拒绝此前已经静默出错的输入；对合规的紧密行调用方行为完全不变。但若未来出现新的调用方以 page 化或 padding 语义传入张量，会被立即拒绝，这是设计意图，但依赖调用方在 CI 中尽早暴露。
  • 性能（极低）：hasattr(torch.ops._C, ...) 位于每次 attn_res 调用的热路径上，CPython 属性查找有一定开销；相对该分派条件既有的多条件判断可忽略，后续可用模块级缓存（启动期确认一次）进一步消除。
  • 兼容性（正面）：CUDA 12.x + Blackwell 从必然 AttributeError 崩溃变为正常回退 Triton；CUDA >= 13 + SM100 行为不变。
  • 覆盖缺口（中）：两条新分支（断言被触发、hasattr 为假）都没有新增自动化测试，错误路径依赖人工验证；鉴于改动行数少且语义直白，风险可控，但建议后续补充'构造 padded 输入断言 raise'的用例。
    • 影响范围限定在 Kimi-K3 模型 + SM100 设备 + AttnRes 内核这条路径，对仓库其他模型与通用执行路径无影响。
  • 影响：对用户：Kimi-K3 + SM100（B200）用户受益最大——此前 row-padded 输入会产生静默的错误输出（数据正确性风险），现在会得到可定位的显式 RuntimeError；CUDA 12.x 的 Blackwell 构建则从启动/调用即崩溃变为自动回退 Triton，可用性明显改善。对系统：改动极小且聚焦，两个文件共 +11/-1，无部署、配置或 schema 变更；运行时性能基本无感。对团队：确立了'内核入口负责前置条件校验、分派侧负责符号是否存在'的防御模式，Kimi-K3 内核维护者可参考；同时本次改动为跨 CUDA 版本构建差异的处理提供了样板。
  • 风险标记：缺少断言触发路径的自动化测试, 热路径新增 hasattr 属性查找（微开销）, 静默损坏转为显式报错，可能暴露上游调用方传参问题, 仅覆盖 Kimi-K3 + SM100 路径

关联脉络

• PR #50656 [Kimi-K3] Add option to shard the shared expert instead of replicating: 同模型线（vllm/models/kimi_k3/nvidia/）的并行配置演进，说明 Kimi-K3 在持续获得 Blackwell 专属内核与配置支持，与本次 AttnRes 分派加固同属 Kimi-K3 增强方向。
• PR #50818 [Kimi-K3] Migrate FlashKDA to PyTorch stable ABI: 与本次改动同处 csrc/libtorch_stable/kimi_k3/ 基础设施线，体现 Kimi 系列内核向 torch::stable API 与 stable ABI 收敛的迁移趋势，二者共同构成 Kimi 内核稳定化脉络。