fla-org/flash-linear-attention · 上手攻略

仓库：fla-org/flash-linear-attention
链接：https://github.com/fla-org/flash-linear-attention
分类：ai
作者：Jay
更新：2026-07-13

是什么

Flash Linear Attention（简称 fla）是一个专注于高效线性注意力（Linear Attention）机制实现的开源库，底层基于 Triton 编写，在 NVIDIA / AMD / Intel 硬件上均做过硬件对齐验证。

它的核心价值在于：把学术界和工业界近年来提出的各种次二次复杂度（Subquadratic）序列建模架构——线性注意力、状态空间模型（SSM）、稀疏注意力、混合架构等——以生产级质量的 Triton kernel 和 PyTorch 层接口统一实现，降低研究者和工程师的使用门槛。

fla 已被集成到 Qwen3-Next（GDN）等生产模型中，有实际落地验证。

解决什么问题

标准 Transformer 的注意力机制对序列长度是 O(n²) 复杂度，在长上下文场景（长文本、语音、视频帧序列）下面临内存和计算瓶颈。fla 提供了多种替代方案：

不换框架换层：可以用 fla.layers 里的线性注意力层直接替换标准 MultiHeadAttention，代码侵入极小
不用自己写 Triton：Triton kernel 开发门槛高，fla 已经写好了各种硬件高效的 kernel，直接调用即可
多架构统一实现：RetNet、GLA、Mamba2、RWKV6/7、Gated DeltaNet 等原本分散在各个仓库的架构，在 fla 里统一维护
跨硬件兼容：同一套代码在 NVIDIA / AMD / Intel GPU 上均可运行

快速安装

⚠️ v0.5 重要变化：pip install flash-linear-attention 不再附带 torch/triton，需手动选择 backend。

# CUDA（最常用，一行搞定）
pip install flash-linear-attention[cuda]

# ROCm（AMD GPU）
pip install --index-url https://download.pytorch.org/whl/rocm7.2 torch
pip install flash-linear-attention[rocm]

完整 backend 选择参考 INSTALL.md，包括： - [xpu] — Intel GPU（XPU 后端） - [npu] — 华为昇腾 NPU - [cpu] — 纯 CPU（无需 GPU）

从源码安装（如需修改或用最新未发布版本）：

git clone https://github.com/fla-org/flash-linear-attention.git
cd flash-linear-attention
pip install -e .

核心用法

1. Token Mixing 层（直接替换 MultiHeadAttention）

fla.layers 提供多种 token mixing 层，可直接替换标准 attention：

import torch
from fla.layers import MultiScaleRetention  # RetNet 的多尺度 retention 层

batch_size, num_heads, seq_len, hidden_size = 32, 4, 2048, 1024
device, dtype = 'cuda:0', torch.bfloat16

# 替换标准 attention 的用法
retnet = MultiScaleRetention(
    hidden_size=hidden_size,
    num_heads=num_heads
).to(device=device, dtype=dtype)

x = torch.randn(batch_size, seq_len, hidden_size).to(device=device, dtype=dtype)
y, *_ = retnet(x)
print(y.shape)  # torch.Size([32, 2048, 1024])

fla.layers 支持的层包括：

层名	对应模型/论文
`MultiScaleRetention`	RetNet (MSR)
`GatedLinearAttention`	GLA (Gated Linear Attention)
`Based`	Based Linear Attention
`DeltaNet`	DeltaNet
`HGRN` / `HGRN2`	Hierarchically Gated RNN
`RWKV6` / `RWKV7`	RWKV-6/7 Eagle & Finch
`DeltaProduct`	DeltaProduct
`MLA`	DeepSeek-V2 Multi-head Latent Attention

2. 完整模型（兼容 HuggingFace Transformers）

用 HuggingFace Transformers 接口加载 fla 内置的模型：

from fla.models import GLAConfig
from transformers import AutoModelForCausalLM

# 从默认配置初始化 GLA 模型
config = GLAConfig()
model = AutoModelForCausalLM.from_config(config)

# 打印模型结构
print(model)

fla.models 支持的完整模型：GLA、Based、RWKV6/7、Mamba2、Samba、YOCO、GSA 等。

