FlashQLA rebuilds gated delta rule attention as fused TileLang kernels

The scope is one operator, not a model

FlashQLA is described as a high-performance linear attention kernel library built on TileLang, and the scope is narrow in a useful way. It applies operator fusion and performance optimization to the forward and backward passes of GDN Chunked Prefill, where GDN is the gated delta rule formulation used by the Qwen3.5 and Qwen3.6 family of models.

The Python package exposes a single operator, chunk_gated_delta_rule, in both a fused high-level form and a split low-level form with separate forward and backward calls. The kernel code itself lives in the flash_qla/ directory, with benchmark scripts under benchmark/, profiling scripts under profile/, a tests/ directory and a tools/ directory excluded from the distributed package.

The distinction from a general attention library matters for adoption. This is not a drop-in torch.nn.functional replacement for scaled dot product attention, and it is not a model implementation. It is a faster path for one specific operator, so the value depends entirely on whether your workload runs that operator.

The claims are stated by the authors themselves. The README says FlashQLA achieves a 2-3x forward speedup and a 2x backward speedup over the FLA Triton kernel across multiple scenarios on NVIDIA Hopper and Blackwell, and says the gains are particularly pronounced in pretraining scenarios and edge-side agentic inference.

The hardware floor is SM90

The requirements are narrow and worth reading before anything else. FlashQLA needs SM90, SM100, SM103, SM120 or SM121, which means Hopper and Blackwell class datacentre parts, and nothing older. It needs CUDA 12.8 or above and PyTorch 2.8 or above.

That is a hard boundary rather than a preference. A team on A100s, on consumer Ada cards, on Apple silicon or on an older driver stack cannot run these kernels at all, and there is no fallback path described in the documentation. The kernel design depends on features that only exist in the newer architecture generations, including the warpgroup specialization that the README describes as manually implemented to overlap data movement, Tensor Core computation and CUDA Core computation.

The setup metadata narrows the software side further. The package declares python_requires >=3.10, install_requires of torch>=2.8 and tilelang>=0.1.12, and it is not zip safe because it ships a coefs/ directory of CSV files as package data.

A useful detail for anyone pinning versions in a production environment is the version suffix. setup.py reads the QLA_VERSION_SUFFIX environment variable, and if it is empty it appends a plus sign and the short git hash of the current commit to the package version, so a source build is identifiable by the commit it came from.

Installing from PyPI or building against a local TileLang

Installation is a pip install of a package named flash-qla.

pip install flash-qla

The alternative is a source build, which is what you need if you want to change the kernels rather than just call them.

git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .

The verbose flag is not decoration here. The package builds TileLang kernels, and the build step is where compilation problems surface. The project also carries a .pre-commit-config.yaml and a pytest.ini, and the Makefile-free layout means the test command is invoked directly through pytest.

The unit test is a single file.

python -m pytest tests/test_gdr_unit.py -v

That test_gdr_unit.py name is a good clue about what is actually covered. There is one unit test file for the gated delta rule operator, which is what the library provides. If you are evaluating it for numerical agreement with an existing implementation, that file is where the reference behaviour lives, and the repository does not describe a broader correctness suite.

The high-level call hides the context parallelism cache

The high-level API is one function call with the tensors the operator needs.

import torch
from flash_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, for variable-length sequences
)

The signature carries the gated delta rule inputs directly: query and key tensors with a head dimension H_q and key dimension K, value tensors with head dimension H_v and value dimension V, the gate g and the beta tensor, both shaped per head along the sequence axis. The gate is what makes the operator a delta rule rather than plain linear attention, and it is also what the parallelism strategy exploits.

Two optional inputs matter for serving. initial_state accepts a previous recurrent state of shape [B, H_v, K, V], and output_final_state returns that state so a decoding loop can carry it forward between steps. cu_seqlens handles variable-length sequences packed into a single batch, which is the shape a training pipeline produces when it concatenates samples.

The low-level API exists for the case where the forward and backward passes cannot be fused, and it is where the internal state becomes visible to the caller.

from flash_qla import chunk_gated_delta_rule_fwd, chunk_gated_delta_rule_bwd

# Forward
g, A, o, h, final_state, cp_cache = chunk_gated_delta_rule_fwd(
    q, k, v, g, beta, scale=scale, initial_state=h0, cu_seqlens=cu_seqlens
)

# Backward
dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
    q, k, v, g, beta, A, do, dht=dht, scale=scale, initial_state=h0,
    cu_seqlens=cu_seqlens, cp_cache=cp_cache,
)

