前言

大模型推理和训练中有一道硬门槛：当序列长度翻倍，标准注意力机制的显存开销翻四倍。对于一张 80GB 显存的昇腾 NPU 卡，序列长度达到 64K 时标准注意力就会吃掉大部分显存。到了 128K 的上下文长度，一张卡根本放不下一个注意力层。CANN 在面向 Transformer 大模型的第二层 AOL（Ascend Operator Library）中构建了 ops-transformer 仓库，它的 FlashAttentionScore 算子通过分块 tiling 计算和中间结果重算策略，把注意力机制的显存复杂度从 O(N平方) 压到 O(N)，对应的显存从 16GB 级别降到了 KB 级别。这个跨越让 128K、256K、甚至更长上下文的推理和训练在单卡上成为工程现实。

ops-transformer 是整个 CANN 软件栈中专攻 Transformer 架构中高阶算子的一环。它位于 opbase 基础算子库之上，ATB 加速库之下——opbase 提供矩阵乘、卷积、归一化这些底层原语，ops-transformer 组装注意力、MoE、融合通信等高级模式，而 ATB 负责编排完整 Transformer 层的计算图。仓库内部包含 FlashAttentionScore、SparseFlashAttention、GatherPAKVCache、MoEComputeExpertTokens、GroupedMatMulSwiGLUQuantV2 以及 MC2 通算融合算子群（MatmulAlltoAll、AttentionToFFN、FFNToAttention）。这些算子处理的共同挑战是同一个问题：在 Transformer 大模型的尺度下，如何把计算和内存的浪费压到最小。

平方级显存的精确来源

Scaled Dot-Product Attention 的数学公式简朴到只有三步：用 Q 和 K 的点积算出每个 token 之间的相关性分数，过 softmax 归一化到概率分布，再按概率对 V 做加权求和。

用矩阵语言表达就是：

S = Q @ K^T / sqrt(d_k)
P = Softmax(S)
O = P @ V

隐藏在这条式子后面的显存放大效应来自第一步和第二步之间的中间矩阵 S。S 的形状是 [batch, heads, N, N]，其中 N 是序列长度。对于 batch=1, heads=32, N=4096，S 占用 32 x 4096 x 4096 x 2 字节（bfloat16）约等于 1GB。N 到 8192 时变成 4GB。N 到 16384 时变成 16GB。N 到 65536 时变成 256GB——单卡显存物理上不可能容纳。

这个爆炸不仅来自 S 本身。标准实现在计算过程中至少需要分配以下内存区域：

第一块区域存放 S（N x N x batch x heads）。第二块区域存放 S 做 softmax 后得到的 P（同样大小的 N x N）。第三块区域在反向传播时需要保存 S 用于梯度计算（前向训练时）。也就是说，在训练场景下标准注意力的显存占用至少是 3 x N x N x batch x heads x sizeof(dtype)。batch 和 heads 都是乘积因子，但在大模型实践中 batch 通常为 1（在线推理）或较小值（单 batch 训练），因此序列长度 N 是主导因子。

没有 FlashAttention 时，工程团队能做什么？压低 batch size、减小 head 数、用 Ring Attention 让多卡分担、或者限制序列长度。这些都是在"平方级"约束下的权宜之计——你不能真正突破 O(N平方)，只能在常数因子上做修正。

解决思路的物理基础：SRAM 比 HBM 快几十倍

昇腾 NPU 的存储层次和 GPU 类似：片上 SRAM（也称共享内存或 L1 缓存）是容量小但带宽极高的快速存储，片外 HBM（高带宽内存）是容量大但带宽慢两个数量级的慢速存储。具体数量级上，HBM 的带宽在 TB/s 级别，而 SRAM 的带宽在几十 TB/s 级别。访问一次 HBM 的能耗和延迟相当于几十次 SRAM 访问。

