perf(purejava): let the hidden-state prefill reach the batched path (15.5x)

[bsbodden]

What was wrong

prefill has consulted the batched execution plan since it was written. prefillHiddenState never did.

// LlamaForwardPass:817  prefill(Session, …)          → HAS the branch
if (batchedPrefill && tokens.length > 1) { return prefillSessionBatched(session, tokens); }

// LlamaForwardPass:851  prefillHiddenState(Session, …) → NO branch, token-by-token loop

Both variants were sequential, so every hidden-state consumer paid decode rates for work the batched path already knew how to do — embeddings through GgufEmbeddingBackend, and any consumer reading a state off a prompt.

Measured

Granite 4.1 3B Q4_K_M (SHA-256 662b0626…, catalog-verified), Hetzner ccx33, 8 vCPU AMD EPYC-Milan, AVX2, Temurin 25.0.4.1, rust-ffm kernels. Identical input, same process.

prefillHiddenState	128 tokens	rate	512 tokens	rate
before	66,538 ms	1.9 tok/s	—	—
after	4,283 ms	29.9 tok/s	13,973 ms	36.6 tok/s
prefill (logits, batched), reference	4,819 ms	26.6 tok/s	13,444 ms	38.1 tok/s

15.5x at 128 tokens. It now runs slightly faster than the logits route at that length and reaches parity by 512 — the expected shape, since it skips the vocabulary projection on every row but the last.

The hypothesis that was refuted

vectors.gguf.pollMillis was the leading explanation, since vectors 0.1.22 warns that a caller running its own pool beside the persistent executor should set it to zero. Ablated in both directions:

pollMillis	64-token hidden prefill
0	34,385 ms (1.86 tok/s)
5 (default)	34,198 ms (1.87 tok/s)

Identical within noise. It explains none of the gap. The cause is memory bandwidth: a single-token forward streams the whole ~2 GB of Q4_K_M weights to produce one row, where a batched block streams it once for the lot.

The change

Both variants advance every position but the last through the existing qualified batched path, then take one ordinary step for the final token so it still yields a hidden state rather than a vocabulary projection. The batched call's logits are discarded; what is wanted is the cache it leaves. No new machinery. The final full sweep remains as a floor a deeper change could remove.

Tests

Taking the branch is asserted through a counter rather than inferred from timing, and in both directions — a prompt with nothing to batch must leave the counter alone, because a branch never taken looks exactly like a branch that did not matter.

That counter earned itself immediately. The first version of these tests used an all-F32 nano fixture, and F32 is absent from TensorOps.supportsBatchedMatmul, so the fixture could not batch at all and the three equivalence tests were passing vacuously — comparing the sequential path against itself. Only the counter assertion failed. The fixture now uses Q8_0 projections at the dimensions Q8_0 alignment requires.

Also pinned: hidden state, key/value cache contents, and the next generated token all match running the prompt one token at a time, within the repository's existing SIMD_REDUCTION_TOLERANCE.

The execution plan is pinned rather than defaulted, since defaultPlan consults RuntimeFingerprint and a test resting on it would assert one thing on a workstation and another on a benchmark host.

681 backend-java tests pass, no regressions. Spotless and SpotBugs clean.

🤖 Generated with Claude Code