前言

CANN 作为昇腾NPU 的基础计算框架，其算子生态的完善程度直接决定了昇腾NPU 上大模型工作负载的实际性能表现。Transformer 架构的 Attention 机制是当前大语言模型推理与训练的核心瓶颈。标准 Softmax Attention 的计算复杂度为 O(n²)，在序列长度不断增长的场景下，显存占用和计算耗时呈二次方膨胀。Flash Attention 通过对 Attention 计算过程的数学重排，实现了在不牺牲精度的前提下将显存占用从 O(n²) 降低到 O(n)，同时大幅减少 HBM（High Bandwidth Memory）的读写次数。这套算法最初面向 GPU 设计，而昇腾NPU 的硬件体系结构与 GPU 存在显著差异——例如 Cube 矩阵计算单元与 Vector 通用计算单元的协同方式、L1 Cache 的层级管理、以及统一内存空间中的地址映射策略。ops-transformer 算子库在 CANN 框架中提供了 Flash Attention 的昇腾原生实现，将其核心思路适配到昇腾 AI Core 的算力单元上，使大模型推理任务在昇腾NPU 上获得与 GPU 持平甚至更优的 Attention 计算性能。ops-transformer 并非简单地将 GPU 上的 CUDA 实现移植到昇腾平台，而是基于 CANN 的算子编程模型重新设计了计算流水线，充分利用了昇腾 AI Core 的 Cube-Vector 协同架构特征。

ops-transformer 是 CANN 生态中专注于 Transformer 系列算子的开源算子库，涵盖了 Flash Attention、Paged Attention、Rope 位置编码、SwiGLU 激活函数等关键组件。本文将从工程实践角度拆解 ops-transformer 中 Flash Attention 算子的实现逻辑，探讨其在昇腾 NPU 上的内存分块策略、算子融合方案以及与主流框架的集成方式。

Flash Attention 的核心思想

理解 ops-transformer 的实现，需要先厘清 Flash Attention 的数学基础。标准 Attention 计算可以分解为三个步骤：计算 Q 与 K 的点积得到注意力分数矩阵 S，对 S 按 softmax 归一化得到概率矩阵 P，再将 P 与 V 相乘得到输出 O。朴素实现会显式生成 n×n 规模的 S 和 P 矩阵，这两者都需要写入 HBM 并在后续步骤中重新读出。

Flash Attention 的关键洞察在于：softmax 的计算具有分块可分解性。具体来说，全局 softmax 可以通过在线算法（online softmax）在分块上逐步累积，而不需要一次性获得所有分数。这意味着可以将 Q、K、V 按行或列切分为较小的块（tile），逐块在 SRAM（昇腾上对应 Unified Buffer / L1 Cache）中完成局部计算，仅将最终结果写回 HBM，从而消除了中间 n×n 矩阵的 HBM 读写开销。

在昇腾 NPU 上，这套思路的落地需要面对几个额外的工程约束。昇腾 AI Core 的计算单元分为 Cube 单元（用于矩阵乘法，即 MatMul 操作）和 Vector 单元（用于逐元素运算如 softmax、缩放等）。Flash Attention 的每个分块步骤都包含 MatMul 和逐元素运算的交替执行，因此如何调度 Cube 与 Vector 之间的数据流转、如何在 Unified Buffer 中合理分配工作空间，直接决定了最终性能。

ops-transformer 的算子注册与接口设计

ops-transformer 遵循 CANN 的算子开发规范，通过 OpRegister 机制将算子注册到框架层。Flash Attention 算子在 ops-transformer 中的注册方式如下：

// WHY: 通过 OpRegister 将 FlashAttention 算子注册到 CANN 框架，使其可被
// Ascend Graph 或 MindSpore 等上层框架调用。dtype 参数控制支持的输入精度，
// 常见的组合为 float16 + float32（compute_precision）以保证数值稳定性。
#include "op_host/flash_attention_host.h"