标准注意力在这个层次上的行为是：两个矩阵乘法（Q @ K^T 和 P @ V）之间出现了一个巨大的中间张量 S，它必须穿过 HBM 的带宽瓶颈——先在乘法结束后由 Cube Unit 把 S 写入 HBM，再由 Vector Unit 从 HBM 读回做 softmax，再将 softmax 结果 P 写回 HBM，末尾由乘法单元读回做加权求和。一条计算链上出现三次 HBM 全量写入和三次全量读取。

FlashAttention 的核心洞察极其直接：如果在 SRAM 上完成 S→P→O 的整条计算链，只在末尾将结果 O 写回 HBM，这三次 HBM 全量写入和三次全量读取就压缩为一次写入。代价是 SRAM 容量有限——装不下完整的 N x N 矩阵——必须分块处理，且在分块边界上需要额外的合并计算。

用一个非技术类比的例子来理解：你在一间大图书馆（HBM）里做文献综述。需要翻阅 100 本书并交叉引用。方案 A 是把这 100 本书全部搬到大桌上（HBM），在桌子上摊开查阅。桌子很大能放下所有书，但图书馆到桌子的搬书路径很慢。方案 B 是每次只搬 3 本书到自己的小书桌（SRAM）上查阅，查完放回去再搬下一批，如果发现某一本之前读过的书需要重新引用，就去书架重新翻一次。方案 B 在搬书上花的时间较少（因为每次搬的距离短），但可能会因为需要重复翻阅而浪费一些精力。实际测试表明，当"搬书的代价远高于重复翻阅的代价"时，方案 B 总体更快——FlashAttention 就是在这个条件下成立的。

分块 Tiling 的计算骨架

FlashAttention 把 Q、K、V 在序列维度上分成固定大小的 tile（块）。块的大小由昇腾 NPU 的 SRAM 容量决定。每个 Q tile 遍历所有 KV tile，在 SRAM 上计算注意力分数、softmax、加权求和，据此累加到输出缓冲区。当一个 Q tile 处理完所有 KV tile 后，将累加的最终输出写回 HBM。

以下代码抽象了 ops-transformer 中 FlashAttentionScore 算子的核心 tiling 循环。真实的实现可能使用更底层的 TBE DSL 或 Ascend C 编写，但控制流与此一致：

def flash_attention_forward_tiled(Q, K, V, block_Br=128, block_Bc=128):
    """
    FlashAttention 前向计算——分块 tiling 核心循环

    Q, K, V: [batch, heads, seq_len, head_dim]
    block_Br: Q 在序列维度的 tile 大小
    block_Bc: K/V 在序列维度的 tile 大小

    显存分配总量: O(batch * heads * seq_len * head_dim)
    不包含 O(seq_len^2) 的中间张量
    """
    B, H, N, d = Q.shape
    Tr = ceil_div(N, block_Br)
    Tc = ceil_div(N, block_Bc)

    # 输出缓冲区，驻留在 HBM
    O = torch.zeros(B, H, N, d, device=Q.device)
    # 每行的全局最大值（用于在线 softmax）
    row_max = torch.full((B, H, N, 1), -float('inf'), device=Q.device)
    # 每行的全局归一化和
    row_sum = torch.zeros(B, H, N, 1, device=Q.device)

    for i in range(Tr):
        q_i = Q[:, :, i*block_Br:(i+1)*block_Br, :]   # [B, H, Br, d]

        # 初始化当前 tile 的在线 softmax 统计量
        m_prev = row_max[:, :, i*block_Br:(i+1)*block_Br, :]
        l_prev = row_sum[:, :, i*block_Br:(i+1)*block_Br, :]
        o_i = O[:, :, i*block_Br:(i+1)*block_Br, :]

        for j in range(Tc):
            k_j = K[:, :, j*block_Bc:(j+1)*block_Bc, :]   # [B, H, Bc, d]
            v_j = V[:, :, j*block_Bc:(j+1)*block_Bc, :]   # [B, H, Bc, d]

            # 第一步：在 Cube Unit 上计算 S = q_i @ k_j^T / sqrt(d)
            # S 只存在于 SRAM，不写回 HBM
            S_ij = torch.matmul(q_i, k_j.transpose(-2, -1)) / sqrt(d)   # [B, H, Br, Bc]

            # 第二步：在 Vector Unit 上做在线 softmax
            m_ij = S_ij.max(dim=-1, keepdim=True).values   # [B, H, Br, 1]
            P_ij = torch.exp(S_ij - m_ij)                  # [B, H, Br, Bc]
            l_ij = P_ij.sum(dim=-1, keepdim=True)          # [B, H, Br, 1]

            # 第三步：将当前块的结果与前序块合并
            m_new = torch.maximum(m_prev, m_ij)
            rescale_prev = torch.exp(m_prev - m_new)
            rescale_cur = torch.exp(m_ij - m_new)

            o_i = o_i * rescale_prev + (P_ij / l_ij * rescale_cur) @ v_j
            l_prev = l_prev * rescale_prev + l_ij * rescale_cur
            m_prev = m_new

        # 当前 Q tile 处理完所有 KV tile 后，归一化并写回 HBM
        O[:, :, i*block_Br:(i+1)*block_Br, :] = o_i / l_prev
        row_max[:, :, i*block_Br:(i+1)*block_Br, :] = m_prev
        row_sum[:, :, i*block_Br:(i+1)*block_Br, :] = l_prev

    return O, row_max, row_sum

