下面是一份"照着读就能写"的清单，全部基于仓库实际行号（未改动任何文件）。

一、PYBIND 到底写在哪、怎么写

1. kernels/common/pack.cuh（30 行，整个仓库的绑定"语法糖"只有一行）

• pack.cuh:15 → #define torch_pybinding_func(f) m.def(#f, &f, #f)
  • 含义：把 C++ 函数 f 注册进 pybind 模块 m，名字 = C++ 函数名本身（#f 是字符串化），Python 侧就出现 lib.<函数名>。
  • 必须在 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 的 {} 里调用，因为宏展开后用到变量名 m。
• pack.cuh:13 LDST128BITS、pack.cuh:17 WARP_SIZE：写新 kernel 时直接复用，别重造。

2. kernels/hgemm/hgemm.cu:884-934（这是唯一需要抄的"顶层实验接口"C++ 侧）

884  #define CHECK_T(x) TORCH_CHECK(x.is_cuda() && x.is_contiguous(), ...)
886  #define binding_tiled_func_gen(name)  ...        # 生成 host 侧 launcher
919-922  binding_tiled_func_gen(hgemm_naive / bcf / bcf_dbf / bcf_dbf_rw);
924  extern void hgemm_cublas(torch::Tensor, torch::Tensor, torch::Tensor);
928  PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { torch_pybinding_func(...); }

binding_tiled_func_gen 宏体（886-917 行）就是你的新 kernel 与 Python 之间唯一的契约，逐行看它做了什么：

1. CHECK_T(a/b/c) → 要求 cuda + contiguous（你的 A2/A2L 传 b_t 也要 contiguous，host 侧 .t().contiguous() 即可）。
2. M=a.size(0), K=a.size(1), N=b.size(1) → 注意 B 的形状约定是 [K, N]；A2/A2L 的 B 是 [N, K]，所以 launcher 里 N/K 的取法要变（N=b.size(0), K=b.size(1)），别照抄。
3. BM=BN=128, threads=256, dim3 blocks_per_grid((N+BN-1)/BN, (M+BM-1)/BM) → grid 顺序是 x=N 方向、y=M 方向。
4. stream = at::cuda::getCurrentCUDAStream() → 必须在当前 stream 上跑。
5. if (a.dtype()==torch::kHalf) ... else ... → fp16/bf16 分派。
6. name##_kernel<128,128,32><<<grid, 256, 0, stream>>>(...) → 模板实参在这里写死。

⚠️ 你要做 STAGES=3（48KB dynamic smem）时，这段宏不够用：还要在 launch 前加 cudaFuncSetAttribute(kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes) 并把 smem 大小作为 <<<...>>> 第 3 个参数传入。 现成写法参考（虽然那两个文件是 TMA、SM80 不可用，但只抄这段 opt-in 写法即可）：

• kernels/flash_attn/flash_attn.cu:595（cudaFuncSetAttribute + <<<..., smem_size, stream>>>）
• kernels/flash_decode/flash_decode.cu:637
• kernels/softmax/softmax.cu:498
• dynamic smem 声明范例：kernels/flash_attn/flash_attn.cu:282（extern __shared__ __align__(128) uint8_t smem_buf[];）

3. kernels/hgemm/cublas.cu（"多 .cu 编成一个模块"的写法）

• 全文只有一个普通函数 void hgemm_cublas(torch::Tensor a, torch::Tensor b, torch::Tensor c)（36-70 行），故意不含 PYBIND11_MODULE。
• 它靠 hgemm.cu:924 的 extern void hgemm_cublas(...); 声明 + torch_pybinding_func(hgemm_cublas) 注册。
• ⇒ 你的 ablation.cu 想复用 cublas.cu，只需在 test.py 的 sources 里同时列上 ../hgemm/cublas.cu，并在 ablation.cu 里写同样的 extern 声明，两个 .cu 由 ninja 一起编译成一个 .so。

