Fused recurrent GDN decode + SGLang DFlash verify kernel (TileLang, SM90)#1
Merged
Conversation
Design for a forward/inference-only fused recurrent decode kernel for GDN, sibling to the chunked-prefill kernels: - exact single-token recurrence (post-update read), no kkt_solve/A/cumsum - single-role threads=256, register-resident fp32 state, gemm_v1 contractions - V-column-split occupancy (decode analogue of auto-CP); K-split rejected (breaks fusion) - fused multi-step (q_len 1..D) + ragged per-sequence accepted length - host-side l2norm; flash_qla-native low+high level API - GQA-group head-batched server variant as a jit-factory specialization - test/benchmark plan + prototype-gated feasibility items Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ed contract New verify-first SGLang integration spec (sibling to the decode spec): - shared infra A: in-kernel paged gather/scatter (state_indices), bf16 pool + fp32-register accumulation, CUDA-graph-safe no-alloc/no-sync entry, state_v_first=True (V-major) SGLang layout, host-side gating primary (in-kernel fused gating prototype-gated on TileLang log2/rsqrt) - verify V1 (recurrent, per-token intermediates native, no-commit) + V2 (chunk-o + honest per-token-state extraction cost) + T=12 benchmark-to-decide - SGLang integration contract, API signatures, bit-identity test plan, feasibility gates, confirm-with-user items Decode-spec fixes (grounded against the repo): - rank-1 update is the transpose_A gemm-into-fragment (fused_fwd.py:204), NOT the scalar-broadcast FMA at :197; the two-vector FMA is prototype-gated (§11.H) - every step is M=1 (root gemm gate); ragged contract resolved to cu_seqlens [N+1] flattened varlen for verify; D up to 12 exceeds the 1..8 design point Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…face User confirmed: DFlash (public upstream linear-chain speculative verify) is the target; DDTree (tree variant) is not needed. Remove the tree seam / retrieve_parent_token from the contract + API signatures, pin DFlash as the target scheme, and add a DFlash verify-entry parity confirm-item. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…loyment facts) - intermediate_state_indices is dense arange (never -1); gate the per-token ibuf write on the POOL slot mask (state_indices[b]>=0), not on cache_slot>=0 (which never fires and would leave garbage in padded ibuf rows) - "host-side" gating/l2norm = PyTorch (not TileLang), runs INSIDE SGLang's captured graph (capture-safe), NOT hoisted outside capture - scope boundary: kernel owns only the GDN SSM state; conv state (causal_conv1d_update) runs upstream in gdn_backend.forward_extend - resolve deployment facts: bf16 pool; intermediate buffer [num_layers, num_slots+1, 12, HV, V, K] single-layer slice (num_cache_slots = num_slots+1, cache_steps=12); PAD_SLOT_ID=-1 / slot>=0 guard / slot 0 valid; DFlash width-1 (TOPK=1), wrapper accepts-and-ignores retrieve_parent_token - fix §10-Gate cross-refs to §9; FLA-reference test ref to §8 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… plan Phase 0 (4 probes: TileLang prims, gemm_v1 M=1, runtime T.serial(L), state_v_first transpose store) + Phase 1 (decode_recur torch reference pinned to chunk path; fused_recurrent package scaffold; single-role core kernel body; low-level wrapper; kernel-vs-reference sweep incl ragged/GQA/g=0/low-occupancy). Defers infra A / V1 / V2 to follow-up plans informed by the gate outcomes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…alidated on H100)
Core single-role decode kernel + flash_qla-native API, validated on an H100:
- recurrence per (sequence, V-head, V-column-tile): decay -> kS -> v_new -> rank-1
-> o, post-update read; raw g (no cumsum); fp32 state, bf16 q/k/v/o
- GEMM-FREE design: state kept [block_DV, DK] (V-major) in an fp32 fragment; the two
GEMVs are reduce_sum over the last (DK) dim, the rank-1 is a T.Parallel outer product.
Avoids tensor-core M=16 padding + warp-partition/fragment-layout friction entirely
and keeps the state fp32 throughout (more accurate than a bf16-operand gemm path).
- V-column-split occupancy ladder (block_DV in {128,64,32} from B*H vs 0.7*SM); threads
scaled to the tile; multistep q_len=1..D via T.serial(L); ragged per-CTA L=seqlens[b]
- new decode_recur torch reference (pinned to the chunk path at chunk_size=64) +
19-case pytest suite: D in {1,8}, GQA {1:1,1:4,MQA}, h0 on/off, g=0 SWA, ragged,
low- and high-occupancy (block_DV 32/64/128). All pass at 0.02 rel.
- Phase-0 feasibility probes under tests/probes/ documenting the gate outcomes
(gemm_v1 M=1 fails; TileLang log2/rsqrt/exp lower on SM90; reduce_sum(dim=-1) works).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ated on H100) Gemm-free GDN verify kernel + low-level graph-safe entry, validated on an H100: - in-kernel paged gather/scatter via state_indices (slot<0 => skip); V-major bf16 pool [num_slots,H,V,K] -- direct (no transpose), natural for the [block_DV,DK] state - per-token intermediate states written to a caller buffer [num_cache_slots,steps,H,V,K], gated by the POOL slot mask (intermediate_state_indices is dense, never -1) - no-commit mode (disable_state_update): skip the final pool scatter (verify); commit path writes the final state back for decode - varlen cu_seqlens [N+1] (B=1 flattened): per-request token range + runtime T.serial(L) - fp32-accum state, bf16 pool, host-side gating (g/beta pre-activated, q/k pre-l2normed) - verify_ref torch reference + 7-case suite: no-commit/commit, ragged, GQA, V-major distinct-gate (transpose guard), -1 slot skip. All pass at 0.02 rel. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…raph safety - recurrent_gated_delta_rule_verify: DFlash-style entry doing host-side (PyTorch, capture-safe) sigmoid gating g=-exp(A_log)*softplus(a+dt_bias), beta=sigmoid(b) (allow_neg_eigval flag) + qk-l2norm, then the paged V-major verify kernel - gdn_sigmoid_gate helper; export verify entry points at the package + top level - tests (pass on H100): wrapper gating matches a hand-built reference; the low-level entry is CUDA-graph capturable and graph-replay == eager Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… validated on H100) - fused_recurrent_gdr_verify_gated_fwd: takes raw (a,b,A_log,dt_bias) and computes g=-exp(A_log)*softplus(a+dt_bias), beta=sigmoid(b), and qk-l2norm INSIDE the kernel (TileLang log/log1p/exp/rsqrt/sigmoid all lower on SM90 per Gate 4). One fewer launch / no gating intermediate tensors. allow_neg_eigval flag (beta x2). - recurrent_gated_delta_rule_verify gains fuse_gating (default False=host gating). - test: in-kernel gating matches the host-gating reference at <=0.03 rel (GQA 1:1, 1:4). Full suite: 30 passed (19 decode + 11 verify). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… notes - benchmark/bench_recurrent_gdr.py: wall time + achieved HBM bandwidth across regimes. H100 results: server (N>=64,H=32) bandwidth-bound at 60-67% peak; single-request TP8 latency-bound (~6%, as predicted). Confirms V2 cannot beat V1 (same per-token writes). - spec status notes: record the gemm-free pivot, the H100 gate outcomes, what's built and validated (V1 verify + in-kernel gating + graph safety), and that the head-batched variant and V2 are intentionally not built. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tract) Review (architecture/quality/test lens) findings acted on: - IMPORTANT: add negative-control tests that prove the suite discriminates wrong answers -- decode step-order (pre- vs post-update read) and GQA mapping (h//grp vs h%Hk), and verify V-major transpose (ibuf must differ from a K/V-swapped reference). decode_recur gains read_pre_update/gqa_mod flags for this. - IMPORTANT: capture-safety -- the host-gating path uses @torch.compile l2norm + allocs (not capture-safe). Document that fuse_gating=True (in-kernel g/beta + l2norm, no PyTorch ops) is the CUDA-graph-safe entry, and add a graph capture+replay test over the gated path. - MINOR: verify wrapper gains retrieve_parent_token=None (accept-and-ignore, DFlash width-1) and a static cache_steps shape assert; recurrent_gated_delta_rule output_final_state defaults True (spec-aligned, sensible for a stateful decode op). - 34 tests pass on H100 (was 30). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Apples-to-apples vs netra-server's vendored FLA fused_sigmoid_gating_delta_rule_update (target_verify mode): in-kernel gating + l2norm, bf16 V-major pool, per-token intermediate states, no-commit, retrieve_parent_token=None. Parity gate (o + ibuf within bf16 tol) green. H100 results: FlashQLA wins the latency-bound regime decisively -- single-request/TP (batch-1 speculative verify) 2.7-3.3x, small-batch up to 3.5x (FLA has a ~53us fixed- overhead floor, num_warps=1; FlashQLA ~15us). FLA edges ahead ~15% at large batch (N>=64) where both are bandwidth-bound (FLA ~2.0-2.15 TB/s vs FlashQLA ~1.8). Loads the FLA kernel by file path to skip the heavy sglang/__init__ import chain. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…(H100) Memory-bound: occupancy beats bigger tiles. Replaced the block_DV ladder (128/64/32 @ threads=block_DV*2) with the autotuned sweet spot block_DV=64 (2 V-tiles) @ threads=128, falling to 32 for the low-CTA tail. block_DV=128 was occupancy-starved (64 fp32/thread state -> low occupancy) and is dropped; threads>128 over-subscribes. H100 vs-FLA after tuning (was FlashQLA 0.81-0.88x at large batch): large batch N>=64: now 0.99-1.02x (PAR/ahead), 2.0-2.18 TB/s (was 1.8) small batch N=8 T=4: 1.80x -> 2.57x single-request/TP: held at ~3.3x faster So FlashQLA is now >= FLA across all regimes. benchmark/tune_verify.py records the sweep. 34 tests pass (no regression). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Tidy the three superpowers design docs to record the as-built reality without losing the design-exploration rationale: - decode spec: §3 gemm-free pivot, §7 head-batch not built, §11 gates all resolved on H100; status note records the autotuned block_DV=64/threads=128 tile and the large-batch compute-bound finding. - verify spec: status note records the large-batch optimization investigation (compute-bound, not write-bound; fp8 1.18x/+5% err and replay-tape 0.66x both rejected; 2-3x at large batch physically impossible) + PR #1. - plan: as-built banner (gemm-free, threads=128, gates resolved, V1 shipped). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Implements section 7 of the decode spec, adapted to the as-built gemm-free
path. One CTA owns (sequence, K/Q head-group hg, V-tile) and processes all
grp = H//Hg V-heads h = hg*grp+i that share K/Q head hg, loading q/k once.
State is row-stacked S[grp*block_DV, DK]; row gv -> head-band gv//block_DV,
channel v0 + gv%block_DV.
- tilelang_fused_recurrent_gdr_fwd_hb + dispatch; forceable `head_batch` flag
threaded through recurrent_gated_delta_rule (auto OFF, restricted grp in {2,4},
threads = grp*128 <= 512, block_DV from the post-collapse B*Hg grid).
- tests/test_head_batch_gdr.py: 21 tests (grp 2/4, block_DV 32/64, use_h0,
ragged, equals-per-head, rejects grp not in {2,4}) + a within-group band-swap
negative control (new decode_recur band_perm flag) that gqa_mod can't express.
- benchmark/bench_head_batch.py.
Measured on H100: NOT a win. grp=2 = 0.98-0.99x (neutral), grp=4 = 0.74-0.87x
(regression) -- the row-stack trades CTA count (B*H -> B*Hg) and uses bigger
512-thread CTAs for only the q/k LOAD dedup (sub-1% on this memory-bound kernel).
So head_batch defaults OFF and is exposed only for experimentation.
TileLang limits learned (recorded in the spec sec 7): every [M] fragment must be
accessed over the FULL Parallel(M,...) range (partial-offset writes break
InverseAffineIterMap); and a shared [grp] band read by gv//block_DV does not
lower -- so per-head gating can't be deduped, which is why the gated-verify
head-batch was not built (it cannot overcome the row-stack penalty).
55 tests pass (head-batch + decode + verify) on H100.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ify) The in-kernel-gated verify kernel recomputes g/beta + qk-l2norm inside the per-token hot loop, once per (token, V-head, V-tile) -> grp*n_vt-redundant. A verify-first probe (benchmark/probe_h1_ceiling.py) measured this in-loop cost at 14-16% of runtime at T=12 large batch; the naive torch pre-pass is ~145us (l2norm @torch.compile + allocs), so the dedup must be a fused TileLang kernel. Add tilelang_gdr_verify_prepass: one CTA per token computes l2norm for all Hk K-heads via a parallel [Hk,128] reduce (writing q_n/k_n directly from a full-range T.Parallel -- the head-batch global-write idiom, no per-row [1,DK] copy) and g=-exp(A_log)*softplus(a+dt_bias) (RAW log-decay) + beta=mul*sigmoid(b) for all Hv V-heads as one contiguous [1,Hv] row. It feeds the unchanged host-gated main kernel, whose hot loop then has zero transcendentals. Wired into recurrent_gated_delta_rule_verify(fuse_gating=True) via should_use_prepass(); prepass=True/False forces it. Persistent unbounded (never-evicting) scratch cache -> CUDA-graph capture-safe after warmup. Measured (H100, eager + CUDA-graph, parity 0.003-0.006, bit-faithful to the in-kernel-gated kernel): +7-24% in the bandwidth-bound T>=4 regime, stable across runs -- N=256,T=12 1.14x (-208us); N=64,T=12 1.18x; N=16,T=12 1.22-1.24x. No regression elsewhere: the conservative gate (T_avg>=4 and N*H*(1+T_avg)>=3000) routes single-request / small-batch / T=1 / launch-noise-flaky boundary regimes to the original in-kernel-gated kernel (zero pre-pass-fired losses in any run). Two iterate-loop learnings (verify-first caught both): the first pre-pass impl (serial Hk loop of single-row [1,128] reduces) was a ~6x slowdown -> the parallel [Hk,128] reduce fixed it; and Hopper copy-layout inference rejects an fp32 contiguous extent <128 bits, so the gate write must be the full [1,Hv] row, not a per-K-head [1,grp] slice (failed for grp<4). Tests: tests/test_prepass_gdr.py (parity vs gdn_sigmoid_gate+l2norm, raw-g negative control, prepass+host-gated == in-kernel-gated, prepass-path CUDA-graph capture); 67 decode+verify+prepass+head_batch tests pass. Bench: bench_prepass.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… sizes)
H2 verify-first profiling (the reduction-floor probe) found the no-write floor
at ~0.7% of HBM peak -> the recurrence is compute/latency-bound and the K-
reduction is unmoved by threads/tile tuning (H2's stated lever is a null). But
the floor isolate surfaced a much bigger, hidden lever: the FINAL-STATE WRITE.
The decode kernel stores ht[bb,bh,jk,v0+jv] = S[jv,jk] into the K-major
[B,H,DK,DV] contract -- a transpose of the S[v,k] register state. At block_DV<128
that transpose is catastrophically uncoalesced (~0.4 TB/s): it costs ~+111% over
the no-write floor (1175->2480us at B=256,T=12). At block_DV=128 (full-V tile,
n_vt=1) it coalesces -> write tail ~+4%. Net: block_DV=128 is ~2x faster at EVERY
batch size measured (1.35x @ B=1, 1.8x @ B=2-4, 3.0x @ B=8-16, 2.0-2.5x @
B=32-256; H in {8,16,32}), with the final state BIT-IDENTICAL to the reference.
The as-built ladder picked block_DV=64 uniformly and called 128 "occupancy-
starved" -- but that was tuned on the VERIFY kernel (V-major pool store, no
transpose, where 64 + more CTAs wins). Tile choice depends on the OUTPUT layout:
the decode K-major write needs the full-128 tile to coalesce; the verify V-major
store does not (it keeps block_DV=64).
Dispatch: fused_recurrent_gdr_fwd uses block_DV=128 when output_final_state (the
decode default), else the 64/32 occupancy ladder (perf-neutral there, no write).
Verify kernel unchanged. 67 decode+verify+prepass+head_batch tests pass (incl.
GQA, ragged seqlens, h0, negative controls). Probes:
benchmark/probe_h2_{reduction,verify_blockdv,blockdv_crossover}.py.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ated)
Measures the H1 win against the FLA/Triton baseline directly in one harness.
Result (H100, Hk16/Hv32, parity ~0.005 vs FLA): H1 lifts FlashQLA verify from
at-par (1.00-1.03x vs FLA, variant A) to a 1.13-1.18x lead at T=12 (the SGLang
draft length) and 1.09-1.10x at T=4, across N in {64,128,256}.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Differential attribution (benchmark/probe_sweep_attribution.py) of the post-H1 verify main kernel: at N=256,T=12 the no-ibuf recurrence floor is 1198us (86% of 1402us total); the per-token ibuf write adds only 204us (14.5%) at an impossible 15.8 TB/s -> the writes OVERLAP the compute-bound recurrence and are largely hidden. So the kernel is compute-bound on the serial reduce_sum K-reductions, not store-bound. Swept the remaining levers, recorded nulls in the verify spec so they are not re-attempted: (b) recurrence elementwise-pass fusion (4 [block_DV,DK] sweeps -> 2) is bit-identical + correct (67 tests) but PERF-NEUTRAL (reductions dominate, S is register-resident) -> reverted; (c) ibuf-store vectorization = drop (V-major store already coalesced, <=14.5% marginal, mostly hidden); (d) decode/verify store-coalescing audit = done (the only bad write, the decode K-major ht transpose, was fixed by block_DV=128; verify stores are V-major); (e) killing H1's 2nd-launch tax to win at small batch needs cross-module conv-epilogue fusion (SGLang-owned) or a risky 2-grid megakernel -- the regime gate already prevents regression. Net: the two shipped wins (H1 verify +8-24%; decode block_DV=128 ~2x) capture the tractable in-kernel headroom; the recurrence-reduction floor is the hard limit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…00->384) The regime gate was eager-calibrated: it routed single-request / small-batch T=12 verify to the in-kernel variant A because the prepass's 2nd launch is a ~13-15us eager net loss (N=1/T=12 eager 0.60x). But production runs the verify inside a CUDA graph, where the 2nd launch shrinks to a graph node and the tax VANISHES. Dedicated graph calibration (bench_prepass.py bench_graph, _time_graph 50+ warmup, H100, Hk=16/Hv=32): every T=12 point N>=1 WINS 1.08-1.24x under graphs -- N=1/T=12 (work=416) is 1.18x (vs 0.60x eager). The only graph losses are work<=320 (T=4 small-N, non-verify). Lowered MIN_WORK to 384 (floor just below N=1/T=12 work=416) so the prepass fires across the ENTIRE T=12 verify path under graphs. Tradeoff: eager small-N now pays the tax -- acceptable for the CUDA-graph deployment (steady state always captured; warmup untimed). bench_prepass.py gains graph-crossover calibration (work + gate decision + MISCAL per point). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
FLASHQLA_PREPASS_MIN_WORK env var overrides the gate work threshold (default 384). Set =3000 to reproduce the old eager-conservative gate on the SAME image -> isolates the gate re-tune from run-to-run serving variance in a parallel A/B. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a forward/inference-only fused recurrent (decode/verify) path for Gated Delta Rule, sitting beside the existing chunked-prefill kernels. Built for the SGLang DFlash speculative-verify use case. Validated end-to-end on a Modal H100 — 34 tests green — and benchmarked against this fork's FLA/Triton verify kernel.
What's in it
recurrent_gated_delta_rule) — gemm-free: state kept[block_DV, DK](V-major) in an fp32 fragment; the two GEMVs arereduce_sumover K; the rank-1 is aT.Parallelouter product. Avoids the tensor-coreM=1/ warp-partition walls and keeps the state fp32.recurrent_gated_delta_rule_verify, low-levelfused_recurrent_gdr_verify_fwd/_gated_fwd):state_indices(-1skip)[num_slots, Hv, V, K](matchesSGLANG_MAMBA_SSM_DTYPE; native layout, no transpose)disable_state_update)cu_seqlens(B=1 flattened), GQA (Hg ≤ H)fuse_gating=True)Speed vs FLA/Triton (H100, parity-gated, bf16 pool, per-token states)
The 1:1 baseline is this fork's
fused_sigmoid_gating_delta_rule_update(target-verify mode).FLA carries a ~53 µs fixed floor (
num_warps=1); FlashQLA floors at ~15 µs, so it wins the latency-bound regime that dominates speculative decode. At large batch both are compute-bound on the identical recurrence (the K-reductions) — measured: state writes are only ~21% of runtime — so parity there is the physical optimum, not a gap (confirmed by ruling out fp8 / replay-tape, which target the non-bottleneck writes).Validation
tests/test_decode_gdr.py,tests/test_verify_gdr.py): D∈{1,8}, GQA {1:1,1:4,MQA}, h0 on/off, g=0 SWA, ragged, low/high occupancy, no-commit/commit, V-major distinct-gate (transpose guard),-1slot skip, in-kernel gating, CUDA-graph replay, and negative controls (step-order / GQA-mapping discrimination).decode_recur/verify_refpinned against the chunk path atchunk_size=64.Files / how to use
flash_qla/ops/gated_delta_rule/fused_recurrent/(kernels + wrappers; SM90-gated)benchmark/bench_recurrent_gdr.py,benchmark/bench_vs_fla.py,benchmark/tune_verify.pydocs/superpowers/specs/,docs/superpowers/plans/Integration: drop
recurrent_gated_delta_rule_verify(..., fuse_gating=True)intogdn_backend.target_verify(the pool/ibuf are already the same V-major layout as the FLA path; softplus 1.0/20.0;allow_neg_eigval=False). Requires SM90 (Hopper) + TileLang 0.1.8.🤖 Generated with Claude Code