FlashQLA: A CP- and Backward-Friendly Fused Linear Attention Kernel Library for GDN

FlashQLA has been released. It's a TileLang-based linear attention kernel library that jointly optimizes the forward and backward passes of GDN Chunked Prefill, achieving 2-3x speedup in forward and 2x speedup in backward compared to FLA Triton on NVIDIA Hopper. The efficiency gains are especially notable in pretraining and edge-side agentic inference.

As GDN spread from Qwen3-Next to Qwen3.5/3.6, GDN block overhead began to stand out noticeably at 397A17B / 122A10B / 35B / 27B-class models and 256K+ context lengths.

The existing FLA GDN Chunked Prefill had two bottlenecks.

• A memory-bound structure with frequent HBM round-trips, requiring continuous reads and writes of intermediate states such as K, V, W, U, S.
• Limited thread block counts due to recursive state, leading to low SM utilization in small models, small batches, and TP environments.

Rather than fully fusing everything, the solution splits the forward pass into two fused kernels with CP preprocessing in between, taking CP (context parallelism) into account. Parallelism is determined by the number of chunks N and chunks per rank L, set as L = λ√N, and in practice CP is enabled only when batch_size * num_heads <= 40 or batch_size * num_heads <= 56 && seq_len >= 8192. Outside these ranges, the existing path of splitting v_head_dim to raise parallelism by 2-4x is retained.

Warmup using gate decay was also added. According to the article, sliding-window properties were observed in 60-80% of linear attention heads, and warming up just 6-8 chunks brought state error below noise level. In this case, sufficient accuracy was obtained without computing the correction matrix M, and CP preprocessing handled both the path that computes both M and S and the path that computes only S within a single fused kernel. Warmup length is determined by a separate kernel that gathers gate statistics, at nearly no cost.

The implementation follows a warpgroup specialization pattern where 1 producer and 3 consumer warpgroups collaborate via shared memory and mbarrier. The forward pass overlaps computation of V', S, and O, while the backward pass recomputes S using the earlier CP preprocessing kernel, then hides memory traffic with a long compute chain instead of multi-stage pipelining due to on-chip resource constraints. Afterward, bwd_dv, bwd_dhu, bwd_dqkwg, and bwd_wy are fused together to reduce additional redundant accesses.

Benchmarks were measured against FLA 0.5.0, Triton 3.5.1, FlashInfer 0.6.9, and TileLang 0.1.8. On H200, the 397B/122B TP8 1x32768 scenario achieved 0.310ms, outperforming FlashInfer 1.653ms and FLA 0.913ms. Consistent improvements were also seen for 397B/122B TP4, 27B TP2, 2B/0.8B TP1, and Sym h32, with the benefit of intra-card AutoCP growing larger as TP increases. Code and benchmarks are available at github.com/QwenLM/FlashQLA.