cuLA: hand-written CUDA kernels for linear attention

Why linear attention needs its own kernels

Linear attention mechanisms reformulate standard attention to use linear-time state updates instead of quadratic pairwise interactions, which makes them well suited to long-context LLM workloads. The quadratic cost of pairwise attention is what limits sequence length, and the linear formulation replaces it with a running state.

Recent variants push the idea further. The README names GLA, KDA, GDN and Lightning Attention, and attributes their improvements to gating, delta-style updates and chunkwise decomposition. Each of those is a different way of controlling how information flows into and out of the recurrent state, and each produces a different kernel problem.

cuLA provides hand-tuned CUDA implementations of these variants, targeting NVIDIA Blackwell (SM10X) and Hopper (SM90) GPUs. The two implementation routes are named in the project title: the CuTe DSL, NVIDIA's Python DSL for writing GPU kernels, and CUTLASS C++ for the pieces that need it.

The positioning is explicit. cuLA is designed as a submodule of flash-linear-attention, sharing the same interface, so adopting it requires only a one-line import change. The example in the README shows exactly that, swapping `fla.ops.kda` for `cula.kda`. The reason it is not merged in yet is maintenance: cuLA is currently developed as a standalone library, with the end goal being to reach these kernels through FLA, and the README says FLA already has a kernel dispatch mechanism in place so integration will be ready soon.

The early stage warning is the first thing to read

The README carries an explicit warning, and an engineer should read it as a statement of what the project is not yet.

cuLA is in its early development phase. Many kernels still have significant room for optimization, and the API may evolve. The project asks for contributions, whether performance tuning, new algorithm support, bug fixes or architectural improvements.

Two practical consequences follow from that wording. The first is that the public function signatures you write against today are not a contract, so a project that depends on cuLA should expect to adjust call sites on upgrade. The second is that the benchmark numbers, however attractive, describe a moving target. A kernel that has significant room for optimization will get faster, and a benchmark table you use for a purchase decision today will not describe the library in six months.

The repository layout supports that reading. There are several benchmark documents at the root, one per hardware configuration, including `BENCHMARK_GB200.md`, `BENCHMARK_GB200_CUDA_130.md`, `BENCHMARK_GB300.md` and `BENCHMARK_H200.md`, plus decode-specific ones for KDA on GB200 and H203E. There is a `benchmarks/` directory for the harness, a `csrc/` directory for C++ sources, a `cula/` directory for the Python package, `docs/` for the per-pipeline documents, `scripts/`, `tests/` and a `third_party/` directory holding the flash-linear-attention checkout. A `REPO_LAYOUT.md` and a `RECOMMENDED_CODING_STYLE.md` sit beside them, along with a `CITATION.cff` and a `.coderabbit.yaml`.

The last push to the tree was on 2026-10-19, the repository is not archived, and the release history is short: v0.2.0 on 2026-08-12 and v0.1.0 on 2026-04-02.

The hardware requirements are the real installation constraint

cuLA supports Hopper (SM90) and Blackwell (SM10X) GPUs, and the requirements are spelled out in one line: Python 3.12+, CUDA Toolkit 12.9+ for SM10X support, NVCC 12.9+, and PyTorch 2.9.1+.

Those are hard floors, and the PyTorch one is a matching requirement rather than a version range. The README notes that the PyTorch CUDA version must match your system CUDA Toolkit version, and the README points to `nvcc --version` and `python -c "import torch; print(torch.version.cuda)"` to check it.

A mismatch between the PyTorch build's CUDA and the toolkit on the machine is the most likely first failure, and it is worth checking before anything else because the fix is reinstalling PyTorch against the right index URL rather than debugging the kernels.

The build script reinforces the point. `setup.py` queries the visible CUDA devices through torch, reads each device's compute capability, and reports which of the SM100, SM103 and SM90 targets it found. The logic for each target is that `CULA_BUILD_ALL_ARCHS` enables everything unconditionally, an explicit environment variable is honoured if set, and otherwise a target is disabled when no matching GPU is detected. In other words, a machine without a Blackwell card will not build Blackwell kernels unless you ask for them, which is a sensible default and also a reason the fat wheel needs an explicit build flag.