CANN_OP_REG(FlashAttention)
    .Input("query")
    .Input("key")
    .Input("value")
    .Output("attention_out")
    .Attr("scale_factor", "float", "1.0")
    .Attr("head_num", "int", "1")
    .Attr("causal_mask", "bool", "false")
    .Attr("num_kv_heads", "int", "0")
    .Attr("compute_precision", "string", "float32")
    .DataType("query", ge::DT_FLOAT16)
    .DataType("key", ge::DT_FLOAT16)
    .DataType("value", ge::DT_FLOAT16);

上述注册声明了 Flash Attention 算子的输入输出张量和关键属性。其中 num_kv_heads 用于支持 GQA（Grouped Query Attention）和 MQA（Multi-Query Attention），当 num_kv_heads 小于 head_num 时自动启用 GQA 模式。causal_mask 控制是否应用因果注意力掩码，在自回归解码场景中需要置为 true。scale_factor 定义了注意力分数缩放系数，通常为 1/√d_k，但某些场景下可以通过此参数注入自定义缩放逻辑。

ops-transformer 的 Flash Attention 算子同时支持向前计算和反向传播，反向计算同样是分块执行的，利用前向过程中的缓存数据（S、O 的分块）来避免额外的 HBM 访问。训练场景下这种双向支持至关重要，因为标准实现中反向传播的梯度计算同样需要 O(n²) 的显存开销。

分块计算与内存管理策略

ops-transformer 中 Flash Attention 的核心实现围绕分块（tiling）策略展开。分块的核心目标是在有限的 Unified Buffer 空间内，尽可能大地加载 Q、K、V 的子块，减少分块轮次（outer loop 的迭代次数），从而最大化 Cube 单元的利用率。

昇腾 AI Core 的 Unified Buffer 容量有限（典型为数十 KB 到数 MB，取决于芯片型号），而 Q、K、V 张量在长序列场景下可能达到数百 MB。ops-transformer 通过一个 tiling 模块动态计算最优分块大小，该模块综合考虑以下因素：序列长度、头维度、Unified Buffer 剩余可用空间、Cube 单元对输入矩阵行列的约束（昇腾 MatMul 要求矩阵的行列满足特定的对齐和规模要求）。

