metal: GPU attention for S>4 prefill (single command buffer)

[RDouglasSharp]

Summary

The Metal fused-attention path (coli_metal_attn_decode) was capped at S<=4, so it only ran for decode / MTP. During prefill (S>4) every layer's attention fell back to the CPU even with COLI_METAL=1. This lifts the cap and runs large-S prefill attention on the GPU, in a single command buffer, cutting prefill attention time ~4× while producing output consistent with the CPU reference. It is off by default (COLI_METAL_PREFILL=1 to opt in), since GPU prefill can pick a different token on near-tie logits (the #622 family) and a greedy stream is not guaranteed bit-identical to the CPU. CPU and CUDA paths are unchanged.

Motivation

On Apple Silicon, COLI_METAL=1 accelerates decode attention and the routed experts, but prefill attention was left on the CPU (the gate required S<=4). For any real prompt that means a large slice of the cold-prefill wall runs on the CPU while the experts run on the GPU.

What changed

• c/colibri.c — the fused-attention gate no longer hard-caps at S<=4. It can engage for any S when Metal absorption is active, controlled by a new COLI_METAL_PREFILL env — default 0, which keeps the old S≤4 cap and leaves S>4 on the CPU; =1 opts in. Also adds a cumulative timer to the [prefill] progress line.
• docs/metal.md — documents the flag: what it buys (~4× prefill attention) and what it costs (near-tie divergence, [Bug]: Metal prefill GEMM is not token-exact vs CPU on near-tie logits (S >= GEMM_MIN); decode unaffected #622).
• c/backend_metal.mm — coli_metal_attn_decode now handles arbitrary S. Projections → attention core → output GEMV run in one command buffer, ordered by memory barriers — the same structure the S≤4 path already used and which is token-exact vs the CPU. A single a_score dispatch is S*AHEADS*T threads; it is guarded under ~2³⁰ and falls back to the CPU for very large prompts (in-encoder row-chunking to keep those on the GPU is a follow-up).

The CPU and CUDA code paths are untouched.

Performance

GLM-5.2 int4, Apple Silicon (M-series), 544-token prompt, greedy, default settings. CPU attention is timed per stage; on the GPU the three stages run fused in one command buffer and are reported as a single total.

	CPU	GPU
prefill attention	35.9 s (projection/RoPE 16.1 s · score·softmax·value 12.0 s · output 6.3 s)	9.0 s (8.6 s GPU kernel)
prefill wall	96.4 s	67.3 s

Prefill attention is ~4× faster on the GPU, trimming the prefill wall ~30% at this length. Prefill is dominated by streaming experts from disk, so the attention share — and the win — grows with S. Decode is unaffected.

Correctness

Compared GPU prefill (COLI_METAL_PREFILL=1) against the pure-CPU reference (COLI_METAL=0), both at the model's default settings, greedy decoding:

• Natural prompts (e.g. a ~540-token README excerpt): the GPU output stays close to and semantically consistent with the CPU — same opening, same reading of the input. Where individual tokens differ they are argmax near-ties (equally-valid tokens), not errors.
• Highly repetitive prompts (e.g. a pangram repeated 40×): an early token can flip and greedy cascades it. Same near-ties, amplified — GPU/CPU int4 accumulation rounds differently and GLM-int4 sits on argmax ties.

This is the documented #100 / #622 sensitivity, identical in kind to the prefill GEMM and the existing S≤4 Metal and MTP paths: every emitted token remains a valid argmax and the continuation stays correct, it just isn't guaranteed to be the same stream. Because of exactly this, the flag ships off by default.

Flags / behavior

• COLI_METAL_PREFILL=0 (default) — S>4 prefill attention stays on the CPU; bit-exact vs the CPU, no behavior change.
• COLI_METAL_PREFILL=1 — opt in to S>4 prefill attention on the GPU (the ~4× above; may diverge on near-tie logits, [Bug]: Metal prefill GEMM is not token-exact vs CPU on near-tie logits (S >= GEMM_MIN); decode unaffected #622).
• Prompts above the single-dispatch thread cap fall back to the CPU automatically.
• Documented in docs/metal.md.

Note (separate, pre-existing)

While validating this, we found that forcing ABSORB=1 routes S>4 prefill through the CPU absorbed path, which produces incorrect output on long prompts. That path predates this PR and is off by default, so this change doesn't touch it — but as a side effect, S>4 absorbed attention now runs on the correct GPU kernel instead. Flagging in case it's worth a separate look.

[JustVugg]

The measurement is good and the finding behind it is real: the S<=4 gate meant prefill attention sat on the CPU while the experts ran on the GPU, which is a strange split for anyone running COLI_METAL=1. 35.9 s → 9.0 s on the attention, 30% off the prefill wall at 544 tokens, single command buffer with the same structure the S<=4 path already uses. All of that I want.

One change requested: make COLI_METAL_PREFILL default to 0.

Why, in your own words

an early token can flip and greedy cascades it
the GPU output stays close to and semantically consistent with the CPU

That is a change to the emitted token stream, shipped on by default. And the front page of this project says:

no SLA on speed, and a hard guarantee on semantics

with the accompanying line that placement changes speed, never the model's answers. A default that can alter which tokens come out contradicts the one promise we make unconditionally — the guarantee is not "the continuation stays sensible", it is "you get the same model".

Your #100 argument is fair, and it is also why this needs the opposite default

You are right that this is the same kind of sensitivity as the existing S<=4 Metal path and MTP, and right that every emitted token remains a valid argmax. Nobody is calling this a correctness bug.

But look at how the project treats the other members of that family:

• [Bug]: Metal prefill GEMM is not token-exact vs CPU on near-tie logits (S >= GEMM_MIN); decode unaffected #622 is open right now: "Metal prefill GEMM is not token-exact vs CPU on near-tie logits (S >= GEMM_MIN); decode unaffected". This PR expands exactly that surface, so turning it on by default would enlarge a bug we have not closed.
• MTP is disabled when acceptance does not repay verification, and Greedy decoding is not reproducible: MTP (3/5 prompts) and the CUDA expert tier (2/5) both change the output #100 — the issue you cite — was closed by making non-reproducibility visible and controllable, not by accepting it.
• EXPERT_BUDGET is quarantined in colibri.c for being faster at the cost of semantics, with the reasoning left in the source as a reminder.
• The two most recent Metal PRs, metal: grouped-int4 (fmt=4) GEMV support + two latent fmt=4 fixes #457 and inkling: shared experts to the GPU, 1.75 to 2.46 tok/s on Apple Silicon #757, both led with "token-identical through eos" as the headline claim rather than the throughput.

Every one of those points the same way: speed that changes the output is opt-in here.

What I am asking for

1. COLI_METAL_PREFILL defaults to 0. Same patch, same measurements, one character.
2. A line in docs/metal.md saying what the flag buys and what it costs — "~4× prefill attention; output may differ from CPU on argmax near-ties ([Bug]: Metal prefill GEMM is not token-exact vs CPU on near-tie logits (S >= GEMM_MIN); decode unaffected #622)". People with a 544-token prompt and a Mac will turn it on; they should know what they are turning on.
3. Optionally, and only if it is cheap: have the engine say when it engages, the way [OMP] and [CUDA] lines do. A run whose output is not CPU-reproducible should carry that fact in its own log, so a benchmark posted here is self-describing.

If you would rather argue for default-on, the argument that would move me is a measurement rather than a principle: run the same prompt set both ways and report how often the streams actually diverge. If it is one prompt in fifty on natural text, that is a different conversation from a coin flip — and nobody has that number today, including me. Your repetitive-pangram case suggests the rate is input-dependent, which is exactly the sort of thing worth quantifying once rather than each of us guessing.

Two other things worth saying since they were easy to miss in a PR this well-written:

The 2³⁰ thread guard with a CPU fallback for very large prompts is the kind of thing that gets left out and then found by a user with a 4,000-token prompt. Thank you for putting it in.

And separating the CPU attention timings by stage (projection/RoPE 16.1 s · score·softmax·value 12.0 s · output 6.3 s) is what makes the 4× claim checkable rather than assertable.

[JustVugg]

Status note so this does not read as stalled: held on one character, and nothing else.

COLI_METAL_PREFILL defaulting to 1 ships a path that, by your own description, can change which tokens come out:

an early token can flip and greedy cascades it

The front page promises no SLA on speed, and a hard guarantee on semantics, and #622 — "Metal prefill GEMM is not token-exact vs CPU on near-tie logits" — is open right now. Default-on would enlarge an unclosed bug.

Everything else in the PR I want: 35.9 s → 9.0 s on prefill attention, single command buffer, the 2^30 thread guard with its CPU fallback, and the per-stage CPU timings that make the 4× checkable rather than assertable. None of that is in question.

Flip the default to 0, add a line to docs/metal.md saying what the flag buys and what it costs, and it merges.

And the standing offer: if you would rather argue for default-on, bring the divergence rate on natural prompts rather than a principle. One in fifty is a different conversation from a coin flip, your repetitive-pangram case suggests it is input-dependent, and nobody has that number — including me. A measurement would settle it either way, and we would know something today we do not.

[RDouglasSharp]

Understood and done. The default is now off.