cuLA — CUDA Linear Attention

High-performance CUDA kernels for linear attention variants, written entirely in CuTe DSL.

Introduction

Linear attention mechanisms reformulate standard attention to use linear-time state updates instead of quadratic pairwise interactions, making them well suited for long-context LLM workloads. Recent variants such as GLA, KDA, GDN, and Lightning Attention further improve expressiveness with gating, delta-style updates, and chunkwise decomposition.

cuLA provides hand-tuned CuTeDSL implementations of these linear attention variants, targeting NVIDIA Blackwell (SM10X) and Hopper (SM90) GPUs. The runtime is a standalone kernel library: it does not import or vendor FLA, Triton, or a C++/CUDA extension. Complete operators and deliberately exposed atomic kernels are available through cula.kernels.

⚠️ Early Stage: cuLA is in its early development phase. Many kernels still have significant room for optimization, and the API may evolve. We warmly welcome contributions from the community — whether it's performance tuning, new algorithm support, bug fixes, or architectural improvements. Every contribution helps push the boundaries of linear attention on modern GPUs!

Installation

cuLA supports both Hopper (SM90) and Blackwell (SM10X) GPUs.

Requirements (Hopper & Blackwell): Python 3.12+, CUDA Toolkit 12.9+, PyTorch 2.8.0+, and a matching nvidia-cutlass-dsl installation. cuLA does not build a C++/CUDA extension.

Note: The PyTorch CUDA version must match your system CUDA Toolkit version. Check with nvcc --version and python -c "import torch; print(torch.version.cuda)".

Pre-built Wheels

Platform-independent wheels containing the CuTeDSL sources are available on GitHub Releases:

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

Replace <TAG> with the release tag (e.g., v0.2.0) and <VERSION> with the version (e.g., 0.2.0). Or download the .whl file directly from the Releases page and install it with pip install <filename>.whl.

Build from Source

Clone cuLA & dependencies:

git clone https://github.com/inclusionAI/cuLA.git

Install PyTorch and CuTeDSL:

For CUDA 12.9:

pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu129
pip install 'nvidia-cutlass-dsl>=4.6.2,<4.9'

For CUDA 13:

pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu130
pip install 'nvidia-cutlass-dsl[cu13]>=4.6.2,<4.9'

Install cuLA & dependencies:

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

Build the platform-independent wheel:

python -m build --wheel

Quick Start

Lightning Attention Prefill — Hopper (SM90)

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 SM90 CuTe DSL backend supports fixed and packed variable-length prefill, optional recurrent state, GVA head mapping, and persistent packed scheduling. It uses BF16 Q/K/V, FP32 decay/state, head dimension 128, and chunk size 64. See the SM90 Lightning prefill pipeline for the warp-group roles, WGMMA dataflow, pipeline stages, recurrent-state placement, and fixed/packed scheduling modes.

KDA (Kimi Delta Attention) prefill

import torch
from cula.kernels.kda import kda_prefill

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

q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16)
k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16)
v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16)
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)

# Forward-only standalone operator
o, final_state = kda_prefill(
    q=q, k=k, v=v, g=g, beta=beta,
    A_log=A_log,
    dt_bias=dt_bias,
    initial_state=init_state,
    output_final_state=True,
    use_qk_l2norm_in_kernel=True,
    use_gate_in_kernel=True,
    safe_gate=True,
    lower_bound=-5.0,
)

print(f'Output shape: {o.shape}')             # [2, 2048, 32, 128]
print(f'Final state shape: {final_state.shape}')  # [2, 32, 128, 128]

Notes:

• kda_prefill is forward-only. The modular SM100 training primitives are exposed separately rather than hidden behind an FLA autograd wrapper.
• safe_gate=True is required by this path.
• beta supports both float32 and bfloat16; initial_state must be float32.
• cu_seqlens (for variable-length sequences) must be int32.

Usage

See USAGE.md for detailed usage examples and notes. See the standalone kernel inventory for the public kernel boundary and the overlap assessment against FlashInfer and cuDNN. The optional cuDNN hybrid-attention runtime documents packed index pooling, memory-bounded DSA selection, and CP layout ownership without precision downgrades.

Benchmarks

Benchmarks run on a single NVIDIA GB200/H200 GPU with PyTorch 2.9.1, Triton 3.5.1.

Historical result files use FLA as an external baseline. FLA is not a runtime or source-tree dependency.

Blackwell (SM10X)

See BENCHMARK_GB200_CUDA_130.md tested with CUDA 13.0 for detailed results.

Hopper (SM90)

See BENCHMARK_H200.md for CuTe DSL FlashKDA results on an H200 141GB with CUDA 12.9.

Highlights:

• KDA Modular Forward (Blackwell): avg 1.33x speedup on fixed-length, avg 1.35x on variable-length (18 configs, uniform/skewed/random).
• Lightning Attention Prefill (Blackwell): up to 2.08x speedup (B=2).
• Lightning Attention Varlen (Blackwell): avg 1.47x speedup across 126 configs (uniform/skewed/random).
• FlashKDA Prefill (Hopper): avg 2.72x speedup over FLA across 28 fixed-length and variable-length configs, up to 7.56x.
• FlashKDA Intracard CP (Hopper): 4.29x geo-mean speedup over serial FlashKDA on 28 CP-engaged long-sequence configs, up to 7.83x.

The commands below run the currently maintained benchmark entry points. The Blackwell KDA command is an intra-kernel microbenchmark; it does not reproduce the historical full-KDA headline numbers above, which used different shapes and included the complete operator pipeline.

The comparison commands below require FLA to be installed separately; FLA is not imported by the cuLA runtime package.

# Blackwell (SM10X)
python benchmarks/bench_kda_chunk_intra.py
python benchmarks/bench_lightning_attn_prefill.py --modes no_state varlen
python benchmarks/bench_la_decode_vs_fla.py --heads 64 --head-dim 128

# Hopper (SM90)
python benchmarks/bench_kda_sm90_prefill.py --mode both
python benchmarks/bench_kda_sm90_cp.py
python benchmarks/bench_kda_bwd_wy_dqkg_sm90.py --mode both --heads 32

Tests

# Time-bounded default suite (representative CUDA shapes, about five minutes on GB200)
python -m pytest -q

# Extended CUDA parameter sweeps only
python -m pytest -m cula_slow -q

# Default tests plus every extended CUDA case
python -m pytest -m cula_full -q

# Optional reference tests (require a separately installed FLA)
# SM90 CuTeDSL two-kernel prefill + intracard CP
python -m pytest tests/test_kda_sm90_prefill_vs_fla.py tests/test_kda_sm90_intracard_cp.py -v
# Tests for Lightning Attention prefill on SM100
python tests/test_lightning_sm100_prefill.py
# Tests for the SM90 Lightning public dispatch, semantics, and kernel structure
python -m pytest tests/test_lightning_attn_prefill_dispatch.py tests/test_lightning_attn_prefill_sm90.py -v
# Tests for Lightning Attention decode
python -m pytest tests/test_lightning_decode.py -v

Sample test output

tests/test_kda_e2e_compare_fla.py::test_safe_gate_chunk[B1-T63-H1-D128-...]    PASSED
tests/test_kda_e2e_compare_fla.py::test_safe_gate_chunk[B2-T500-H3-D128-...]   PASSED
tests/test_kda_e2e_compare_fla.py::test_safe_gate_chunk[B2-T1000-H3-D128-...]  PASSED
tests/test_kda_e2e_compare_fla.py::test_safe_gate_chunk[B3-T1024-H4-D128-...]  PASSED
tests/test_kda_e2e_compare_fla.py::test_safe_gate_chunk[B4-T1024-H4-D128-...]  PASSED
tests/test_kda_e2e_compare_fla.py::test_safe_gate_chunk[B4-T2048-H8-D128-...]  PASSED
tests/test_kda_e2e_compare_fla.py::test_safe_gate_chunk_varlen[...]             PASSED
...
======================= 17 passed in 40.95s =======================

CUDA kernel tuning is significantly more labor-intensive than Triton — contributions from the open-source community are warmly welcomed!

Repository Layout

See REPO_LAYOUT.md for the full directory structure and a summary of each component.

Roadmap

• Stabilize the standalone atomic-kernel contracts and publish per-kernel benchmarks.
• Polynomial approximation to mitigate the exponential bottleneck, as in Flash-Attentiton-4.
• Larger chunk size and 2-CTA on SM10X for improved throughput.
• Continuous optimization via agentic methods such as AVO.
• Support for more algorithms.
• Small B/H/S optimizations.
• Support for BF16 beta input.

Train

• Modular KDA Forward (SM10X)

  • kda chunk intra
  • chunk gated delta h
  • recompute wu
  • chunk fwd o

• Modular GDN Forward / Backward Kernels (compatible with Kimi CP)

• Backward pass optimizations.

• Kernel-level compute-communication overlapping for CP linear attention kernels (via nvshmem)

Inference

• Lightning prefill kernel (SM90 & SM10X)

• Lightning decode kernel (SM90 & SM10X)

• Fused KDA prefill kernel (SM90)

• Fused KDA prefill kernel (SM10X)

• KDA and Lightning Attention MTP decode/verify support

• More aggressive fusion of small neighboring kernels like cumsum for inference scenarios.

Acknowledgements

This project is inspired by flash-linear-attention, CUTLASS, CuTe DSL, FlashInfer, Flash-Attention, and FlashMLA. We thank FLA-org and NVIDIA for their great work.

Citation

If you find cuLA useful, please cite it using the metadata in our CITATION.cff file:

@software{cula2026,
  title   = {cuLA: CUDA Linear Attention},
  author  = {Chaofan Yu, Bowen Zeng, Hao Chen, Zhe Yang, Zhiqiang Zhang, Huan Li and Jun Zhou},
  year    = {2026},
  url     = {https://github.com/InclusionAI/cuLA}
}