QwenLM/FlashQLA

high-performance linear attention kernel library built on TileLang

What it is

FlashQLA is a high‑performance kernel library that implements the Gated‑Delta‑Rule (GDN) linear‑attention algorithm used in Qwen‑style large language models. It is written on top of the TileLang DSL and ships as a Python package (flash‑qla). The library provides both a convenient high‑level function (chunk_gated_delta_rule) and low‑level forward/backward kernels that can be called from PyTorch code.

Why it matters

Linear‑attention mechanisms replace the quadratic cost of classic soft‑max attention with a computation that scales linearly with sequence length, making very long contexts feasible. FlashQLA speeds up the forward pass by 2‑3× and the backward pass by ~2× compared with the reference Flash Linear Attention (FLA) Triton kernels, especially on NVIDIA Hopper/Blackwell GPUs. Faster kernels mean cheaper pre‑training and lower latency for inference‑heavy or edge‑side agents that rely on Qwen‑style models.

Core features

Feature What it does
Gate‑driven automatic intra‑card context parallelism Detects when the GDN gate’s exponential decay allows the work to be split across multiple CUDA warps/SMs, automatically enabling parallelism inside a single GPU card (TP, long‑seq, small‑head settings).
Hardware‑friendly algebraic reformulation Rewrites the GDN forward/backward equations to reduce Tensor‑Core, CUDA‑Core and SFU usage while keeping numerical precision.
TileLang‑fused warp‑specialized kernels Uses the TileLang DSL to generate a handful of fused kernels that balance data movement, Tensor‑Core matmuls and CUDA‑Core ops, rather than a single monolithic kernel.
Dual‑CP (inter/intra two‑level context parallelism) (v0.1.3) Adds a second level of parallelism across GPUs (inter‑card) and within a GPU (intra‑card).
SM120 backward path (v0.1.3) Adds support for the newest NVIDIA SM120 architecture in the backward kernel.
Plug‑and‑play backend for flash‑linear‑attention (v0.1.2) Can be used as a drop‑in replacement for the FLA GDN kernel via the standard FLA API.

How to get started

# Install the pre‑built wheel (requires CUDA 12.8+, PyTorch 2.8+)
pip install flash-qla

Or build from source:

git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .

Simple usage example

import torch
from flash_qla import chunk_gated_delta_rule

# q, k, v, g, beta are tensors shaped as described in the README
out, final_state = chunk_gated_delta_rule(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=scale,
    initial_state=init_state,
    output_final_state=True,
    cu_seqlens=cu_seqlens,
)

For custom training loops you can call the separate forward and backward functions (*_fwd / *_bwd) and feed the returned cache objects back into the backward pass.

When to use it

  • You are training or fine‑tuning a Qwen‑style model that uses GDN linear attention.
  • Your workloads run on recent NVIDIA GPUs (H100, H200, Blackwell) and you need every ounce of speed for long‑sequence pre‑training or low‑latency inference.
  • You already use PyTorch and want a drop‑in kernel that works with the existing flash-linear-attention API.

Limitations / open points

  • Only supports NVIDIA GPUs with SM 90‑SM 121 and CUDA 12.8+. No CPU or AMD support.
  • The library focuses on the GDN Chunked Prefill variant; other linear‑attention formulations are not covered.
  • Backward‑pass speedups are tied to the dual‑CP implementation; on older GPUs you may see less benefit.
  • Requires the TileLang runtime; users must be comfortable compiling TileLang‑generated kernels.

Community & maintenance

  • The project is maintained by the QwenLM team (see the citation list). Recent releases (v0.1.1‑v0.1.3) show active development.
  • Issues and pull requests can be filed on the GitHub repo; the README links to benchmark scripts and profiling utilities for developers.

All information above is taken directly from the repository’s README; no external speculation is added.

Related

  • Dispatch
    FlashQLA: CP-/Bwd-Friendly Fused Linear Attention Kernels for GDN

    Qwen has open-sourced FlashQLA, a high-performance linear attention kernel library built on TileLang that achieves 2-3x forward and 2x backward speedups for Gated Delta Network (GDN) layers on NVIDIA Hopper GPUs.

  • Project
    deepseek-ai/FlashMLA

    FlashMLA is DeepSeek’s open‑source library of ultra‑fast attention kernels (sparse and dense) for LLM inference. It provides CUDA and Ascend implementations that work with FP8/FP4‑quantized KV caches, a fused kernel that merges RoPE and casting, and APIs compatible with the popular flash‑attention package. The kernels achieve >1 000 TFlops on modern GPUs and ~95 % of peak on Ascend NPUs, dramatically speeding up inference of DeepSeek‑V4.1 and similar models.

  • Project
    Relaxed-System-Lab/Flash-Sparse-Attention

    Flash‑Sparse‑Attention (FSA) is a Triton‑based, drop‑in PyTorch module that implements an optimized version of Native Sparse Attention. By reordering loops and splitting the work into three specialized kernels, it cuts memory traffic and removes atomic adds, giving 2‑3× speed‑ups on long‑sequence LLM training and inference on modern NVIDIA GPUs. It works with PyTorch ≥ 2.4, transformers ≥ 4.45, and supports fp16/bf16 on Ampere/Hopper hardware.

  • Project
    flashinfer-ai/flashinfer

    FlashInfer is a GPU‑kernel library (with a Python package) that provides ultra‑fast, low‑precision attention, GEMM, MoE, and sampling operations for LLM inference. It automatically selects the best backend (FlashAttention‑2/3, cuDNN, CUTLASS, TensorRT‑LLM) for a wide range of NVIDIA GPUs, from Turing up to Blackwell, and is already used by popular inference servers like vLLM and TensorRT‑LLM.

  • Project
    flagos-ai/FlagAttention

    FlagAttention is a Triton‑based library offering fast, memory‑efficient attention kernels (FlashAttention, Piecewise, paged, sparse, quantised, recurrent, etc.) for training and inference of large models across many accelerators. Stable dense operators are ready for production; alpha variants provide experimental features like KV‑cache paging, block‑sparse decoding, and INT8 quantisation.