CONTRIBUTING

Contributing

Thank you for your interest in contributing to Flash Linear Attention! All pull requests are super welcomed and greatly appreciated.

* Report Bugs
* Ask Questions
* Core Principles
* Setup Development Environment
* Prerequisites
* Setup
* Lint Check
* Test Locally
* Project Structure
* Code Style
* Copyright Header
* Formatting and Linting
* Docstrings and Comments
* Prose and Markdown
* Naming Conventions
* Triton Kernels
* PyTorch Operators
* Adding a New Operator
* Adding a New Model
* Testing
* Running Tests
* Writing Tests
* NaN Memory Poisoning
* Benchmarking
* Submit Pull Requests
* Commit Message Convention
* PR Description
* CI Pipeline
* Review Checklist
* Environment Variables
* License

Report Bugs

If you run into any weird behavior while using fla, feel free to open a new issue! Please run a search before opening a new issue, to make sure that someone else hasn't already reported or solved the bug you've found.

Any issue you open should include:

- A minimal code snippet that reproduces the bug.
- A clear explanation of what the issue is.

Ask Questions

Please ask questions in issues or on Discord. Check FAQs.md first for common questions.

Core Principles

Read these before changing any kernel — they are the bar every PR is held to.

1. Match the reference numerically. Every optimized kernel must agree with its naive reference within assert_close tolerance. Pure refactors and other non-computational changes (rewrites, fused paths, autotune tweaks) must leave outputs and gradients unchanged — verify before vs. after, don't assume.
2. Find the root cause before patching. Don't land band-aid fixes. If a change appears to help but you can't explain why, keep digging.
3. Reuse over duplication. Check fla/ops/common/ and existing operators before writing new kernels; unify shared code paths instead of copying per-operator variants.
4. Audit every callsite when touching shared code. Renaming a symbol, changing a config field, or editing a common kernel/component means updating all of its uses in one pass — not one spot at a time. Changes in fla/ops/ or fla/modules/ ripple up to fla/layers/ and fla/models/: check those consumers and decide explicitly whether the public interface needs to change. See Triton Kernels for the kernel-level checklist.
5. Protect battle-tested paths; keep diffs minimal. Changes to converged kernels or public APIs can silently break user code or checkpoints. Change only what the fix or feature needs, plus light incidental cleanups — don't revert or rewrite working code just because it could be cleaner (note it as optional in review instead). Flag risky changes, and when in doubt, ask.

Setup Development Environment

Prerequisites

- Python >= 3.10
- PyTorch >= 2.7.0
- A GPU with Triton support (NVIDIA, AMD, or Intel)

Setup

1. Fork flash-linear-attention (fork) on GitHub and clone the repository.

bash

git clone [email protected]:<your username>/flash-linear-attention.git
    cd flash-linear-attention    git remote add upstream [email protected]:fla-org/flash-linear-attention.git

2. Install in development mode with a backend extra (cuda / rocm / xpu / npu / cpu):

bash

pip install -e '.[cuda,test]'

For non-CUDA backends, install the matching torch + triton flavor from the PyTorch index first (see INSTALL.md), then run the editable install with the matching extra (e.g. .[rocm,test]).

> [!TIP]
> If the install fails, double-check that your PyTorch version matches your local CUDA toolkit and that nvcc is available in your PATH.

3. Setup the pre-commit hooks:

bash

pip install pre-commit
    pre-commit install

Lint Check

To check the linting, run:

bash

pre-commit run --all-files

Test Locally

bash

pytest tests/

Project Structure

text

fla/ ├── layers/ # PyTorch attention layer implementations ├── ops/ # Triton kernel operators (the core of the project) │ ├── common/ # Shared kernels reused across operators │ └── <op_name>/ # Each operator in its own directory │ ├── __init__.py │ ├── naive.py # Reference implementation in pure PyTorch │ ├── chunk.py # Chunk-based implementation │ ├── parallel.py # Parallel Triton kernel implementation │ ├── fused_recurrent.py # Fused recurrent implementation │ └── README.md # (optional) Mathematical derivations ├── models/ # Full language model definitions (config + modeling) ├── modules/ # Utility modules (norms, feature maps, rotary, etc.) └── utils.py # Global utilities and decorators

