融合 Attention Online Softmax 设计【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills定位说明本文档描述的是FlashAttention 融合场景中的 Online Softmax 设计涉及QK^T、running max/sum、P×V、O_acc、causal mask 和 workspace。当前templates/目录下的七个 Kernel未提供与此设计完全对应的融合 Attention Kernel。如需独立的 Online Softmax 实现不含QK^T/P×V参见 softmax_v2_ara_online.md——两者只共享在线 max/sum 数学基础不是同一个实现。1. 优化目标在 FlashAttention 等长序列场景中Softmax 的输入是逐 tile 生成的 $QK^T$ score规模为 $S \times S$。Naive 实现先计算完整矩阵再逐行做 Softmax需要 $O(S^2)$ 内存长序列下 SRAM 放不下。本设计将 Softmax 从全量计算后归一化改为逐 Tile 计算维护 running max/sum的增量算法。核心思想在 S2 方向分 Tile 计算 $Q_i K_j^T$跨 Tile 维护当前见过的最大值 $m$ 和累积和 $l$输出 $O$ 增量更新。内存复杂度从 $O(S^2)$ 降至 $O(S)$。指标naiveoptimized收益中间内存占用$O(S^2)$需完整 $QK^T$ 矩阵$O(S)$仅当前 tile running 状态长序列下内存大幅降低max/sum 计算双 pass先 max 再 sum单 passtile 内即时更新天然 Safe Softmax输出更新最后一次性 prob × V每 tile 增量加权累加无需显式分配 prob 矩阵适用场景FlashAttention / Sparse FlashAttention 等融合 Attention 算子中内嵌的 Online Softmax。2. 架构概览2.1 存储层级与数据流GM (Global Memory) │ │ MTE2 (Q_i, K_j, V_j 从 GM → L1/UB) ▼ ┌─────────────────────────────────────────────────────────────────────────────┐ │ SRAM / UB │ │ ┌──────────┐ ┌─────────────┐ ┌──────────────┐ ┌─────────────────┐ │ │ │ Q_i Tile │ × │ K_j Tile^T │ │ S_ij Tile │ │ O_i (output) │ │ │ │ [B,S1,D]│ │ [D,S2tile]│ │ [B,S1,S2tile]│ → │ [B,S1,D] │ │ │ └──────────┘ └─────────────┘ └──────────────┘ └─────────────────┘ │ │ │ ▲ │ │ │ ┌─────────────────────────────┐ │ │ │ └───────→│ Online Softmax (Vector PIPE)│───────┘ │ │ │ running m: max so far │ │ │ │ running l: sum of exp │ │ │ │ O_acc: weighted output │ │ │ └─────────────────────────────┘ │ └─────────────────────────────────────────────────────────────────────────────┘ │ │ MTE3 (O_i, m_final, l_final → GM / workspace) ▼ GM (Output)2.2 Tiled 数据流外循环沿 S2 维度分 Tile$j 0, 1, \dots, N_{tile}-1$K、V 以统一块大小 $B$ 划分。每 Tile 内计算 $S_{ij} Q_i K_j^T / \sqrt{d}$然后做 Online Softmax 更新 running 状态。跨 Tile 状态传递$m$running max和 $l$running sum作为向量在 Tile 间传递。Ascend 场景下通常需 GM workspace 保存跨 S2 tile 的中间态。输出 O 增量更新每轮用当前 tile 的 Softmax 权重 $P_{ij}$ 对 $V_j$ 做加权按修正因子 rescale 后累加到 $O_i$。2.3 核心数学Online Softmax 单 Pass 公式对当前 Q block $Q_i$ 与 K block $K_j$ 计算Step 1 — Score$$S_{ij} Q_i \times K_j^T / \sqrt{d}$$Step 2 — 局部 max 与更新 running max$$m_{new} \max(m_{old}, \ \text{rowmax}(S_{ij}))$$Step 3 — 概率指数中间变量$$P_{ij} \exp(S_{ij} - m_{new})$$Step 4 — 更新 running sum$$l_{new} l_{old} \times \exp(m_{old} - m_{new}) \text{rowsum}(P_{ij})$$Step 5 — 输出 O 增量更新$$O_i O_i \times \frac{l_{old} \times \exp(m_{old} - m_{new})}{l_{new}} \frac{P_{ij} \times V_j}{l_{new}}$$等价概括形式m_new max(m_old, tile_max) sum_new sum_old * exp(m_old - m_new) tile_sum_exp(tile - m_new)最终 Softmax$\text{softmax}(x) \exp(x - m_{final}) / sum_{final}$2.4 事件同步模型事件类型含义用途MTE2_VMTE2 搬运完成 → 允许 Vector 读取Q/K/V tile 数据就绪V_MTE3Vector 完成 → 允许 MTE3 写回O 增量计算完成可写 GMV_VVector 完成 → 允许 Vector 继续Online softmax 内部依赖3. 关键参数配置// Host 侧 TilingData struct OnlineSoftmaxTiling { uint32_t B; // tile 块大小统一划分 K、V通常 64 uint32_t D; // head dimension uint32_t seqLen; // 序列长度 S uint32_t s2TileNum; // S2 方向 tile 数量 ceil(S / B) }; // Kernel 侧 running 状态每行一个 TBufQuePosition::VECCALC runningMaxBuff; // m: running max TBufQuePosition::VECCALC runningSumBuff; // l: running sum TBufQuePosition::VECCALC outAccBuff; // O_acc: 加权输出累加器3.1 Tile 大小选取原则参数典型值说明$B$64 / 128统一 tile 块大小划分 K、V 沿 S2 维度。需 fit UBD64 / 128head dimension由模型决定Ascend Cube 粒度对齐约束Tile 大小必须对齐 Cube 单元 matmul 粒度通常为16×16 或 32×32。若 $B$ 未对齐CUBE 核心需额外 padding。B align_up(preferred_B, cube_granularity) // typically 16 or 323.2 内存预算Online Softmax 将内存复杂度从 $O(S^2)$ 降至 $O(S)$但 score tile$B \times B$、Q tile、K/V tile 同时存在于 UB 时需确保总占用小于 UB 容量。典型配置下约 65KB安全。4. 核心计算循环4.1 naive 版本优化前// 阶段 1完整计算 QK^TS×S 矩阵需要 O(S^2) 空间 for (uint32_t i 0; i S; i) { for (uint32_t j 0; j S; j) { scoreGm[i * S j] ComputeScore(QGm, KGm, i, j); } } // 阶段 2逐行 Softmax需两次遍历先求 max再求 exp/sum for (uint32_t i 0; i S; i) { float maxVal -INFINITY; for (uint32_t j 0; j S; j) { maxVal max(maxVal, scoreGm[i * S j]); } float sum 0.0f; for (uint32_t j 0; j S; j) { sum exp(scoreGm[i * S j] - maxVal); } for (uint32_t j 0; j S; j) { probGm[i * S j] exp(scoreGm[i * S j] - maxVal) / sum; } } // 阶段 3prob × V for (uint32_t i 0; i S; i) { for (uint32_t d 0; d D; d) { float acc 0.0f; for (uint32_t j 0; j S; j) { acc probGm[i * S j] * VGm[j * D d]; } OGm[i * D d] acc; } }4.2 optimized 版本优化后Tiled Online Softmax// Init分配 running 状态 buffer每行一个 pipe-InitBuffer(runningMaxBuff, B * sizeof(COMPUTE_T)); // m pipe-InitBuffer(runningSumBuff, B * sizeof(COMPUTE_T)); // l pipe-InitBuffer(outAccBuff, B * D * sizeof(COMPUTE_T)); // O_acc LocalTensorCOMPUTE_T mUb runningMaxBuff.GetCOMPUTE_T(); LocalTensorCOMPUTE_T lUb runningSumBuff.GetCOMPUTE_T(); LocalTensorCOMPUTE_T oUb outAccBuff.GetCOMPUTE_T(); // 初始化 running 状态 Duplicate(mUb, FLOAT_NEG_INF, B); // m -inf Duplicate(lUb, 0.0f, B); // l 0 Duplicate(oUb, 0.0f, B * D); // O 0 // MTE2 加载 Q_i tile固定在 UB LocalTensorT qUb LoadQTile(QGm, i * B, B, D); // S2 方向分 tile 循环Online Softmax 核心 for (uint32_t j 0; j s2TileNum; j) { // 1. MTE2 加载 K_j, V_j tiles LocalTensorT kUb LoadKTile(KGm, j * B, B, D); LocalTensorT vUb LoadVTile(VGm, j * B, B, D); // 2. 计算 S_ij Q_i * K_j^T (CUBE / Vector) LocalTensorCOMPUTE_T sUb ComputeScore(qUb, kUb, B, D); // 3. Online Softmax 更新公式见 2.3 OnlineSoftmaxUpdate(sUb, vUb, mUb, lUb, oUb, B, D); // 4. 可选MTE3 将中间结果写回 GM workspace跨 tile 状态传递 } // 所有 tile 结束后O_final O_acc / l_final ElemwiseDiv(oUb, lUb, B, D); // MTE3 写回最终输出 WriteBackOGm(OGm, oUb, i * B, B, D);4.3 Online Softmax Update 伪代码Vector 核心void OnlineSoftmaxUpdate(LocalTensorCOMPUTE_T sUb, // [B, B] score tile S_ij LocalTensorT vUb, // [B, D] V tile V_j LocalTensorCOMPUTE_T mUb, // [B] running max LocalTensorCOMPUTE_T lUb, // [B] running sum LocalTensorCOMPUTE_T oUb, // [B, D] output acc O_i uint32_t B, uint32_t D) { for (uint32_t r 0; r B; r) { // Step 1: m_new max(m_old, rowmax(S_ij)) float m_old mUb[r]; float m_local ReduceMax(sUb[r * B], B); float m_new max(m_old, m_local); // Step 2: P_ij exp(S_ij - m_new) // Step 3: l_new l_old * exp(m_old - m_new) rowsum(P_ij) float l_old lUb[r]; float scale exp(m_old - m_new); float sum_exp 0.0f; LocalTensorCOMPUTE_T pUb tempBuff.GetCOMPUTE_T(); for (uint32_t c 0; c B; c) { float p_val exp(sUb[r * B c] - m_new); pUb[c] p_val; sum_exp p_val; } lUb[r] l_old * scale sum_exp; // Step 4: O_i O_i * (l_old * scale / l_new) P_ij * V_j / l_new float rescale_o (l_old * scale) / lUb[r]; float rescale_p 1.0f / lUb[r]; for (uint32_t d 0; d D; d) { float pv 0.0f; for (uint32_t c 0; c B; c) { pv pUb[c] * vUb[c * D d]; } oUb[r * D d] oUb[r * D d] * rescale_o pv * rescale_p; } mUb[r] m_new; } }5. 从 naive 到 online_softmax 的关键修改点修改项naive优化前online_softmax优化后计算范式先完整算 $QK^T$再 Softmax再 ×V逐 tile 计算增量更新 running 状态内存复杂度$O(S^2)$需完整矩阵$O(S)$仅当前 tile running m/l/Omax/sum 计算双 pass先 max 再 sum单 passtile 内即时更新数值稳定性减 max 需额外遍历天然 Safe Softmax每步减当前 m_new输出更新最后一次性 prob × V每 tile 增量加权累加 $P_{ij} \cdot V_j$ 到 O跨 tile 状态无全量后一次性处理running m、running l、O_acc跨 tile 传递6. 注意事项 / 约束数值稳定性必须维护 running max。$m_{new} \max(m_{old}, m_{local})$所有指数计算以 $m_{new}$ 为基准减。O_acc 的修正因子。当 $m_{new} m_{old}$ 时$O_{old}$ 和 $l_{old}$ 需乘 $e^{m_{old} - m_{new}}$scale 必须在更新前计算。与 GM workspace 的关系。Flash Attention 通常需要 workspace 保存跨 tile 的 running m/l。本优化是算法层面的内存降低workspace 属于工程层面的跨调用状态传递。last tile 的边界处理。当 $S$ 不是 $B$ 整数倍时最后一个 tile 的有效列数小于 $B$避免 padding 值干扰。精度与性能的平衡。$m$ 和 $l$ 建议用 FP32 维护即使输入/输出是 FP16。Ascend Cube-Vector 负载均衡。在 Ascend 910B 上$QK^T$ 运行在 CUBE 核心softmax 和 $P \times V$ 运行在 Vector 核心优化时需兼顾两端。Workspace MTE3 写回是显著开销。running 状态跨 S2 tile 传递通常需要 GM workspaceMTE3 向 workspace 的写回是性能瓶颈之一。Causal Mask 的 Block-level Skip。Decoder 场景下 causal mask 可在 tile 级别做粗粒度跳过若 $s2_start s1_end$SKIP整个 block若 $s2_end \le s1_start$FULLblock否则PARTIALblock需 element-level masking理论跳过率$(n-1)/(2n) \to 50%$。实际$S32K$ 约 45%$S2K$ 约 37%$S512$ 约 22%。KV-CachePrefill vs Decode 的差异。Prefill 阶段$Q$ 为完整序列需完整 tiling。Decode 阶段$Q$ 长度为 1无需在 $Q$ 维度 tiling仅沿 KV 序列维度分 tile完全 memory-bound。Sliding Window Attention。若使用滑动窗口仅 attend 最近 $W$ 个 tokenvalid 条件为 $i - W j \le i$。Block-level 可额外跳过 $s2_end s1_start - W$ 的 block。preTokens / nextTokens 参数。Ascend Attention API 通过preTokens和nextTokens定义 attention 窗口等效于 band mask。配置时需确保与 causal / sliding window 逻辑一致。7. 选型决策与自检清单7.1 选型决策if (融合 Attention 算子 输入序列长度 S 512): → 启用 online_softmaxtiled 实现 → B 64 或 128对齐 Cube 粒度 → running 状态用 FP32 else: → 标准 Safe Softmax 即可S 小全量矩阵 fit SRAM8. AscendC Kernel 优化实现上述设计在 Python 层逐 tile 实现时受限于 Python → NPU 的调度开销和独立的 kernel launch延迟数据不能反映算法真实优势。生产级优化需在 AscendC 层面将 QK^T、Online Softmax 更新、PV 乘积全部融合到单个 kernel 内。8.1 从 Python 层到 AscendC 的关键跨越维度Python 层实现AscendC Fused Kernel调度开销16~128 次 Python 循环 kernel launch单次 kernel launch零 Python 开销Score 计算torch.matmul独立调用内联MulReduceSumVector PIPESoftmaxtorch.max/torch.exp/torch.sumReduceMax→Adds→Exp→ReduceSumP Vtorch.matmul独立调用内联标量累加V 为行优先列访问需 strided并行度batch 维度串行Multi-block每个 AICore 处理一个或多个 Q rowCausal SkipPythonif判断Block-levelcontinue跳过整 tile8.2 AscendC Kernel 架构每个 AI Core Block 处理一个或多个 Q rows ├─ Load Q row [D] → UB ├─ Cast FP16 → FP32, Muls(scale) ├─ Init: m-inf, l0, O_acc0 ├─ For each S2 tile j: │ ├─ Causal check: tileStart s1Idx ? skip │ ├─ MTE2: Load K tile [actualB, D] → UB │ ├─ Vector: Cast K→FP32, Mul(q, k_slice), ReduceSum → score[actualB] │ ├─ MTE2: Load V tile [actualB, D] → UB │ ├─ Causal: mask upper part to -inf (if tile crosses diagonal) │ ├─ Vector: ReduceMax(score) → m_local │ ├─ Vector: Adds(score, -m_new), Exp → P[actualB] │ ├─ Vector: ReduceSum(P) → sum_exp │ ├─ Scalar: update m, l, rescale factors │ ├─ Vector: Muls(O_acc, rescale_o) │ └─ ScalarVector: P V column-by-column, accumulate to O_acc ├─ Cast O_acc → FP16 ├─ MTE3: Write O row → GM └─ MTE3: Write L l_final → GM8.3 Vector API 替换映射// Score: Q K^T (原为标量循环) Cast(kFloat, kLocal, ..., actualB * D); for (b 0; b actualB; b) { Mul(tmp, qFloat, kFloat[b * D], D); // Vector 逐元素乘 ReduceSum(sumBuf, tmp, sumBuf, D); // Vector 规约求和 score.SetValue(b, sumBuf.GetValue(0)); } // Softmax: max → exp → sum (原为标量循环) ReduceMax(maxBuf, score, maxBuf, actualB); // Vector 规约求最大 Adds(tmp, score, -mNew, actualB); // Vector 逐元素加 Exp(tmp, tmp, actualB); // Vector 指数 ReduceSum(sumBuf, tmp, sumBuf, actualB); // Vector 规约求和 // O rescale (原为标量循环) Muls(oAcc, oAcc, rescaleO, D); // Vector 逐元素乘标量8.4 多核并行策略总 row 数totalRows batch × num_heads × S1Block 分配blockDim min(totalRows, 8)匹配 Ascend 910B 的 8 AICore每 Block 工作量rowsPerCore ceil(totalRows / blockDim)Flat index 分解row → (batchIdx, headIdx, s1Idx)8.5 UB 内存预算优化后Buffer用途D64, tileB128D128, tileB128qBufQ row (FP16)128 B256 BkBuf/vBufK/V tile (FP16)2 × 16 KB2 × 32 KBqFloatBufQ row (FP32)256 B512 BkvFloatBufK/V float (时间复用)32 KB64 KBscoreBuf/pBufScore / Prob (FP32)2 × 512 B2 × 512 BoAccBufOutput accumulator256 B512 BtmpBuf/reduceBufVector workspace2 × 512 B2 × 512 B总计~100 KB~135 KBUB 容量 192 KB两种配置均安全。K/V float buffer 时间复用节省 32~64 KB。8.6 预期性能指标Python 层 OnlineAscendC FusedS512慢于 naive 28x预计快 1.2~1.5xS4K慢于 naive 1.4x预计快 1.5~3xS32K慢于 naive 25x预计快 2~5x(causalskip)内存O(S)O(S)与 Python 层相同关键结论AscendC kernel 消除了 Python 调度开销使 Online Softmax 的带宽优势省去 O(S²) 中间矩阵搬运真正转化为延迟收益。7.2 自检清单running m 初始化为-infrunning l 初始化为0使用m_new max(m_old, m_local)更新$scale e^{m_{old} - m_{new}}$ 在更新 l 和 O 之前计算O_acc 和 running l 使用 FP32即使输入输出为 FP16Last tile 边界处理有效列数 min(B, seqLen - j * B)Causal mask 场景全 masked tile 正确跳过更新精度校验通过与 naive Safe Softmax 对比误差 1e-5FP32或 1e-3FP16AscendC kernel 使用 Vector API 替代标量循环GetValue/SetValue仅在必要处使用Multi-block 并行正确分配 Q rows无重复/遗漏UB 内存总量 192 KBAscend 910B9. 与独立 ARA Online Softmax 的关系本文档描述的融合 Attention Online Softmax 与 softmax_v2_ara_online.md 中的独立 ARA Online Softmax共享相同的在线 max/sum 数学基础但不是同一个实现特性融合 Attention Online Softmax本文档独立 ARA Online Softmax输入来源Attention 内部逐 tile 生成QK^TGM 中已存在的完整数据score 消费单次 score tile 消费生成即消费两次读取输入在线 maxsum 一遍输出一遍P×V 融合是P×V融合到循环中否O_acc是增量加权累加否输出完整 Softmax 概率完整概率矩阵通常不生成生成完整 Softmax 概率causal mask支持 block-level skip不涉及workspace需要跨 tile 状态传递不需要两者不得合并成同一个 Kernel 实现说明。本文档是融合 Attention 的设计参考当前目录未提供与之完全对应的融合 Attention Kernel。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN/cannbot-skills：融合Attention Online Softmax设计
发布时间：2026/7/15 18:00:08
融合 Attention Online Softmax 设计【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills定位说明本文档描述的是FlashAttention 融合场景中的 Online Softmax 设计涉及QK^T、running max/sum、P×V、O_acc、causal mask 和 workspace。当前templates/目录下的七个 Kernel未提供与此设计完全对应的融合 Attention Kernel。如需独立的 Online Softmax 实现不含QK^T/P×V参见 softmax_v2_ara_online.md——两者只共享在线 max/sum 数学基础不是同一个实现。1. 优化目标在 FlashAttention 等长序列场景中Softmax 的输入是逐 tile 生成的 $QK^T$ score规模为 $S \times S$。Naive 实现先计算完整矩阵再逐行做 Softmax需要 $O(S^2)$ 内存长序列下 SRAM 放不下。本设计将 Softmax 从全量计算后归一化改为逐 Tile 计算维护 running max/sum的增量算法。核心思想在 S2 方向分 Tile 计算 $Q_i K_j^T$跨 Tile 维护当前见过的最大值 $m$ 和累积和 $l$输出 $O$ 增量更新。内存复杂度从 $O(S^2)$ 降至 $O(S)$。指标naiveoptimized收益中间内存占用$O(S^2)$需完整 $QK^T$ 矩阵$O(S)$仅当前 tile running 状态长序列下内存大幅降低max/sum 计算双 pass先 max 再 sum单 passtile 内即时更新天然 Safe Softmax输出更新最后一次性 prob × V每 tile 增量加权累加无需显式分配 prob 矩阵适用场景FlashAttention / Sparse FlashAttention 等融合 Attention 算子中内嵌的 Online Softmax。2. 架构概览2.1 存储层级与数据流GM (Global Memory) │ │ MTE2 (Q_i, K_j, V_j 从 GM → L1/UB) ▼ ┌─────────────────────────────────────────────────────────────────────────────┐ │ SRAM / UB │ │ ┌──────────┐ ┌─────────────┐ ┌──────────────┐ ┌─────────────────┐ │ │ │ Q_i Tile │ × │ K_j Tile^T │ │ S_ij Tile │ │ O_i (output) │ │ │ │ [B,S1,D]│ │ [D,S2tile]│ │ [B,S1,S2tile]│ → │ [B,S1,D] │ │ │ └──────────┘ └─────────────┘ └──────────────┘ └─────────────────┘ │ │ │ ▲ │ │ │ ┌─────────────────────────────┐ │ │ │ └───────→│ Online Softmax (Vector PIPE)│───────┘ │ │ │ running m: max so far │ │ │ │ running l: sum of exp │ │ │ │ O_acc: weighted output │ │ │ └─────────────────────────────┘ │ └─────────────────────────────────────────────────────────────────────────────┘ │ │ MTE3 (O_i, m_final, l_final → GM / workspace) ▼ GM (Output)2.2 Tiled 数据流外循环沿 S2 维度分 Tile$j 0, 1, \dots, N_{tile}-1$K、V 以统一块大小 $B$ 划分。每 Tile 内计算 $S_{ij} Q_i K_j^T / \sqrt{d}$然后做 Online Softmax 更新 running 状态。跨 Tile 状态传递$m$running max和 $l$running sum作为向量在 Tile 间传递。Ascend 场景下通常需 GM workspace 保存跨 S2 tile 的中间态。输出 O 增量更新每轮用当前 tile 的 Softmax 权重 $P_{ij}$ 对 $V_j$ 做加权按修正因子 rescale 后累加到 $O_i$。2.3 核心数学Online Softmax 单 Pass 公式对当前 Q block $Q_i$ 与 K block $K_j$ 计算Step 1 — Score$$S_{ij} Q_i \times K_j^T / \sqrt{d}$$Step 2 — 局部 max 与更新 running max$$m_{new} \max(m_{old}, \ \text{rowmax}(S_{ij}))$$Step 3 — 概率指数中间变量$$P_{ij} \exp(S_{ij} - m_{new})$$Step 4 — 更新 running sum$$l_{new} l_{old} \times \exp(m_{old} - m_{new}) \text{rowsum}(P_{ij})$$Step 5 — 输出 O 增量更新$$O_i O_i \times \frac{l_{old} \times \exp(m_{old} - m_{new})}{l_{new}} \frac{P_{ij} \times V_j}{l_{new}}$$等价概括形式m_new max(m_old, tile_max) sum_new sum_old * exp(m_old - m_new) tile_sum_exp(tile - m_new)最终 Softmax$\text{softmax}(x) \exp(x - m_{final}) / sum_{final}$2.4 事件同步模型事件类型含义用途MTE2_VMTE2 搬运完成 → 允许 Vector 读取Q/K/V tile 数据就绪V_MTE3Vector 完成 → 允许 MTE3 写回O 增量计算完成可写 GMV_VVector 完成 → 允许 Vector 继续Online softmax 内部依赖3. 关键参数配置// Host 侧 TilingData struct OnlineSoftmaxTiling { uint32_t B; // tile 块大小统一划分 K、V通常 64 uint32_t D; // head dimension uint32_t seqLen; // 序列长度 S uint32_t s2TileNum; // S2 方向 tile 数量 ceil(S / B) }; // Kernel 侧 running 状态每行一个 TBufQuePosition::VECCALC runningMaxBuff; // m: running max TBufQuePosition::VECCALC runningSumBuff; // l: running sum TBufQuePosition::VECCALC outAccBuff; // O_acc: 加权输出累加器3.1 Tile 大小选取原则参数典型值说明$B$64 / 128统一 tile 块大小划分 K、V 沿 S2 维度。需 fit UBD64 / 128head dimension由模型决定Ascend Cube 粒度对齐约束Tile 大小必须对齐 Cube 单元 matmul 粒度通常为16×16 或 32×32。若 $B$ 未对齐CUBE 核心需额外 padding。B align_up(preferred_B, cube_granularity) // typically 16 or 323.2 内存预算Online Softmax 将内存复杂度从 $O(S^2)$ 降至 $O(S)$但 score tile$B \times B$、Q tile、K/V tile 同时存在于 UB 时需确保总占用小于 UB 容量。典型配置下约 65KB安全。4. 核心计算循环4.1 naive 版本优化前// 阶段 1完整计算 QK^TS×S 矩阵需要 O(S^2) 空间 for (uint32_t i 0; i S; i) { for (uint32_t j 0; j S; j) { scoreGm[i * S j] ComputeScore(QGm, KGm, i, j); } } // 阶段 2逐行 Softmax需两次遍历先求 max再求 exp/sum for (uint32_t i 0; i S; i) { float maxVal -INFINITY; for (uint32_t j 0; j S; j) { maxVal max(maxVal, scoreGm[i * S j]); } float sum 0.0f; for (uint32_t j 0; j S; j) { sum exp(scoreGm[i * S j] - maxVal); } for (uint32_t j 0; j S; j) { probGm[i * S j] exp(scoreGm[i * S j] - maxVal) / sum; } } // 阶段 3prob × V for (uint32_t i 0; i S; i) { for (uint32_t d 0; d D; d) { float acc 0.0f; for (uint32_t j 0; j S; j) { acc probGm[i * S j] * VGm[j * D d]; } OGm[i * D d] acc; } }4.2 optimized 版本优化后Tiled Online Softmax// Init分配 running 状态 buffer每行一个 pipe-InitBuffer(runningMaxBuff, B * sizeof(COMPUTE_T)); // m pipe-InitBuffer(runningSumBuff, B * sizeof(COMPUTE_T)); // l pipe-InitBuffer(outAccBuff, B * D * sizeof(COMPUTE_T)); // O_acc LocalTensorCOMPUTE_T mUb runningMaxBuff.GetCOMPUTE_T(); LocalTensorCOMPUTE_T lUb runningSumBuff.GetCOMPUTE_T(); LocalTensorCOMPUTE_T oUb outAccBuff.GetCOMPUTE_T(); // 初始化 running 状态 Duplicate(mUb, FLOAT_NEG_INF, B); // m -inf Duplicate(lUb, 0.0f, B); // l 0 Duplicate(oUb, 0.0f, B * D); // O 0 // MTE2 加载 Q_i tile固定在 UB LocalTensorT qUb LoadQTile(QGm, i * B, B, D); // S2 方向分 tile 循环Online Softmax 核心 for (uint32_t j 0; j s2TileNum; j) { // 1. MTE2 加载 K_j, V_j tiles LocalTensorT kUb LoadKTile(KGm, j * B, B, D); LocalTensorT vUb LoadVTile(VGm, j * B, B, D); // 2. 计算 S_ij Q_i * K_j^T (CUBE / Vector) LocalTensorCOMPUTE_T sUb ComputeScore(qUb, kUb, B, D); // 3. Online Softmax 更新公式见 2.3 OnlineSoftmaxUpdate(sUb, vUb, mUb, lUb, oUb, B, D); // 4. 可选MTE3 将中间结果写回 GM workspace跨 tile 状态传递 } // 所有 tile 结束后O_final O_acc / l_final ElemwiseDiv(oUb, lUb, B, D); // MTE3 写回最终输出 WriteBackOGm(OGm, oUb, i * B, B, D);4.3 Online Softmax Update 伪代码Vector 核心void OnlineSoftmaxUpdate(LocalTensorCOMPUTE_T sUb, // [B, B] score tile S_ij LocalTensorT vUb, // [B, D] V tile V_j LocalTensorCOMPUTE_T mUb, // [B] running max LocalTensorCOMPUTE_T lUb, // [B] running sum LocalTensorCOMPUTE_T oUb, // [B, D] output acc O_i uint32_t B, uint32_t D) { for (uint32_t r 0; r B; r) { // Step 1: m_new max(m_old, rowmax(S_ij)) float m_old mUb[r]; float m_local ReduceMax(sUb[r * B], B); float m_new max(m_old, m_local); // Step 2: P_ij exp(S_ij - m_new) // Step 3: l_new l_old * exp(m_old - m_new) rowsum(P_ij) float l_old lUb[r]; float scale exp(m_old - m_new); float sum_exp 0.0f; LocalTensorCOMPUTE_T pUb tempBuff.GetCOMPUTE_T(); for (uint32_t c 0; c B; c) { float p_val exp(sUb[r * B c] - m_new); pUb[c] p_val; sum_exp p_val; } lUb[r] l_old * scale sum_exp; // Step 4: O_i O_i * (l_old * scale / l_new) P_ij * V_j / l_new float rescale_o (l_old * scale) / lUb[r]; float rescale_p 1.0f / lUb[r]; for (uint32_t d 0; d D; d) { float pv 0.0f; for (uint32_t c 0; c B; c) { pv pUb[c] * vUb[c * D d]; } oUb[r * D d] oUb[r * D d] * rescale_o pv * rescale_p; } mUb[r] m_new; } }5. 从 naive 到 online_softmax 的关键修改点修改项naive优化前online_softmax优化后计算范式先完整算 $QK^T$再 Softmax再 ×V逐 tile 计算增量更新 running 状态内存复杂度$O(S^2)$需完整矩阵$O(S)$仅当前 tile running m/l/Omax/sum 计算双 pass先 max 再 sum单 passtile 内即时更新数值稳定性减 max 需额外遍历天然 Safe Softmax每步减当前 m_new输出更新最后一次性 prob × V每 tile 增量加权累加 $P_{ij} \cdot V_j$ 到 O跨 tile 状态无全量后一次性处理running m、running l、O_acc跨 tile 传递6. 注意事项 / 约束数值稳定性必须维护 running max。$m_{new} \max(m_{old}, m_{local})$所有指数计算以 $m_{new}$ 为基准减。O_acc 的修正因子。当 $m_{new} m_{old}$ 时$O_{old}$ 和 $l_{old}$ 需乘 $e^{m_{old} - m_{new}}$scale 必须在更新前计算。与 GM workspace 的关系。Flash Attention 通常需要 workspace 保存跨 tile 的 running m/l。本优化是算法层面的内存降低workspace 属于工程层面的跨调用状态传递。last tile 的边界处理。当 $S$ 不是 $B$ 整数倍时最后一个 tile 的有效列数小于 $B$避免 padding 值干扰。精度与性能的平衡。$m$ 和 $l$ 建议用 FP32 维护即使输入/输出是 FP16。Ascend Cube-Vector 负载均衡。在 Ascend 910B 上$QK^T$ 运行在 CUBE 核心softmax 和 $P \times V$ 运行在 Vector 核心优化时需兼顾两端。Workspace MTE3 写回是显著开销。running 状态跨 S2 tile 传递通常需要 GM workspaceMTE3 向 workspace 的写回是性能瓶颈之一。Causal Mask 的 Block-level Skip。Decoder 场景下 causal mask 可在 tile 级别做粗粒度跳过若 $s2_start s1_end$SKIP整个 block若 $s2_end \le s1_start$FULLblock否则PARTIALblock需 element-level masking理论跳过率$(n-1)/(2n) \to 50%$。实际$S32K$ 约 45%$S2K$ 约 37%$S512$ 约 22%。KV-CachePrefill vs Decode 的差异。Prefill 阶段$Q$ 为完整序列需完整 tiling。Decode 阶段$Q$ 长度为 1无需在 $Q$ 维度 tiling仅沿 KV 序列维度分 tile完全 memory-bound。Sliding Window Attention。若使用滑动窗口仅 attend 最近 $W$ 个 tokenvalid 条件为 $i - W j \le i$。Block-level 可额外跳过 $s2_end s1_start - W$ 的 block。preTokens / nextTokens 参数。Ascend Attention API 通过preTokens和nextTokens定义 attention 窗口等效于 band mask。配置时需确保与 causal / sliding window 逻辑一致。7. 选型决策与自检清单7.1 选型决策if (融合 Attention 算子 输入序列长度 S 512): → 启用 online_softmaxtiled 实现 → B 64 或 128对齐 Cube 粒度 → running 状态用 FP32 else: → 标准 Safe Softmax 即可S 小全量矩阵 fit SRAM8. AscendC Kernel 优化实现上述设计在 Python 层逐 tile 实现时受限于 Python → NPU 的调度开销和独立的 kernel launch延迟数据不能反映算法真实优势。生产级优化需在 AscendC 层面将 QK^T、Online Softmax 更新、PV 乘积全部融合到单个 kernel 内。8.1 从 Python 层到 AscendC 的关键跨越维度Python 层实现AscendC Fused Kernel调度开销16~128 次 Python 循环 kernel launch单次 kernel launch零 Python 开销Score 计算torch.matmul独立调用内联MulReduceSumVector PIPESoftmaxtorch.max/torch.exp/torch.sumReduceMax→Adds→Exp→ReduceSumP Vtorch.matmul独立调用内联标量累加V 为行优先列访问需 strided并行度batch 维度串行Multi-block每个 AICore 处理一个或多个 Q rowCausal SkipPythonif判断Block-levelcontinue跳过整 tile8.2 AscendC Kernel 架构每个 AI Core Block 处理一个或多个 Q rows ├─ Load Q row [D] → UB ├─ Cast FP16 → FP32, Muls(scale) ├─ Init: m-inf, l0, O_acc0 ├─ For each S2 tile j: │ ├─ Causal check: tileStart s1Idx ? skip │ ├─ MTE2: Load K tile [actualB, D] → UB │ ├─ Vector: Cast K→FP32, Mul(q, k_slice), ReduceSum → score[actualB] │ ├─ MTE2: Load V tile [actualB, D] → UB │ ├─ Causal: mask upper part to -inf (if tile crosses diagonal) │ ├─ Vector: ReduceMax(score) → m_local │ ├─ Vector: Adds(score, -m_new), Exp → P[actualB] │ ├─ Vector: ReduceSum(P) → sum_exp │ ├─ Scalar: update m, l, rescale factors │ ├─ Vector: Muls(O_acc, rescale_o) │ └─ ScalarVector: P V column-by-column, accumulate to O_acc ├─ Cast O_acc → FP16 ├─ MTE3: Write O row → GM └─ MTE3: Write L l_final → GM8.3 Vector API 替换映射// Score: Q K^T (原为标量循环) Cast(kFloat, kLocal, ..., actualB * D); for (b 0; b actualB; b) { Mul(tmp, qFloat, kFloat[b * D], D); // Vector 逐元素乘 ReduceSum(sumBuf, tmp, sumBuf, D); // Vector 规约求和 score.SetValue(b, sumBuf.GetValue(0)); } // Softmax: max → exp → sum (原为标量循环) ReduceMax(maxBuf, score, maxBuf, actualB); // Vector 规约求最大 Adds(tmp, score, -mNew, actualB); // Vector 逐元素加 Exp(tmp, tmp, actualB); // Vector 指数 ReduceSum(sumBuf, tmp, sumBuf, actualB); // Vector 规约求和 // O rescale (原为标量循环) Muls(oAcc, oAcc, rescaleO, D); // Vector 逐元素乘标量8.4 多核并行策略总 row 数totalRows batch × num_heads × S1Block 分配blockDim min(totalRows, 8)匹配 Ascend 910B 的 8 AICore每 Block 工作量rowsPerCore ceil(totalRows / blockDim)Flat index 分解row → (batchIdx, headIdx, s1Idx)8.5 UB 内存预算优化后Buffer用途D64, tileB128D128, tileB128qBufQ row (FP16)128 B256 BkBuf/vBufK/V tile (FP16)2 × 16 KB2 × 32 KBqFloatBufQ row (FP32)256 B512 BkvFloatBufK/V float (时间复用)32 KB64 KBscoreBuf/pBufScore / Prob (FP32)2 × 512 B2 × 512 BoAccBufOutput accumulator256 B512 BtmpBuf/reduceBufVector workspace2 × 512 B2 × 512 B总计~100 KB~135 KBUB 容量 192 KB两种配置均安全。K/V float buffer 时间复用节省 32~64 KB。8.6 预期性能指标Python 层 OnlineAscendC FusedS512慢于 naive 28x预计快 1.2~1.5xS4K慢于 naive 1.4x预计快 1.5~3xS32K慢于 naive 25x预计快 2~5x(causalskip)内存O(S)O(S)与 Python 层相同关键结论AscendC kernel 消除了 Python 调度开销使 Online Softmax 的带宽优势省去 O(S²) 中间矩阵搬运真正转化为延迟收益。7.2 自检清单running m 初始化为-infrunning l 初始化为0使用m_new max(m_old, m_local)更新$scale e^{m_{old} - m_{new}}$ 在更新 l 和 O 之前计算O_acc 和 running l 使用 FP32即使输入输出为 FP16Last tile 边界处理有效列数 min(B, seqLen - j * B)Causal mask 场景全 masked tile 正确跳过更新精度校验通过与 naive Safe Softmax 对比误差 1e-5FP32或 1e-3FP16AscendC kernel 使用 Vector API 替代标量循环GetValue/SetValue仅在必要处使用Multi-block 并行正确分配 Q rows无重复/遗漏UB 内存总量 192 KBAscend 910B9. 与独立 ARA Online Softmax 的关系本文档描述的融合 Attention Online Softmax 与 softmax_v2_ara_online.md 中的独立 ARA Online Softmax共享相同的在线 max/sum 数学基础但不是同一个实现特性融合 Attention Online Softmax本文档独立 ARA Online Softmax输入来源Attention 内部逐 tile 生成QK^TGM 中已存在的完整数据score 消费单次 score tile 消费生成即消费两次读取输入在线 maxsum 一遍输出一遍P×V 融合是P×V融合到循环中否O_acc是增量加权累加否输出完整 Softmax 概率完整概率矩阵通常不生成生成完整 Softmax 概率causal mask支持 block-level skip不涉及workspace需要跨 tile 状态传递不需要两者不得合并成同一个 Kernel 实现说明。本文档是融合 Attention 的设计参考当前目录未提供与之完全对应的融合 Attention Kernel。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
)
