From 7f9a695ae4ddda063ce9c0b85387e40cd5fdc75f Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 13 Jul 2026 17:04:12 +0200 Subject: [PATCH 1/2] feat(llm-agent): thread PrefillStrategy through generateUntilStop and 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). --- .../apps/kllama/agent/GenerateExtensions.kt | 52 +++++--- .../sk/ainet/apps/kllama/chat/AgentLoop.kt | 16 ++- .../apps/kllama/java/GenerationConfig.kt | 28 ++++- .../ainet/apps/kllama/java/KLlamaSession.kt | 6 +- ...GenerateUntilStopPrefillEquivalenceTest.kt | 113 ++++++++++++++++++ 5 files changed, 188 insertions(+), 27 deletions(-) create mode 100644 llm-runtime/kllama/src/jvmTest/kotlin/sk/ainet/apps/kllama/GenerateUntilStopPrefillEquivalenceTest.kt diff --git a/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/agent/GenerateExtensions.kt b/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/agent/GenerateExtensions.kt index ed8ee172..72057b58 100644 --- a/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/agent/GenerateExtensions.kt +++ b/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/agent/GenerateExtensions.kt @@ -2,6 +2,7 @@ package sk.ainet.apps.kllama.agent import kotlin.random.Random import sk.ainet.apps.llm.InferenceRuntime +import sk.ainet.apps.llm.PrefillStrategy import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.FloatArrayTensorData import sk.ainet.lang.types.DType @@ -22,12 +23,20 @@ import sk.ainet.lang.types.DType * @param random Random generator for sampling. * @param onToken Optional callback invoked for each generated token. * @param decode Optional function to decode a token ID to a string. - * @param onPrefill Optional callback invoked after each prompt token is fed - * through `forward`, with `(done, total)` where `done` counts the tokens - * already processed (1-based) and `total` is `prompt.size`. Useful for - * showing progress: prefill is autoregressive (one `forward` per prompt - * token) and on a CPU-only runtime can dominate wall time before the first - * generated token appears. + * @param onPrefill Optional callback invoked as the prompt is ingested, with + * `(done, total)` where `done` counts the tokens already processed and + * `total` is `prompt.size`. Under [PrefillStrategy.Autoregressive] it fires + * once per token; under [PrefillStrategy.Batched] once per chunk. Useful + * for showing progress: on a CPU-only runtime prefill can dominate wall + * time before the first generated token appears. + * @param prefillStrategy How to ingest the prompt. The default + * [PrefillStrategy.Autoregressive] preserves the historical one-token-per- + * forward behavior. [PrefillStrategy.Batched] feeds the prompt through + * [InferenceRuntime.forwardBatched] in chunks — typically 3–10× faster on + * long prompts. Equivalence with the autoregressive path is pinned by + * `BatchedPrefillEquivalenceTest` / `GenerateUntilStopPrefillEquivalenceTest` + * (the historical position-collapse and head-reshape divergences are fixed; + * see `swapSeqHeadDims` in MultiHeadAttention and the batch-aware RoPE). */ public fun InferenceRuntime.generateUntilStop( prompt: IntArray, @@ -37,25 +46,30 @@ public fun InferenceRuntime.generateUntilStop( random: Random = Random.Default, onToken: ((Int) -> Unit)? = null, decode: ((Int) -> String)? = null, - onPrefill: ((Int, Int) -> Unit)? = null + onPrefill: ((Int, Int) -> Unit)? = null, + prefillStrategy: PrefillStrategy = PrefillStrategy.Autoregressive ): GenerateResult { - // Feed prompt tokens through the model. NOTE: previously this used - // `forwardBatched(prompt)` for ~5–10× faster prefill, but - // `OptimizedLLMRuntime.forwardBatched` has known correctness issues - // at N>1 in our DSL (the batched prefill path produces different - // logits than the equivalent autoregressive sequence — see - // `position_collapse_bug.md` and the post-RoPE-fix smoke-test - // results: model output goes from a clean `<|tool_call>` to garbage - // prose). Reverted to autoregressive prefill until forwardBatched - // is brought to parity with autoregressive on real-model output. if (prompt.isEmpty()) { return GenerateResult(emptyList(), "", false) } val total = prompt.size var lastLogits: Tensor? = null - for ((i, tokenId) in prompt.withIndex()) { - lastLogits = forward(tokenId) - onPrefill?.invoke(i + 1, total) + when (prefillStrategy) { + is PrefillStrategy.Autoregressive -> { + for ((i, tokenId) in prompt.withIndex()) { + lastLogits = forward(tokenId) + onPrefill?.invoke(i + 1, total) + } + } + is PrefillStrategy.Batched -> { + var i = 0 + while (i < total) { + val end = minOf(i + prefillStrategy.maxBatch, total) + lastLogits = forwardBatched(prompt.copyOfRange(i, end)) + i = end + onPrefill?.invoke(i, total) + } + } } lastLogits ?: return GenerateResult(emptyList(), "", false) diff --git a/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/chat/AgentLoop.kt b/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/chat/AgentLoop.kt index c1ec2929..dbb236a2 100644 --- a/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/chat/AgentLoop.kt +++ b/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/chat/AgentLoop.kt @@ -3,6 +3,7 @@ package sk.ainet.apps.kllama.chat import kotlin.random.Random import sk.ainet.apps.kllama.agent.GenerateResult import sk.ainet.apps.llm.InferenceRuntime +import sk.ainet.apps.llm.PrefillStrategy import sk.ainet.apps.kllama.agent.generateUntilStop import sk.ainet.lang.types.DType @@ -13,12 +14,19 @@ import sk.ainet.lang.types.DType * @param maxTokensPerRound Maximum tokens to generate in each round. * @param temperature Sampling temperature. * @param random Random generator for sampling. + * @param prefillStrategy How each round's prompt is ingested. The agent loop + * re-renders and re-processes the FULL conversation every round, so on + * CPU-only runtimes prefill dominates round latency — + * [PrefillStrategy.Batched] typically cuts it 3–10× on long chat + * templates. Default stays [PrefillStrategy.Autoregressive] for unchanged + * behavior. */ public data class AgentConfig @kotlin.jvm.JvmOverloads constructor( val maxToolRounds: Int = 5, val maxTokensPerRound: Int = 512, val temperature: Float = 0.7f, - val random: Random = Random.Default + val random: Random = Random.Default, + val prefillStrategy: PrefillStrategy = PrefillStrategy.Autoregressive ) /** @@ -125,7 +133,8 @@ public class AgentLoop( random = config.random, onToken = { tokenId -> listener?.onToken(decode(tokenId)) }, decode = decode, - onPrefill = { done, total -> listener?.onPrefillProgress(done, total) } + onPrefill = { done, total -> listener?.onPrefillProgress(done, total) }, + prefillStrategy = config.prefillStrategy ) lastResponse = result.text @@ -247,7 +256,8 @@ public class AgentLoop( random = config.random, onToken = { tokenId -> listener?.onToken(decode(tokenId)) }, decode = decode, - onPrefill = { done, total -> listener?.onPrefillProgress(done, total) } + onPrefill = { done, total -> listener?.onPrefillProgress(done, total) }, + prefillStrategy = config.prefillStrategy ) lastResponse = result.text diff --git a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/GenerationConfig.kt b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/GenerationConfig.kt index ef69d992..0f4a9b9c 100644 --- a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/GenerationConfig.kt +++ b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/GenerationConfig.kt @@ -1,5 +1,7 @@ package sk.ainet.apps.kllama.java +import sk.ainet.apps.llm.PrefillStrategy + /** * Configuration for text generation, used by KLlamaJava facades. * @@ -8,18 +10,20 @@ package sk.ainet.apps.kllama.java * GenerationConfig config = GenerationConfig.builder() * .maxTokens(256) * .temperature(0.7f) + * .batchedPrefill(64) // opt-in: 3-10x faster prompt ingestion * .build(); * ``` */ public class GenerationConfig private constructor( public val maxTokens: Int, - public val temperature: Float + public val temperature: Float, + public val prefillStrategy: PrefillStrategy ) { public companion object { @JvmStatic public fun builder(): Builder = Builder() - /** Default configuration: 256 max tokens, temperature 0.8. */ + /** Default configuration: 256 max tokens, temperature 0.8, autoregressive prefill. */ @JvmStatic public fun defaults(): GenerationConfig = Builder().build() } @@ -27,6 +31,7 @@ public class GenerationConfig private constructor( public class Builder { private var maxTokens: Int = 256 private var temperature: Float = 0.8f + private var prefillStrategy: PrefillStrategy = PrefillStrategy.Autoregressive public fun maxTokens(maxTokens: Int): Builder { this.maxTokens = maxTokens @@ -38,6 +43,23 @@ public class GenerationConfig private constructor( return this } - public fun build(): GenerationConfig = GenerationConfig(maxTokens, temperature) + /** + * Ingest the prompt via chunked `forwardBatched` calls instead of one + * `forward` per token — typically 3–10× faster prefill on long + * prompts, numerically equivalent (see the prefill equivalence tests). + */ + @JvmOverloads + public fun batchedPrefill(maxBatch: Int = 64): Builder { + this.prefillStrategy = PrefillStrategy.Batched(maxBatch) + return this + } + + /** Set an explicit [PrefillStrategy] (Kotlin-friendly variant). */ + public fun prefillStrategy(strategy: PrefillStrategy): Builder { + this.prefillStrategy = strategy + return this + } + + public fun build(): GenerationConfig = GenerationConfig(maxTokens, temperature, prefillStrategy) } } diff --git a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/KLlamaSession.kt b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/KLlamaSession.kt index fac24515..abc20748 100644 --- a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/KLlamaSession.kt +++ b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/java/KLlamaSession.kt @@ -52,7 +52,8 @@ public class KLlamaSession( maxTokens = config.maxTokens, eosTokenId = eosTokenId, temperature = config.temperature, - decode = { tokenizer.decode(it) } + decode = { tokenizer.decode(it) }, + prefillStrategy = config.prefillStrategy ) return result.text } @@ -75,7 +76,8 @@ public class KLlamaSession( eosTokenId = eosTokenId, temperature = config.temperature, onToken = { tokenId -> tokenConsumer.accept(tokenizer.decode(tokenId)) }, - decode = { tokenizer.decode(it) } + decode = { tokenizer.decode(it) }, + prefillStrategy = config.prefillStrategy ) return result.text } diff --git a/llm-runtime/kllama/src/jvmTest/kotlin/sk/ainet/apps/kllama/GenerateUntilStopPrefillEquivalenceTest.kt b/llm-runtime/kllama/src/jvmTest/kotlin/sk/ainet/apps/kllama/GenerateUntilStopPrefillEquivalenceTest.kt new file mode 100644 index 00000000..e2cb1a5a --- /dev/null +++ b/llm-runtime/kllama/src/jvmTest/kotlin/sk/ainet/apps/kllama/GenerateUntilStopPrefillEquivalenceTest.kt @@ -0,0 +1,113 @@ +package sk.ainet.apps.kllama + +import java.nio.file.Path +import kotlin.io.path.exists +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlinx.coroutines.runBlocking +import sk.ainet.apps.kllama.agent.GenerateResult +import sk.ainet.apps.kllama.agent.generateUntilStop +import sk.ainet.apps.llm.OptimizedLLMMode +import sk.ainet.apps.llm.OptimizedLLMRuntime +import sk.ainet.apps.llm.PrefillStrategy +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.model.QuantPolicy +import sk.ainet.lang.types.FP32 +import sk.ainet.models.llama.LlamaNetworkLoader + +/** + * Verifies that `generateUntilStop(prefillStrategy = Batched)` produces the + * same greedy token sequence as the autoregressive default. This is the + * agent-loop-level counterpart to [PrefillStrategyEquivalenceTest] (which + * covers llm-core's `generate`): `generateUntilStop` is what [sk.ainet.apps.kllama.chat.AgentLoop] + * and `KLlamaSession.generate` actually call, and it historically kept a + * hardcoded autoregressive prefill long after the batched path reached + * parity — this test pins the wiring so it cannot silently regress. + * + * Both the single-chunk and the forced multi-chunk (`maxBatch` < prompt + * length) paths are covered in ONE model pairing per test to keep the peak + * footprint at two loaded models per JVM. + * + * Skipped if the model is not present. + */ +class GenerateUntilStopPrefillEquivalenceTest { + + companion object { + private val MODEL_PATH = Path.of( + System.getProperty("user.home"), + ".lmstudio/models/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", + "tinyllama-1.1b-chat-v1.0.Q8_0.gguf" + ) + } + + private suspend fun greedyRun(prefillStrategy: PrefillStrategy, promptTokens: IntArray, maxTokens: Int): GenerateResult { + val ctx = DirectCpuExecutionContext() + val model = LlamaNetworkLoader.fromGguf( + randomAccessProvider = { JvmRandomAccessSource.open(MODEL_PATH.toString()) }, + quantPolicy = QuantPolicy.DEQUANTIZE_TO_FP32 + ).load(ctx) + val runtime = OptimizedLLMRuntime( + model = model, ctx = ctx, + mode = OptimizedLLMMode.DIRECT, dtype = FP32::class + ) + return runtime.generateUntilStop( + prompt = promptTokens, + maxTokens = maxTokens, + eosTokenId = 2, // llama + temperature = 0f, + random = Random(0), + prefillStrategy = prefillStrategy + ) + } + + @Test + fun `Batched prefill yields identical greedy generateUntilStop output`() { + if (!MODEL_PATH.exists()) { + println("[skip] Model not at $MODEL_PATH") + return + } + runBlocking { + val tokenizer = JvmRandomAccessSource.open(MODEL_PATH.toString()).use { source -> + GGUFTokenizer.fromRandomAccessSource(source) + } + val promptTokens = tokenizer.encode("The capital of France is") + val maxTokens = 12 + + val auto = greedyRun(PrefillStrategy.Autoregressive, promptTokens, maxTokens) + val batched = greedyRun(PrefillStrategy.Batched(maxBatch = 64), promptTokens, maxTokens) + + println("[diag] auto : ${auto.tokens}") + println("[diag] batched : ${batched.tokens}") + assertEquals(auto.tokens, batched.tokens, + "Batched generateUntilStop diverges from autoregressive at temperature=0") + assertEquals(auto.stoppedByEos, batched.stoppedByEos) + } + } + + @Test + fun `Chunked batched prefill yields identical greedy generateUntilStop output`() { + if (!MODEL_PATH.exists()) { + println("[skip] Model not at $MODEL_PATH") + return + } + runBlocking { + val tokenizer = JvmRandomAccessSource.open(MODEL_PATH.toString()).use { source -> + GGUFTokenizer.fromRandomAccessSource(source) + } + val promptTokens = tokenizer.encode("The capital of France is Paris and") + val maxTokens = 8 + + val auto = greedyRun(PrefillStrategy.Autoregressive, promptTokens, maxTokens) + // maxBatch < prompt length forces the multi-chunk path, i.e. + // forwardBatched with a non-empty KV cache (seqKV > seqQ). + val chunked = greedyRun(PrefillStrategy.Batched(maxBatch = 3), promptTokens, maxTokens) + + println("[diag] auto (chunked) : ${auto.tokens}") + println("[diag] batched (chunked): ${chunked.tokens}") + assertEquals(auto.tokens, chunked.tokens, + "Chunked batched generateUntilStop diverges from autoregressive at temperature=0") + } + } +} From 66c4db337459c61770fa719e0011db839aee9d7d Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 13 Jul 2026 18:01:57 +0200 Subject: [PATCH 2/2] fix(build): expose llm-core as api where PrefillStrategy leaks into public signatures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- llm-agent/build.gradle.kts | 5 ++++- llm-runtime/kllama/build.gradle.kts | 4 +++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/llm-agent/build.gradle.kts b/llm-agent/build.gradle.kts index 438ed1ec..ac63b566 100644 --- a/llm-agent/build.gradle.kts +++ b/llm-agent/build.gradle.kts @@ -45,7 +45,10 @@ kotlin { commonMain.dependencies { implementation(project.dependencies.platform(project(":llm-bom"))) implementation(libs.skainet.lang.core) - implementation(project(":llm-core")) + // api, not implementation: AgentConfig.prefillStrategy exposes + // llm-core's PrefillStrategy in a public signature — consumers + // must be able to resolve it (see #226). + api(project(":llm-core")) implementation(libs.kotlinx.serialization.json) } diff --git a/llm-runtime/kllama/build.gradle.kts b/llm-runtime/kllama/build.gradle.kts index 120c9a28..8151a144 100644 --- a/llm-runtime/kllama/build.gradle.kts +++ b/llm-runtime/kllama/build.gradle.kts @@ -65,7 +65,9 @@ kotlin { implementation(project(":llm-inference:llama")) implementation(project(":llm-inference:qwen")) implementation(project(":llm-agent")) - implementation(project(":llm-core")) + // api, not implementation: GenerationConfig.prefillStrategy exposes + // llm-core's PrefillStrategy in a public signature (see #226). + api(project(":llm-core")) implementation(libs.skainet.lang.core) implementation(libs.skainet.compile.core) implementation(libs.skainet.backend.cpu)