Batched prefill never re-enabled in generateUntilStop / agent loop after position-collapse fixes

[michalharakal]

Summary

generateUntilStop() in llm-agent — the entry point used by AgentLoop, JavaAgentLoop, and KLlamaSession.generate() — still hardcodes autoregressive prefill (one forward() per prompt token) with a comment pointing at the old position_collapse_bug.md divergence. That revert is stale: the underlying correctness work is done and tested, but the batched path was never re-enabled on the agent/session route.

What is already fixed (verified on develop @ a119dfb)

• Head/seq reshape divergence (the bd3eb9c-era root cause): fixed by swapSeqHeadDims in MultiHeadAttention — the detailed post-mortem lives in the code comment.
• RoPE position collapse: both applyRoPEInterleaved (Llama family) and the split-half variants are batch-aware — row s rotates at position + s.
• Cache-aware causal masking: DefaultCpuOps.scaledDotProductAttention aligns bottom-right (maxKi = seqKV - seqQ + qi), so chunked prefill with a warm KV cache masks correctly.
• PrefillStrategy infrastructure: llm-core's generate() already exposes PrefillStrategy.Autoregressive | Batched(maxBatch) with chunked forwardBatched prefill.
• Equivalence evidence (TinyLlama 1.1B Q8_0, FP32 dequant, M-series CPU):
  • BatchedPrefillEquivalenceTest: full-prompt batched prefill matches the autoregressive baseline at the last position, max_abs_diff = 5.7e-6, identical argmax.
  • PrefillStrategyEquivalenceTest (both methods, incl. forced multi-chunk maxBatch=3): identical greedy token sequences.
  • New GenerateUntilStopPrefillEquivalenceTest (both methods): identical greedy sequences through the exact agent-loop entry point.
  • Note: these model-loading tests want a quiet machine — two concurrent 4.4 GB FP32 loads alongside another large JVM reproducibly killed the test JVM before any assertion ran.

The gap

llm-agent's generateUntilStop() never got the prefillStrategy parameter, so every agent-loop and Java-facade consumer is pinned to one-forward-per-token prefill. On CPU-only runtimes prefill dominates wall time — and the agent loop re-processes the full conversation every round, multiplying the cost.

Real-world impact

Daily-StandAPP (JavaLand 2026 demo, Llama 3.2 1B Q8_0 embedded via JavaAgentLoop): 87 s to first token for a ~700-token tool-calling template on an M-series CPU; standup summarisation prompts run 1700+ tokens. With the documented 3–10× batched-prefill speedup this drops to roughly 10–25 s.

Proposed fix (PR follows)

Opt-in, default behavior unchanged:

• generateUntilStop(...) gains prefillStrategy: PrefillStrategy = Autoregressive; Batched ingests the prompt in maxBatch chunks via forwardBatched, reporting onPrefill per chunk.
• AgentConfig gains prefillStrategy, threaded through both AgentLoop call sites (JavaAgentLoop inherits it via AgentConfig — no facade change needed).
• GenerationConfig gains a Java-friendly batchedPrefill(maxBatch) builder knob, threaded through both KLlamaSession.generate overloads.
• New GenerateUntilStopPrefillEquivalenceTest pins greedy equivalence (single-chunk and forced multi-chunk) through the exact entry point the agent loop uses, so the wiring cannot silently regress again.

Follow-up candidates (out of scope here):

• Consider flipping the default to Batched after a release of soak time.
• JavaAgentLoop.Builder.listener(AgentListener) — the richer structured events (onToolCalls, onPrefillProgress, …) are wired internally but not exposed to Java consumers.