外层循环 i 遍历 Q 的 tile，内层循环 j 遍历 K/V 的 tile。每个内层迭代中，最关键的约束是 S_ij（大小 Br x Bc）必须能在 SRAM 中放下。Br 和 Bc 的乘积不得超过 SRAM 减去 Q、K、V tile 占用后的剩余容量。在昇腾 NPU 上，Br 和 Bc 通常在 64 到 256 之间，head_dim 则固定在 64 到 256 的范围内。S_ij 在 SRAM 上被创建、用于 softmax 计算、据此立即丢弃——一字节都不写回 HBM。在线 softmax 的 rescale 计算是融合这段循环的关键：每次内层循环出口时，o_i 和 l_prev 中就累积了从 j=0 到当前 j 的所有 KV tile 的正确加权结果和归一化因子。当 j 循环完毕，o_i / l_prev 就等价于把完整 N x N 注意力矩阵算一次的结果。

注意每轮内层循环只会向 SRAM 中加载 Q_i 和 K_j、V_j，而 P_ij 和 rescale 中间变量都完全在 SRAM 生命周期内被消费。整段内层循环对外只产生一次 HBM 写入：循环结束后把 O_i 的归一化结果写回。

在线 Softmax 的数学原理拆解

标准 softmax 定义为：

softmax(x)_i = exp(x_i - max(x)) / sum(exp(x_j - max(x)))

分母需要对整行所有元素求和，分子需要对每个元素做指数运算。在标准实现中，这要求你完整的 x 向量在内存中连续可用——这正是 FlashAttention 想避免的情形。

在线 softmax 的核心想法是：把分母的求和拆成可累加的部分。定义两个状态量——m 为已处理部分的最大值，l 为已处理部分的指数和。当处理第 t 个块时：

def online_softmax_step(m_prev, l_prev, x_t):
    """
    在线 softmax 的单步更新

    参数:
        m_prev: 前 t-1 个块的最大值 (标量)
        l_prev: 前 t-1 个块归一化后的指数和 (标量)
        x_t:    当前块的值 (向量)

    返回:
        m_cur, l_cur: 更新后的全局统计量
    """
    # 当前块的最大值和指数和
    m_t = max(x_t)
    s_t = sum(exp(x_t - m_t))

    # 与前序统计量合并
    m_cur = max(m_prev, m_t)

    # 关键步骤：rescaling
    # 如果 m_cur == m_prev，前序块的指数值不变，当前块整体缩小 exp(m_t - m_cur)
    # 如果 m_cur == m_t，前序块的指数值整体缩小 exp(m_prev - m_cur)
    # 两种情况合并为统一的变尺度操作
    l_cur = l_prev * exp(m_prev - m_cur) + s_t * exp(m_t - m_cur)

    return m_cur, l_cur