The package metadata says `requires-python = ">=3.10"` while the README asks for Python 3.12+, and the ruff configuration targets `py312`. The README requirement is the one to follow, since the kernels are written against a newer Python than the packaging metadata declares.

Installing a prebuilt fat-binary wheel

Pre-built fat-binary wheels covering SM90, SM100 and SM103 are published on the GitHub Releases page. The Linux wheels target `manylinux_2_28`, so they need glibc 2.28 or newer.

The install command carries a version placeholder, a CUDA tag and a URL built from the release tag.

    pip install "cuda-linear-attention==<VERSION>+<CUDA_TAG>" -f https://github.com/inclusionAI/cuLA/releases/expanded_assets/<TAG>

The substitutions are given as examples: `<TAG>` is the release tag such as `v0.2.0`, `<VERSION>` is the base version such as `0.2.0`, and `<CUDA_TAG>` is the PyTorch CUDA build tag such as `cu129` or `cu130`. You can also download the `.whl` file directly from the Releases page and install it with `pip install <filename>.whl`.

The distribution name is `cuda-linear-attention` rather than `cula`, which is worth noting since the import name and the package name differ. The `pyproject.toml` records the name as `cuda-linear-attention` with the description `cuLA CUDA extension`, the version is dynamic and comes from setuptools_scm, and the dependencies are `nvidia-cutlass-dsl` with a version range that excludes 4.5.0, `apache-tvm-ffi` and `packaging`. The licence is Apache-2.0.

That cutlass-dsl constraint is worth a second look if you already have the DSL installed for other work, since the range pins it to a narrow band and excludes one release in the middle.

Building from source, and the fat wheel for all architectures

Building from source takes three steps, and the first one is a git clone with submodules, since the CUTLASS dependency is vendored rather than fetched by pip.

git clone https://github.com/inclusionAI/cuLA.git
git submodule update --init --recursive

PyTorch is installed against a specific CUDA index, and the version is pinned rather than floating.

pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu129

Then the package itself and its dependency, with the flash-linear-attention checkout installed in editable mode for benchmark reproduction, and the build isolation turned off so the local PyTorch is used.

# Install flash-linear-attention for benchmark repro
pip install -e third_party/flash-linear-attention

# Install cuLA
pip install -e . --no-build-isolation

Building a wheel that covers all three architectures takes one environment variable and the build module.

CULA_BUILD_ALL_ARCHS=1 python -m build --wheel --no-isolation

That flag matters if you are producing a wheel to distribute rather than to install locally. Without it, the build script disables any architecture for which it cannot see a matching GPU, so a build machine with an H200 produces an SM90 wheel and nothing else.

A first call: Lightning Attention prefill on Hopper

The quick start runs on an SM90 GPU, and the API takes tensors in the usual batched layout.

import torch
from cula.lightning import lightning_attn_fwd

B, T, H, D = 2, 4096, 64, 128
q = torch.randn(B, T, H, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)
decay = torch.linspace(0.0, 0.04, H, device="cuda", dtype=torch.float32)

output, final_state = lightning_attn_fwd(
    q,
    k,
    v,
    decay,
    scale=D**-0.5,
    output_final_state=True,
)

The function returns the output and, when `output_final_state=True`, the final recurrent state. The decay tensor is the gating term, and it is float32 while the query, key and value tensors are bfloat16.

The README then describes what the SM90 CuTe DSL backend actually supports: fixed and packed variable-length prefill, optional recurrent state, GVA head mapping, and persistent packed scheduling. The numerics are fixed in the same paragraph: BF16 for Q, K and V, FP32 for decay and state, head dimension 128, and chunk size 64. Those are not defaults you can change, they are the configuration the kernel is built for, and a workload with a different head dimension or a different chunking need is out of scope for this path.

The document `docs/lightning_attn_sm90.md` is the reference for the internals, covering the warp-group roles, the WGMMA dataflow, the pipeline stages, where the recurrent state is placed, and the fixed and packed scheduling modes. That level of detail is what you would want before tuning the kernel for your own sequence shapes.

