内存访问模式优化指南【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills触发条件当 Triton Agent 在优化 kernel 内存访问模式时遇到以下任一场景应参考本文档kernel 中存在%取余运算用于边界处理循环内存在多条tl.load且存在可重排空间循环内使用更新指针地址kernel 入参中存在运行时不变的固定值参数核心知识4 个优化模式总览#优化模式核心问题优化手段性能收益1变量取余 → Mask 替代%导致标量化破坏向量化访存用 mask 显式处理边界保持连续地址2x - 10x2Load 指令重排序有依赖的 load 阻塞无依赖的 load将无依赖的 load 提前发射与上次 store 并行1.05x - 1.20x3避免循环内更新指针产生 RAW 依赖阻塞流水线用基地址 偏移量独立计算地址1.05x - 1.20x4入参静态化运行时参数阻止编译期常量折叠tl.constexpr启用编译期优化视场景而定模式 1变量取余 → Mask 替代触发条件kernel 中使用%对地址索引取余如(pid * BLOCK offset) % N使用取余处理矩阵尾块padding 区域的越界索引原理在昇腾 NPU 上模运算%通过标量 ALU 执行导致向量运算退化为逐元素标量计算。取余后的地址不连续破坏向量化内存访问模式向量宽度如 128/256 bit无法充分利用性能可能下降 5-10 倍。核心思路移除%取余保持连续地址计算改用mask显式标记有效元素mask本身是向量化的比较操作不破坏向量化访存。特性使用%取余使用 Mask地址计算标量化逐个元素向量化SIMD内存访问离散、不连续连续、对齐边界处理隐式通过取余显式通过 mask性能低可能慢 5-10 倍高充分利用向量宽度代码对比优化前使用%取余offset_xm (pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M)) % M offset_wn (pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N)) % N offset_k tl.arange(0, BLOCK_SIZE_K) x_ptrs x_ptr (offset_xm[:, None] * K offset_k[None, :]) w_ptrs w_ptr (offset_k[:, None] * N offset_wn[None, :]) accumulator tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): x tl.load( x_ptrs, maskoffset_k[None, :] K - k * BLOCK_SIZE_K, other0.0, care_paddingFalse ) w tl.load( w_ptrs, maskoffset_k[:, None] K - k * BLOCK_SIZE_K, other0.0, care_paddingFalse ) accumulator tl.dot(x, w) x_ptrs BLOCK_SIZE_K w_ptrs BLOCK_SIZE_K * N优化后使用 Mask 替代%offset_xm pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offset_wn pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) offset_k tl.arange(0, BLOCK_SIZE_K) a_ptrs_base x_ptr (offset_xm[:, None] * K offset_k[None, :]) b_ptrs_base w_ptr (offset_k[:, None] * N offset_wn[None, :]) msk_m offset_xm M msk_n offset_wn N accumulator tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): x_ptrs a_ptrs_base k * BLOCK_SIZE_K w_ptrs b_ptrs_base k * BLOCK_SIZE_K * N x tl.load( x_ptrs, maskmsk_m[:, None] (offset_k[None, :] K - k * BLOCK_SIZE_K), other0.0, ) w tl.load( w_ptrs, maskmsk_n[None, :] (offset_k[:, None] K - k * BLOCK_SIZE_K), other0.0, ) accumulator tl.dot(x, w)关键点移除%操作offset_xm和offset_wn不再取余保持连续地址添加边界 maskmsk_m offset_xm M和msk_n offset_wn N在循环外计算一次组合 mask使用位与而非and逻辑与组合多个 mask 条件and是 Python 逻辑操作符返回单个布尔值不适用于 tensor maskmask 计算位置边界 mask 在循环外计算避免循环内重复计算不要遗漏维度确保所有需要边界处理的维度都有对应的 mask模式 2Load 指令重排序触发条件循环内存在多条tl.load指令其中一条 load 与上一次循环的 store 存在数据依赖如 load 同一地址而另一条 load 无此依赖原理编译器不会修改用户 load 指令的顺序。当循环内存在多条 load 时如果排在前的 load 因等待上一次循环的 store 而阻塞排在后的无依赖 load 也无法提前发射导致串行执行。核心思路将无数据依赖的 load 提前到有依赖的 load 之前使其可与上一次循环的 store 并行执行充分利用流水线并行能力。执行时序对比优化前load B 在前阻塞 load A: 迭代 i: load B → load A → calc → store O → store B 迭代 i1: load B(等store B) → load A(等load B) → ... 优化后load A 在前与 store B 并行: 迭代 i: load A → load B → calc → store O → store B 迭代 i1: load A(与store B并行) → load B → ...优化前优化后load B 等待上一次 store B 完成后才能执行load A 无需等待可以提前发射load A 必须等 load B 完成后才能执行load A 可以与上一次循环的 store B 并行执行串行执行并行度低并行执行并行度高代码对比优化前load B 在前阻塞 load Atriton.jit def AB_load_kernel( A, B, B_index, O, B_DIM: tl.constexpr, HEAD_NUM: tl.constexpr, HEAD_DIM: tl.constexpr, ): i_n tl.program_id(0) i_range tl.arange(0, B_DIM) for i in range(HEAD_NUM): p_A A i * HEAD_DIM i_n * B_DIM i_range p_O O i * HEAD_DIM i_n * B_DIM i_range p_B_index B_index i idx_B tl.load(p_B_index) p_B B idx_B b_B tl.load(p_B) b_A tl.load(p_A) b_B tl.sum(b_A) b_O b_A * b_B tl.store(p_O, b_O) idx_B tl.load(p_B_index) p_B B idx_B tl.store(p_B, b_B)优化后load A 在前与 store B 并行triton.jit def AB_load_kernel( A, B, B_index, O, B_DIM: tl.constexpr, HEAD_NUM: tl.constexpr, HEAD_DIM: tl.constexpr, ): i_n tl.program_id(0) i_range tl.arange(0, B_DIM) for i in range(HEAD_NUM): p_A A i * HEAD_DIM i_n * B_DIM i_range p_O O i * HEAD_DIM i_n * B_DIM i_range p_B_index B_index i b_A tl.load(p_A) idx_B tl.load(p_B_index) p_B B idx_B b_B tl.load(p_B) b_B tl.sum(b_A) b_O b_A * b_B tl.store(p_O, b_O) idx_B tl.load(p_B_index) p_B B idx_B tl.store(p_B, b_B)关键点识别依赖关系分析循环内每条 load 与上一次循环 store 之间是否存在数据依赖无依赖 load 前置将与上次 store 无依赖的 load 提前到有依赖的 load 之前编译器不重排Triton 编译器不会自动调整 load 顺序必须由开发者手动重排不影响语义重排仅改变 load 的发射时机不改变计算逻辑模式 3避免循环内更新指针触发条件循环内使用ptr offset更新指针地址循环迭代次数较多 4的场景原理ptr offset产生读后写依赖RAW当前迭代的指针值依赖前一次迭代的写入结果。这导致编译器难以重排指令阻塞流水线无法充分利用昇腾 NPU 的多级流水架构。核心思路将地址的增量更新转换为绝对地址计算——ptr base i * offset每次迭代独立计算地址消除迭代间的 RAW 依赖。方面使用使用基地址 偏移量数据依赖存在 RAW 依赖阻塞流水线无依赖各次迭代独立计算指令调度编译器难以重排指令编译器可自由重排提升并行度流水优化难以充分利用 NPU 多级流水更好地匹配昇腾流水架构向量化可能阻碍向量化优化有利于 SIMD 向量化代码对比优化前使用更新指针offset_xm pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offset_wn pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) offset_k tl.arange(0, BLOCK_SIZE_K) a_ptrs_base x_ptr (offset_xm[:, None] * K offset_k[None, :]) b_ptrs_base w_ptr (offset_k[:, None] * N offset_wn[None, :]) msk_m offset_xm M msk_n offset_wn N x_ptrs a_ptrs_base w_ptrs b_ptrs_base accumulator tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): x tl.load( x_ptrs, maskmsk_m[:, None] (offset_k[None, :] K - k * BLOCK_SIZE_K), other0.0, ) w tl.load( w_ptrs, maskmsk_n[None, :] (offset_k[:, None] K - k * BLOCK_SIZE_K), other0.0, ) accumulator tl.dot(x, w) x_ptrs BLOCK_SIZE_K w_ptrs BLOCK_SIZE_K * N优化后使用基地址 偏移量offset_xm pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offset_wn pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) offset_k tl.arange(0, BLOCK_SIZE_K) a_ptrs_base x_ptr (offset_xm[:, None] * K offset_k[None, :]) b_ptrs_base w_ptr (offset_k[:, None] * N offset_wn[None, :]) msk_m offset_xm M msk_n offset_wn N accumulator tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): x_ptrs a_ptrs_base k * BLOCK_SIZE_K w_ptrs b_ptrs_base k * BLOCK_SIZE_K * N x tl.load( x_ptrs, maskmsk_m[:, None] (offset_k[None, :] K - k * BLOCK_SIZE_K), other0.0, ) w tl.load( w_ptrs, maskmsk_n[None, :] (offset_k[:, None] K - k * BLOCK_SIZE_K), other0.0, ) accumulator tl.dot(x, w)关键点基地址在循环外计算a_ptrs_base和b_ptrs_base在循环外计算一次避免重复计算循环内独立计算地址ptr base k * offset消除迭代间 RAW 依赖不要丢失基地址确保偏移量计算是base offset而非仅offset多维地址同理w_ptrs b_ptrs_base k * BLOCK_SIZE_K * N不同维度使用对应步长模式 4入参静态化tl.constexpr触发条件kernel 入参中存在运行时不变的固定值参数典型参数BLOCK_SIZEBLOCK_M、BLOCK_N、BLOCK_K、STRIDEstride_m、stride_n等如果不确定参数是否固定应询问用户是否可设为tl.constexpr原理将固定数值的入参声明为tl.constexpr编译器可在编译期进行常量折叠constant folding和常量传播constant propagation生成更优的指令序列。未声明为tl.constexpr的参数即使在运行时是固定值编译器也无法利用其常量性质进行优化。优化效果启用编译时常量折叠帮助编译器进行更 aggressive 的常量传播减少运行时分支判断开销可能影响 tiling 策略和内存布局选择代码对比优化前stride 为普通入参triton.jit def kernel( A, B, C, M, N, K, stride_am, stride_an, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): offset_am pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offset_an tl.arange(0, BLOCK_SIZE_K) a_ptrs A offset_am[:, None] * stride_am offset_an[None, :] * stride_an # ...优化后stride 声明为tl.constexprtriton.jit def kernel( A, B, C, M, N, K, stride_am: tl.constexpr, stride_an: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): offset_am pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offset_an tl.arange(0, BLOCK_SIZE_K) a_ptrs A offset_am[:, None] * stride_am offset_an[None, :] * stride_an # ...关键点只有运行时不变的参数才适合tl.constexpr要求参数在编译期已知动态参数不可使用优先对性能敏感参数使用BLOCK_SIZE、stride 等直接影响内存访问模式和 tiling 策略的参数应优先考虑不确定时询问用户如果不确定某个参数是否在所有调用中都是固定值应询问用户确认声明位置tl.constexpr在函数签名中通过类型注解声明格式为param_name: tl.constexpr910_95 特别注意向量宽度有限910_95 的 Vector Core 单次处理 256 bit32 Bytes标量化退化对性能影响尤为严重。模式 1Mask 替代%在 910_95 上收益最大可能达到 5-10x 加速。UB 容量仅 248KB910_95 UB 容量较小248KB 256KB - 8KB 预留向量化访存对 UB 利用率至关重要。取余导致的离散访存会浪费宝贵的 UB 带宽。多级流水架构910_95 采用 Cube 2 Vector 架构流水线并行能力强但依赖正确调度。模式 2Load 重排序和模式 3基地址 偏移量能更好地匹配昇腾流水架构减少流水线气泡。L0C 容量 256KB910_95 的 L0C 为 256KB相比 910B 的 128KB 翻倍矩阵乘法 accumulator 可充分利用。配合模式 4tl.constexpr编译器可更精确地规划 L0C tiling 策略。对齐要求910_95 L1 对齐 32B、L0C 对齐 512B、UB 对齐 32B。连续地址访问模式 1 模式 3有助于满足对齐约束避免硬件异常。tl.constexpr与 tiling910_95 的 UB/L1/L0C 容量有限tl.constexpr使编译器能在编译期确定 BLOCK_SIZE从而精确计算 tiling 分块是否适配各级缓存。相关文档链接00-hardware-quick-ref.md — Ascend910_95 硬件速查手册【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考