💥 Flash Linear Attention brings together hardware-efficient building blocks, training-ready layers, and components for modern sequence models, spanning linear attention, sparse attention, state space models, and hybrid LLM architectures. All implementations are platform-agnostic and verified on NVIDIA, AMD, and Intel hardware. Pull requests are welcome!
News
- [2026-11] 🔀 Add dynamic short convolution to
fla (paper).
- [2026-11] 🧱 Add HC, mHC, and iHC implementations to
fla (paper).
- [2026-10] 🧠 Add Gated Slot Attention 2 (GSA2) implementation to
fla (paper).
- [2026-10] 🛰️ Add Gated KalmaNet (GKA) implementation to
fla (paper) - a fading-memory layer that solves a test-time ridge regression at every token.
- [2026-10] 🌊 Add CyFA (Cyclic Flow Attention) implementation to
fla - a Linear RNN that partitions memory by relative time through learned cyclic transport.
- [2026-10] 🎉 Add Stick-Breaking Attention kernels and models to
fla (paper).
- [2026-07] 🐈 Add CAT (Compress and Attend Transformer) implementation to
fla (paper) - a meta-sequence mixer that unlocks test-time control of inference costs.
- [2026-07] 🧱 Add a Gluon backend for AttnRes.
- [2026-07] 🚀 Add FlashQLA backend for Gated DeltaNet.
- [2026-06] 🔭 Add Parallax implementation to
fla (paper).
- [2026-06] 🧮 Add Preconditioned Gated DeltaNet (PGDN) and Preconditioned KDA (PKDA) to
fla (paper) — curvature-aware preconditioning of the linear recurrence via an ATK preconditioner.
- [2026-06] 🧱 Add Wall attention implementation to
fla (blog).
- [2026-05] 🚪 Add Gated DeltaNet 2 (GDN-2) implementation to
fla (paper).
- [2026-05] 🦅 Add Raven implementation to
fla (repo).
- [2026-05] 🚀 Add YOCO (You Only Cache Once) implementation to
fla.
- [2026-05] ⚡ Add fused AttnRes support to
fla (paper).
- [2026-04] 🐍 Add Mamba3 implementation to
fla (paper).
- [2026-04] 🧱 Add MoBA (Mixture of Block Attention) implementation to
fla, with FlashMoBA backend support.
- [2026-04] 🧱 Add TileLang backend support for selected kernels.
- [2026-04] 🎯 Add GPT-OSS-style attention sink support to
fla’s attention kernels.
- [2026-03] 🚀 Add Context Parallel support for KDA and GDN, enabling efficient distributed training across sequence dimension.
Older news
- [2025-10] 🌘 Add Kimi Delta Attention (KDA) implementation to
fla (paper).
- [2025-09] 🌲 Add DeltaFormer implementation to
fla (paper).
- [2025-09] 🐻 Thrilled to announce that GDN has been integrated into Qwen3-Next. Check out their blog post for more info!
- [2025-08] 🌲 Add Log-Linear Attention implementation to
fla (paper).
- [2025-08] 🎓 Add MoM implementation to
fla (paper).
- [2025-07] 🐳 Add MLA implementation to
fla (paper).
- [2025-07] 🛣️ Add PaTH Attention implementation to
fla (paper).
- [2025-06] 🎉 Add MesaNet implementation to
fla (paper).
- [2025-06] 🐍 Add Comba implementation to
fla (paper).
- [2025-05] 🎉 Add Rodimus* implementation to
fla (paper).
- [2025-04] 🎉 Add DeltaProduct implementation to
fla (paper).
- [2025-04] 🎉 Add FoX implementation to
fla (paper).
- [2025-03]
We have changed the default initializer_range to the magic 🐳 0.006 The initializer_range was rolled back to the default value of 0.02. For actual training, we recommend trying both.
- [2025-02] 🐳 Add NSA implementations to
fla. See kernels here.
- [2025-01] 🔥 We are migrating to
torchtitan-based training framework. Check out the flame repo for more details.
- [2025-01] 🦅 Add RWKV7 implementations (both kernels and models) to
fla.
- [2024-12] Add
flash-bidirectional-attention to fla-org (repo).
- [2024-12] 🎉 Add Gated DeltaNet implementation to
fla (paper).
- [2024-12] 🚀
fla now officially supports kernels with variable-length inputs.
- [2024-11] The inputs are now switched from head-first to seq-first format.
- [2024-11] 💥
fla now provides a flexible way for training hybrid models.
- [2024-10] 🔥 Announcing
flame, a minimal and scalable framework for training fla models. Check out the details here.
- [2024-09]
fla now includes a fused linear and cross-entropy layer, significantly reducing memory usage during training.
- [2024-09] 🎉 Add GSA implementation to
fla (paper).
- [2024-05] 🎉 Add DeltaNet implementation to
fla (paper).
- [2024-05] 💥
fla v0.1: a variety of subquadratic kernels/layers/models integrated (RetNet/GLA/Mamba/HGRN/HGRN2/RWKV6, etc., see Models).
- [2023-12] 💥 Launch
fla, offering a collection of implementations for state-of-the-art linear attention models.
Models
Installation
torch lives in a backend extra ([cuda] / [rocm] / [xpu] / [npu] / [cpu]). CUDA is one command; other backends are two so torch (and the right triton flavor that torch pulls transitively) come from the PyTorch wheel index instead of PyPI:
# CUDA
pip install flash-linear-attention[cuda]
# ROCm
pip install --index-url https://download.pytorch.org/whl/rocm7.2 torch
pip install flash-linear-attention[rocm]
See INSTALL.md for the full backend table, XPU / NPU (Ascend) / CPU flows, source installs, and the --no-deps path for torch pre-release / triton-nightly.
Choose a backend extra to install the appropriate torch and triton dependencies. Bare pip install flash-linear-attention does not install them.
Usage
Token Mixing
We provide “token mixing” linear attention layers in fla.layers for you to use.
You can replace the standard multihead attention layer in your model with other linear attention layers.
Example usage is as follows:
>>> import torch
>>> from fla.layers import MultiScaleRetention
>>> batch_size, num_heads, seq_len, hidden_size = 32, 4, 2048, 1024
>>> device, dtype = 'cuda:0', torch.bfloat16
>>> 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)
>>> y.shape
torch.Size([32, 2048, 1024])
We provide the implementations of models that are compatible with 🤗 Transformers library.
Here’s an example of how to initialize a GLA model from the default configs in fla:
>>> from fla.models import GLAConfig
>>> from transformers import AutoModelForCausalLM
>>> config = GLAConfig()
>>> model = AutoModelForCausalLM.from_config(config)
Click to expand config and model structure
>>> config
GLAConfig {
"attn": null,
"attn_mode": "chunk",
"bos_token_id": 1,
"clamp_min": null,
"conv_size": 4,
"elementwise_affine": true,
"eos_token_id": 2,
"expand_k": 0.5,
"expand_v": 1,
"feature_map": null,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": true,
"hidden_act": "swish",
"hidden_ratio": 4,
"hidden_size": 2048,
"initializer_range": 0.02,
"intermediate_size": null,
"max_position_embeddings": 2048,
"model_type": "gla",
"norm_eps": 1e-06,
"num_heads": 4,
"num_hidden_layers": 24,
"num_kv_heads": null,
"tie_word_embeddings": false,
"transformers_version": "4.50.1",
"use_cache": true,
"use_gk": true,
"use_gv": false,
"use_output_gate": true,
"use_short_conv": false,
"vocab_size": 32000
}
>>> model
GLAForCausalLM(
(model): GLAModel(
(embeddings): Embedding(32000, 2048)
(layers): ModuleList(
(0-23): 24 x GLABlock(
(attn_norm): RMSNorm(2048, eps=1e-06)
(attn): GatedLinearAttention(
(q_proj): Linear(in_features=2048, out_features=1024, bias=False)
(k_proj): Linear(in_features=2048, out_features=1024, bias=False)
(v_proj): Linear(in_features=2048, out_features=2048, bias=False)
(g_proj): Linear(in_features=2048, out_features=2048, bias=False)
(gk_proj): Sequential(
(0): Linear(in_features=2048, out_features=16, bias=False)
(1): Linear(in_features=16, out_features=1024, bias=True)
)
(o_proj): Linear(in_features=2048, out_features=2048, bias=False)
(g_norm_swish_gate): FusedRMSNormGated(512, eps=1e-06, activation=swish)
)
(mlp_norm): RMSNorm(2048, eps=1e-06)
(mlp): GatedMLP(
(gate_proj): Linear(in_features=2048, out_features=5632, bias=False)
(up_proj): Linear(in_features=2048, out_features=5632, bias=False)
(down_proj): Linear(in_features=5632, out_features=2048, bias=False)
(swiglu_linear): SwiGLULinear()
)
)
)
(norm): RMSNorm(2048, eps=1e-06)
)
(lm_head): Linear(in_features=2048, out_features=32000, bias=False)
)
Fused Modules
We offer a collection of fused modules in fla.modules to facilitate faster training:
Rotary Embedding: rotary positional embeddings as adopted by the Llama architecture, a.k.a., Transformer++.Norm Layers:
RMSNorm, LayerNorm and GroupNormRMSNormLinear, LayerNormLinear and GroupNormLinear to reduce memory usage of intermediate tensors for improved memory efficiency.
Norm Layers with Gating: combine norm layers with element-wise sigmoid or swish gating, as used by RetNet/GLA.Cross Entropy: faster Triton implementation of cross entropy loss.Linear Cross Entropy: fused linear layer and cross entropy loss to avoid the materialization of large logits tensors. Also refer to implementations by mgmalek and Liger-Kernel.Linear KL Divergence: fused linear layer and KL divergence loss in a similar vein as CE loss.
You can set fuse_linear_cross_entropy in the model configuration to enable or disable the fused linear cross entropy loss.
This fused implementation is more memory-efficient but may reduce numerical precision. Due to this trade-off, it is disabled by default.
If you enable this feature and encounter training instability (e.g., loss divergence), we recommend disabling it to see if the issue is resolved.
Generation
After pretraining, the model can generate text with the 🤗 text generation APIs.
In the following, we give a generation example:
>>> import fla
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
>>> name = 'fla-hub/gla-1.3B-100B'
>>> tokenizer = AutoTokenizer.from_pretrained(name)
>>> model = AutoModelForCausalLM.from_pretrained(name).cuda()
>>> input_prompt = "Power goes with permanence. Impermanence is impotence. And rotation is castration."
>>> input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids.cuda()
>>> outputs = model.generate(input_ids, max_length=64)
>>> tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
We also provide a simple script here for benchmarking the generation speed.
Simply run it by:
$ python -m benchmarks.benchmark_generation \
--path 'fla-hub/gla-1.3B-100B' \
--repetition_penalty 2. \
--prompt="Hello everyone, I'm Songlin Yang"
Prompt:
Hello everyone, I'm Songlin Yang
Generated:
Hello everyone, I'm Songlin Yang.
I am a 20 year old girl from China who is currently studying in the United States of America for my Master degree and also working as an English teacher at school here on campus since last summer (1st semester). My main goal to be able do well with this course so that we can have
Prompt length: 10, generation length: 64
Total prompt processing + decoding time: 4593ms
All of the pretrained models currently available can be found in fla-hub.
>>> from huggingface_hub import list_models
>>> for model in list_models(author='fla-hub'): print(model.id)
Hybrid Models
fla provides a flexible method to incorporate standard attention layers into existing linear attention models.
This is easily achieved by specifying the attn argument in the model configuration.
The original dictionary form applies one shared attention specification to every listed layer.
For example, to create a 2-layer Samba model with one Mamba layer followed by one local attention layer, using a sliding window size of 2048:
>>> from fla.models import SambaConfig
>>> from transformers import AutoModelForCausalLM
>>> config = SambaConfig(num_hidden_layers=2)
>>> config.attn = {
'layers': [1],
'num_heads': 18,
'num_kv_heads': 18,
'qkv_bias': False,
'rope_theta': 10000.,
'window_size': 2048
}
>>> model = AutoModelForCausalLM.from_config(config)
Click to expand config and model structure
>>> config
SambaConfig {
"attn": {
"layers": [
1
],
"num_heads": 18,
"num_kv_heads": 18,
"qkv_bias": false,
"rope_theta": 10000.0,
"window_size": 2048
},
"bos_token_id": 1,
"conv_kernel": 4,
"eos_token_id": 2,
"expand": 2,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": true,
"hidden_act": "swish",
"hidden_ratio": 4,
"hidden_size": 2304,
"initializer_range": 0.02,
"intermediate_size": 4608,
"max_position_embeddings": 2048,
"model_type": "samba",
"norm_eps": 1e-05,
"num_hidden_layers": 2,
"pad_token_id": 0,
"rescale_prenorm_residual": false,
"residual_in_fp32": false,
"state_size": 16,
"tie_word_embeddings": false,
"time_step_floor": 0.0001,
"time_step_init_scheme": "random",
"time_step_max": 0.1,
"time_step_min": 0.001,
"time_step_rank": 144,
"time_step_scale": 1.0,
"transformers_version": "4.50.1",
"use_bias": false,
"use_cache": true,
"use_conv_bias": true,
"vocab_size": 32000
}
>>> model
SambaForCausalLM(
(backbone): SambaModel(
(embeddings): Embedding(32000, 2304)
(layers): ModuleList(
(0): SambaBlock(
(mixer_norm): RMSNorm(2304, eps=1e-05)
(mixer): Mamba(
(conv1d): Conv1d(4608, 4608, kernel_size=(4,), stride=(1,), padding=(3,), groups=4608)
(in_proj): Linear(in_features=2304, out_features=9216, bias=False)
(x_proj): Linear(in_features=4608, out_features=176, bias=False)
(dt_proj): Linear(in_features=144, out_features=4608, bias=True)
(out_proj): Linear(in_features=4608, out_features=2304, bias=False)
)
(mlp_norm): RMSNorm(2304, eps=1e-05)
(mlp): GatedMLP(
(gate_proj): Linear(in_features=2304, out_features=6144, bias=False)
(up_proj): Linear(in_features=2304, out_features=6144, bias=False)
(down_proj): Linear(in_features=6144, out_features=2304, bias=False)
(swiglu_linear): SwiGLULinear()
)
)
(1): SambaBlock(
(mixer_norm): RMSNorm(2304, eps=1e-05)
(mixer): Attention(
(q_proj): Linear(in_features=2304, out_features=2304, bias=False)
(k_proj): Linear(in_features=2304, out_features=2304, bias=False)
(v_proj): Linear(in_features=2304, out_features=2304, bias=False)
(o_proj): Linear(in_features=2304, out_features=2304, bias=False)
(rotary): RotaryEmbedding(dim=128, base=10000.0, interleaved=False, pos_idx_in_fp32=True)
)
(mlp_norm): RMSNorm(2304, eps=1e-05)
(mlp): GatedMLP(
(gate_proj): Linear(in_features=2304, out_features=6144, bias=False)
(up_proj): Linear(in_features=2304, out_features=6144, bias=False)
(down_proj): Linear(in_features=6144, out_features=2304, bias=False)
(swiglu_linear): SwiGLULinear()
)
)
)
(norm_f): RMSNorm(2304, eps=1e-05)
)
(lm_head): Linear(in_features=2304, out_features=32000, bias=False)
)
To use different attention settings at different depths, pass a list of specifications. For example, this six-layer Samba model uses local attention at layers 1 and 3, full attention at layer 5, and the native Mamba mixer at layers 0, 2, and 4:
>>> config = SambaConfig(
... num_hidden_layers=6,
... attn=[
... {
... 'layers': [1, 3],
... 'num_heads': 18,
... 'num_kv_heads': 18,
... 'qkv_bias': False,
... 'rope_theta': 10000.,
... 'window_size': 2048,
... },
... {
... 'layers': [5],
... 'num_heads': 18,
... 'num_kv_heads': 18,
... 'qkv_bias': False,
... 'rope_theta': 10000.,
... 'window_size': None,
... },
... ],
... )
>>> model = AutoModelForCausalLM.from_config(config)
Each specification is normalized independently. Layers omitted from the plan retain the model’s native linear-attention, recurrent, or state-space mixer.
During inference, you DO NOT need to revise anything for generation!
The model will produce output as-is, without any need for additional configurations or modifications.
Training
We provide a minimal framework called 🔥 flame built on top of torchtitan, for efficient training of fla models.
Check out the GLA example for more details.
Evaluation
We support benchmark evaluation through lm-evaluation-harness with python -m evals.harness, and perplexity evaluation with python -m evals.ppl.
See the evaluation guide for setup, single- and multi-GPU commands, RULER benchmarks, and perplexity evaluation settings.
Benchmarks
We compare our Triton-based implementations (chunk_retention, chunk_gla, chunk_gdn) with CUDA-based FlashAttention2 across various shape configurations.
These tests were conducted on a single NVIDIA GB200 GPU (CUDA 12.9, PyTorch 2.9.0).
# you might have to first install `fla` via `pip install -e .` to enable its import
$ python -m benchmarks.ops.run --op chunk_retention chunk_gla chunk_gdn flash_attn
=================================================================================
Machine: NVIDIA GB200 | CUDA 12.9 | PyTorch 2.9.0+cu129.msh
=================================================================================
fwd B T H D op main[0a484709](ms)
-----------------------------------------------------------------------
1 8192 96 128 chunk_retention 0.787
chunk_gla 1.765
chunk_gdn 1.265
flash_attn 3.753
-----------------------------------------------------------------------
2 16384 16 128 chunk_retention 0.792
chunk_gla 1.445
chunk_gdn 1.029
flash_attn 5.035
-----------------------------------------------------------------------
4 2048 16 128 chunk_retention 0.559
chunk_gla 0.514
chunk_gdn 0.753
flash_attn 0.346
-----------------------------------------------------------------------
4 4096 64 128 chunk_retention 0.997
chunk_gla 2.251
chunk_gdn 1.581
flash_attn 2.560
-----------------------------------------------------------------------
8 1024 8 64 chunk_retention 0.425
chunk_gla 0.358
chunk_gdn 0.631
flash_attn 0.157
-----------------------------------------------------------------------
8 2048 32 256 chunk_retention 1.174
chunk_gla 2.897
chunk_gdn 1.831
flash_attn 1.408
=================================================================================
fwdbwd B T H D op main[0a484709](ms)
-----------------------------------------------------------------------
1 8192 96 128 chunk_retention 2.618
chunk_gla 7.670
chunk_gdn 4.738
flash_attn 15.371
-----------------------------------------------------------------------
2 16384 16 128 chunk_retention 2.122
chunk_gla 5.984
chunk_gdn 3.616
flash_attn 19.960
-----------------------------------------------------------------------
4 2048 16 128 chunk_retention 1.047
chunk_gla 1.434
chunk_gdn 2.085
flash_attn 0.902
-----------------------------------------------------------------------
4 4096 64 128 chunk_retention 3.459
chunk_gla 10.216
chunk_gdn 5.964
flash_attn 10.815
-----------------------------------------------------------------------
8 1024 8 64 chunk_retention 0.898
chunk_gla 1.707
提示：代码过长已自动截断
Citation
If you find this repository helpful, please cite our work:
@software{yang2024fla,
title = {FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism},
author = {Yang, Songlin and Zhang, Yu},
url = {https://github.com/fla-org/flash-linear-attention},
month = jan,
year = {2024}
}
@misc{chen2026attnres,
title = {Attention Residuals},
author = {Chen, Guangyu and Zhang, Yu and Su, Jianlin and Xu, Weixin and Pan, Siyuan and Wang, Yaoyu and Wang, Yucheng and Chen, Guanduo and Yin, Bohong and Chen, Yutian and Yan, Junjie and Wei, Ming and Zhang, Y. and Meng, Fanqing and Hong, Chao and Xie, Xiaotong and Liu, Shaowei and Lu, Enzhe and Tai, Yunpeng and Chen, Yanru and Men, Xin and Guo, Haiqing and Charles, Y. and Lu, Haoyu and Sui, Lin and Zhu, Jinguo and Zhou, Zaida and He, Weiran and Huang, Weixiao and Xu, Xinran and Wang, Yuzhi and Lai, Guokun and Du, Yulun and Wu, Yuxin and Yang, Zhilin and Zhou, Xinyu},
year = {2026},
eprint = {2603.15031},
archiveprefix = {arXiv},
primaryclass = {cs.CL}
}
@misc{zhang2025kda,
title = {Kimi Linear: An Expressive, Efficient Attention Architecture},
author = {Zhang, Yu and Lin, Zongyu and Yao, Xingcheng and Hu, Jiaxi and Meng, Fanqing and Liu, Chengyin and Men, Xin and Yang, Songlin and Li, Zhiyuan and Li, Wentao and Lu, Enzhe and Liu, Weizhou and Chen, Yanru and Xu, Weixin and Yu, Longhui and Wang, Yejie and Fan, Yu and Zhong, Longguang and Yuan, Enming and Zhang, Dehao and Zhang, Yizhi and T. Liu, Y. and Wang, Haiming and Fang, Shengjun and He, Weiran and Liu, Shaowei and Li, Yiwei and Su, Jianlin and Qiu, Jiezhong and Pang, Bo and Yan, Junjie and Jiang, Zhejun and Huang, Weixiao and Yin, Bohong and You, Jiacheng and Wei, Chu and Wang, Zhengtao and Hong, Chao and Chen, Yutian and Chen, Guanduo and Wang, Yucheng and Zheng, Huabin and Wang, Feng and Liu, Yibo and Dong, Mengnan and Zhang, Zheng and Pan, Siyuan and Wu, Wenhao and Wu, Yuhao and Guan, Longyu and Tao, Jiawen and Fu, Guohong and Xu, Xinran and Wang, Yuzhi and Lai, Guokun and Wu, Yuxin and Zhou, Xinyu and Yang, Zhilin and Du, Yulun},
year = {2025},
eprint = {2510.26692},
archivePrefix = {arXiv},
primaryClass = {cs.CL}
}
@inproceedings{yang2025path,
title = {PaTH Attention: Position Encoding via Accumulating Householder Transformations},
author = {Yang, Songlin and Shen, Yikang and Wen, Kaiyue and Tan, Shawn and Mishra, Mayank and Ren, Liliang and Panda, Rameswar and Kim, Yoon},
booktitle = {Proceedings of NeurIPS},
year = {2025}
}
@inproceedings{yang2024gdn,
title = {Gated Delta Networks: Improving Mamba2 with Delta Rule},
author = {Yang, Songlin and Kautz, Jan and Hatamizadeh, Ali},
booktitle = {Proceedings of ICLR},
year = {2025}
}
@inproceedings{yang2024deltanet,
title = {Parallelizing Linear Transformers with the Delta Rule over Sequence Length},
author = {Yang, Songlin and Wang, Bailin and Zhang, Yu and Shen, Yikang and Kim, Yoon},
booktitle = {Proceedings of NeurIPS},
year = {2024}
}
@inproceedings{zhang2024gsa,
title = {Gated Slot Attention for Efficient Linear-Time Sequence Modeling},
author = {Zhang, Yu and Yang, Songlin and Zhu, Ruijie and Zhang, Yue and Cui, Leyang and Wang, Yiqiao and Wang, Bolun and Shi, Freda and Wang, Bailin and Bi, Wei and Zhou, Peng and Fu, Guohong},
booktitle = {Proceedings of NeurIPS},
year = {2024}
}
@inproceedings{qin2024hgrn2,
title = {HGRN2: Gated Linear RNNs with State Expansion},
author = {Qin, Zhen and Yang, Songlin and Sun, Weixuan and Shen, Xuyang and Li, Dong and Sun, Weigao and Zhong, Yiran},
booktitle = {Proceedings of COLM},
year = {2024}
}
@inproceedings{yang2024gla,
title = {Gated Linear Attention Transformers with Hardware-Efficient Training},
author = {Yang, Songlin and Wang, Bailin and Shen, Yikang and Panda, Rameswar and Kim, Yoon},
booktitle = {Proceedings of ICML},
year = {2024}
}
Acknowledgements
We extend our gratitude to Bitdeer and Moonshot AI for their support in maintaining and powering our project infrastructure.
评论