marcsun13/ggml-gated-delta-net

ggml attention kernels

The layers a decode step actually spends its time in, as single kernels, taken from llama.cpp's ggml rather than reimplemented.

This is the companion to gguf-kernels and deliberately separate from it: that package computes on packed GGUF blocks and holds two ops (dequantize, mul_mat_vec), and it should stay that small. This one holds the sequence-mixing layers, which have nothing to do with quantization.

Ops

op	signature
gated_delta_net	(q, k, v, g, beta, state) -> (out, final_state)
flash_attn	(q, k, v, mask, scale) -> out
supports_gated_delta_net	(head_dim) -> bool
supports_flash_attn	(n_q, head_dim_k, head_dim_v) -> bool

gated_delta_net is one dispatch for a whole run of gated-delta-rule steps — the recurrence behind Qwen3-Next and Qwen3.5's linear-attention layers, which eager torch spells out in ~200 ops per layer.

from kernels import get_kernel

k = get_kernel("marcsun13/ggml-attn-kernels", version=1)
out, state = k.gated_delta_net(q, k_, v, g, beta, state)

Shapes, in torch order:

tensor	shape
q, k, v	(n_seqs, n_tokens, n_heads, head_dim)
g, beta	(n_seqs, n_tokens, n_heads)
state	(n_seqs, n_heads, head_dim, head_dim)

Two things about the contract are worth reading twice, because both are silent if you get them wrong:

• q and k must already carry one head per value head. Upstream maps a value head to a key head with i21 % ne01, which is the tiled convention; passing them pre-expanded makes that the identity, so your own head order is the one that applies rather than ggml's.
• state is indexed [value][key], transposed relative to the k-outer-v product that builds it, because upstream stores it that way to keep a thread's row contiguous. final_state comes back in the same layout, so a caller that stores what the op returns never transposes: a fresh state is zeros, which is symmetric. Mixing this op with a torch fallback that holds [key][value] means a 2 MB copy per layer per token, which is most of the win.

The output carries the kernel's own 1/sqrt(head_dim) scaling. g is the log-domain gate; the kernel exponentiates it. Ask supports_gated_delta_net(head_dim) rather than assuming: upstream covers a state row with 32 threads times head_dim/32 values each, so the head dim must be a multiple of 32, and only widths up to 128 are instantiated.

flash_attn

ggml's Metal flash attention, vector variant only — the one upstream itself picks for decode (n_q < 20). supports_flash_attn says no to anything wider, so a caller falls back to torch's own attention for prefill.

tensor	shape
q	(n_seqs, n_heads, n_q, head_dim)
k, v	(n_seqs, n_heads_kv, n_kv, head_dim)
mask	(n_seqs, 1, n_q, n_kv) additive, or None

Returns (n_seqs, n_q, n_heads, head_dim) — tokens before heads, which is what SDPA gives after its own .transpose(1, 2), so a caller usually wants precisely this and no further permute.

Grouped-query attention is native: pass k and v with n_heads_kv heads and do not expand them. That copy disappears rather than moving — for Qwen3.5-4B at 512 cached positions it is 16 MB a layer.

The mask is cast to f16 internally, as the kernel requires. This kernel accumulates in f16, which is the precision llama.cpp runs its attention at by default; it matches torch's f32 math path to about 2e-4, enough to change a greedy path after a few dozen tokens. That is a choice to make knowingly — see Measured.

Devices

backend	torch	targets
Metal	2.12, 2.13	aarch64-darwin

CUDA is a [kernel.*] section and a directory away — ggml ships gated_delta_net.cu too — and the bindings' schema does not move when it is added.

Measured

Qwen3.5-4B-Q4_K_M on an M2 Max, one token, 32 value heads, head_dim 128, against the torch recurrence in transformers:

	host µs	device µs	ms/token over 24 layers
torch recurrence	98.2	163.5	3.92
this kernel	3.8	16.6	0.40

End to end in eager generate(), 128 tokens, five interleaved rounds in one process, reproduced twice:

	tok/s	speedup	decode step, wall	decode step, host
torch rules	41.24	1.000×	22.23 ms	16.68 ms
this kernel	44.71	1.084×	19.71 ms	11.51 ms

(the other run: 41.85 → 44.98, 1.075×). Byte-identical completion — all 139 tokens.

flash_attn, 16 query heads over 4 kv heads, head_dim 256, one query:

n_kv	math SDPA + repeat_kv	this kernel
141	117.4 µs	34.5 µs	3.4×
512	348.1 µs	31.6 µs	11×

Nearly flat in cache length where SDPA scales with it. End to end in eager generate(), five interleaved rounds:

config	tok/s	speedup
baseline	41.99	1.000×
gated_delta_net	45.24	1.078×
flash_attn	46.52	1.108×
both	48.01	1.143×

Unlike the delta rule, flash attention changes the greedy path: on this prompt the completion diverges at generated token 68 into an equally coherent continuation. f16 accumulation is why, and it is what llama.cpp does too — but it means this op is not bit-compatible with torch's attention.

Where the kernels come from

vendor/ is a pinned subset of llama.cpp; vendor/UPSTREAM records the revision. The Metal backend compiles ggml-metal.metal into the embedded metallib and dispatches kernel_gated_delta_net_f32_<nsg> out of it, specialised through MTLFunctionConstantValues exactly as ggml specialises it.

Dispatches are encoded into torch's own current command buffer, on the MPS stream's serial queue. Both halves of that matter: encoding into torch's buffer means these kernels join the work already queued instead of each becoming its own submission, and doing it on the stream's queue is what stops it racing torch's commit. Skipping the second is how a Metal kernel ends up dying with failed assertion 'A command encoder is already encoding to this command buffer'.

Updating

python vendor.py --rev <llama.cpp commit>   # re-vendor, updates vendor/UPSTREAM
nix run .#build-and-copy                    # rebuild every variant into build/

For iterating without nix, localbuild/build.sh <python> builds the same sources with the system toolchain and leaves the metallib on disk; point GGML_ATTN_METALLIB at it and torch.ops.load_library at the .so. Tests run against either:

localbuild/build.sh python
export GGML_ATTN_METALLIB=$PWD/localbuild/out/ggml-metal.metallib
export GGML_ATTN_LOCAL_LIB=$PWD/localbuild/out/ggml_attn_kernels_local.so
pytest tests/ -q

See gguf-kernels/SKILL.md for the build.toml traps, the LFS-pointer trap, and what the Hub and the kernels client each require before get_kernel will load a repo.