# WHY: 分块参数的动态计算是 Flash Attention 性能的关键。block_size 的选择
# 需要平衡 Cube 矩阵乘法的效率与 Unified Buffer 的容量限制。
# 过小的 block 会导致 Cube 单元计算不饱和，过大的 block 会超出 Buffer 容量。
# 以下为 tiling 逻辑的简化示意，实际实现包含更多硬件相关的约束检查。
def compute_tiling_params(seq_len, head_dim, ub_size, num_heads):
    # Cube 单元对矩阵维度的对齐要求
    cube_align = 16
    # 预留给中间结果的 Buffer 空间（softmax 中间量、输出缓存）
    reserve_size = head_dim * 2 * 2  # float16, 两份缓存

    available = ub_size - reserve_size
    # Q 按行分块，K 按列分块
    q_block = min(seq_len, available // (head_dim * 2))
    q_block = (q_block // cube_align) * cube_align

    k_block = min(seq_len, available // (head_dim * 2))
    k_block = (k_block // cube_align) * cube_align

    return q_block, k_block

在实际执行中，Q 的每一块被固定在 Unified Buffer 中（因为它要和所有 K 块配对计算），而 K 和 V 按相同的列分块从 HBM 逐块加载。对于每个 Q 块和 K 块的组合，Cube 单元计算局部注意力分数，Vector 单元执行在线 softmax 的累积更新，最终将完成的输出块写回 HBM。这种 Q-outer、K-inner 的分块遍历顺序是 ops-transformer 选择的默认策略，它确保了每个 Q 块只需加载一次，而 K、V 块的重复加载次数取决于 Q 的分块数。遍历顺序的选择并非随意——如果改为 K-outer、Q-inner，则 K 块只需加载一次，但在线 softmax 的累积需要在 Q 维度上展开，这与 Flash Attention 的分块 softmax 算法不太契合，因此 ops-transformer 采用了以 Q 为外层循环的策略。

因果掩码与 ALiBi 的硬件友好实现

在自回归语言模型的推理中，因果掩码（causal mask）确保每个位置只能关注自身及之前的位置。标准实现通过一个 n×n 的布尔矩阵在 softmax 之前将非法位置的分数设为负无穷来实现。但 Flash Attention 的分块执行模式使得这种全局掩码矩阵不再可行——无法一次性获得完整的 n×n 矩阵。

ops-transformer 的解决方案是在分块内部动态计算掩码。对于给定的 Q 块（行范围 [i, i+bq)）和 K 块（列范围 [j, j+bk)），因果掩码的条件是 i + bq - 1 < j，即 Q 块的所有行都严格在 K 块的列之前。当满足这个条件时，整个分块的注意力分数可以直接跳过（设为零）。当 Q 块和 K 块部分重叠时，需要在分块内部对非法位置逐元素施加负无穷掩码。这些操作在 Vector 单元中完成，通过向量化的比较和赋值指令实现，开销很小。

ALiBi（Attention with Linear Biases）是一种替代绝对位置编码的方案，通过在注意力分数上叠加一个与距离成正比的线性偏置项来引入位置信息。ops-transformer 将 ALiBi 偏置的计算融合到 Cube 输出之后、softmax 之前的 Vector 单元操作中。由于 ALiBi 偏置只依赖于 Q 和 K 的位置索引，它可以在分块计算时根据当前块的行列索引动态生成，无需预分配额外的 n×n 偏置矩阵。这种延迟计算的方式进一步节省了 HBM 带宽。

GQA 与 MQA 的支持路径

分组查询注意力（Grouped Query Attention）和多查询注意力（Multi-Query Attention）是当前大模型推理中的主流优化手段，通过让多个 Q 头共享同一组 K、V 头来降低 K、V 矩阵的规模。ops-transformer 的 Flash Attention 通过 num_kv_heads 参数原生支持这一特性。

在 GQA 模式下，K 和 V 的头数少于 Q 的头数，每个 K/V 头被多个 Q 头共享。ops-transformer 的处理方式是在分块遍历时引入一个额外的循环维度：对于每对 Q 头和 K/V 头的映射关系，将共享同一 K/V 头的 Q 头的计算合并到一次 K/V 加载中。这意味着 K 和 V 的每个块只加载一次，但参与多次 Q-K 的点积计算（对应不同的 Q 头），从而在不增加 HBM 访问的情况下完成所有 Q 头的注意力计算。

这种实现方式相比简单的"复制 K/V 头"方案有显著优势——它避免了 K/V 数据的冗余存储和重复加载，在 KV Cache 压力巨大的长上下文推理场景中尤为重要。MQA 是 GQA 的特例（num_kv_heads = 1），由相同的代码路径处理，无需特殊分支。

Paged Attention 的集成

在推理服务的 KV Cache 管理中，ops-transformer 还提供了 Paged Attention 算子的实现。Paged Attention 的核心思想是将 KV Cache 按固定大小的页进行分页管理，类似操作系统的虚拟内存机制，从而解决 KV Cache 碎片化问题。在大规模并发推理服务中，不同请求的序列长度差异巨大，预分配完整长度的 KV Cache 会导致严重的显存浪费。Paged Attention 通过按需分配物理页来解决这一问题——每个请求的逻辑 KV Cache 由一组页组成，页的大小通常对应一个 token 块的 K 和 V 数据，物理页在显存中可以不连续存储。当新的 token 需要追加 KV Cache 时，系统只需分配一个新的物理页并通过页表建立映射即可，无需预留完整的序列长度空间。这种机制使显存利用率从传统方案的平均不足一半提升到接近饱和，在批量推理场景下可以将并发服务能力提升一倍以上。ops-transformer 的 Paged Attention 实现与 Flash Attention 共享底层的基础设施，包括分块策略和 Cube/Vector 协同调度逻辑。差异在于数据读取方式：标准 Flash Attention 从连续内存读取 K、V，而 Paged Attention 需要通过页表将逻辑块索引映射到物理页地址，然后从可能不连续的物理内存中加载数据。ops-transformer 通过在分块循环中插入地址翻译步骤来实现这一点，地址翻译本身在 Vector 单元中以批量方式完成，对 Cube 计算的影响被控制到最小。此外，ops-transformer 还实现了 Copy-on-Write 语义来支持 KV Cache 的共享与分裂，这在 beam search 和 parallel sampling 场景中非常有用——多个候选序列可以共享相同前缀的 KV Cache 页，仅在分叉点之后才分配独立的物理页。

ops-transformer 的 Paged Attention 与 Flash Attention 共享底层的基础设施，包括分块策略和 Cube/Vector 协同调度逻辑。差异在于数据读取方式：标准 Flash Attention 从连续内存读取 K、V，而 Paged Attention 需要通过页表（page table）将逻辑块索引映射到物理页地址，然后从可能不连续的物理内存中加载数据。ops-transformer 通过在分块循环中插入地址翻译步骤来实现这一点，地址翻译本身在 Vector 单元中以批量方式完成，对 Cube 计算的影响被控制到最小。

算子融合与计算图优化

CANN 框架的一个重要优化维度是算子融合——将多个连续算子合并为一个复合算子，减少中间结果的 HBM 读写。ops-transformer 中的 Flash Attention 天然支持多种融合模式。

在推理场景中，Flash Attention 通常与之前的线性投影（Q/K/V 的投影矩阵乘法）和之后的投影（Output 投影）进行融合。ops-transformer 提供了融合版的算子变体，将 QKV 线性投影、Flash Attention 和 Output 投影打包为一个复合算子。这种融合带来的收益是显著的：线性投影的输出不再需要写回 HBM，而是直接留在 Unified Buffer 或 L2 Cache 中供 Flash Attention 消费，省去了三次 HBM 写入和三次 HBM 读取。

// WHY: 算子融合将 QKV 投影、Flash Attention、Output 投影合并为一个
// 复合算子，消除中间结果的 HBM 读写。在推理场景中，融合算子可以将
// 计算图中的三个独立 kernel 调用合并为一个，减少 kernel launch 开销
// 和数据搬运次数，对延迟敏感的在线推理服务意义重大。
// 以下为融合算子注册的结构示意。
CANN_OP_REG(FlashAttentionFusedQKV)
    .Input("hidden_states")
    .Input("q_proj_weight")
    .Input("k_proj_weight")
    .Input("v_proj_weight")
    .Input("o_proj_weight")
    .Input("kv_cache")       // 可选，用于增量解码
    .Output("output")
    .Attr("head_num", "int")
    .Attr("head_dim", "int")
    .Attr("num_kv_heads", "int")
    .Attr("causal_mask", "bool")
    .DataType("hidden_states", ge::DT_FLOAT16)
    .DataType("q_proj_weight", ge::DT_FLOAT16);

融合算子的实现复杂度远高于单独算子，因为它需要在内部管理更长的计算流水线，并处理不同计算阶段之间的数据依赖关系。ops-transformer 通过精心设计的 Buffer 管理策略来解决这个问题——将 Unified Buffer 划分为多个区域，分别用于输入缓存、中间计算和输出累积，通过双缓冲（double buffering）技术实现计算与数据加载的重叠。

效率对比

以下是 ops-transformer 中 Flash Attention 算子与标准 Attention 实现的关键效率指标对比（基于典型的大模型推理场景，具体数值因模型配置和硬件型号而异）：

指标	使用前（标准 Attention）	使用后（Flash Attention）
注意力分数矩阵显存占用	O(n²)，需在 HBM 中完整存储	O(n)，仅保留分块中间结果
HBM 读写次数（前向）	约 4 次完整的 QKV 矩阵读写	显著降低，中间结果在片上缓存完成
长序列（4K 以上）推理延迟	受限于 HBM 带宽瓶颈，延迟较高	延迟明显降低，接近理论计算下界
KV Cache 与注意力计算的数据搬运	KV Cache 需多次完整回读	按需分块加载，减少无效搬运
GQA/MQA 场景的 KV Cache 占用	相同（取决于实现是否支持）	通过共享头机制有效降低 KV Cache 规模
因果掩码处理	需要 n×n 掩码矩阵，额外显存开销	分块内动态计算，无额外显存占用

需要注意，上述对比中的具体加速比受多种因素影响：序列长度（越长 Flash Attention 优势越明显）、头维度（影响分块效率和 Cube 利用率）、批次大小（影响 HBM 带宽的竞争程度）以及昇腾芯片的具体型号和驱动版本。在实际部署中，建议针对目标模型和硬件配置进行基准测试，以获得精确的性能数据。

与主流推理框架的集成

ops-transformer 的 Flash Attention 算子已经集成了多个主流推理框架的适配层。在 Ascend 的推理加速库中，通过 CANN 的 OpAPI 接口调用 ops-transformer 提供的算子；在基于 MindSpore 的模型部署中，ops-transformer 的算子被注册为 MindSpore 的自定义算子，可直接在计算图中使用。这种多框架的适配策略使 ops-transformer 能够覆盖从训练到推理的全流程场景，开发者无需针对不同框架分别实现 Attention 算子。

对于增量解码（incremental decoding）场景，ops-transformer 支持将新的 token 对应的 K、V 向量追加到 KV Cache 中，然后仅对新增的查询位置执行 Flash Attention 计算。这种增量模式下的分块策略与全量计算有所不同——由于 K、V 的序列长度在不断增长，tiling 模块需要动态调整 K/V 的分块大小，确保在 KV Cache 逐渐增大时仍然保持高效的片上缓存利用率。在增量解码中，Q 的维度通常为 1（单 token 推理），此时 Flash Attention 的主要工作是对整个 KV Cache 执行一次注意力计算。ops-transformer 在这种退化场景下的优化重点是减少 KV Cache 的 HBM 读取开销，通过预取和缓存复用策略来缓解长序列解码时的延迟问题。

在 vLLM 等第三方推理框架的昇腾适配分支中，ops-transformer 的 Paged Attention 算子被用于管理虚拟内存式的 KV Cache，使 vLLM 的 PagedAttention 机制能够在昇腾NPU 上直接运行，无需修改上层的调度逻辑。这种即插即用的适配能力是 ops-transformer 架构设计的重要目标——通过标准化算子接口，让上层框架的注意力调度算法（如 vLLM 的 block manager、TensorRT-LLM 的 KV cache manager）能够无缝对接昇腾硬件。开发者在使用这些框架时，只需确保 ops-transformer 的算子库已正确安装并注册到 CANN 运行时，即可获得与 GPU 版本对等的 Paged Attention 性能表现。

RoPE 位置编码与 SwiGLU 激活

ops-transformer 除了 Attention 系列算子外，还提供了 RoPE（Rotary Position Embedding）和 SwiGLU 激活函数的昇腾实现。这两个算子通常与 Flash Attention 配合使用，构成完整的 Transformer 前向传播路径。在 ops-transformer 的算子编排中，RoPE、Flash Attention 和 SwiGLU 这三类算子的执行顺序和内存流转方式经过了精心设计，以最小化整体计算图的 HBM 访问次数。例如，Q 和 K 的 RoPE 编码可以与 QKV 线性投影的输出融合，在数据尚未写回 HBM 之前就完成旋转编码，然后再进入 Flash Attention 的分块计算流程。这种跨算子的融合编排是 ops-transformer 相对于简单逐算子调用方案的核心优势之一。

RoPE 通过对 Q 和 K 施加旋转变换来编码位置信息，其计算可以分解为若干次逐元素运算和简单的矩阵变换。ops-transformer 将 RoPE 的计算融合到 QKV 投影之后、Flash Attention 之前，在 Vector 单元中以流水线方式完成。由于 RoPE 的计算仅依赖于位置索引和头维度，它可以完全在片上完成，不引入额外的 HBM 访问。

SwiGLU 是当前大模型普遍采用的激活函数，其公式为 SwiGLU(x) = x * SiLU(gate_proj(x))，其中 SiLU 为带有 sigmoid 门控的线性激活。ops-transformer 的 SwiGLU 实现将两次矩阵乘法（gate_proj 和 up_proj）和逐元素运算融合为一个复合算子，利用 Cube 单元的高吞吐来加速矩阵乘法部分，Vector 单元处理门控计算和逐元素乘法。在昇腾NPU 上，SwiGLU 的融合实现将原本需要三次独立算子调用的工作合并为一次 kernel launch，显著减少了中间结果的 HBM 写入。gate_proj 的输出不需要写回 HBM 再读取，而是直接在 Unified Buffer 中完成 SiLU 门控和逐元素乘法后，再由 Cube 单元执行最终的 down_proj 投影。这种深度融合的实现方式要求 ops-transformer 对 Unified Buffer 的空间进行精细的分区管理，确保各阶段的中间结果在片上内存中不发生覆盖冲突。

实践中的调优建议

在实际部署中，开发者可以通过以下几个方面进一步优化 ops-transformer Flash Attention 的性能表现。

首先，合理设置 compute_precision 属性。在 float16 输入的场景下，将 compute_precision 设为 float32 可以在 softmax 计算中获得更好的数值稳定性，代价是部分中间计算需要在 Vector 单元中以更高精度执行。对于大多数推理场景，float16 计算已经足够，但在训练或涉及极大注意力分数的场景下，float32 计算精度可以避免梯度溢出问题。

其次，关注批次大小与序列长度的组合对分块效率的影响。当批次较小而序列较长时，单个请求的 HBM 带宽竞争较少，Flash Attention 的带宽节约优势更加明显。当批次较大时，多个请求的 KV Cache 可能竞争 L2 Cache 空间，此时需要评估是否通过控制并发请求数来维持每个请求的分块效率。

另外，对于增量解码场景，需要注意 KV Cache 的内存布局。ops-transformer 支持连续内存布局和分页内存布局两种模式。连续布局在短序列场景下访问效率更高（更好的空间局部性），而分页布局在长序列和流式解码场景下能有效管理显存碎片。选择哪种布局取决于具体的应用场景和显存资源约束。

在多卡分布式推理中，ops-transformer 的 Flash Attention 与 CANN 的分布式通信算子配合工作。当使用张量并行时，Q 的投影分布在多张卡上，K 和 V 也相应切分。ops-transformer 需要在分块计算前通过 AllGather 操作收集完整的 K、V 块，或者在分块循环内嵌入通信步骤。后者（通信与计算重叠）是更优的实现方式，但复杂度更高，需要 CANN 框架层面的支持。

总结

ops-transformer 在 CANN 生态中承担了 Transformer 核心算子的昇腾原生实现任务。Flash Attention 作为其中最复杂的算子之一，其实现涉及分块策略的动态计算、Cube 与 Vector 单元的协同调度、因果掩码和 ALiBi 的分块内处理、GQA/MQA 的共享头优化，以及与线性投影的算子融合等多个工程维度。理解这些实现细节，有助于开发者在昇腾NPU 上针对大模型推理和训练任务进行更精细的性能调优。从整体架构来看，ops-transformer 的设计哲学是将 Transformer 中所有与注意力相关的计算环节作为一个整体来优化，而非孤立地改进单个算子。

---

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