rescaling 操作不是近似，是精确的数学等效变换。e^(x_i - m_prev) 乘以 e^(m_prev - m_cur) 等于 e^(x_i - m_cur)，即把以 m_prev 为基准的指数值整体变尺度到以 m_cur 为基准。这个变换是可逆的、无损的（忽略浮点精度）。因此无论分成多少块合并，最终的 softmax 结果都和一次性全局计算完全一致。这解释了为什么 FlashAttention 在数学上正确——它不是在逼近全局 softmax，而是在分步计算完全相同的全局 softmax，只不过利用 rescaling 实现了流水线式合并。

在线 softmax 的信息论含义是：你不需要全量存储位置之间的配对分数来获得正确的概率分布。你只需要每行的最大值和一个累加和。这正是"从 O(N平方) 到 O(N)"的信息论根源——要准确描述一个概率分布，不一定要存储所有元素，只需要足够多的统计量。Softmax 恰好有 m 和 l 两个统计量就能无损重构其归一化结果。

反向传播的重算策略及其合理性

训练场景下，FlashAttention 前向计算完就丢弃了 S 矩阵。反向传播需要 dQ、dK、dV 三个梯度张量，它们都依赖于 S 矩阵的各个元素。标准实现直接从 HBM 缓存中读取 S，FlashAttention 必须重新计算。

反向传播的重算不是在全部序列维度上重算，而是和分块 tiling 相同的方式逐块重算。每次只重算一个块大小 Br x Bc 的注意力分数，利用前向保留的 row_max 和 row_sum 恢复出正确的注意力权重 P_ij，据此计算梯度并累加到 dQ、dK、dV 的对应 tile 中。

def flash_attention_backward_tiled(Q, K, V, O, dO, row_max, row_sum,
                                    block_Br=128, block_Bc=128):
    """
    FlashAttention 反向传播——重算策略

    前向阶段只保留了 row_max 和 row_sum（各 O(N) 大小），
    因此 S 矩阵和 P 矩阵必须在这段函数中重新计算。

    dO 是损失对输出 O 的梯度，从上游层传入。
    """
    B, H, N, d = Q.shape
    Tr = ceil_div(N, block_Br)
    Tc = ceil_div(N, block_Bc)

    dQ = torch.zeros_like(Q)
    dK = torch.zeros_like(K)
    dV = torch.zeros_like(V)

    for i in range(Tr):
        q_i = Q[:, :, i*block_Br:(i+1)*block_Br, :]      # [B, H, Br, d]
        do_i = dO[:, :, i*block_Br:(i+1)*block_Br, :]    # [B, H, Br, d]

        # 读取前向保存的统计量
        m_i = row_max[:, :, i*block_Br:(i+1)*block_Br, :]
        l_i = row_sum[:, :, i*block_Br:(i+1)*block_Br, :]

        for j in range(Tc):
            k_j = K[:, :, j*block_Bc:(j+1)*block_Bc, :]
            v_j = V[:, :, j*block_Bc:(j+1)*block_Bc, :]

            # 重算 S_ij（前向阶段已经算过一次）
            S_ij = torch.matmul(q_i, k_j.transpose(-2, -1)) / sqrt(d)

            # 从统计量 m_i 和 l_i 还原 P_ij
            P_ij = torch.exp(S_ij - m_i) / l_i

            # 梯度计算链
            # dV += P_ij^T @ dO_i
            dV[:, :, j*block_Bc:(j+1)*block_Bc, :] += \
                torch.matmul(P_ij.transpose(-2, -1), do_i)

            # dP = dO_i @ v_j^T
            dP = torch.matmul(do_i, v_j.transpose(-2, -1))

            # softmax 反向传播公式：dS = P * (dP - sum(P * dP))
            dS = P_ij * (dP - torch.sum(P_ij * dP, dim=-1, keepdim=True))

            # dQ_i += dS @ k_j, dK_j += dS^T @ q_i
            dQ[:, :, i*block_Br:(i+1)*block_Br, :] += \
                torch.matmul(dS, k_j)
            dK[:, :, j*block_Bc:(j+1)*block_Bc, :] += \
                torch.matmul(dS.transpose(-2, -1), q_i)

    return dQ, dK, dV

