GEMM、Triton 与 TileLang：从分块代码到 Roofline 性能账本
整理日期：2026-10-07。范围来自本次实际可检索对话、可见提问摘录与已保存材料；原回答不完整处由本次教学推导补充，以下算例均非设备实测。技术来源核查于2026-10-07。
问题从哪里开始
9月24日追问矩阵AB的行/列切分、7个元素怎么均分；10月上旬的提问摘录进一步涉及 Triton program_id、mask、TileLang 的64×64×32分块，以及GEMM/GEMV何时受带宽限制。共同问题是：逻辑张量如何分给执行实例，数据如何复用，哪一项成本决定速度？
先把三类维度分开
对A[M,K]、B[K,N]，C[M,N]=AB。M、N决定输出元素，K决定每个输出的归约长度。BM/BN/BK是软件分块；grid是程序实例的组织；thread/warp是设备执行组织。它们有映射关系，不能直接画等号。
Triton的program_id(axis)返回对应grid轴的程序实例编号，axis=0不意味着“只处理一维张量”。一维grid可以用pid//num_n、pid%num_n映射二维输出块；二维grid也可以分别取得pid_m、pid_n。一个program处理一个tile，元素最终如何分布到线程/寄存器由编译配置与布局决定。向量加教程、矩阵乘教程
索引、mask与other分别保护什么
元素加的逻辑索引是offset=pid×BLOCK_SIZE+arange(0,BLOCK_SIZE)，mask=offset<n。教学例n=1000、BLOCK_SIZE=256：4个program，最后一块只有232个有效元素。load mask保护读取；store mask保护写入，两者不能互相替代。
你的判断“反正最后不写，other=0还有什么用”在独立逐元素加法中成立：若无效lane不影响有效lane，且相同mask禁止无效写回，填充值不会改变有效结果。但这不是通用结论。归约sum需要0，max常需负无穷；softmax的masked score通常用负无穷；GEMM的K尾块需要补0，避免无效乘积污染有效输出。有效/无效lane之间一旦有归约或交换，填充值就进入正确性契约。
二维指针依stride构造：A_ptr+i×stride_am+k×stride_ak。切片、转置与非连续张量不能仅凭shape假设行主连续。输出M/N尾块和K归约尾块要分别检查。
手算TileLang 64×64×32的成本
来自你的代码摘录：BM=BN=64、BK=32，threads=128，A/B shared为FP16，C fragment为FP32。这里只做逻辑存储估算：
- A tile：64×32×2=4096字节；B tile同为4096字节。
- 单套A/B共8192字节；若真实生成两套缓冲约16384字节，另加padding、屏障和临时量。num_stages与实际分配须看编译结果，不能只乘一个参数。
- C逻辑累加量64×64×4=16384字节。fragment分布在执行线程中，不代表每线程拥有完整C，也不能当成另一块shared。
- 一次K tile做2×64×64×32=262144 FLOPs，按一次FMA计2个FLOPs。忽略C写回及跨tile缓存时，A/B载入对应32 FLOP/byte。
128线程平均对应32个FP32累加元素只是均匀分布的教学账本；实际寄存器数量还包括索引、地址与中间值，并受布局和编译器影响。TileLang官方GEMM示例
增大BM/BN提高复用，也可能增加寄存器/shared占用、降低驻留块数，并增加尾块浪费。BK主要改变归约步数与流水资源；不是越大越快。T.Pipelined表达流水意图，是否重叠应由生成代码和时间线确认。
Roofline为什么没有统一矩阵尺寸门槛
定义I=FLOPs/指定存储层的传输字节，理想上界P≤min(P_compute,B×I)，转折点I*=P_compute/B。必须匹配同一设备、数据类型、稀疏口径和时间单位。用FP16矩阵峰值比较FP32标量路径会得出错误判断。
对beta=0、理想一次读取A/B并写一次C的普通GEMM，教学近似：
它是理想流量模型；实际HBM重复载入、cache复用、padding和额外转换会改变分母。
GEMV令N=1，矩阵A的读取占主导时I≈2/s_A，FP16约1 FLOP/byte，容易受带宽限制；batch多个向量可复用同一权重，从而改变强度。因此不是“GEMM一定算力受限、GEMV绝对带宽受限”。小shape还可能被launch/调度开销限制。NVIDIA矩阵乘性能指南
教学设备若计算峰值50 TFLOP/s、带宽1 TB/s，转折点为50 FLOP/byte；这不是你显卡或NPU的规格。32 FLOP/byte的tile局部账本也不能直接替代整个kernel的HBM算术强度。
如何验证一次优化
先用包含非整除维度、非连续stride、不同K、边界值的参考对照验证正确性；明确dtype、累加精度及容差。然后固定设备/版本/shape，预热、排除JIT，正确处理异步计时，比较成熟库基线。
每次只改变tile、warp数、stage或布局中的一个变量，记录耗时、误差、寄存器、shared、spill、实际访存与矩阵指令。Autotune会多次执行候选；若kernel会累加或修改输入，需配置恢复/重置，不能把一次性的状态更新直接拿来重复调参。Autotune官方API
最后接入模型，看该算子占关键路径多少。局部kernel快2倍不保证服务快2倍：量化转换、通信、排队和CPU后处理可能占主导。未完成上述测量前，不记录“已获得加速”。
与其他笔记的联系
先读存储与同步：CUDA 与昇腾数据流对照：存储层级、分块与异步流水；矩阵跨rank切分见DP2、TP8 与 MoE EP：两节点部署为何不一定是两份独立模型；低精度数据流见W8A8C8 与 GPQA：把量化位宽、数值误差和模型能力分开看。
复习验收：不看答案解释“为什么K尾块的0会影响正确性，M/N尾块的store mask又在保护什么”；给一个tile变大反而慢的可验证假设。