Kernels:

phanerozoic
/

gated-deltanet

Kernel card Files Files and versions

xet

Community

phanerozoic commited on 17 days ago

Commit

1e4acca

verified ·

1 Parent(s): 1c890a4

Card: standardized form with hero

Browse files

Files changed (2) hide show

CARD.md +91 -115
README.md +91 -115

CARD.md CHANGED Viewed

@@ -5,95 +5,27 @@ license: apache-2.0
 # gated-deltanet
-The gated delta rule in recurrent form, loadable through `kernels`. 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. At 48 such layers, 48 heads and a 128 x 128 state the entire
-recurrent memory is **144 MB regardless of sequence length**, where attention
-over the same context would need gigabytes.
-Reference: transformers' `torch_recurrent_gated_delta_rule`. That layer is
-decorated `@use_kernel_forward_from_hub("Qwen3_5GatedDeltaNet")`, so this is
-directly substitutable for it.
-## How it works
-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 is what separates this from plain linear attention. The state is
-corrected by the difference between the value and what the state already holds
-for that key, rather than having the value added blindly, so writing the same
-key twice does not double its contribution. Two gates control it: `exp(g)`
-decides how much of the past survives, `beta` how strongly the correction lands.
-The state is the whole cost, 128 x 128 floats per head, 64 KiB, and a
-straightforward implementation reads and writes it several times per token. Here
-one block owns one (batch, head) pair and holds the entire state in its threads'
-registers for the whole sequence: a thread owns one value column and a slice of
-the key rows. The state is read from memory once at the start and written once
-at the end, whatever the length, and every intermediate token touches registers
-only. Prefill and decode are therefore the same kernel, differing only in how
-many tokens the loop runs.
-## Measured
-RTX 6000 Ada (48 GB), torch 2.10 + CUDA 12.6, 48 heads, 128 x 128 state, against
-the eager reference on the same inputs.
-| batch | tokens | eager | this | speedup |
-|---|---|---|---|---|
-| 1 | 1 | 0.435 ms | 0.028 ms | **15.3x** |
-| 1 | 64 | 7.451 ms | 0.347 ms | **21.4x** |
-| 1 | 512 | 61.114 ms | 2.787 ms | **21.9x** |
-| 4 | 256 | 40.268 ms | 2.468 ms | 16.3x |
-The eager path is a Python loop over timesteps, each step launching several
-tensor operations over the full state; this kernel runs the loop on the device
-with the state resident in registers.
-State footprint, Bonsai 27B geometry:
-| | |
-|---|---|
-| per layer, per sequence | 3.0 MB |
-| all 48 linear layers | **144 MB, constant in context length** |
-## Correctness
-Verified on RTX 6000 Ada 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, both 64 and 128 dimensions, with
-  and without the in-kernel L2 normalization of queries and keys.
-- **Chunked equals whole.** Running a 96-token sequence in six 16-token calls,
-  carrying the state, is `torch.equal` to one call. Prefill and decode are the
-  same operation.
-- **Token by token equals whole.** Decoding 48 tokens one at a time is
-  `torch.equal` to the batched pass, state included.
-- **An initial state is respected**, matching the reference given the same
-  start.
-- **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.
-### A stability note, not a kernel limit
-The delta rule is stable only while `beta * |k|^2` stays near or below 2. Raw
-N(0,1) keys at dimension 128 give roughly 65, and the recurrence genuinely
-diverges: the reference reaches 1.3e26 on such input and this kernel reproduces
-that to five significant figures. Real models avoid it by L2 normalizing the
-keys, which is what `use_qk_l2norm` does. If you see the state explode, the
-input is outside the rule's stable region and both implementations will agree
-that it is.
 ## Usage
@@ -111,8 +43,9 @@ 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, so pass the
-pre-exponential gate.
 ## API
@@ -122,36 +55,79 @@ pre-exponential gate.
 | `new_state(batch, heads, key_dim, value_dim)` | a zeroed `[B, H, Dk, Dv]` fp32 state |
 | `GatedDeltaNet(batch, heads, key_dim, value_dim, use_qk_l2norm)` | module holding the state across steps, one per layer |