反向传播相比前向多了一倍的矩阵乘法——每个内层迭代中，不仅要重算 S_ij 和 P_ij，还要计算三个梯度分量（dV、dQ、dK）的各自矩阵乘法。这意味着反向传播的总 FLOPs 是前向的约 2.5 到 3 倍。但是重算节省下的 HBM 读写量是 N平方 级别的，而额外增加的运算量是 N x d 级别的（d 远小于 N）。在长序列下，N平方 >> N x d，因此重算的算术开销相对于 HBM 带宽节省可以忽略。这个计算与通信的交换比决定了 FlashAttention 在长序列下占优而在短序列下可能劣化的行为模式。在 ops-transformer 仓库中，FlashAttentionScore 有内部决策逻辑来判断是否启用 FlashAttention 路径还是回退到标准注意力路径，判断依据就是序列长度和芯片规格的比值。

分块参数的物理约束

Block size（分块大小）不是一个自由参数，它由昇腾 NPU 的 SRAM 容量直接决定。设 SRAM 可用容量为 M_flops（字节），则需要满足：

Br * d + Bc * d + Br * Bc < M_sram

其中 Br * d 是 Q tile，Bc * d 是 K 或 V tile（加载一个即可复用），Br * Bc 是注意力分数 tile。对于 head_dim=128，Br=Bc=128 约需 32KB 空间。对于 head_dim=256，Br=Bc=64 约需 20KB。

过大的 block size 会导致 SRAM 溢出，导致算子退化为使用 HBM 的 fallback 路径。过小的 block size 则增加 tiling 循环次数和重算开销。ops-transformer 的 FlashAttentionScore 根据输入序列长度和设备能力自动选择合适的块大小：

// FlashAttentionScore 块大小选择逻辑（抽象版本）
struct TileConfig {
    int Br;  // Q tile size
    int Bc;  // K/V tile size
};

TileConfig select_tile_size(
    int seq_len,
    int head_dim,
    int sram_bytes,
    int num_heads
) {
    // 预估 SRAM 中同时驻留的数据量
    // q_tile + k_tile + v_tile + s_tile + p_tile + temp_buffer
    int head_bytes = head_dim * sizeof(bfloat16);
    int sram_per_head = sram_bytes / num_heads;

    // 候选块大小组合
    vector<pair<int,int>> candidates = {
        {256, 256}, {256, 128}, {128, 128},
        {128, 64},  {64, 64},   {64, 32}
    };

    // 从大到小选择第一个能放下的组合
    for (auto [Br, Bc] : candidates) {
        int needed = Br * head_bytes        // Q tile
                   + Bc * head_bytes        // K tile
                   + Bc * head_bytes        // V tile
                   + Br * Bc * 2;           // S + P tile (shared)

        if (needed <= sram_per_head) {
            return {Br, Bc};
        }
    }

    // 最保守的 fallback
    return {32, 32};
}

选择逻辑的优先级是"越大越好+能放下"。大 tile 减少循环次数和重算频次，但受限于物理 SRAM。注意在代码中 S 和 P 共享同一片 SRAM 区域（S 计算并 softmax 后可以覆写为 P），因此计为 Br * Bc * 2 而非 Br * Bc * 4。在昇腾 NPU 上，Cube Unit 计算完 S 后，Vector Unit 可以直接在原地做 softmax 得到 P——这不仅节省 SRAM，还省了数据搬移。性能最优的 block size 需要通过 profiling 确定，但上述选择逻辑保证了不会因 SRAM 溢出而触发 HBM fallback。

MC2：通信融合的下一个层级

FlashAttention 解决了注意力计算内部的 HBM 读写问题。但在张量并行（Tensor Parallelism）场景下，注意力层后面还跟着一个跨卡 AlltoAll 通信——把注意力输出的分布从"每卡算完整头"切换到"每卡持有序列片段"。这个通信操作需要把输出张量从 HBM 读出来、拆包、发送、在接收端重组、再写回 HBM。如果 FlashAttention 把输出压到了 1x 写，而 AlltoAll 又把它读出来再写回去，融合就做了半套。

