ggml-attn

ggml's flash attention, as a torch op and as a transformers attention implementation. The kernel is llama.cpp's, not a reimplementation of it — the shader is vendored from ggml and dispatched with the same function constants upstream uses.

What it buys over the SDPA fallback on MPS, where torch has no fused attention and drops to the math path: the kernel is nearly flat in cache length where math SDPA scales with it, and grouped-query attention is native, so repeat_kv's copy disappears rather than moving.

Coverage

ggml has two implementations of flash_attn_ext and picks between them by query count — ggml_metal_op_flash_attn_ext_use_vec() selects the vector path for ne01 < 20. Only that one is ported here, which is the path upstream itself uses for decode. Wider q — real prefill — falls through to torch's own attention inside flash_attn_forward, so this is usable as a model's single attention implementation.

n_q	path
1–19	this kernel
≥ 20	torch.nn.functional.scaled_dot_product_attention

Head dims must match between k and v and be a multiple of 32. Ask supports_flash_attn rather than assuming; a second backend will cover a different set.

Ops

from kernels import get_kernel

k = get_kernel("kernels-community/ggml-attn")

out = k.flash_attn(q, k_, v, mask=None, scale=None)

	shape
q	(n_seqs, n_heads, n_q, head_dim) f32
k, v	(n_seqs, n_heads_kv, n_kv, head_dim) f32 — do not expand to n_heads, GQA is native
mask	(n_seqs, 1, n_q, n_kv) additive, or None. Cast to f16 internally
returns	(n_seqs, n_q, n_heads, head_dim) f32 — tokens before heads

The output layout is what SDPA gives after its own .transpose(1, 2), so a caller usually wants exactly this and no further permute.

Causality is the mask

ggml_attn_ext takes no is_causal argument. In ggml the mask is the causality, and a null mask means attend to everything — that is a real mode, used for bidirectional attention. llama.cpp never hits this because it always builds a KQ mask into its graph.

transformers is the other convention: sdpa carries causality in its own is_causal flag, so the mask is dropped as an optimisation whenever it would be plain-causal — and that flag never reaches an attention function. flash_attn_forward therefore rebuilds the mask when it gets None with n_q > 1. Calling the raw flash_attn op yourself, you own that: pass a mask, or get bidirectional attention.

This is silent when wrong. On Qwen3.5-4B a non-causal 16-token prefill moved the last-position logits by 1.44 while leaving the greedy top-5 unchanged, so a text-level comparison does not catch it.

With transformers

model = AutoModelForCausalLM.from_pretrained(
    ...,
    attn_implementation="kernels-community/ggml-attn:flash_attn_forward",
    allow_all_kernels=True,
)

The package publishes mask_implementation = "sdpa", which transformers reads when it registers the kernel. That matters: without it a kernel-provided attention inherits flash-attention-2 semantics, where no mask is built at all because the kernel is assumed to handle causality and padding itself. This one does neither — it takes a dense additive mask, and it cannot reconstruct padding it never receives.

Measured

Qwen3.5-4B-Q4_K_M on an M2 Max, 16 query heads over 4 kv heads, head_dim 256, f32. End to end through generate(), 128 tokens, arms alternated within one process and the median of per-rep ratios reported — between-process tok/s on this machine varies by more than the effect being measured.

prompt	sdpa	this kernel
16	52.33 tok/s	53.52	1.023×
1024	29.91 tok/s	35.83	1.198×

The gain grows with cache depth, which is the whole point — a 16-token prompt never gets past ~140 cached positions, the cheap end of SDPA's curve, and is this kernel's worst case rather than its typical one.

Against torch's f32 math attention at decode shapes (n_q = 1), maximum absolute difference: 5e-05 at n_kv = 141, 9e-05 at 512. The kernel accumulates in f16, which is what llama.cpp's attention does by default, so it is not bit-identical to torch — close, but a caller comparing exact logits should know.

Devices

backend	torch	targets
Metal	2.12, 2.13	aarch64-darwin

CUDA is a [kernel.*] section and a directory away — ggml ships fattn*.cu, and unlike its Metal backend the CUDA tree is many translation units, so that section can cherry-pick files rather than compiling everything. The bindings' schema does not move when it is added.

Where the kernel comes from

vendor/ is a pinned subset of llama.cpp; the revision is in vendor/UPSTREAM. Refresh it with python vendor.py --rev <sha>.

The whole of ggml-metal.metal is compiled, not just the flash-attention part. That is not a choice — ggml's Metal backend is a single 11k-line shader whose kernels share ggml-common.h, the type tables and the dequantize helpers, so there is nothing to select. Unused kernels cost 6.8 MB of metallib and nothing at runtime, since Metal builds pipelines lazily per kernel.

Three things the dispatch does that a torch.mps.compile_shader caller cannot:

• MTLFunctionConstantValues, so a kernel is specialised the way ggml specialises it instead of the source being rewritten before compiling.
• encoding into torch's current command buffer, so these dispatches join the work already queued rather than each becoming its own submission.
• doing that on the MPS stream's own serial queue via dispatch_sync(stream->queue()), which is what keeps it legal — encoding from elsewhere trips 'A command encoder is already encoding to this command buffer'.

Building

nix run .#build-and-copy          # the published artifacts, into build/
localbuild/build.sh               # ~20s, system toolchain, for iterating
pytest tests/                     # 47 cases

localbuild/build.sh leaves the metallib on disk rather than embedding it and stages a loadable variant under localbuild/variants/, which the tests pick up when build/ is absent.