-`q`, `k` are `[B, T, H, Dk]`; `v` is `[B, T, H, Dv]`; `g`, `beta` are
-`[B, T, H]`. Returns `[B, T, H, Dv]`.
 ## Requirements and limits
-- NVIDIA GPU with compute capability 8.0+.
-- fp32 arithmetic. The state is a long product chain and low precision would
-  compound across the sequence; the reference computes in fp32 for the same
-  reason.
-- `Dv` must divide the 256-thread block, `Dk` must be divisible by
-  `256 / Dv`, and `Dk / (256 / Dv)` must be at most 64. The 128 x 128 geometry
-  this targets satisfies all three; 64 x 64 also works.
-- `Dk` at most 256.
-- **Recurrent form only.** Long prefills would run faster in the chunked
-  formulation, which turns the sequential scan into matrix products at the cost
-  of a more complex derivation. This kernel keeps the exact recurrence, which is
-  what makes chunked and token-by-token bitwise identical, and it is already
-  20x the eager path.
-- Grouped heads are the caller's business: a model with more value heads than
-  key heads repeats the 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

 # 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.
+![The 128 by 128 recurrent state blooms, wipes under the retention gate, and re-blooms while the memory line stays flat against attention's climbing KV curve](https://huggingface.co/kernels/phanerozoic/gated-deltanet/resolve/main/media/hero.gif)
+*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
 `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
 | `new_state(batch, heads, key_dim, value_dim)` | a zeroed `[B, H, Dk, Dv]` fp32 state |
 | `GatedDeltaNet(batch, heads, key_dim, value_dim, use_qk_l2norm)` | module holding the state across steps, one per layer |
+## 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:
+| batch | tokens | eager | this | speedup |
+|---|---|---|---|---|
+| 1 | 1 | 0.435 ms | 0.028 ms | 15.3x |
+| 1 | 64 | 7.451 ms | 0.347 ms | 21.4x |
+| 1 | 512 | 61.114 ms | 2.787 ms | 21.9x |
+| 4 | 256 | 40.268 ms | 2.468 ms | 16.3x |
+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

README.md CHANGED Viewed

@@ -5,95 +5,27 @@ license: apache-2.0
 # gated-deltanet
-The gated delta rule in recurrent form, loadable through `kernels`. 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. At 48 such layers, 48 heads and a 128 x 128 state the entire
-recurrent memory is **144 MB regardless of sequence length**, where attention
-over the same context would need gigabytes.
-Reference: transformers' `torch_recurrent_gated_delta_rule`. That layer is
-decorated `@use_kernel_forward_from_hub("Qwen3_5GatedDeltaNet")`, so this is
-directly substitutable for it.
-## How it works
-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 is what separates this from plain linear attention. The state is
-corrected by the difference between the value and what the state already holds
-for that key, rather than having the value added blindly, so writing the same
-key twice does not double its contribution. Two gates control it: `exp(g)`
-decides how much of the past survives, `beta` how strongly the correction lands.
-The state is the whole cost, 128 x 128 floats per head, 64 KiB, and a
-straightforward implementation reads and writes it several times per token. Here
-one block owns one (batch, head) pair and holds the entire state in its threads'
-registers for the whole sequence: a thread owns one value column and a slice of
-the key rows. The state is read from memory once at the start and written once
-at the end, whatever the length, and every intermediate token touches registers
-only. Prefill and decode are therefore the same kernel, differing only in how
-many tokens the loop runs.
-## Measured
-RTX 6000 Ada (48 GB), torch 2.10 + CUDA 12.6, 48 heads, 128 x 128 state, against
-the eager reference on the same inputs.
-| batch | tokens | eager | this | speedup |
-|---|---|---|---|---|
-| 1 | 1 | 0.435 ms | 0.028 ms | **15.3x** |
-| 1 | 64 | 7.451 ms | 0.347 ms | **21.4x** |
-| 1 | 512 | 61.114 ms | 2.787 ms | **21.9x** |
-| 4 | 256 | 40.268 ms | 2.468 ms | 16.3x |
-The eager path is a Python loop over timesteps, each step launching several
-tensor operations over the full state; this kernel runs the loop on the device
-with the state resident in registers.
-State footprint, Bonsai 27B geometry:
-| | |
-|---|---|
-| per layer, per sequence | 3.0 MB |
-| all 48 linear layers | **144 MB, constant in context length** |
-## Correctness
-Verified on RTX 6000 Ada 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, both 64 and 128 dimensions, with
-  and without the in-kernel L2 normalization of queries and keys.
-- **Chunked equals whole.** Running a 96-token sequence in six 16-token calls,
-  carrying the state, is `torch.equal` to one call. Prefill and decode are the
-  same operation.
-- **Token by token equals whole.** Decoding 48 tokens one at a time is
-  `torch.equal` to the batched pass, state included.
-- **An initial state is respected**, matching the reference given the same
-  start.
-- **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.
-### A stability note, not a kernel limit
-The delta rule is stable only while `beta * |k|^2` stays near or below 2. Raw
-N(0,1) keys at dimension 128 give roughly 65, and the recurrence genuinely
-diverges: the reference reaches 1.3e26 on such input and this kernel reproduces
-that to five significant figures. Real models avoid it by L2 normalizing the
-keys, which is what `use_qk_l2norm` does. If you see the state explode, the
-input is outside the rule's stable region and both implementations will agree
-that it is.
 ## Usage
@@ -111,8 +43,9 @@ 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, so pass the
-pre-exponential gate.
 ## API
@@ -122,36 +55,79 @@ pre-exponential gate.
 | `new_state(batch, heads, key_dim, value_dim)` | a zeroed `[B, H, Dk, Dv]` fp32 state |
 | `GatedDeltaNet(batch, heads, key_dim, value_dim, use_qk_l2norm)` | module holding the state across steps, one per layer |
-`q`, `k` are `[B, T, H, Dk]`; `v` is `[B, T, H, Dv]`; `g`, `beta` are
-`[B, T, H]`. Returns `[B, T, H, Dv]`.
 ## Requirements and limits
-- NVIDIA GPU with compute capability 8.0+.
-- fp32 arithmetic. The state is a long product chain and low precision would
-  compound across the sequence; the reference computes in fp32 for the same
-  reason.
-- `Dv` must divide the 256-thread block, `Dk` must be divisible by
-  `256 / Dv`, and `Dk / (256 / Dv)` must be at most 64. The 128 x 128 geometry
-  this targets satisfies all three; 64 x 64 also works.
-- `Dk` at most 256.
-- **Recurrent form only.** Long prefills would run faster in the chunked
-  formulation, which turns the sequential scan into matrix products at the cost
-  of a more complex derivation. This kernel keeps the exact recurrence, which is
-  what makes chunked and token-by-token bitwise identical, and it is already
-  20x the eager path.
-- Grouped heads are the caller's business: a model with more value heads than
-  key heads repeats the 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

 # 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.
+![The 128 by 128 recurrent state blooms, wipes under the retention gate, and re-blooms while the memory line stays flat against attention's climbing KV curve](https://huggingface.co/kernels/phanerozoic/gated-deltanet/resolve/main/media/hero.gif)
+*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
 `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
 | `new_state(batch, heads, key_dim, value_dim)` | a zeroed `[B, H, Dk, Dv]` fp32 state |
 | `GatedDeltaNet(batch, heads, key_dim, value_dim, use_qk_l2norm)` | module holding the state across steps, one per layer |
+## 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:
+| batch | tokens | eager | this | speedup |
+|---|---|---|---|---|
+| 1 | 1 | 0.435 ms | 0.028 ms | 15.3x |
+| 1 | 64 | 7.451 ms | 0.347 ms | 21.4x |
+| 1 | 512 | 61.114 ms | 2.787 ms | 21.9x |
+| 4 | 256 | 40.268 ms | 2.468 ms | 16.3x |
+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