![\"\"](\"https://qianwen-res.oss-cn-beijing.aliyuncs.com/flashqla/flashqla.png\")
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.
\nToday 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.
\nKey highlights of this release:
\nGate-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.
\nHardware-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.
\nTileLang 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.
\nFlashQLA code and benchmarks are open-sourced at github.com/QwenLM/FlashQLA.
\nLet us first review the forward computation flow of GDN Chunked Prefill, taking chunk index i as an example:
\nIgnoring gate preprocessing and CP, each step of this flow corresponds to one kernel in FLA. This flow has two main efficiency problems:
\nchunk_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:
\nHowever, 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.
\nBased 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:
\nWe 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.
Revisiting the GDN recurrence:
\nSᵢ₊₁ = αᵢSᵢ(I−βᵢkᵢkᵢ^ᵀ)+βᵢvᵢkᵢ^ᵀ
\nFor αᵢ∈(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.
\nTherefore, 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:
\n| \n | C0 | \nC1 | \nC2 | \nC3 | \nC4 | \nC5 | \nC6 | \nC7 | \nC8 | \nC9 | \nC10 | \nC11 | \nC12 | \n
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| R1 | \nO | \nO | \nO | \nO | \nO | \n\n | \n | \n | \n | \n | \n | \n | \n |
| R2 | \n\n | \n | \n | X | \nX | \nO | \nO | \nO | \nO | \n\n | \n | \n | \n |
| R3 | \n\n | \n | \n | \n | \n | \n | \n | X | \nX | \nO | \nO | \nO | \nO | \n
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.
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.
\nIn the forward pass, the three consumer warp groups compute V’, S, and O respectively, overlapping computation and memory traffic through a ping-pong structure.
\n| \n | WG3 | \nWG2 | \nWG1 | \nWG0 | \n\n | \n |
|---|---|---|---|---|---|---|
| BAR 0 | \nLDQ | \nLDγ | \nSTO | \nγ, γ_(C−1)γ⁻¹ | \nTCP = Q K^ᵀ | \n\n |
| BAR 1 | \nLDK | \nLDβ | \nSTSᵢ | \nTCU = K Sᵢ | \nΓ = L(γ I γ⁻¹) A_γ = Γ ⊙ A P_γ = sΓ ⊙ P | \nSᵢ₊₁ = γ_(C−1) Sᵢ | \n
| BAR 2 | \nLDV | \n\n | W = β (V − γ U) | \nTCO = Q Sᵢ | \n\n | \n |
| BAR 3 | \nLDA | \n\n | TCV^Δ = A_γ W | \nO = sγ O | \n\n | \n |
| BAR 4 | \n\n | V’ = γ_(C−1)γ⁻¹ V^Δ | \nTCO = O + P_γ V^Δ | \n\n | \n | \n |
| BAR 5 | \n\n | \n | \n | TCSᵢ₊₁ = Sᵢ₊₁ + K^ᵀ V′ | \n\n | \n |
Notes:
\nAs 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:
\n| \n | WG3 | \nWG2 | \nWG1 | \nWG0 | \n\n | \n |
|---|---|---|---|---|---|---|
| BAR 0 | \nLDK | \nLDγ | \n\n | γ_(C−1)γ⁻¹ | \nTCX = A^ᵀ K | \n\n |
| BAR 1 | \nLDV | \nLDβ | \nSTSᵢ | \nTCU = K Sᵢ Y = −γ_(C−1)γ⁻¹ V + γ_(C−1) U | \nX = −β X | \nSᵢ₊₁ = γ_(C−1) Sᵢ | \n
| BAR 2 | \nLDA | \n\n | γ^π = γ^π γ_(C−1) | \nγ^π = γ^π γ_(C−1) | \nTCSᵢ₊₁ = Sᵢ₊₁ + X^ᵀ Y | \n\n |
| BAR 3 | \n\n | TCZ^L = K M^L TCM^L = M^L + X^ᵀ Z^L | \nTCZ^R = K M^R TCM^R = M^R + X^ᵀ Z^R | \n\n | \n | \n |
Notes:
\nIn 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.
| \n | WG3 | \nWG2 | \nWG1 | \nWG0 | \n\n |
|---|---|---|---|---|---|
| BAR 00 | \nST$dK$ | \n\n | \n | TC P=QK^ᵀ | \nγ, γ_(C−1)γ⁻¹ | \n
| BAR 01 | \n\n | TC dV’=KdSᵢ₊₁ dV’=γ_(C−1)γ⁻¹dV′ | \nΓ=γ I γ⁻¹ P_γ=sL(Γ)⊙ P | \ndSᵢ = γ_(C−1) dSᵢ₊₁ | \n\n |
| BAR 02 | \n\n | TC dV’=dV’+P_γ^ᵀ dO | \nA_β = A β A_γ = Γ ⊙ A_β | \n\n | \n |
| BAR 03 | \n\n | \n | TCU=KSᵢ | \n\n | \n |
| BAR 04 | \n\n | TC dV=A_γ^ᵀ dV’ | \nW=V−γ U | \ndγ_(C−1)=∑ Sᵢ ⊙ dSᵢ₊₁ | \n\n |
| BAR 05 | \nST$dV$ | \nLDV | \ndV_γ = −γ dV dγ = ∑ᵢ dV_γ ⊙ U | \nTC dA_γ = dV’Wᵀ TC V’=A_γ W | \n\n |
| BAR 06 | \n\n | \n | TC dP_γ = dO V’^ᵀ | \n\n | \n |
| BAR 07 | \n\n | LDK | \nTC dK=V’dSᵢ₊₁^ᵀ | \ndA_β = Γ ⊙ dA_γ dγ = dγ + ∑ᵢ dP_γ ⊙ L(P) dγ = dγ − ∑ⱼ dP_γ ⊙ L(P) dP = sL(Γ)⊙ dP_γ | \n\n |
| BAR 08 | \n\n | dK=γ_(C−1)γ⁻¹dK dγ_(C−1)=∑ K ⊙ dK dγ = −∑ᵢ K ⊙ dK | \nTC dQ=dOSᵢᵀ | \n\n | \n |
| BAR 09 | \n\n | LDQ | \nTC dK=dK+dV_γ Sᵢ^ᵀ | \ndQ=sγ dQ dγ = ∑ Q ⊙ dQ | \n\n |
| BAR 10 | \n\n | LDS | \n\n | TC dQ=dQ+dPK | \n\n |
| BAR 11 | \nST$dQ$ | \n\n | \n | dγ = dγ + ∑ᵢ dA_β ⊙ A β dγ = dγ − ∑ⱼ dA_β ⊙ A β dβ = ∑ⱼ dA_β ⊙ A dA=dA_β β | \nTC dSᵢ = dSᵢ + K^ᵀ dV_γ | \n
| BAR 12 | \n\n | TC dK=dK+dP^ᵀ Q | \n\n | \n | \n |
| BAR 13 | \n\n | \n | TC dA = −A^ᵀ dA A^ᵀ TC A_T = KK^ᵀ | \ndO_γ=sγ dO | \n\n |
| BAR 14 | \n\n | LD$dO$ LDA | \n\n | dβ = dβ + ∑ᵢ dA ⊙ A_T dA_T = β dA dA_S = dA_T + dA_T^ᵀ | \nTC dS₀ = dS₀ + Q^ᵀ dO_γ | \n
| BAR 15 | \n\n | TC dK=dK+dA_S K | \n\n | \n | \n |
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.
\n![\"FlashQLA](\"https://qianwen-res.oss-cn-beijing.aliyuncs.com/flashqla/fwd_bwd_latency_comparison.png\")
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.
\nSelected H200 single-layer forward results:
\n| Model / TP | \nSeqlen | \nh_(qk) | \nhᵥ | \nFlashQLA | \nFlashInfer | \nFLA | \nvs FLA | \nvs FI | \n
|---|---|---|---|---|---|---|---|---|
| 397B/122B TP8 | \n1x32768 | \n2 | \n8 | \n0.310ms | \n1.653ms | \n0.913ms | \n2.95× | \n5.33× | \n
| 397B/122B TP8 | \n1x16384 | \n2 | \n8 | \n0.184ms | \n0.833ms | \n0.465ms | \n2.53× | \n4.53× | \n
| 397B/122B TP8 | \n24576+8192 | \n2 | \n8 | \n0.302ms | \n1.242ms | \n0.767ms | \n2.54× | \n4.11× | \n
| 397B/122B TP4 | \n1x32768 | \n4 | \n16 | \n0.486ms | \n1.654ms | \n1.250ms | \n2.57× | \n3.40× | \n
| 397B/122B TP4 | \n1x16384 | \n4 | \n16 | \n0.292ms | \n0.832ms | \n0.623ms | \n2.13× | \n2.85× | \n
| 27B TP2 | \n1x32768 | \n8 | \n24 | \n0.659ms | \n1.616ms | \n1.564ms | \n2.37× | \n2.45× | \n
| 2B/0.8B TP1 | \n1x32768 | \n16 | \n16 | \n0.493ms | \n1.640ms | \n1.285ms | \n2.60× | \n3.33× | \n
| Sym h32 | \n1x32768 | \n32 | \n32 | \n0.877ms | \n1.554ms | \n1.952ms | \n2.23× | \n1.77× | \n
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.
\nFlashQLA exposes both a high-level API matching FLA’s signature and low-level fwd/bwd entry points:
\n\n\n\n, = \nRequirements: SM90, CUDA 12.8+, PyTorch 2.8+. Install:
\n\n && \nFlashQLA is inspired by Flash Linear Attention, FlashInfer and TileLang. We thank these communities for the reference implementations.
\nIf FlashQLA is useful for your research, please cite:
\n\n