KDA on Blackwell, where the import change happens

The second quick start covers KDA, Kimi Delta Attention, on Blackwell, and the whole point of it is the import line.

import torch
from cula.kda import chunk_kda  # <-- one-line change from fla.ops.kda

B, T, H, K, V = 2, 2048, 32, 128, 128
device = 'cuda'

q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16, requires_grad=True)
k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16, requires_grad=True)
v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16, requires_grad=True)
g = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) * 0.1   # gate (log space)
beta = torch.randn(B, T, H, device=device, dtype=torch.bfloat16).sigmoid()
A_log = torch.randn(H, device=device, dtype=torch.float32) * 0.01
dt_bias = torch.zeros(H * K, device=device, dtype=torch.float32)
init_state = torch.zeros(B, H, K, V, device=device, dtype=torch.float32)

The call passes a good many keyword arguments, and the ones that are not obvious from the names are documented in the notes. `safe_gate=True` is required to get TensorCore acceleration, which makes it less of an option and more of a prerequisite. `beta` accepts both float32 and bfloat16, while `initial_state` must be float32. The `cu_seqlens` tensor for variable-length sequences must be int32.

The gate tensor `g` is in log space, and it is multiplied by a small constant in the example, which is the kind of detail that changes results if you get it wrong.

The claim behind all this is that a swap of the import path gets you a different kernel implementation with the same interface. If that holds for your workload, the adoption cost is genuinely small, and the FLA integration will remove even the import change. If it does not hold, because the semantics of the arguments differ subtly, the cost is a debugging session across two libraries.

What the benchmark numbers do and do not tell you

The benchmarks section states its conditions before any numbers: they run on a single NVIDIA GB200 or H200 GPU, with PyTorch 2.9.1 and Triton 3.5.1, against flash-linear-attention v0.5.0 as the baseline.

The reported highlights are specific. KDA Modular Forward on Blackwell averages a 1.33x speedup on fixed-length shapes and 1.35x on variable-length across 18 configurations spanning uniform, skewed and random distributions. Lightning Attention Prefill on Blackwell reaches up to 2.08x at B=2. Lightning Attention variable-length on Blackwell averages 1.47x across 126 configurations. On Hopper, FlashKDA Prefill averages 2.72x over FLA across 28 fixed-length and variable-length configurations, with a maximum of 7.56x, and FlashKDA Intracard CP reports a 4.29x geometric mean.

Three things about how to read this. The baseline is FLA v0.5.0, a specific pinned release, not the current tip of that project, and the two libraries will move independently. The results are single GPU, so they say nothing about how the kernels behave when a model is sharded across devices, where the communication pattern dominates. And the distributions are named, uniform, skewed and random, which is a reasonable spread, but the shapes are chosen by the benchmark author.

The detailed results live in the per-hardware documents rather than in the README, and the README points at the right one for each case: `BENCHMARK_GB200_CUDA_130.md` for Blackwell tested with CUDA 13.0, and `BENCHMARK_H200.md` for the CuTe DSL FlashKDA results on an H200 with 141GB of memory and CUDA 12.9. The `USAGE.md` file has the usage examples and notes beyond the two quick starts.

Editorial conclusion

Adopt cuLA if you are running long-context models on Hopper or Blackwell hardware, you already depend on flash-linear-attention, and the import change it promises is worth the early stage risk. Do not adopt it on Hopper hardware below the stated requirements, on a CUDA toolkit older than 12.9, or for a workload that a Triton implementation already handles well enough. Before integrating, read the SM90 Lightning prefill pipeline document in `docs/` to see the warp-group roles and dataflow, and check that the PyTorch CUDA build on your machines matches the system toolkit with `nvcc --version` and `python -c "import torch; print(torch.version.cuda)"`, since a mismatch is the failure the README warns about first.

Frequently asked questions

Official sources

1. inclusionAI/cuLA on GitHub
2. Issues
3. License: Apache-2.0
4. README
5. Releases