cp.async / ldmatrix 消融实验 —— 初步规划

目标：设计一套逐步叠加优化的消融实验（ablation），量化每一项优化（ldmatrix vs ld.shared、cp.async vs 同步 load vs TMA）到底带来了多少性能提升。

相关代码位置：kernels/hgemm/、kernels/hgemm_sm120/、kernels/sgemm/、kernels/flash_attn/、kernels/flash_decode/。 kernels_dsl/ 中无相关内容（仅有 triton/tilelang 的 vector_add、softmax toy 示例）。

0. 现状盘点

0.1 仓库中已有的 cp.async / ldmatrix 素材

0.2 可直接复用的 PTX 宏

kernels/hgemm/hgemm.cu:22-40、hgemm_sm120/hgemm_sm120.cu:26-55、sgemm/sgemm_tf32.cu:22-35：

#define CP_ASYNC_CG(dst, src)  asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n" ...)
#define CP_ASYNC_COMMIT_GROUP()      // cp.async.commit_group
#define CP_ASYNC_WAIT_GROUP_0()      // cp.async.wait_group 0
#define LDMATRIX_X4(R0,R1,R2,R3,PTR)   // ldmatrix.sync.aligned.m8n8.x4.shared.b16
#define LDMATRIX_X2(R0,R1,PTR)         // ldmatrix.sync.aligned.m8n8.x2.shared.b16
#define LDMATRIX_X2_TRANS(R0,R1,PTR)   // ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16

hgemm_sm120 里另有 cp_async_wait_group<N>()（N = 0..3），以及完整 MBarrier/TMA 封装：

• mbarrier_init / mbarrier_expect_tx / mbarrier_arrive / mbarrier_wait（hgemm_sm120/hgemm_sm120.cu:70-108）
• cp_async_bulk_tensor_2d / 3d / 4d（TMA 通道）

0.3 已有的对照数据（可直接当参考 baseline）

kernels/hgemm_sm120/readme.md，4096³ bf16：

hgemm_cublas              4.259 ms  32.27 TF
hgemm_bcf_dbf_rw          4.042 ms  34.00 TF   ← cp.async 路径
hgemm_k_stages            4.131 ms  33.27 TF   ← cp.async 多级流水
hgemm_tma_r_k_stages_64   4.288 ms  32.05 TF   ← TMA 路径
hgemm_tma_r_k_stages_32   4.006 ms  34.31 TF   ← TMA 路径

kernels/sgemm/readme.md，4096³ TF32：

sgemm_tf32_bt_swizzle_dbf      8.723 ms  15.76 TF   (非 cp.async：普通 load + st.shared)
sgemm_tf32_swizzle_bcf_dbf     8.276 ms  16.61 TF   (cp.async)

⚠️ 这组差异里 swizzle / dbf / grid-swizzle 三个因素是混在一起的，不能直接归因到 cp.async —— 这正是需要做受控消融的原因。

0.4 结论

目前仓库里没有现成的"cp.async vs ldmatrix 对比实验"框架：现有代码都是两者串联使用（cp.async/TMA 做 gmem→smem，ldmatrix 做 smem→reg），而不是互斥方案的 A/B 对照。

1. TMA 是什么

TMA = Tensor Memory Accelerator，Hopper (SM90) 引入的硬件 DMA 引擎，Blackwell (SM100/SM120) 保留并增强。它负责"把一整块多维 tile 从 global 搬到 shared"。

1.1 它解决什么问题

1.2 三个组成部分

1) TensorMap 描述符（host 端创建，128B 结构体）

// kernels/hgemm_sm120/hgemm_sm120.cu:731
cuTensorMapEncodeTiled(&tmap, type, /*rank=*/2, global_address,
                       globalDim, globalStrides, boxDim, elementStrides,
                       CU_TENSOR_MAP_INTERLEAVE_NONE,
                       swizzle,                    // ← 硬件 swizzle
                       CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
                       CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

globalDim/globalStrides 描述整个张量，boxDim 描述"每次搬多大一块"。在 kernel 里通过 __grid_constant__ const CUtensorMap 传参（hgemm_sm120.cu:556）。

2) Kernel 端发起

// kernels/hgemm_sm120/hgemm_sm120.cu:114
cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes
    [mbar], [tmap], [smem_ptr], fast_coord, slow_coord];

注意：cp.async.bulk.tensor（TMA 通道）与 cp.async.cg（LSU async copy 通道）不是一回事，两者互不冲突。

3) 完成跟踪：mbarrier

mbarrier_init(&mbar[i], 1);
mbarrier_expect_tx(&mbar[i], tx_bytes);        // 声明"我要等 N 字节"
cp_async_bulk_tensor_2d(&mbar[i], &tma_a, ...); // 发货
mbarrier_wait(&mbar[compute_stage], wait_phase);// 收货

比 cp.async.wait_group N 精确得多：wait_group 是"等到第 N 组之前都完成"，容易多等或少等；mbarrier 是"这 32KB 真的到位才放行"，多级流水（STAGES=3/4）时不浪费带宽。

