Thread PrefillStrategy through generateUntilStop and the agent loop

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)
         }
 

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 <T : DType> InferenceRuntime<T>.generateUntilStop(
     prompt: IntArray,
@@ -37,25 +46,30 @@ public fun <T : DType> InferenceRuntime<T>.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<T, Float>? = 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)
 

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<T : DType>(
                 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<T : DType>(
                 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

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)

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,25 +10,28 @@ 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()
     }
 
     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)
     }
 }

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
     }

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<FP32, Float>(ctx)
+        val runtime = OptimizedLLMRuntime(
+            model = model, ctx = ctx,
+            mode = OptimizedLLMMode.DIRECT, dtype = FP32::class
+        )
+        return runtime.generateUntilStop(
+            prompt = promptTokens,
+            maxTokens = maxTokens,
+            eosTokenId = 2, // llama </s>
+            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")
+        }
+    }
+}