inclusionAI/cuLA

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

What it solves

cuLA addresses the performance bottlenecks of linear attention mechanisms in long-context LLM workloads. While linear attention reduces the quadratic complexity of standard attention to linear-time updates, implementing these efficiently on modern GPUs requires deep hardware-level optimization to fully leverage the compute capabilities of the latest NVIDIA architectures.

How it works

cuLA provides hand-tuned CUDA kernels written using the CuTe DSL and CUTLASS C++. It implements high-performance variants of linear attention, including GLA, KDA, GDN, and Lightning Attention. These kernels are specifically optimized for NVIDIA Blackwell (SM10X) and Hopper (SM90) GPUs, utilizing features like WGMMA dataflow and persistent packed scheduling to maximize throughput during both prefill and decode phases.

Who it’s for

It is designed for AI researchers and engineers working with long-context LLMs who need maximum inference and training performance on NVIDIA Hopper or Blackwell hardware. It is also intended as a high-performance backend for the flash-linear-attention (FLA) library.

Highlights

  • Hardware Optimized: Specifically tuned for NVIDIA Blackwell (SM10X) and Hopper (SM90) architectures.
  • Significant Speedups: Demonstrates substantial performance gains over FLA Triton implementations, including up to 7.56x speedup for FlashKDA prefill on Hopper.
  • Broad Algorithm Support: Implements multiple linear attention variants such as KDA (Kimi Delta Attention) and Lightning Attention.
  • Seamless Integration: Designed to be a drop-in replacement for FLA kernels with a simple import change.

Related

  • Project
    MoonshotAI/FlashKDA

    A high-performance CUDA kernel implementation of Kimi Delta Attention (KDA) built on CUTLASS for NVIDIA SM90+ GPUs, serving as an optimized backend for the flash-linear-attention library.

  • Project
    QwenLM/FlashQLA

    A high-performance linear attention kernel library that accelerates Gated Delta Rule (GDN) Chunked Prefill, providing significant speedups for training and inference on NVIDIA Hopper and Blackwell GPUs.

  • Project
    Tencent/hpc-ops

    A production-grade operator library for LLM inference that provides highly optimized CUDA kernels for attention, MoE, and sampling to reduce serving latency.

  • Project
    fla-org/flash-linear-attention

    A library of hardware-efficient, training-ready implementations for linear attention, state space models, and hybrid LLM architectures across multiple hardware platforms.

  • Project
    deepseek-ai/FlashMLA

    A library of optimized attention kernels for Multi-head Latent Attention (MLA), powering DeepSeek-V3 and V3.2 with support for sparse and dense attention patterns.