perf(metal): pin_prefix on Gemma 4 E4B is slower than re-prefilling (4.7 s pin, no gain after)

[geisten]

Summary

geist_session_pin_prefix on Gemma 4 E4B Q4_K_M with the Metal backend makes requests slower, not faster.

HELIO kept one session per figure. Each session pinned its constant system turn once (BOS plus about 40 % of the prompt). After that, every request called geist_session_reset (which truncates to the pin) and prefilled only the rest, in 64-token chunks via geist_session_prefill_tokens.

M1 Max, geistlib 7a0b421, --ai-bench 20 across three figures:

	P50	P95	prefill P50	prefill P95	peak RSS
no pinning (full prompt every time)	1940 ms	2417 ms	1566 ms	2016 ms	9.93 GiB
pinning	2000 ms	5174 ms	1737 ms	4795 ms	11.23 GiB

• Pinning itself costs about 4.7 s per session on Metal. A cold full prefill of the whole prompt takes about 2.1 s.
• After pinning, prefilling the shorter remainder is not faster than prefilling the whole prompt.
• On NVIDIA Vulkan the same code gives −11 % P50 (1536 → 1360 ms).

Requests

• Check whether pin_prefix takes a slow path on Metal (token by token instead of batched).
• Check why prefilling after a pinned prefix costs as much as the full prompt (does reset re-run the prefix?).

Related: #548 (snapshot/restore), #577 (memory).

[geisten]

Likely the same root cause, fixed in #594: a prefill batch whose row count was off the GEMM tile (32 rows Q4_K, 64 Q6_K) ran every tile through slow kernels, and 2–31 rows cost a full 64-row pass. The remainder after a pinned prefix is exactly such a ragged batch. Not re-measured with pinning (HELIO dropped it); left open for that.

[geisten]

I went through the code (read only, nothing measured, since there's no Metal here).

pin_prefix and reset are cheap on Metal:

• Pin: transformer_pin_prefix (arch_ops.c:763-802) is a normal batched prefill (prefill_text_batch_inner, in chunks of m_max = 256 on Metal) plus prefix_length = kv_len. There is no token-by-token step, no KV download or copy, and no extra sync. The only exception is KIVI, which is off by default.
• Reset: transformer_session_reset (forward/step.c:462) only sets kv_len = prefix_length. Nothing is recomputed.
• Gemma 4 specifics: sliding-window layers mask by q_pos over a linear cache, shared-KV layers read the owner's cache, and PLE is computed per chunk. None of these invalidates the pin.

The cost therefore comes from the forward pass, depending on the row count of each chunk:

1. Ragged rows in the Q4_K/Q6_K GEMM. This was fixed in metal: one copy of load_from_memory weights; ragged prefill batches on the fast GEMM (#577, #578) #594 and matches both symptoms. The pin is a single chunk with a ragged row count, and the remainder after reset is ragged too.
2. Still open after metal: one copy of load_from_memory weights; ragged prefill batches on the fast GEMM (#577, #578) #594: metal_attention uses the flash prefill kernel only when q_rows % 8 == 0 (src/backends/metal/ops.c:3730). Every other chunk falls back to metal_encode_attention_rows. That hits the pin chunk and the remainder's tail, and gets more expensive with a larger kv_len, i.e. a longer pin.
3. The first session might pay one-time warm-up costs: lazy pipelines, and page-in of the no-copy-wrapped weights.

Suggested A/B on the M1 Max (I can't measure it here): at HEAD, pin + remainder with a prefix padded to a multiple of 64 tokens versus unpadded, and each of those with GEIST_METAL_FLASH=0. That separates causes 1 and 2. If cause 2 dominates, the fix is a flash kernel that masks the row tail. I can build it, but only verify it here as code without Metal.

---

Generated by Claude Code