Qwen系列模型线性注意力高性能优化实践

随着大模型序列长度从几万向百万级迈进，softmax 注意力 O(N²) 的内存与计算开销已成强约束。近年来，以 Qwen 和 Kimi 为代表的主流模型逐步转向"线性注意力 + softmax 注意力"的混合架构，但线性注意力的高效 GPU 实现仍面临并行度不足、访存冗余与 SFU 负载过重等系统性难题。

本文整理自阿里巴巴 Qwen 团队算法工程师章程瑞东在 AICon 2026 深圳站的分享《面向 Qwen 系列模型线性注意力的高性能优化实践》。他详细介绍了 FlashQLA——一个面向 GDN（Gated Delta Network）线性注意力的高性能 GPU 算子库，围绕 FlashQLA 的定位与背景、序列并行方案、滑动窗口性质利用、硬件友好代数改写以及 warp specialization 实现等核心技术展开，并给出性能基准与生态集成进展。

以下是演讲实录（经 InfoQ 进行不改变原意的编辑整理）。

线性注意力的核心动机来自一个朴素的矩阵连乘观察。softmax 注意力的计算流程是先对 Q 与 K 做矩阵乘法，再对每一行做 softmax 归一化，最后与 V 相乘。即便 FlashAttention 避免了显式实例化 N×N 的 P 矩阵，它仍然需要在 P 上完成逐行且带数据依赖的 softmax，因此计算量依旧是 O(N²d)。

但如果我们暂时忽略 softmax 这个非线性，仅从矩阵连乘角度重新审视：Q 与 K 的形状都是 sequence length 乘以 head dim，先算 QKᵀ 得到一个 N×N 矩阵，再与 V 相乘，总复杂度是 2N²d；反过来，如果先算 KᵀV，则得到 d×d 矩阵，再与 Q 相乘，两次矩阵乘法的复杂度都变为 Nd²。也就是说，只要能够绕开中间那个逐行 softmax，注意力计算的时间复杂度就存在从 O(N²d) 降到 O(Nd²) 的空间。

沿着这个思路，一个最直接的做法是保留因果下三角矩阵、去掉 softmax。因果性来自时序依赖，不能丢弃；softmax 则可以移除。这样我们得到一种递推形式：不再一次性计算整条序列的 KᵀV，而是一个 token 一个 token 地累加，每个 token 的 K 与 V 外积被写入一个 d×d 的状态矩阵，然后由该状态矩阵递推到下一步，并最终计算输出。这就是线性注意力最基本的时间 O(Nd²) 逻辑。

这个递推公式在结构上类似 RNN。它的好处是时间复杂度随序列长度线性增长，但代价同样是明显的：必须从前一个 token 的 state 才能获得下一个 token 的 state，数据依赖非常强。从 2022 年到 2025 年，学术界围绕这一框架衍生出大量线性注意力变体。其中第一个真正具备 scale up 能力的是 2025 年初由杨松林提出的 Gated Delta Network，即 GDN。GDN 的 state 更新公式演进得相当复杂，新引入了两个关键参数：更新门 beta 和遗忘门 alpha。其中遗忘门 alpha 在后续优化中扮演了非常关键的角色。目前来看，GDN 是表达力和效率兼顾得较好的线性注意力方案，相比原始 linear attention 有明显提升。此后还有不少 DeltaNet 变体，例如 Kimi K3 的 KDA。

Qwen 与 Kimi 目前主导的混合架构有一个共同设计思路：将大部分注意力层替换为线性注意力，线性注意力在算法上有一个明显缺陷，即无论序列多长，都必须把所有信息压缩到 d×d 的状态空间内，信息承载量天然有限，因此很难做好极长程召回。为了弥补这一点，模型会在每隔几层插入原始 softmax attention。Qwen 系列采用三层 GDN 加一层 softmax attention 的三比一结构，Kimi K3 则是三层 KDA 加一层 softmax attention。

在 GDN 模块的常规 prefill 中，直接按递推公式逐 token 计算会导致一个问题：公式中的 k 与 v 都是向量，计算依赖向量点积和向量外积，无法高效利用 GPU 上的 tensor core。解决思路是把递推步长从单 token 扩展到矩阵乘法单元的输入尺寸。这里需要引入一个数学变换，即 WY representation。变换后的递推从 token-by-token 变为 chunk-by-chunk，chunk size 通常取 64。

