1. 引言：Attention 为什么需要"Flash"

Transformer 的核心是自注意力（Self-Attention）。给定查询 $Q$、键 $K$、值 $V$，标准注意力计算为：

$$S=\frac{Q{K}^{\top }}{\sqrt{d}},P=\text{softmax}(S),O=PV$$

其中 $Q\in {\mathbb{R}}^{N\times d}$，$K,V\in {\mathbb{R}}^{N\times d}$，$N$ 为序列长度，$d$ 为头维度。

看起来只是三次矩阵乘法，但问题在于中间矩阵 $S$ 和 $P$ 的显存开销。$S$ 的大小是 $N\times N$，当 $N=8192$ 时，单个头就需要 $8192\times 8192\times 2\text{B}\approx 128\text{MB}$（FP16），多头、多层叠加后显存迅速爆炸。更关键的是，即便显存够用，HBM 与 SRAM 之间的数据搬运也会成为瓶颈——GPU 的算力增长远快于显存带宽增长，Attention 这类访存密集型算子很容易被"喂不饱"。

Flash Attention 的核心洞察正是：Attention 的瓶颈不在算力，而在显存访问。它通过分块（Tiling）与在线 Softmax（Online Softmax）把计算全部放进片上 SRAM 完成，避免 $N\times N$ 中间矩阵落盘，从而同时获得"省显存"和"更快"两个收益。

本文先讲通用 Flash Attention 的优化思想并给出 CUDA 示例，再重点剖析昇腾（Ascend）NPU 上的实现思路，最后对比两者的共性与差异。

2. 通用 Flash Attention 的核心优化思想

2.1 标准 Attention 的访存问题

先看朴素实现（PyTorch 风格伪代码）：

def naive_attention(Q, K, V):
    S = Q @ K.transpose(-2, -1) / sqrt(d)   # 写 HBM：N×N
    P = softmax(S, dim=-1)                  # 读 HBM + 写 HBM：N×N
    O = P @ V                               # 读 HBM：N×N
    return O

一次前向至少产生 4 次 $N\times N$ 级别的 HBM 读写。以 $N=4096$、$d=64$、FP16 为例，$S$ 约 32MB，读写带宽压力巨大。这就是 Flash Attention 要消灭的对象。

2.2 分块计算（Tiling）

把 $Q$ 按行切成 $B_r$ 大小的块，$K,V$ 按行切成 $B_c$ 大小的块。外层遍历 $Q$ 块，内层遍历 $K/V$ 块，每次只把一小块数据搬进 SRAM：

• $Q_i$：$B_r\times d$
• $K_j,V_j$：$B_c\times d$
• 累加器 $O_i$：$B_r\times d$

这样片上内存占用是 $O(B_r\cdot d+B_c\cdot d)$，与 $N$ 无关。

2.3 在线 Softmax（Online Softmax）

分块后无法一次性拿到整行 $S$ 的最大值，softmax 不能直接算。Flash Attention 用在线更新的方式维护每行的 running max $m$ 和 running sum $\ell$：

对第 $j$ 个 $K/V$ 块：

$$m_i^{\text{new}}=\max (m_i,\underset{j}{\max }{S}_{ij})$$

$${\ell }_i^{\text{new}}=e^{m_i-m_i^{\text{new}}}\cdot {\ell }_i+\sum _j{e}^{S_{ij}-m_i^{\text{new}}}$$

$$O_i^{\text{new}}=\frac{e^{m_i-m_i^{\text{new}}}\cdot {\ell }_i\cdot O_i+e^{S_{ij}-m_i^{\text{new}}}{V}_j}{{\ell }_i^{\text{new}}}$$

关键点：当 max 更新时，用缩放因子 $e^{m_i-m_i^{\text{new}}}$ 修正历史累加结果，保证结果与一次性 softmax 数学等价。这就是 Flash Attention 能"边算边归一化"的根本。

2.4 重计算（Recomputation）与反向传播

