Cube 核算子优化路径【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills原则1首先是模式选择不同模式对应不同的优化路径要选择对应的优化路径进行参考。原则2专门的调优实验要比总结的规律重要总结的规律只有当按照实验做完没有收益时再考虑。1. 模式识别与路径选择两种编程模式模式计算指令标识特征适用场景Expert 模式T.mma显式分配 L0A/L0B/L1 手动 flag 同步 auto_syncFalse追求极致性能Developer 模式T.gemm_v0L1 分配 auto_syncTrue 无 L0 层显式控制快速开发、已有实现适配识别方法检查 kernel 代码中的关键特征# Expert 模式标识 T.alloc_L0A(...) # 显式分配 L0A T.alloc_L0B(...) # 显式分配 L0B T.mma(A_L0, B_L0, C_L0) # 使用 T.mma intrinsic T.set_flag(...) # 手动 flag 同步 T.wait_flag(...) # 手动 flag 同步 # Developer 模式标识 T.gemm_v0(A_L1, B_L1, C_L0) # 使用 T.gemm_v0 高层抽象 # 无 L0A/L0B 显式分配 # 无手动 flag 同步依赖 auto_sync路径选择识别到 T.mma → Expert 路线§2 识别到 T.gemm_v0 → Developer 路线§3 其他 cube 算子FA、SparseFA 等同理识别2. Expert 模式优化路径T.mma2.1 优化路径4 步Step 1: 最简 T.mma单缓冲 → 选当前最优 tilingblock_KK_L1128打满 L0详见 §5 硬件约束速查 / §6 Tiling 参数说明 → 性能0.64x Step 2: 整体双缓冲 三级K分块不可分割 → 重新评估 tilingblock_K128 双缓冲时 L0B 溢出 → 调整为 block_K64, K_L1256 block_K × 4 → 性能0.95x48.8% 变体预取双缓冲L0 层预取重叠❌ 不推荐 → 在整体双缓冲基础上循环外预取 kk0循环内预取 kk1 → copy[kk1] 和 mma[kk] 重叠执行 → 无 FC0.93x比整体双缓冲还差flag 同步开销暴露 → 有 FC1.01x与整体双缓冲FC 持平无净收益 → 结论预取重叠收益被 flag 同步开销抵消不推荐采用 Step 3: Fixed Core → tiling 不变只改 launch 方式按物理核数 launch循环处理 tile → 性能1.01x5.7% → 注意小矩阵tile 数 ≤ 20可能劣化 Step 4: Swizzle实验性劣化不采用 → 加 T.use_swizzle 重映射 tile 到 core 的分配顺序 → 性能0.99x-1.8% → 结论整体劣化不默认采用但后续算子仍需实验验证2.2 双缓冲策略双缓冲与三级K分块的耦合关系双缓冲和三级K分块是不可分割的优化单元必须同时实现没有双缓冲三级K是纯串行K_L1 切成 block_K 的小块每次 L1→L0 搬运和 MMA 计算串行执行多切只会多搬运没有重叠。实测单独加三级K单缓冲性能反而下降 19.6%。没有三级K双缓冲的 L0 ping-pong 不生效block_KK_L1 时kL0split1buffer 1 从未使用白白付出双缓冲同步开销。两者同时存在才有意义三级K 提供 L1→L0 的内层循环loop_kk ≥ 2双缓冲让内层循环的搬运和计算重叠。整体双缓冲 vs 预取双缓冲整体双缓冲gemm_intrinsic.pyL0 buffer 数 S22但 copy 和 MMA串行执行每个 kk 迭代copy[kk] → wait → mma[kk]无 FC0.95x / 有 FC1.01x预取双缓冲example_gemm_hiperef.pyL0 buffer 数 S22copy 和 MMA重叠执行循环外预取 kk0循环内预取 kk1每个 kk 迭代copy[kk1] ∥ mma[kk]逐步对比实验20 case 全量3 次取平均版本无 FC有 FC有 FC Swizzle整体双缓冲0.95x1.01x0.99x预取双缓冲0.93x1.01x0.97xFixed Core 在预取双缓冲上的收益9.2%大部分 case 正向Case 8 38.8%Swizzle 在预取双缓冲上的收益-4.1%整体劣化仅 2 case 正向结论预取双缓冲无净收益不推荐采用。原因分析纯预取双缓冲无 FC反而更差0.93x 0.95x。预取多出的 flag 同步开销循环外预取 kk0 循环内预取 kk1 的 wait/set flag在按 tile 数 launch 时暴露抵消了 copy/mma 重叠的理论收益。有 FC 时两者持平1.01x 1.01x。FC 省掉的 ~200us 固定开销淹没了预取的同步开销但预取重叠的收益仍被同步开销抵消没有净增益。Swizzle 在预取双缓冲上劣化更严重-4.1% vs -1.8%。预取双缓冲的调度更复杂Swizzle 打乱 tile 顺序对 cache 局部性的负面影响更大。最优方案整体双缓冲 Fixed Core无 Swizzle即 Step 3性能 1.01x。预取双缓冲代码模式仅供参考不推荐使用# 循环外预取 kk0 - L0[0] T.wait_flag(mte2, mte1, k % S1) T.wait_flag(m, mte1, 0) T.copy(A_L1[k % S1, 0, 0], A_L0[0, :, :]) T.copy(B_L1[k % S1, 0, 0], B_L0[0, :, :]) T.set_flag(mte1, m, 0) for kk in T.serial(loop_kk): # 预取 kk1 - L0[(kk1)%S2]与 mma[kk] 重叠 if kk loop_kk - 1: T.wait_flag(m, mte1, (kk 1) % S2) T.copy(A_L1[k % S1, 0, (kk 1) * block_K], A_L0[(kk 1) % S2, :, :]) T.copy(B_L1[k % S1, (kk 1) * block_K, 0], B_L0[(kk 1) % S2, :, :]) T.set_flag(mte1, m, (kk 1) % S2) if kk loop_kk - 2: T.set_flag(mte1, mte2, k % S1) # 计算 kk数据已在上一轮预取就绪 T.wait_flag(mte1, m, kk % S2) T.mma(A_L0[kk % S2, :, :], B_L0[kk % S2, :, :], C_L0, initT.And(k 0, kk 0)) T.set_flag(m, mte1, kk % S2)2.3 Tiling 推导方法前提本节涉及的 L0C/L0A/L0B 容量约束和参数定义详见 §5 硬件约束速查 和 §6 Tiling 参数说明。Step 1: 选 block_M, block_N, block_K对于 fp16 fp32 累加L0C: block_M × block_N ≤ 32768 (128KB / 4B) L0A: block_M × block_K ≤ 32768 (64KB / 2B) L0B: block_K × block_N ≤ 32768 (64KB / 2B)目标最大化 L0C 利用率block_M × block_N 接近 32768单缓冲下的最优选择(128, 256, 128): L0C 32768 (100%), L0A 16384 (50%), L0B 32768 (100%)双缓冲下的最优选择(128, 256, 64): L0C 32768 (100%), L0A 8192 (25%), L0B 16384 (50%)双缓冲后 L0B 64 × 256 × 2 × 2 64KB (100%)刚好用满Step 2: 加双缓冲三级K时重新评估L0A: block_M × block_K × sizeof(dtype) × 2 ≤ 64KB L0B: block_K × block_N × sizeof(dtype) × 2 ≤ 64KB调整逻辑单缓冲下 block_K128 打满 L0B双缓冲后 L0B 翻倍溢出需将 block_K 减半到 64。Step 3: 选 K_L1K_L1 应该是 block_K 的倍数越大搬运效率越高。对于 (128, 256, 64)K_L1 256: L1 (128×256 256×256) × 2 192KB ✅推荐K_L1 256 block_K × 4每次 L1 加载做 4 轮 MMA什么时候需要重新评估 tiling触发条件需要检查可能的调整加双缓冲三级KS12, S22L0A/L0B 是否溢出block_K 减半加 L1 双缓冲S12L1 是否溢出K_L1 减小改 dtype如 fp16→int8分形限制、内存占用所有参数重新推导改 accum_dtypeL0C 占用变化block_M × block_N 上限变化2.4 优化案例逐步优化过程Step 1: 最简 T.mma单缓冲block_M 128, block_N 256, block_K 128, K_L1 128 性能0.64xStep 2: 整体双缓冲 三级K分块block_K: 128 → 64, K_L1: 128 → 256 性能0.95x48.8%Step 2 变体: 预取双缓冲❌ 不推荐性能无 FC 0.93x / 有 FC 1.01x与整体双缓冲持平Step 3: Fixed Core性能1.01x5.7%Step 4: Swizzle❌ 劣化不采用性能0.99x-1.8%最终 tilingblock_M 128, block_N 256, block_K 64, K_L1 256, S1 2, S2 2性能验证Step优化内容tiling 变化平均加速比Step 1最简 T.mma单缓冲block_K128, K_L11280.64xStep 2整体双缓冲 三级Kblock_K64, K_L12560.95xStep 2 变体预取双缓冲 ❌ 不推荐不变无 FC 0.93x / 有 FC 1.01xStep 3Fixed Core不变1.01x最优Step 4Swizzle劣化不采用不变0.99x2.5 常见陷阱Expert 专属陷阱 1: L0C 溢出L0C 容量限制详见 §5 硬件约束速查。# 错误block_M × block_N × 4 128KB block_M 256; block_N 256 # L0C 256KB 128KB ❌ # 正确 block_M 128; block_N 256 # L0C 128KB ✅陷阱 2: 加双缓冲后未重新评估 tilingL0B 容量限制详见 §5 硬件约束速查。# 单缓冲选 block_K128加双缓冲后 L0B 溢出 # L0B (双缓冲) 128×256×2×2 128KB 64KB ❌ # 正确block_K 减半到 64 # L0B (双缓冲) 64×256×2×2 64KB ✅陷阱 3: 双缓冲与三级K分开实现# 错误先单独加三级K单缓冲→ 性能下降 19.6% # 错误先单独加双缓冲block_KK_L1→ buffer 1 从未使用 # 正确双缓冲 三级K同时实现 # block_K64, K_L1256, S12, S22 → 性能提升 48.8%陷阱 4: 预取双缓冲看起来更优但实测无净收益# 整体双缓冲copy 和 MMA 串行 for kk in T.serial(loop_kk): T.copy(...) T.set_flag(mte1, m, kk % S2) T.wait_flag(mte1, m, kk % S2) # 阻塞等待 T.mma(...) T.set_flag(m, mte1, kk % S2) # 性能无 FC 0.95x / 有 FC 1.01x ← 推荐 # 预取双缓冲copy 和 MMA 重叠删除了 wait_flag # 但多出的 flag 同步开销抵消了 copy/mma 重叠收益 # 性能无 FC 0.93x / 有 FC 1.01x ← 不推荐陷阱 5: Swizzle 不是万能的# 错误默认使用 Swizzle → 整体劣化 -1.8% # 正确实验验证劣化就不采用 # 正向 case 特征tile 数 18~32计算时间 20us # 劣化 case 特征tile 数 500 或 tile 数 183. Developer 模式优化路径T.gemm_v03.1 gemm_v0 特性与限制特性说明影响auto_syncTrue编译器自动插入barrier_all()同步开销大于手动 flagL0 层双缓冲以kL0Size为单位进行切分内部实现L1到L0的双缓冲优化内嵌同步内部实现双缓冲的同时也引入了内部隐式同步算子外部无法实现l1层的双缓冲3.2 可调参数参数定义和容量约束详见 §6 Tiling 参数说明 和 §5 硬件约束速查。参数说明默认值调优建议block_MM 维度 tile 大小128固定受 L0C 约束block_NN 维度 tile 大小256固定受 L0C 约束K_L1GM→L1 搬运粒度128需 kL0Size 才能触发 ping-pongkL0SizeL0 层切片大小内部参数128需小于K_L1非前端参数可去往src/tl_templates/ascend/common.h的constexpr uint32_t kL0Size进行修改⚠️框架级改动提示若按本条修改了 common.h完成后必须告知用户kL0Size位于框架文件src/tl_templates/ascend/common.h并非算子前端参数。若 agent 依据本 skill 修改了该文件请在完成后主动向用户输出以下说明此改动是框架级的会影响所有调用gemm_v0的算子不限于当前算子理论上该改动只会让性能更好、不会劣化理论上只要生成算子即可修改 common.h 属临时手段该处已向 common.h 负责人提交 issue后续框架会对此处进行优化届时算子侧无需再改 common.h。3.3 优化路径以及实测数据优化路径可参考如下调优实验三版逐步优化实验由于gemmv0的特性与限制developer模式的乘法算子调优空间较小集中再tiling策略与核调度策略优化路径为基础实现-tiling策略调优-核调度调优实验数据如下。版本K_L1kL0SizeLaunch平均加速比说明原始版64128tile 数0.44xping-pong 失效实验A12864tile 数0.51x (15.9%)K_L1/kL0Size 调优ping-pong 生效实验B1286420核 grid-stride0.52x (2.0%)Fixed Core有效收益实验A既然无法实现l1阶同步那就要充分利用l0阶同步kL0Size一定要小于K_L1又由于硬件限制kL0Size改为64K_L1是其两倍使双缓冲ping-pong生效。实验B由于大部分case都是大shape分核后tile数远超核数故选择FixedCore进行优化。结论kL0Size 必须小于 K_L1否则 L0 层 ping-pong 失效而K_L1以128为最优所以将kL0Size改为64使用FC优化榨干最后的性能收益3.4 常见陷阱Developer 专属陷阱 1: kL0Size ≥ K_L1 导致 ping-pong 失效# gemm_v0 内部common.h 默认 kL0Size128若前端 K_L1128 # kL0split ceil(K_L1 / kL0Size) ceil(128/128) 1 → ping-pong 不生效 # 正确方案 A编辑 common.h 改 kL0Size64前端 K_L1128 # → kL0split2ping-pong 生效实验A0.51x # 正确方案 B升级到 Expert 模式用 T.mma 替代 gemm_v0⚠️框架级改动提示若按正确方案 A修改了 common.h完成后必须告知用户kL0Size位于框架文件src/tl_templates/ascend/common.h并非算子前端参数。若 agent 依据本 skill 修改了该文件请在完成后主动向用户输出以下说明此改动是框架级的会影响所有调用gemm_v0的算子不限于当前算子理论上该改动只会让性能更好、不会劣化理论上只要生成算子即可修改 common.h 属临时手段该处已向 common.h 负责人提交 issue后续框架会对此处进行优化届时算子侧无需再改 common.h。4. 共享优化技术4.1 Fixed Core原理按物理核数launch每个核循环处理多个 tileL1/L0 buffer 只分配一次被复用。相比按 tile 数 launch节省了重复的 buffer 分配/释放开销。收益规律Fixed Core 的收益本质上是省掉一笔固定开销约 200us这笔开销在总时间中的占比决定了收益百分比。控制变量实验tile 数2048K 变化Expert 模式K每 tile 计算量Step 2 (us)Step 3 (us)节省 (us)收益102433M小764.50552.22212.2827.8%204867M中1254.421052.06202.3616.1%4096134M大2262.422086.24176.187.8%8192268M很大4388.344276.58111.762.5%各模式适用性模式Fixed Core 效果说明Expert (T.mma)✅ 有效5.7%手动 flag 同步开销小FC 收益能体现Developer (T.gemm_v0)✅ 有效收益2.0%auto_sync 的 barrier_all 部分抵消但仍有净增益代码规范模板Fixed Core 模板core_num 20 # 910B 物理核数 T.prim_func def main(A, B, C): # 1. 固定核数core_num with T.Kernel(core_num, is_npuTrue) as (cid, _): # ... buffer申请 ... with T.Scope(C): # 2. grid-stride 循环处理所有 tile for i in T.serial(T.ceildiv(m_num * n_num, core_num)): # 3. 计算当前 tile ID cid_task i * core_num cid # 4. 尾块守卫tile 数不整除 core_num 时 if cid_task m_num * n_num: bx cid_task // n_num by cid_task % n_num # ... tile 计算 ...关键规范点T.Kernel(core_num)固定核数 launch以实际物理核数为准if cid_task m_num * n_num尾块守卫tile 数不整除 core_num 时避免越界4.2 Swizzle原理T.use_swizzle重映射 tile 到 core 的分配顺序让相邻核访问不同的 A/B 区域减少 L2 cache 冲突。# 简单映射相邻核访问同一块 AL2 cache 冲突 核0 → tile(0,0) 核1 → tile(0,1) → 都访问 A[0:128] # Swizzle 映射off3相邻核访问不同 A 区域 核0 → tile(0,0) 核1 → tile(1,1) → 访问 A[0:128] 和 A[128:256]实验结果Expert 模式劣化不采用整体表现Step 4 (Fixed Core Swizzle) 平均 0.99x比 Step 3 (Fixed Core) 的 1.01x 还差-1.8%。逐 case 分析CaseShapetilesStep 3 (us)Step 4 (us)Swizzle 收益11024³3216.1813.9014.1%51024³ bf163216.0813.9213.4%121×3584×46081831.3028.588.7%48192³20484202.534366.12-3.9%15512×4096²51296.76114.55-18.4%171009×1024²3213.1415.30-16.4%1964³13.364.98-48.2%劣化原因分析小矩阵tile 数 ≤ 32Swizzle 打乱了原本紧凑的 tile 调度增加 cache miss大矩阵tile 数 500计算时间主导Swizzle 的 cache 优化收益被淹没中等矩阵tile 数 100~500收益不稳定依赖具体 shape正向 case 特征tile 数适中18~32、计算时间较短 20us、L2 cache 冲突是主要瓶颈各模式适用性模式Swizzle 效果说明Expert (T.mma)❌ 劣化-1.8%整体劣化仅 3 case 正向Developer (T.gemm_v0)❌ 劣化与 FC 一起测试整体劣化使用建议条件建议原因默认不采用整体劣化后续算子调优实验验证不同算子可能有不同效果tile 数 18~32 且计算时间短可尝试正向 case 的特征tile 数 500不采用计算主导cache 优化无感结论Swizzle 不是万能的需要根据具体算子和 shape 实验验证。如果实验发现劣化就不采用。5. 硬件约束速查Ascend 910B内存层级容量层级容量用途L0C128 KB输出累加器C 矩阵L0A64 KBA 矩阵寄存器L0B64 KBB 矩阵寄存器L1~512 KB共享缓存A_L1, B_L1分形限制最小 tile 维度dtypeL0A (M×K)L0B (K×N)L0C (M×N)fp16/bf16M≥16, K≥16K≥16, N≥16M≥16, N≥16int8M≥16, K≥32K≥32, N≥16M≥16, N≥16fp32M≥16, K≥8K≥8, N≥16M≥16, N≥16对齐要求存储单元对齐UB / L132 ByteL0A / L0B512 ByteL0C64 Byte6. Tiling 参数说明L0 层参数MMA 计算单元block_M: 输出 tile 的 M 维度大小block_N: 输出 tile 的 N 维度大小block_K: L1→L0 一次搬运的 K 维度大小MMA 粒度L1 层参数共享缓存K_L1: GM→L1 一次搬运的 K 维度大小搬运粒度参数关系容量数值详见 §5 硬件约束速查。L0 层约束block_M, block_N, block_K 互相制约 L0C: block_M × block_N × sizeof(accum) ≤ 128KB L0A: block_M × block_K × sizeof(dtype) ≤ 64KB L0B: block_K × block_N × sizeof(dtype) ≤ 64KB L1 层约束K_L1 独立选择 L1: (block_M × K_L1 K_L1 × block_N) × sizeof(dtype) ≤ 512KB K_L1 应该是 block_K 的倍数7. 通用公式本节使用的 L0C/L0A/L0B 容量参数详见 §5 硬件约束速查。给定 dtype 和 accum_dtype快速计算最优 tilingdef compute_optimal_tiling(dtype_size, accum_size, L0C_capacity128*1024, L0AB_capacity64*1024, use_double_bufferFalse): 计算最优 tiling 参数。 约束联立 L0C: block_M × block_N × accum_size ≤ L0C_capacity 尽量打满 L0A: block_M × block_K × dtype_size × buf ≤ L0AB_capacity L0B: block_K × block_N × dtype_size × buf ≤ L0AB_capacity Args: dtype_size: 输入数据类型大小字节如 fp162, int81 accum_size: 累加器数据类型大小字节如 fp324 L0C_capacity: L0C 容量字节默认 128KB L0AB_capacity: L0A/L0B 容量字节默认 64KB use_double_buffer: 是否使用双缓冲 Returns: (block_M, block_N, block_K, K_L1) 注block_K 对应 Developer 模式 gemm_v0 内部的 kL0Size Developer 模式无前端 block_K需在 common.h 中调整 kL0Size 与之对应。 FRACTAL 16 # 分形对齐fp16 M/N/K ≥ 16 且为 16 倍数向下取整不会爆容量 # Step 1: 从 L0C 约束推导 block_M, block_N尽量打满 L0C max_l0c_elems L0C_capacity // accum_size # fp32: 128K/4 32768 # 采用 M:N 1:2GEMM 常用打满 L0C 同时不撑爆 L0A2·block_M² ≤ max_l0c_elems block_M int((max_l0c_elems / 2) ** 0.5) block_M max(FRACTAL, (block_M // FRACTAL) * FRACTAL) # 对齐 → 128 block_N min(2 * block_M, max_l0c_elems // block_M) block_N max(FRACTAL, (block_N // FRACTAL) * FRACTAL) # 对齐 → 256 # Step 2: 从 L0A/L0B 约束推导 block_K受 block_M、block_N 双向约束 buffer_multiplier 2 if use_double_buffer else 1 max_block_K_from_L0A L0AB_capacity // (block_M * dtype_size * buffer_multiplier) max_block_K_from_L0B L0AB_capacity // (block_N * dtype_size * buffer_multiplier) block_K min(max_block_K_from_L0A, max_block_K_from_L0B) block_K max(FRACTAL, (block_K // FRACTAL) * FRACTAL) # 对齐 # Step 3: 选 K_L1block_K 的倍数 K_L1 block_K * 2 return block_M, block_N, block_K, K_L1 # 示例fp16 fp32 累加 # 单缓冲 bm, bn, bk, kl1 compute_optimal_tiling(2, 4, use_double_bufferFalse) print(f单缓冲: block_M{bm}, block_N{bn}, block_K{bk}, K_L1{kl1}) # 输出单缓冲: block_M128, block_N256, block_K128, K_L1256 # 双缓冲 bm, bn, bk, kl1 compute_optimal_tiling(2, 4, use_double_bufferTrue) print(f双缓冲: block_M{bm}, block_N{bn}, block_K{bk}, K_L1{kl1}) # 输出双缓冲: block_M128, block_N256, block_K64, K_L1128【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考