具体分四步：第一步对每个 chunk 内的 K 矩阵做 K Kᵀ，得到 chunk size × chunk size 的矩阵并求逆；第二步利用逆矩阵对 K 与 V 做中间投影，得到 W 和 U；第三步执行 chunk-by-chunk 的状态递推；第四步根据递推得到的 S 计算输出。值得注意的是，除了第三步之外，第一步、第二步和第四步不存在数据依赖，不同 chunk 的工作可以独立并行。但第三步必须拿到前一个 chunk 的 state 才能计算后一个 chunk 的 state，这是并行度的关键瓶颈。

FLA 仓库的现有实现中，这四步各对应一个独立算子，中间会实例化 W、U、S 三个变量。它们只在递推过程中作为临时变量存在，推理框架并不需要保留，但仍然会写回 HBM，再被下一步算子读取，造成纯 HBM 带宽开销。另外还有一个参数 gamma，表示每个 chunk 内 gate 的累乘，gate 就是前面提到的 alpha。

这个四步前向流程有三个主要痛点。第一是并行度不足。第三步要求每个 chunk 的 state 依赖前一个 chunk，因此每个 attention head 只能开出一个 thread block（CTA），对应一个 SM。以 Qwen 小模型为例，线性注意力头数可能只有 32，而 H100 或 GB 系列卡往往有一百多个 SM，GPU 利用率严重不足。这种并行度限制还会对 TP tensor parallel 优化带来挑战。第二是状态传递开销。W、U、S 三个中间变量需要写回 HBM 又被后续算子读取，但它们并不进入推理框架的持久状态，实质上是在白白占用带宽。第三个问题是 FLA 中使用的 Triton 算子的能力上限。Triton 无法精细调控算子内部，也没有提供手写 warp specialization 的接口。当多个矩阵乘法被合并到同一个算子内，并且还要做 chunk-by-chunk 递推时，这就非常像传统矩阵乘法的流水线，Triton 很难做到理想的流水线编排。因此我们需要更底层的工具来人工编排合理的流水线，以打满 GPU 的各类硬件单元。

要解决并行度不足，一个自然的思路是让序列并行化。此前 FLA 仓库中已有 GDN 序列并行的前序工作。核心问题可以这样表述：假设我们在序列中间切一刀，让后半部分从零初始状态开始递推，那么递推出的状态一定是错的。我们要问的是，如何矫正这个错误？

设 S 为正确状态序列，S* 为从错误初始状态开始递推得到的状态序列。对两者递推公式做减法，可以得到一个重要性质：相邻两个 chunk 的状态差异之间存在一个矩阵乘法关系，下一步的 difference 等于上一步的 difference 乘以一个 dK×dK 的矩阵。进一步利用矩阵乘法的结合律，我们可以为每一个从错误状态递推出的 chunk 独立计算矫正矩阵。也就是说，不需要先把矫正矩阵乘到 S 上，而是可以先把每个 chunk 对应的矫正矩阵计算出来，最后再用一次矩阵乘法完成矫正。这个观察到 dK×dK 矩阵的递推可以被"结合"到一起，是序列并行的理论前提。

但直接为每个 chunk 都存储一个 dK×dK 矫正矩阵代价太大。以 dK=128、chunk size=64 为例，每个 chunk 存一个 128×128 矩阵，相当于把之前通过算子融合省掉的状态传递开销又补了回来。实际实现中，我们采取的是"计算换访存"策略：每个 CP rank 只需要保存一个矫正矩阵。具体做法是每个 CP rank 内先独立计算各 chunk 的错误状态 S* 与矫正矩阵 M，但只保留该 rank 最后一个 chunk 的结果；然后在 CP size 长度上做一次状态依赖的递推，依次获得每个 rank 的精确矫正；最后重新计算一遍真正的 S 和 O。这个方案意味着 S 被算了两遍，第一遍是 S*，第二遍才是真正的 S，存在一定计算浪费，但显存上非常划算。这套方案逻辑上并不区分卡内序列并行与卡间序列并行，FLA 仓库中目前只支持卡间序列并行，但其原理完全一致，算法层面没有必要刻意区分二者。