4. kernels/sgemm/sgemm.cu:526-541（另一个 PYBIND 例子，确认"一函数一名字"的对应关系）

• sgemm.cublas / sgemm_naive / sgemm_tiling / sgemm_at_tiling / sgemm_at_bcf_swizzling / sgemm_at_bcf_swizzling_rw / sgemm_at_bcf_swizzling_dbf_rw / sgemm_tf32_*
• ⇒ 这就是"阶梯式消融命名法"的现成范式：名字本身就是阶梯表的一列。你的 A1/A2/A2L/A3 就照这个命名（例如 abl_A1_ld_shared、abl_A2_ld_shared_bt、abl_A2L_ldmatrix_bt、abl_A3_ldmatrix），Python 侧 lib.abl_A1_ld_shared(...) 直接对应。

二、.py 顶层实验接口怎么写

1. kernels/hgemm/test.py（这就是你要模仿的模板，逐块看）

2. kernels/hgemm/test_for_ncu.py（ncu 版）

• 与 test.py 几乎相同，差别：warmup=0, rep=1，且不计时不校验，只为了让 ncu 采样一次。sources 写成相对路径 "hgemm.cu"（因为 ncu 脚本通常在同目录跑）。
• kernels/hgemm/run_ncu.sh（1 行）：ncu -k regex:"gemm|Kernel" --set full -o profile_res -f python test_for_ncu.py ⇒ 你的 ablation/ 里照抄成 ncu -k regex:"ablation" --set full -o prof -f python test_for_ncu.py。

3. kernels/sgemm/test.py + kernels/sgemm/readme.md（表格与命名范式）

• sgemm/test.py 的 load(name="gemm_lib", sources=[sgemm.cu, cublas.cu, sgemm_tf32.cu]) → 一个模块里放多个 .cu 的实例。
• sgemm/readme.md → 结果表格式（每档一行 ms + speedup + tflops 的 yaml 块）→ 你的 kernels/ablation/readme.md 就照这个格式填"表 2"。

三、"实验里调用后续要修改的 kernel 的代码原理"

这条链要读的是 kernels/hgemm/hgemm.cu 的 kernel 主体（它是你的实验基座，A3 = 现有 bcf_dbf 的读路径）：

调用链全景（文档 §6 的代码依据）：

test.py: lib = load(sources=[...])          # ninja 现场编译 → ablation_lib.so
      ↓ import 即执行 PYBIND11_MODULE → m.def("abl_A1", ...)  (pack.cuh:15)
      ↓ lib.abl_A1(a,b,c)
   C++ launcher (hgemm.cu:886 宏)  → 检查 contiguous → 取 M/N/K → grid/block → stream
      ↓ kernel<<<grid,256,smem,stream>>>(...)
   device: grid swizzle → cp.async gmem→smem → wait_group + __syncthreads
           → ldmatrix/ld.shared smem→reg → mma → 写回 C
      ↓
   benchmark(): warmup10 → sync → perf_counter → rep200 → sync → TFLOPS

四、"新增一组需要测试的 kernel" 具体要写哪些、写在哪

