bwd-kernelA-beginner-guide-full
反向 A Kernel 优化全解析（小白版，从零讲起）
日期：2026-07-21
GPU：H20-3e（sm_90a，Hopper）。
本文对象 = 反向传播的 Kernel A（串行 dstate 反向扫描），当前冠军源码
csrc/smxx/bwd_a_dstate_scan_s12.cuh（= 优化日志 docs/20260715-bwd-kernelA-optimization-log.md
里的 s12 = 阶段 11）。姊妹篇是 Kernel B 的 docs/bwd-kernelB-beginner-guide-full.md。
目标读者：不懂 cute、swizzle、寄存器、occupancy、bank conflict、coalescing、Tensor Core 的人。 所有术语从第一性原理讲明白，再讲每一步优化在干什么、为什么有用。
上游材料：
docs/20260715-bwd-kernelA-optimization-log.md（Kernel A 从基线到 s12 的全部 11 个阶段）- Kernel B 的优化史见
docs/20260714/0716/0719-*.md，小白版见bwd-kernelB-beginner-guide-full.md。 - 惯用法源头
csrc/smxx/fwd_kernel2.cuh（前向 K2：Kernel A 的寄存器 MMA / swizzle / STSM 全部照抄它）。
术语扫盲（GPU 硬件层级、三种内存、occupancy、bank conflict、coalescing、向量化访存、cp.async、 Tensor Core / MMA / cute / swizzle / LDSM、warp shuffle、ncu / roofline / stall reason）在 Kernel B 小白版
bwd-kernelB-beginner-guide-full.md第 0 部分已从零讲透。本文默认你已读过那一节， 只在 Kernel A 特有的地方（MOVM_T寄存器转置、split-V、INTER vs SW32、4-way bank 地板）重新展开。
第 0 部分：Kernel A 是什么，在整个反向里处于什么位置
0.1 一句话定义
Kernel A = _flash_kda_bwd_dstate_scan（源码 L166-689），grid (N, H[, VSplit])。
它干一件事：沿时间反向，把"状态梯度" dS[D,D] 一个 chunk 一个 chunk 地往回传（D=128，
一个 chunk=CHUNK=16 个 token）。
打个比方：前向时状态 S 像滚雪球，从第 0 个 chunk 滚到最后一个，每个 chunk 都在它上面叠一层。
反向时要算"最终损失对每一层的敏感度"，就得倒着把这个雪球拆回去 —— 从最后一个 chunk 的
dS 出发，一步步推到第 0 个 chunk。这条倒着推的链就是 Kernel A 的串行扫描
（源码主循环 for (int t = t_tiles - 1; t >= 0; --t)，L288）。
0.2 为什么它是一块独立的 kernel（拆分背景）
反向原本有个巨大的融合 kernel _flash_kda_bwd_recurrence，把两件逻辑上分开的事塞一起
（源码注释 L50-67 白纸黑字）：
- 串行部分：反向 scan，把
dS[D,D]从后往前传（chunk 之间有依赖，天生串行）。 - 大规模并行部分：每个 tile 独立算一堆 per-token 梯度（dv/dbeta/dq/dk/dg…），它们唯一的 跨-chunk 输入就是"进入该 tile 时的 dS"。
拆分方案：Kernel A 只跑串行 scan，并且对每个 tile 把"进入时的 dS"快照存进一块 per-tile
缓冲（fp32，[H*total_tiles, D, D]）；Kernel B 再用一个 CTA/tile 完全并行地读这个快照 +
前向存的中间量，复现所有 per-token 梯度。用"多写一趟 dS 快照到显存"换"把并行部分彻底解耦、
各自打满占用率"。 本文只讲 Kernel A（串行 scan 那半）。
0.3 Kernel A 每个 tile 算什么：5 个 GEMM + 1 个 combine
每个 chunk（tile）里，Kernel A 做 5 个小矩阵乘（GEMM）加一个逐元素"粘合"（源码注释 L81-87）：
| GEMM | 公式 | 形状（TN，fp32 累加） | 干什么 |
|---|---|---|---|
| G1 | dU = Mqk^T @ do |
M=C N=Dv K=C | do 经注意力矩阵回传 |
| G2 | dU += kr @ dS |
M=C N=Dv K=D | dS 在这里进来（唯一读入 dS 的 GEMM） |
| G3 | dvcorr = INV^T @ dU |
M=C N=Dv K=C | 解逆矩阵得 v 修正梯度 |
| G4 | T1 = qd^T @ do |
M=D N=Dv K=C | 组装新 dS 的第一项 |
| G5 | T1 -= kd^T @ (beta·dvcorr) |
M=D N=Dv K=C | 组装新 dS 的第二项 |
| — | dS = exp2(gT)·dS + T1（combine，bf16 存回） |
逐元素 | 更新 dS，传给上一个 chunk |
"TN"是什么：MMA 硬件约定 sC[M,N] = sA[M,K] @ sB[N,K]^T —— A 正着放、B 转置放、沿 K 收缩。
所以表里每个 GEMM 都要指明 M/N/K 三个维度各是谁。
combine 是链的关键（源码 L656-672）：dS_new = exp2(gT)·dS_old + T1。exp2(gT) 是这个 chunk 的
门控衰减（gate decay），gT 是每个 feature 一个值 → 对角缩放（每行 k 乘自己的 exp2(gT[k])）。
算完 dS_new 就是"传给上一个 chunk 的 dS"，写回 smem，下一轮循环（t-1）接着用。
0.4 c 轴旋转 —— 为什么需要"寄存器内转置 MOVM_T"
这是 Kernel A 最"绕"的一点，务必讲清。看 G1/G3：
- G1 产出
dU，形状[c, v]（M=c，行是 token c，列是状态列 v）。 - G3 要用
dU，但它沿 c 收缩（INV^T @ dU，K=c）→ 需要dU的转置[v, c]当 B 操作数。
朴素做法：把 dU 写进 smem，再按转置读回来 —— 一来一回浪费（smem 往返 + 一个 barrier）。
硬件有条专门指令 SM75_U32x1_MOVM_T：直接在 32 个线程的寄存器内部把一个小矩阵块转置，
不碰 smem。于是（源码 L590-599）：dU 累加器（本就在寄存器）→ cast bf16 → MOVM_T 转置 → 直接当 G3 的 B 操作数。G3→G5 之间的 bdv（=−beta·dvcorr）也用第二次 MOVM_T（L624-632）转置留在
寄存器。这就是注释里反复出现的 “c 轴旋转 / MOVM_T 寄存器内转置”，源自前向 K2 的 u_acc 惯用法。
MOVM_T到底是什么（小白版）：ldmatrix/stmatrix那一族里，有个变体能把一个 warp 寄存器里 持有的 8×8 bf16 小块"就地转置"—— 输入是"每线程持有矩阵的某几个元素"的 fragment 布局，输出是 转置后的 fragment 布局，全程在寄存器。省掉"写 smem→同步→按转置读回"整套动作。
0.5 Kernel A 的基线诊断（为什么值得优化）
基线 per-launch µs（bench_kernelA.py，3 个配置，日志 §0）：
| 配置 | 基线 |
|---|---|
| B32/T8192 | 107252 |
| B8/T32768 | 122213 |
| B1/T131072 | 90464 |
基线 ncu（日志 §0）：
- 理论占用率 12.5%，被 shared memory 卡死 → 1 block/SM。
- 动态 smem/block = 124.48 KB（放 2 块要 249 > 227KB 上限，放不下）。
- 延迟受限：每 tile 串行 5 个 cooperative_gemm + ~8 个
__syncthreads，warp 太少 （Active Warps/Sched 2.01）掩盖不了 barrier/GEMM 延迟。 - 首要 stall = MIO throttle 31.4%；DRAM 仅 ~4%（非带宽受限）。
- 标注的最大杠杆：把 smem 砍到 ~113KB 以下放下 2 block/SM。“大象” =
T1_smem（D×Dv fp32，VSplit=1 时 64KB）。
端到端占比：Kernel A ≈ 总 bwd 时间的 30%（日志 §4 明确纠正了早先"1-2%"的错估）。所以把它 砍 50%+ ≈ 端到端 −15%，值得追。
第 1 部分：第一场硬仗 —— 消灭 T1_smem，占用率 12.5%→50%（阶段 1-2）
整个 Kernel A 优化史分两大战役：先靠"寄存器融合"打赢占用率（阶段 1-2，本部分），再靠 "swizzle/padding"打赢 bank conflict（阶段 3-11，第 2 部分）。这个顺序不是随意的 —— 先修占用率 （让 SM 上住满 warp），bank conflict 才会浮出水面成为新瓶颈；反过来先修 conflict，会被低占用率 的延迟盖住看不出收益。
1.1 病根：fp32 中间缓冲吃光了 smem
基线的 5 个 GEMM 是"cooperative_gemm"（cute 提供的、整个 block 协作的 GEMM）。它们把中间结果
dU、dvcorr、T1 全都写进 fp32 的 smem 缓冲再读回来。其中最大的一块是
T1_smem（D×Dv fp32 = VSplit=1 时 128×128×4 = 64KB）—— 日志称之为**“大象”**。
124.48KB 的 smem 里，光这头大象就占一半，直接把"能不能塞下第 2 个 block"给否了。
1.2 先讲透两件事：大象是怎么来的、又是怎么被干掉的（零 cute 术语版）
小白最容易卡在这两句话：“干掉大象 T1_smem” 和 “每个 warp 无内部同步、无跨 warp 共享”。 下面用完全不带 cute 的大白话，从"一群人砌墙"的比喻讲清楚。
(a) 一个矩阵乘有两种分工方式
Kernel A 每个 tile 要算 5 个矩阵乘。一个矩阵乘 = 算一块输出矩阵 C = A×B。 一个 block 有 256 个线程（8 个 warp，每 warp 32 线程）。这 256 人算一块 C，有两种分工：
-
方式一：全员合砌一面墙（旧代码，“cooperative_gemm”）。 256 个人一起砌同一面大墙（同一块输出 C）。既然大家都往同一面墙上贴砖，这面墙就必须放在 所有人都够得着的公共场地 = shared memory（smem）。T1 这块输出是
[128, 128]的 fp32 矩阵 = 128×128×4 字节 = 64KB。这 64KB 就是那头大象T1_smem—— 它之所以在 smem， 纯粹是因为"全员合砌一面墙，墙必须放公共场地"。 -
方式二：把墙切成 8 段，一人一段（新代码，“每 warp 一个 MMA”）。 把这面 128 列宽的大墙，按列切成 8 段、每段 16 列，8 个 warp 一人包一段：
warp0 只砌第 0–15 列，warp1 只砌第 16–31 列……各砌各的、互不重叠。 既然每个 warp 只砌自己 那一小段（16 列 × 128 行的一条），这一小段的中间结果小到能全部塞进这 32 个线程自己的口袋 （寄存器）里，根本不用放到公共场地 smem。→ 大象没了。
【图 1】旧"全员合砌一面墙" vs 新"一人一段"
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25┌─────────── 旧代码 (cooperative_gemm) ───────────┐
│ │
│ 256 个线程 ──全都往同一面墙贴砖──▶ │
│ │
│ ┌────────────────────────────┐ │
│ │ T1_smem [128 × 128] │ │
│ │ fp32 = 64KB │ ◀── 🐘 大象│
│ │ （必须放公共场地 smem， │ 住在 smem│
│ │ 因为 256 人共写一块） │ │
│ └────────────────────────────┘ │
└──────────────────────────────────────────────────┘
┌─────────── 新代码 (每 warp 一个 MMA) ────────────┐
│ T1 / dS (128 行 k × 128 列 v) 按列切 8 段: │
│ │
│ 列: 0–15 |16–31|32–47|48–63|...|96–111|112–127 │
│ ┌────┬────┬────┬────┬───┬─────┬──────┐ │
│ 行 │ w0 │ w1 │ w2 │ w3 │...│ w6 │ w7 │ │
│ k │ 寄 │ 寄 │ 寄 │ 寄 │ │ 寄 │ 寄 │ │
│ 0–127 │ 存 │ 存 │ 存 │ 存 │...│ 存 │ 存 │ │
│ │ 器 │ 器 │ 器 │ 器 │ │ 器 │ 器 │ │
│ └────┴────┴────┴────┴───┴─────┴──────┘ │
│ 每 warp 只管自己 16 列 → 塞进自己寄存器 → │
│ smem 里一头大象都没有 ✔ │
└──────────────────────────────────────────────────┘
(b) 为什么"寄存器"就等于"省 smem、升占用率"
-
smem 是每个 SM 一池固定的公共场地（~228KB），住在这个 SM 上的所有 block 共用。 一个 block 占的 smem 越大，这个 SM 能同时住的 block 就越少 → 活跃 warp 越少 → 一个 warp 卡住时没别人顶上 → SM 空转。
-
寄存器是每个线程自己的口袋，不占 smem 那池公共场地。只要口袋没爆（本例 64 寄存器/线程，0 spill）， 把中间结果放寄存器就等于凭空省下那 64KB smem。
-
账一算：smem/block 从 124KB 掉到 50KB → 一个 SM 能住的 block 从 1 个变 4 个（
228÷50≈4）→ 活跃 warp ×4 → 占用率 12.5%→50%。这就是 −37% 的来源:不是算得更快,是 SM 不空转了。【图 3】一个 SM 的 smem 池（~228KB）能住几个 block
1
2
3
4
5
6
7
8
9
10
11
12
13
14旧: 每 block 吃 124KB 新: 每 block 吃 50KB
┌──────── SM smem 池 228KB ────┐ ┌──────── SM smem 池 228KB ────┐
│ ┌──────────────────────────┐ │ │ ┌────────┐ ┌────────┐ │
│ │ block #1 124KB │ │ │ │block#1 │ │block#2 │ │
│ │ │ │ │ │ 50KB │ │ 50KB │ │
│ └──────────────────────────┘ │ │ └────────┘ └────────┘ │
│ ……剩 104KB 塞不下第 2 个…… │ │ ┌────────┐ ┌────────┐ │
│ （浪费） │ │ │block#3 │ │block#4 │ │
│ │ │ │ 50KB │ │ 50KB │ │
└──────────────────────────────┘ │ └────────┘ └────────┘ │
1 block/SM 占用率 12.5% └──────────────────────────────┘
↑ 这 1 个 block 的 warp 一卡 4 block/SM 占用率 50%
住,整个 SM 就空转 ↑ 一个 block 卡住,还有 3 个
block 的 warp 顶上,SM 不空转
© 为什么"切成 8 段"顺便消掉了几乎所有 __syncthreads
先说 __syncthreads 是什么：它是一道"全体 256 线程都到齐才能继续"的栅栏。什么时候必须用它?
当线程 A 往 smem 写了个东西、线程 B 要读它时 —— B 必须等 A 写完,否则读到旧值。栅栏很贵:
快的线程要干等慢的线程到齐。
-
旧代码(全员合砌一面墙)为什么栅栏多:5 个矩阵乘是接力的(G1 的产出 dU 喂给 G2,G2 喂 G3……)。 每次接力都要经过公共场地 smem:G1 全员把 dU 写进 smem → 一道栅栏(等所有人写完)→ G2 全员从 smem 读 dU。每一次"经 smem 交接"就是一道栅栏,一条链下来 ~8 道。
-
新代码(一人一段)为什么不用栅栏:warp0 从头到尾(G1→G2→G3→G4→G5→combine)只在自己那 16 列上 折腾,全程在自己口袋(寄存器)里接力。它读的旧 dS 是自己那 16 列、写的新 dS 也是自己那 16 列, 从不去读 warp1、warp2 产出的任何东西。既然 warp 之间没有"你写我读"的交接,自然谁也不用等谁 → 中间一道栅栏都不需要。这就是"无跨 warp 共享 → 无内部
__syncthreads"。
(c-2) 关键:为什么偏偏能"按列切"?——被切的那根轴不能被求和
有人会问:那我按行切成 8 段行不行?不行。 但准确的原因不是"列可以、行不可以"这条通用规律, 而是这条特定的链里,v 列轴从头到尾都没被"求和"掉,而 k 行轴在 G2 被求和掉了。先分清两种轴:
- 自由轴(free):直接落到输出上、不参与求和。比如
C[c,v] = A[c,k] @ B[k,v]里的c和v。 - 收缩轴(contract):被
Σ求和掉、不出现在输出里。上面的k。
规律:切一根自由轴 → 各段天然独立、最后拼接即可;切一根收缩轴 → 每段都要读别人的数据来求和 →
必须跨 warp 共享 + 栅栏。 dS/T1 的形状是 [D_k 行, D_v 列],我们切的是 v 列。逐 GEMM 看两根轴的角色:
| 步骤 | 运算 | v 轴(列) | k 轴(行) |
|---|---|---|---|
| G1 | dU = Mqk^T @ do |
自由(do 的 v 直落) | — |
| G2 | dU += kr @ dS |
自由 | 收缩(dS 的 k 被 Σ 掉) |
| G3 | dvcorr = INV^T @ dU |
自由 | — |
| G4 | T1 = qd^T @ do |
自由 | 自由(输出行) |
| G5 | T1 -= kd^T @ bdv |
自由 | 自由(输出行) |
| combine | dS = exp2(gT)·dS + T1 |
自由(逐元素) | 自由(逐元素) |
看最右那列 v:从 G1 到 combine 全是"自由",一次都没被求和。 所以 warp0 只要拿 do、dS、bdv 的
第 0–15 列,就能独立算出输出的第 0–15 列,谁也不用问 → 列切成立、最后拼接。而 k 行轴在 G2 被整根
求和(dU[c,v] = Σ_k kr[c,k]·dS[k,v]):若按行切,warp0 只有 dS 的 0–15 行,算 G2 却要 dS 全 128 行
求和 → 必须去读别的 warp 的行 → 跨 warp 共享 + 落 smem + 加栅栏,独立性就没了。
【图 4】按列切(v,成立) vs 按行切(k,G2 就崩)
1 | ┌──────── 按列切 v (新代码,成立) ────────┐ |
一句话:是"被切的轴在链里是否被收缩"决定的,不是"列/行"这个名字。 若链里哪一步出现了
Σ_v(对 v 求和),那即使按列切也照样要跨 warp 归约。本链恰好 v 全程自由、k 在 G2 收缩 → 只能沿 v(列)切。
【图 5】整条链每个矩阵在哪里(smem 全宽共享 / warp 寄存器切分)
先说清三根轴是什么(否则 c/k/v 光看字母看不懂)——这是一个 chunk 内的反向计算,涉及三根轴:
1 | c = chunk 内的 token 数 = CHUNK = 16。一个 chunk 里有 16 个时间步(token),c 就数这 16 个。 |
再约定三种"位置标签"(看懂链路的关键):
1 | 【S 】smem·全宽·共享 : 存在 shared memory, 宽度满 128, 8 个 warp 共读同一份, 不切分。 |
贴着标签走一遍整条链(warp_w 视角, S = 它的列区间 [16w,16w+16), c=16, D_k=D_v=128):
1 | 步骤 输出 = A(左, 全在 smem) @ B(带 v, 按列切) 收缩轴(被Σ的轴) |
【图 6】决定性的 G2:dU = kr @ dS(收缩 k)——标注每块在哪
1 | ● 按列切 v(新代码, 成立): 切的是"自由轴"v |
-
那结尾那一道栅栏(源码 L676)干嘛的:8 个 warp 各砌完自己那段后,下一个 tile 开头的"快照"步骤 (源码 L377-393)会把整块 dS(全 128 列)一次性读出来存到显存 —— 这一步是 256 线程不分 warp 地扫全表, 会读到"别的 warp 刚写的列"。所以必须在这里同步一次,保证 8 个 warp 的写全部落地、对所有人可见, 下一个 tile 才能安全地读整块 dS。整条链只剩这一道结尾栅栏。
【图 2】旧代码 ~8 道栅栏 vs 新代码 1 道
1
2
3
4
5
6
7
8
9
10
11
12
13旧 (全员合砌,接力都过 smem): 新 (一人一段,全程在寄存器):
──────────────────────────── ────────────────────────────
G1 全员写 dU→smem warp0: G1→G2→G3→G4→G5→combine ┐
══ 栅栏1(等所有人写完)══ (只碰自己那 16 列,全在寄存器) │ 8 个 warp
G2 全员读 dU / 写 dvcorr warp1: G1→G2→G3→G4→G5→combine │ 并行跑,
══ 栅栏2 ══ (只碰自己那 16 列) │ 互不读写
G3 ... ... │ 对方 →
══ 栅栏3 ══ warp7: G1→G2→G3→G4→G5→combine ┘ 零栅栏
G4 ...
══ 栅栏4 ══ ─────────────────────────────
G5 ... ══ 唯一栅栏(L676) ══
══ ... 一条链 ~8 道 ══ ↑ 发布整块 dS,给下一 tile
combine ... 的"快照"(扫全 128 列)读
一句话:大象 = “全员合砌一面墙,墙必须放 smem 公共场地”。把墙按列切成 8 段一人一段,同时办成 三件事 ——(1)每段小到能塞进寄存器 → 大象没了、省 64KB、占用率翻 4 倍;(2)warp 之间不共享数据 → 中间零栅栏;(3)结尾只留一道栅栏发布整块 dS 给下一 tile。
1.3 具体怎么搬(照抄前向 K2 的惯用法)
关键洞察（日志 §0 策略）：前向 K2 根本没有 D×D 的 fp32 累加器 —— 它的状态更新全程是 “寄存器 MMA + STSM 写 swizzled bf16 smem”。把这套惯用法搬到 Kernel A，就能把 dU/dvcorr/bdv/T1 全部留在寄存器，fp32 中间缓冲一个不剩。
阶段 1（日志 §1）：先把 C×Dv 链（G1→G2→dU→G3→dvcorr→bdv）搬进寄存器。做法：
- 用单 warp MMA 惯用法（
MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>+Tile<16,16,16>）， 每个 warp 拥有一个 16 宽的 v-blockvb = warp_id（源码 L516-520 现在的形态）。 dU在寄存器累加（G1Mqk^T@do，再 G2 循环kr@dS），cast bf16，MOVM_T转置成 B 操作数， G3INV^T@dU在寄存器，beta·dvcorr也在寄存器算完。- 消除：
dU_smem/dvcorr_smem（fp32）、dU_bf暂存循环、3 个__syncthreads。 - 本阶段 G4/G5 仍用 cooperative_gemm（大象 T1 留到阶段 2 处理）。
- 结果：−6.7% / −6.9% / −11.9%。gradcheck L0 PASS，cos 0.99993，0 spill。
阶段 2（日志 §2，大赢）：把寄存器链延伸到 D×Dv 链（G4/G5/combine），干掉大象。看现在的源码 L634-673 —— 每个 m-block 里：
- G4
T1 = qd^T @ do、G5T1 -= kd^T @ bdv都是寄存器 MMA（T1_acc是 MMA 累加器，在寄存器）。 - combine
dS = exp2(gT)·dS_old + T1在寄存器完成：LDSM 读旧 dS 入 fragment（L658），与寄存器T1_acc融合（L661-669，gk0/gk1是gT按 M-row 分的对角缩放），STSM 写回（L672）。T1 全程不落 smem。 - G5 的减号靠预取负 beta 实现：
b = -beta（源码 L614-615-beta_smem[...]），让寄存器gemm（C += A@B）自然得到T1 = qd@do − kd@(beta·dvcorr)。
为什么每个 warp 能无内部同步地跑完整条链：见上面 §1.2© 的"一人一段"解释 —— 每个 warp 独占
它那 16 宽的 v 列，G2 只读 dS[vb,·]、combine 只写 dS[vb,·]，warp 之间不共享数据 → 中间零
__syncthreads，只在循环末尾一道 barrier（源码 L676）把整块 dS 发布给下一个 tile。
- 消除的 smem：
dU_smem/dvcorr_smem/T1_smem（fp32）+dU_bf/bdv_bf（bf16），全没了。 - 新 smem =
dS[D·Dv]bf16 +(D+CHUNK+…)fp32 + 几块 bf16 MMA 操作数。启动器KDA_LAUNCH_K2A的smem_公式相应删掉 fp32 里的CHUNK·Dv + D·Dv两项。
1.4 阶段 2 的战果：占用率翻两番
| 指标 | 基线 | 阶段 2 |
|---|---|---|
| 动态 smem/block | 124.48 KB | 50.75 KB |
| Block Limit Shared Mem | 1 | 4 |
| 理论占用率 | 12.5% | 50% |
| 实测占用率 | ~16% | 32.3% |
| Active Warps/Sched | 2.01 | 5.17 |
结果：66301 / 76413 / 57098 = 相对基线 −38% / −37.5% / −37%，三配置统一下降（不只 B1）， 与占用率跃升吻合。gradcheck L0 PASS，cos 0.99993。所有 VSplit 变体 64 regs、0 spill。
但瓶颈转移了：首要 stall 仍是 MIO throttle（现 16.3 cyc = 51%），L1/TEX 吞吐 96.69% 钉死。 瓶颈从"占用率受限"转成"shared-memory-pipe 受限" —— 也就是 bank conflict 登场，引出第 2 部分。
小白记住：阶段 1-2 的招 = “把 fp32 中间缓冲从 smem 挪进寄存器”。收益不是来自算得更快， 而是腾出 smem → 一个 SM 从住 1 个 block 变住 4 个 → 延迟被藏住。这是典型的"占用率优化"。 前提是这些中间量本来就只在 warp 内用、不跨 warp 共享，才能安全地纯寄存器化。
第 2 部分：第二场硬仗 —— swizzle 消灭 bank conflict（阶段 3a-3d）
占用率修好后，L1/TEX 吞吐钉在 96.69% = smem 访存管线打满。病根是 bank conflict：MMA 要按
"fragment 布局"从 smem 读小矩阵块，但 smem 若朴素行主序摆放，这种读法会撞 bank，硬件排队。
本部分的主线：用 swizzle 布局让每个操作数的 ldmatrix 读无冲突。
2.1 先失败一次：普通布局的 LDSM（阶段 3a）
阶段 2 的 ncu 假设：转置视图上的通用标量 copy 无法向量化 → 每个 16×16 tile 约 8 次窄 LDS →
MIO throttle。ncu 直接建议用 ldmatrix（更少更宽的载入）。于是阶段 3a 把全部 6 个操作数载入
换成 ldmatrix tiled-copy（源码现在 L539-552 那套 make_tiled_copy_A/B/C 机器）。
证明了一件事：LDSM_T 在普通、非 swizzled 的 MN-major 视图上就能用，gradcheck bit 一致 ——
正确性不需要 swizzle。但结果只有 −1.5% / −1.2% / −2.6%，不理想。
ncu 揭示原因（日志 §3）：Warp Cycles/Issued 从 32 涨到 46（更差！），两个首要 stall
Short Scoreboard 40.3% + MIO Throttle 32.1% = 72% 全是 smem-pipe。结论：kernel 不是受指令
数量限制，而是 shared-memory bank conflict 受限。 普通布局的 ldmatrix 砍了指令数，但 bank
冲突（例如 sQd 各行相隔 256B → 别名到同一组 bank → ldmatrix.trans 串行化）让 L1/TEX 仍钉死 97.5%。
👶 逐句拆开这段 ncu 诊断（小白版）
① Warp Cycles/Issued 32→46(更差)。 这个指标 = 平均每发射一条指令、一个 warp 要耗多少周期,
越大越堵(warp 越多时间在干等而非干活)。换 ldmatrix 后指令总数是砍了,但每条更慢 → 净负 →
直接推翻"减少指令数就会快"的假设。
② 两个首要 stall = 72% 全在 smem 管线。 ncu 会告诉你 warp 卡住时在等什么:
- Short Scoreboard 40.3% = 在等一个短延迟片上访问的结果回来 —— 几乎总是 smem 读(scoreboard 记分板追踪"数据到了没",卡这儿 = 在等 smem 取数)。
- MIO Throttle 32.1% = MIO 是所有 smem 指令(LDS/LDSM/STS)走的管线,Throttle = 管线已排满,新的 smem 指令发不进去、只能等。
- 两者合 72% 全是 smem 相关 → 病根是 smem bank conflict,不是指令条数、不是算力。所以砍指令治不了病。
③ 机制:“sQd 各行相隔 256B → 别名到同一组 bank → ldmatrix.trans 串行化”。 拆三段:
- 各行相隔 256B:
sQd是普通 [D,C] bf16,一行 128 个 bf16 = 128×2 = 256 字节 → 相邻行起点每隔 256B。 - 别名到同一组 bank:bank 一圈 = 32 bank × 4 字节 = 128 字节。256B = 正好 2 整圈(256/128=2) → 隔整数圈 = 回到完全相同的 bank。于是不同行的对应元素 bank 号一模一样 = “别名(alias)”:不同行却映射 到同一组 bank。ldmatrix 要跨多行取一个 8×8 小块,这些行 bank 全重合 → 挤同几个 bank → 冲突。
- ldmatrix.trans 串行化:带转置的 ldmatrix 本该一条指令并行搬一个小块进寄存器,但 bank 全重合 → 硬件无法一拍完成 → 只能一组 bank 一组 bank 地串行搬 = 串行化,一条并行指令被拉成好几拍。
④ “L1/TEX 仍钉死 97.5%”。 L1/TEX 是含 smem/MIO 的存取单元,97.5% = 它几乎全程在忙 —— 但不是在高效 搬数据,而是在反复串行处理 bank 冲突(忙于重试,不是干活)。所以换了 ldmatrix、砍了指令数,L1/TEX 因 bank 冲突照样钉死 97.5%、warp 全卡等 smem → 净负。这正是逼出下一步(s3b/3c 上 swizzle)的直接证据。
教训：
ldmatrix本身不解决 bank conflict —— 它只是把"很多窄 LDS"合成"一条宽 ldmatrix"。 如果那条宽 ldmatrix 读的 smem 布局仍撞 bank，照样串行化。必须配 swizzle 才有用。
2.2 分诊：只 swizzle "R1 产出"的操作数（阶段 3b-3c）
关键分诊（日志 §4）：不是所有操作数都能 swizzle。
- R1 产出的操作数（qd/kd/kr）：由 R1 逐元素写进 smem → 写地址可自由重排 → swizzle 友好。
- cp.async 载入的操作数（do/mqk/inv）：向量化 16B 载入会被 swizzle 打碎，且每 tile 只用一次 （1 次 LDSM → 冲突代价可忽略）→ 保持普通布局。
👶 "向量化 16B 载入会被 swizzle 打碎"是什么意思
先看两拨操作数进 smem 的方式完全不同:
qd/kd/kr是 R1 在 kernel 里逐个元素算出来、逐个写进 smem 的 → 写到哪个地址完全自由。do/mqk/inv是用cp.async从显存整块搬进 smem 的 → 一次搬 16 字节 = 连续 8 个 bf16(硬件 向量化要求,搬运的最小单位就是"连续 16B",不可拆)。
swizzle 要做的事 = 打乱写入地址,让相邻元素散到不同 bank(即要求 8 个元素别再连续摆):
1 | cp.async 只会这样连续落地: [a0 a1 a2 a3 a4 a5 a6 a7] ← 一条 16B,必须连续 |
这两个要求直接打架:要把 8 个元素散开,就得把这一条 16B 拆成 8 次零散写 → cp.async 的"一次搬 16B"
向量化就没了,退化成慢速逐元素搬运。"被打碎"就是这个意思:swizzle 的散地址会把 16B 向量化拆成 8 次
小搬运,毁掉 cp.async 的性能。 而 qd/kd/kr 本就是逐元素写、没有"必须连续 16B"的包袱 → swizzle 随便加。
再补一刀"只用一次":就算不 swizzle、do/mqk/inv 保持普通布局会有点 bank 冲突,但它们每个 tile 只被
LDSM 读一次(用完即弃) → 一次读的冲突代价小到可忽略;而 qd/kd/kr 在 5 个 GEMM 里被反复读,冲突
被放大 → 值得 swizzle。结论:冲突代价小(只读 1 次)+ swizzle 代价大(打碎 cp.async)→ do/mqk/inv 保持
普通布局最划算。 这也正是后来 s7(硬去 swizzle 它们)回归、被回退的原因。
什么是 R1：Kernel A 不直接读前向存的 qd/kd/kr（那是 bf16、有精度损失），而是在片上用 fp32
重新算一遍（reconstruct）。源码 L407-493 就是 R1：从 raw q/k + gc + gT 重建 qd/kd/kr。因为是
逐元素现算，写到哪个 smem 地址完全自由 → 可以直接写进 swizzled 布局。
阶段 3b — swizzle qd/kd（源码 L118-120 QdKdSw = Layout_MN_INTER_Atom over [D,C]）：
R1（生产者）写 sQd(col, r_row)，G4/G5（消费者）读同一个布局对象 → 由构造一致（无论 swizzle
怎么置换，写读同一对象就对得上，正确性与置换无关）。qd/kd 在 G4/G5 各用 8 次（8 个 m-block）→
冲突代价大。
| 配置 | 阶段 3a | 阶段 3b | Δ vs 3a | vs 基线 |
|---|---|---|---|---|
| B32/T8192 | 65338 | 57531 | −11.9% | −46.4% |
| B8/T32768 | 75484 | 66173 | −12.3% | −45.9% |
| B1/T131072 | 55634 | 55595 | −0.1% | −38.5% |
阶段 3c — swizzle kr（源码 L124-126 KrSw = Layout_K_INTER_Atom over [C,D]）：
kr 是 A 操作数、K-major、G2 里用 8 次 → LDSM_N。同样生产者/消费者共享布局对象。
| 配置 | 阶段 3b | 阶段 3c | Δ vs 3b | vs 基线 |
|---|---|---|---|---|
| B32/T8192 | 57531 | 53614 | −6.8% | −50.0% |
| B8/T32768 | 66173 | 61687 | −6.8% | −49.5% |
| B1/T131072 | 55595 | 54834 | −1.4% | −39.4% |
两个"8×循环"操作数（qd/kd、kr）承载了大部分冲突代价。注意 B1 几乎不动（−0.1%/−1.4%）—— 因为 B1 是 N=1 占用率饥饿（grid=128，1 block/SM），瓶颈是并行度不是冲突，swizzle 帮不上它。 这个"B32/B8 赢、B1 持平"的模式会在后面每个 conflict 阶段重复出现。
2.3 最后一个操作数：dS 的转置对偶 swizzle（阶段 3d）
qd/kd/kr 都 swizzle 后，ncu 显示冲突仍高（op_st.sum 617M → 91% store wavefront 是纯冲突）。
唯一剩下的普通操作数 —— dS（最大的 [128,128]，每 tile 用 8 次）—— 是主导冲突源。
dS 被两个正好互为转置的方向用（源码注释 L128-139）：
- G2 读它当 B 操作数
[v,k]（源码 L529DSb）； - combine 读/写它当 C 累加器
[k,v]（源码 L530DSc）。
这跟前向 K2 的 State/TransposedState 完全同构。所以照抄它的**转置对偶（transpose-dual）**招：
DSb=Layout_K_INTER_Atomover [Dv,D]（B 操作数，LDSM_N，源码 L140-142）DSc=Layout_MN_INTER_Atomover [D,Dv]（C 累加器，LDSM_T读 /STSM_T写，源码 L143-145）
妙处：K_INTER 和 MN_INTER 两个 atom 互为转置，tile_to_shape 一致铺贴 → DSb[v,k] 和
DSc[k,v] 命中同一个物理字节（源码注释 L134-136 说明：方形 VS1 和矩形 VS>1 都成立）。于是 dS
只存一份物理数据，两个方向都能无冲突读/写。全部 5 个 dS 站点（init/snapshot/G2-B/combine-C/
ds_out）都走这两个视图，一致性自动保证。cosize == D·Dv（INTER atom 无 padding）→ smem 不变。
| 配置 | 阶段 3c (s4) | 阶段 3d (s5) | Δ vs 3c | vs 基线 |
|---|---|---|---|---|
| B32/T8192 | 53614 | 41366 | −22.8% | −61.4% |
| B8/T32768 | 61687 | 50656 | −17.9% | −58.5% |
| B1/T131072 | 54834 | 54277 | −1.0% | −40.0% |
dS swizzle 是 conflict 战役里最大的单项赢（−18~23%）—— 因为 dS 是最大操作数、每 tile 复用 8 次。 Kernel A 累计相对基线 = −61%/−58%/−40%。
小白记住 transpose-dual：一个矩阵在两个 GEMM 里被"正着乘"和"转置着乘"。用 INTER 类 swizzle， 同一份字节 既能 K_INTER 正读、又能 MN_INTER 转置读 → 存一次，喂两个方向，零转置、零 padding。 这是 cute swizzle 最优雅的用法，Kernel A（dS）和 Kernel B（sin/dS）都靠它。
第 3 部分：写侧 + gmem 延迟（阶段 4/s6、6/s8、7/s9）
阶段 3d 把读侧 bank conflict 基本清干净了，但还有两个战场：R1 写 swizzled smem 的 store 冲突， 以及 gmem 载入延迟。
3.1 阶段 4/s6 —— 用 STSM 写 R1 的产出，消灭 380M store 冲突
3d 之后载入已 swizzled，但 qd/kd/kr 写入其 swizzled smem 的 R1 生产者 store 仍是通用标量
copy(...) = 16-way STS bank 冲突（ncu op_st 仍高）。修法 = 用 SM90_U32x4_STSM_N/SM90_U16x8_STSM_T
（前向 K2 store 惯用法）写这些 store，让生产者写入与 swizzled 消费者布局一致 → 成为真正的 stmatrix。
看源码 L485-491：R1 phase-2 把 qd/kd 用 STSM_T、kr 用 STSM_N 写进各自 swizzled 目标。正确性由
构造保证（store 和 load 同一布局对象）。
| 配置 | 阶段 3d (s5) | 阶段 4 (s6) | Δ vs s5 |
|---|---|---|---|
| B32/T8192 | 41366 | 35403 | −14.4% |
| B8/T32768 | 50656 | 44868 | −11.4% |
| B1/T131072 | 54277 | 51306 | −5.5% |
读写对称：swizzle 只做一半（读）等于没做 —— 写侧仍撞 bank 会拖后腿。
ldmatrix(读)/stmatrix(写) 要成对用同一 swizzled 布局，才彻底 conflict-free。
👶 逐步拆开讲（小白版）
一块 swizzled smem 有"两端":写进去的生产者、读出来的消费者。
- 消费者端(读):MMA 用
ldmatrix(LDSM)读它 —— s3~s5(swizzle 那几步)已经把这端配好,读 0 冲突。 - 生产者端(写):R1 算出
qd/kd/kr后,得先写进这块 swizzled smem —— 这端 s6 之前还没改。
病根:读端修好了,写端还在用"通用标量 copy"往 swizzled 地盘上写。
R1 phase-2 算完 qd/kd,原先用普通逐元素 copy(...)(一线程写几个元素)store 进 smem。但目标 smem 是
swizzled 布局,而普通 STS(scalar store) 对地址的摆法跟"swizzled 消费者读时的摆法"对不上 →
32 个 lane 的写地址挤进同一/少数几个 bank → 16-way STS 冲突(ncu op_st/store-conflict 一直高)。
即:读端(LDSM)已无冲突,写端(STS)还在拖后腿。
修法:把"写"也换成 stmatrix(STSM),让写和读用同一个 layout 对象。
STSM = stmatrix,是 ldmatrix 的孪生写指令。关键规则:swizzle 只有在"写和读用同一套 swizzle 变换"
时才无冲突 —— 读端 ldmatrix 带 swizzle,写端就必须用配套的 stmatrix(同一 swizzle)。所以 s6 把 R1
生产者 store 从"通用 copy"换成前向 K2 早就在用的 STSM 惯用法(源码 L485-491):
qd、kd用SM90_U16x8_STSM_T(带转置_T—— 因为 qd/kd 后面要以转置形态喂 MMA);kr用SM90_U32x4_STSM_N(不转置_N)。
这样生产者写入用的 layout 对象和消费者读取用的 layout 对象是同一个 swizzled layout(同一个
CuTe layout 变量):写按它铺、读按它取 → bank 完美错开 → STS 冲突消失,qd/kd/kr 从"半吊子 swizzle
(读好写差)"升级成真·stmatrix 通路(读写两端都 swizzle)。
为什么"正确性由构造保证"(免验证)。 swizzle 本质是打乱物理地址,但只要 store 和 load 引用同一个 layout 对象,"store 把元素 X 按打乱规则塞到某物理 bank"和"load 按同一规则去那个 bank 取 X"就是一个 置换 + 它的逆 —— 绕一圈,load 必然拿回原来的 X。swizzle 是对上层透明的地址重排,写读同源则数据语义 恒等。所以 s6 不改任何数值、只改物理摆放 → 无需重新验证正确性。
和 s11 的分工别搞混:s11 修的是 R1 输入 buffer q_sm/k_sm/gc_sm 的读冲突(pad stride);
s6 修的是 R1 输出 buffer qd/kd/kr 的写冲突(STSM)。R1 一头(读输入)一尾(写输出),两处不同战场。
3.2 阶段 5/s7 —— ❌ 生产者 pass swizzle mqk/inv/do：回归，已回退
s6 后想顺手也 swizzle 三个 cp.async 载入的操作数（mqk/inv/do）。plan (b) 试法：保持 row-major cp.async 到 raw buffer，再加一个"每 tile 一次的生产者 pass"把 raw→swizzled 孪生 buffer 复制，消费者 切到 swizzled buffer。由构造 bit 精确。
结果：正确但更慢 —— +3.1% / +3.2% / +14.5%，已回退。
为何输（日志 §7）：s7 时绑定 stall 仍是 gmem long_scoreboard(2.39) —— 805M 的 shared-load 冲突
数量虽大，但藏在 gmem 延迟后面（不在关键路径）。而生产者 pass 增加了一次全 block 复制 +
1 个 __syncthreads/tile → 净损失；B1（占用率饥饿、延迟受限）被额外 sync 打击最狠（+14.5%）。
第一条"红鲱鱼"：ncu 报 805M shared-load 冲突，看着吓人，但它不在关键路径（被 gmem 延迟盖住）。 追一个不在关键路径上的指标 = 白费力气还倒退。 必须先看"绑定 stall 是谁"。
3.3 阶段 6/s8 —— cp.async 预取 R1 的 q/k/gc，碾碎 gmem long_scoreboard
s7 失败暴露了真瓶颈 = gmem 延迟（long_scoreboard 2.39）。s8 转而攻它：把 R1 的输入
（raw q/k + ws_gc）用 cp.async 异步预取进 smem，让 gmem 载入和上一 tile 的 GEMM 链重叠。
看源码 L303-328 的两组 cp.async设计（很讲究）：
- Group A（q/k/gc）：R1 消费 → 先 commit（L328），R1 phase-1 前
wait<1>（L395）→ 它的延迟只 藏在"快照存储"后面。 - Group B（inv/mqk/do）：GEMM 链消费 → 后 commit（L357），R1 之后
wait<0>（L496）→ 它的延迟 藏在"快照 + 整个 R1"后面。 - FIFO 顺序：A 先提交 ⇒
wait<1>排空 A（留 B≤1 未完成），wait<0>稍后排空 B。
套路精髓：先发一堆异步搬运 + commit，不等、马上干别的（存快照 L377-393、跑 R1），等真要用
时才 wait → gmem 搬运时间被计算藏住。
同时修了 varlen 的 12B 寄存器 spill：把每 CTA 一次的 varlen tile_base 累加（int64）抽成
__noinline__ 的 kda_varlen_tile_base（源码 L40-46），让这段冷 prologue 有独立寄存器帧、不和热
GEMM body 抢寄存器 → 全 6 个变体 0-spill。
| 配置 | 阶段 4 (s6) | 阶段 6 (s8) | Δ vs s6 |
|---|---|---|---|
| B32/T8192 | 35403 | 35058 | −1.0% |
| B8/T32768 | 44868 | 43173 | −3.8% |
| B1/T131072 | 51306 | 45911 | −10.5% |
ncu：前提翻转（日志 §8）：long_scoreboard 塌缩 2.39→0.39（B32）。gmem 延迟消失后，B32/B8 上 新的首要 stall 变成 short_scoreboard（等 smem LDS/LDSM）= 由仍未 swizzle 的 mqk/inv/do 喂养 —— 这翻转了 s7 的结论：现在 swizzle mqk/inv/do 瞄准的是真正的关键路径了。
3.4 阶段 7/s9 —— swizzled cp.async scatter mqk/inv/do（plan a）：WIN
s8 让 mqk/inv/do 的 swizzle 值得做了 —— 但只通过 plan (a)（无生产者 pass）：把 MN_INTER swizzle
直接折进 cp.async LOAD，把每个 16B chunk scatter 到它的 swizzled smem 槽，消费者 LDSM_T 直接
触发真正的 ldmatrix.trans。看源码 L330-357：
- 加了
MqkSw/DoSwswizzled 布局（源码 L158-163）； - LOAD 用
copy16(&sMqk_sw(c,j), &mqk_tile[i])做 scatter（L340-341），do 尾部在 swizzled 位置清零 （L362-369）。 - 无生产者 pass、无额外
__syncthreads、无 smem 变化（INTER atom cosize == 普通大小）—— 正是拖垮 s7 plan (b) 的那些开销全都不存在。由构造 bit 精确（scatter 和 load 同一 swizzled 布局）。
| 配置 | 阶段 6 (s8) | 阶段 7 (s9) | Δ vs s8 |
|---|---|---|---|
| B32/T8192 | 35058 | 34163–34393 | −2.2% |
| B8/T32768 | 43173 | 42325–42356 | −1.9% |
| B1/T131072 | 45911 | 45667 | −0.5% |
Kernel A 相对 pre-阶段-1 基线累计 = −68% / −65% / −50%。 收益如预测精准落地：B32/B8 （short_scoreboard 受限、冲突被消灭）赢；B1 几乎不动（barrier/占用率受限）。
s7 vs s9 的对比是全篇最重要的一课：同一个优化目标（swizzle mqk/inv/do），换个实现方式， 从 +14.5% 变成 −2.2%。 差别在：s7 加了一趟额外的 block-复制 + barrier；s9 把 swizzle 白嫖进 本来就要做的 cp.async scatter，零额外开销。"做什么"对了不够，“怎么做（有无额外 sync/pass）” 才决定输赢，尤其在 latency-bound 场景，任何多余的 barrier 都是纯亏。
第 4 部分：终局 —— 打破 bf16 的 4-way bank 地板（阶段 8-11）
这是全篇智力密度最高的一段。s9 后 ncu 说"还有 1.89B conf_ld"，但追这个数字的两次尝试 （s10 用更大 swizzle、s12-分析用 XOR-swizzle）都被证明是死路，真正的赢来自一个反直觉的 “奇数半-stride + 更小拷贝粒度”。
4.1 阶段 8/s10 —— ❌ dS INTER→SW32 真 xor-swizzle：回归，已回退
s9 ncu 确认 B32/B8 仍有 conf_ld ~1.89B。诊断：每个 INTER atom 都是 Swizzle<0,4,3> = 零个
xor 位 = 只交织、不分散 bank。真正分散 bank 的 atom 是 SW32/64/128（Swizzle<1/2/3,4,3>）。
s10 就把最大操作数 dS 从 K_INTER/MN_INTER 升级成 K_SW32/MN_SW32。
static_assert(cosize == D*Dv) 成立 → SW32 无 padding，smem 不变，占用率不是变量。gradcheck
PASS → SW32 也是 transpose-coherent。正确性 OK，但性能全线回归 +11~12%，已回退。
为何输 —— 关键线索是 B1（+12.5%，和 B32/B8 一样）：B1 不受冲突限制（barrier/占用率受限），
却同样严重回归 → 损失不是冲突权衡。SW32 给每次访问加了一个固定的额外 XOR 地址运算
（INTER 的 Swizzle<0> 会被编译器优化掉，SW32 不会），把热 GEMM 循环推过寄存器预算 → 4B spill。
教训（贯穿 s10 和后面的 s12-分析）：对我们的 16 宽 ldmatrix 子 tile，INTER 已经够了。 INTER 的"16B-run 交织"就足以让
ldmatrix.trans触发；SW32 的额外分散在这里几乎买不到冲突缓解， 却在每次访问都花地址运算 + 寄存器 → 净亏。不要把操作数"升级"到 SW32/64/128。 1.89B 的 conf_ld 计数是真的，但不能靠更大的 swizzle atom 廉价移除（和 s7 是同类红鲱鱼）。
4.2 阶段 9/s11 —— 源码级归因：真正的冲突源不是 MMA 操作数
s10 死路逼出源码级重新归因：对 s9 跑 text-mode ncu（--page source），把 1.99B excessive
shared wavefront 按 SASS 指令拆开（日志 §11）：
LDS.U16= 1.24B (62%) = R1 对q_sm/k_sm（bf16）的标量读，LDS(32位) = 336M (17%) = R1 的gc_sm（fp32），LDGSTS(cp.async) = 214M (11%)，LDSM/STSM（每个 MMA 操作数）= 0 冲突！
惊人结论：79% 的冲突是 R1 根本没被 swizzle 碰过 的 token-strided 标量读，而所有 ldmatrix 操作数 早已无冲突。 这正是为什么 s10（对 dS 的 LDSM 上 SW32）是纯代价 —— 它去追的那个操作数根本没冲突。
机制（日志 §11，务必看懂）：q_sm/k_sm/gc_sm 是 [C, D] row-major，行 stride = D = 128。
R1 读 q_sm[c*D + kp]：c（token）在 32 个 lane 间变化，kp（feature）在 fragment 槽内固定。
bf16 bank = ((c*128 + kp)/2) % 32 = (kp/2)%32 → 与 c 完全无关 → 32 个 token 全命中同一个 bank
→ 最坏 32-way 冲突。
修法（s11）：只 pad 这 3 个 R1 buffer 的行 stride —— 它们不是 MMA 操作数（qd/kd/kr 才是），
所以 padding 对 GEMM 链不可见。QK_PAD = D+8 = 136（bf16）、GC_PAD = D+4 = 132（fp32），
两者仍是 16B cp.async chunk 的倍数。
👶 上面这段太密，逐步拆开讲（小白版，5 步）
第 0 步 · 什么是 bank(存储体)。 shared memory 被切成 32 个 bank(编号 0~31),像 32 个收银窗口。
一个 warp 有 32 个 lane(线程)。黄金规则:一条读指令里,32 个 lane 若分别落到 32 个不同 bank →
一拍办完;若有 N 个 lane 挤进同一个 bank → 排队办 N 次 = N-way 冲突,慢 N 倍。bank 按地址算:
硬件按 4 字节一格编号,bank = (字节地址/4) % 32。一格 4 字节 = 2 个 bf16,所以对 bf16 元素:
bank = (元素下标/2) % 32。
第 1 步 · 这三个 buffer 长什么样。 q_sm/k_sm/gc_sm 是 [C, D] 的 row-major 矩阵:
C=CHUNK=16(token=行)、D=128(feature=列),元素 [c,kp] 的下标 = c*D+kp = c*128+kp。
行 stride = 128(跨一行,下标跳 128):
1 | feature kp → 0 1 2 ... 127 |
第 2 步 · R1 这条读,32 个 lane 读哪些格子。 R1 读 q_sm[c*128+kp],关键是谁变谁不变:
c(token)在 32 个 lane 间变化(lane0 读 token0、lane1 读 token1……),kp(feature)在 fragment
槽内固定(同一拍大家读同一列)。所以 32 个 lane 读的是 q_sm[0*128+kp], q_sm[1*128+kp], ...
—— 同一列 kp、不同行 c = 矩阵里竖着的一列:
1 | kp 固定这一列 ↓ 这 32 个 lane 竖着读同一列: |
第 3 步 · 算 bank,灾难发生。 代入 bank=(下标/2)%32:
1 | lane0 (c=0): 下标 kp → (kp/2)%32 |
每跨一个 token 下标 +128,而 128/2=64,64%32=0 —— 加的是 32 的整数倍,对 bank 毫无影响。
于是 bank=(kp/2)%32 与 c 完全无关 → 32 个 lane 全撞进同一个 bank → 32-way 冲突(全 kernel 79%)。
病根一句话:行 stride=128 是 32 的倍数(128/2=64 是 32 的倍数),竖读一列时行与行 bank 号完美重合。
第 4 步 · s11 的修法:pad 行 stride 打破整除。 既然病根是"stride 是 32 的倍数",就把 stride 撑成 不是 32 的倍数——给每行尾部塞几个废元素:
QK_PAD = D+8 = 136(bf16)。验算:跨一个 token +136,136/2=68,68%32=4→ 相邻 token 的 bank 错开 4 → token 铺开覆盖多个 bank → 冲突大降。GC_PAD = D+4 = 132(fp32,每格 1 元素)。132%32=4,同样错开。
为什么偏选 136/132?两个约束同时满足:(1) 打破 %32 整除(错开量≠0);(2) 每行仍是 16 字节
(cp.async 最小搬运块)的整数倍——136 bf16 = 272B = 16×17 ✔、132 fp32 = 528B = 16×33 ✔。既错开又保 cp.async 对齐。
第 5 步 · 为什么 pad “对 GEMM 链不可见”(最易漏)。 MMA 的 ldmatrix 操作数吃的是 qd/kd/kr,
不是 q_sm/k_sm/gc_sm。后三者只是 R1 的临时输入 buffer,padding 只加在它们每行尾巴,R1 用
新 stride 136/132 索引即可,读到的有效数据一模一样,尾部废元素根本不读。所以 padding 停在 R1 内部,
GEMM 看到的 qd/kd/kr 一个字节没变 → 零风险。这也顺带解释为什么 s10 是死路:它去给 dS(一个
ldmatrix 操作数)上 swizzle,而那个操作数本来就 0 冲突 → 追错了目标。
| 配置 | 阶段 7 (s9) | 阶段 9 (s11) | Δ vs s9 |
|---|---|---|---|
| B32/T8192 | 34163 | 30418–30643 | −10.3% |
| B8/T32768 | 42325 | 37819–37831 | −10.6% |
| B1/T131072 | 45667 | 42153–42159 | −7.7% |
ncu：conf_ld 1.89B→1.16B（−39%），short_scoreboard 2.53→1.20（−53%）。新冠军。
但 —— s11 的 gc（fp32）冲突塌缩 −90%，q/k（bf16）却几乎没动（1.24B→1.007B，仅 −19%）。为什么？
引出 s12。
4.3 阶段 10/s12-分析 —— ⏭️ XOR-swizzle q/k：写代码前就被否决
显而易见的下一步 = XOR-swizzle q/k buffer。但它在动手前就被分析否决，因为可证明打不破 s11 的地板 （日志 §12）：
- q/k 冲突是 token 变、feature 固定：32 lane 在同一 feature 读 32 个不同 token（行）。跨 token 的
bank 展开 =
((Δc)·STRIDE/2) % 32。完整 32-bank 分散需要STRIDE/2与 32 互质（奇数）， 即STRIDE ≡ 2 (mod 4)。 - 但 16B cp.async 写强制
STRIDE是 8 的倍数 →STRIDE/2是 4 的倍数 → 最多 8 个不同 bank = 4-way，硬地板。s11 pad 到 136：136/2=68 仍是 4 的倍数 → 卡在 4-way（这就是只买到 −19% 的原因）。 - XOR-swizzle 受同一约束：cutlass 式 XOR 置换整个 16B（8×bf16）块，读地址两项都是 4 的倍数 → 仍 ≤8 bank = 4-way。XOR 无法分散到"16B-coalescing 粒度"以下。
绑定约束是 coalescing 单元，不是 swizzle atom（和 s10 同类死路）。只要还用 16B 写，bf16 的 token-strided 读就永远卡在 4-way。 想破地板，必须换更小的写粒度。
4.4 阶段 11/s12 —— ✅ 奇数半-stride 130 + 4B cp.async：打破地板（当前冠军）
兑现 §12 末尾预言的"唯一破地板的办法"—— 放弃 16B 写单元（源码 L200-206、L313-327 的注释是全篇 最精华的解释）：
- q/k 存在
QK_PAD = D+2 = 130（源码 L207）：STRIDE/2 = 65与 32 互质 → 32 lane 跨 token 读 散布到全部 32 个 bank = 无冲突（不再是 4-way）。 - 4B（2×bf16）cp.async.ca（源码 L318-323
copy4）：130 不是 8 的倍数，16B 写会 misalign；改用 4B 拷贝（col 步进 2 → 目标字节 4 对齐）。注意cp.async.cg只支持 16B，4B 路径必须用cp.async.ca（cache-all L1+L2，支持 4/8/16B）—— 这就是源码 L27-30 新增copy4helper 的原因。 代价 = q/k 的 LDGSTS 数量 4×（16B→4B）。 - gc 不变（
GC_PAD = D+4 = 132+ 16B copy，源码 L208、L324-327）：s11 已把 gc 冲突塌缩 −90%， 不值得为它多付 4× LDGSTS。 - 所有 q/k 读点（源码 L416-417 l2norm、L466-467/480 的 qd/kd/kr、L370-374 tail-zero）经
QK_PAD常量自动跟随；launcher 的 bf16 smem 项从2·CHUNK·(D+8)改成更小的2·CHUNK·(D+2)。占用率不变。
| cfg | s11 | s12 | Δ vs s11 |
|---|---|---|---|
| B32 T8192 | 30418~30643 | 30123~30288 | −1.0% |
| B8 T32768 | 37819~37831 | 37455 | −1.0% |
| B1 T131072 | 42153~42159 | 40355~40401 | −4.2% |
为什么 B1 收益最大（−4.2%），B32/B8 只 −1%（日志 §13）：R1 在 B1 的关键路径上（B1 warp 少， 没法用并行掩盖 LDS 冲突深度）→ 去掉 4-way→无冲突直接缩短关键路径（B1 short_scoreboard 1.14→0.59）。 B32/B8 占用率充足，4× LDGSTS 的代价大体抵消了冲突消除的收益 → 只 −1%。这正好命中最重要的配置 （B1 = 对 FLA 差距最大处）。gradcheck L0 PASS，cos 0.99993，0-spill（VS1/2/4 = 80/106/106 regs）。
s12 是一次"精算过的赌博"：明知 4× LDGSTS 会在富并行配置上吃掉收益，但赌"延迟绑定的 B1 会净赚"。 赌赢了，且赢在最该赢的地方。教训：当一个指标（bf16 token 读冲突）被 coalescing 粒度锁死时， 唯一出路是换粒度（16B→4B）+ 互质 stride 把读散满 32 bank —— 代价（更多 LDGSTS）在延迟绑定处才划算。
4.5 当前冠军状态
- 冠军 =
bwd_a_dstate_scan_s12.cuh（bwd_launch.cu第 12 行 include 它）。 - 相对基线累计 ≈ −71.8% / −69.4% / −55.4%（B32/B8/B1；per-launch us 30.1k / 37.5k / 40.4k）。
- 历史快照保留作记录：s4(到3c)/s5(3d)/s6(STSM)/s7(回归)/s8(cp.async 预取)/s9(swizzled cp.async)/ s10(SW32 回归)/s11(pad R1，前冠军)。
- 尚未 promote 为正式
bwd_a_dstate_scan.cuh（待用户确认）。
第 5 部分：split-V —— Kernel A 专属的并行度杠杆
前面每个 conflict 阶段 B1 都持平，因为 B1（B=1, N=1）并行度饥饿：grid = N×H = 128 个 block， 但 H20 有 78 个 SM 且每 SM 能住多块 → block 太少，SM 填不满。这跟 bank conflict 正交，swizzle 帮不上。
split-V（源码 L216-217、模板参数 VSplit） 是给 B1 补并行度的招：把状态的列轴 v 切给多个 CTA。
blockIdx.z = v_split 决定这个 CTA 负责哪 Dv = D/VSplit 列（源码 L189 constexpr int Dv = D/VSplit）。
所有按 v 索引的 buffer（dS/T1/dU/do）宽度变 Dv；只按收缩轴 k 索引的（mqk/inv/qd/kd/kr）保持全宽 D。
VSplit=1 时 Dv==D，所有布局塌回原始 monolithic 视图（bit-exact）。
- 好处：grid 从 (N,H) 变 (N,H,VSplit) → block 数 ×VSplit → B1 从 128 block 变 256/512 → 填满 SM。
- 代价：每个 v-band 的 CTA 都要重算一遍 R1（qd/kd/kr 是全宽、与 v 无关）→ VSplit 越大重复越多。
- 基线时 split-V 自动派发：B1 大赢 −33%，饱和 grid（B32/B8）上中性（那里 block 本就够多）。
- VSplit ∈ {1,2,4} 由 grid 占用率在 launcher 里自动决定（
ctas_base = N*H < sm_count时开大）。
小白记住：conflict 优化（swizzle/pad）修的是"每个 warp 跑多快"；split-V 修的是"有没有足够多 warp 填满 SM"。B32/B8 缺的是前者，B1 缺的是后者。 这就是为什么整个 conflict 战役 B1 一直持平 —— 它的病在别处（并行度），得用 split-V / grid remap 这类正交杠杆治。开放杠杆：VSplit>4 / grid remap。
第 6 部分：一句话教训总结（小白记这几条就够）
-
先修占用率，再修 conflict（第 1→2 部分）。占用率没修好时，bank conflict 被延迟盖住看不出； 反过来也一样。阶段 2 把 fp32 中间量搬进寄存器、干掉大象
T1_smem→ 1→4 block/SM → −37%。 -
swizzle 要读写成对、且只对"可自由重排写地址"的操作数做（第 2-3 部分）。R1 逐元素产出的 qd/kd/kr/dS 可以 swizzle（
ldmatrix+stmatrix成对）；cp.async 载入的 mqk/inv/do 要么保持普通、 要么把 swizzle 白嫖进 cp.async scatter（s9 plan-a），绝不加额外的生产者 pass（s7 plan-b 因此 +14.5%）。 -
transpose-dual：一份字节喂两个方向（dS 的 DSb/DSc，第 2.3）。互为转置的 K_INTER/MN_INTER atom 命中同一物理字节 → 存一次、正读+转置读都无冲突、零 padding、smem 不变。
-
INTER 对 16 宽子 tile 已够，别升级到 SW32/64/128（s10，第 4.1）。更大 swizzle atom 每次访问多花 XOR 地址运算 + 寄存器 → spill → 全线回归（连不受冲突限制的 B1 都退 = 铁证不是冲突权衡）。
-
别追不在关键路径上的 ncu 指标（s7 的 805M、s10/s12-分析的 1.89B conf_ld，第 3.2/4.1/4.3）。 先看"绑定 stall 是谁"。s7 时是 gmem long_scoreboard，追 shared conflict 白费；s8 crush 掉 gmem 后 short_scoreboard 才浮上来，那时 swizzle mqk/inv/do 才值得。
-
源码级归因胜过按 atom 瞎试（s11，第 4.2）。text-mode ncu 把冲突拆到 SASS 指令，发现 79% 冲突是 R1 没被 swizzle 碰过 的 token-strided 标量读（
LDS.U16），所有 ldmatrix 操作数早已 0 冲突 —— 这才指向"pad R1 raw buffer"这个正确修法。 -
bf16 token-strided 读的 4-way 地板由 16B-coalescing 锁死，破它必须换更小写粒度（s12，第 4.3-4.4）。 奇数半-stride 130（65 与 32 互质）+ 4B cp.async.ca → 读散满 32 bank。代价 4× LDGSTS 在延迟绑定的 B1 净赚（−4.2%），在富并行 B32/B8 大体抵消（−1%）—— 精算过的赌博，赢在最该赢的 B1。
-
构建陷阱（必记）：
s*都是被bwd_launch.cu#include的 header。改 header 不touch csrc/smxx/bwd_launch.cu就不重编，stale.so给假结果。改任何 header 后必须touch bwd_launch.cu再 build。A/B 测试永远同 session、单变量（日志 §0 曾差点把变化误判成 caching-allocator 的抖动）。
附：优化日志阶段 vs 冠军源码的对应关系速查
| 阶段 | 招 | 源码行（s12） | 收益（B32/B8/B1） |
|---|---|---|---|
| 1 | 寄存器融合 C×Dv 链 | L564-632 | −6.7/−6.9/−11.9% |
| 2 | 寄存器融合 D×Dv 链，消灭 T1_smem | L634-673 | −38/−37.5/−37%（占用率 12.5→50%） |
| 3a | 普通布局 LDSM（证明正确性，收益小） | L539-552 | −1.5/−1.2/−2.6% |
| 3b | swizzle qd/kd（QdKdSw MN_INTER） |
L118-120,443-444 | −11.9/−12.3/−0.1% |
| 3c | swizzle kr（KrSw K_INTER） |
L124-126,445 | −6.8/−6.8/−1.4%（累计 −50%） |
| 3d | dS transpose-dual swizzle（DSb/DSc） | L128-145,529-530 | −22.8/−17.9/−1.0%（累计 −61%） |
| 4/s6 | STSM 写 R1 产出，消 380M store 冲突 | L485-491 | −14.4/−11.4/−5.5% |
| 5/s7 | ❌ 生产者 pass swizzle mqk/inv/do | （已回退） | +3.1/+3.2/+14.5% |
| 6/s8 | cp.async 预取 R1 q/k/gc + varlen spill 修复 | L303-328,40-46 | −1.0/−3.8/−10.5% |
| 7/s9 | swizzled cp.async scatter mqk/inv/do（plan-a） | L330-357,158-163 | −2.2/−1.9/−0.5% |
| 8/s10 | ❌ dS INTER→SW32（每访问多花 XOR → spill） | （已回退） | +10.8/+12.2/+12.5% |
| 9/s11 | pad R1 raw q/k/gc buffer（破 token-strided） | （前冠军 stride 136） | −10.3/−10.6/−7.7% |
| 10/s12-分析 | ⏭️ XOR-swizzle q/k（被 source-ncu 否决） | — | （4-way 地板锁死） |
| 11/s12 | q/k 奇半 stride-130 + 4B cp.async.ca | L200-208,313-327 | −1.0/−1.0/−4.2%（当前冠军） |
当前状态：冠军 = Kernel A s12 + Kernel B l8（后者见 bwd-kernelB-beginner-guide-full.md）。
Kernel A 相对基线累计 −71.8% / −69.4% / −55.4%。B1（低 N）下 Kernel A 串行 dstate scan 是和 FLA
差距的主因；conflict 轴已耗尽（INTER 最优、SW32 净负、bf16 q/k 读已打到无冲突），剩余开放杠杆：
- B1 并行度 —— VSplit>4（VS8 模板实例化 + 派发）/ grid remap；B1 仍 barrier/占用率受限（12.5%）。
- 计算侧 —— Compute SoL 现 64%；内存 stall 进一步下降后，GEMM 链本身成为下一个天花板。
- context-parallel（CP） —— 沿时间维切段给 B=1 补并行度（见
docs/multi_card_cp.md+cp_bwd.py）， 与 kernel 内优化正交。
附二：本文档 vs 源码/日志的补充说明
本文档与 docs/20260715-bwd-kernelA-optimization-log.md 完全一致（Kernel A 只有一份日志，不像
Kernel B 有三份分散的日志需要合并补充）。相对日志，本文档的增量是：
- 从第一性原理讲清了
MOVM_T寄存器内转置、c 轴旋转（第 0.4）—— 日志只当惯用法带过。 - 把 split-V 单列一部分讲透（第 5）—— 日志散在 §0 和各阶段的"B1 持平"注释里。
- 把"两次红鲱鱼死路（s7 的 805M、s10 的 SW32）+ 最终 s12 破地板"串成一条归因方法论主线 （第 3.2 / 4.1 / 4.3-4.4 + 教训 5-7）—— 这是 Kernel A 优化史最有价值的思维过程，日志按时间线记录， 本文档按"教训"重组，便于小白吸收。
- 标题: bwd-kernelA-beginner-guide-full
- 作者: 鱿鱼圈
- 创建于 : 2026-07-21 23:50:00
- 更新于 : 2026-09-12 22:26:40
- 链接: https://yuyanqi.com/2026/07/21/bwd-kernelA-beginner-guide-full/
- 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。