解决了并行度问题之后，是否就可以直接做 fully fused 算子？答案是否定的。因为 fully fused 算子中有些计算与序列并行前序依赖相关，有些则无关。最关键的相关量是矩阵求逆。这个求逆在 GPU 上并不算便宜，而计算序列并行矫正矩阵时要使用矩阵求逆之后的结果 W。如果把求逆放到后续算子中，那么在计算序列并行前序数据依赖 S* 和 M 时，就必须自己重新算一遍求逆，这非常不划算。因此我们选择在四步方案的第一步之后切一刀：矩阵求逆单独拎出来，后面三步可以融合，而序列并行相关操作则插在这两部分之间。

卡内序列并行与卡间序列并行有一个重要差别：卡间并行度是人为固定的，而卡内并行时，我们并不知道一个给定长度的序列应该切成几份。这个数量理论上是可以估算的。回到前述三步并行方案：第一步与第三步在每个 CP rank 之间可以独立运行，其递推长度为每个 rank 内的 chunk 数 L；第二步的递推长度是 N/L，即 CP size。在做 one wave 简化模型时，第一步和第三步的时间正比于 L，第二步的时间正比于 N/L，因此总时间 T = λ₁L + λ₂N/L。求极小值即可得到最优 L，它一定出现在根号 N 的线性关系上。λ₁ 与 λ₂ 的具体数值来自算子在特定硬件上的实际性能，可以通过调优获得。目前我们还在优化这个模型，one wave 简化过于粗糙，后续会针对每个算子引入更多参数，获得更精确的序列并行方案估计。

并行方案虽然可以插入，但插入本身会带来 overhead，因此需要决定何时启用。一个显然的条件是总投数足够小时触发，即 GPU 的 SM 明显没有被占满。但还有一个容易被忽略的隐藏条件：如果序列本身不够长，插入序列并行后的递推也没有什么意义。这里并不是简单的 and 或 or 关系。若以 batch size 乘以 head number 为横轴、以序列长度为纵轴，启用 auto CP 的合理边界应该是一条二次曲线。当前实现做了一个简化的梯形切分，后续会替换为更精确的方案。

GDN 本身还有一个此前未被充分利用的数学性质，即滑动窗口性质。观察原始 GDN 递推公式：alpha 是 0 到 1 之间的实数，beta 也是 0 到 1 之间的实数，同时 K 是经过 L2 norm 的，每个 token 的 K 模长受约束。因此 I 减去 beta K Kᵀ 这个矩阵的行列式有明确数值范围。而每一步累加都会乘上一个 alpha，这意味着每个 state 对后续 state 的影响力会随着 alpha 指数递减。具体来说，S_t 对 S_{t+1} 的影响被 alpha 衰减一次，对 S_{t+8} 的影响则被衰减八次。当窗口足够长时，早期 state 对当前状态的影响可以视为零。也就是说，只要一个 head 上的 alpha 不恒为 1，就一定会存在一个滑动窗口，窗口大小由 alpha 分布和 gate 数据分布共同决定。

我们对训练出的模型做了统计，大约有 60% 到 80% 的线性注意力头的 alpha 不恒为 1。这个比例很合理：GDN 中加入的 gate 参数 alpha 正是为了提升表达力。如果 gate 不发挥作用，就退化为表达力不足的 DeltaNet。换句话说，只要训练让 gate 发挥了作用，滑动窗口现象就会出现。对于这些具备滑动窗口性质的 head，序列并行的方案可以大幅简化：不需要从第 0 个状态开始递推，而只需对每个 chunk 的起点向前回溯一个足够长的 window size，直接获得误差足够小的起始状态。这与 sliding window attention 的序列并行思路一致。

更有利的是，这个滑动窗口方案与前面的序列并行方案可以复用算子。序列并行方案中一定需要计算 S*，而对于滑动注意力头，S* 本身就是精确的 S，可以直接使用，不再需要计算 M 矩阵和第二步矫正。第一步算子可以完全复用。在决定哪些 head 采用滑动窗口优化时，需要额外启动一个轻量 kernel 扫描每个 head 的 gate 数据分布。对 GDN 而言，gate 是每个 token 上的一个实数，扫描成本与 gate 累加相当，overhead 可以忽略。

围绕 GDN 的递推过程，我们做了两类硬件友好的代数改写：一类针对 SFU（special function unit）的负载，另一类针对矩阵连乘的时间复杂度。

