bwd-kernelB-beginner-guide-full
反向 B Kernel 优化全解析（小白版，从零讲起）
日期：2026-07-20
GPU：H20-3e（sm_90a，Hopper）。
本文对象 = 反向传播的 Kernel B（split 成 B1/B2/B3/B0 四个），当前冠军源码
csrc/smxx/bwd_b_split_fuse_b2_b3_sinfix_a1d.cuh（= 文档 l8 + 两项文档未记录的补充
sinfix / a1d）。
目标读者：不懂 cute、swizzle、寄存器、occupancy、bank conflict、coalescing 的人。 所有术语从第一性原理讲明白，再讲每一步优化在干什么、为什么有用。
上游材料：
docs/20260714-bwd-kernelB-optimization-log.md（融合 Kernel B 时代的优化）docs/20260716-bwd-splitB-L1-vectorization-log.md（split-B 的 L1 向量化轮）docs/20260719-bwd-splitB-occupancy-experiments-log.md（l5→l13 占用率实验 + 冠军状态）- Kernel A 的优化史见
docs/20260715-bwd-kernelA-optimization-log.md（冠军 = s12）。
第 0 部分：小白扫盲 —— 先把术语讲明白
所有优化都围绕下面这几个概念转。先懂它们，后面才看得懂。
GPU 的硬件层级（自上而下）
- GPU 里有很多 SM（Streaming Multiprocessor，流多处理器）。H20 有 78 个 SM。 可以把 SM 想成 78 个独立的“小核”，各干各的。
- 每个 SM 一次能装载好几个 block（线程块）。一个 block 是
<<<grid, block>>>启动的 一组线程，保证在同一个 SM 上、能共享一块高速内存。 - 每个 block 里有很多 thread（线程）。本文的 kernel 每 block 256 线程。
- 32 个线程绑成一个 warp（线程束）。warp 是硬件真正的调度单位 —— 32 线程锁步执行 同一条指令（SIMD）。256 线程 = 8 个 warp。
三种内存（快 → 慢）
- register（寄存器）：每线程私有、最快，就是变量
float x。数量有限（每 SM 65536 个 32-bit）。一个线程用太多寄存器，SM 能同时跑的线程就变少。 - smem（shared memory，共享内存）：一个 block 内所有线程共享的片上高速内存，几十~两百
KB。比寄存器慢一点，比显存快几十倍。代码里
extern __shared__就是它。本文 kernel 的 兵家必争之地。 - gmem（global memory，全局内存/显存）：几十 GB，但慢。读一次要几百个时钟周期。
occupancy（占用率）—— 贯穿全文的核心指标
occupancy = 一个 SM 上实际同时活跃的 warp 数 ÷ 硬件上限。
为什么重要？GPU 靠“人多”藏延迟：当 warp A 在等一次慢速 gmem 读（几百周期）时，SM 立刻 切去执行 warp B、C、D…… 只要有别的 warp 能跑，那几百周期就被藏住了，SM 不空转。
- occupancy 高 = SM 上住了很多 warp = 总有活干 = 延迟被藏住。
- occupancy 低 = SM 上没几个 warp = 一个 warp 卡住，整个 SM 干等 = 慢。
什么限制 occupancy？ 每个 SM 的 smem/寄存器固定。一个 block 用的 smem 越多，一个 SM 能塞 的 block 就越少。这是量子化（阶梯式）的：
- H20 每 SM 可用动态 smem ≈ 233472 字节（228KB），每 block 另占 1024B 保留。
- 一个 block 用 76KB →
233472 ÷ 76KB ≈ 3→ 3 block/SM。 - 砍到 ≤57KB → 4 block/SM。
- 注意是整数：76KB 和 68KB 都是“3 block”，占用率相同；必须跨过 57KB 的坎才变 4。 所以“省了 8KB 但没跨阶梯”= 0 收益。
这就是文档反复说的 “3-block / 37.5%”：3 block × 256 线程 = 768，硬件上限 2048， 768/2048 = 37.5%。
bank conflict（bank 冲突）—— smem 的坑
smem 硬件上切成 32 个 bank（存储体），恰好对应一个 warp 的 32 线程。 理想：一个 warp 的 32 线程访问 32 个不同 bank → 一拍搞定（conflict-free）。
bank 编号规则：bank = (地址 ÷ 4字节) % 32。连续 32 个 float 正好铺满 32 bank。
冲突怎么发生：warp 里多个线程访问落在同一个 bank（但不同地址）→ 硬件排队一个个来。
文档里 “stride-8 → 4-way bank conflict”：线程按 8 个 float 的间隔访问 (0,8,16,24,32,…)，
% 32 后只落在 {0,8,16,24} 这 4 个 bank → 32 线程挤 4 bank → 4 路冲突 → 慢 4 倍。
解法：让 warp 内线程访问连续地址（铺满 32 bank），或用 padding（填充） / swizzle 把地址错开。
coalescing（合并访问）—— gmem 的坑（bank conflict 的 gmem 版）
一个 warp 的 32 线程访问 gmem 时，硬件按 128 字节的 cache line（缓存行） 为单位搬数据。
- coalesced（合并）：32 线程读连续的数（如 32 个 bf16 = 64 字节），一两次搬运搞定。
- uncoalesced（未合并）：32 线程读分散地址（如每线程隔 128 个 float），每个线程命中 新 cache line，要搬 32 次，浪费 96% 带宽。文档 “stride-128 uncoalesced” 是灾难写法。
“73% excessive global sectors” = ncu 报告：实际搬的 gmem 数据量比理论多 73%，全是浪费。
向量化访存：uint4 / float4 / LDS.128 / STS.128
一条指令搬多少字节？
- 标量：
LDG.E.U16一次 2 字节（一个 bf16）；LDS.32一次 4 字节（一个 float）。 - 向量化：
float4= 4 float = 16B；uint4= 16B = 8 bf16。对应硬件指令 LDS.128 / STS.128 （smem 读/写 128 bit）、LDG.128（gmem 读）。
为什么快：搬同样 128 字节，标量发 8 条指令，float4×2 只发 2 条 → 指令数、访存事务数 ÷4。
当 kernel 卡在“访存指令太多”（L1 吞吐打满）时直接提速。前提：地址 16 字节对齐、数据连续。
cp.async —— 异步搬运（Ampere/Hopper 新特性）
普通“gmem→smem”要经寄存器中转，且线程必须等读完成。cp.async（cp.async.cg.shared.global）
让线程发一条“把这 16 字节从 gmem 搬到 smem”的指令后立刻返回去干别的，搬运后台进行，
用时再 wait。代码里：
1 | copy16() → 发一条 16 字节异步搬运 |
好处：把慢速 gmem 搬运和计算重叠。本文套路：先 copy16 发一堆搬运 + commit，不等、
马上跑 R1（R1 只依赖直接从 gmem 读的数据），R1 跑完搬运也差不多好，再 wait → 搬运时间被
R1 的计算藏住。
Tensor Core / MMA / cute / swizzle / LDSM —— 张量核那一套
最“黑话”的部分，拆开讲：
- Tensor Core（张量核）：SM 里专做小矩阵乘的硬件单元。一条 MMA（Matrix-Multiply-
Accumulate）指令让一个 warp（32 线程）协作算一个小矩阵块乘法（如
SM80_16x8x16= 算一个 16×8 输出 =（16×16）×（16×8）），一条指令顶几百次标量乘加。 - operand / fragment（操作数 / 片段）：MMA 要求参与运算的小矩阵按特定方式分散在 32 线程 的寄存器里（每线程持有矩阵的几个元素）。这种“矩阵切碎分给 32 线程”的布局叫 fragment。
- LDSM（
ldmatrix）：一条专门指令，把 smem 里的小矩阵块一次、按 fragment 布局灌进 32 线程寄存器，正好喂 MMA。比 32 线程自己一个个LDS快得多。代码里SM75_U32x4_LDSM_N（正 向）、SM75_U16x8_LDSM_T（转置）就是它。_N=原方向；_T=加载时顺便转置。文档 “ldsm=0” = 某 kernel 一次 ldmatrix 都没用、全是慢速标量 LDS。 - cute：CUTLASS 里一套描述“数据在内存怎么摆放”“哪个线程负责哪个元素”的 C++ 抽象。
make_tensor / Layout / local_tile / partition_C / cooperative_gemm全是 cute。理解成：用类型 系统把易错的“矩阵布局”在编译期算好，纯编译期、无运行开销。 - swizzle（异或置换）：一种打乱 smem 地址的固定数学变换，专消 bank conflict。LDSM 要按
fragment 读 smem，若 smem 朴素行主序摆放，LDSM 的读法会撞 bank；swizzle 通过对地址做异或
变换把撞同一 bank 的元素散到不同 bank，让 LDSM conflict-free。代码里
GMMA::Layout_K_INTER_Atom、tile_to_shape(...)、SwStateLayout、SwSinLayout就是定义 swizzle 布局。关键：同一份物理数据，写时按 swizzle 摆、读时也按 swizzle 算地址，两边对上即可。 K_INTERvsMN_INTER/ 转置对偶（transpose-dual）：一个矩阵在矩阵乘里可能被“正着乘” 也可能“转置着乘”。K_INTER适合正读（LDSM_N），MN_INTER适合转置读（LDSM_T）。妙处： 对 INTER 类 swizzle，同一份物理字节既能K_INTER正读、又能MN_INTER转置读 —— 存一次， 喂两个方向的 MMA。这就是代码反复出现的 “transpose-dual over the SAME physical bytes”， 省一份 smem + 省一次转置。
warp shuffle 归约（__shfl_xor_sync）
warp 的 32 线程各算了一个部分和，要加成总和。老办法写 smem 再读回（慢、占 smem）。
__shfl_xor_sync 让 warp 内线程直接交换寄存器值，不经 smem：
1 | for (delta = 16; delta >= 1; delta >>= 1) |
树形归约：第一轮线程 i 和 i^16 交换相加，第二轮 i^8…… 5 轮（16→8→4→2→1）后每个线程 都持有 32 个数的总和。log2(32)=5 步搞定，不碰 smem。dA_log、dbeta、dgT-dot 归约全用它。
wall / roofline / ncu（性能诊断）
- ncu（Nsight Compute）：kernel 剖析器，告诉你卡在哪。
- wall（墙）：真正的瓶颈。ncu “L1/shared 85.8%” = L1/smem 吞吐用到硬件上限 85.8% → 这是墙，想快得减少 smem 访存流量。“Compute 42.6%” = 算力才用一半 → 算力不是墙。
- stall reason（停顿原因）：
short_scoreboard=等 smem 读（LDS）；long_scoreboard=等 gmem 读（LDG）；barrier=等__syncthreads；mio_throttle=访存/转换指令发射队列满（如 bf16 打包 F2FP 太多）。 - 诊断主线（整场战役的灵魂）：这条 split-B 路径是 occupancy/latency-bound。只有两类优化 真提速：(a) 提高占用率（跨 smem 阶梯）；(b) 直击真墙、砍其访存流量。其余一切（纯 issue-rate 微调、cp.async 加流水、GEMM 链重写）一律 ~0% 或负 —— 那些延迟本来就被占用率藏住了。
第 1 部分：为什么要把一个大 kernel 拆成四个？（全局设计）
背景：反向里有个巨大的融合 kernel _flash_kda_bwd_grads_fused，把一个 tile 需要的所有
中间量塞在 smem（176KB）。176KB > 233KB/2，所以一个 SM 只能塞 1 个 block，occupancy 仅
12.5%。block 里 warp 一卡住整个 SM 就干等 —— 严重 latency-bound。
Plan B（照抄 FLA）：拆成 4 个小 kernel，每个只留自己那步的 smem（~70KB），中间结果写回 gmem 传给下一个。smem 小 → 一个 SM 塞 3 block → occupancy 37.5%（×3）→ 延迟被藏。用“多写几趟 gmem 中间量”换“占用率翻三倍”，在 latency-bound 场景净赚。
四个 kernel（都是 grid=(总tile数, H)，每 tile 一个 block，tile 间无依赖 → 天然并行）：
| kernel | 干什么 | smem | 墙 |
|---|---|---|---|
| B1 | dS 侧：dvcorr、dkr、dv | ~74KB，dS 常驻 | Tensor Core 依赖延迟（健康） |
| B2 | sin 侧：dqd_base、dkd_base、dbeta；折叠了 B0 dgT-dot | ~68KB，sin 常驻 | gmem 延迟（dS 流） |
| B0 | dgT-dot：p[k]=Σ_d dS[k,d]·sin[d,k] |
无 smem，流式 | 已死，融进 B2 |
| B3 | chunk 空间收尾 + 内联 K1（L2norm 反向 + 反向 cumsum） | ~95KB | fp32 epilogue 的 FFMA / L1（最大热点） |
第 2 部分：逐 kernel 优化解析
B1 —— dS 侧（算 dvcorr、dkr、dv）
它算什么：
dkr[c,k] = Σ_d U[c,d]·dS[k,d]（U 与状态梯度 dS 收缩 d 维）dU[c,d] = Σ_c' Mqk[c',c]·do[c',d] + Σ_k kr[c,k]·dS[k,d]（两个 GEMM 累加）dvcorr[c,d] = Σ_c' INV[c',c]·dU[c',d]dv = dvcorr · beta
优化 1：寄存器驻留的 dU 链 + MOVM_T 寄存器内转置（源码 L289-408，文档 0719§1）
朴素做法：算完 dU 写进 smem 再读回喂下一个 GEMM（G4），一来一回浪费。
冠军做法：让 dU 结果一直待在寄存器（MMA 累加器 dacc 本就在寄存器）。G4 需要 dU 的转置，
硬件有条 SM75_U32x1_MOVM_T 能在寄存器内部直接转置小矩阵块（不经 smem）。于是
dacc → 转 bf16 → MOVM_T 转置 → 直接当 G4 的 B 操作数。省掉 dU 的 smem 缓冲（-7KB）+ 标量
读回 + 一次 gmem 加载。
优化 2：dS 的 swizzle 转置对偶（L170-179）
dS（[128,128] 大矩阵）被两个方向用：dkr 正读（收缩 d），G2 转置读（收缩 k）。用 INTER 转置对偶：
SwStateLayout（K_INTER，正读）与 SwStateLayoutT（MN_INTER，转置读）共享同一份物理字节。
dS 只存一次，两个方向 LDSM 都能读且 conflict-free。
优化 3：fp32 C-staging tile 行填充 stride D+4（L156-165）
dkr/dvcorr 的 MMA 结果存进 [16,128] fp32 smem tile。若行 stride=128（% 32 = 0），同列每行落
同一 bank → ~75% 冲突。改 stride D+4=132（% 32 = 4），每行错开 4 bank → 8 行落 8 bank →
conflict-free。用 +4 而非 +1 保证行字节数 132×4=528B 仍 16B 对齐，cute 向量化 C-store 不崩。
代价：多分配 512B smem（余量内，占用率不变）。同 session A/B ~0.2-0.5%，bit-identical。
优化 4（补充，文档只泛讲）：cp.async 分组重叠 R1（L209-287）
先把所有 bf16 GEMM 操作数（dS/U/INV/Mqk/do）用 cp.async 发出 + commit，不等，马上跑 R1
（重建 kr，只依赖直接从 gmem 读的 raw k/gc）。R1 算完再 wait<0> —— gmem 搬运被 R1 藏住。
优化 5：R1 内部向量化 + warp-shuffle（L247-285）
R1 重建 kr = kn·exp2(-gc)·gT，每线程负责 8 连续元素（EPT=8）：
- k 读：8 连续 bf16 = 一次
uint4(16B) 合并加载（原 8 次 stride 标量LDG.E.U16）。 - k 的 L2 范数：
__shfl_xor_sync树形归约求Σk²，不碰 smem。 - gc 读：8 连续 fp32 = 两次
float4合并加载。 - kr 写：凑成
alignas(16) BF16[8]=16B，一条 STS.128 写到 swizzle 基址。
B1 的墙：HMMA short_scoreboard 18.6%，占用率 37% —— 健康的 Tensor Core 依赖延迟。
文档 0719§7 指出下一步杠杆是把 fp32 staging 也 alias 到死掉的 bf16 tile、砍 smem 跨 4-block，但
需先 ncu 确认 B1 是 smem-limited 而非 register-limited。
B2 —— sin 侧（算 dqd_base、dkd_base、dbeta；折叠 B0）
它算什么：
dqd = do @ S_in^T、dkd = dvcorr @ S_in^T（收缩 sin 的行）kdS = kd @ S_in^T（dbeta 路径，收缩 sin 的列）dbeta_vcorr[c] = Σ_d dvcorr[c,d]·(v[c,d] − kdS[c,d])- 折叠进来的 B0：
p[k] = Σ_d dS[k,d]·sin[d,k]
优化 1（★补充，文档完全没写 —— 文件名里的 “sinfix”）：sin 双方向 swizzle 修正 （L470-483, L592-599） 这是一个正确性 bug 的修复，不只是性能。三个 GEMM 里：
- dqd/dkd 需收缩 sin 的行（用
SwSinLayoutT/ MN_INTER / LDSM_T，转置读）； - kdS 需收缩 sin 的列（用
SwSinLayout/ K_INTER / LDSM_N，正读）。
旧代码的错（注释 L476-478, L597-599 白纸黑字）：图省事三个 GEMM 全用一个 K_INTER 读，统一
收缩 sin 的列。这对 kdS 恰好正确，但 dqd/dkd 被悄悄转置了 —— 只要初始状态 S_in ≠ 0，
dq 和门控梯度（gate grads）就算错。dv/dbeta 因走 kdS 那个恰好对的方向而没暴露。sinfix 给
dqd/dkd 换上转置对偶的 SwSinLayoutT，三个 GEMM 各读各的正确方向 —— 仍是同一份物理字节存
一次、喂两个方向。这是“文档不全、需补充”的头号项。
优化 2：kdS 复用已死 bf16 区（smem union，跨 2→3 block）（L461-465）
kdS_smem 不申请新空间，而是别名（alias）到 do_bf —— 因为 do_bf/dvcorr_bf 只在 GEMM 的
kk 循环里被读，循环结束就“死了”。复用这块死内存 → 省一整块 fp32 CD tile → smem 78.9→68.1KB →
跨过 2→3 block 阶梯（占用率 25%→37.5%）。代价：循环后加一个 __syncthreads()（L672），防某
warp 抢先往 kdS 写、覆盖了另一个还没读完 do_bf 的慢 warp。“占用率优化 cash 到 %”的典型。
优化 3：dvcorr 窄化的向量化（L522-542，文档 Opt#1/#3）
dvcorr 从 gmem(fp32) 读进来，要同时（a）存进 fp32 smem 给 dbeta、（b）转 bf16 存进 swizzle 区给
GEMM。逐元素做：gmem 读 stride-8 未合并、bf16 写 stride-8 四路冲突。改成每线程 8 连续元素：gmem
float4×2 合并读、fp32 smem float4×2 STS.128 写、bf16 一条 uint4 STS.128 写到 swizzle 基址。
文档说这是 B2 里最大的单项收益。
优化 4：折叠的 B0 dgT-dot，复用片上 sin（L677-702，文档 0719§0）
原独立 B0 要从 gmem 重读一遍 all_states（它 L1 打满的元凶）。但 sin 已常驻 B2 的 sin_bf
（sin[d,k]=sin_bf(d,k)，swizzle 地址算得到）。折进 B2 后 dgT-dot 只需从 gmem 流式读 dS（合并读：
相邻 lane 读相邻 d），sin 直接用片上 → 干掉整条 all_states 重读。归约用 warp-per-k：8 warp × 16 个
k，每个 k 的 128 维点积由 32 lane 分摊（每 lane 4 个），__shfl_xor 树形归约、lane0 写。
B0 为什么本身“死了”（L37-40, L733-773）：把 B0 折进 B1（dS 常驻）试过 → 反而 -0.8%。因为 B0 是延迟受限的串行依赖归约，不是带宽受限；让 dS 常驻（stride-D 的 bank 冲突 smem）省不到东西， 且它的尾部追加没法和 B1 的 Tensor Core GEMM 重叠。折进 B2（sin 常驻）才对 —— 因为 B0 要的正好 是 sin。教训（文档 0719§4.3 l12）：从 gmem 取的那个操作数，应该是能合并读的那个。
B2 的墙：IMAD/LDG long_scoreboard 14.2%，是每 tile 32KB 的 dS gmem 流延迟 —— latency/
occupancy-bound。所以 §4.1（l9 给它加 cp.async 流水）和 §4.4（l13 向量化 dS 读）都是负收益：
延迟本来就被占用率藏住了，加流水只增开销；向量化还打乱了 sin 的 swizzle bank pattern 引入新冲突。
B3 —— chunk 空间收尾 + 内联 K1（最大热点，33.5%）
最大、优化最密集的 kernel。做 chunk 空间的最终梯度组装 + 把原本独立的 K1 kernel（L2norm 反向
- 门控反向 cumsum）内联进来，省一次 kernel 启动和一趟 gmem。核心铁律：整个 epilogue 必须 全程 fp32（gate-locked），不能用 bf16。
铁律来源（文档 0719§4.2 l11）：门控梯度经由
dkn_signed = dkd·e^gc − dki·e^−gc − dkr·e^gt·gc，这是个强抵消表达式（dkd、dki 是 O(1)，相减
抵消到 O(0.01)）。bf16 操作数的 ~4e-3 误差在抵消里存活，再被 K1 反向 cumsum 相干放大 →
dg 从 2.6e-3 爆到 6.9e-2（炸 26 倍）。所以 B3 所有优化都绕开数值、只动布局/访存。
优化 1：l7 smem union + R1b 重算（★补充，L838-851, L1014-1055，文档 0719§2）
epilogue 需 3 个 fp32 CD 缓冲（qd_f/kd_f/ki_f，24KB）。但此时 6 个 bf16 CD tile
（do_bf/U_bf/… 24KB）在 G5/G6 的 __syncthreads 之后全死了。让 qd_f/kd_f/ki_f 别名到这些
死 bf16 tile（6 bf16 CD = 24KB = 正好 3 fp32 CD）。别名后原 R1 算的值被覆盖怎么办 → R1b：在
G5/G6 barrier 之后，用完全相同的 fp32 输入 + 相同运算顺序把 qd_f/kd_f/ki_f 重算一遍填进
别名空间 → bit-identical。q/k/gc 从 gmem 重新合并读（延迟被藏），L2norm 倒数从 R1 缓存的
q_inv_smem/k_inv_smem 拿。把 b3 峰值 smem 100.75KB→68.75KB → 跨过 2→3 block/SM。
“用一次重算换占用率”。
优化 2：l8 epilogue float4 向量化（当前 B 冠军的招牌）（L1101-1142，文档 0719§3）
行级 PC 采样发现：b3 全部停顿的 63% 集中在 epilogue 那 4 行 —— 3 个 fp32 chunk 空间 GEMM。
旧写法：每线程只拥有 1 个输出 idx=c·D+k，内层 K=16 循环做 stride 标量 LDS（操作数 stride128、
mask stride16），零复用，打满 mio/LDS 管线。
修法（保持全 fp32 → bit-identical）：每线程改拥有 4 个连续 k（idx4，base=idx4·4，因
D%4==0 故 4 个输出共享行 c）：
- 操作数
ki_f/qd_f/kd_f[m*D+k0]变float4LDS.128（smem 事务 ÷4）； - k 无关的 mask 列（
dMqk行/列、dL列·beta）提到每个 c 只读一次存进寄存器数组m_qd[]/m_ki[]/m_dl[]，4 个 k-lane 间复用（原来每输出重读一遍）。 累加顺序按 k-lane 精确保留（acc_ki先算 dMqk@qd、再 beta·dL@kd）→ bit-identical。寄存器 72→74（<85 的 3-block 上限，占用率不变）。这是第一个 cash 真实 % 的非占用率微优化，因为它 直击 b3 的 L1/mio 墙（最热 4 行上 LDS 流量 ÷4）。
优化 3：Opt#2 —— 删掉死掉的 bf16 打包（L963-967, L998-1011）
活跃的 fp32-carry epilogue 读的是 fp32 mask（dMqk_smem/dL_smem，即 G5/G6 MMA 的 fp32 原始
输出）。那些 bf16 副本（qd_bf/dMqk_bf/dL_bf/dLTb_bf）只有诊断路径 B3_SCALAR_EPI_DIAG 才读。
把它们的存储 #ifdef 掉，省一堆 F2FP.BF16 打包（缓解 mio-throttle）+ STS + 整条 dLTb_bf 的 CC-pass。
优化 4：Opt#7 —— 缓存 L2norm 倒数给 K1 复用（L834-837, L925-927, L1219-1221）
R1 已为每行算了 q_inv/k_inv（L2norm 倒数）。内联的 K1 需要同样倒数，且 R1 与 K1 用完全相同
的 tid→行 映射。于是 R1 把倒数写进 q_inv_smem/k_inv_smem（每行只写一次），K1 直接读 → 省掉
K1 里重复的 q_sq/k_sq shfl 归约 + 两次 rsqrt。
优化 5：Opt#8 —— K1 里 5 个 fp32 scratch 读的向量化（L1232-1250）
K1 每行读 5 个 fp32 中间量（dqd/dkd/dU/dki/dkr）。旧写法 stride-8 标量 LDS → 4 路冲突。用宏
LDV_ 把每个变成 float4×2 LDS.128，连续读 → conflict-free。
优化 6（★补充，文档没细讲）：l7c 列交错 conflict-free 布局（L1019-1054, L1180-1186）
K1 里每线程本来拥有一行内连续 8 列 [my_col, my_col+8)，这样 float4 STS 的 8-thread phase 把
bank 塌缩成 2 组 → 2× 存储 wavefront。改成每线程拥有两个不相连的 float4 块：[c0,c0+4) 和
[c1,c1+4)（c1=c0+64）。这样每次 float4 LDS/STS.128 的 8-thread phase 恰好铺满 32 bank →
conflict-free。逐元素写 [row,col] 和行归约分区无关 → bit-identical。R1b 也用了 EPT=4 交错。
优化 7：内联 K1 + dkd finalize 折叠 + dg 反向 cumsum（L1164-1339）
- dkd finalize 折叠（L1164-1166, L1258-1260）：
dkd = beta·(−dkd_base + s2)不再单独走一趟 smem 读写，而是折进 K1 phase1，算 dkn 时顺手读dkd_smem+dU_smem(s2) 现算 → 省一整趟 CD 读写 + 一个__syncthreads。 - 反向 cumsum → dg（L1298-1326）：门控梯度需沿 chunk 内 token 反向累加（
rev_sum）。每线程 负责一列，从row=15倒着加到row=0，途中算 dg、累加 dA_log/ddt_bias。 - dA_log 的 warp-shuffle + 单原子加（L1327-1330）：
dA_log是整个 head 的一个标量，256 线程都 有贡献。先 warp 内__shfl_xor归约成每 warp 一个值，只让 lane0 做一次atomicAdd→ 原子从 256 次降到 8 次。ddt_bias 每列一个 → 每列一次 atomicAdd。
优化 8（★补充，文件名 “a1d”）：原生 [T,H,D] 输出布局（L412-414, L1189-1193, L1299-1333）
文档没写的第二大补充。梯度输出（dq/dk/dv/dg/dbeta）直接写 native [T,H,D] 布局（token 步长
= H·D），而非旧的 [H,T,D] 中间暂存再转置。省掉一次布局转换 pass。代价：phantom 行
（row ≥ actual_len，varlen 里越界 token）在 native 布局没有对应地址，直接 skip（旧 [H,T,D]
暂存会往那写 0，native 里那是越界地址不能写）—— 见 L1320-1323 注释。
B3 的墙：FFMA short_scoreboard 29.1%，fp32 epilogue，gate-locked 不能 bf16。它是最大的 B
kernel（占 bwd 33.5%）。文档 0719§7 下一步是给 epilogue 加 ILP（idx4→idx8，更多独立 FFMA 链藏
延迟），盯 85-reg 3-block 上限。
第 3 部分：文档“不全、需补充”的优化（集中列出）
核对源码后，以下几项代码里有、三份文档没（充分）记录：
- ★ sinfix —— sin 双方向 swizzle 的正确性修复（B2，L470-483, L592-599）。不是性能优化而是
bug fix：旧代码三个 sin-GEMM 统一收缩列，导致
S_in≠0时 dqd/dkd 被悄悄转置、dq 和门控 梯度算错。用转置对偶（SwSinLayoutT/LDSM_T）给 dqd/dkd 换正确方向。三份文档均未提，最 需补讲。 - ★ a1d —— 原生 [T,H,D] 输出布局（B1 dv、B3 dq/dk/dg/dbeta）。梯度直接写 native（token 步长 H·D），省掉 [H,T,D]→native 的转置趟；phantom 行因无地址而 skip。文档只覆盖到 l8，未含 a1d。
- cp.async 分组 + 不等待重叠 R1 的精确套路（B1/B2/B3 都有，B3 是双 group A/B，L869-886）。 文档 0719§4.1 只从负面角度提过 cp.async（l9 给 dgT 加流水是负收益），但主路径里 cp.async 藏 gmem 搬运是保留的赢法，这点文档没正面讲。
- l7c 列交错 conflict-free 布局（B3 R1b 和 K1，L1019-1025, L1180-1186）。文档 0719§0 提 l7 是 “union 到 3-block”，但“两个不相连 float4 块 c0/c1=c0+64 让 phase 铺满 32 bank”只在源码注释里。
- 寄存器内
MOVM_T转置的 dU 链（B1，L384-407）。文档 0719§1 一句“寄存器驻留 + MOVM_T” 带过，没讲 MOVM_T 是“寄存器内小矩阵转置指令、省 smem 往返”这层语义。 - Opt#7 L2norm 倒数跨 R1/K1 缓存、Opt#8 K1 scratch 向量化、dkd finalize 折叠进 K1（B3）： L1-vectorization 文档只是编号出现，没展开“为什么能复用（tid→行映射相同）、折叠省了什么（一趟 CD 读写 + 一个 barrier）”。
第 4 部分：一句话教训总结（小白记这几条就够）
- 这条路是 occupancy/latency-bound。能提速的只有两类：
- (a) 砍 smem 跨过 block/SM 阶梯（B2 kdS union 2→3 block、B3 l7 union）；
- (b) 直击某个真墙、砍它的访存流量（B3 l8 在最热 4 行把 LDS ÷4）。
- 其余一切（cp.async 加流水、GEMM 链重写、给 latency-bound 的读做向量化）一律 ~0% 或负 —— 那些延迟本来就被占用率藏住了。
- “从 gmem 取的那个操作数，必须是能合并读（coalesced）的那个”（§4.3 l12：dgT-dot 折进 B1 时 sin 变 stride-128 未合并 → +34%；折进 B2 用片上 sin、只流式读合并的 dS → 赢）。
- “向量化时别打乱和它同一次乘法里 co-read 的 swizzled-smem 操作数的 bank pattern”（§4.4 l13： 把 dS 的 per-lane 分配改连续，顺带打乱了 sin_bf 的 swizzle bank pattern → 引新冲突 → +5~8%）。 float4/128-bit 向量化只在“LDS 事务数本身就是墙”时赢（B3 l8），对 latency-bound 的 B2 dgT 无效。
- B3 epilogue 必须全程 fp32（gate-locked）：门控梯度
dkn_signed强抵消 + 反向 cumsum 相干放大， bf16 会炸 gate 梯度 26~76 倍。所有 B3 优化都绕开数值、只动布局/访存。 - 构建陷阱（必记）：
s*/l*都是被bwd_launch.cu#include的 header。改 header 不touch csrc/smxx/bwd_launch.cu就不重编，stale.so会给“两臂看起来一样”的假结果。改任何 bwd header 后必须touch bwd_launch.cu再 build，并核实 A/B 两臂确实不同。 - A/B 测试永远同 session、单变量（§4.5：曾差点把 Kernel A 的变化误判成 caching-allocator 的 −10%）。
附：三份文档 vs 冠军源码的对应关系速查
| 优化 | 出处文档 | 源码行 | 收益 |
|---|---|---|---|
| 拆 4 kernel（占用率） | 0716 | 全文件结构 | −12~16% |
| B1 寄存器 dU 链 + MOVM_T | 0719 §1 | L289-408 | ~0%（省 smem） |
| B1 fp32 tile 行填充 D+4 | 源码注释 | L156-165 | 0.2~0.5% |
| B2 kdS union（2→3 block） | 0719 §5 | L461-465 | 占用率 cash |
| B2 dvcorr 向量化 | 0716 Opt#1/#3 | L522-542 | B2 最大项 |
| B2 折叠 B0（复用片上 sin） | 0719 §0 | L677-702 | 干掉 all_states 重读 |
| B2 sinfix（正确性） | 无（补充） | L470-483 | bug fix |
| B3 l7 smem union + R1b 重算 | 0719 §2 | L838-851,1014-1055 | 跨 2→3 block（+~1%） |
| B3 l8 epilogue float4 | 0719 §3 | L1101-1142 | −2~3%（B 冠军招牌） |
| B3 Opt#2 删死 bf16 打包 | 0716 Opt#2 | L963-1011 | −0.5% |
| B3 Opt#7 L2norm 倒数缓存 | 0716 Opt#7 | L925-927,1219 | issue 省 |
| B3 Opt#8 K1 scratch 向量化 | 0716 | L1232-1250 | 消 4 路冲突 |
| B3 l7c 列交错 conflict-free | 仅源码注释（补充） | L1019-1025,1180 | 消 2× wavefront |
| B3 dkd finalize 折叠进 K1 | 源码注释 | L1164,1258 | 省一趟 CD + 一 barrier |
| dA_log warp-shuffle 单原子 | 0714 | L1327-1330 | 256→8 原子 |
| a1d 原生 [T,H,D] 输出 | 无（补充） | L412,1189,1299 | 省转置趟 |
当前状态：冠军 = Kernel A s12 + Kernel B l8（本文件即 l8 + sinfix + a1d）。bwd-only 实测
B32/B8/B1 = 242/254/196ms（vs FLA 241/244/123）。B1 低 N 下 Kernel A 串行 dstate scan 占 52% 是
和 FLA 1.60× 差距的主因 —— 文档 0719§7 明确指出：同等力气，攻 Kernel A 的 V-dim grid-split
（修低 N 占用率悬崖）期望回报 > B 上任何剩余杠杆。
附录 B：深挖问答完整原文（一字不落）
以下为六个深挖问题 + 全部追问的完整原始问答，逐字保留（bank 冲突怎么数、SPMD 藏延迟、swizzle/XOR、K_INTER/MN_INTER 转置对偶、D+4 填充、三处向量化、uint4/STS.128、两行
using声明、Swizzle<B,M,S>源码、gemmvscooperative_gemm、cute::transform），文末附print_layout真实运行输出。
问题 1:为什么"32 线程挤 4 bank"叫"4 路冲突"?你的理解对不对?
你的理解完全正确:冲突路数(serialization factor)的严格定义 = 同一个 bank 上同时有几个线程访问不同地址。不是看用了几个 bank,而是看"最挤的那个 bank 上压了几个线程"。
那"4-way"这个说法哪来的?我把 stride-8 的访问画出来你就懂了。
代码里那个会冲突的存储,线程映射是 r_row = tid/16, r_col = (tid%16)8,即一行 16 个线程,每个线程负责 8 个连续元素,但逐个标量存(for i: smem[r_rowD + r_col + i])。
在某一步 i 固定时,一行 16 个线程写的地址是 r_row*128 + r_col + i。bank 编号 = 地址 % 32: bank = (r_col + i) % 32, 其中 r_col = (tid%16)*8
tid%16: 0 1 2 3 4 5 6 7 8 9 10 11 … r_col: 0 8 16 24 32 40 48 56 64 72 80 88 … %32: 0 8 16 24 0 8 16 24 0 8 16 24 … bank: [0] [8][16][24][0] [8][16][24][0] [8][16][24]… 16 个线程只落在 4 个 bank:{0, 8, 16, 24},每个 bank 上挤了 4 个线程(16÷4=4)。
→ 严格按你说的"同一 bank 几个线程"来数,就是 每个 bank 4 个线程 = 4 路冲突,慢 4 倍。“4-way” 和"只用了 4 个 bank"在这个 stride-8 + 16 线程/行 的场景下恰好数值一致,所以两种说法都对、指的同一件事。
▎ 注:如果按整个 32 线程 warp(跨 2 行)算,{0,8,16,24} 每 bank 就是 8 个线程。文档按"每行 16 线程"这个粒度叫它 4-way。关键点你抓对了:看的是每个 bank 上的线程数,不是 bank 总数。
修复:让线程访问连续地址(见问题 6 的 float4 向量化)。连续 32 个 float 正好铺满 32 个 bank,每 bank 1 个线程 → conflict-free。
问题 2:SPMD 下每个线程干一样的活,怎么会 warp A 跑 LDG、warp B 跑 MMA?
这是最关键的一个误解,拆清楚:
SPMD/SIMT 是"所有 warp 跑同一份代码",但不是"所有 warp 在同一时刻跑到同一行"。
每个 warp 有自己独立的 PC(程序计数器),它们在同一份代码里各自走到不同位置、彼此异步。warp 调度器每个周期挑一个"当前 ready(不在等数据)"的 warp 发射它的下一条指令。
关键机制叫 latency hiding via TLP(线程级并行藏延迟): 时间轴 → warp A: … LDG(发出) ──────等数据(几百周期)──────> MMA … warp B: … MMA MMA LDG(发出) ──等──> … warp C: … LDG ──等──> MMA ADD …
- warp A 在某时刻卡在"等 LDG 返回",它这几百周期什么也发不出。
- 但 warp B 早就过了它自己的 LDG(它比 A 提前几步),此刻正好走到 MMA。
- 调度器发现 A 不 ready、B ready → 这个周期发 B 的 MMA。
所以不是"给不同 warp 写了不同代码",而是同一份代码里的多个 warp 在时间上错开了(drift out of sync),自然形成"有的在等 LDG、有的在算 MMA"的交叠。 这是全自动的,不需要你写任何 if(warpid)。
你说的 warp specialization(warp-spec) 是另一回事:那是你故意写 if (warpid < 4) 当生产者搬数据; else 当消费者算 让不同 warp 扮演不同角色。基础的藏延迟不需要 warp-spec —— 只要 SM 上住着足够多的 warp(= 高 occupancy),它们自然错峰,延迟就被填满了。这也正是"占用率高才 能藏延迟"的物理原因。
补充两点:
- 单个 warp 内部也能藏一点(ILP):LDG 是"发出即返回"的非阻塞发射,warp 发完 LDG 可以继续发几条不依赖该数据的独立指令,直到真正要用那个数据才卡住。
- cp.async 把这件事做到极致:它显式地让"搬运"和"使用"解耦,搬运在后台跑,warp 继续算别的(见问题 6 前的 R1 重叠)。
问题 3:为什么朴素行主序的 smem 被 LDSM 读会撞 bank?swizzle 怎么修?
先讲 LDSM(ldmatrix)怎么读 smem。它是给 Tensor Core 喂矩阵的指令:一条 ldmatrix 让 warp 的 32 个线程各提供一个地址,每个地址指向 smem 里一个 16 字节(8 个 bf16)的行片段;硬件按这些地址把一个小矩阵灌进 32 线程的寄存器(fragment)。
朴素行主序为什么撞:假设一个 [16,16] bf16 tile,行主序摆放。一行 = 16 个 bf16 = 32 字节 = 8 个 word(4 字节)。第 r 行的起始地址 = r8 word。LDSM 让 32 个线程分别指向这 16 行的行首(每行被两个线程分),bank = 行首 word % 32: 行 r: 0 1 2 3 4 5 6 7 8 … r8: 0 8 16 24 32 40 48 56 64 … %32: 0 8 16 24 0 8 16 24 0 … bank: [0] [8][16][24][0] [8][16][24][0] … 16 行的行首只落在 {0,8,16,24} 4 个 bank → 4 路冲突。根因:行宽(32 字节 = 8 word)恰好是 32 bank 的因子,导致每隔 4 行 bank 就绕回来撞上。
swizzle 怎么修:swizzle = 存数据时对地址做一个固定的异或(XOR)置换,把本来撞同一 bank 的行"错位"到不同 bank。核心思路: 物理地址 = 基址 XOR (行号的某几位 << 偏移) 效果:第 r 行的行首不再是老实的 r*8,而是被 XOR 打散。设计目标就是让 16 个行首 word 的 %32 变成 {0,8,16,24, 4,12,20,28, …} 这样铺满 32 个 bank、每 bank 1 个 → conflict-free。
直观理解:异或是一个双射置换(不丢数据、不重叠),它只是把"逻辑上第 r 行第
c 列的元素"换个物理 bank 存。因为写的时候按 swizzle
摆、读的时候(LDSM)也按同一个 swizzle
算地址,两边对得上,数据不会错,但物理 bank 分布被打散了。代码里
tile_to_shape(GMMA::Layout_K_INTER_Atom
问题 4:K_INTER vs MN_INTER 转置对偶(transpose-dual),为啥省一份 smem + 省一次转置?用在哪?
先说要解决的痛点。 Tensor Core 的矩阵乘 C = A @ B,硬件对 A 和 B 在 smem 里的摆放方向有要求:
- 算 dS @ x 时,dS 要按正常方向(行 = 行)喂进去。
- 算 dSᵀ @ y 时,同一个 dS 却要按转置方向(行 = 列)喂进去。
同一块数据 dS,一会儿要正着读,一会儿要转置读。朴素做法有两条烂路:
- 存两份:一份正的 dS,一份转置的 dSᵀ。→ 多花一倍 smem(smem 是 split-B 抢 occupancy 的命根子,多一份可能从 3-block 掉到 2-block)。
- 存一份 + 运行时转置:用一次额外的 smem→smem 搬运把 dS 转成 dSᵀ。→ 多一趟 LDS/STS,还常常撞 bank。
两条路都要么费 smem、要么费指令。
transpose-dual 的解法:让"同一堆物理字节"既能正读又能转置读。
关键在 LDSM(ldmatrix)有两个变体:
- LDSM_N(normal):把 smem 的 16×16 块按正常方向装进 fragment。
- LDSM_T(transposed,汇编里 .trans):读同一块 smem 字节,但硬件在装进 fragment 时顺手把行列换了。
问题是:LDSM_T 要 conflict-free,smem 的 swizzle 摆法必须和 LDSM_N 想要的摆法兼容。cute 用两种 atom layout 描述这件事:
- GMMA::Layout_K_INTER_Atom → 配 LDSM_N(K 方向内层交错,正读友好)。
- GMMA::Layout_MN_INTER_Atom → 配 LDSM_T(MN 方向内层交错,转置读友好)。
这就是"对偶": 你写一次 smem(用一种 swizzle),然后声明两个 cute view 指向同一块 smem —— 一个用 K_INTER layout(给 LDSM_N,正读),一个用 MN_INTER layout(给 LDSM_T,转置读)。两个 view、一份数据。
物理 smem(只存一份 dS,写一次)
│
┌───────────┴───────────┐
K_INTER view MN_INTER view
(LDSM_N 正读) (LDSM_T 转置读)
│ │
GEMM1: dS @ x GEMM2: dSᵀ @ y
代码里用在哪(champion .cuh):
- B1,dS 侧(L170-179):SwStateLayout(K_INTER)+ SwStateLayoutT(MN_INTER)一对。dS 存一份,dkr 那条 GEMM 正读,后面需要 dSᵀ 的 GEMM 转置读。
- B2,sin 侧(L470-483):SwSinLayout(K_INTER,给 kdS 那条列方向 GEMM)+ SwSinLayoutT(MN_INTER,给 dqd/dkd 那两条行方向 GEMM)。
顺带提 sinfix 的教训:老代码曾经让 B2 里三条 GEMM 全用 K_INTER 读同一块 sin。结果 dqd/dkd 那两条本该转置读的,被悄悄读成了转置版——数值不对,S_in≠0 时污染 dq 和 gate。sinfix 就是给这两条补上 MN_INTER view。这正是 transpose-dual 用错(少了一个 view)的反面教材。
一句话总结:同一份物理字节,靠 LDSM_N/LDSM_T + K_INTER/MN_INTER 两个 cute view,一次存储同时喂正读和转置读两种 GEMM,既不存两份也不显式转置。
问题 5:fp32 C-staging tile 行填充 stride D+4,优化前后 bank 布局
场景:B1 里算完 dS 相关结果,要先落到一块 fp32 的 smem 暂存 tile(C-staging),形状 [16, 128](16 行 × 128 列,一列对应 D 的一维)。Tensor Core 的 C 输出是按列写的:一次 MMA 的 8 个线程,写的是同一列的 8 个行。
bank 复习:smem 32 个 bank,bank = (字地址) % 32,fp32 一个字 = 4B。撞 bank 看的是"同一时刻访问同一 bank 的线程数"。
优化前:stride = 128(紧凑,无填充)
元素 [r, c] 的字地址 = r128 + c。 bank = (r128 + c) % 32。因为 128 % 32 = 0,所以 r*128 % 32 = 0,于是:
bank([r,c]) = c % 32 ← 跟行号 r 完全无关!
MMA 写"同一列 c 的 8 个行"(r=0…7),它们的 bank 全是 c%32 —— 8 个线程挤同一个 bank = 8 路冲突。
列 c=0,写 8 行(r=0…7),stride=128: r=0 → addr 0 → bank 0 r=1 → addr 128 → bank 0 ← 128%32=0 r=2 → addr 256 → bank 0 r=3 → addr 384 → bank 0 r=4 → addr 512 → bank 0 r=5 → addr 640 → bank 0 r=6 → addr 768 → bank 0 r=7 → addr 896 → bank 0 ─────────────────────────── 8 个线程 → 全 bank 0 → 8 路冲突(串行 8 拍)
优化后:stride = 132(= D+4,行尾填 4 个 fp32)
元素 [r, c] 的字地址 = r*132 + c。132 % 32 = 4,于是:
bank([r,c]) = (r*4 + c) % 32 ← 现在跟 r 有关了!
同一列 c 的 8 个行,bank 变成 c, c+4, c+8, …, c+28——8 个不同 bank:
列 c=0,写 8 行(r=0…7),stride=132: r=0 → addr 0 → bank 0 r=1 → addr 132 → bank 4 ← 132%32=4 r=2 → addr 264 → bank 8 r=3 → addr 396 → bank 12 r=4 → addr 528 → bank 16 r=5 → addr 660 → bank 20 r=6 → addr 792 → bank 24 r=7 → addr 924 → bank 28 ─────────────────────────── 8 个线程 → bank {0,4,8,12,16,20,24,28} → 8 个不同 bank → conflict-free(1 拍)
为啥是 +4 而不是 +1? 只要 stride 对 32 不整除,就能把行错开,+1 也能散 bank。但这里 C-staging 的写是 cute 向量化 store(一次 float4 = 16B)。行 stride 必须是 16B 对齐:132 * 4B = 528B,528 % 16 = 0 ✅ 对齐;129 * 4B = 516B,516 % 16 = 4 ✗ 不对齐,float4 store 会崩。所以选 +4(4 个 fp32 = 16B):既打散 bank(132%32=4≠0),又保住 16B 对齐让向量化 store 能用。这是"消 bank 冲突"和"保向量化对齐"的双重约束下的最小填充。
代价:每行多 4 个 fp32 = 每 tile 多 1644B = 256B smem,微不足道,换来 8 路→无冲突。
问题 6:三个向量化优化展开
优化 3:dvcorr 窄化的向量化(B2,L522-542)
背景:dvcorr 从 gmem 读进来是 fp32,要用在两个地方:
- 累加进 dbeta(需要 fp32 精度)→ 存到 fp32 smem。
- 喂给一条 GEMM(Tensor Core 吃 bf16)→ 转成 bf16 存到 swizzled smem。
“窄化”(narrowing)就是指 fp32 → bf16 这个位宽变窄的转换。
朴素做法的两个坑:
- 每个线程只管 1 列 → gmem 读是 stride-8 散读(不连续,uncoalesced,一次 cache line 128B 只用到几个字节)。
- bf16 逐个散写到 swizzled smem → 落点乱,撞 bank。
优化做法:让每个线程负责 8 个连续列,一口气向量化处理:
线程 t 负责列 [8t, 8t+8):
① gmem 读:8 个 fp32 = 32B = 2 次 float4(LDG.128) → 连续 → 完美 coalesced
② 写 fp32 smem(给 dbeta):2 次 float4 STS.128 → 8 个 fp32 一次两拍写完
③ 窄化 fp32→bf16:8 个 fp32 → 8 个 bf16 = 16B → 打包成 1 个 uint4,一次 STS.128 写到 swizzle 的 “8-run base”(每 8 个 bf16 一组的对齐基址) → 一次写完 8 个,落在同一 swizzle 组,conflict-free
收益:gmem 从散读变连续读(coalesced);smem 写从 8 次散写变 2×float4 + 1×uint4 向量化写;指令数和 transaction 数都 ÷4 量级。核心是"一个线程管 8 个连续元素",把散的标量访问全变成 128-bit 向量访问。
优化 6(★文档没细讲):l7c 列交错 conflict-free 布局(B3,L1180-1186)
背景:B3 epilogue 里,每个线程要处理 128 列(D)中的一部分。老布局让每个线程管连续 8 列;l7c 改成每个线程管两块不连续的 float4。
老布局(连续 8 列):线程 tid,my_col = (tid%16)*8,管 [my_col, my_col+8)。 一次 float4(4 个 fp32)覆盖的 bank:
16 个线程(tid%16 = 0…15),每人 8 列: 线程0: 列 0-7 → 头 float4 覆盖 bank 0-3 线程1: 列 8-15 → 头 float4 覆盖 bank 8-11 线程2: 列16-23 → 头 float4 覆盖 bank16-19 线程3: 列24-31 → 头 float4 覆盖 bank24-27 线程4: 列32-39 → 头 float4 覆盖 bank 0-3 ← 和线程0 撞! … 这一拍 float4 只落在 {0-3, 8-11, 16-19, 24-27} = 16 个 bank → 32 个 bank 只用一半 → 要 2 个 wavefront(2 拍)
l7c 布局(列交错):每个线程拿两块不连续的 float4: c0 = (tid%16)*4,c1 = c0 + 64(即 [c0,c0+4) 和 [c0+64, c0+68))。
16 个线程,头 float4 用 c0 = (tid%16)*4: 线程0: c0=0 → bank 0-3 线程1: c0=4 → bank 4-7 线程2: c0=8 → bank 8-11 线程3: c0=12 → bank12-15 … 线程7: c0=28 → bank28-31 → 8 个线程的头 float4 = stride-4 铺满 bank 0-31 → 32 个 bank 全用上 → conflict-free(1 拍) 第二块 c1=c0+64 同理铺满另一半工作。
为啥结果不变? 这只是换了线程↔列的分工(partition),每个线程算的还是那些 列的同样计算,只是分配方式变了。计算与分区无关(partition-invariant)→ bit-identical,纯粹是访存布局优化。
一句话:老布局连续 8 列让 float4 只铺半个 bank 组(2 拍);l7c 让 8 个线程的 float4 以 stride-4 铺满全部 32 bank(1 拍),再用 c1=c0+64 拿第二块。
优化 2:l8 epilogue float4 向量化(当前 B 冠军的招牌,B3,L1101-1142)
背景:B3 epilogue 有 3 条 chunk-space 的 GEMM:
- dqd = dMqk @ ki
- dki = dMqkᵀ @ qd + β·dLᵀ @ kd
- s2 = dL @ ki
三条的收缩维(K 维)都是 16(chunk 内 16 个 token)。
老做法(标量):
- 每个线程只算 1 个输出。
- K=16 的累加循环里,每拍用标量 LDS 去 smem 取一个 operand,stride-128 散读。
- 每个 k 都重算一遍那些"跟 k 无关的 mask"(比如三角 mask、边界判断),白算 16 遍。
新做法(float4 向量化 + 提取不变量):
① 每个线程算 4 个连续 k 的贡献(一次处理 4 拍): operand 用 float4 LDS.128 一次取 4 个 fp32 → LDS transaction ÷4
② 把"跟 k 无关"的 mask 提到寄存器,每个输出列 c 只算一次: m_qd[], m_ki[], m_dl[] ← 每 c 算一遍 然后在 4 个 k-lane 里复用 → 省掉 16× 重复 mask 计算
③ 累加顺序严格保持和老代码一样 → bit-identical(全程 fp32,不碰 gate 精度红线)
为啥叫"招牌":epilogue 是 B3 最热的收尾段,3 条 GEMM 全在这。把标量 LDS 换成 float4 LDS.128(访存 transaction ÷4)+ mask 提到寄存器复用(算术指令大减),实测这一招是 split-B 从 l7→l8 的关键降幅(occupancy 日志里 l8 float4 epilogue −2~3%),而且顺序不变、全 fp32 → 数值 bit-identical,不触碰 B3 那条"epilogue 必须 fp32、否则 dg 爆 26-76×"的精度红线。
这 6 个问题就全部展开完了。串一下核心直觉:问题 1-3 是"bank 冲突到底怎么数、GPU 怎么靠多 warp 藏延迟、swizzle 怎么打散冲突"的底层机制;问题 4-5 是两个具体的 smem 布局技巧(转置对偶省 smem、D+4 填充打散列写冲突);问题 6 是三个"把标量散访存改成 128-bit 向量访存 + 复用不变量"的向量化套路。所有优化的共同母题只有一句话:让访存连续/铺满 32 bank + 一次搬 128 bit + 别重复算不变的东西,同时死守 B3 epilogue 的 fp32 精度红线。
问题一和问题六的float4怎么就把conflict消灭了？我理解只有l7c 布局(列交错):每个线程拿两块不连续的 float4这样才行。另外问题六优化我还是不理解
● 好问题,你抓到了我前面讲得不精确的地方。你的判断是对的:float4 本身并不消灭 bank 冲突。 让我先把两件被我混在一起的事拆开,这是理解 Q6 的钥匙。
关键:向量化 和 消 bank 冲突 是两件正交的事
┌───────────────┬──────────────────┬─────────────────┬────────────┐ │ │ 是什么 │ 靠什么实现 │ 收益 │ ├───────────────┼──────────────────┼─────────────────┼────────────┤ │ 向量化(float4 │ 一条指令搬 128 b │ │ transactio │ │ ) │ it(4×fp32)而不是 │ LDS.128/STS.128 │ n 数量 ÷4 │ │ │ 32 bit │ │ │ ├───────────────┼──────────────────┼─────────────────┼────────────┤ │ │ 同一次 │ 地址布局(swizzl │ 每次 trans │ │ 消 bank 冲突 │ transaction 里 │ e / 列交错 / │ action 从 │ │ │ 32 个 bank 别撞 │ D+4 填充) │ N 拍变 1 │ │ │ │ │ 拍 │ └───────────────┴──────────────────┴─────────────────┴────────────┘
float4 只管"一次搬多少 bit",不管"搬的地址撞不撞 bank"。 一个 float4 访问照样能撞 bank,一个标量访问照样能不撞。你说的完全对——只有 l7c 那种列交错布局才是真正在消冲突,float4 只是顺便把次数降下来。
为什么会让人误以为 float4 消了冲突? 因为硬件处理一次 warp 级 LDS.128 时,32 个线程各要 16B,共 512B,会被拆成 4 次 transaction(每次 128B = 32 bank × 4B,服务 8 个线程)。在这 8 个线程这一组内,它们的 float4 是否铺满 32 个 bank,才决定这次 transaction 是 1 拍还是多拍。所以:
▎ float4 把"要多少次 transaction"定死了(÷4),但每次 transaction 撞不撞,还得看那 8 个线程的地址布局。布局才是消冲突的主角,float4 是配角。
重讲 Q6 三个优化(这次标清楚:哪部分省次数、哪部分消冲突)
优化 3:dvcorr 窄化向量化
这里根本没有 bank 冲突问题要解决——因为一半是 gmem。拆清楚三步各自在干嘛:
① gmem 读 8 个 fp32: 问题 = coalescing(gmem 没有 bank,只有 128B cache line) 散读 → 一条 cache line 只用几个字节 → 浪费带宽 float4 连续读 → coalesced → 带宽用满 【这步靠 float4:省的是 gmem 带宽,不是 bank】
② 写 fp32 smem(给 dbeta):2×float4 STS.128 【这步靠 float4:省 transaction 数,布局本身已对齐,不撞】
③ 窄化 fp32→bf16,打包成 uint4 写 swizzle smem: 不撞 bank 是因为写到"8-run base"(swizzle 布局) float4 只是让这 8 个 bf16 一次写完 【消冲突靠 swizzle 布局,不是靠 uint4】
所以 dvcorr 的收益 = gmem coalesced(①)+ 写次数 ÷4(②③),冲突本来就靠既有 swizzle 布局解决了。这里 float4 的功劳是省次数/省带宽,不是消冲突。 我前面说的"conflict-free"是布局给的,不是 float4 给的。
优化 6:l7c 列交错 —— 这个才是真·消冲突
这个优化唯一目的就是消 bank 冲突,float4 在这里反而是"被布局服务的对象"。
回顾一下:硬件处理 float4 访问时,一组 8 个线程 = 一次 transaction。这 8 个线程的 float4 要铺满 32 个 bank 才 1 拍。
老布局(连续 8 列),线程 t 的头 float4 起点 = 8t: 8 个线程头 float4 起点: 0, 8, 16, 24, 32, 40, 48, 56 对 32 取模(bank 组): 0, 8, 16, 24, 0, 8, 16, 24 ↑ ↑ ↑ ↑ 线程4/5/6/7 和 线程0/1/2/3 撞! 一次 transaction 只铺到 {bank0-3, 8-11, 16-19, 24-27} = 半个 smem → 2 路冲突 → 2 拍
l7c(列交错),线程 t 的头 float4 起点 = 4t(c0=(tid%16)*4): 8 个线程头 float4 起点: 0, 4, 8, 12, 16, 20, 24, 28 每个 float4 占 4 个连续 bank: 线程0 → bank 0-3 线程1 → bank 4-7 … 线程7 → bank 28-31 → 8 个 float4 = stride-4 严丝合缝铺满 32 bank → 1 拍 conflict-free 第二块 c1=c0+64 处理剩下一半的列,同理铺满。
看出来了吗: 老布局和 l7c 用的都是 float4(次数一样),差别只在起点 stride——8 vs 4。stride-4 让 8 个线程的 float4 拼成完整 32 bank,stride-8 让它们叠成半个。这就证明了:消冲突是布局(stride-4 交错)干的,和 float4 无关。
而且因为只是换了"哪个线程算哪列"的分工,计算内容一字不差 → bit-identical。这是纯访存布局优化。
优化 2:l8 epilogue float4 —— 主要省次数,不是消冲突
这个优化的名字带 float4,但它的收益几乎全在"省 transaction 数 + 省重复计算",不在消冲突(operand 的 smem 布局本来就是好的)。三条 GEMM 收缩维 K=16:
老做法(标量): 每个线程算 1 个输出,K=16 循环 每拍标量 LDS 取 1 个 operand → 16 次 LDS 每个 k 都重算一遍跟 k 无关的 mask → 白算 16 遍
新做法: ① 每线程一次处理 4 个连续 k: float4 LDS.128 一次取 4 个 operand → LDS 次数 16 → 4(÷4) 【省次数,不是消冲突】 ② 跟 k 无关的 mask(三角/边界)提到寄存器,每列 c 只算 1 次, 4 个 k-lane 复用: → 省掉 16× 重复算术 【省算术指令,和访存无关】 ③ 累加顺序不变、全 fp32 → bit-identical,不碰 gate 精度红线
为什么是"招牌"? 不是因为它消了什么冲突,而是因为 epilogue 是 B3 最热的收尾段(3 条 GEMM 全在这),把"16 次标量 LDS + 16 遍重复 mask"压成"4 次 float4 LDS + 1 遍 mask 复用",指令数直接砍一大截,实测 l7→l8 降 2~3%。它的价值是吞吐(少发指令),不是延迟(消冲突)。
一句话收尾(纠正版)
- float4 = 省 transaction 次数(和省 gmem 带宽),从不单独消 bank 冲突。
- 消 bank 冲突 = 地址布局的事:l7c 的 stride-4 列交错、D+4 行填充、swizzle,才是消冲突的主角。
- Q6 三个里:dvcorr 靠 float4 省带宽/次数(冲突交给既有 swizzle);l8 epilogue 靠 float4 省次数 + 寄存器复用 mask 省算术;只有 l7c 是真正在改布局消冲突——你的理解完全正确,是我前面把"float4 省次数"和"布局消冲突"糊在一起了。
✻ Baked for 1m 35s
❯ 窄化 fp32→bf16:8 个 fp32 → 8 个 bf16 = 16B → 打包成 1 个 uint4,一次 STS.128 写到 swizzle 的 “8-run base”(每 8 个 bf16 一组的对齐基址) → 一次写完 8 个,落在同一 swizzle 组,conflict-free 是啥意思，为啥牵扯到uint4
● 先拆开"uint4"到底是什么
uint4 不是什么高深东西,就是一个 128 位(16 字节)的寄存器容器 = 4 个 uint32 拼在一起。它存在的唯一意义:让编译器发一条 STS.128(一次写 128 bit 到 smem)。
对照一下位宽:
bf16 = 2 字节 8 个 bf16 = 16 字节 = 128 bit = 正好 1 个 uint4
所以"8 个 fp32 → 窄化成 8 个 bf16 → 塞进 1 个 uint4"这句话的意思就是:
▎ 这个线程手上算出的 8 个 bf16 结果,拼成一个 16 字节的寄存器包(uint4),然后一条指令(STS.128)整包写进 smem —— 而不是发 8 条 STS.16(每次写 1 个 bf16)。
收益之一(省次数): 8 条标量存 → 1 条向量存,指令数 ÷8。
"8-run base / 落在同一 swizzle 组 / conflict-free"是什么意思
这牵扯到 swizzle 的粒度。这块 bf16 operand smem 用的是 cute 的 K_INTER swizzle 布局,它的 XOR 地址置换是以 16 字节(128 bit)为一个"原子格"来定义的。也就是说:
- swizzle 把地址按 每 16B 一格 打散到不同 bank 区。
- 每一个对齐的 16B 格,整体落在一个 conflict-free 的目标位置(一组 bank 里),这是 swizzle 设计时就保证好的(为的是后面 LDSM/Tensor Core 读它时不撞)。
"8-run base"就是指:每 8 个 bf16(=16B=一个 swizzle 原子格)的对齐起始地址。你手上这 8 个 bf16 恰好构成一个完整的原子格。
为什么必须用 uint4(而不是逐个写)?
情况 A:逐个写 8 条 STS.16 每个 2 字节地址各自过一遍 swizzle XOR → 8 个地址可能被打散到不同 bank 区,拆成多拍 → 而且和别的线程的碎片交错,容易撞 bank
情况 B:打包成 uint4,一条 STS.128 写整个 16B 格 硬件把这 16B 当成"一个对齐的 128-bit 事务", 整包丢进 swizzle 给它安排好的那一格 → 天然落在同一个 swizzle 组里,不跨格、不撞 bank → 1 拍写完
换句话说:swizzle 是按 16B 一格设计的,你就得按 16B 一格去写,才能踩在它的节奏上。 uint4 = 正好一格,STS.128 = 正好一次写一格 —— 对齐了 swizzle 的粒度,所以既是 1 条指令,又天然 conflict-free。如果你偏要按 2B(逐 bf16)去写,就跨在格子的碎片上,swizzle 的对齐前提被打破,可能撞 bank。
用一张字节图看清楚
假设某个 8-run 格的基址 = 字节 1024(16B 对齐):
逻辑上这 8 个 bf16 要写到: 字节 1024,1026,1028,1030,1032,1034,1036,1038 (共 16 字节)
uint4 一次 STS.128: [1024, 1040) 这整块 16B 一次性写入 ↓ 对 smem bank:1024/4=256, 256%32=0 → 覆盖 bank 0,1,2,3(4 个字 = 8 个 bf16) 这个线程占 4 个 bank;同组另外 7 个线程各占相邻 4 个 bank → 8 线程 × 4 bank = 32 bank 铺满 → 1 拍 conflict-free
对比逐个 STS.16:8 条指令,每条 2 字节,swizzle 对 2B 粒度的置换没保证铺满,反而可能两个碎片撞进同一 bank。
一句话
-
uint4 出现,纯粹是因为 8 个 bf16 = 16B = 128 bit,是能触发一条 STS.128 的最小容器 → 省指令。
-
conflict-free 不是 uint4 给的,是 swizzle 按 16B 一格设计给的;uint4 的作用是让你"按整格(16B)对齐地写",踩准 swizzle 的粒度,从而既省次数又不撞 bank。
-
反过来,这也再次印证上一条结论:消冲突是布局(swizzle)的功劳,uint4/STS. 128 只是把写次数压到 1 条并保证按格对齐。两件事配合,缺一不可——布局定" 撞不撞",向量宽度定"几条指令"。
先看这两行在干什么(一句话)
using SwStateLayout =
decltype(tile_to_shape(GMMA::Layout_K_INTER_Atom
它们各自定义了一个"地址计算公式"的类型,名字叫 SwStateLayout 和 SwStateLayoutT。这两个公式描述:一块 [D,D]=[128,128] 的 bf16 数据在 smem 里,逻辑坐标 (行,列) 应该落到哪个物理字节。区别是:一个给"正读"用, 一个给"转置读"用,但指向同一块物理字节(这就是 Q4 说的转置对偶)。
下面逐词拆。
- Layout 是什么 —— cute 的核心概念
在 cute 里,Layout = 一个从"逻辑坐标"到"物理偏移"的函数。你可以把它想成一个纯函数:
layout(行, 列) → 这个元素在内存里的第几个字节(其实是第几个元素)
比如最普通的行主序 [128,128],stride=(128,1),那 layout(r,c) = r*128 + c。这就是一个 Layout。cute 把"形状 shape"和"步长 stride"打包成一个 Layout 对象,layout(r,c) 就自动帮你算偏移。
为什么要抽象成对象? 因为后面 swizzle(bank 打散)会让这个映射不再是简单的 r*128+c,而是一堆 XOR 位运算。你不想手写那些位运算,就让 cute 生成好一个 Layout,你只管 sdS(k, d0)(见源码 L218-222),它内部替你算出 swizzle 后的真实地址。所以:
▎ Layout = 帮你把"我要第 (r,c) 个元素"翻译成"它在 smem 的真实字节位置"的翻译器,swizzle 的复杂位运算全藏在里面。
- decltype(…) 和 using … = —— 纯编译期,零运行成本
- using SwStateLayout = 某类型; 就是 C++ 的类型别名(给类型起个短名字),和 typedef 一样。
- decltype(表达式) = “取这个表达式的类型”,不真正执行它。
cute 的 Layout 全是编译期常量类型(shape/stride 都是 Int<128> 这种编译期整数)。所以 tile_to_shape(…) 这个调用在编译时就被算成了一个具体的 Layout 类型,decltype 把那个类型抠出来,using 给它起名 SwStateLayout。
关键:运行时 GPU 上不会执行任何"构造 layout"的代码——地址公式在编译时就 定死了,跑起来只有几条位运算算偏移。这就是我一直说的"cute 零运行成本":全是编译期模板元编程。
- GMMA::Layout_K_INTER_Atom
{} —— 一小块"swizzle 样板"
这是一个**"原子"布局(atom)**,是 swizzle 的最小重复单元。拆词:
- GMMA = Hopper Tensor Core(GMMA 指令)配套的布局工具包。
- Layout_K_INTER_Atom = 一个为 Tensor Core 正读优化过的、conflict-free
的小块 swizzle 模板。
- K_INTER:沿 K 方向(收缩维)内层交错(INTER-leave)。回忆 Q4:正读(LDSM_N)友好。
- MN_INTER(第二行):沿 MN 方向内层交错,转置读(LDSM_T)友好。
= 元素是 bf16(2 字节)。swizzle 的粒度和位宽有关,所以要告诉它数据类型。 - 末尾 {} = 构造一个该类型的实例(编译期空对象,只为把类型传进 tile_to_shape)。
这个 atom 有多大? 它只是一小块(比如 8×某 的样板),描述"在这么小的范围内,元素怎么摆才能让 Tensor Core 读的时候不撞 bank"。它不是整块 [128,128]——整块要靠下一步 tile_to_shape 把这个小样板平铺开。
▎ 类比:atom 是一块"防冲突瓷砖"的花纹,tile_to_shape 把这块瓷砖铺满整个 [128,128] 地板。
- make_shape(Int
{}, Int {}) —— 目标整块的形状
- make_shape(a, b) = 造一个形状 (a, b)。
- Int
{} = 编译期整数 128(D=128)。用 Int<128> 而不是普通 int 128,是为了让形状在编译期就是常量(前面说的零运行成本前提)。
所以这是在说:我要的整块是 [128, 128]。
- tile_to_shape(atom, shape, order) —— 把小瓷砖铺满整块
tile_to_shape 就是那个"铺瓷砖"的函数:
tile_to_shape( 防冲突瓷砖花纹(atom), 要铺满的地板形状[128,128], 铺的方向 ) → 一个覆盖整块 [128,128] 的、带 swizzle 的完整 Layout
它保证:整块里每一个 16B 小格都按 atom 的防冲突花纹摆放,拼起来覆盖完整的 [128,128]。生成的这个 Layout 你 sdS(k,d0) 一调用,它就吐出 swizzle 后的真实字节偏移。
- LayoutLeft{} vs LayoutRight{} —— 铺的方向(列主序 vs 行主序)
这是两行最关键的区别,也是"转置对偶"落地的地方:
- LayoutLeft{} = 列主序(左边的维度变得最快,像 Fortran)。第一行用它。
- LayoutRight{} = 行主序(右边的维度变得最快,像 C)。第二行用它。
配合 atom 方向:
第一行 SwStateLayout = K_INTER atom + LayoutLeft → 正读友好(LDSM_N) 第二行 SwStateLayoutT = MN_INTER atom + LayoutRight → 转置读友好(LDSM_T)
为什么同一块物理字节能被两个 Layout 描述? 因为 INTER 系列 atom 的巧妙之处:K_INTER + LayoutLeft 和 MN_INTER + LayoutRight 这一对,对同一堆物理字节,给出的是互为转置的逻辑视图。也就是:
- 你只写一次数据(用 SwStateLayout 写,源码 L218-222)。
- dkr 那条 GEMM 用 SwStateLayout 正读(L316)。
- G2 那条 GEMM 用 SwStateLayoutT 转置读(L339 dST_t … SwStateLayoutT{})。
两个 Layout、一份字节 → 省一份 smem、省一次显式转置。这就是 Q4 讲的转置对偶,这两行就是它的声明。
- 最后那行 static_assert
static_assert(cute::cosize_v
- cosize = 这个 Layout 实际占用的物理元素个数(最大偏移+1)。
- DD = DD = 128128 = 16384。
- 这行在编译期检查:swizzle 之后没有多占空间(cosize == DD,没 padding)。注释 L174 说"128-wide INTER 不需要 padding"——如果哪天改坏了导致 cosize≠DD,编译直接报错,防止悄悄多吃 smem 把 occupancy 拖垮。
串起来:两行的完整白话翻译
// 第一行:给 [128,128] 的 dS 造一个"正读友好、conflict-free"的 smem 地址公式, // 用 K_INTER 防冲突瓷砖铺满,列主序 → 取名 SwStateLayout(给 dkr 正读) using SwStateLayout = …K_INTER… LayoutLeft;
// 第二行:给同一块 [128,128] 造一个"转置读友好"的地址公式, // 用 MN_INTER 瓷砖铺满,行主序 → 取名 SwStateLayoutT(给 G2 转置读) // 它和第一行指向同一堆物理字节,互为转置视图 using SwStateLayoutT = …MN_INTER… LayoutRight;
一句话总结: 这两行是纯编译期的"地址翻译器"定义。Layout_K_INTER_Atom/MN_INTER_Atom 是防 bank 冲突的 swizzle 瓷砖花纹,tile_to_shape 把它铺满 [128,128],LayoutLeft/LayoutRight 决定正读还是转置读。二者共享同一份物理字节 → 实现 Q4 的转置对偶(存一份、喂两种 GEMM 方向)。你调 sdS(k,d0) 时,swizzle 的位运算全在里面替你算好,运行时零额外开销。
一、先看两个 atom 的"本体"(8×8 小样板)
print_layout 打出的表格 = 逻辑坐标 (行,列) → 物理偏移(第几个 bf16)。
Layout_K_INTER_Atom
col: 0 1 2 3 4 5 6 7
row 0 → 0 1 2 3 4 5 6 7 row 1 → 8 9 10 11 12 13 14 15 row 2 → 16 17 18 … … row 7 → 56 57 … 63 偏移公式:offset(r,c) = 8*r + c。这就是普普通通的行主序 8×8 块。
Layout_MN_INTER_Atom
col: 0 1 2 3 4 5 6 7
row 0 → 0 8 16 24 32 40 48 56 row 1 → 1 9 17 25 … … row 7 → 7 15 23 … 63 偏移公式:offset(r,c) = r + 8*c。这是列主序 8×8 块 = 第一个的精确转置。
▎ ⚠️ 重要纠正: 这两个 INTER atom 根本没有 XOR swizzle(源码里是 Swizzle<0,4,3>,第一个参数 B=0 = 零位交换)。它们就是朴素的行主序 / 列主序 8×8 瓷砖,互为转置。我前面几轮说"swizzle 用 XOR 打散 bank"——那是 SW32/SW64/SW128 变体才有的事,INTER 没有。下面第三节给你看真正的 XOR 长啥样。
二、铺满 [16,16] 后:转置对偶的"眼见为实"
tile_to_shape 把 8×8 瓷砖平铺成 [16,16](真 kernel 是 [128,128],规律一样,只是大 8 倍)。看两张表的同一个坐标:
SwStateLayout(K_INTER + LayoutLeft) —— 左上角 8×8 块存偏移 0–63,行主序: (0,0)=0 (0,8)=128 ← 列块往右跳到 128 (8,0)=64 (8,8)=192 ← 行块往下跳到 64
SwStateLayoutT(MN_INTER + LayoutRight) —— 同样的物理槽位,但转置视图: (0,0)=0 (0,8)=64 (8,0)=128 (8,8)=192
程序最后那行断言把它证死了: transpose-dual check SwStateT(i,j) == SwState(j,i) : TRUE cosize(SwState)=256 cosize(SwStateT)=256 D*D=256
含义:
- SwStateT(i,j) 和 SwState(j,i) 永远指向同一个物理偏移 → 两个 Layout 是同一堆 256 个物理字节的"正读视图"和"转置视图"。这就是 Q4 转置对偶的本质,现在你亲眼看到映射表相等。
- cosize = 256 = 16×16 → 没有 padding,转置视图没多占一个字节。这正是源码 L179 那句 static_assert(cosize == DD) 在编译期守的东西。
所以那两行 using 的作用就是:给同一块 dS smem 造两个地址翻译器,一个正读(喂 dkr GEMM)、一个转置读(喂 G2 GEMM),存一份、喂两向。
三、对照:真正带 XOR 的 swizzle(SW128)长什么样
为了让你看清"INTER 无 XOR"和"SW 有 XOR"的差别,我又打印了 Layout_K_SW128_Atom,头部是 Sw<3,3,3>(真 XOR):
row 0: 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 | … ← 原样 row 1: 72 73 74 75 76 77 78 79 | 64 65 66 67 68 69 70 71 | … ← 每8个一组被 XOR 打乱顺序! row 2:144 … 152 … 128 … ← 又是另一种打乱 row 3:216 217 … 208 209 …
看 row 1:本该是 64 65 66 … 71 | 72 … 79,却变成 72…79 | 64…71 —— 相邻两个 8-组被 XOR 换了位。row 2、row 3 的打乱模式又各不相同。这才是 XOR swizzle:用行号去异或列地址,把本会撞同一 bank 的元素错开。
为什么 dS/sin 用 INTER 而不用这个 SW? 因为 dS/sin 是 128 宽(D=128)的 K-major 收缩,在这个宽度下 INTER 的 8-元素交错平铺本身就已经 conflict-free(LDSM 的 8 线程各取一个 16B 行片,天然落在不同 bank),不需要再加 XOR。只有更窄或不同访问模式才需要 SW 的 XOR 去补救。源码注释 L174 那句 “no padding for 128-wide INTER” 就是这个意思。
四、一句话总结
- SwStateLayout / SwStateLayoutT 的本体 = 一对 8×8 行主序 / 列主序瓷砖,tile_to_shape 铺满 [128,128]。
- 二者是同一物理字节的转置对偶(程序验证 T(i,j)==(j,i),cosize 都=256 无 padding)—— 这就是 Q4"存一份喂两向"的实现。
- 纠正:INTER atom 没有 XOR swizzle(Swizzle<0>),conflict-free 靠的是 8-元素交错平铺;真正的 XOR 在 SW128(Sw<3,3,3>)那种变体里,本 kernel 的 dS/sin 用不到。
● 源码全在 cutlass/include/cute/atom/mma_traits_sm90_gmma.hpp。我已经读到了,直接给你 贴出来并逐行解析。
一、原始定义(bit 单位,L74-84)
// M|N-major GMMA layouts in units of bits using Layout_MN_INTER_Atom_Bits = ComposedLayout<Swizzle<0,4,3>, smem_ptr_flag, Layout<Shape< _128,_8>,Stride<_1, _128>>>; using Layout_MN_SW32_Atom_Bits = ComposedLayout<Swizzle<1,4,3>, smem_ptr_flag, Layout<Shape< _256,_8>,Stride<_1, _256>>>; using Layout_MN_SW64_Atom_Bits = ComposedLayout<Swizzle<2,4,3>, smem_ptr_flag, Layout<Shape< _512,_8>,Stride<_1, _512>>>; using Layout_MN_SW128_Atom_Bits = ComposedLayout<Swizzle<3,4,3>, smem_ptr_flag, Layout<Shape<_1024,_8>,Stride<_1,_1024>>>;
// K-major GMMA layouts in units of bits using Layout_K_INTER_Atom_Bits = ComposedLayout<Swizzle<0,4,3>, smem_ptr_flag, Layout<Shape<_8, _128>,Stride< _128,_1>>>; using Layout_K_SW32_Atom_Bits = ComposedLayout<Swizzle<1,4,3>, smem_ptr_flag, Layout<Shape<_8, _256>,Stride< _256,_1>>>; using Layout_K_SW64_Atom_Bits = ComposedLayout<Swizzle<2,4,3>, smem_ptr_flag, Layout<Shape<_8, _512>,Stride< _512,_1>>>; using Layout_K_SW128_Atom_Bits = ComposedLayout<Swizzle<3,4,3>, smem_ptr_flag, Layout<Shape<_8,_1024>,Stride<_1024,_1>>>;
二、按元素类型特化(upcast,L86-104)
// M|N-major layouts in units of Type
template
// K-major layouts in units of Type
template
逐部分解析
- ComposedLayout<Swizzle<…>, smem_ptr_flag, Layout<…>>
ComposedLayout<F, offset, G> = 函数复合:先用底层 Layout G 把坐标 → 偏移,再用 Swizzle F 对偏移做 XOR 变换。即 最终偏移 = F(G(坐标))。中间的 smem_ptr_flag(其实是 _0 + 一个标记)只是个零偏移占位,标注"这是 smem 指针"。
- Swizzle<B, M, S> —— XOR 的三个参数
这是真正决定"打不打散、怎么打散"的东西。定义在 cute/swizzle.hpp,语义:
- B(NumBits) = 参与 XOR 的位数(打散多少条"通道")。
- M(MBase) = 被打散的最小单元大小 = 2^M 个元素。这里 M=4 → 2^4=16(对应 bit 单位下 16 bit = 一个 bf16?不,是 base 粒度)。
- S(SShift) = XOR 时的移位量。
XOR 规则大致是:swizzled = addr ^ ((addr >> S 的某几位) << M)。关键就看 B:
┌───────┬────────────────┬──────┬──────────────────────────────────────────┐ │ atom │ Swizzle<B,4,3> │ B 值 │ 效果 │ ├───────┼────────────────┼──────┼──────────────────────────────────────────┤ │ INTER │ Swizzle<0,4,3> │ B=0 │ 零位 XOR = 恒等,不打散(就是纯行/列主序) │ ├───────┼────────────────┼──────┼──────────────────────────────────────────┤ │ SW32 │ Swizzle<1,4,3> │ B=1 │ XOR 1 位 → 2 路打散,swizzle 周期 32 字节 │ ├───────┼────────────────┼──────┼──────────────────────────────────────────┤ │ SW64 │ Swizzle<2,4,3> │ B=2 │ XOR 2 位 → 4 路打散,周期 64 字节 │ ├───────┼────────────────┼──────┼──────────────────────────────────────────┤ │ SW128 │ Swizzle<3,4,3> │ B=3 │ XOR 3 位 → 8 路打散,周期 128 字节 │ └───────┴────────────────┴──────┴──────────────────────────────────────────┘
这就印证了上一轮的结论:INTER 的 Swizzle<0> = 没有 XOR;SW32/64/128 的 B=1/2/3 才是真打散,数字 32/64/128 = swizzle 的周期字节数(= 2^(B) * 2^M * … ,本质是打散跨度)。
- 底层 Layout<Shape<…>, Stride<…>> —— K-major vs MN-major 的区别
看 bit 单位的 shape/stride,**M/K 谁在内层(stride=1)**决定 major:
- K 系列:Shape<_8, _128>, Stride<_128, _1> → 第二维(K,128 bit=8 bf16)stride=1 是内层连续 → K-major(收缩维连续,正读 LDSM_N 友好)。
- MN 系列:Shape<_128, _8>, Stride<_1, _128> → 第一维(MN)stride=1 内层连续 → MN-major(转置读 LDSM_T 友好)。
两者互为转置——正是转置对偶的根。
SW32→SW64→SW128 逐级放大:第二维从 _128→_256→_512→_1024 bit,swizzle 周期也 1×→2×→4×→8× 增大,配合更大的 tile 做更宽的无冲突。
- upcast<sizeof_bits
::value>(…) —— 从 bit 折算到元素
原始定义是按 bit 的(通用于任何精度)。upcast
- Layout_K_INTER_Atom_Bits 的 Shape<_8,_128>(bit)→ 除以 16 → Shape<_8,_8>(bf16)。
- 这正是上一轮 print_layout 打出的 (_8,_8)😦_8,_1)。
对 fp8(8 bit)会 upcast<8> → Shape<_8,_16>,更宽;对 fp32(32 bit)→
Shape<_8,_4>。同一套 bit 定义,upcast
按精度自动折算成元素粒度,这就是为什么模板参数是
想看 SW 参数怎么实现的,再往下挖
Swizzle 模板本体和它的 XOR operator() 在:
● Searched for 1 pattern, read 1 file (ctrl+o to expand)
● 这就是 XOR 的真身。Swizzle<B,M,S> 的核心(cute/swizzle.hpp L54-79):
template <int BBits, int MBase, int SShift = BBits> struct Swizzle { static constexpr int num_bits = BBits; // B:参与 XOR 的位数 static constexpr int num_base = MBase; // M:最小不动单元 = 2^M 元素 static constexpr int num_shft = SShift; // S:移位量
using bit_msk = cute::constant<int, (1 << num_bits) - 1>; //
(1<<B)-1 using yyy_msk = cute::constant<int, bit_msk{} << (num_base + max(0,num_shft))>; using zzz_msk = cute::constant<int, bit_msk{} << (num_base - min(0,num_shft))>; using msk_sft = cute::constant<int, num_shft>;
// 核心:把地址高段的 B 位 XOR 到低段去
static apply(Offset const& offset) {
return offset ^ shiftr(offset & yyy_msk{}, msk_sft{}); // ZZZ ^= YYY
}
operator()(offset) const { return apply(offset); }
};
逐字解析 apply
一个物理偏移 offset(bit 单位),按位分成三段:ZZZ | YYY | 底部 M 位。
- yyy_msk 圈出高段 YYY(要拿去 XOR 的那 B 位)。
- offset & yyy_msk 取出 YYY;shiftr(…, S) 右移 S 位,对齐到低段 ZZZ 的位置。
- offset ^ (…) 把 YYY XOR 进 ZZZ。
直觉: 用地址的高位段(≈ 行号)去异或低位段(≈ 列号),从而"行号越大,列被翻转得越多" → 把本会落同一 bank 的元素错位。这正是上一轮 SW128 打印里看到的现象: row 0: 0 1 2 3 4 5 6 7 (行号高位=0 → 不翻) row 1: 72…79 | 64…71 (行号高位=1 → 相邻 8-组对调,B=1 bit 翻转)
为什么 B=0(INTER)= 恒等,不打散
看 bit_msk = (1<<num_bits)-1:
- B=0 → bit_msk = (1<<0)-1 = 0 → yyy_msk = 0 → apply 变成 offset ^ shiftr(offset & 0, S) = offset ^ 0 = offset。原样返回,零变换。 这就是数学上证明的"INTER 没有 XOR"。
- B=1(SW32)→ bit_msk=1,翻 1 位 → 2 路对调。
- B=2(SW64)→ bit_msk=3,翻 2 位 → 4 路重排。
- B=3(SW128)→ bit_msk=7,翻 3 位 → 8 路重排(SW128 打印里 8-组内 8 种模式)。
static_assert(abs(S) >= B) 保证移位量不小于翻转位数,YYY 段和 ZZZ 段不会重叠自撞。
完整链条串起来
Layout_K_INTER_Atom
Layout_K_SW128_Atom
- Swizzle<B,M,S>(swizzle.hpp:54)= XOR 变换,B 定打不打散(0=不打)。
- Layout<Shape,Stride>(bit 定义,mma_traits_sm90_gmma.hpp:74-84)= 底层行列摆放,K/MN 决定 major(谁 stride=1)。
- ComposedLayout = 先底层布局定位、再 Swizzle XOR。
- upcast<sizeof_bits
>(:88-104)= 把 bit 定义按精度折算成元素粒度。
这四层拼出你在 champion kernel 里看到的 Layout_K_INTER_Atom
gemm vs cooperative_gemm — 从源码看本质区别
两者根本区别就一句话:操作数在哪里,谁负责搬运。
- 线程级 gemm(mma, D, A, B, C) — 只做一条 MMA 指令
gemm.hpp:178-198,最底层的 dispatch:
gemm(MMA_Atom
关键:模板约束 is_rmem
- 数据怎么从 smem 进寄存器(你得自己 copy/ldmatrix)
- K 维怎么循环(你自己写外层 for)
- 累加器怎么初始化 / 结果怎么写回 smem(你自己 epilogue)
所以你在 bwd_b_split.cuh:333-343 看到的手动 stcA/stcB/stcBT(LDSM_N/LDSM_T)+ gemm(thr_mma, …) 就是在手工喂这个裸 MMA。
- cooperative_gemm(thread_idx, tiled_mma, α, sA, sB, β, sC) — 整条流水线
cooperative_gemm.hpp:371,操作数是 smem tensor(sA/sB/sC)。它内部(fast path cooperative_gemm_no_predication, L281-356)把三件事全包了:
// L304-307: 自己开寄存器 fragment Tensor tCrA = thr_mma.partition_fragment_A(sA); Tensor tCrB = thr_mma.partition_fragment_B(sB); // L312-330: 用 SmemCopyOpA/B(默认 DefaultCopy→自动挑 ldmatrix)做 smem→reg auto smem_tiled_copy_A = make_tiled_copy_A(…); copy(smem_tiled_copy_A, tCsA(,,0), …); // PREFETCH k=0 // L337-355: 整个 K 循环 + 每步 gemm for (k_block …) { copy(… k_next …); // 预取下一 k cute::transform(tCrAi, tCrA, sA_load_op); // 类型转换 gemm(thr_mma, tCrA(,,k_block), tCrB(,,k_block), tCrC); // ← 调用第1种! } // 然后 epilogue_no_predication: axpby + 写回 sC
注意 L354:cooperative_gemm 内部就是循环调用那个线程级 gemm。它是线程级 gemm 的"全自动外壳":自动 ldmatrix + 自动 K 循环 + 自动 clear 累加器(L411 clear(tCrC))+ 自动 epilogue 写回 smem。整块线程通过 thread_idx 协同(get_thread_slice)。
何时用哪个
┌──────────────────────┬────────────────────┬──────────────────────────────────┐ │ 维度 │ cooperative_gemm │ 线程级 gemm │ ├──────────────────────┼────────────────────┼──────────────────────────────────┤ │ 操作数 │ smem(sA,sB,sC) │ 寄存器 fragment │ ├──────────────────────┼────────────────────┼──────────────────────────────────┤ │ 搬运 │ 自动(内含 ldmatrix │ 手动 │ │ │ + copy) │ │ ├──────────────────────┼────────────────────┼──────────────────────────────────┤ │ K │ 全包 │ 全自己写 │ │ 循环/累加器/epilogue │ │ │ ├──────────────────────┼────────────────────┼──────────────────────────────────┤ │ 代码量 │ 一行 │ 一堆 copy + for │ ├──────────────────────┼────────────────────┼──────────────────────────────────┤ │ │ │ 要融合的链:累加器常驻寄存器跨多 │ │ 适用 │ 孤立的、算完就写回 │ GEMM、共享同一个 │ │ │ smem 的 GEMM │ B-load、transpose-dual、自定义 │ │ │ │ epilogue │ └──────────────────────┴────────────────────┴──────────────────────────────────┘
这正好对应 bwd_b_split.cuh:dkr 用 cooperative_gemm(L294,注释"独立、写到不相干的 smem buffer"→省事);dU 链用线程级 gemm(L333+,累加器 dacc 跨 G1→G2 常驻寄存器、B-load 复用、LDSM_T 转置对偶)——这些融合技巧 cooperative_gemm 的黑盒做不到,因为它每次都强制 reg→smem 写回、每个 GEMM 独立开累加器。
一句话:孤立 GEMM 用 cooperative_gemm 省事;要跨 GEMM 融合(寄存器常驻/共享 load/转置对偶/定制 epilogue)必须降到线程级 gemm 自己手搓。
cute::transform 是干啥用的
tensor_algorithms.hpp:80-114,它就是 CuTe 版的 std::transform —— 对一个 tensor 逐元素套用一个函数:
// 一元:原地或写到另一个 tensor transform(Tensor& t, UnaryOp op) { for i: t(i) = op(t(i)); } // L84 transform(Tensor in, Tensor& out, op) { for i: out(i) = op(in(i)); } // L107 // 二元:两个输入一个输出 transform(in1, in2, out, BinaryOp op) { for i: out(i)=op(in1(i),in2(i));} // L140
就是个 CUTE_UNROLL 的逐元素循环,没有任何魔法。在 GEMM 里它专门干"喂给 tensor-core 前的逐元素变换",最常见就是类型转换:
- cooperative_gemm.hpp:350-351:cute::transform(tCrAi, tCrA, sA_load_op) —— 把 smem 里的原始类型(比如 bf16)按 sA_load_op 转成 MMA 要的计算类型。注释明说"relying on the compiler to remove in case of identity ops":如果不需要转换(identity),编译器直接删掉,零开销。
- cooperative_gemm.hpp:145,150:epilogue 里对 C 做 load/store 变换(sC_load_op/sC_store_op)。
在你的 bwd_b_split.cuh 里,线程级 gemm 的融合链手写 cute::transform(…) 就是干同一件事:把 fp32 的累加器逐元素 cast 成 bf16,再当作下一个 GEMM 的 B 操作数喂进去(tensor-core 吃 bf16 操作数)。因为线程级 gemm 不带这个自动转换,所以要手动补一句 transform 做 dtype 收窄。
简言之:cute::transform = 逐元素 map,在 GEMM 上下文里 = 在 smem→寄存器之后、MMA 之前(或结果写回之前)做类型转换/缩放等逐元素处理,identity 时被编译器优化掉。
附：print_layout 真实运行输出（K_INTER / MN_INTER / SwStateLayout / SwStateLayoutT）
1 | === RUN === |
- 标题: bwd-kernelB-beginner-guide-full
- 作者: 鱿鱼圈
- 创建于 : 2026-07-20 23:50:00
- 更新于 : 2026-09-12 22:26:40
- 链接: https://yuyanqi.com/2026/07/20/bwd-kernelB-beginner-guide-full/
- 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。