flashinfer.gdn_decode

Gated Delta-Rule decode-side kernels used in Mamba-2 / GDN-style sequence models. These functions consume a pre-built KV / state cache and run the recurrent gated delta-rule update for the current decode step.

gated_delta_rule_decode(q, k, v, state, ...)	Gated Delta Rule Decode kernel (K-major layout, no transpose needed).
gated_delta_rule_decode_pretranspose(q, k, ...)	Gated Delta Rule Decode kernel for single-token generation.
gated_delta_rule_mtp(q, k, v, initial_state, ...)	Gated Delta Rule MTP kernel (Multiple Token Processing).
gated_delta_rule_replayssm_commit(...[, ...])	Commit an accepted ReplaySSM prefix into all layers' checkpoints.

Rubin T=1 bf16-state decode (SM107)

On Rubin (compute capability 10.7), gated_delta_rule_decode_pretranspose routes T=1 calls with a bfloat16 state, K=V=128 and B * HV >= 512 to a Rubin-tuned kernel under backend="auto". It requests each 64-row state tile at once into shared memory, prefetches a later tile into L2 and issues fewer instructions; outputs and the updated state are bitwise equal to the generic kernel. backend="rubin" forces it (any batch size, Rubin only) and backend="flashinfer" keeps the generic kernel.

uGPU localization: Rubin is built from two dies, each with its own SMs and memory. Allocate the state pool with ugpu_state_pool() (layout [pool_size, HV, 2, V // 2, K], half t of every state in die t’s memory) and a T=1 call on Rubin runs half t on die t’s SMs, so no state byte crosses the die-to-die link, whatever the batch order or slot indices. Results are bitwise equal to the one-launch kernel. This targets CUDA-graph captured decode (make one eager call before capture, which creates the per-die green contexts); eager calls with B * HV < 8192 are host-bound and run the two halves back to back on the current stream.

ugpu_state_pool(pool_size, num_v_heads[, ...])

Allocate a bfloat16 GDN state pool for uGPU localization on Rubin.

ReplaySSM speculative verification (SM100/SM103)

Use gated_delta_rule_mtp(cache_replayssm=True, disable_state_update=True) with FP32 128x128 checkpoints. For each layer, pass contiguous views into all-layer BF16 raw K/V windows and FP32 log-decay/beta windows. The windows are indexed by the same pool slots as the checkpoints. Verify produces outputs without changing the checkpoint or storing a full state for every token.

Once the accepted lengths are known, call gated_delta_rule_replayssm_commit once across all layers. It replays only the accepted prefix and updates each live checkpoint in place. Zero acceptance leaves the checkpoint unchanged. The next verify overwrites the raw windows; no circular cache or deferred state update is retained between rounds. Use identical Q/K normalization settings for verify and commit.

The tcgen05 verify specialization uses TMA and TF32 tensor cores for T=4..8, B=1..256, BF16 output, int32 indices, contiguous FP32 checkpoints, normalized Q/K and an even value-head/key-head ratio. Other supported SM100 inputs use the FP32 MTP fallback. TF32 contraction changes numerical reduction precision. See the API docstrings for shape, index uniqueness, padding and optional tracked-checkpoint contracts.

Commit defaults to backend="auto". With normalized K and no tracking, T=8 uses tcgen05; T=6/7 uses tcgen05 when B>1 and layers * B * HV >= 256. Other cases use SIMT. Explicit backend="tcgen05" supports T=4..8. The heuristic uses tensor metadata, not device-resident accepted lengths; short accepted prefixes may favor backend="simt". SIMT also retains FP32 contractions when checkpoint accuracy is preferred over tensor-core throughput.