Qwen团队开源FlashQLA优化GDN线性注意力GPU性能

阿里巴巴 Qwen 团队算法工程师章程瑞东分享了面向 Gated Delta Network (GDN) 线性注意力的高性能 GPU 算子库 FlashQLA，该工具作为 FLA 官方后端已开源，旨在解决混合架构中线性注意力部分在长序列下的内存与计算开销问题。 FlashQLA 针对现有实现中的并行度不足、访存冗余与 SFU 负载过重三大痛点进行了底层重构。核心突破在于利用 GDN 的滑动窗口性质（约 60%-80% 的 head alpha 不恒为 1）简化了卡内序列并行方案，并设计了基于 TileLang 的手写 Warp Specialization 流水线。通过将数据搬运、Tensor Core 计算与 SFU 计算重叠，显著提升了 HBM 带宽利用率，使得前向推理获得 2-3 倍加速，反向传播获得约 2 倍加速。 该算子库目前支持 Hopper 及 Blackwell 家族（SM100-SM121），已在千问内部训练和推理链路接入。其接口与 FLA 兼容，新增 auto CP 开关以自动决策序列并行切分策略。对于依赖确定性推理或训推一致的场景，用户可手动关闭 auto CP。未来计划推广卡间序列并行以支持 PD 分离的超长序列推理。 事实：FlashQLA 是 FLA 官方后端，提供与 chunk_gated_delta_rule 一致的 API。 推断：在 TP 并行度较大或模型较小的场景下，auto CP 生效带来的加速效果尤为显著。 猜测：随着对算子性能建模的细化，one wave 简化模型将被更精确的方案取代，进一步优化序列并行拆分决策。