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

Qwen has released FlashQLA, a high-performance linear attention kernel library designed to optimize Gated Delta Network (GDN) layers. Built on TileLang, FlashQLA achieves a 2-3× forward speedup and a 2× backward speedup over the FLA Triton kernel on NVIDIA Hopper GPUs, specifically benefiting pretraining and edge-side agentic inference.

Optimizing GDN Chunked Prefill

FlashQLA addresses two primary efficiency bottlenecks in the Gated Delta Network (GDN) Chunked Prefill process found in the original FLA implementation:

  1. Memory-Bound Kernels: The standard flow repeatedly reads and writes intermediate variables ($W, U, S$) to High Bandwidth Memory (HBM), creating significant overhead.
  2. Low GPU Utilization: The recurrent nature of State Space Model (SSM) states limits the number of simultaneous thread blocks to batch_size * num_heads. In scenarios with small models, small batches, or Tensor Parallelism (TP), this leads to idle GPU Streaming Multiprocessors (SMs).

To resolve these conflicting issues, FlashQLA avoids a fully-fused kernel (which would fail in small-batch scenarios) and instead splits the forward computation into two fused kernels with Context Parallelism (CP) preprocessing steps inserted between them.

Key Technical Innovations

Gate-Driven Automatic Intra-Card Context Parallelism (AutoCP)

FlashQLA implements an automatic intra-card CP mechanism under TP, long-sequence, and small-head-count settings to improve SM utilization. It uses a mathematical model to determine the optimal degree of parallelism ($L = \lambda \sqrt{N}$), where $N$ is the number of chunks and $L$ is the number of chunks per CP rank.

To further reduce overhead, FlashQLA exploits the exponential decay property of the GDN gate. For heads where the gate $\alpha_i \in (0,1)$, the influence of previous states decays exponentially. FlashQLA uses a "warmup" process (typically 6–8 chunks) to drive the state error below the noise floor, allowing the system to discard the expensive correction term $M$ matrix computation and directly obtain an accurate sub-sequence $S_0$.

TileLang Warp-Specialized Kernels

FlashQLA utilizes TileLang to implement warpgroup-specialized kernels. This architecture employs one producer warpgroup and three consumer warpgroups within the same SM, exchanging data via shared memory and synchronizing with mbarriers.

  • Forward Pass: Three consumer warpgroups compute $V'$, $S$, and $O$ respectively, using a ping-pong structure to overlap computation and memory traffic.
  • CP Preprocessing: A single fused kernel handles both the original $M$ and $S$ computation and the lighter sliding-window warmup approach.
  • Backward Pass: FlashQLA fuses bwd_dv, bwd_dhu, bwd_dqkwg, and bwd_wy into a single kernel. Due to on-chip resource constraints, it relies on a long compute chain to hide memory traffic rather than multi-stage pipelining.

Performance Benchmarks

Benchmarks conducted on NVIDIA H200 GPUs against FLA Triton and FlashInfer baselines show significant gains, particularly as the Tensor Parallelism (TP) degree increases.

Model / TP Seqlen $h_{qk}$ $h_v$ FlashQLA FlashInfer FLA vs FLA vs FI
397B/122B TP8 1x32768 2 8 0.310ms 1.653ms 2.95×
397B/122B TP4 1x32768 4 16 0.486ms 1.654ms 2.57×
27B TP2 1x32768 8 24 0.659ms 1.616ms 2.37×
2B/0.8B TP1 1x32768 16 16 0.493ms 1.640ms 2.60×

Implementation and Requirements

FlashQLA provides a high-level API compatible with FLA's signature and low-level entry points for forward and backward passes.

System Requirements:

  • Hardware: NVIDIA SM90 (Hopper)
  • Software: CUDA 12.8+, PyTorch 2.8+

The code and benchmarks are available at github.com/QwenLM/FlashQLA.

Sources