反向传播需要 $P$ 矩阵，但前向不保存 $P$。Flash Attention 在反向时用保存的 $O$、$m$、$\ell$ 重新计算 $S$ 和 $P$，用少量额外算力换取大量显存节省。由于 Attention 是访存瓶颈而非算力瓶颈，这笔交易非常划算。

2.5 一个简化的 CUDA 示例

下面给出一个教学性质的简化 CUDA 实现，聚焦"分块 + 在线 softmax"的骨架（省略了 warp 级 shuffle、异步拷贝、边界处理等工程细节，真实实现见官方 flash_attn 仓库）：

// 简化版 Flash Attention 前向：每个 block 负责 Q 的一行块
// 假设 d <= HEAD_DIM_MAX，Bc 为 K/V 分块大小
template <int Br, int Bc, int D>
__global__ void flash_attn_fwd(const half* Q, const half* K, const half* V,
                               half* O, int N, float scale) {
    // 共享内存：Q 块、K 块、V 块
    __shared__ half sQ[Br][D];
    __shared__ half sK[Bc][D];
    __shared__ half sV[Bc][D];

    int q_start = blockIdx.x * Br;          // 当前 Q 块起始行
    int tid = threadIdx.x;

    // 1. 载入 Q 块到 SRAM
    for (int i = tid; i < Br * D; i += blockDim.x)
        sQ[i / D][i % D] = Q[(q_start + i / D) * D + i % D];

    // 2. 初始化在线 softmax 状态（每个 Q 行一份）
    float m[Br], l[Br], acc[Br][D];
    for (int r = 0; r < Br; ++r) {
        m[r] = -INFINITY;
        l[r] = 0.f;
        for (int c = 0; c < D; ++c) acc[r][c] = 0.f;
    }

    // 3. 遍历 K/V 块
    for (int kv_start = 0; kv_start < N; kv_start += Bc) {
        // 载入 K/V 块
        for (int i = tid; i < Bc * D; i += blockDim.x) {
            sK[i / D][i % D] = K[(kv_start + i / D) * D + i % D];
            sV[i / D][i % D] = V[(kv_start + i / D) * D + i % D];
        }
        __syncthreads();

        // 计算 S = Q_i * K_j^T * scale，并做在线 softmax 更新
        for (int r = 0; r < Br; ++r) {
            float row_max = m[r];
            float s[Bc];
            for (int c = 0; c < Bc; ++c) {
                float dot = 0.f;
                for (int k = 0; k < D; ++k)
                    dot += __half2float(sQ[r][k]) * __half2float(sK[c][k]);
                s[c] = dot * scale;
                row_max = fmaxf(row_max, s[c]);
            }
            // 缩放因子修正历史累加
            float corr = expf(m[r] - row_max);
            float row_sum = l[r] * corr;
            for (int c = 0; c < Bc; ++c) {
                float p = expf(s[c] - row_max);
                row_sum += p;
                for (int k = 0; k < D; ++k)
                    acc[r][k] = acc[r][k] * corr + p * __half2float(sV[c][k]);
            }
            m[r] = row_max;
            l[r] = row_sum;
        }
        __syncthreads();
    }

    // 4. 归一化并写回 O
    for (int r = 0; r < Br; ++r)
        for (int k = 0; k < D; ++k)
            O[(q_start + r) * D + k] = __float2half(acc[r][k] / l[r]);
}

这段代码体现了 Flash Attention 的三个灵魂：

1. sQ/sK/sV 放在 __shared__：数据只在 SRAM 中流转，$N\times N$ 的 $S$ 从未落盘；
2. m[r]、l[r]、acc[r][k] 是 running 状态：对应在线 softmax 的 max、sum、输出累加器；
3. corr = expf(m[r] - row_max)：max 变化时对历史结果做指数修正，保证数学等价。

真实的高性能实现还会叠加：warp 级 MMA（Tensor Core）、cp.async 异步流水、寄存器分块、exp2 替代 exp、以及针对 causal mask 的块级跳过等。