先看 SFU 负载。在 GDN 计算公式中，小 gamma 表示 alpha 的累乘，而大 Gamma 是一个 chunk 内的小 gamma 外积矩阵。形式上，大 Gamma 的第 s 行第 t 列元素等于 gamma_i_s 除以 gamma_i_t。如果直接在数值上计算小 gamma 分之一会存在溢出风险，尤其当 chunk size 较大时，gamma 很容易趋近于零，BF16 精度可能无法表示。更安全的做法是显式展开，利用 log alpha 存储、以加法替代乘法：gamma_i_s / gamma_i_t 等于 exp(g_i_s - g_i_t)。于是，计算这个 C×C 矩阵的每一项都需要一次 EXP 运算，共需 C² 次 EXP。这正是 SFU 负载的来源。

关键观察是：在矩阵求逆过程中，大 Gamma 可以被提取出来。原求逆公式中，大 Gamma 位于 K Kᵀ 之后。由于大 Gamma 的对角线恒为 1，可以把它与单位矩阵合并，进而将大 Gamma 提到外层。

形式上，大 Gamma 可以写为小 gamma 乘以单位矩阵再乘以小 gamma 分之一的形式，在下三角矩阵求逆时，这个对角缩放可以整体提出。提出后的求逆结果记为 A'。接着在计算 W 和 U 时，再把大 Gamma 乘回 A' 上。由于 W、U、S 在后续算子中本来就要计算大 Gamma，这样做避免了在两个算子中分别计算 O(C²) 次 EXP 的问题，显著降低了 SFU 单元的压力。

第二类改写围绕矩阵连乘。以序列并行转移矩阵 M 为例：原始 M 是若干 128×128 矩阵的累乘，每一次矩阵乘法的时间复杂度是 128×128×128，实际很大。利用前述大 Gamma 提取技巧，可以得到不包含大 Gamma 的求逆矩阵 A'，重新推导 S 的递推公式，就能把递推公式中每 chunk 最后一个 gamma 除以所有 gamma 的项省略掉。此时递推所需变量可以整理为 X、Y、Z 形式，同时适用于 M。直觉上 M 与 S 非常相似，一定存在一种以 S 的计算复杂度来计算 M 的方式。整理后，M 递推中每 chunk 内 gate 累乘项被提到括号外。括号内部的 X 是求逆后的矩阵乘以 beta K Kᵀ，数值范围有保证；Z 是 K 乘以 M，同样有保证。因此整个 M 递推的数值范围是可控的，加上 gamma 衰减项后，M 矩阵本身就体现出随 gate 累乘而衰减的性质。这实际上也从另一条路径验证了 GDN 的滑动窗口性质。

前向递推与反向递推也有类似的矩阵连乘优化。前向方面，由于 W 与 U 包含同类项，可以将 alpha 归并，先算 V 减 S，再乘 gamma，从而减少矩阵乘法次数。这个方案在早先的 fully fused GDN 中已经存在，我们做的主要是搬运与整合。反向递推的优化比较零碎，其中一个代表性工作是：K Kᵀ 梯度反传回 K 时，需要计算 A 乘以 K 再加上 DA 的转置乘以 K，这里可以把 DA 与 DA 的转置合并同类项，再与 K 相乘，从而降低矩阵乘法负载。虽然对三十二位浮点矩阵做 shared memory 转置会对原本紧张的 shared memory 带宽提出更高要求，但在整个反向流程中，该合并对后续 gate 与 beta 梯度的 reduce 操作同样有利，因此是一个整体优化的选择。

经过代数改写和融合之后，后三步算子在计算量上仍小于访存量。以我们的自定义单位衡量，一个 chunk 内 64×64 的 BF16 矩阵访存量记作一个单位，Q、K、V、A 加上输出 O 的总访存量大约是 18 个单位；而 H200 上 WGMMA 的 tensor core 计算带宽等效映射后，矩阵乘法计算量约为 30 个单位。也就是说，计算量并不明显低于 HBM 访存量，但要注意这里直接以 HBM 带宽计算，不考虑 L2 命中。对 GDN prefill 来说，L2 命中确实较低，因为每块 Q 和 K 用完后不会像普通矩阵乘法那样被不同 SM 重复利用。总体而言这仍然是一个 memory bound 算子。

基于此，我们设计了乒乓式 two-stage 流水线。传统 multi-stage 方案往往是为了利用 L2 缓存，但 GDN 算子本身 L2 命中率低，two-stage 已经足够。寄存器方面，64K 个 128 位寄存器平均分到四个 warp group，每个线程约 128 个寄存器，相当紧张。尤其是 S 矩阵为 128×128，摊到每个 warp group 后每个线程正好占满 128 个三十二位寄存器。shared memory 同样紧张，Q、K、V 等有 head dim 维度的输入若做双缓冲，总计约 9 个单位，每个单位对应 28KB，刚好容纳乒乓结构，但中间变量需要做极致的 shared memory 复用。

