QwenLM/FlashQLA

它是什么

FlashQLA 是一个高性能内核库，实现了 Qwen 风格大型语言模型中使用的 Gated‑Delta‑Rule（GDN）线性注意力算法。它基于 TileLang DSL 编写，并以 Python 包（flash‑qla）的形式发布。该库既提供了便捷的高级函数（chunk_gated_delta_rule），也提供了可从 PyTorch 代码调用的低级前向/反向内核。

为什么重要

线性注意力机制将经典 softmax 注意力的二次方成本替换为随序列长度线性扩展的计算，使得超长上下文成为可能。与参考的 Flash Linear Attention（FLA）Triton 内核相比，FlashQLA 将前向传递速度提升了 2‑3 倍，反向传递速度提升了约 2 倍，尤其是在 NVIDIA Hopper/Blackwell GPU 上。更快的内核意味着更低的预训练成本和更低的推理延迟，这对于依赖 Qwen 风格模型的推理密集型或边缘侧代理尤为重要。

核心特性

特性	功能描述
门控驱动的自动卡内上下文并行	检测 GDN 门控的指数衰减是否允许将工作负载拆分到多个 CUDA warp/SM 上，自动在单个 GPU 卡内启用并行（适用于 TP、长序列、小头设置）。
硬件友好的代数重构	重写 GDN 前向/反向方程，以减少 Tensor‑Core、CUDA‑Core 和 SFU 的使用，同时保持数值精度。
TileLang 融合的 warp 专用内核	使用 TileLang DSL 生成少量融合内核，以平衡数据移动、Tensor‑Core 矩阵乘法和 CUDA‑Core 操作，而不是使用单个大型内核。
双 CP（卡间/卡内两级上下文并行） (v0.1.3)	增加了跨 GPU（卡间）和 GPU 内部（卡内）的第二级并行。
SM120 反向路径 (v0.1.3)	在反向内核中增加了对最新 NVIDIA SM120 架构的支持。
即插即用的 flash‑linear‑attention 后端 (v0.1.2)	可通过标准 FLA API 作为 FLA GDN 内核的即插即用替代品。

如何开始

# 安装预构建的 wheel（需要 CUDA 12.8+、PyTorch 2.8+）
pip install flash-qla

或者从源码构建：

git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .

简单使用示例

import torch
from flash_qla import chunk_gated_delta_rule

# q、k、v、g、beta 是形状如 README 中描述的张量
out, final_state = chunk_gated_delta_rule(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=scale,
    initial_state=init_state,
    output_final_state=True,
    cu_seqlens=cu_seqlens,
)

对于自定义训练循环，您可以调用单独的前向和反向函数（*_fwd / *_bwd），并将返回的缓存对象传回反向传播。

何时使用

您正在训练或微调使用 GDN 线性注意力的 Qwen 风格模型。
您的工作负载运行在最新的 NVIDIA GPU（H100、H200、Blackwell）上，并且需要极致速度来进行长序列预训练或低延迟推理。
您已经在使用 PyTorch，并希望使用一个与现有 flash-linear-attention API 兼容的即插即用内核。

限制 / 待解决问题

仅支持 SM 90‑SM 121 的 NVIDIA GPU 和 CUDA 12.8+。不支持 CPU 或 AMD。
该库专注于 GDN 分块预填充变体；其他线性注意力公式未涵盖。
反向传递加速依赖于双 CP 实现；在较旧的 GPU 上可能收益较小。
需要 TileLang 运行时；用户必须能够编译 TileLang 生成的内核。

社区与维护

该项目由 QwenLM 团队维护（参见引用列表）。最近的版本（v0.1.1‑v0.1.3）显示出积极的开发状态。
问题（Issues）和拉取请求（Pull Requests）可以在 GitHub 仓库中提交；README 中提供了基准测试脚本和性能分析工具的链接，供开发者使用。

以上所有信息均直接取自仓库的 README，未添加任何外部推测。