Consolidate VLM sequential engine: shared KV cache, token box, chunked prefill (B1+B2)

swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift

@@ -108,6 +108,7 @@ public final class GrowingNDArrayState: SyncStateHandler {
             newCapacity = min(newCapacity * 2, maxCapacity)
         }
 
+        let previousCapacity = currentCapacity
         for (i, name) in stateNames.enumerated() {
             let desc = descriptors[i]
             let newShape = desc.shape.map { $0 < 0 ? newCapacity : $0 }
@@ -119,6 +120,7 @@ public final class GrowingNDArrayState: SyncStateHandler {
         }
 
         currentCapacity = newCapacity
+        CLILogger.log("KV cache grew: \(previousCapacity) -> \(newCapacity)")
         return true
     }
 

swift/Sources/CoreAILanguageModels/InferenceEngines/ChunkedPrefill.swift

@@ -0,0 +1,52 @@
+// Copyright 2026 Apple Inc.
+//
+// Use of this source code is governed by a BSD-3-clause license that can
+// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+import CoreAI
+
+/// Shared chunked-prefill loop for the sequential engines.
+///
+/// Splits `tokens` into chunks of `chunkSize`, holding back the final `heldBack` tokens
+/// (see `prefillHeldBackTokens`) for the logits-producing pass. Each chunk is handed to
+/// `processChunk`, which returns that chunk's logit buffer — or `[]` when the chunk produces
+/// no logits (e.g. a KV-only prefill-graph chunk). Returns the last-token logits.
+///
+/// - `CoreAISequentialEngine` passes `heldBack = prefillHeldBackTokens(hasPrefillGraph:)` and a
+///   closure that routes non-held-back chunks through the prefill graph (returning `[]`) when one
+///   exists, otherwise through `main`.
+/// - `CoreAISequentialVLMEngine` passes `heldBack = 0` and a closure that always runs the decoder,
+///   so the loop collapses to a plain chunked pass.
+func runChunkedPrefill(
+    tokens: ArraySlice<Int32>,
+    chunkSize: Int,
+    heldBack: Int,
+    vocabSize: Int,
+    processChunk: (_ chunk: ArraySlice<Int32>, _ isHeldBack: Bool) async throws -> [LogitsScalarType]
+) async throws -> [LogitsScalarType] {
+    let plan = prefillChunkSizes(tokenCount: tokens.count, chunkSize: chunkSize, heldBack: heldBack)
+    let trailing = tokens.count - plan.reduce(0, +)
+    let total = plan.count + (trailing > 0 ? 1 : 0)
+
+    var lastLogits: [LogitsScalarType] = []
+    var remaining = tokens
+    var index = 0
+
+    for size in plan {
+        let chunkEnd = remaining.startIndex + size
+        let chunk = remaining[remaining.startIndex..<chunkEnd]
+        CLILogger.log("Prefill chunk \(index + 1)/\(total): \(chunk.count) tokens")
+        let logits = try await processChunk(chunk, false)
+        if !logits.isEmpty { lastLogits = logits }
+        remaining = remaining[chunkEnd...]
+        index += 1
+    }
+
+    // Held-back tail (the tokens `prefillChunkSizes` left for the logits-producing pass).
+    if !remaining.isEmpty {
+        CLILogger.log("Prefill chunk \(index + 1)/\(total): \(remaining.count) tokens (held back)")
+        lastLogits = try await processChunk(remaining, true)
+    }
+
+    return lastTokenLogits(from: lastLogits, vocabSize: vocabSize)
+}

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift

@@ -6,7 +6,6 @@
 import CoreAI
 import CoreAIShared
 import Foundation
-import Synchronization
 
 // MARK: - Prefill Strategy
 
@@ -72,14 +71,14 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
     public private(set) var lastPrefixHitCount: Int = 0
 
     // Track in-flight generation via token (replaces simple bool lock)
-    private let _activeToken = Mutex<GenerationToken?>(nil)
+    private let tokenBox = GenerationTokenBox()
 
-    public var isBusy: Bool { _activeToken.withLock { $0 != nil } }
+    public var isBusy: Bool { tokenBox.isBusy }
 
     /// Clear the engine's active token if it matches the given token.
     /// Called by the iterator when generation finishes or is cancelled.
     func clearTokenIfActive(_ token: GenerationToken) {
-        _activeToken.withLock { if $0 === token { $0 = nil } }
+        tokenBox.clearIfActive(token)
     }
 
     // MARK: - Init