3. 昇腾 NPU 上的 Flash Attention 实现思路

昇腾 AI 处理器（如 Ascend 910B）的硬件架构与 GPU 差异很大，直接照搬 CUDA 写法行不通。理解昇腾实现，先要理解它的"三件套"。

3.1 昇腾的硬件抽象：Cube、Vector、Scalar

昇腾 AI Core 内部有三类计算单元：

单元	职责	类比
Cube	矩阵乘（MatMul），基于 3D Cube 结构	类似 Tensor Core
Vector	向量运算（exp、max、add、mul 等）	类似 CUDA Core 的 SIMT 部分
Scalar	标量控制、地址计算、循环	类似控制流单元

此外还有 Unified Buffer（UB） 作为片上高速缓存，以及 L1/L0 缓冲。数据从 HBM 搬到 UB 需要显式指令（DataCopy），这与 GPU 的隐式 cache 模型不同——昇腾要求程序员显式管理数据搬运。

3.2 编程模型：Ascend C 与 Tiling

昇腾主流的高性能算子开发语言是 Ascend C（基于 C++ 的扩展），核心概念：

• AI Core 并行：多个 AI Core 同时执行，通过 GetBlockIdx() 区分当前核；
• Tiling 策略：Host 侧根据 shape 计算切分参数，Device 侧按 tiling 执行；
• 多核切分：把总计算量显式划分并绑定到有限的 AI Core 上，每个核循环消化多个基本块；
• 流水并行：通过 TPipe/TQue 管理多级缓冲，实现 CopyIn → Compute → CopyOut 的流水重叠。
  一个 Ascend C 算子的典型骨架：

// Ascend C Flash Attention 骨架（简化，突出多核切分与持久化循环）
// Host 侧 tiling 结构体：决定每个 AI Core 处理哪些 Q 块
struct FlashAttnTiling {
    int32_t qBlockNum;      // Q 方向基本块总数
    int32_t kvBlockNum;     // K/V 方向基本块总数
    int32_t qBlockSize;     // Br
    int32_t kBlockSize;     // Bc
    int32_t headDim;        // d
    int32_t coreNum;        // 实际使用的 AI Core 数（<= 物理核数）
    int32_t qBlockPerCore;  // 每个核分到的 Q 块数（向上取整）
};

class KernelFlashAttention {
public:
    __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR o,
                                const FlashAttnTiling& tiling) {
        this->tiling = tiling;
        qGm.SetGlobalBuffer((__gm__ half*)q);
        kGm.SetGlobalBuffer((__gm__ half*)k);
        vGm.SetGlobalBuffer((__gm__ half*)v);
        oGm.SetGlobalBuffer((__gm__ half*)o);

        // 双缓冲队列：CopyIn 与 Compute 重叠
        pipe.InitBuffer(qQueue, 2, tiling.qBlockSize * tiling.headDim * sizeof(half));
        pipe.InitBuffer(kQueue, 2, tiling.kBlockSize * tiling.headDim * sizeof(half));
        pipe.InitBuffer(vQueue, 2, tiling.kBlockSize * tiling.headDim * sizeof(half));
        pipe.InitBuffer(accBuf, tiling.qBlockSize * tiling.headDim * sizeof(float));
        pipe.InitBuffer(mBuf, tiling.qBlockSize * sizeof(float));
        pipe.InitBuffer(lBuf, tiling.qBlockSize * sizeof(float));
    }

    __aicore__ inline void Process() {
        // 关键：手动多核切分。当前核只处理 [qBegin, qEnd) 区间的 Q 块
        int32_t coreIdx = GetBlockIdx();
        int32_t qBegin  = coreIdx * tiling.qBlockPerCore;
        int32_t qEnd    = min(qBegin + tiling.qBlockPerCore, tiling.qBlockNum);

        // 持久化循环：一个 AI Core 依次消化分配给它的多个 Q 块
        for (int32_t qi = qBegin; qi < qEnd; ++qi) {
            CopyInQ(qi);                 // DataCopy: GM -> UB（异步）
            InitOnlineSoftmaxState();    // m = -inf, l = 0, acc = 0

            // causal 场景下，K/V 只需遍历到对角线所在块
            int32_t kvEnd = GetKvBlockEnd(qi);
            for (int32_t kj = 0; kj < kvEnd; ++kj) {
                CopyInKV(kj);            // DataCopy: GM -> UB（异步）
                ComputeS();              // Cube: S = Q_i * K_j^T * scale
                OnlineSoftmaxUpdate();   // Vector: rowmax / exp / sum + 修正
                AccumulateO();           // Cube: acc += P * V_j
            }
            NormalizeAndCopyOut(qi);     // acc / l -> O
        }
    }

