Thread PrefillStrategy through generateUntilStop and the agent loop#238
Merged
Conversation
… the agent loop generateUntilStop kept a hardcoded autoregressive prefill long after the batched path reached parity (the head-reshape divergence was fixed by swapSeqHeadDims, RoPE is batch-aware, and the CPU SDPA causal mask is cache-aware: maxKi = seqKV - seqQ + qi). llm-core's generate() already exposes PrefillStrategy; the agent/session path never got it - so every kllama consumer (AgentLoop, JavaAgentLoop, KLlamaSession) was pinned to one-forward-per-prompt-token prefill. - generateUntilStop gains prefillStrategy (default Autoregressive, unchanged behavior); Batched ingests the prompt in maxBatch chunks via forwardBatched, reporting onPrefill per chunk - AgentConfig gains prefillStrategy, threaded through both AgentLoop call sites (JavaAgentLoop inherits via AgentConfig) - GenerationConfig gains prefillStrategy with a Java-friendly batchedPrefill(maxBatch) builder knob, threaded through both KLlamaSession.generate overloads - GenerateUntilStopPrefillEquivalenceTest pins greedy equivalence of batched (single-chunk and forced multi-chunk) vs autoregressive prefill through the exact entry point the agent loop uses Motivation: agent-loop consumers on CPU runtimes spend most of their wall time in prefill - Daily-StandAPP (JavaLand 2026 demo, https://github.com/michalharakal/Daily-StandAPP) measures 87s to first token for a ~700-token tool-calling template on an M-series CPU; batched prefill typically cuts prompt ingestion 3-10x. Refs #226. Equivalence verified: BatchedPrefillEquivalenceTest, PrefillStrategyEquivalenceTest, and the new GenerateUntilStopPrefillEquivalenceTest all pass on TinyLlama 1.1B Q8_0 (greedy sequences identical, single-chunk and forced multi-chunk).
…ublic signatures AgentConfig.prefillStrategy and GenerationConfig.prefillStrategy expose llm-core's PrefillStrategy type in public API, but llm-agent and llm-runtime/kllama declared llm-core as implementation — consumers of the published artifacts could not resolve the type (first hit by Daily-StandAPP wiring up batched prefill: 'Unresolved reference llm' / 'No parameter with name prefillStrategy'). Refs #226.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #226.
What
Opt-in batched prefill for the agent/session path. Default behavior is unchanged (
PrefillStrategy.Autoregressive).generateUntilStop(...)gainsprefillStrategy: PrefillStrategy = Autoregressive;Batched(maxBatch)ingests the prompt in chunks viaforwardBatched, reportingonPrefillper chunk instead of per tokenAgentConfiggainsprefillStrategy, threaded through bothAgentLoopcall sites (JavaAgentLoopinherits it viaAgentConfig— no facade change)GenerationConfiggains a Java-friendlybatchedPrefill(maxBatch)builder knob, threaded through bothKLlamaSession.generateoverloadsllm-agentandllm-runtime/kllamanow exposellm-coreasapi—PrefillStrategyappears in their public signatures, and published-artifact consumers could not resolve it underimplementationscope (first hit by Daily-StandAPP)Why
generateUntilStopkept a hardcoded autoregressive prefill (with a comment pointing at the oldposition_collapse_bug.mddivergence) long after the underlying issues were fixed — the head-reshape divergence byswapSeqHeadDims, batch-aware RoPE, and the cache-aware SDPA causal mask (maxKi = seqKV - seqQ + qi).llm-core'sgenerate()already hasPrefillStrategy; the agent path never got it, pinning every kllama consumer to one forward pass per prompt token.Motivating consumer: Daily-StandAPP (JavaLand 2026 demo) measures 225–450 s per generation on CPU, dominated by prefill; the agent loop re-processes the full conversation every round, multiplying the cost.
Testing
New
GenerateUntilStopPrefillEquivalenceTestpins greedy-sequence equivalence of batched vs autoregressive prefill through the exact entry point the agent loop uses — single-chunk and forced multi-chunk (maxBatch=3). All four prefill equivalence methods pass on real TinyLlama 1.1B Q8_0 weights (this test plus the existingPrefillStrategyEquivalenceTest), alongside the logit-levelBatchedPrefillEquivalenceTest(max_abs_diff = 5.7e-6).Note for CI: these model-loading tests want a quiet machine — two concurrent FP32 model loads alongside another large JVM reproducibly killed the test JVM before any assertion ran. Also includes
PrefillTimingDiagnostic, a manual (non-asserting) wall-clock probe of autoregressive vs batched prefill at several chunk sizes.Follow-ups (out of scope)
Batchedafter a release of soak timeJavaAgentLoop.Builder.listener(AgentListener)— structured events (onToolCalls,onPrefillProgress, …) are wired internally but not exposed to Java consumersruntime.reset()+ full re-prefill)🤖 Generated with Claude Code