Add Hopper (SM90) GDN fused-recurrent decode + SGLang verify kernels (TileLang)

[rifkybujana]

Summary

Adds fused, warp-specialized TileLang kernels for the Gated Delta Rule (GDN) fused-recurrent path on NVIDIA Hopper (SM90) — the decode forward and the SGLang speculative-decode verify path — as a drop-in faster alternative to the FLA Triton recurrent/verify kernels. Complements the existing chunked-prefill kernels with the single-step / draft-verify regime used by speculative decoding.

Scope: this PR is the kernel library + correctness tests only (no benchmarks, design docs, or feasibility probes).

What's new

• Decode (recurrent_gated_delta_rule / fused_recurrent_gdr_fwd): gemm-free GDN recurrence — state kept [block_DV, DK] fp32 in registers, the two GEMVs via T.reduce_sum over K, rank-1 via a T.Parallel outer product (no tensor-core M-padding constraint). GQA, ragged seqlens, optional initial/final state. On the final-state path block_DV=128 coalesces the K-major transposed store (~2× at all batch sizes; block_DV<128 is uncoalesced).
• Verify (recurrent_gated_delta_rule_verify): paged V-major bf16 state-pool gather/scatter (state_indices, -1 = skip), per-token intermediate states, no-commit, varlen cu_seqlens, host- and in-kernel fused gating (g = -exp(A_log)·softplus(a+dt_bias), β = sigmoid(b), qk-l2norm), CUDA-graph safe. A regime-gated gating+l2norm dedup pre-pass gives +8–24% at large batch (1.13–1.18× vs FLA at T=12, the DFlash draft length); single-request stays on the fast in-kernel path (~3.3× vs FLA's num_warps=1 floor).

Correctness

67 tests (tests/test_{decode,verify,prepass,head_batch}_gdr.py) against a reference recurrence and the FLA kernel — GQA, ragged batches, host/in-kernel gating, the pre-pass path, CUDA-graph capture/replay, and negative controls (step order, GQA mapping, V-major transpose). The accuracy check runs the kernel 1000× asserting ≤2% drift (a deliberate race/nondeterminism detector).

Requirements

SM90+, CUDA 12.8+, PyTorch 2.8+, tilelang==0.1.8. Dispatch hard-gates on cc == "9.0".

🤖 Generated with Claude Code

[minatoyukinaa]

FlashQLA is designed for prefill, why decode kernel is added here?

[rifkybujana]

@minatoyukinaa We've been working on our fork of SGLang. The speculative-decode verify step runs the GDN recurrence over the draft tokens and needs per-token intermediate states, which the chunk/prefill kernel doesn't produce, so we added the recurrent decode/verify kernel to cover that path instead of falling back to FLA's Triton. In our H100 benchmarks it comes out ~1.13-1.18x faster than FLA's verify at draft length 12 (large batch) and ~3.3x in the single-request latency-bound case.

[Starmys]

Thank you for your contribution! Would you mind providing the benchmark script at your convenience?

[rifkybujana]

Thanks @Starmys ! Sure thing.

benchmark/bench_recurrent_gdr.py

single-request / TP (N=1, T=12):    3.6–3.9× faster
small batch         (N=8, T=12):    1.3× faster
large batch         (N≥64, T≥4):    ~parity (HBM-bound, both ~2.2 TB/s)

I've also got an apples-to-apples vs-FLA script can push that too if it's more useful. Let me know!