FlashQLA优化Qwen线性注意力的GPU算子实践

随着大模型序列长度迈向百万级，传统 softmax 注意力的 O(N²) 计算与显存开销日益突出。Qwen、Kimi 等模型因此采用线性注意力与 softmax 注意力混合的架构。面向其中的 Gated Delta Network（GDN），FlashQLA 从并行度、显存访问和 GPU 流水线等方面重新设计高性能算子。

线性注意力通过调整矩阵乘法顺序，将复杂度从 O(N²d) 降至 O(Nd²)，但递推状态必须逐步传递，难以充分利用 GPU。GDN 在此基础上加入更新门 beta 和遗忘门 alpha，兼顾表达能力与效率。Qwen 使用三层 GDN 加一层 softmax attention 的“三比一”结构，Kimi K3 则采用三层 KDA 加一层 softmax attention，以弥补线性注意力长程信息承载能力有限的问题。

GDN 的常规 prefill 通过 WY representation 将逐 token 递推改为 chunk-by-chunk 处理，chunk size 通常为 64。流程包括计算每个 chunk 的 K Kᵀ 并求逆、生成 W 与 U、执行状态递推，以及依据状态 S 计算输出。除状态递推外，其余步骤可以并行，但现有实现会将 W、U、S 写入 HBM 后再读取，造成额外带宽消耗；每个 chunk 内 gate alpha 的累乘结果 gamma 也需参与计算。

原有方案主要受三方面限制：状态递推让每个 attention head 通常只能对应一个 CTA，Qwen 小模型可能只有 32 个头，而 H100 或 GB 系列 GPU 拥有上百个 SM，导致利用率不足并增加 TP 优化难度；中间状态往返 HBM 带来冗余访存；Triton 又难以精细控制融合算子和 warp specialization，无法编排理想的矩阵乘流水线。

FlashQLA 通过序列并行提升并行度。将序列切分后，后半段从错误初始状态递推得到的 S* 可与正确状态 S 做差，并利用差异按 chunk 传播的矩阵关系构造矫正矩阵 M。为控制显存，dK=128、chunk size=64 时不为每个 chunk 保存 128×128 矩阵，而是每个 CP rank 仅保留末端结果，再沿 CP size 递推完成校正，最后重新计算真正的 S 和输出。该方案以重复计算换取访存节省，S 会被计算两遍，且卡内、卡间序列并行遵循同一原理。

序列并行并不能直接支持完全融合。矩阵求逆成本较高，且计算 S* 和 M 需要使用求逆后的 W，因此 FlashQLA 将求逆单独保留，随后融合其余步骤，并把序列并行相关处理插入两者之间。针对卡内并行时切分数量不固定的问题，方案以每个 rank 内的 chunk 数 L 建模，第一、三步耗时近似与 L 成正比，第二步与 N/L 成正比，总时间可写为 T=λ₁L+λ₂N/L，通过平衡两项选择切分规模。

整体上，FlashQLA围绕 GDN 的代数结构进行硬件友好改写，结合序列维度分块、状态矫正、滑动窗口性质复用和 warp specialization，减少中间状态传输及显存带宽压力，并为 Qwen 等混合注意力模型提供更适合 GPU 的推理算子实现。