FlashQLA - A high-performance linear attention operator library open-sourced by Tongyi Labs

What is FlashQLA?

FlashQLA is an open-source, high-performance linear attention operator library based on TileLang, developed by Tongyi Labs. Through operator fusion, gate-driven in-card sequence parallelism, and Warp-Specialized optimization, FlashQLA achieves 2–3× forward and 2× backward speedups compared to FLA Triton on Hopper, covering models from 2B to 397B, and improving pre-training and edge inference efficiency. FlashQLA requires an SM90, CUDA 12.8+, and PyTorch 2.8+ environment.

Main functions of FlashQLA

• High-performance linear attention operator libraryDeep optimization for the attention layer of the entire Qwen Gated Delta Network (GDN) series.
• Operator fusion accelerationThe forward and reverse processes of GDN Chunked Prefill are reasonably fused and their performance optimized.
• Full specification model coverageSupports multiple specifications from 2B to 397B, covering TP1 to TP8 scenarios.
• Two-level API interfaceProvides a high-level API for aligning FLA signatures, as well as underlying fwd/bwd entry points.
• Variable-length sequence supportBuilt-in varlen variable-length sequence processing capability, adapting to real training and inference data distributions.

FlashQLA Technical Principles

• TileLang Warp-Specialized KernelThe key fused kernel is built based on TileLang, and warpgroup specialization is used to achieve the overlap of data transfer, Tensor Core computation and CUDA Core computation.
• Automated In-Card Sequence Parallelism (AutoCP)By leveraging the exponential decay property of the GDN gate, in-card sequence parallelism is automatically enabled in scenarios such as TP, long sequences, and small head counts, thereby improving GPU SM utilization.
• Sliding window warmup mechanismFor linear attention heads with decay properties, the initial state of the subsequence can be accurately obtained with only 6–8 chunks of warmup, discarding the calculation of the correction matrix M.
• Hardware-friendly algebra rewriting: Perform algebraic transformations on the forward and reverse processes of GDN Chunked Prefill to effectively reduce the overhead of Tensor Core, CUDA Core and SFU without affecting numerical accuracy.
• A compromise architecture that balances memory access and parallelismThe computation process is split into two fused kernels and CP preprocessing is inserted in between to avoid the problem of low GPU utilization in small batch/TP scenarios of fully-fused kernels.

How to use FlashQLA

• Environmental inspectionThe hardware was confirmed to be NVIDIA SM90 (Hopper architecture) and the software environment met the requirements of CUDA 12.8+ and PyTorch 2.8+.
• Installation and DeploymentClone the FlashQLA repository from GitHub and compile and install it using pip.
• Module importImporting in Python chunk_gated_delta_rule function.
• Data preparationPrepare the input tensors q, k, v and the gate parameters g and beta, and ensure that the shape of each tensor meets the interface requirements.
• Perform calculation: call chunk_gated_delta_rule It also passes in the corresponding parameters to obtain the output result O and the final state.
• Advanced configurationTo process variable-length sequences, you can pass in... cu_seqlens Parameters; if status continuation is required, you can pass in the following parameters. initial_state.
• Automatic optimizationAutoCP sequence parallelism will be automatically triggered based on batch size and sequence length, without the need for manual configuration.

Key information and usage requirements for FlashQLA

• PublisherQwenTeam
• Open source addressgithub.com/QwenLM/FlashQLA
• Hardware RequirementsNVIDIA SM90 (Hopper architecture, such as H200)
• Software RequirementsCUDA 12.8+, PyTorch 2.8+
• Support ModelQwen3.5 / Qwen3.6 series (head dim covers 64 to 8, corresponding to TP1 to TP8)
• acceleration effectForward 2–3×, Backward 2× (compared to FLA Triton Kernel)

FlashQLA's core advantages

• A compromise architecture that balances memory access and parallelismThe computation is split into two fused kernels with CP preprocessing inserted in between. This avoids the problem of low GPU utilization in small batch/TP scenarios of fully-fused kernels. By reasonably splitting the kernels, the memory access overhead of HBM repeatedly reading and writing intermediate variables is reduced.
• AutoCP automatic activation mechanismOnly in batch_size × num_heads ≤ 40 or batch_size × num_heads ≤ 56 且 seq_len ≥ 8192 It automatically triggers in-card sequence parallelism to avoid unnecessary redundant calculations and adaptively balances parallelism with memory access costs.
• Sliding window warmup mechanismUsing the exponential decay property of the GDN gate, only 6–8 chunks of warmup are needed to accurately obtain the initial state of the subsequence for 60–80% of linear attention heads, directly discarding the calculation of the correction matrix M, which greatly reduces the CP preprocessing overhead.
• Warp-Specialized calculates overlapBased on the TileLang warpgroup specialization design, it enables producer and consumer warpgroup collaboration within the same SM, and uses a ping-pong structure to cover data transfer and Tensor Core/CUDA Core computation.
• Hardware-friendly algebra rewritingIt performs algebraic transformations and simplifications on the forward and reverse processes, effectively reducing the hardware overhead of Tensor Cores, CUDA Cores, and SFUs without affecting numerical accuracy.

FlashQLA project address

• Project official websitehttps://qwen.ai/blog?id=flashqla
• GitHub repositoryhttps://github.com/QwenLM/FlashQLA

FlashQLA's Competitive Product Comparison

Application scenarios of FlashQLA

• Large Model Pre-trainingCovers the full range of Qwen models, including 397B, 122B, 35B, and 27B, supports training with 256K long contexts, and significantly reduces the computational and time overhead of the attention layer in end-to-end training.
• End-side agentic inferenceFor chunked prefill scenarios with batch size=1 and small-sized models (e.g., 2B / 0.8B), AutoCP improves GPU utilization under small batch size and accelerates real-time response of the edge agent.
• Large-scale online deploymentIn Tensor Parallelism (TP) scenarios, this addresses the GPU utilization bottleneck caused by insufficient batch size due to chunked prefill when processing long input sequences from coding agents, thereby improving service throughput.
• General GDN / Linear Attention Architecture AccelerationApplicable to any LLM training and inference based on Gated Delta Network or Linear Attention architecture, providing a high-performance out-of-the-box operator replacement solution.