MC2（Matmul Communication Fusion）系列算子在 ops-transformer 中把矩阵乘法和 AlltoAll 融合成一个算子。AttentionToFFN 把注意力输出的规整和 FFN 输入的通信融合，FFNToAttention 反过来把 FFN 输出和注意力输入之间的通信融合。MatmulAlltoAll 则是一个更通用的融合模式。

下面是 MatmulAlltoAll 的流水线调度抽象。真实实现中使用的是 Ascend C 的 Pipe 机制做双缓冲和计算通信重叠：

// MC2: MatMul AlltoAll 流水线融合调度
// 将全连接层计算与跨卡通信流水线化
void MatMulAlltoAll_pipelined(
    const Tensor& input,        // [B, S, H]
    const Tensor& weight,       // [H, H']
    const CommGroup& comm,      // 通信域配置
    Tensor& output,
    int tile_size = 128
) {
    int num_tiles = input.shape[1] / tile_size;
    int hidden = input.shape[2];

    // 分配两个双缓冲区：一个被 Cube 写入，一个被 DMA/通信读取
    Tensor buf_a = alloc_local(tile_size, hidden, SRAM);
    Tensor buf_b = alloc_local(tile_size, output.shape[2], SRAM);

    Tensor* write_buf = &buf_a;
    Tensor* send_buf = &buf_b;

    for (int i = 0; i < num_tiles; i++) {
        // 阶段1: Cube Unit 计算 input[i] @ weight
        // 结果写入 write_buf
        matmul_async(input.slice(i * tile_size, tile_size, 1),
                     weight,
                     *write_buf);

        // 阶段2: 把前一个已经计算完成的 send_buf 发出 alltoall
        // 通信操作和当前 tile 的 matmul 在时间上重叠
        if (i > 0) {
            alltoall_async(*send_buf,
                           comm,
                           output.slice((i-1) * tile_size, tile_size, 1));
        }

        // 管道切换：当前计算缓冲区和发送缓冲区互换角色
        std::swap(write_buf, send_buf);
    }

    // 处理末尾一个 tile
    alltoall_sync(*send_buf, comm,
                  output.slice((num_tiles-1) * tile_size, tile_size, 1));
}

标准实现中，matmul 先把整块结果写完 HBM，据此 alltoall 再从 HBM 读数发送。中间张量的大小是 [B, S, H’]，和 FlashAttention 的输出大小在同一量级。MC2 把这个中间张量分成了 tile，每个 tile 的 matmul 结果在 SRAM 上即被通信单元读取并发出，中间张量一字节都不写回 HBM。双缓冲技术的要点是 Cube Unit 向 write_buf 写入的同时，DMA 正从 send_buf 中读取并发出通信——计算和通信完全流水线重叠，没有闲置等待。MC2 和 FlashAttention 在没有直接依赖关系的情况下采用了相同的设计哲学：在 SRAM 上完成尽可能多的工作，把 HBM 的写回量压到最低。

GatherPAKVCache 的显存管理逻辑

长序列推理中 KV Cache 是另一个显存大户。对于 MHA（多头注意力），KV Cache 大小为 2 x N x num_kv_heads x head_dim x batch x layers。在 128K 序列、32 层、32 KV 头的场景下，KV Cache 轻松超过 30GB。GatherPAKVCache 是 ops-transformer 中的 KV Cache 管理算子，它的功能是在推理的每个 step 中从分块存储的 KV Cache 中取出当前注意力计算需要的部分。

