执行摘要

本 PR 对 Kimi K3 的注意力残差内核 csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu 做延迟优化，四项改动叠加后微基准在 1-2048 token 区间延迟下降 9%-27%，8×B300 端到端 TPOT/TTFT 普遍改善。改动只涉及一个文件（+84/-39），三位维护者直接 APPROVED，属于低风险、量化收益明确的性能优化。

功能与动机

Kimi K3 的注意力残差融合内核在 decode 与 prefill 中频繁执行，直接影响 TPOT/TTFT。PR body 列出的四项改进：

• Use vectorized loads when populating q_cache
• Treat B and N as compile-time constants
• Set NC = 3 for low batch size path
• Unroll the main loop over num_chunks

微基准显示 Token=1 时延迟从 5.23ms 降至 4.16ms（-25.72%），Token=1024 时从 22.74ms 降至 19.05ms（-19.37%）。

实现拆解

1. 编译期参数化 B 与 N（attn_res_fwd_online_v2_kernel / launch_fwd）
   - 模板从 <int H, int NC, ...> 扩展为 <int H, int N, int NC, int B, ...>，num_chunks 变为 constexpr，这是后续循环展开的前提。
   - launch_fwd 与内核签名同步移除运行时 N、B 参数，调用点只传 T 与各 stride。

2. q_cache 向量化填充
   - H==7168 特化分支与通用分支均改用 int2/int4 向量读取 rms_w/res_w，经 __nv_bfloat1622float2 成对计算 rms * res，减少标量加载与乘加指令密度。

3. 低 batch 路径 NC=3
   - kimi_k3_attn_res 通用分支用 switch(num_blocks)（case 1 / case 2 / default）分派到不同 NSRC 编译期模板参数，并对每组调用 launch_fwd<7168, NSRC, 3, ...>，使 chunk 数、每 chunk 源数与前缀 chunk 判断全部编译期化。
   - 长 prefill 分支（blocks==8 且 tokens>=4096）保留 NC=2 双源双驻留 CTA 配置。

4. 主循环展开与内存序配套
   - 对 for (ci...) 主循环加 #pragma unroll。
   - 给 mbarrier_wait / mbarrier_arrive 的内联汇编补 "memory" clobber，避免编译器优化破坏 barrier 与访存之间的顺序语义。

5. 测试与验证
   - 复用 tests/models/kimi_k3/test_attn_res.py，未新增测试文件。
   - 提供微基准、8×B300 端到端 TPOT/TTFT 以及 gsm8k 精度对比。

csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu

唯一变更文件，包含全部四项延迟优化：B/N 模板化、q_cache 向量化、NC=3 低 batch 分派、num_chunks 展开，以及 mbarrier 汇编 memory clobber 配套。

// 内核模板新增 N、B 编译期参数，使 num_chunks 在编译期确定，
// 从而允许主循环完全展开；运行时参数表只保留 T 与各 stride。
template <int H, int N, int NC = N_CHUNK_DEFAULT, int B = 1,
          bool RELEASE_TMEM = false, bool HAS_DELTA = false,
          bool HAS_OUTPUT_NORM = false, bool OUTPUT_NORM_IN_SMEM = false>
__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
    const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
    const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
    const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int T,
    int block_stride_m, int block_stride_r, float rms_eps,
    const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
  constexpr int num_chunks = (N + N_CHUNK - 1) / N_CHUNK; // 编译期常量
​
  // q_cache 填充：int2/int4 向量读入后按 __nv_bfloat162 成对转 float
  // 并累乘，替代逐元素标量加载，减少指令数与访存次数。
  int4 rms_v = *reinterpret_cast<const int4*>(rms_w + h_base);
  int4 res_v = *reinterpret_cast<const int4*>(res_w + h_base);
  auto* rms2 = reinterpret_cast<__nv_bfloat162*>(&rms_v);
  auto* res2 = reinterpret_cast<__nv_bfloat162*>(&res_v);
  #pragma unroll
  for (int k = 0; k < 4; k++) {
    float2 rf = __bfloat1622float2(rms2[k]);
    float2 sf = __bfloat1622float2(res2[k]);
    q_cache[si * VEC + 2 * k] = rf.x * sf.x;
    q_cache[si * VEC + 2 * k + 1] = rf.y * sf.y;
  }
​
  // 主循环全展开；循环展开依赖上面的 constexpr num_chunks。
  #pragma unroll
  for (int ci = 0; ci < num_chunks; ci++, gci++) {
    int ns = ci * N_CHUNK;
    int an = min(N_CHUNK, N - ns);
    // ... 每 chunk 的 mbarrier 同步与计算主体 ...
  }
}

评论区精华

• 三位维护者（gau-nernst、zyongye、ZJY0516）直接 APPROVED，无实质代码讨论。
• claude[bot] 因 fork 来源自动跳过审查，未产生有效建议。
• 值得注意的实现决策是 mbarrier_wait/mbarrier_arrive 内联汇编补 "memory" clobber，这是循环展开后保证 barrier 与访存顺序语义的必要配套。

风险与影响

• 并发正确性：mbarrier 汇编新增 memory clobber，方向是收紧编译器约束，风险低，但属于并发语义改动，值得在后续内核变更中留意。
• 数值精度：gsm8k strict-match 从 0.9636 降至 0.9591（-0.45%），未提供误差分析，可能为数值噪声，但建议在后续精度回归中持续观察。
• 模板膨胀：B/N 编译期化与 switch(num_blocks) 分派会增加实例化数量与二进制体积；极端 num_blocks 走 default 分支时可能回到接近基线的性能。
• 覆盖度：仅一个既有测试，未覆盖 B>1、长 prefill 与低 batch 的组合矩阵。

关联脉络

• 与 #51249（Kimi-Linear packed_modules_mapping 修复）同属 kimi_k3 模型线的连续完善。
• 与 #51149（Interns2mobius 支持）同属共享 MoE 与 kimi 系架构的性能优化潮。
• 与 #50411（设备侧 RMS 归一化前移）一样体现了“把可编译期化的计算前移/常量折叠 + 向量化访存”这类 kernel 性能优化范式在 vLLM 中的常态化应用。