建议新建目录 kernels/ablation/，完全不动 kernels/hgemm/*（保持历史数据可复现），需要写 6 个文件：

1) kernels/ablation/ablation.cu —— 唯一的 C++ 主体

• 模板 kernel：template <int LOAD_MODE, int READ_MODE, int B_LAYOUT, int STAGES, int SWIZZLE, typename T> __global__ void hgemm_ablation_kernel(...)
  • smem：dynamic（extern __shared__ __align__(128) char smem_raw[];）+ cudaFuncSetAttribute（抄 flash_attn.cu:595 的写法），因为 STAGES=3 → 48KB、STAGES=4 → 64KB 静态 smem 放不下。
  • device 主体照抄 hgemm.cu:58-201（naive）的骨架，只把 ①cp.async 装载段、②ldmatrix 段换成 if constexpr (READ_MODE==...) 三套寻址（文档 §3.4 已给出具体地址公式）。
  • 搬运模式 LOAD_MODE（cp.async vs ldg+st.shared）保留参数位，为阶梯 B 预留。
• 绑定层：
  • ABL_BIND(NAME, LM, RM, BL, ST, SW) 宏（仿 hgemm.cu:886 的 binding_tiled_func_gen），但不能照抄它的宏体——你要加 cudaFuncSetAttribute、要按 B_LAYOUT 决定 B 的形状取法（N=b.size(0),K=b.size(1) vs N=b.size(1),K=b.size(0)）、要传 smem 大小。
  • extern void hgemm_cublas(torch::Tensor, torch::Tensor, torch::Tensor); + torch_pybinding_func(hgemm_cublas);（复用 baseline，注意 cublas.cu 期望 B 是 [K,N]）。
  • PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { torch_pybinding_func(abl_A1); ... }。
  • 命名建议沿用 sgemm 阶梯范式：abl_A1_ld_shared / abl_A2_ld_shared_bt / abl_A2L_ldmatrix_bt / abl_A3_ldmatrix。

2) kernels/ablation/test.py —— 顶层实验接口（PyBIND 的消费端）

• 抄 kernels/hgemm/test.py 全文，改：load(name="ablation_lib", sources=[ablation.cu, ../hgemm/cublas.cu], ...)。
• test_4096() 里：torch.matmul baseline → lib.hgemm_cublas → 四个 lib.abl_* 各 benchmark + diff_check。
• A2/A2L 额外：b_t = b.t().contiguous()（计时之外生成，且要连续分配以免引入额外 D2D 抖动），并 diff_check 时用 b 而非 b_t 算参考值。
• 建议再加一个"档位矩阵"打印函数，直接输出文档表 2 的 Markdown 行。

3) kernels/ablation/test_for_ncu.py + 4) kernels/ablation/run_ncu.sh

• 抄 hgemm/test_for_ncu.py（warmup=0, rep=1）与 hgemm/run_ncu.sh（ncu -k regex:"ablation" --set full -o prof -f python test_for_ncu.py）。

5) kernels/ablation/micro.cu + micro_test.py（文档 §4.3 强烈建议）

• 同一模板加 DO_MMA 开关：DO_MMA=0 只做 fragment 载入并 acc ^= reg 防 DCE（纯读路径吞吐），DO_MMA=1 加 16 次 mma。
• micro_test.py 抄 test.py 的 benchmark，sources=[micro.cu]。

6) kernels/ablation/readme.md

• 抄 kernels/sgemm/readme.md 的格式：勾选清单（每档一行说明）+ yaml 结果块（ms / speedup / tflops）+ 复现命令（export TORCH_CUDA_ARCH_LIST=...）。

五、动手前必须记住的 5 个"坑"（都来自上述文件）

1. 模块名要唯一：name="ablation_lib" 与 hgemm_lib 不同，否则 torch 扩展缓存会串（缓存在 ~/.cache/torch_extensions）。
2. sources 必须列全：.cu 增删后要改 sources（hgemm 目录 vs sgemm 目录的 name 分别是 hgemm_lib / gemm_lib，同一 .so 内多文件靠这个 list 合并）。
3. extern __shared__ + cudaFuncSetAttribute 必须成对，否则 >48KB 直接 launch 失败。
4. B 的形状契约随 B_LAYOUT 变：A1/A3 是 [K,N]，A2/A2L 是 [N,K]，launcher 与 Python 侧要同时改，且 CHECK_T 的 contiguous 要求不能破。
5. TORCH_CUDA_ARCH_LIST 必须 export，且测性能前 nvidia-smi -lgc/-lmc 锁频、测完 -rgc/-rmc 解锁，否则噪声大于 A1↔A2 的 1~5% 效应量。

需要的话，下一步我可以按上面这份清单，把 kernels/ablation/ 的 6 个文件骨架写出来（仍然不改动 kernels/hgemm/*）。