tests/ ├── conftest.py # Pytest config with NaN memory poisoning ├── ops/ # Operator tests ├── layers/ # Layer tests ├── models/ # Model tests └── modules/ # Module tests

Code Style

Copyright Header

Every source file should begin with the following header:

python

Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li

#
This source code is licensed under the MIT license found in the

LICENSE file in the root directory of this source tree.

For a list of all contributors, visit:

  https://github.com/fla-org/flash-linear-attention/graphs/contributors

A CI workflow (check-header.yml) enforces this automatically.

Formatting and Linting

We use Ruff for linting and autopep8 for formatting. Pre-commit hooks run both automatically.

Key rules:
- Max line length: 127 characters
- Target Python version: 3.10+
- Import sorting: isort-compatible via Ruff (fla as first-party)
- Type hints: Use modern syntax (X | None instead of Optional[X], list[str] instead of List[str])
- Use TYPE_CHECKING for imports only needed at type-check time
- Multi-line calls: don't wrap a call that fits within the line limit. When a call does overflow, use a hanging indent with one keyword argument per line, not several per line.

Docstrings and Comments

Comments and docstrings are hints for other readers, not a chain of thought. Give the reader what the code cannot say for itself, in as few words as possible — correct, simple, and with no narration of your reasoning.

Use a two-line hanging format for Args: / Returns: entries: a name (type, Optional): header line, then the description and Default: on the next indented line(s).

python

Args:
    hidden_size (int, Optional):
        The hidden size of the input. Default: 2048.
    use_output_gate (bool, Optional):
        Whether to apply a gated RMSNorm on the attention output. Default: False.

Capitalize Optional (not optional), put the default as Default: <value> (not "Defaults to ..."), and wrap True / False / None in backticks. See fla/layers/gla.py::GatedLinearAttention for the canonical example.

