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

Introduction Following the release of Qwen3-Next, Gated Delta Network (GDN) has become the workhorse attention layer across the Qwen family — from Qwen3-Next-80B-A3B all the way to the subsequent Qwen3.5 / Qwen3.6 series. As models scale to 397A17B / 122A10B / 35B / 27B and context windows stretch beyond 256K, the overhead of the GDN block in end-to-end training and inference has become non-negligible. Today we officially open-source FlashQLA: a high-performance linear attention kernel library built on TileLang. FlashQLA applies reasonable operator fusion and performance optimization to the forward and backward passes of GDN Chunked Prefill, achieving 2-3× forward speedup and 2× backward speedup over the FLA Triton kernel across multiple scenarios on NVIDIA Hopper. The efficiency gains are particularly pronounced in pretraining scenarios and edge-side agentic inference. Key highlights of this release: Gate-driven automatic intra-card context parallelism. By exploiting the exponential decay property of the GDN gate, FlashQLA automatically enables intra-card CP under TP, long-sequence, and small-head-count settings, improving GPU SM utilization. Hardware-friendly algebraic reformulation. We reformulate the forward and backward flows of GDN Chunked Prefill to a certain extent, effectively reducing Tensor Core, CUDA Core, and SFU overhead without sacrificing numerical precision. TileLang fused warp-specialized kernels. Rather than following the step-by-step decomposition into independent kernels, nor fusing the entire computation flow into a single kernel, we take CP and backward requirements into account, use TileLang to build several key fused kernels, and manually implement warpgroup specialization to overlap data movement, Tensor Core computation, and CUDA Core computation. FlashQLA code and benchmarks are open-sourced at github.com/QwenLM/FlashQLA. Key Problems in FLA GDN Chunked Prefill Let us first review the forward computation flow of GDN Chunked Prefill, taking chunk index i as an example: Aᵢ ← (I+StrictLower( diag(βᵢ)(Γᵢ ⊙ KᵢKᵢ^ᵀ) ))⁻¹ {Wᵢ ← Aᵢdiag(βᵢ)diag(γᵢ)Kᵢ; Uᵢ ← Aᵢdiag(βᵢ)Vᵢ {Vᵢ’ ← Uᵢ−WᵢSᵢ; Sᵢ₊₁ ← γ_(i,C−1)Sᵢ + Kᵢ^ᵀdiag((γ_(i,C−1))/γᵢ)Vᵢ’ Oᵢ ← diag(γ)QᵢSᵢ + (Lower(Γᵢ) ⊙ QᵢKᵢ^ᵀ)Vᵢ′ Ignoring gate preprocessing and CP, each step of this flow corresponds to one kernel in FLA. This flow has two main efficiency problems: Most of the above are memory-bound kernels. The flow repeatedly reads K, V and other data, while W, U, S as intermediate variables must be written to HBM and then read by the next kernel, incurring significant memory access overhead. The recurrent nature of the SSM state means that the corresponding third step chunk_gated_delta_rule_fwd_kernel can only launch batch_size * num_heads thread blocks simultaneously, resulting in low GPU utilization in small-model, small-batch, or TP scenarios. The solutions to these two problems are contradictory. For the first problem, the most intuitive solution is to write a fully-fused kernel, where all data is accessed only once and all intermediate variables are kept on-chip. When batch_size * num_heads is large enough, this is certainly optimal. However, such a solution obviously runs into the second problem: for edge-side inference with small models and batch_size=1, or for large-model online deployments with TP where long-sequence inputs from coding agents etc. cannot launch a large enough batch for chunked prefill, the speedup of a fully-fused kernel over the original FLA implementation is limited. The earliest solution to the second problem comes from how DeltaNet does context parallelism, which splits a long sequence into multiple sub-sequences, uses S₀=0 to parallelize the recurrence, and then computes an additional M matrix to correct the recurrent results. This scheme was later optimized to insert a step before the recurrence kernel to compute the S₀ of each sub-sequence, and has now been merged into the FLA repository. For CP rank j, the specific preprocessing flow is: {S^∗_(j,i+1) ← γ_(j,i,C−1)S^∗_(j,i) + K_(j,i)^ᵀ diag((γ_(j,i,C−1))/γ_(j,i)) V’_(j,i); M_(j,i+1) ← ( γ_(j,i,C−1) I −K_(j,i)^ᵀ diag((γ_(j,i,C−1))/γ_(j,i)) W_(j,i) ) M_(j,i) S_(j,0) ← S^∗_(j,0) + M_(j,0)S_(j−1,0) However, this CP scheme also has its drawbacks: first, it introduces significant extra computation, with the time complexity of recurrently computing the M matrix even exceeding that of the S matrix; second, it does not work well with fully-fused kernels, because matrix inversion and other steps must be performed before the S₀ of each sub-sequence can be computed. A Balanced Solution: Fusing Kernels While Enabling Intra-Card CP Based on the two problems above, a compromise solution can be derived: split the GDN Chunked Prefill forward computation into two fused kernels, inserting CP-related preprocessing steps between them. After some transformations and simplifications, the following computation flow is obtained: Aᵢ ← (I+StrictLower( diag(βᵢ)KᵢKᵢ^ᵀ))⁻¹ CP Preprocess 2.1. {X_(j,i) ← −β_(j,i) A_(j,i)’^ᵀ K_(j,i); Y_(j,i) ← γ_(j,i,C−1) K_(j,i) S^∗_(j,i) − diag((γ_(j,i,C−1))/γ_(j,i)) V_(j,i); Z_(j,i) ← K_(j,i) M_(j,i); S^∗_(j,i+1) ← γ_(j,i,C−1)S^∗_(j,i) + X_(j,i)^ᵀ Y_(j,i); M_(j,i+1) ← γ_(j,i,C−1) ( M_(j,i) + X_(j,i)^ᵀ Z_(j,i) ) 2.2. S_(j,0) ← S^∗_(j,0) + M_(j,0)S_(j−1,0) {Vᵢ^Δ ← Vᵢ − diag(γᵢ)KᵢSᵢ; Vᵢ’ ← (Γᵢ ⊙ Aᵢ)diag(βᵢ)Vᵢ^Δ; Sᵢ₊₁ ← γ_(i,C−1)Sᵢ + Kᵢ^ᵀdiag((γ_(i,C−1))/γᵢ)Vᵢ’; Oᵢ ← diag(γᵢ)QᵢSᵢ + (Lower(Γᵢ) ⊙ QᵢKᵢ^ᵀ)Vᵢ’ We also designed a simple mathematical model to automatically determine the degree of parallelism. Let N be the number of chunks in a sequence and L be the number of chunks per CP rank. It is easy to see that the runtime of steps 2.1 and 3 is proportional to L, while the runtime of step 2.2 is proportional to N/L; therefore we can choose L=λ √N to minimize total time, where λ is a coefficient composed of batch_size, num_heads, and other hyperparameters. In production, intra-card CP is not always needed. Following the original FLA implementation, step 3 can also increase parallelism by 2-4× via splitting v_head_dim, at the cost of redundant memory access to Q and K. Based on measured data, we enable CP only when batch_size * num_heads <= 40 or batch_size * num_heads <= 56 && seq_len >= 8192. Further Optimization via Gate Decay Revisiting the GDN recurrence: Sᵢ₊₁ = αᵢSᵢ(I−βᵢkᵢkᵢ^ᵀ)+βᵢvᵢkᵢ^ᵀ For αᵢ∈(0,1), the influence of each Sᵢ on subsequent states decays exponentially, giving it a sliding-window property. For a sufficiently long window size of W, starting computation from S_(i−W)=0 can obtain the accurate Sᵢ, without the need to start from S₀. We refer to this process as warmup. On real data, we find that αᵢ is not constantly 1 on 60–80% of linear attention heads, and 6–8 chunks of warmup are sufficient to drive the Sᵢ error below the noise floor. Therefore, for linear attention heads with the sliding-window property, we can design a lighter CP preprocessing flow that discards the computation of the correction term M and directly obtains an equally accurate sub-sequence S₀ through warmup: C0 C1 C2 C3 C4 C5 C6 C7 C8 C9 C10 C11 C12 R1 O O O O O R2 X X O O O O R3 X X O O O O X denotes warming up with a zero initial state until the gate has decayed sufficiently, then writing out the S₀ of that CP rank; O denotes subsequent normal recurrent computation. The warmup length for each rank is determined by an independent kernel that collects gate statistics, and the cost of this step is negligible. TileLang Warp-Specialized Kernel We implement FlashQLA in TileLang using a warpgroup-specialization pattern: one producer warpgroup and three consumer warpgroups reside in the same SM, exchange data through shared memory, and synchronize via mbarriers. Forward In the forward pass, the three consumer warp groups compute V’, S, and O respectively, overlapping computation and memory traffic through a ping-pong structure. WG3 WG2 WG1 WG0 BAR 0 LDQ LDγ STO γ, γ_(C−1)γ⁻¹ TCP = Q K^ᵀ BAR 1 LDK LDβ STSᵢ TCU = K Sᵢ Γ = L(γ I γ⁻¹) A_γ = Γ ⊙ A P_γ = sΓ ⊙ P Sᵢ₊₁ = γ_(C−1) Sᵢ BAR 2 LDV W = β (V − γ U) TCO = Q Sᵢ BAR 3 LDA TCV^Δ = A_γ W O = sγ O BAR 4 V’ = γ_(C−1)γ⁻¹ V^Δ TCO = O + P_γ V^Δ BAR 5 TCSᵢ₊₁ = Sᵢ₊₁ + K^ᵀ V′ Notes: S output per chunk is for debugging only; normally only O and the last chunk’s S are output. CP Preprocessing As mentioned earlier, the CP preprocessing splits into two cases: the original approach (computing both M and S) and the sliding-window approach (computing only S). We designed a single fused kernel that handles both: WG3 WG2 WG1 WG0 BAR 0 LDK LDγ γ_(C−1)γ⁻¹ TCX = A^ᵀ K BAR 1 LDV LDβ STSᵢ TCU = K Sᵢ Y = −γ_(C−1)γ⁻¹ V + γ_(C−1) U X = −β X Sᵢ₊₁ = γ_(C−1) Sᵢ BAR 2 LDA γ^π = γ^π γ_(C−1) γ^π = γ^π γ_(C−1) TCSᵢ₊₁ = Sᵢ₊₁ + X^ᵀ Y BAR 3 TCZ^L = K M^L TCM^L = M^L + X^ᵀ Z^L TCZ^R = K M^R TCM^R = M^R + X^ᵀ Z^R Notes: The last two steps of WG1 and WG2 correspond to the M matrix computation and are triggered only when required. S is output per chunk only during backward recomputation. Backward In the backward pass, we reuse the CP preprocessing kernel from the previous section to recompute the S matrix, then fuse bwd_dv, bwd_dhu, bwd_dqkwg, bwd_wy into a single kernel with corresponding algebraic optimizations. Because of on-chip resource constraints, the backward kernel does not use multi-stage pipelining; instead it relies on the long compute chain to hide memory traffic. The full schedule is available in the FlashQLA repo. WG3 WG2 WG1 WG0 BAR 00 ST$dK$ TC P=QK^ᵀ γ, γ_(C−1)γ⁻¹ BAR 01 TC dV’=KdSᵢ₊₁ dV’=γ_(C−1)γ⁻¹dV′ Γ=γ I γ⁻¹ P_γ=sL(Γ)⊙ P dSᵢ = γ_(C−1) dSᵢ₊₁ BAR 02 TC dV’=dV’+P_γ^ᵀ dO A_β = A β A_γ = Γ ⊙ A_β BAR 03 TCU=KSᵢ BAR 04 TC dV=A_γ^ᵀ dV’ W=V−γ U dγ_(C−1)=∑ Sᵢ ⊙ dSᵢ₊₁ BAR 05 ST$dV$ LDV dV_γ = −γ dV dγ = ∑ᵢ dV_γ ⊙ U TC dA_γ = dV’Wᵀ TC V’=A_γ W BAR 06 TC dP_γ = dO V’^ᵀ BAR 07 LDK TC dK=V’dSᵢ₊₁^ᵀ dA_β = Γ ⊙ dA_γ dγ = dγ + ∑ᵢ dP_γ ⊙ L(P) dγ = dγ − ∑ⱼ dP_γ ⊙ L(P) dP = sL(Γ)⊙ dP_γ BAR 08 dK=γ_(C−1)γ⁻¹dK dγ_(C−1)=∑ K ⊙ dK dγ = −∑ᵢ K ⊙ dK TC dQ=dOSᵢᵀ BAR 09 LDQ TC dK=dK+dV_γ Sᵢ^ᵀ dQ=sγ dQ dγ = ∑ Q ⊙ dQ BAR 10 LDS TC dQ=dQ+dPK BAR 11 ST$dQ$ dγ = dγ + ∑ᵢ dA_β ⊙ A β dγ = dγ − ∑ⱼ dA_β ⊙ A β dβ = ∑ⱼ dA_β ⊙ A dA=dA_β β TC dSᵢ = dSᵢ + K^ᵀ dV_γ BAR 12 TC dK=dK+dP^ᵀ Q BAR 13 TC dA = −A^ᵀ dA A^ᵀ TC A_T = KK^ᵀ dO_γ=sγ dO BAR 14 LD$dO$ LDA dβ = dβ + ∑ᵢ dA ⊙ A_T dA_T = β dA dA_S = dA_T + dA_T^ᵀ TC dS₀ = dS₀ + Q^ᵀ dO_γ BAR 15 TC dK=dK+dA_S K Benchmark We benchmarked FlashQLA against the FLA Triton and FlashInfer baseline (FLA 0.5.0, Triton 3.5.1, FlashInfer 0.6.9, TileLang 0.1.8) on the head configurations used by the Qwen3.5 / Qwen3.6 family — hᵥ ∈ 64, 48, 32, 24, 16, 8, corresponding to TP1 through TP8. Specifically, the forward (FWD) benchmarks measure single-kernel latency for different models and TP settings under varying batch lengths, while the backward (BWD) benchmarks examine the relationship between total token count within a batch and latency during a single update step. Selected H200 single-layer forward results: Model / TP Seqlen h_(qk) hᵥ FlashQLA FlashInfer FLA vs FLA vs FI 397B/122B TP8 1x32768 2 8 0.310ms 1.653ms 0.913ms 2.95× 5.33× 397B/122B TP8 1x16384 2 8 0.184ms 0.833ms 0.465ms 2.53× 4.53× 397B/122B TP8 24576+8192 2 8 0.302ms 1.242ms 0.767ms 2.54× 4.11× 397B/122B TP4 1x32768 4 16 0.486ms 1.654ms 1.250ms 2.57× 3.40× 397B/122B TP4 1x16384 4 16 0.292ms 0.832ms 0.623ms 2.13× 2.85× 27B TP2 1x32768 8 24 0.659ms 1.616ms 1.564ms 2.37× 2.45× 2B/0.8B TP1 1x32768 16 16 0.493ms 1.640ms 1.285ms 2.60× 3.33× Sym h32 1x32768 32 32 0.877ms 1.554ms 1.952ms 2.23× 1.77× The speedup grows with TP degree because FlashQLA improves SM utilization via intra-card AutoCP in the exact regimes — TP sharding and small head number — where the baseline leaves SMs idle. Usage FlashQLA exposes both a high-level API matching FLA’s signature and low-level fwd/bwd entry points: import torch from qla import chunk_gated_delta_rule o, final_state = chunk_gated_delta_rule( q=q, # [B, T, H_q, K] k=k, # [B, T, H_q, K] v=v, # [B, T, H_v, V] g=g, # [B, T, H_v] beta=beta, # [B, T, H_v] scale=scale, initial_state=initial_state, # optional, [B, H_v, K, V] output_final_state=True, cu_seqlens=cu_seqlens, # optional, varlen support ) Requirements: SM90, CUDA 12.8+, PyTorch 2.8+. Install: git clone https://github.com/QwenLM/FlashQLA.git cd FlashQLA && pip install -v . Acknowledgments FlashQLA is inspired by Flash Linear Attention, FlashInfer and TileLang. We thank these communities for the reference implementations. Citation If FlashQLA is useful for your research, please cite: @misc{flashqla2026, title = {FlashQLA: Flash Qwen Linear Attention}, author = {Zhang, Chengruidong and Lin, Xi and Jiang, Huiqiang and Wang, Zekun and Li, Xiao and Cao, Yizhong and Zhuang, Bohan and Men, Rui and Zhang, Jianwei and Zheng, Bo and Lin, Junyang and Liu, Dayiheng and Zhou, Jingren}, year = {2026}, publisher = {GitHub}, howpublished = {\url{https://github.com/QwenLM/FlashQLA}} }