3. 用预训练模型推理

从 HuggingFace fla-hub 下载预训练权重：

from transformers import AutoModelForCausalLM, AutoTokenizer

# RWKV7 模型示例（需 transformers>=4.48.0）
model = AutoModelForCausalLM.from_pretrained(
    'fla-hub/rwkv7-2.9B-world',
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
    'fla-hub/rwkv7-2.9B-world',
    trust_remote_code=True
)

text = "The future of AI is"
inputs = tokenizer([text], return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

4. Triton 原语（算子）调用

fla/ops/ 下提供底层 Triton kernel，可单独调用：

from fla.ops.gated_delta_rule import chunk_flash_gated_deltanet

# 直接调用底层算子做序列建模
# 具体 API 参考各 op 目录下的 README

5. 训练框架（flame）

fla 配套的轻量训练框架 flame（基于 torchtitan）：

git clone https://github.com/fla-org/flame.git
cd flame
# 参考 flame 仓库的 README 进行分布式训练配置

典型适用场景

场景	说明
长上下文语言模型	用 GLA / RWKV7 / Mamba2 等替代标准 Transformer，降低 O(n²) 内存开销
高效序列建模研究	在新论文复现中直接调用 `fla.layers`，无需从零写 Triton kernel
硬件高效推理	Triton 写死的 kernel 在 NVIDIA/AMD 上接近硬件上限
多硬件部署	同一套代码在 NVIDIA / AMD / Intel GPU 上均可运行
状态空间模型研究	Mamba2 / Mamba3 / RWKV 系列统一实现，方便对比

坑与注意

CUDA 版本要求：使用 CUDA backend 需要 PyTorch 与 CUDA 版本匹配；建议通过 pip install flash-linear-attention[cuda] 自动拉取兼容的 torch。
v0.5 breaking change：老用户注意 bare pip install flash-linear-attention 不再附带 torch，必须显式选择 [cuda] / [rocm] / [xpu] / [npu] / [cpu] 其中之一。
Triton 依赖严格：部分 kernel 需要特定版本的 Triton，pip install -e . 从源码安装时确保 triton 版本正确。
生产部署需验证：虽然 GDN 已用于 Qwen3-Next，但部分新加入的架构（如 2026 年的 Wall Attention、Parallax）尚未经过大规模生产验证，使用前建议自行做精度对齐测试。
RWKV7 需要 trust_remote_code：从 fla-hub 加载 RWKV7 模型必须 trust_remote_code=True，且 transformers 版本需 >= 4.48.0（建议同时升级 transformers）。
混合模型训练文档偏少：虽然 fla 支持 hybrid 模型（线性注意力 + 标准 attention 混合），但训练相关文档较少，有需求建议参考 flame 仓库的 examples。

与同类对比

特性	`fla` (Flash Linear Attention)	`mamba` (Mamba)	`flash-attn`	`transformers` 内置
架构覆盖	30+ 种（RetNet/GLA/Mamba2/RWKV7…）	主要 Mamba 系列	仅标准 attention	标准 Transformer
Triton kernel	✅ 全部自写	第三方	✅ 自写	❌
多硬件支持	NVIDIA/AMD/Intel	NVIDIA 为主	NVIDIA 为主	通用
HuggingFace 兼容	✅ 模型层全支持	有限	❌ 底层算子	✅
训练框架	flame（torchtitan）	有独立包	❌	PyTorch 原生
生产落地	Qwen3-Next GDN	大量生产	广泛使用	N/A

一句话总结：fla 是目前覆盖最广的线性注意力 + SSM 开源实现库，30+ 架构统一维护，HuggingFace 无缝对接，适合研究探索和生产落地双场景。

fla-org/flash-linear-attention · 上手攻略

fla-org/flash-linear-attention · 上手攻略

是什么

解决什么问题

快速安装

核心用法

1. Token Mixing 层（直接替换 MultiHeadAttention）

2. 完整模型（兼容 HuggingFace Transformers）

3. 用预训练模型推理

4. Triton 原语（算子）调用

5. 训练框架（flame）

典型适用场景

坑与注意

与同类对比

推荐结论

相关 · 按语义匹配