Keep inline comments restrained, especially in Triton kernels: shape annotations (e.g. # [BL, BD]) plus at most a one-line "why" for genuinely non-obvious tricks. Avoid multi-line derivations and narration that just restates the next line — math derivations belong in the operator's README.md, the PR description, or a single pointer, not inline.

Put explanatory comments on their own line above the code they describe, not trailing it — write # why on the line above x = f(), not x = f() # why. Start the comment text with a lowercase letter (# guard against overflow, not # Guard against overflow), and wrap a multi-line comment at clause boundaries like other prose. Reserve inline trailing comments for terse shape / type annotations like # [BL, BD].

Comments and docstrings must not go stale: when a change makes one factually wrong — a renamed symbol, a changed default, a removed code path — update or delete it in the same commit. An outdated comment is worse than none. If you can't tell whether a comment is still true, keep it and say so in the PR description; don't delete a "why" comment you merely can't verify. Fixing stale content means fixing the words, not reformatting the surrounding comment or docstring style.

Beyond narration that restates the next line, these comment patterns are banned:

- Banner blocks (##### ... #####): section boundaries should be visible from the code structure; use a blank line.
- Commented-out code: delete it — git has the history. Exception: commented-out configurations deliberately kept as documented alternatives (e.g. known-good autotune configs for a future dtype) may stay if a one-line comment says why they are kept.
- Personal asides (# XY: remove this?): a name in a comment is not an owner. Convert it to a TODO with an anchor, or delete it.
- Anchorless TODOs: a TODO must name when it can be acted on — a link to a tracking issue (this repo or upstream), a version bound (TODO: drop once we require triton>=3.5), or an externally checkable event. This applies to TODOs in docstrings too; see fla/ops/utils/op.py::safe_dot for the pattern.

Never treated as excess comments: the license header required by scripts/check_header.py, a one-line attribution with a URL for adapted code, shape/dtype annotations, and NOTE: / WARNING: prefixes on a genuine "why" comment.

Prose and Markdown

Don't hard-wrap prose at an arbitrary short column — this covers Markdown files, Python docstrings (including Args: / Returns: descriptions), and comment paragraphs. Either keep a paragraph on a single line, or break only at sentence or clause boundaries (after a ., ,, ;, or —), never mid-clause. In Python files the 127-character limit still applies, so wrap a docstring or comment at a clause boundary before it reaches the limit. Format Markdown tables with aligned columns so the | separators line up; table rows are exempt from the line limit.

Naming Conventions

| Entity | Convention | Example |
| --------------- | ------------------ | ----------------------------------------- |
| Classes | PascalCase | GatedDeltaNet, LinearAttention |
| Functions | snake_case | chunk_delta_rule, fused_recurrent_gla |
| Constants | UPPER_SNAKE_CASE | FLA_CI_ENV, SUPPORTS_AUTOTUNE_CACHE |
| Private helpers | Leading underscore | _guarded_empty, _is_called_from_fla |

Triton Kernels

- Kernel functions use @triton.jit with do_not_specialize=['T'] for the sequence-length argument.
- Use tl.constexpr for compile-time constants (block sizes, flags like USE_INITIAL_STATE).
- Write block accesses as explicit offset vectors (offset + tl.arange) with
plain tl.load / tl.store. Masked loads must cover every dimension that
can overrun at any call site, with other= where masked lanes matter; assert
any divisibility you rely on. Do not use tl.make_block_ptr / tl.advance:
deprecated upstream and removed in triton main. (backends/triton_ascend/ is
exempt — triton-ascend still requires block pointers.)
- tl.make_tensor_descriptor (TMA) is an opt-in optimization for hot-path
tiles on Hopper and newer, not a default substitute for block access. It
requires 16-byte-aligned bases and stride multiples, a stride-1 innermost
dim, no transposed blocks, and a registered allocator for device-side
descriptors. Use it when a per-kernel benchmark in the PR shows it pays off,
e.g. behind a flag as in fla/ops/utils/solve_tril.py.
- Treat program IDs and grid-derived indices as potentially narrow integers.
Cast them to tl.int64 before multiplying by sizes, strides, or sequence
offsets. This is especially important for non-first grid dimensions on NVIDIA
and for non-NVIDIA backends where every grid dimension may be narrow.
- Keep all tensor address arithmetic in tl.int64: block bases, varlen offsets,
strides, sequence positions, and element offsets must not rely on int16 or
int32 overflow behavior.
- Gate autotune configs with autotune_cache_kwargs for cache support.
- Kernel naming: <op>_fwd_kernel_<suffix> / <op>_bwd_kernel_<suffix>.
- When renaming a symbol or adding/moving a parameter, sweep every site in one pass: the tensor, its b_ value, the p_ block pointer, comments, and — across the forward/backward kernels, host wrappers, and autograd Function — every signature, launch, return tuple, and save_for_backward/saved_tensors list. Keyword-argument order at each call site must match the parameter order in the signature.

PyTorch Operators

- Wrap public-facing ops with the @input_guard decorator to ensure tensor contiguity.
- Use @autocast_custom_fwd / @autocast_custom_bwd for mixed-precision support.
- Provide a reference (naive) implementation in naive.py for testing.

Adding a New Operator

When adding a new operator under fla/ops/<op_name>/:

1. Create the directory with an __init__.py that exports the public API.
2. Write a naive implementation (naive.py) in pure PyTorch. This serves as the ground-truth reference for testing.
3. Implement the optimized kernel(s) in chunk.py, parallel.py, and/or fused_recurrent.py.
4. Reuse shared kernels from fla/ops/common/ where possible (e.g., chunk_fwd_o, chunk_gated_delta_rule_fwd_h).
5. Add tests in tests/ops/test_<op_name>.py (see Testing below).
6. (Optional) Add a README.md with mathematical derivations.

Adding a New Model

Each model lives under fla/models/<model_name>/ with:

- configuration_<model_name>.py — Config class extending PretrainedConfig
- modeling_<model_name>.py — Model, PreTrainedModel, and ForCausalLM classes
- __init__.py — Auto-registration with transformers

Testing

Every change to fla/ops/ or fla/modules/ must add or update the matching test under tests/, and a new operator must ship with a naive reference to compare against. Correctness is checked by strict numerical comparison against that reference — forward outputs and gradients — so a change that lacks a test, or only checks the forward pass, is not complete.

Running Tests

bash

Run all tests

pytest tests/
Run a specific test file

pytest tests/ops/test_delta.py
Run a specific test

pytest tests/ops/test_delta.py::test_chunk -v

Writing Tests

Tests compare optimized (Triton) implementations against reference (naive/recurrent) implementations. Follow this pattern:

python

import pytest
import torch
from fla.ops.your_op import chunk_your_op, fused_recurrent_your_op
from fla.utils import assert_close, device, device_platform

@pytest.mark.parametrize(
    ('B', 'T', 'H', 'D', 'dtype'),
    [
        pytest.param(test, id="B{}-T{}-H{}-D{}".format(test))
        for test in [
            (1, 63, 1, 64, torch.float16),
            (2, 1000, 4, 128, torch.float16),
        ]
    ],
)
def test_chunk(B: int, T: int, H: int, D: int, dtype: torch.dtype):
    torch.manual_seed(42)
    q = torch.randn(B, T, H, D, dtype=dtype).to(device).requires_grad_(True)
    k = torch.randn(B, T, H, D, dtype=dtype).to(device).requires_grad_(True)
    v = torch.randn(B, T, H, D, dtype=dtype).to(device).requires_grad_(True)
    do = torch.rand_like(v)
    # Triton implementation
    tri = chunk_your_op(q.clone(), k.clone(), v.clone())
    (tri * do).sum().backward()
    tri_dq, tri_dk, tri_dv = q.grad, k.grad, v.grad
    q.grad = k.grad = v.grad = None
    # Reference implementation
    ref = fused_recurrent_your_op(q.clone(), k.clone(), v.clone())
    (ref * do).sum().backward()
    ref_dq, ref_dk, ref_dv = q.grad, k.grad, v.grad    assert_close('o', ref, tri, 0.006)
    assert_close('dq', ref_dq, tri_dq, 0.006)
    assert_close('dk', ref_dk, tri_dk, 0.006)
    assert_close('dv', ref_dv, tri_dv, 0.006)

Key guidelines:

- Always use torch.manual_seed(42) for reproducibility.
- Use assert_close from fla.utils for numerical comparison with relative tolerance.
- Test both forward and backward passes by computing gradients.
- Use device from fla.utils for device-agnostic tests.
- Parametrize with diverse shapes including non-power-of-2 sequence lengths (e.g., 63, 100, 2000).
- Skip unsupported platforms with @pytest.mark.skipif(device_platform == 'intel', ...) when needed.
- Include test IDs in parametrize for readable output.

Naming and structure. Name the file tests/ops/test_<op>.py, and name each test after the implementation entry point it exercises — test_chunk, test_fused_recurrent, test_parallel — mirroring the functions in fla/ops/<op>/. Distinguish a genuinely different code path with a short suffix (test_chunk_varlen, test_fused_recurrent_state_v_first). Prefer adding a new shape, dtype, or flag as a @parametrize case on an existing test rather than writing a new function; only add a new function when the path or purpose is clearly different — varlen vs. dense, a specific feature flag, or a separate entry point. See tests/ops/test_gla.py and tests/ops/test_gdn.py for the pattern.

NaN Memory Poisoning

The test suite (conftest.py) automatically replaces torch.empty with NaN-filled tensors for tests/ops/ and tests/modules/. This catches bugs where uninitialized memory is accidentally used. You don't need to do anything special — just be aware that your kernels must fully initialize all output tensors.

Benchmarking

Any change that can affect performance — a new or rewritten kernel in fla/ops/ or fla/modules/, an autotune or backend tweak — should come with before/after numbers in the PR, measured on the same hardware and workload. [Perf] PRs must include them.

Benchmark only against a green test gate. A kernel that runs faster but fails its tests/ops/test_<op>.py (forward, backward, and NaN-poisoned init) is not an improvement, so confirm correctness first — see Testing.

Op microbenchmark — times forward and forward+backward across a shape sweep, and compares against a git ref (it builds a throwaway worktree, so your working tree is untouched):

bash

python -m benchmarks.ops.run --op chunk_gla --base main   # one op vs. main
python -m benchmarks.ops.run --list                       # registered ops

New ops are registered in benchmarks/ops/registry.py.

Correctness-gated driver — runs the op's pytest as a frozen gate, then benchmarks, and refuses to report a speedup on a red gate. Use it as the per-iteration command when optimizing a kernel; the fla-optimization-loop agent skill drives the full loop:

bash

python -m benchmarks.ops.verify --op chunk_gla --base main

Model-level throughput and generation:

bash

python benchmarks/benchmark_training_throughput.py --name kda --batch_size 2 --seq_len 8192 [--varlen]
python benchmarks/benchmark_generation.py --name kda

For profiling (Nsight Compute, hot-instruction analysis), see the fla-nvidia-performance agent skill. Report throughput (tokens/s or iters/s) and, when relevant, peak memory, and flag any shape or backend that regressed and why.

Submit Pull Requests

Once your change is implemented, tested, and (if it touches performance) benchmarked, open a pull request against main.

NOTE

Please include tests with every pull request if applicable!

- Keep the scope focused: one PR should do one thing. If you have multiple unrelated changes, please split them into separate PRs.
- Use Draft PRs: feel free to open a draft early for design feedback or work-in-progress discussion.
- Read AGENTS.md and .agents/skills/fla-mr-readiness first: they cover the PR checklist, test-plan requirements, and benchmark evidence standards expected of every pull request.
- No busywork PRs: don't open standalone PRs for typos or isolated style tweaks; fold them into a related substantive change instead.

Commit Message Convention

Use a prefix tag in square brackets to categorize your change. Here are some common examples:

| Tag | Usage | Example |
| ------------ | -------------------------- | ------------------------------------------------- |
| [Fix] | Bug fixes | [Fix] Guard checkpoint weight re-initialization |
| [Misc] | Miscellaneous | [Misc] Upgrade minimum PyTorch requirement |
| [Docs] | Documentation | [Docs] Update CP README |
| [CI] | CI/CD changes | [CI] Fix skip-test check failing on fork PRs |
| [Test] | Test additions or fixes | [Test] Add varlen backward gradient checks |
| [Perf] | Performance optimizations | [Perf] Fuse gate multiplication in delta rule |
| [Refactor] | Code refactoring | [Refactor] Unify chunk kernel entry points |
| [Ops] | General operator changes | [Ops] Refactor common chunk reduction utilities |
| [Model] | Model architecture changes | [Model] Add RoPE scaling to GLA config |
| [Layer] | Layer-level changes | [Layer] Normalize initial state initialization |
| [Attn] | Attention-related changes | [Attn] Add sliding window attention support |
| [GDN] | Gated Delta Net | [GDN] Add fused gate kernel |
| [KDA] | Kimi Delta Attention | [KDA] Fix illegal memory access in backward |
| [CP] | Context Parallel | [CP] Enable KCP for DPLR |
| [Conv] | Convolution | [Conv] Fix int32 overflow in varlen conv kernel |
| [CE] | Cross Entropy | [CE] Add logit softcapping support |

If your change doesn't fit any of the above, [Misc]/[chore] is the safe default.

PR Description

Lead with what changed and why, at a high level — describe the behavior or capability, not a file-by-file walkthrough. Include:

- Summary: the change and its motivation, stated up front. Keep it concise; reviewers read the diff for details.
- Test plan: how you verified it (commands run, hardware used).
- Breaking changes (if any): list any API changes that are not backward compatible, and describe the migration path.

See recent PRs for examples.

CI Pipeline

When you submit a PR, the following checks run automatically:

- Linting — Ruff + autopep8 via pre-commit
- License header check — Ensures copyright headers are present
- GPU tests — On NVIDIA H100/A100/4090 and Intel B580 (when available)
- Benchmarks — Performance regression checks; results are posted automatically as a PR comment

Add [skip test] to your commit message to skip GPU tests for documentation-only changes. For [Perf] changes, include before/after numbers in the PR — see Benchmarking.

Review Checklist

Before submitting, please go through the following checklist:

- Code follows the project's style conventions.
- Copyright header is present on all new files.
- Changes to fla/ops/ or fla/modules/ add or update the matching test in tests/.
- Tests pass locally (pytest tests/ops/test_<your_op>.py).
- New operators include a naive reference implementation.
- Both forward and backward passes are tested.
- Gradient correctness is verified against a reference implementation.
- Pre-commit hooks pass (pre-commit run --files <your_files>).

Environment Variables

See ENVs.md for a full list.

License

By contributing, you agree that your contributions will be licensed under the MIT License.

---

INSTALL

Installation

fla ships as two PyPI packages: fla-core (kernels in fla/ops,
fla/modules, fla/utils) and flash-linear-attention (everything in
fla/layers + fla/models, plus fla-core as a dep). Both follow the same
backend-extras layout.

Pick a backend

torch lives in a backend extra, not in the base deps, so wheel metadata is
the same across backends. The triton flavor either ships in the extra
(cuda / cpu / npu) or comes transitively from torch when you source it
from the matching PyTorch wheel index (rocm / xpu).

| Backend | Extra | Wheel index | triton flavor |
| ------- | -------- | ------------------------------------------ | ------------------------------------------------ |
| CUDA | [cuda] | https://download.pytorch.org/whl/cu128 | triton (PyPI) |
| ROCm | [rocm] | https://download.pytorch.org/whl/rocm7.2 | pulled by torch (pytorch-triton-rocm / triton-rocm) |
| XPU | [xpu] | https://download.pytorch.org/whl/xpu | pulled by torch (pytorch-triton-xpu) |
| NPU | [npu] | https://triton-ascend.osinfra.cn/pypi/simple | triton-ascend |
| CPU | [cpu] | https://download.pytorch.org/whl/cpu | triton (PyPI, import-only) |

From PyPI

CUDA can use a single command since triton lives on PyPI:

pip install flash-linear-attention[cuda]

For ROCm / XPU / CPU, do it in two steps so torch (and the matching triton
flavor that torch pulls transitively) come from the PyTorch wheel index
instead of letting the resolver mix and match (pip docs are explicit that
there is no priority across configured indices). This mirrors the
AMD-recommended pattern:

ROCm

pip install --index-url https://download.pytorch.org/whl/rocm7.2 torch
pip install flash-linear-attention[rocm]
XPU

pip install --index-url https://download.pytorch.org/whl/xpu torch
pip install flash-linear-attention[xpu]
CPU

pip install --index-url https://download.pytorch.org/whl/cpu torch
pip install flash-linear-attention[cpu]

For nightly torch, swap whl/<backend> for whl/nightly/<backend> and add --pre.

Ascend NPU

NPUs use triton-ascend, not
upstream triton. Since triton is in backend extras (not base deps), the
old "install fla, then pip uninstall triton, then install triton-ascend"
dance is no longer needed.

1. install CANN 9.0.0 + source set_env.sh

2. install torch / torch_npu / triton-ascend, then fla with the npu extra

pip install torch==2.7.1 torch_npu==2.7.1 torchvision==0.22.1
pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple
pip install flash-linear-attention[npu]

The [npu] extra pins torch==2.7.1, torch_npu==2.7.1, torchvision==0.22.1,
and triton-ascend==3.2.1 (CANN 9.0.0 stack tested in CI). Install
triton-ascend with --extra-index-url as shown above.

From source

pip uninstall fla-core flash-linear-attention -y
CUDA

pip install -U "git+https://github.com/fla-org/flash-linear-attention#egg=flash-linear-attention[cuda]"
Non-CUDA: install backend torch + triton from the PyTorch index first

(see the per-backend block above), then run the same git+ install with the

matching extra ([rocm] / [xpu] / [npu] / [cpu]).

Or with submodules:

git submodule add https://github.com/fla-org/flash-linear-attention.git 3rdparty/flash-linear-attention
ln -s 3rdparty/flash-linear-attention/fla fla

Behavior change vs. pre-v0.5

Before v0.5, pip install flash-linear-attention resolved a CUDA-built
torch + triton from the default PyPI index even on ROCm / XPU / NPU
machines, which silently overlaid the wrong wheels. Now the base install
contains no torch / triton at all: you pick a backend extra. Bare
pip install flash-linear-attention no longer imports.

Notes

- Already have a working torch for your backend? pip install -e .[rocm]
(or the matching extra) leaves it alone because the torch>=2.7.0 pin is
satisfied.
- For AMD GPUs the [rocm] extra pulls pytorch-triton-rocm. For Intel GPUs
the [xpu] extra pulls pytorch-triton-xpu. See FAQs for
backend-specific issues.

Skipping the dep resolver

torch pre-release / triton-nightly setups can sidestep resolution
entirely:

pip install transformers einops
pip uninstall fla-core flash-linear-attention -y
pip install -U --no-deps git+https://github.com/fla-org/flash-linear-attention

---

README

[](https://huggingface.co/fla-hub) [](https://discord.gg/vDaJTmKNcS)

</div>

<p>
💥 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!
</p>

--------

* News
* Models
* Installation
* Usage
* Token Mixing
* Fused Modules
* Generation
* Hybrid Models
* Training
* Evaluation
* Benchmarks
* Citation
* Star History
* Acknowledgements

News

- [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 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.
- [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).

<details>
<summary>Older news</summary>

- [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&ast; 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.

</details>

Models

Installation

[](https://github.com/fla-org/flash-linear-attention/actions/workflows/nvidia-h100.yml)

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.

NOTE

Behavior change vs. pre-v0.5: bare pip install flash-linear-attention no longer pulls torch / triton. Pick a backend extra. This fixes ROCm / XPU / NPU users silently getting CUDA wheels.

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)

<details>
<summary>Click to expand config and model structure</summary>

text

/ Detailed source-code truncated for AI context efficiency. /

</details>

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 GroupNorm
* RMSNormLinear, 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.

IMPORTANT

You can control using fuse_linear_cross_entropy in the model configuration to enable/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

Upon successfully pretraining a model, it becomes accessible for generating text using 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)

<details>
<summary>Click to expand config and model structure</summary>

text

/ Detailed source-code truncated for AI context efficiency. /

</details>

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

The lm-evaluation-harness library allows you to easily perform (zero-shot) model evaluations.
Follow the steps below to use this library:

1. Install lm_eval following their instructions.

2. Run evaluation with:

$ MODEL='fla-hub/gla-1.3B-100B'
$ python -m evals.harness --model hf \
    --model_args pretrained=$MODEL,dtype=bfloat16 \
    --tasks wikitext,lambada_openai,piqa,hellaswag,winogrande,arc_easy,arc_challenge,boolq,sciq,copa,openbookqa \
    --batch_size 64 \
    --num_fewshot 0 \
    --device cuda \
    --show_config

We've made fla compatible with hf-style evaluations, you can call evals.harness to finish the evaluations.
Running the command above will provide the task results reported in the GLA paper.

3. Multi-GPU Evaluation with Hugging Face accelerate 🚀

To perform data-parallel evaluation (where each GPU loads a separate full copy of the model), we leverage the accelerate launcher as follows:

$ MODEL='fla-hub/gla-1.3B-100B'
$ accelerate launch -m evals.harness --model hf  \
    --model_args pretrained=$MODEL,dtype=bfloat16,trust_remote_code=True  \
    --tasks wikitext,lambada_openai,piqa,hellaswag,winogrande,arc_easy,arc_challenge,boolq,sciq,copa,openbookqa \
    --batch_size 64  \
    --num_fewshot 0  \
    --device cuda  \
    --show_config  \
    --trust_remote_code

4. 📏 RULER Benchmark suite

The RULER benchmarks are commonly used for evaluating model performance on long-context tasks.
You can evaluate fla models on RULER directly using lm-evaluation-harness. RULER is only available in a relatively recent version of lm-evaluation-harness, so make sure you have the latest version installed.

text

git clone --depth 1 https://github.com/EleutherAI/lm-evaluation-harness
cd lm-evaluation-harness
pip install -e .

Then, install the necessary dependencies for RULER:

pip install lm_eval["ruler"]

and run evaluation by (e.g., 32k contexts):

$ accelerate launch -m evals.harness \
    --output_path $OUTPUT \
    --tasks niah_single_1,niah_single_2,niah_single_3,niah_multikey_1,niah_multikey_2,niah_multikey_3,niah_multiquery,niah_multivalue,ruler_vt,ruler_cwe,ruler_fwe,ruler_qa_hotpot,ruler_qa_squad \
    --model_args pretrained=$MODEL,dtype=bfloat16,max_length=32768,trust_remote_code=True \
    --metadata='{"max_seq_lengths":[4096,8192,16384,32768]}' \
    --batch_size 2 \
    --show_config  \
    --trust_remote_code

If a GPU can't load a full copy of the model, please refer to this link for FSDP settings.

TIP

If you are using lm-evaluation-harness as an external library and can't find (almost) any tasks available, before calling lm_eval.evaluate() or lm_eval.simple_evaluate(), simply run the following to load the library's stock tasks:

``py


>>> from lm_eval.tasks import TaskManager; TaskManager().initialize_tasks()

`

`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).

text

/ Detailed source-code truncated for AI context efficiency. /

Citation

If you find this repository helpful, please cite our work:

text

/ Detailed source-code truncated for AI context efficiency. /

Star History

[](https://github.com/fla-org/flash-linear-attention/stargazers)

[](https://star-history.com/#fla-org/flash-linear-attention&Date)

Acknowledgements

We extend our gratitude to Bitdeer and Moonshot AI for their support in maintaining and powering our project infrastructure.

---

CONTRIBUTING Copy Section

Contributing

Table of Contents Copy Chapter

Report Bugs Copy Chapter

Ask Questions Copy Chapter

Core Principles Copy Chapter

Setup Development Environment Copy Chapter

Prerequisites Copy Section

Setup Copy Section

Lint Check Copy Section

Test Locally Copy Section

Project Structure Copy Chapter

Code Style Copy Chapter

Copyright Header Copy Section

Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li

This source code is licensed under the MIT license found in the

LICENSE file in the root directory of this source tree.

For a list of all contributors, visit:

https://github.com/fla-org/flash-linear-attention/graphs/contributors

Formatting and Linting Copy Section

Docstrings and Comments Copy Section

Prose and Markdown Copy Section

Naming Conventions Copy Section

Triton Kernels Copy Section

PyTorch Operators Copy Section

Adding a New Operator Copy Chapter

Adding a New Model Copy Chapter

Testing Copy Chapter

Running Tests Copy Section

Run all tests

Run a specific test file

Run a specific test

Writing Tests Copy Section

NaN Memory Poisoning Copy Section

Benchmarking Copy Chapter

Submit Pull Requests Copy Chapter

Commit Message Convention Copy Section

PR Description Copy Section

CI Pipeline Copy Section

Review Checklist Copy Section

Environment Variables Copy Chapter

License Copy Chapter

INSTALL Copy Section

Installation

Pick a backend Copy Chapter

From PyPI Copy Chapter

ROCm

XPU

CPU

Ascend NPU Copy Chapter

1. install CANN 9.0.0 + source set_env.sh

2. install torch / torch_npu / triton-ascend, then fla with the npu extra

From source Copy Chapter

CUDA

Non-CUDA: install backend torch + triton from the PyTorch index first

(see the per-backend block above), then run the same git+ install with the

matching extra ([rocm] / [xpu] / [npu] / [cpu]).

Behavior change vs. pre-v0.5 Copy Chapter

Notes Copy Chapter

Skipping the dep resolver Copy Chapter

README Copy Section

News Copy Chapter

Models Copy Chapter

Installation Copy Chapter

CUDA

ROCm

Usage Copy Chapter

Token Mixing Copy Section

Fused Modules Copy Section

Generation Copy Section

Hybrid Models Copy Section

Training Copy Chapter

Evaluation Copy Chapter

Benchmarks Copy Chapter

Citation Copy Chapter

Star History Copy Chapter

Acknowledgements Copy Chapter

Explore Related Technologies & Topics

CONTRIBUTING

Table of Contents

Report Bugs

Ask Questions

Core Principles

Setup Development Environment

Prerequisites

Setup

Lint Check

Test Locally

Project Structure

Code Style

Copyright Header

Formatting and Linting

Docstrings and Comments

Prose and Markdown

Naming Conventions

Triton Kernels

PyTorch Operators

Adding a New Operator

Adding a New Model

Testing

Running Tests

Writing Tests

NaN Memory Poisoning

Benchmarking

Submit Pull Requests

Commit Message Convention

PR Description

CI Pipeline

Review Checklist

Environment Variables

License

INSTALL

Pick a backend

From PyPI

Ascend NPU

From source

Behavior change vs. pre-v0.5

Notes

Skipping the dep resolver

README

News

Models

Installation

Usage

Token Mixing

Fused Modules

Generation

Hybrid Models

Training

Evaluation

`Benchmarks`

Citation

Star History

Acknowledgements