蚂蚁 inclusionAI 开源 cuLA：面向线性注意力变体的 CUDA 内核库

inclusionAI; cuLA

Repository navigation

NotificationsYou must be signed in to change notification settings; Fork69 (69); Star555 (555)You must be signed in to star a repository

CUDA kernels for linear attention variants, written in CuTe DSL and CUTLASS C++.

Repository files navigation

README; Apache-2.0 license

High-performance CUDA kernels for linear attention variants, written in CuTe DSL and CUTLASS C++.

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 CUDA implementations of these linear attention variants, targeting NVIDIA Blackwell (SM10X) and Hopper (SM90) GPUs. It is designed as a submodule of flash-linear-attention (FLA), sharing the same interface — adopting cuLA requires only a one-line import change. For ease of maintenance, cuLA is currently developed as a standalone library; the end goal is for users to seamlessly access these kernels through FLA. Since FLA already has a kernel dispatch mechanism in place, integration will be ready soon.

⚠️ 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+ (SM10X support), NVCC 12.9+, PyTorch 2.9.1+

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

Pre-built fat-binary wheels (SM90 + SM100 + SM103) are available on GitHub Releases. Linux wheels target manylinux_2_28 and require glibc 2.28 or newer:

Replace with the release tag (e.g., v0.2.0), with the base version (e.g., 0.2.0), and with your PyTorch CUDA build tag (e.g., cu129 or cu130). Or download the .whl file directly from the Releases page and install it with pip install .whl.

Install cuLA & dependencies:

Build fat wheel (SM90 + SM100 + SM103):

Quick Start

Lightning Attention Prefill — Hopper (SM90)

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.

Just change the import:

Notes:

safe_gate=True is required to leverage TensorCore acceleration; 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.

Benchmarks

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

FLA baseline: flash-linear-attention v0.5.0.

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

To reproduce the benchmark suites directly:

Tests

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

Integrate into flash-linear-attention via FLA's kernel dispatch mechanism; 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

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); MTP 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:

CUDA kernels for linear attention variants, written in CuTe DSL and CUTLASS C++.