QwenLM/FlashQLA

概要

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

重要性

线性注意力机制用随序列长度线性扩展的计算取代了经典 soft‑max 注意力的二次成本，使得超长上下文成为可能。与参考的 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 兼容的即插即用内核。

限制 / 待解决问题

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

社区与维护

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

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