1.3 约束

• global 地址需 16B 对齐
• box 最内维字节数必须匹配 swizzle 宽度（128B 模式要求最内维恰为 128B）
• tensor rank ≤ 5
• descriptor 需在 device 生命周期内有效（__grid_constant__ 传参）
• 只支持整 tile copy，不支持任意 gather

1.4 仓库中的位置

2. 消融实验设计

核心思路：把优化拆成两条独立的阶梯，每条阶梯逐档叠加，每档只改变一个变量，从而把增量收益归因清楚。

固定基线：GEMM 形状 M=N=K=4096、bf16、BM=BN=128、BK=32、256 线程 / 8 warp、warp tile 64×32、mma.m16n8k16。

2.1 阶梯 A：smem → reg（ldmatrix vs ld.shared）

对应 ldmatrix 这一项。需要三档，且必须给 ld.shared 变体最有利布局，否则测到的是布局差异而非指令效率。

A3 即现有 hgemm/hgemm.cu:131/142 的写法。

关键：A3 相比 A1/A2 省的是 LSU 指令数和地址计算指令，不省带宽（数据搬运量相同）。

2.2 阶梯 B：gmem → smem（cp.async vs 同步 load vs TMA）

2.3 ⚠️ 三个会让结论作废的坑

坑 1：B1 单 buffer 下 cp.async 几乎无收益，甚至更慢。 若按 B0→B1→B2 顺序测，会得出"cp.async 没用"的错误结论。cp.async 的价值在于释放 LSU 和 issue slot，只有当有第二级 buffer 承接、让搬运算与 mma 重叠时才能兑现。

✅ 正确阶梯顺序：同步 load + 双缓冲 → cp.async + 双缓冲 → TMA + 多级缓冲。 这样每一步的增量都归因清晰。

坑 2：在最终 kernel 上测 ldmatrix vs ld.shared，差异会被淹没。 若 kernel 被 mma 吞吐或 DRAM 带宽卡住，smem→reg 段的差异在 wall-clock 上看不见。

✅ 建议额外做一个 smem→reg 微基准（只 load 不 mma，循环 N 次），直接测 LSU 吞吐差异； 同时用 ncu 采集：

smsp__inst_executed_pipe_lsu_mem_shared_op_ld.sum
l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum
sm__inst_executed_pipe_lsu.avg.pct_of_peak_sustained_active

坑 3：只测一种 shape 会被误导。

• K 小（如 M=N=4096, K=256）→ latency-bound，cp.async / 多级流水收益最大
• K 大（K=8192）→ compute-bound，所有搬运优化收益都被 mma 吞吐盖住

✅ 建议扫 K ∈ {256, 512, 1024, 4096, 8192} × M/N ∈ {512, 4096}，画热力图。

另外两条硬性要求：

• 锁频：nvidia-smi -lgc 3050（hgemm_sm120/readme.md 已有此做法），否则重测噪声大于效应
• TMA 的硬件 swizzle 与软件 swizzle 不能同时开：TMA 变体的 smem 布局须与 cp.async 变体等价（都是 XOR 128B），否则测到的是 swizzle 实现差异
• 精度：cp.async / TMA 都不改数值，ldmatrix 与 ld.shared 也只重排不换算（ldmatrix.trans 仅重排），因此正确性可直接 diff 对比

3. 代码组织建议

现在 hgemm.cu 里有 4 个 kernel 几乎完全重复（naive / bcf / bcf_dbf / bcf_dbf_rw），继续加变体会爆炸。

建议新建 kernels/ablation/，用一个模板参数化的 kernel 覆盖全组合：

template <int LOAD_MODE,   // 0=ldg+st.shared, 1=cp.async, 2=TMA
          int READ_MODE,   // 0=ld.shared, 1=ldmatrix
          int B_LAYOUT,    // 0=Bs[BK][BN], 1=Bs[BN][BK]
          int STAGES,      // 1,2,3
          int SWIZZLE>     // 0/1
__global__ void hgemm_ablation_kernel(...);

用 if constexpr 分支（编译器消掉未走分支，零运行时开销），变体命名 abl_L{lm}_R{rm}_bl{bl}_s{s}_sw{sw}。

4. 待办清单

• 新建 kernels/ablation/，实现模板化 kernel（LOAD_MODE / READ_MODE / B_LAYOUT / STAGES / SWIZZLE）
• 实现 smem→reg 微基准（隔离出 ldmatrix 收益，避免被 mma 吞吐掩盖）
• 写 test.py：扫多个 shape、锁频、输出 Markdown 表格（每档 TFLOPS + 相对上一档的增量百分比）
• 写 ncu 指标采集脚本（stall reason / LSU 指令数 / bank conflict / 带宽利用率）
• 补 readme.md 记录结论表格

5. 附：参考指标（ncu）

已有 profiling 脚手架：kernels/hgemm/run_ncu.sh、kernels/hgemm/test_for_ncu.py （用法：ncu -k regex:"gemm|Kernel" --set full -o profile_res -f python test_for_ncu.py）。