Qwen 团队开源 FlashQLA：面向 Gated Delta Network 的融合线性注意力算子库
FlashQLA：面向 Gated Delta Network 的融合线性注意力算子库
Qwen 团队开源 FlashQLA，一个基于 TileLang 的 GDN 线性注意力算子库，在 NVIDIA Hopper 上前向比 FLA Triton 内核快 2-3 倍、反向快 2 倍。
FlashQLA：面向 GDN 的 CP/反向友好融合线性注意力内核
引言
随着 Qwen3-Next 的发布，门控 Delta 网络（GDN）已成为 Qwen 系列中承担主力作用的注意力层——从 Qwen3-Next-80B-A3B 一直到后续的 Qwen3.5 / Qwen3.6 系列。随着模型规模扩展到 397A17B / 122A10B / 35B / 27B，上下文窗口延伸至 256K 以上，GDN 模块在端到端训练和推理中的开销已变得不可忽视。
今天我们正式开源 FlashQLA：一个基于 TileLang 构建的高性能线性注意力内核库。FlashQLA 对 GDN Chunked Prefill 的前向与反向传播进行了合理的算子融合与性能优化，在 NVIDIA Hopper 上的多个场景中，相比 FLA Triton 内核实现了 2-3 倍前向加速和 2 倍反向加速。这一效率提升在预训练场景和边缘侧智能体推理中尤为显著。
本次发布的核心亮点：
- 门控驱动的自动卡内上下文并行。利用 GDN 门控的指数衰减特性，FlashQLA 在 TP、长序列和小头数设置下自动启用卡内 CP，提升 GPU SM 利用率。
- 硬件友好的代数重构。我们在一定程度上重构了 GDN Chunked Prefill 的前向与反向流程，在不牺牲数值精度的前提下有效降低了 Tensor Core、CUDA Core 和 SFU 开销。
- TileLang 融合 warp 特化内核。我们没有采用逐步分解为独立内核的方式，也没有将整个计算流程融合进单个内核，而是兼顾 CP 与反向需求，使用 TileLang 构建若干关键融合内核，并手动实现 warpgroup 特化，以重叠数据搬运、Tensor Core 计算和 CUDA Core 计算。
FlashQLA 的代码与基准测试已在 github.com/QwenLM/FlashQLA 开源。
FLA GDN Chunked Prefill 中的关键问题
我们先回顾 GDN Chunked Prefill 的前向计算流程，以 chunk 索引 \(i\) 为例：
- \(A_i \gets \left(I+\mathrm{StrictLower}\left( \mathrm{diag}(\beta_i)(\Gamma_i \odot K_iK_i^\intercal) \right)\right)^{-1}\)
- \(\left\{\begin{aligned} W_i &\gets A_i\mathrm{diag}(\beta_i)\mathrm{diag}(\gamma_i)K_i \\ U_i &\gets A_i\mathrm{diag}(\beta_i)V_i \end{aligned}\right.\)
- \(\left\{\begin{aligned} V_i’ &\gets U_i-W_iS_i \\ S_{i+1} &\gets \gamma_{i,C-1}S_i + K_i^\intercal\mathrm{diag}\left(\frac{\gamma_{i,C-1}}{\gamma_i}\right)V_i’ \end{aligned}\right.\)
- \(O_i \gets \mathrm{diag}({\gamma})Q_iS_i + \left(\mathrm{Lower}(\Gamma_i) \odot Q_iK_i^\intercal\right)V_i'\)
忽略门控预处理和 CP，该流程的每一步在 FLA 中都对应一个内核。该流程存在两个主要效率问题：
- 上述大部分是内存受限内核。该流程反复读取 \(K\)、\(V\) 等数据，而 \(W\)、\(U\)、\(S\) 作为中间变量必须写入 HBM，再由下一个内核读取，带来显著的内存访问开销。
- SSM 状态的循环特性意味着对应的第三步
chunk_gated_delta_rule_fwd_kernel只能同时启动batch_size * num_heads个线程块，导致在小模型、小批量或 TP 场景下 GPU 利用率偏低。
这两个问题的解决方案是相互矛盾的。对于第一个问题，最直观的解决方案是编写一个全融合内核，其中所有数据只被访问一次，所有中间变量都保留在片上。当 batch_size * num_heads 足够大时，这当然是最优的。然而，这样的方案显然会遇到第二个问题：对于小模型和 batch_size=1 的边缘侧推理，或者对于采用 TP 的大模型在线部署，其中来自编码代理等的长序列输入无法为分块预填充启动足够大的批次，全融合内核相对于原始 FLA 实现的加速是有限的。
第二个问题的最早解决方案来自 DeltaNet 处理上下文并行性的方式，它将长序列拆分为多个子序列，使用 \(S_0=0\) 来并行化递归，然后计算一个额外的 \(M\) 矩阵来校正递归结果。该方案后来被优化为在递归内核之前插入一个步骤来计算每个子序列的 \(S_0\)，现已合并到 FLA 仓库中。对于 CP 秩 \(j\)，具体的预处理流程为：
- \(\left\{\begin{aligned} S^\ast_{j,i+1} &\gets \gamma_{j,i,C-1}S^\ast_{j,i} + K_{j,i}^\intercal \mathrm{diag}\left(\frac{\gamma_{j,i,C-1}}{\gamma_{j,i}}\right) V’_{j,i} \\ M_{j,i+1} &\gets \left( \gamma_{j,i,C-1} I -K_{j,i}^\intercal \mathrm{diag}\left(\frac{\gamma_{j,i,C-1}}{\gamma_{j,i}}\right) W_{j,i} \right) M_{j,i} \end{aligned}\right.\)
- \(S_{j,0} \gets S^\ast_{j,0} + M_{j,0}S_{j-1,0}\)
然而，这种 CP 方案也有其缺点：首先，它引入了大量的额外计算，递归计算 \(M\) 矩阵的时间复杂度甚至超过了 \(S\) 矩阵；其次，它不能很好地与全融合内核配合，因为必须先执行矩阵求逆等步骤，才能计算每个子序列的 \(S_0\)。
一种平衡的解决方案：在融合内核的同时启用卡内 CP
基于上述两个问题，可以得出一种折中方案：将 GDN 分块预填充前向计算拆分为两个融合内核，并在它们之间插入与 CP 相关的预处理步骤。经过一些变换和简化后，得到以下计算流程：
- \(A_i \gets \left(I+\mathrm{StrictLower}\left( \mathrm{diag}(\beta_i)K_iK_i^\intercal\right)\right)^{-1}\)
- CP 预处理
- 2.1. \(\left\{\begin{aligned} X_{j,i} &\gets -\beta_{j,i} A_{j,i}’^\intercal K_{j,i} \\ Y_{j,i} &\gets \gamma_{j,i,C-1} K_{j,i} S^\ast_{j,i} - \mathrm{diag}\left(\frac{\gamma_{j,i,C-1}}{\gamma_{j,i}}\right) V_{j,i} \\ Z_{j,i} &\gets K_{j,i} M_{j,i} \\ S^\ast_{j,i+1} &\gets \gamma_{j,i,C-1}S^\ast_{j,i} + X_{j,i}^\intercal Y_{j,i} \\ M_{j,i+1} &\gets \gamma_{j,i,C-1} \left( M_{j,i} + X_{j,i}^\intercal Z_{j,i} \right) \end{aligned}\right.\)
- 2.2. \(S_{j,0} \gets S^\ast_{j,0} + M_{j,0}S_{j-1,0}\)
- \(\left\{\begin{aligned} V_i^\Delta &\gets V_i - \mathrm{diag}(\gamma_i)K_iS_i \\ V_i’ &\gets \left(\Gamma_i \odot A_i\right)\mathrm{diag}(\beta_i)V_i^\Delta \\ S_{i+1} &\gets \gamma_{i,C-1}S_i + K_i^\intercal\mathrm{diag}\left(\frac{\gamma_{i,C-1}}{\gamma_i}\right)V_i’ \\ O_i &\gets \mathrm{diag}({\gamma_i})Q_iS_i + \left(\mathrm{Lower}(\Gamma_i) \odot Q_iK_i^\intercal\right)V_i’ \end{aligned}\right.\)
我们还设计了一个简单的数学模型来自动确定并行度。设 \(N\) 为序列中的分块数，\(L\) 为每个 CP rank 的分块数。容易看出，步骤 2.1 和步骤 3 的运行时间与 \(L\) 成正比，而步骤 2.2 的运行时间与 \(\frac NL\) 成正比；因此我们可以选择 \(L=\lambda \sqrt N\) 来最小化总时间，其中 \(\lambda\) 是由 batch_size、num_heads 及其他超参数构成的系数。
在生产环境中，并非总是需要卡内 CP。沿用原始 FLA 实现，步骤 3 还可以通过拆分 v_head_dim 将并行度提高 2-4 倍，代价是对 Q 和 K 的冗余内存访问。根据实测数据，我们仅在 batch_size * num_heads <= 40 或 batch_size * num_heads <= 56 && seq_len >= 8192 时启用 CP。
通过门控衰减进一步优化
重新审视 GDN 递推式：
\[S_{i+1} = \alpha_iS_i(I-\beta_ik_ik_i^\intercal)+\beta_iv_ik_i^\intercal\]
对于 \(\alpha_i\in(0,1)\)，每个 \(S_i\) 对后续状态的影响呈指数衰减，从而具有滑动窗口特性。对于足够长的窗口大小 \(W\)，从 \(S_{i-W}=0\) 开始计算即可得到准确的 \(S_i\)，无需从 \(S_0\) 开始。我们将这一过程称为 warmup。在真实数据上，我们发现 \(\alpha_i\) 在 60–80% 的线性注意力头上并不恒为 1，且 6–8 个分块的 warmup 足以将 \(S_i\) 的误差降至噪声底以下。
因此，对于具有滑动窗口特性的线性注意力头，我们可以设计一种更轻量的 CP 预处理流程，舍弃修正项 \(M\) 的计算，直接通过 warmup 获得同样准确的子序列 \(S_0\)：
| C0 | C1 | C2 | C3 | C4 | C5 | C6 | C7 | C8 | C9 | C10 | C11 | C12 | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| R1 | O | O | O | O | O | ||||||||
| R2 | X | X | O | O | O | O | |||||||
| R3 | X | X | O | O | O | O |
X 表示以零初始状态进行 warmup，直到门控充分衰减，然后写出该 CP rank 的 \(S_0\)；O 表示后续正常的循环计算。每个 rank 的 warmup 长度由一个独立的 kernel 收集门控统计信息来确定，该步骤的开销可以忽略不计。
TileLang Warp 专用化 Kernel
我们使用 warpgroup 专用化模式在 TileLang 中实现 FlashQLA：一个生产者 warpgroup 和三个消费者 warpgroup 驻留在同一个 SM 中，通过共享内存交换数据，并通过 mbarrier 进行同步。
前向
在前向传播中，三个消费者 warpgroup 分别计算 \(V’\)、\(S\) 和 \(O\)，通过 ping-pong 结构重叠计算与内存访问。
| WG3 | WG2 | WG1 | WG0 | |||
|---|---|---|---|---|---|---|
| WG3/0 | WG3/1 | WG3/2 | ||||
| BAR 0 | LD\(Q\) | LD\(\gamma\) | ST\(O\) | \(\gamma, \gamma_{C-1}\gamma^{-1}\) | TC\(P = Q K^\intercal\) | |
| BAR 1 | LD\(K\) | LD\(\beta\) | ST\(S_i\) | TC\(U = K S_i\) | \(\Gamma = L(\gamma I \gamma^{-1})\) \(A_\gamma = \Gamma \odot A\) \(P_\gamma = s\Gamma \odot P\) | \(S_{i+1} = \gamma_{C-1} S_i\) |
| BAR 2 | LD\(V\) | \(W = \beta (V - \gamma U)\) | TC\(O = Q S_i\) | |||
| BAR 3 | LD\(A\) | TC\(V^\Delta = A_\gamma W\) | \(O = s\gamma O\) | |||
| BAR 4 | \(V’ = \gamma_{C-1}\gamma^{-1} V^\Delta\) | TC\(O = O + P_\gamma V^\Delta\) | ||||
| BAR 5 | TC\(S_{i+1} = S_{i+1} + K^\intercal V'\) | |||||
备注：
- 每个分块输出的 \(S\) 仅用于调试；通常只输出 \(O\) 和最后一个分块的 \(S\)。
CP 预处理
如前所述，CP 预处理分为两种情况：原始方法（同时计算 \(M\) 和 \(S\)）和滑动窗口方法（仅计算 \(S\)）。我们设计了一个融合 kernel 来同时处理这两种情况：
| WG3 | WG2 | WG1 | WG0 | |||
|---|---|---|---|---|---|---|
| WG3/0 | WG3/1 | WG3/2 | ||||
| BAR 0 | LD\(K\) | LD\(\gamma\) | \(\gamma_{C-1}\gamma^{-1}\) | TC\(X = A^\intercal K\) | ||
| BAR 1 | LD\(V\) | LD\(\beta\) | ST\(S_i\) | TC\(U = K S_i\) \(Y = -\gamma_{C-1}\gamma^{-1} V + \gamma_{C-1} U\) | \(X = -\beta X\) | \(S_{i+1} = \gamma_{C-1} S_i\) |
| BAR 2 | LD\(A\) | \(\gamma^\pi = \gamma^\pi \gamma_{C-1}\) | \(\gamma^\pi = \gamma^\pi \gamma_{C-1}\) | TC\(S_{i+1} = S_{i+1} + X^\intercal Y\) | ||
| BAR 3 | TC\(Z^L = K M^L\)
TC\(M^L = M^L + X^\intercal Z^L\) | TC\(Z^R = K M^R\)
TC\(M^R = M^R + X^\intercal Z^R\) | ||||
备注：
- WG1 和 WG2 的最后两步对应 \(M\) 矩阵计算，仅在需要时触发。
- \(S\) 仅在反向重计算期间按 chunk 输出。
反向
在反向传播中，我们复用上一节的 CP 预处理 kernel 来重新计算 \(S\) 矩阵，然后将 bwd_dv 、 bwd_dhu 、 bwd_dqkwg 、 bwd_wy 融合为单个 kernel，并施加相应的代数优化。由于片上资源限制，反向 kernel 不使用多级流水线；而是依靠长计算链来隐藏内存流量。完整调度可在 FlashQLA 仓库中获取。
| WG3 | WG2 | WG1 | WG0 | |
|---|---|---|---|---|
| BAR 00 | ST\(dK\) | TC\(P=QK^\intercal\) | \(\gamma, \gamma_{C-1}\gamma^{-1}\) | |
| BAR 01 | TC\(dV’=KdS_{i+1}\) \(dV’=\gamma_{C-1}\gamma^{-1}dV'\) | \(\Gamma=\gamma I \gamma^{-1}\)
\(P_\gamma=sL(\Gamma)\odot P\) | \(dS_i = \gamma_{C-1} dS_{i+1}\) | |
| BAR 02 | TC\(dV’=dV’+P_\gamma^\intercal dO\) | \(A_\beta = A \beta\)
\(A_\gamma = \Gamma \odot A_\beta\) | ||
| BAR 03 | TC\(U=KS_i\) | |||
| BAR 04 | TC\(dV=A_\gamma^\intercal dV’\) | \(W=V-\gamma U\) | \(d\gamma_{C-1}=\sum S_i \odot dS_{i+1}\) | |
| BAR 05 | ST\(dV\) | LD\(V\) | \(dV_\gamma = -\gamma dV\)
\(d\gamma = \sum_i dV_\gamma \odot U\) | TC\(dA_\gamma = dV’W^T\)
TC\(V’=A_\gamma W\) |
| BAR 06 | TC\(dP_\gamma = dO V’^\intercal\) | |||
| BAR 07 | LD\(K\) | TC\(dK=V’dS_{i+1}^\intercal\) | \(dA_\beta = \Gamma \odot dA_\gamma\)
\(d\gamma = d\gamma + \sum_i dP_\gamma \odot L(P)\) \(d\gamma = d\gamma - \sum_j dP_\gamma \odot L(P)\) \(dP = sL(\Gamma)\odot dP_\gamma\) | |
| BAR 08 | \(dK=\gamma_{C-1}\gamma^{-1}dK\)
\(d\gamma_{C-1}=\sum K \odot dK\) \(d\gamma = -\sum_i K \odot dK\) | TC\(dQ=dOS_i^T\) | ||
| BAR 09 | LD\(Q\) | TC\(dK=dK+dV_\gamma S_i^\intercal\) | \(dQ=s\gamma dQ\)
\(d\gamma = \sum Q \odot dQ\) | |
| BAR 10 | LD\(S\) | TC\(dQ=dQ+dPK\) | ||
| BAR 11 | ST\(dQ\) | \(d\gamma = d\gamma + \sum_i dA_\beta \odot A \beta\)
\(d\gamma = d\gamma - \sum_j dA_\beta \odot A \beta\) \(d\beta = \sum_j dA_\beta \odot A\) \(dA=dA_\beta \beta\) | TC\(dS_i = dS_i + K^\intercal dV_\gamma\) | |
| BAR 12 | TC\(dK=dK+dP^\intercal Q\) | |||
| BAR 13 | TC\(dA = -A^\intercal dA A^\intercal\)
TC\(A_T = KK^\intercal\) | \(dO_\gamma=s\gamma dO\) | ||
| BAR 14 | LD\(dO\)
LD\(A\) | \(d\beta = d\beta + \sum_i dA \odot A_T\)
\(dA_T = \beta dA\) \(dA_S = dA_T + dA_T^\intercal\) | TC\(dS_0 = dS_0 + Q^\intercal dO_\gamma\) | |
| BAR 15 | TC\(dK=dK+dA_S K\) | |||
基准测试
我们在 Qwen3.5 / Qwen3.6 系列所使用的 head 配置下，将 FlashQLA 与 FLA Triton 和 FlashInfer 基线（FLA 0.5.0、Triton 3.5.1、FlashInfer 0.6.9、TileLang 0.1.8）进行了基准测试——\(h_v \in {64, 48, 32, 24, 16, 8}\) ，分别对应 TP1 到 TP8。
具体而言，前向（FWD）基准测试测量不同模型和 TP 设置在不同 batch 长度下的单 kernel 延迟，而反向（BWD）基准测试则考察单次更新步骤中一个 batch 内的总 token 数与延迟之间的关系。
部分 H200 单层前向结果：
| 模型 / TP | Seqlen | \(h_{qk}\) | \(h_v\) | FlashQLA | FlashInfer | FLA | 对比 FLA | 对比 FI |
|---|---|---|---|---|---|---|---|---|
| 397B/122B TP8 | 1x32768 | 2 | 8 | 0.310ms | 1.653ms | 0.913ms | 2.95× | 5.33× |
| 397B/122B TP8 | 1x16384 | 2 | 8 | 0.184ms | 0.833ms | 0.465ms | 2.53× | 4.53× |
| 397B/122B TP8 | 24576+8192 | 2 | 8 | 0.302ms | 1.242ms | 0.767ms | 2.54× | 4.11× |
| 397B/122B TP4 | 1x32768 | 4 | 16 | 0.486ms | 1.654ms | 1.250ms | 2.57× | 3.40× |
| 397B/122B TP4 | 1x16384 | 4 | 16 | 0.292ms | 0.832ms | 0.623ms | 2.13× | 2.85× |
| 27B TP2 | 1x32768 | 8 | 24 | 0.659ms | 1.616ms | 1.564ms | 2.37× | 2.45× |
| 2B/0.8B TP1 | 1x32768 | 16 | 16 | 0.493ms | 1.640ms | 1.285ms | 2.60× | 3.33× |
| Sym h32 | 1x32768 | 32 | 32 | 0.877ms | 1.554ms | 1.952ms | 2.23× | 1.77× |
加速比随 TP 程度增长，因为 FlashQLA 通过卡内 AutoCP 在基线会让 SM 空闲的特定场景——TP 分片和小 head 数——中提升了 SM 利用率。
用法
FlashQLA 同时提供与 FLA 签名一致的高层 API 以及底层的 fwd/bwd 入口：
python
import torchfrom qla import chunk_gated_delta_rule
o, final_state = chunk_gated_delta_rule( q=q, # [B, T, H_q, K] k=k, # [B, T, H_q, K] v=v, # [B, T, H_v, V] g=g, # [B, T, H_v] beta=beta, # [B, T, H_v] scale=scale, initial_state=initial_state, # optional, [B, H_v, K, V] output_final_state=True, cu_seqlens=cu_seqlens, # optional, varlen support)
要求：SM90、CUDA 12.8+、PyTorch 2.8+。安装：
bash
git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA && pip install -v .
致谢
FlashQLA 受 Flash Linear Attention、FlashInfer 和 TileLang 启发。我们感谢这些社区提供的参考实现。
引用
如果 FlashQLA 对你的研究有帮助，请引用：
bibtex
@misc{flashqla2026, title = {FlashQLA: Flash Qwen Linear Attention}, author = {Zhang, Chengruidong and Lin, Xi and Jiang, Huiqiang and Wang, Zekun and
Li, Xiao and Cao, Yizhong and Zhuang, Bohan and Men, Rui and Zhang, Jianwei and
Zheng, Bo and Lin, Junyang and Liu, Dayiheng and Zhou, Jingren}, year = {2026}, publisher = {GitHub}, howpublished = {\url{https://github.com/QwenLM/FlashQLA}}}
来源：Qwen Blog · qwen.ai