### CONTRIBUTING # Contributing Thank you for your interest in contributing to Flash Linear Attention! All pull requests are super welcomed and greatly appreciated. ## Table of Contents * [Report Bugs](#report-bugs) * [Ask Questions](#ask-questions) * [Core Principles](#core-principles) * [Setup Development Environment](#setup-development-environment) * [Prerequisites](#prerequisites) * [Setup](#setup) * [Lint Check](#lint-check) * [Test Locally](#test-locally) * [Project Structure](#project-structure) * [Code Style](#code-style) * [Copyright Header](#copyright-header) * [Formatting and Linting](#formatting-and-linting) * [Docstrings and Comments](#docstrings-and-comments) * [Prose and Markdown](#prose-and-markdown) * [Naming Conventions](#naming-conventions) * [Triton Kernels](#triton-kernels) * [PyTorch Operators](#pytorch-operators) * [Adding a New Operator](#adding-a-new-operator) * [Adding a New Model](#adding-a-new-model) * [Testing](#testing) * [Running Tests](#running-tests) * [Writing Tests](#writing-tests) * [NaN Memory Poisoning](#nan-memory-poisoning) * [Benchmarking](#benchmarking) * [Submit Pull Requests](#submit-pull-requests) * [Commit Message Convention](#commit-message-convention) * [PR Description](#pr-description) * [CI Pipeline](#ci-pipeline) * [Review Checklist](#review-checklist) * [Environment Variables](#environment-variables) * [License](#license) ## Report Bugs If you run into any weird behavior while using `fla`, feel free to open a new [issue](https://github.com/fla-org/flash-linear-attention/issues)! 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](https://github.com/fla-org/flash-linear-attention/issues) or on [Discord](https://discord.gg/vDaJTmKNcS). Check [FAQs.md](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](#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](https://github.com/fla-org/flash-linear-attention/fork)) on GitHub and clone the repository. ```bash git clone git@github.com:/flash-linear-attention.git cd flash-linear-attention git remote add upstream git@github.com: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](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`](https://pre-commit.com) 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 ``` fla/ ├── layers/ # PyTorch attention layer implementations ├── ops/ # Triton kernel operators (the core of the project) │ ├── common/ # Shared kernels reused across operators │ └── / # 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](https://docs.astral.sh/ruff/) for linting and [autopep8](https://github.com/hhatto/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: ` (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: `_fwd_kernel_` / `_bwd_kernel_`. - 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//`: 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_.py` (see [Testing](#testing) below). 6. **(Optional)** Add a `README.md` with mathematical derivations. ## Adding a New Model Each model lives under `fla/models//` with: - `configuration_.py` — Config class extending `PretrainedConfig` - `modeling_.py` — Model, PreTrainedModel, and ForCausalLM classes - `__init__.py` — Auto-registration with `transformers` Register your model in `fla/models/__init__.py` for auto-discovery. ## 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_.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//`. 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_.py` (forward, backward, and NaN-poisoned init) is not an improvement, so confirm correctness first — see [Testing](#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](https://github.com/fla-org/flash-linear-attention/pulls?q=is%3Apr+is%3Amerged) 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](#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_.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 `). ## Environment Variables See [ENVs.md](ENVs.md) for a full list. ## License By contributing, you agree that your contributions will be licensed under the [MIT License](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: ```sh 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](https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/docs/install/installrad/native_linux/install-triton.html): ```sh # 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/` for `whl/nightly/` and add `--pre`. ## Ascend NPU NPUs use [`triton-ascend`](https://github.com/triton-lang/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. ```sh # 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 ```sh 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: ```sh 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](FAQs.md) for backend-specific issues. ## Skipping the dep resolver `torch` pre-release / `triton-nightly` setups can sidestep resolution entirely: ```sh 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
Flash Linear Attention
[](https://huggingface.co/fla-hub) [](https://discord.gg/vDaJTmKNcS)

💥 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](#news) * [Models](#models) * [Installation](#installation) * [Usage](#usage) * [Token Mixing](#token-mixing) * [Fused Modules](#fused-modules) * [Generation](#generation) * [Hybrid Models](#hybrid-models) * [Training](#training) * [Evaluation](#evaluation) * [Benchmarks](#benchmarks) * [Citation](#citation) * [Star History](#star-history) * [Acknowledgements](#acknowledgements) ## News - [2026-07] 🧱 Add a [Gluon](https://triton-lang.org/main/getting-started/tutorials/gluon/) backend for [AttnRes](fla/ops/attnres). - [2026-07] 🚀 Add [FlashQLA](https://github.com/QwenLM/FlashQLA) backend for [Gated DeltaNet](fla/ops/gated_delta_rule). - [2026-06] 🔭 Add Parallax implementation to `fla` ([paper](https://arxiv.org/abs/2605.29157)). - [2026-06] 🧱 Add Wall attention implementation to `fla` ([blog](https://blog.tilderesearch.com/blog/wall-attn)). - [2026-05] 🚪 Add Gated DeltaNet 2 (GDN-2) implementation to `fla` ([paper](https://arxiv.org/abs/2605.22791)). - [2026-05] 🦅 Add Raven implementation to `fla` ([repo](https://github.com/goombalab/raven)). - [2026-05] 🚀 Add [YOCO](https://arxiv.org/abs/2405.05254) (You Only Cache Once) implementation to `fla`. - [2026-05] ⚡ Add fused [AttnRes](fla/ops/attnres) support to `fla` ([paper](https://arxiv.org/abs/2603.15031)). - [2026-04] 🐍 Add Mamba3 implementation to `fla` ([paper](https://arxiv.org/abs/2603.15569)). - [2026-04] 🧱 Add [MoBA](https://arxiv.org/abs/2502.13189) (Mixture of Block Attention) implementation to `fla`, with [FlashMoBA](https://github.com/mit-han-lab/flash-moba) backend support. - [2026-04] 🧱 Add [TileLang](https://github.com/tile-ai/tilelang) backend support for selected kernels. - [2026-04] 🎯 Add [GPT-OSS](https://openai.com/index/introducing-gpt-oss/)-style attention sink support to `fla`'s attention kernels. - [2026-03] 🚀 Add [Context Parallel](fla/ops/cp/README.md) support for KDA and GDN, enabling efficient distributed training across sequence dimension. - [2025-10] 🌘 Add Kimi Delta Attention (KDA) implementation to `fla` ([paper](https://arxiv.org/abs/2510.26692)). - [2025-09] 🌲 Add DeltaFormer implementation to `fla` ([paper](https://arxiv.org/abs/2505.19488v1)). - [2025-09] 🐻 Thrilled to announce that [GDN](fla/ops/gated_delta_rule) has been integrated into Qwen3-Next. Check out their [blog post](https://qwen.ai/blog?id=4074cca80393150c248e508aa62983f9cb7d27cd&from=research.latest-advancements-list) for more info! - [2025-08] 🌲 Add Log-Linear Attention implementation to `fla` ([paper](https://arxiv.org/abs/2506.04761)). - [2025-08] 🎓 Add MoM implementation to `fla` ([paper](https://arxiv.org/abs/2502.13685)).
Older news - [2025-07] 🐳 Add MLA implementation to `fla` ([paper](https://arxiv.org/abs/2405.04434)). - [2025-07] 🛣️ Add PaTH Attention implementation to `fla` ([paper](https://arxiv.org/abs/2505.16381)). - [2025-06] 🎉 Add MesaNet implementation to `fla` ([paper](https://arxiv.org/abs/2506.05233)). - [2025-06] 🐍 Add Comba implementation to `fla` ([paper](https://arxiv.org/abs/2506.02475)). - [2025-05] 🎉 Add Rodimus* implementation to `fla` ([paper](https://arxiv.org/abs/2410.06577)). - [2025-04] 🎉 Add DeltaProduct implementation to `fla` ([paper](https://arxiv.org/abs/2502.10297)). - [2025-04] 🎉 Add FoX implementation to `fla` ([paper](https://arxiv.org/abs/2503.02130)). - [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](fla/ops/nsa). - [2025-01] 🔥 We are migrating to `torchtitan`-based training framework. Check out the [flame](https://github.com/fla-org/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](https://github.com/fla-org/flash-bidirectional-linear-attention)). - [2024-12] 🎉 Add Gated DeltaNet implementation to `fla` ([paper](https://arxiv.org/abs/2412.06464)). - [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](https://github.com/fla-org/flame). - [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](https://arxiv.org/abs/2409.07146)). - [2024-05] 🎉 Add DeltaNet implementation to `fla` ([paper](https://arxiv.org/abs/2102.11174)). - [2024-05] 💥 `fla` v0.1: a variety of subquadratic kernels/layers/models integrated (RetNet/GLA/Mamba/HGRN/HGRN2/RWKV6, etc., see [Models](#models)). - [2023-12] 💥 Launch `fla`, offering a collection of implementations for state-of-the-art linear attention models.
## Models | Year | Model | Paper | | | :---: | :------------------: | :-------------------------------------------------------------------------------------------------------------------------------------------- | :----------------------------------------------------------------------------------------------------- | | 2022 | ABC | [ABC: Attention with Bounded-memory Control](https://arxiv.org/abs/2110.02488) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/abc.py) | | 2023 | RetNet | [Retentive network: a successor to transformer for large language models](https://arxiv.org/abs/2307.08621) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/multiscale_retention.py) | | 2023 | HGRN | [Hierarchically Gated Recurrent Neural Network for Sequence Modeling](https://openreview.net/forum?id=P1TCHxJwLB) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/hgrn.py) | | 2024 | GLA | [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gla.py) | | 2024 | Based | [Simple linear attention language models balance the recall-throughput tradeoff](https://arxiv.org/abs/2402.18668) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/based.py) | | 2024 | Rebased | [Linear Transformers with Learnable Kernel Functions are Better In-Context Models](https://arxiv.org/abs/2402.10644) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/rebased.py) | | 2024 | DeltaNet | [Parallelizing Linear Transformers with Delta Rule over Sequence Length](https://arxiv.org/abs/2406.06484) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/delta_net.py) | | 2024 | HGRN2 | [HGRN2: Gated Linear RNNs with State Expansion](https://arxiv.org/abs/2404.07904) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/hgrn2.py) | | 2024 | RWKV6 | [Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence](https://arxiv.org/abs/2404.05892) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/rwkv6.py) | | 2024 | LightNet | [You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet](https://arxiv.org/abs/2405.21022) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/lightnet.py) | | 2024 | YOCO | [You Only Cache Once: Decoder-Decoder Architectures for Language Models](https://arxiv.org/abs/2405.05254) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/models/yoco) | | 2024 | Mamba2 | [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://arxiv.org/abs/2405.21060) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/models/mamba2) | | 2024 | GSA | [Gated Slot Attention for Efficient Linear-Time Sequence Modeling](https://arxiv.org/abs/2409.07146) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/models/gsa) | | 2024 | MLA | [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://arxiv.org/abs/2405.04434) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/mla.py) | | 2025 | Samba | [Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling](https://arxiv.org/abs/2406.07522) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/models/samba) | | 2025 | Gated DeltaNet | [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://arxiv.org/abs/2412.06464) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/gated_delta_rule) | | 2025 | RWKV7 | [RWKV-7 "Goose" with Expressive Dynamic State Evolution](https://arxiv.org/abs/2503.14456) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/rwkv7) | | 2025 | NSA | [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://arxiv.org/abs/2502.11089) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/nsa) | | 2025 | FoX | [Forgetting Transformer: Softmax Attention with a Forget Gate](https://arxiv.org/abs/2503.02130) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/forgetting_attn) | | 2025 | DeltaProduct | [DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products](https://arxiv.org/abs/2502.10297) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltaproduct.py) | | 2025 | Rodimus* | [Rodimus*: Breaking the Accuracy-Efficiency Trade-Off with Efficient Attentions](https://arxiv.org/abs/2410.06577) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/rodimus.py) | | 2025 | MesaNet | [MesaNet: Sequence Modeling by Locally Optimal Test-Time Training](https://arxiv.org/abs/2506.05233) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/mesa_net.py) | | 2025 | Comba | [Comba: Improving Bilinear RNNs with Closed-loop Control](https://arxiv.org/abs/2506.02475) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/comba.py) | | 2025 | PaTH | [PaTH Attention: Position Encoding via Accumulating Householder Transformations](https://arxiv.org/abs/2505.16381) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/path_attn.py) | | 2025 | MoM | [MoM: Linear Sequence Modeling with Mixture-of-Memories](https://arxiv.org/abs/2502.13685) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/mom.py) | | 2025 | Log-Linear Attention | [Log-Linear Attention](https://arxiv.org/abs/2506.04761) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/log_linear_attn) | | 2025 | DeltaFormer | [Understanding Transformer from the Perspective of Associative Memory](https://arxiv.org/abs/2505.19488v1) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/deltaformer.py) | | 2025 | KDA | [Kimi Linear: An Expressive, Efficient Attention Architecture](https://arxiv.org/abs/2510.26692) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/kda) | | 2025 | MoBA | [MoBA: Mixture of Block Attention for Long-Context LLMs](https://arxiv.org/abs/2502.13189) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/moba.py) | | 2026 | Mamba3 | [Mamba-3: Improved Sequence Modeling using State Space Principles](https://arxiv.org/abs/2603.15569) | [code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/models/mamba3) | | 2026 | Raven | [Raven: High-Recall Sequence Modeling with Sparse Memory Routing](https://github.com/goombalab/raven/blob/main/raven.pdf) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/models/raven) | | 2026 | Gated DeltaNet 2 | [Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention](https://arxiv.org/abs/2605.22791) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/gdn2) | | 2026 | Wall | [Wall Attention: Length Generalization With Diagonal Gates](https://blog.tilderesearch.com/blog/wall-attn) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/wall_attn) | | 2026 | Parallax | [Parallax: Parameterized Local Linear Attention for Language Modeling](https://arxiv.org/abs/2605.29157) | [code](https://github.com/fla-org/flash-linear-attention/tree/main/fla/ops/parallax) | ## 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: ```sh # 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](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: ```py >>> 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`: ```py >>> from fla.models import GLAConfig >>> from transformers import AutoModelForCausalLM >>> config = GLAConfig() >>> model = AutoModelForCausalLM.from_config(config) ```
Click to expand config and model structure ``` /* Detailed source-code truncated for AI context efficiency. */ ```
### Fused Modules We offer a collection of fused modules in `fla.modules` to facilitate faster training: * [`Rotary Embedding`](fla/modules/rotary.py): rotary positional embeddings as adopted by the Llama architecture, a.k.a., Transformer++. * [`Norm Layers`](fla/modules/layernorm.py): * `RMSNorm`, `LayerNorm` and `GroupNorm` * `RMSNormLinear`, `LayerNormLinear` and `GroupNormLinear` to reduce memory usage of intermediate tensors for improved memory efficiency. * [`Norm Layers with Gating`](fla/modules/fused_norm_gate.py): combine norm layers with element-wise sigmoid or swish gating, as used by RetNet/GLA. * [`Cross Entropy`](fla/modules/fused_cross_entropy.py): faster Triton implementation of cross entropy loss. * [`Linear Cross Entropy`](fla/modules/fused_linear_cross_entropy.py): fused linear layer and cross entropy loss to avoid the materialization of large logits tensors. Also refer to implementations by [mgmalek](https://github.com/mgmalek/efficient_cross_entropy) and [Liger-Kernel](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py). * [`Linear KL Divergence`](fla/modules/fused_kl_div.py): 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: ```py >>> 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](benchmarks/benchmark_generation.py) for benchmarking the generation speed. Simply run it by: ```sh $ 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`](https://huggingface.co/fla-hub). ```py >>> 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: ```py >>> 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 ``` /* Detailed source-code truncated for AI context efficiency. */ ```
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: ```py >>> 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`](https://github.com/fla-org/flame) built on top of `torchtitan`, for efficient training of `fla` models. Check out [the GLA example](https://github.com/fla-org/flash-linear-attention/blob/main/examples/training.md) for more details. ## Evaluation The [lm-evaluation-harness](https://github.com/EleutherAI/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](https://github.com/EleutherAI/lm-evaluation-harness/blob/main/README.md). 2. Run evaluation with: ```sh $ 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](evals/harness.py) 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: ```sh $ 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. ``` 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: ```sh pip install lm_eval["ruler"] ``` and run evaluation by (e.g., 32k contexts): ```sh $ 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](https://github.com/EleutherAI/lm-evaluation-harness?tab=readme-ov-file#multi-gpu-evaluation-with-hugging-face-accelerate) 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). ``` /* Detailed source-code truncated for AI context efficiency. */ ``` ## Citation If you find this repository helpful, please cite our work: ``` /* 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](https://www.bitdeer.com/) and [Moonshot AI](https://www.moonshot.ai/) for their support in maintaining and powering our project infrastructure. ---