TileLang kernels for Gated DeltaNet-2 inference and training.
Prefill · cached decode · training forward and backward
Install · Quick start · API · Roadmap · Cite
FlashGDN-2 accelerates Gated DeltaNet-2 with independent channel-wise decay, erase and write gates. It provides TileLang kernels, PyTorch autograd integration, a differentiable reference, and reproducible comparisons against FLA's GDN-2 implementation.
Initial release: measured and tuned on NVIDIA RTX 5080 (SM120, 84 SMs), using BF16/FP16 projections and FP32 state. This is a single-GPU research release; broader GPU support and distributed execution are planned for future releases. The technical report will follow. Until then, please cite this repository.
Against standard FLA GDN-2, on the measured RTX 5080 BF16/D128 workloads:
- Inference prefill: 1.42–2.02× faster.
- Training: 1.10–1.64× forward, 2.07–2.57× backward, and 1.78–2.27× for forward + backward together.
- Functional decode: 0.99–2.13× across the measured batches and head counts. The H32/batch-16 case remains slightly slower than FLA.
These are steady-state operator measurements, not end-to-end model throughput. Functional comparisons include output and final state. Training uses FLA's default recomputation policy. Peak allocation excludes resident inputs, allocator reservation, graph pools, model weights and optimizer state; timing excludes JIT compilation and Python overhead. Methodology and exact results.
For output-only serving, prepare_replay_decode reaches 1.38–1.78× prepared
FLA at measured H32/D128 batches 4–128, including periodic state writes. It uses
additional history memory; batch 1 favors ordinary prepare_decode.
Serving APIs and tradeoffs.
Numerical accuracy and validation
FP64 is the untimed oracle; the accelerated kernels use BF16/FP16 and FP32 state. Tests compare identical quantized inputs, check output and state, and audit all seven training gradients. Decode audits carry each provider's state through 128 changing tokens. Fixed error policies apply equally to FlashGDN-2 and FLA.
The GPU-validated source snapshot passed 2,255 tests, with two multi-GPU skips and six documented numerical expected failures. Extreme-decay and BF16 model-gradient limitations remain; operator tests do not establish model convergence. Precision policies · Training limitations · Validation record.
Tested on Linux with PyTorch 2.9.1+cu128, TileLang 0.1.9,
apache-tvm-ffi 0.1.9, and nvcc 13.2. Install a compatible NVIDIA driver and
CUDA toolkit, with nvcc discoverable through CUDA_HOME or PATH.
git clone https://github.com/Bezdarnost/Flash-GDN-2.git
cd Flash-GDN-2
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu128
python -m pip install '.[cuda]'
python -m scripts.check_environmentFor development, install -e '.[cuda,dev,plots]'. The CPU reference needs only
PyTorch and the base package; package import does not initialize CUDA or import
TileLang. Installation checks.
Prefill a prompt, then decode the next token using its final state:
import torch
from flash_gdn2 import chunk_gdn2, fused_recurrent_gdn2
B, T, H, D = 1, 1024, 16, 128
shape = (B, T + 1, H, D)
q, k = (
torch.nn.functional.normalize(torch.randn(shape, device="cuda"), dim=-1).bfloat16()
for _ in range(2)
)
v = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
g = -torch.rand(shape, device="cuda", dtype=torch.float32) * 0.1 # log-decay
b, w = torch.rand_like(k), torch.rand_like(v) # activated erase/write gates
inputs = (q, k, v, g, b, w)
with torch.inference_mode():
output, state = chunk_gdn2(
*(x[:, :T] for x in inputs),
backend="tilelang", output_final_state=True,
)
next_output, state = fused_recurrent_gdn2(
*(x[:, T:] for x in inputs),
initial_state=state, backend="tilelang", output_final_state=True,
)The operators consume projected Q/K/V and gates. For fixed-length batches,
q,k=[B,T,H,K], v,w=[B,T,HV,V], g,b=[B,T,HV,K], and FP32 state
[B,HV,K,V]. Recurrent inference supports grouped value heads with HV % H == 0;
chunk prefill requires HV=H. Functional calls do not mutate inputs or state.
See the API for supported shapes, layouts and gate conventions.
Choose the mixed-precision training API explicitly:
from flash_gdn2 import mixed_chunk_gdn2
q, k, v, g, b, w = (x.detach().requires_grad_() for x in inputs)
output, final_state = mixed_chunk_gdn2(q, k, v, g, b, w, output_final_state=True)
loss = output.float().square().mean() + final_state.square().mean()
loss.backward()Mixed training supports first-order gradients, FP16/BF16 Q/K/V, half/FP32 gates,
K/V in {64, 128}, equal heads and FP32 initial state on SM120. The stricter
training path is fused_recurrent_gdn2(..., backend="tilelang");
backend="auto" uses the reference when gradients are needed.
Training guide · NVIDIA layer integration.
I plan to extend FlashGDN-2 with:
- GPU-specific tuning and validation for H100/H200, datacenter Blackwell B200/B300/GB200/GB300, and additional RTX 20/30/40/50-series devices.
- Tensor parallelism and context parallelism, broader training shapes, and further prefill, decode and backward optimizations.
- Longer model-training and serving evaluations, including end-to-end throughput and quality checks. AMD GPUs and Apple Silicon are longer-term research targets.
These are planned targets, not current compatibility claims. TileLang does not make a tuned schedule portable automatically. The roadmap describes the validation required for each extension. A dedicated technical report will follow this initial code release; implementation notes and reproducibility records are in other/.
python -m pytest -m 'not cuda' -q # CPU tests
python -m pytest -q # GPU tests require CUDA and TileLang
ruff check flash_gdn2 tests benchmarks scripts examples
python -m buildContributing · Benchmarks and figure reproduction · Support details. FlashQLA and FlashKDA compute different operators; their separate comparisons are not GDN-2 speedups.
Please cite the repository while the technical report is in preparation. CITATION.cff provides metadata for GitHub's Cite this repository feature.
@misc{aman2026flashgdn2,
author = {Aman},
title = {{FlashGDN-2}: TileLang Kernels for Gated DeltaNet-2},
year = {2026},
url = {https://github.com/Bezdarnost/Flash-GDN-2},
note = {Version 0.1.0}
}Please also credit the underlying Gated DeltaNet-2 paper and the relevant upstream work listed in NOTICE.md.
FlashGDN-2 uses the NVIDIA Source Code License-NC, which restricts use to non-commercial research and evaluation. This is a source-available research release, not a permissively licensed open-source package. FLA and NVIDIA KDA attribution notices are retained. All runtime kernels are TileLang; reference repositories are optional validation dependencies.