private:
    __aicore__ inline void OnlineSoftmaxUpdate() {
        // Vector 单元整块运算，避免逐元素循环
        LocalTensor<float> s = sBuf.Get<float>();
        LocalTensor<float> m = mBuf.Get<float>();
        LocalTensor<float> l = lBuf.Get<float>();
        LocalTensor<float> acc = accBuf.Get<float>();

        // 1) 本块行最大值，并与历史 m 合并
        Max(mNew, s, tiling.kBlockSize);          // 块内行 max
        Max(mNew, mNew, m, tiling.qBlockSize);    // 与 running max 合并

        // 2) 修正因子 corr = exp(m_old - m_new)
        Sub(corr, m, mNew, tiling.qBlockSize);
        Exp(corr, corr, tiling.qBlockSize);

        // 3) P = exp(S - m_new)，同时更新 l
        Sub(s, s, mNew, tiling.qBlockSize);
        Exp(s, s, tiling.qBlockSize);
        Mul(l, l, corr, tiling.qBlockSize);       // 历史 sum 修正
        ReduceSum(lNew, s, tiling.kBlockSize);    // 本块 sum
        Add(l, l, lNew, tiling.qBlockSize);

        // 4) acc = acc * corr + P * V（P*V 由 Cube 完成，此处只做修正）
        Muls(acc, acc, corr, tiling.qBlockSize * tiling.headDim);
        Copy(m, mNew, tiling.qBlockSize);         // 更新 running max
    }

    TPipe pipe;
    TQue<QuePosition::VECIN, 2> qQueue, kQueue, vQueue;
    TBuf<QuePosition::VECCALC> accBuf, mBuf, lBuf, sBuf;
    GlobalTensor<half> qGm, kGm, vGm, oGm;
    FlashAttnTiling tiling;
};

Host 侧 tiling 函数负责把「总块数」均匀分给「核数」：

// Host 侧：多核切分，保证每个核负载均匀且不超过物理核数
void FlashAttnTilingFunc(const FlashAttnTilingInput& in, FlashAttnTiling& out) {
    int64_t qBlockNum = (in.seqLen + in.qBlockSize - 1) / in.qBlockSize;
    int64_t kvBlockNum = (in.seqLen + in.kBlockSize - 1) / in.kBlockSize;

    // 核数取 min(物理核数, Q 块数)，避免空核
    int32_t coreNum = std::min(in.physicalCoreNum, (int32_t)qBlockNum);
    out.coreNum = coreNum;
    // 向上取整，保证所有 Q 块都被覆盖
    out.qBlockPerCore = (qBlockNum + coreNum - 1) / coreNum;
    out.qBlockNum = qBlockNum;
    out.kvBlockNum = kvBlockNum;
    out.qBlockSize = in.qBlockSize;
    out.kBlockSize = in.kBlockSize;
    out.headDim = in.headDim;
}

可以看到，整体算法结构与 CUDA 版完全一致：外层 Q 块、内层 KV 块、在线 softmax 状态、累加器归一化。区别在于：

• 数据搬运是显式的 DataCopy（GM → UB），而非依赖 cache；
• 矩阵乘走 Cube 单元，向量运算走 Vector 单元，两者需要流水编排才能重叠；
• 通过 TQue 双缓冲实现 CopyIn 与 Compute 的并行；
• 多核切分由 Host 侧 tiling 显式完成，每个 AI Core 通过 GetBlockIdx() 认领自己的 Q 块区间，并在核内循环消化多个块。