@@ -304,44 +303,33 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
         // The prefill graph produces no logits, so hold the final token back for
         // `function`: it is the one whose logits seed sampling. Without one, nothing is
         // held back and the trailing chunk carries the logits.
-        let floor = prefillHeldBackTokens(hasPrefillGraph: prefillFunction != nil)
-        let plan = prefillChunkSizes(
-            tokenCount: tokens.count, chunkSize: chunkSize, heldBack: floor)
+        let heldBack = prefillHeldBackTokens(hasPrefillGraph: prefillFunction != nil)
 
         let chunkSignpost = InstrumentsProfiler.beginCustomInterval(
             name: "CoreAIClean Chunked Prefill",
-            details: "\(tokens.count) tokens in \(plan.count) chunks of \(chunkSize)"
+            details: "\(tokens.count) tokens, chunkSize \(chunkSize)"
         )
-
-        var lastLogits: [LogitsScalarType] = []
-        var remainingTokens = tokens
-
-        for (chunkIndex, currentChunkSize) in plan.enumerated() {
-            let chunkEnd = remainingTokens.startIndex + currentChunkSize
-            let chunk = remainingTokens[remainingTokens.startIndex..<chunkEnd]
-
-            CLILogger.log(
-                "Chunk \(chunkIndex + 1)/\(plan.count): \(chunk.count) tokens at position \(processedTokenCount)"
+        defer {
+            InstrumentsProfiler.endCustomInterval(
+                name: "CoreAIClean Chunked Prefill",
+                signpostID: chunkSignpost
             )
-
-            if let prefillFn = prefillFunction {
-                try await encodePrefillChunk(chunk, using: prefillFn)
-            } else {
-                lastLogits = try await processTokenBatch(chunk)
-            }
-            remainingTokens = remainingTokens[chunkEnd...]
         }
 
-        if !remainingTokens.isEmpty {
-            lastLogits = try await processTokenBatch(remainingTokens)
+        return try await runChunkedPrefill(
+            tokens: tokens,
+            chunkSize: chunkSize,
+            heldBack: heldBack,
+            vocabSize: config.vocabSize
+        ) { chunk, isHeldBack in
+            // Held-back tail (and every chunk when there is no prefill graph) runs through
+            // `main` for logits; earlier chunks fill the KV cache via the prefill graph.
+            if !isHeldBack, let prefillFn = self.prefillFunction {
+                try await self.encodePrefillChunk(chunk, using: prefillFn)
+                return []
+            }
+            return try await self.processTokenBatch(chunk)
         }
-
-        InstrumentsProfiler.endCustomInterval(
-            name: "CoreAIClean Chunked Prefill",
-            signpostID: chunkSignpost
-        )
-
-        return lastTokenLogits(from: lastLogits, vocabSize: config.vocabSize)
     }
 
     /// Process tokens in chunks, returning ALL position logits (not just last token).
@@ -373,10 +361,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
         inferenceOptions: InferenceOptions
     ) async throws -> GenerationSequence {
         // Cancel any prior generation so its Iterator stops on next poll.
-        _activeToken.withLock {
-            $0?.cancel()
-            $0 = nil
-        }
+        tokenBox.cancelActive()
 
         // Implicit prefix caching: resolve input against history.
         // For hybrid models with recurrent states, we must full-reset on any
@@ -406,7 +391,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
         }
 
         let token = GenerationToken()
-        _activeToken.withLock { $0 = token }
+        tokenBox.install(token)
         return GenerationSequence(
             engine: self,
             input: input,
@@ -421,7 +406,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
     /// Wait for any in-flight generate() Task to finish.
     private func drain() {
         var attempts = 0
-        while _activeToken.withLock({ $0 != nil }) {
+        while tokenBox.isBusy {
             attempts += 1
             if attempts > 5000 {
                 fatalError("Sequential engine drain() timeout — generation Task stuck?")
@@ -431,10 +416,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
     }
 
     public func cancel() async throws {
-        _activeToken.withLock {
-            $0?.cancel()
-            $0 = nil
-        }
+        tokenBox.cancelActive()
     }
 
     public func reset(to tokenIndex: Int) async throws {
@@ -446,10 +428,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
                 "Partial reset is not supported for hybrid models with recurrent state. "
                     + "Use reset(to: 0) and replay the prefix.")
         }
-        _activeToken.withLock {
-            $0?.cancel()
-            $0 = nil
-        }
+        tokenBox.cancelActive()
         internalReset(to: tokenIndex)
     }