流水线分工上，我们按照 warp group 划分职责。在 Hopper 上，WGMMA 指令由整个 warp group 发起，因此需要以 warp group 为粒度来编排。设计中共四个 warp group：一个生产者 warp group 负责通过 TMA 从 HBM 搬运 Q、K、V，同时该 group 也承担消费者角色，负责把输出 O 写回 HBM；剩下三个消费者 warp group 分别承担 V' 计算、O 计算和 S 计算。为什么让 S 独占一个 warp group？因为 S 需要常驻寄存器，独占最合理。V 和 O 的中间临时变量较多，分别放到两个 warp group 可以留出足够空间。通过这种 warp specialization 设计，我们把数据搬运、tensor core 计算、CUDA core 计算以及 SFU 计算全部重叠起来。作为 memory bound 算子，计算负载不会太高，但在 NCU 上可以看到访存负载超过 70%。

最终性能方面，前向获得两到三倍加速，反向获得约两倍加速。反向优化逻辑与前向基本一致，也包括基于滑动窗口的卡内序列并行方案。反向算子虽然 warp specialization 更复杂，但框架相通。这套算子仓库是训练推理通用的，目前千问内部训练和推理链路都已接入。在 TP 并行度较大、或者模型较小的场景中，auto CP 生效后加速效果尤其显著。与 FLA 的 Triton 版本对比，常见场景下均有两到三倍加速，在挑选的有利场景中可以接近三倍。与 FlashInfer 的较早版本相比，TP 较大时我们的方案有更大的压缩空间，这主要得益于解决了 SM 利用率问题。基准覆盖了 Qwen3.5 / Qwen3.6 系列等常见配置，GDN head 配置从最小 8 到最大 128 都有覆盖，不同 head 配置上均能复现加速效果。硬件方面，最初为 Hopper 设计，后续陆续支持了 SM100、SM103、SM120 和 SM121，覆盖整个 Blackwell 家族。

FlashQLA 在工程上对齐了 FLA 的接口约定。它提供与 FLA 签名一致的高级 API，原本使用 chunk_gated_delta_rule 的代码可以直接切换为 from flash_qla import chunk_gated_delta_rule。接口中新增了 auto CP 开关。开启后我们能保证性能不下降，但无法保证确定性推理和训推一致等特殊需求，因此有此类需求的用户可以手动关闭。项目已发布到 PyPI，可以直接 pip install。从 v0.1.2 版本之后，FlashQLA 成为 FLA 的官方后端，也就是说用户安装 FLA 后，只要硬件满足条件，就能直接调用 FlashQLA。社区也帮助我们完成了部分 vLLM 集成。

项目时间线上，2026 年 4 月 24 日首次开源，6 月完成了反向卡内序列并行与 SM100 支持，7 月在社区帮助下完成 SM120 支持，并正式成为 FLA 官方后端。目前我们正在推进的工作包括两方面：一是卡间序列并行的推广。由于序列并行算子在算法层面不分卡内卡间，推广是自然延伸，主要面向 PD 分离的超长序列推理与超长数据预训练。二是新的 auto CP 算法。目前 one wave 模型过于简化，后续会针对算子性能做更细致的建模，以获得更精确的序列并行拆分决策。

总体而言，FlashQLA 的目标是覆盖大模型预训练、端侧推理以及 TP 长序列部署等场景，补全开源模型生态。所有大模型均可覆盖；小模型在端侧推理和单卡训练上有优化；大模型则在开 TP 时的序列并行上有优化，同时支持变长序列。核心贡献有三项：第一是 auto CP 卡内序列并行，并利用了 GDN 此前未被系统性强调的滑动窗口性质；第二是硬件友好的代数改写，包括提取大 Gamma 以降低 SFU 计算量，以及矩阵连乘顺序优化；第三是使用 TileLang 手写了多个 warp specialization 算子。整体性能在 Hopper 与 Blackwell 家族上均可复现，前向与反向加速两到三倍，且开箱即用、作为官方后端存在。最后要感谢 Flash Linear Attention、TileLang、FlashInfer 相关工作人员，以及 GDN 序列并行的前序工作。

会议推荐