def gather_paged_kv_cache(
    key_cache_blocks,      # 分块存储的 K cache
    value_cache_blocks,    # 分块存储的 V cache
    block_tables,          # 逻辑块到物理块的映射表
    seq_len,               # 当前序列长度
    head_indices           # 需要取出的 head 索引（GQA 模式）
):
    """
    Paged KV Cache 的 gather 操作

    将逻辑上连续的 KV Cache 从物理上分片存储的块中
    重新排列为注意力计算所需要的连续布局。
    """
    output_K = []
    output_V = []

    for block_idx in range(ceil_div(seq_len, block_size)):
        # 逻辑块号 -> 物理块号
        phys_block = block_tables[block_idx]

        # 从物理块中读取数据
        k_block = key_cache_blocks[phys_block]
        v_block = value_cache_blocks[phys_block]

        # 如果是 GQA 模式，选择的 head 不是全部的 head
        if head_indices is not None:
            k_block = k_block[:, :, head_indices, :]
            v_block = v_block[:, :, head_indices, :]

        output_K.append(k_block)
        output_V.append(v_block)

    return torch.cat(output_K, dim=2), torch.cat(output_V, dim=2)

分块存储（也称作 paged KV Cache）解决了大序列推理时 KV Cache 的内存碎片问题。连续分配一整块 N x num_heads x head_dim 的内存在序列长度动态增长时容易引发碎片和重新分配。分块存储允许物理块不连续，通过 block_tables 做地址翻译。在推理中，每次 decode 一个新 token，只需要新增一个 KV block，而不是重新分配整个 Cache。GQA/MQA 模式下 KV head 数少于 query head 数，head_indices 参数确保只取出需要的 head 进行注意力计算，避免无效数据传输。

多算子协同下的完整推理路径

FlashAttentionScore、GatherPAKVCache、MC2 这三个算子不是独立工作的。在一个实际的大模型推理过程中，它们被 ATB 加速库编排成一个完整的计算图：

1. GatherPAKVCache 从分块的 KV Cache 中取出当前推理步所需的 K、V 块
2. FlashAttentionScore 在 SRAM 上完成分块注意力计算，输出结果 O
3. MC2 AttentionToFFN 融合算子接收 O，在线做矩阵乘法和 AlltoAll 通信分片
4. 每个分片的 FFN 计算在各自设备上独立完成
5. MC2 FFNToAttention 将 FFN 输出做 AlltoAll 重新分片，进入下一层的注意力计算

这条路径的每一次数据移动都处在"能不搬就不搬、能少搬就少搬、要搬就搬最小块"的原则下。在推理服务的实际部署中，关键是降低 TTFT（Time To First Token）和 TPOT（Time Per Output Token）。FlashAttention 在 prefilling 阶段（TTFT 计算）中消除显存瓶颈，允许更大的 batch size；在 decode 阶段优化单步计算延迟。

使用前 vs 使用后差异

下表从结构化的四个维度对比标准注意力实现和以 FlashAttentionScore 为代表的优化实现：

维度	通用实现	优化实现	差异来源
显存占用	O(N平方)，保存完整的注意力分数矩阵 S 和 softmax 输出 P，反向传播时还多保存一份 S	O(N)，仅保存每行的最大值 m 和归一化因子 l 两个统计量	在线 softmax 以两个标量而非 N 个向量来编码完整的概率分布信息；舍弃 S 和 P 的存储需求，反向传播时用重算替代存储
HBM 读写次数	前向：1次写S + 1次读S + 1次写P + 1次读P + 1次写O = 5次全量HBM访问	前向：分块读取 QKV + 1次写O = 约2次全量等价访问	SRAM 上完成 S→P→O 的全链计算，中间状态不出 SRAM 即不产生 HBM 流量
可扩展性	序列长度 N 翻倍，显存增 4 倍，批大小减半	序列长度 N 翻倍，显存近似翻倍	O(N平方) 与 O(N) 的复杂度差决定了在长序列下的本质差异
数值精度	标准全局 softmax，精度由硬件 FPU 保证	在线 softmax 经多次 rescaling，精度略微下降但可通过 clamp 控制	rescaling 乘法引入微小浮点误差，但序列长度在 128K 内 bfloat16 精度完全满足要求

在各序列长度下的相对表现

下表反映不同序列长度下 FlashAttention 相对标准实现的性能趋势（p 为某模型参数量级，包含推理的矩阵计算和显存访问两部分的综合开销）：

