gated-deltanet

The gated delta rule in recurrent form, loadable through kernels. The reference baseline is transformers' torch_recurrent_gated_delta_rule, which this kernel matches to 5.5e-7 relative and replaces via @use_kernel_forward_from_hub("Qwen3_5GatedDeltaNet").

Three of every four layers in a hybrid-attention model are not attention: they carry a fixed-size state per head and update it once per token, so memory does not grow with context. The catch is speed: written in framework ops the recurrence is a Python loop that reads and writes the whole state several times per token. This kernel holds the entire state in registers for the whole sequence, one block per (batch, head) pair, reading it from memory once and writing it once whatever the length, which makes prefill and decode the same kernel and the eager loop about twenty times slower.

One head's 128 x 128 state filmed live across 2,048 tokens of updates, including periodic retention-gate wipes (g large and negative erases the past, then the state re-blooms). At 48 layers the recurrent memory is 144 MB at any context; attention's KV cache at the same geometry passes 155 GB by 131,072 tokens. Output matches the transformers reference to 1e-6.

Usage

import torch
from kernels import get_kernel

gdn = get_kernel("phanerozoic/gated-deltanet", version=1, trust_remote_code=True)

state = gdn.new_state(batch, heads, key_dim, value_dim)     # [B, H, Dk, Dv] fp32

# prefill and decode are the same call; the state carries
out = gdn.gated_delta_rule(q, k, v, g, beta, state, use_qk_l2norm=True)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. g is the log decay: the kernel applies exp(g), matching the reference. q, k are [B, T, H, Dk]; v is [B, T, H, Dv]; g, beta are [B, T, H]; returns [B, T, H, Dv].

API

Method

Per token the update is a decay, a read, a rank-one correction and a read:

S    <- S * exp(g)      gated decay, one scalar per head
m    <- k^T S           what the state already predicts for this key
d    <- (v - m) * beta  the delta, scaled by the write gate
S    <- S + k d^T       rank-one write
out  <- q^T S           read with the query

The delta rule corrects the state by the difference between the value and what the state already holds for that key, so writing the same key twice does not double its contribution. The state is the whole cost, 64 KiB per head, and here one block owns one (batch, head) pair with the state in its threads' registers for the whole sequence: a thread owns one value column and a slice of the key rows. Every intermediate token touches registers only.

Measured

RTX 6000 Ada, 48 heads, 128 x 128 state, against the eager reference on the same inputs:

State footprint, Bonsai 27B geometry: 3.0 MB per layer per sequence, 144 MB across all 48 linear layers, constant in context length.

Correctness

Verified against transformers 5.13:

• Output and final state agree with the reference to 5.5e-7 relative across batch 1 to 3, 8 to 256 tokens, 2 to 8 heads, 64 and 128 dimensions, with and without the in-kernel L2 normalization.
• Chunked equals whole: a 96-token sequence in six 16-token calls, carrying the state, is torch.equal to one call; decoding token by token equals the batched pass, state included.
• Gate semantics: beta = 0 leaves the state exactly untouched; a large negative g shrinks it by nine orders of magnitude; g = 0 with beta = 0 is the exact identity.
• Deterministic: repeated runs are bitwise identical.
• The delta rule is stable only while beta * |k|^2 stays near or below 2; real models L2-normalize keys (use_qk_l2norm), and on out-of-region input this kernel reproduces the reference's divergence to five significant figures.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+; fp32 arithmetic (the state is a long product chain).
• Dv must divide the 256-thread block, Dk divisible by 256 / Dv, Dk / (256 / Dv) <= 64; Dk <= 256. The 128 x 128 and 64 x 64 geometries satisfy all constraints.
• Recurrent form only: long prefills would run faster in the chunked formulation; this kernel keeps the exact recurrence, which is what makes chunked and token-by-token bitwise identical.
• Grouped heads are the caller's: repeat key and query heads before calling, as the reference layer does. Forward only.

References

Yang et al., "Gated Delta Networks: Improving Mamba2 with Delta Rule" (2024); Yang et al., "Parallelizing Linear Transformers with the Delta Rule over Sequence Length" (2024); Schlag et al., "Linear Transformers Are Secretly Fast Weight Programmers" (2021); the reference implementation in transformers' Qwen3_5GatedDeltaNet.

License

Apache-2.0.