3.3 AI Core 数量受限：Tiling 与调度的本质差异

这是昇腾与 GPU 在底层实现上最容易被忽略、却极其关键的一点：任务分配与调度的责任方不同。

CUDA：硬件自动调度。 启动 Kernel 时指定的 Grid 大小可以远大于物理 SM 数量。硬件调度器会自动把线程块（Block）分发给空闲的 SM，开发者通常不需要关心哪个 Block 跑在哪个 SM 上。一个 SM 在运行时也可以驻留多个 Block，通过 Block 切换实现延迟隐藏。换句话说，Block 是短生命周期的，调度是硬件的事。

昇腾：Host 侧显式绑定。 AI Core 数量有限（通常数十个），且 Grid 大小直接受限于 AI Core 总数。在涉及多核同步的算子中，框架甚至要求逻辑核数必须不大于实际物理核数，否则会异常。因此 Tiling 阶段必须主动做「多核切分」：把总计算量划分成若干块，并显式指定每个 AI Core 处理哪一部分数据、从哪个地址偏移开始读。

每个 AI Core 承担多个 Block 任务。 因为 AI Core 数量少，而 Flash Attention 的计算量（尤其是长序列）往往很大，单个 AI Core 在 kernel 执行期间需要循环处理多个基本块。昇腾的 Tiling 设计明确要求「将核心的数量用满」且「每一个核心被分配的计算量相对均匀」。这意味着 Tiling 算法会计算每个核应处理的基本块数量，然后在 Kernel 内部通过循环让每个 AI Core 依次完成分配给它的多个块。

这本质上是一种持久化核（Persistent Kernel） 的思路：启动与物理核数相当的执行单元，让每个执行单元在内部迭代处理多个逻辑任务，而不是像 CUDA 那样依赖硬件调度器来动态分发大量短生命周期的 Block。

补充：CV 分离带来的额外约束。 昇腾 AI Core 内部 Cube（矩阵）和 Vector（向量）是分离的，这使 Tiling 更复杂。核间切分后还要保证「AIC 和 AIV 之间处理的数据量要符合其对应的算力」，避免一方空闲。所以昇腾的 Tiling 不仅要管「分给哪个核」，还要管「Cube 和 Vector 各分多少」，这是 CUDA 侧不需要考虑的维度。

总结：CUDA 的 Block 调度是「硬件自动分发、SM 动态驻留」；昇腾是「Host 侧 Tiling 显式绑定、AI Core 循环消化」。前者靠硬件隐藏延迟，后者靠均匀切分和持久化循环来压榨有限的核资源。

3.3 关键优化点

（1）Cube 与 Vector 的流水编排

这是昇腾 Flash Attention 性能的核心。Cube 算 $Q{K}^{\top }$ 和 $PV$，Vector 算 softmax，两者若串行执行，硬件利用率只有一半。昇腾通过多缓冲 + 异步流水让第 $j$ 块的 Vector 计算与第 $j+1$ 块的 Cube 计算重叠：

（2）Softmax 的向量化实现

昇腾 Vector 单元支持批量 Exp、Max、ReduceSum 等指令。在线 softmax 的 exp(s - m) 可以整块向量化，避免逐元素循环。同时用 exp2 替代 exp（乘以 ${\log }_{2}e$ 常数）可进一步提速，这与 GPU 上的技巧一致。

（3）分块大小与 UB 容量匹配

UB 容量有限（通常几百 KB 级别），$B_r$、$B_c$ 的选择必须让 $Q$ 块、$K/V$ 块、累加器同时放得下。Tiling 参数通常在 Host 侧根据 d、N、dtype 动态计算，这也是 Ascend C 算子需要写 tiling 函数的原因。

（4）Causal Mask 的块级跳过