序列长度	标准实现开销	FlashAttention 开销	核心差异
1K	基准 (1.0x)	1.0x-1.2x	短序列时标准实现已够用，FlashAttention 的 tiling 循环和 rescaling 开销抵消了带宽节省
4K	约 4.0x	约 1.5x	注意力分数矩阵 S 的大小开始显著，标准实现 HBM 读写成为瓶颈
8K	约 8.0x	约 1.8x	标准实现中 S 矩阵占用显存超过 500MB（bfloat16），显著影响其他层的缓存
32K	约 16.0x	约 2.5x	标准实现显存接近芯片容量上限，各层之间出现内存竞争和碎片化
128K	不可运行（OOM）	约 3.5x	标准实现所需的 N平方显存超过单卡物理上限，FlashAttention 仍以 O(N) 资源稳定运行

标准实现的开销超出芯片容量后表现为显存溢出错误（Out of Memory），而非性能退化。FlashAttention 的开销增长曲线大体上与序列长度呈线性关系，这与 O(N) 复杂度的预测一致。训练场景下 FlashAttention 的收益更大，因为反向传播的重算策略在前向的"少量写"基础上继续压缩显存——前向和反向共享同一套统计量 m、L，不需要额外存储。

数值安全的边界条件

任何浮点算法都面临上溢（数值过大超出 float 表示范围）和下溢（数值过小被截断为零）两种风险。

标准 softmax 通过先减去最大值规避了上溢风险。在线 softmax 保留了这一特性——每次 rescaling 后，所有指数值仍是负指数（因为 m_global >= m_prev 和 m_cur），因此不会上溢。

下溢风险存在于 rescaling 阶段。当 m_global 大幅超过 m_prev 时，rescale_prev = exp(m_prev - m_global) 会非常小。在 float32 中，exp(-100) ≈ 3.7e-44，远低于 float32 的最小正规数 1.2e-38，进入非正规数（subnormal）范围。极端情况下被截断为零。

bfloat16 的精度更差：最小可表示正数约为 1.2e-38，但最小正规数约为 9.2e-39，指数下限约在 e^(-88) 就触及底边界。序列长度超过 256K 时，多块 rescaling 的累积可能触及下溢。

ops-transformer 中 FlashAttentionScore 的应对措施包括：

1. 当 m_global - m_prev 超过一个安全阈值时，不对前序块做 rescaling 而是直接重置前序块的影响（等价于承认前序块在总和中占比可忽略）
2. 在 rescale_factor 上做 clamp 到 dtype.tiny，防止完全归零
3. 在推理精度测试中加入特定边界条件（全 1 输入、全随机大值输入、交替极值输入）的回归测试

这些安全措施使得 FlashAttentionScore 在实际部署中能达到与标准注意力无差别的推理精度。

总结

SparseFlashAttention 在 FlashAttention 基础上增加了稀疏掩码支持——只计算和存储采样后的注意力模式，适用于 Longformer、BigBird 等稀疏注意力架构。它在 tiling 循环的外层增加了一个 mask 过滤步骤，跳过所有被稀疏掩码标记为不需要的 KV 块。这个算子在保存语义长上下文的同时把显存量进一步压低。

MoEComputeExpertTokens 用于 MoE（混合专家）架构的路由与分发。它根据 gate 网络的输出，将每个 token 派发到其对应的 expert 上，据此在一个 batched matmul 中完成所有 expert 的计算。如果不做这个算子，路由后的 token 需要动态重新排列张量，产生大量零碎的内存操作。

GroupedMatMulSwiGLUQuantV2 把分组矩阵乘、SwiGLU 激活函数和量化解压融合为一个算子。SwiGLU 是 LLaMA 系列模型使用的激活函数，包含两个门控矩阵乘加一个 SiLU 激活。标准实现中，矩阵乘结果先写 HBM，再读回做 SiLU 和点乘，再写回。GroupedMatMulSwiGLUQuantV2 在 SRAM 上完成矩阵乘后直接做 SiLU 和激活门控运算，一次写回。

---

仓库地址：https://atomgit.com/cann/ops-transformer