The forward call returns a cp_cache that the backward call consumes, plus the A matrix used by the backward pass. That contract means the two calls cannot be split across processes without carrying that cache yourself, which is a real constraint for anyone planning a distributed backward pass.

Automatic intra-card context parallelism keyed on the gate

The first key feature is gate-driven automatic intra-card context parallelism. The mechanism is stated in one line: by exploiting the exponential decay property of the GDN gate, FlashQLA automatically enables intra-card CP under tensor parallelism, long-sequence and small-head-count settings, improving GPU SM utilization.

That is the interesting technical claim in the project. The gate decays exponentially, which means contributions from far-away tokens shrink predictably, so work can be partitioned across streaming multiprocessors within a card without waiting for a full sequence reduction. Under tensor parallelism, long sequences, or a small number of heads, the alternative is poor SM utilization, because there is not enough parallel work to fill the device.

The second feature is a hardware-friendly algebraic reformulation of the forward and backward flows, which reduces Tensor Core, CUDA Core and SFU overhead without sacrificing numerical precision.

The third explains the kernel structure. Rather than decomposing the computation into a sequence of independent kernels, and rather than fusing the entire flow into a single kernel, the project builds several key fused kernels in TileLang and manually implements warpgroup specialization to overlap data movement with Tensor Core and CUDA Core computation.

The v0.1.3 release extended this further, adding dual-CP inter/intra two-level context parallelism, a better auto-CP partitioning strategy and the SM120 backward path. The v0.1.1 release had already added intra-card sequence parallelism for the backward pass.

The benchmark numbers come from the authors' own script

The performance figures in the README are attributed to the authors, and the comparison setup is spelled out. The benchmark ran FlashQLA against the FLA Triton and FlashInfer baseline, with FLA 0.5.0, Triton 3.5.1, FlashInfer 0.6.9 and TileLang 0.1.8, on the head configurations used by the Qwen3.5 and Qwen3.6 family, with h_k,v in {64, 48, 32, 24, 16, 8}, corresponding to TP1 through TP8.

The two benchmarks measure different things. The forward benchmarks measure single-kernel latency for different models and TP settings under varying batch lengths, while the backward benchmarks examine the relationship between total token count within a batch and latency during a single update step. Raw results are published in benchmark_results_H200.txt and benchmark_results_GB200.txt.

The reproduction path is given, which is more than most projects offer.

# require flash linear attention and flashinfer for comparison
pip install flash_linear_attention==0.5.0 flashinfer-python==0.6.13

python benchmark/bench_gated_delta_rule.py

Note the version discrepancy worth checking before you rely on the numbers. The benchmark pins flashinfer-python 0.6.13 while the comparison baseline in the README is stated as FlashInfer 0.6.9, and the benchmark ran against TileLang 0.1.8 while the package now requires tilelang>=0.1.12. A team that reruns the script on current dependencies may not reproduce the published figures, and should compare against their own baseline rather than assume the ratio holds.

A PyTorch drop-in for the FLA gated delta rule backend

The most practical route into the project is through flash-linear-attention. The v0.1.2 release notes state that FlashQLA now serves as a backend for flash-linear-attention's GDN, providing plug-and-play acceleration through the standard FLA API.

That changes the integration shape considerably. Instead of writing against the flash_qla import, a team already using the FLA interface can keep its call sites and swap the backend. The v0.1.1 notes add that entry function signatures were aligned to the latest flash-linear-attention interface, which is what makes the swap transparent.

The profiling path is separate and exists to compare against the alternative. The README notes that flash linear attention is required for comparison, then gives two invocations, one of them with a --skip-bwd flag.

# require flash linear attention for comparison
pip install flash_linear_attention==0.5.0

python profile/profile_gdr.py --set develop
python profile/profile_gdr.py --set develop --skip-bwd

The licence is MIT, with the project released under the MIT License and the setup metadata recording the same. The repository also asks for a citation, with a BibTeX entry naming the FlashQLA paper and the author list, and it credits Flash Linear Attention, TileLang and FlashInfer as the projects it draws on.

Editorial conclusion

Use FlashQLA if you are training or serving a gated delta rule model on Hopper or Blackwell and already depend on the flash-linear-attention API, since the same interface can be swapped in for the GDN backend. Skip it if your GPUs are older than SM90, if you cannot install CUDA 12.8 and PyTorch 2.8, or if you need a published speedup that is independent of the authors' own benchmark script. Before adopting, run python benchmark/bench_gated_delta_rule.py against your own head configuration and check the tilelang version your environment resolves, since the package requires tilelang>=0.1.12 while the benchmark was run against TileLang 0.1.8.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. QwenLM/FlashQLA on GitHub
4. README
5. Releases