对于因果注意力，当 $K/V$ 块完全位于对角线右侧时，整块跳过计算，直接省掉约一半的 Cube 和 Vector 工作量。昇腾实现中通常在 tiling 阶段就标记出需要计算的块范围。

（5）融合与算子边界

昇腾生态中，Flash Attention 常作为融合算子出现在 MindSpore 或 PyTorch-NPU 的图优化中，与前面的 QKV 投影、后面的输出投影做算子融合，减少 HBM 往返。这与 GPU 上 flash-attn 库的定位类似。

4. 共性与差异分析

4.1 共性：算法层面的"不变量"

无论 CUDA 还是昇腾，Flash Attention 的算法内核完全一致：

优化思想	CUDA 实现	昇腾实现
分块计算	__shared__ 存 Q/K/V 块	UB 存 Q/K/V 块
在线 Softmax	寄存器维护 m/l/acc	Vector 寄存器 + UB 维护
避免 N×N 落盘	中间结果留 SRAM	中间结果留 UB
反向重计算	保存 O/m/l 重算 P	同左
exp2 加速	__expf / exp2f	Vector Exp 指令
Causal 块跳过	块级判断	tiling 阶段标记

核心结论：Flash Attention 是算法级优化，与硬件无关。 它把"访存瓶颈"转化为"片上计算"，这个思想在任何有片上缓存的加速器上都成立。

4.2 差异：硬件架构决定的实现分野

维度	CUDA / GPU	昇腾 / NPU
计算单元	SM 内 Tensor Core + CUDA Core 混合	Cube（矩阵）+ Vector（向量）分离
数据搬运	隐式 cache（L1/L2），cp.async 可选	显式 DataCopy，必须手动管理
片上存储	Shared Memory（可编程）+ 寄存器	UB + L1/L0，容量与层级不同
并行模型	线程/warp/block 三级	AI Core + 流水队列
编程语言	CUDA C++	Ascend C（C++ 扩展）
流水控制	硬件 warp 调度 + 软件 pipeline	显式 TQue 双缓冲编排
生态库	flash-attn、cuDNN	MindSpore 融合算子、CANN 算子库

最本质的差异在于"数据搬运的显式性"：GPU 依赖多级 cache 自动缓存热点数据，程序员主要关注计算；昇腾要求程序员显式 DataCopy 并编排流水，控制力更强但开发门槛更高。这也导致昇腾 Flash Attention 的实现代码中，搬运与计算的流水编排往往比算法本身更占篇幅。

4.3 性能调优的侧重点

• GPU：重点在 warp 级 MMA 排布、寄存器压力、cp.async 流水深度、bank conflict 规避；
• 昇腾：重点在 tiling 参数与 UB 容量匹配、Cube/Vector 流水重叠率、DataCopy 的搬运效率、多核负载均衡。

两者都需要针对具体 shape 做 autotune，但调优的"旋钮"完全不同。

5. 总结

Flash Attention 的价值在于它揭示了一个通用规律：当算子的算术强度不足时，优化访存比优化算力更有效。它通过分块 + 在线 softmax，把 Attention 从"访存密集型"改造成"计算密集型"，从而在 GPU 上获得数倍加速和线性显存占用。

昇腾上的实现证明了这一思想的硬件无关性：算法骨架不变，变的是数据如何搬运、计算单元如何编排。理解 CUDA 版有助于掌握算法本质，理解昇腾版则能体会"显式流水"编程模型的特点。

对于开发者而言，实践建议是：

1. 先吃透在线 softmax 的数学推导，这是所有实现的基石；
2. 在 GPU 上读一遍 flash-attn 源码，理解 warp 级优化；
3. 在昇腾上从 Ascend C 的 tiling 与流水入手，先跑通再调优；
4. 关注 causal mask、变长序列、GQA/MQA 等实际场景的适配。

Flash Attention 之后，Flash Decoding、Paged Attention 等变体不断涌现，但"分块 + 片上计算 + 在线归一化"的内核思想始终贯穿其中。