From f56e84b3d9a68d7b168a4692d8eeabc70067d4ce Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Mon, 14 Sep 2026 14:04:04 -0700 Subject: [PATCH] Consolidate VLM sequential engine: shared KV cache, token box, chunked prefill (B1+B2) B1: Replace the VLM engine's hand-rolled KV cache (keyCache/valueCache/currentKVCapacity plus ensureKVCapacity/copyCache/zeroFill) with the shared StateHandlerFactory and runWithStates path used by the text engine. Growth semantics unchanged (initial 256, 2x). Restore the per-realloc "KV cache grew" log in GrowingNDArrayState. Read back hasNonTruncatableStates and guard reset(to:) against partial reset for recurrent-state models, matching the text engine. B2: Add GenerationTokenBox for the active-token busy/cancel/install bookkeeping both engines duplicated, and runChunkedPrefill() to replace both processChunkedPrompt implementations. The text engine passes a non-zero heldBack and routes non-held-back chunks through the prefill graph; the VLM engine passes heldBack = 0. Adds GenerationTokenBox and chunked-prefill unit tests. --- .../Handlers/StateHandler+NDArray.swift | 2 + .../InferenceEngines/ChunkedPrefill.swift | 52 ++++ .../CoreAISequentialEngine.swift | 75 +++--- .../CoreAISequentialVLMEngine.swift | 226 +++++------------- .../InferenceEngines/GenerationTokenBox.swift | 39 +++ .../ChunkedPrefillTests.swift | 162 +++++++++++++ .../GenerationTokenBoxTests.swift | 111 +++++++++ 7 files changed, 451 insertions(+), 216 deletions(-) create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/ChunkedPrefill.swift create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/GenerationTokenBox.swift create mode 100644 swift/Tests/LanguageModelsTests/ChunkedPrefillTests.swift create mode 100644 swift/Tests/LanguageModelsTests/GenerationTokenBoxTests.swift diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift index 0f2202f6..e56522d6 100644 --- a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift +++ b/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 } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/ChunkedPrefill.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/ChunkedPrefill.swift new file mode 100644 index 00000000..bacb279f --- /dev/null +++ b/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, + chunkSize: Int, + heldBack: Int, + vocabSize: Int, + processChunk: (_ chunk: ArraySlice, _ 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..(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.. 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) } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift index 7cd2f32b..54526733 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift @@ -10,7 +10,6 @@ import CoreAI import CoreAIShared import CoreImage import Foundation -import Synchronization // MARK: - VLM Model Config @@ -43,7 +42,7 @@ public struct VLMModelConfig: InferenceConfiguration, Codable, Sendable { /// /// ## Model Contract /// -/// Manages three model functions (potentially from separate `.aimodel` bundles): +/// Manages four model functions (potentially from separate `.aimodel` bundles): /// /// 1. **Vision encoder** (`encode_image`): /// - Input: `pixel_values` (Float32, shape `[1, 3, H, W]`) @@ -99,8 +98,6 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec // LLM I/O names from descriptor private let embeddingsInputName: String private let positionIdsName: String - private let keyCacheName: String - private let valueCacheName: String private let logitsName: String // LLM descriptors for dynamic shape resolution @@ -110,13 +107,16 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec // MARK: - Persistent State - private var keyCache: NDArray - private var valueCache: NDArray + /// KV cache state, managed by the shared state-handler infrastructure + /// (allocation, 2x growth, copy-on-grow, reset) — identical to `CoreAISequentialEngine`. + private var kvCache: any SyncStateHandler + /// Additional non-KV states (nil for the VLM's two-state KV contract; carried for symmetry). + private var additionalStates: FixedNDArrayState? + /// True if any state is non-truncatable (conv/recurrent). Gates partial `reset(to:)`, + /// matching `CoreAISequentialEngine`. False for every current full-attention VLM. + private let hasNonTruncatableStates: Bool private var logitsArray: NDArray private var cachedLogitsBatchSize: Int - private var currentKVCapacity: Int - private let keyCacheDescriptor: NDArrayDescriptor - private let valueCacheDescriptor: NDArrayDescriptor // Track processed tokens for incremental inference public private(set) var processedTokenCount: Int = 0 @@ -127,11 +127,11 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec // MARK: - Generation Token - private let _activeToken = Mutex(nil) - public var isBusy: Bool { _activeToken.withLock { $0 != nil } } + private let tokenBox = GenerationTokenBox() + public var isBusy: Bool { tokenBox.isBusy } func clearTokenIfActive(_ token: GenerationToken) { - _activeToken.withLock { if $0 === token { $0 = nil } } + tokenBox.clearIfActive(token) } // MARK: - Init @@ -252,8 +252,6 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec // Extract I/O names self.embeddingsInputName = llmDesc.inputNames[0] self.positionIdsName = llmDesc.inputNames[1] - self.keyCacheName = llmDesc.stateNames[0] - self.valueCacheName = llmDesc.stateNames[1] self.logitsName = llmDesc.outputNames[0] // Extract and validate descriptors @@ -279,34 +277,21 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec } self.logitsDescriptor = logitsDesc - // Extract KV cache state descriptors - guard case .ndArray(let keyCacheDesc) = llmDesc.stateDescriptor(of: keyCacheName), - case .ndArray(let valueCacheDesc) = llmDesc.stateDescriptor(of: valueCacheName) - else { - throw InferenceRuntimeError.invalidOutputType("Cannot get KV cache state descriptors") - } - self.keyCacheDescriptor = keyCacheDesc - self.valueCacheDescriptor = valueCacheDesc - - // Allocate KV cache - let isDynamic = keyCacheDesc.shape.contains(where: { $0 < 0 }) - let initialCapacity: Int - if options.kvCacheStrategy == .fixedSize || !isDynamic { - initialCapacity = config.maxContextLength - } else { - initialCapacity = min(256, config.maxContextLength) - } - self.currentKVCapacity = initialCapacity - - let resolvedKeyDesc = keyCacheDesc.resolvingDynamicDimensions( - keyCacheDesc.shape.map { $0 < 0 ? initialCapacity : $0 }) - let resolvedValueDesc = valueCacheDesc.resolvingDynamicDimensions( - valueCacheDesc.shape.map { $0 < 0 ? initialCapacity : $0 }) - self.keyCache = NDArray(descriptor: resolvedKeyDesc) - self.valueCache = NDArray(descriptor: resolvedValueDesc) + // Create KV cache state handler(s) from the LLM descriptor, mirroring + // CoreAISequentialEngine. The shared factory handles allocation, 2x growth, + // copy-on-grow, and reset for both dynamic (growing) and fixed-size KV caches. + let stateHandlers = try StateHandlerFactory.createSyncHandlers( + descriptor: llmDesc, + maxContextLength: config.maxContextLength, + options: options + ) + self.kvCache = stateHandlers.kvCache + self.additionalStates = stateHandlers.additionalStates + self.hasNonTruncatableStates = stateHandlers.hasNonTruncatableStates CLILogger.log( - "VLM KV cache: dynamic=\(isDynamic), initial=\(initialCapacity), key=\(keyCacheDesc.shape) -> \(resolvedKeyDesc.shape)" + "VLM KV cache: capacity=\(stateHandlers.kvCache.currentCapacity), " + + "states=\(stateHandlers.kvCache.stateNames)" ) // Allocate initial logits (1 token) @@ -702,7 +687,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec throw InferenceRuntimeError.invalidState("Cannot process empty embedding batch") } - try ensureKVCapacity(forContextLength: processedTokenCount + batchSize) + _ = try kvCache.ensureCapacity(forContextLength: processedTokenCount + batchSize) let batchSignpost = InstrumentsProfiler.beginCustomInterval( name: "CoreAIVLM Batch", @@ -723,20 +708,15 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec cachedLogitsBatchSize = batchSize } - // Build states (KV cache -- persistent, inout) - var states = InferenceFunction.MutableViews() - states.insert(&keyCache, for: keyCacheName) - states.insert(&valueCache, for: valueCacheName) - - // Build output backings (logits -- written in-place) - var outputViews = InferenceFunction.MutableViews() - outputViews.insert(&logitsArray, for: logitsName) - - // Execute LLM forward pass - _ = try await llmFunction.run( + // Bind KV cache states, build output views, and execute — shared with + // CoreAISequentialEngine via runWithStates (zero-copy state binding). + try await runWithStates( + function: llmFunction, inputs: [embeddingsInputName: embeddings, positionIdsName: positionIds], - states: consume states, - outputViews: consume outputViews + primary: kvCache, + secondary: additionalStates, + outputArray: &logitsArray, + outputName: logitsName ) // Read logits from NDArray @@ -830,27 +810,16 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec tokens: ArraySlice, chunkSize: Int ) async throws -> [LogitsScalarType] { - let totalChunks = (tokens.count + chunkSize - 1) / chunkSize - - var lastLogits: [LogitsScalarType] = [] - var remainingTokens = tokens - var chunkIndex = 0 - - while !remainingTokens.isEmpty { - let currentChunkSize = min(chunkSize, remainingTokens.count) - let chunkEnd = remainingTokens.startIndex + currentChunkSize - let chunk = remainingTokens[remainingTokens.startIndex.. GenerationSequence { - _activeToken.withLock { - $0?.cancel() - $0 = nil - } + tokenBox.cancelActive() let token = GenerationToken() - _activeToken.withLock { $0 = token } + tokenBox.install(token) return GenerationSequence( engine: self, input: input, @@ -889,12 +855,9 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions ) async throws -> GenerationSequence { - _activeToken.withLock { - $0?.cancel() - $0 = nil - } + tokenBox.cancelActive() let token = GenerationToken() - _activeToken.withLock { $0 = token } + tokenBox.install(token) return GenerationSequence( engine: self, input: tokens, @@ -915,15 +878,17 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec precondition( tokenIndex >= 0 && tokenIndex <= processedTokenCount, "reset(to: \(tokenIndex)) out of range [0, \(processedTokenCount)]") + if tokenIndex != 0 && hasNonTruncatableStates { + throw InferenceRuntimeError.invalidState( + "Partial reset is not supported for hybrid models with recurrent state. " + + "Use reset(to: 0) and replay the prefix.") + } if tokenIndex == 0 { - _activeToken.withLock { - $0?.cancel() - $0 = nil - } + tokenBox.cancelActive() let resetSpan = InstrumentsProfiler.beginReset(engine: "CoreAIVLM") processedTokenCount = 0 - zeroFill(&keyCache) - zeroFill(&valueCache) + kvCache.reset() + additionalStates?.reset() resetSpan.end() } else { processedTokenCount = tokenIndex @@ -931,10 +896,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec } public func cancel() async throws { - _activeToken.withLock { - $0?.cancel() - $0 = nil - } + tokenBox.cancelActive() } public func cleanup() { @@ -949,80 +911,8 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec _ = try await processTokenBatch(dummyTokens) // Reset state after warmup processedTokenCount = 0 - zeroFill(&keyCache) - zeroFill(&valueCache) - } - - // MARK: - KV Cache (dynamic growth) - - private func ensureKVCapacity(forContextLength needed: Int) throws { - guard needed > currentKVCapacity else { return } - guard needed <= config.maxContextLength else { - throw InferenceRuntimeError.invalidState( - "Context length \(needed) exceeds maximum \(config.maxContextLength)") - } - - var newCapacity = max(currentKVCapacity, 1) - while newCapacity < needed { newCapacity *= 2 } - newCapacity = min(newCapacity, config.maxContextLength) - - let resolvedKeyDesc = keyCacheDescriptor.resolvingDynamicDimensions( - keyCacheDescriptor.shape.map { $0 < 0 ? newCapacity : $0 }) - let resolvedValueDesc = valueCacheDescriptor.resolvingDynamicDimensions( - valueCacheDescriptor.shape.map { $0 < 0 ? newCapacity : $0 }) - - var newKeyCache = NDArray(descriptor: resolvedKeyDesc) - var newValueCache = NDArray(descriptor: resolvedValueDesc) - _ = newKeyCache.mutableRawView() - _ = newValueCache.mutableRawView() - - try Self.copyCache(from: keyCache, to: &newKeyCache) - try Self.copyCache(from: valueCache, to: &newValueCache) - - CLILogger.log("VLM KV cache grew: \(currentKVCapacity) -> \(newCapacity)") - keyCache = newKeyCache - valueCache = newValueCache - currentKVCapacity = newCapacity - } - - private static func copyCache(from source: NDArray, to destination: inout NDArray) throws { - let srcShape = source.shape - let dstShape = destination.shape - guard let headDim = srcShape.last else { - throw InferenceRuntimeError.invalidState("KV cache has empty shape -- cannot copy") - } - let seqDim = KVCacheFactory.detectSequenceDim(shape: srcShape) - - let numBlocks = srcShape[..` +/// bookkeeping so both engines share one implementation. +final class GenerationTokenBox: Sendable { + private let _token = Mutex(nil) + + /// True while a generation is in flight. + var isBusy: Bool { _token.withLock { $0 != nil } } + + /// Cancel the active token (if any) and clear it. + func cancelActive() { + _token.withLock { + $0?.cancel() + $0 = nil + } + } + + /// Install `token` as the active generation. Does not cancel any prior token — + /// callers that need to supersede an in-flight generation call `cancelActive()` first. + func install(_ token: GenerationToken) { + _token.withLock { $0 = token } + } + + /// Clear the active token only if it is `token`. Called by the iterator when + /// generation finishes or is cancelled, so a newer generation is left untouched. + func clearIfActive(_ token: GenerationToken) { + _token.withLock { if $0 === token { $0 = nil } } + } +} diff --git a/swift/Tests/LanguageModelsTests/ChunkedPrefillTests.swift b/swift/Tests/LanguageModelsTests/ChunkedPrefillTests.swift new file mode 100644 index 00000000..b3775e64 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/ChunkedPrefillTests.swift @@ -0,0 +1,162 @@ +// 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 Foundation +import Testing + +@testable import CoreAILanguageModels + +/// Covers the shared `runChunkedPrefill` loop that both sequential engines call: how a prompt +/// is split into per-chunk `processChunk` calls, which chunks are marked held-back, and which +/// chunk's logits are returned. Uses a recording closure so the loop can be exercised without a +/// live graph. +@Suite("Chunked Prefill") +struct ChunkedPrefillTests { + /// Records the token slices and held-back flags `runChunkedPrefill` hands to `processChunk`. + private final class Recorder { + var calls: [(tokens: [Int32], heldBack: Bool)] = [] + } + + /// One logit row per token, filled with the token's value so `lastTokenLogits` can be checked. + private func rows(for chunk: ArraySlice, vocabSize: Int) -> [LogitsScalarType] { + var out: [LogitsScalarType] = [] + for token in chunk { + out.append(contentsOf: Array(repeating: LogitsScalarType(Int(token)), count: vocabSize)) + } + return out + } + + // MARK: - VLM case (heldBack == 0) + + @Test("With nothing held back, every chunk runs and the last token's logits come back") + func heldBackZeroProcessesEveryChunk() async throws { + let tokens: [Int32] = Array(0..<10) + let vocabSize = 3 + let recorder = Recorder() + + let result = try await runChunkedPrefill( + tokens: tokens[...], + chunkSize: 4, + heldBack: 0, + vocabSize: vocabSize + ) { chunk, isHeldBack in + recorder.calls.append((Array(chunk), isHeldBack)) + return self.rows(for: chunk, vocabSize: vocabSize) + } + + // Contiguous chunks of width 4 covering all 10 tokens, none held back. + #expect(recorder.calls.map(\.tokens) == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]]) + #expect(recorder.calls.allSatisfy { !$0.heldBack }) + + // Last token is 9, so the returned row is [9, 9, 9]. + #expect(result == Array(repeating: LogitsScalarType(9), count: vocabSize)) + } + + @Test("With nothing held back, a single-chunk prompt runs once") + func heldBackZeroSingleChunk() async throws { + let tokens: [Int32] = [5, 6] + let recorder = Recorder() + + let result = try await runChunkedPrefill( + tokens: tokens[...], chunkSize: 8, heldBack: 0, vocabSize: 2 + ) { chunk, isHeldBack in + recorder.calls.append((Array(chunk), isHeldBack)) + return self.rows(for: chunk, vocabSize: 2) + } + + #expect(recorder.calls.map(\.tokens) == [[5, 6]]) + #expect(recorder.calls[0].heldBack == false) + #expect(result == [LogitsScalarType(6), LogitsScalarType(6)]) + } + + // MARK: - Text case (heldBack > 0) + + @Test("With one held back, earlier chunks fill KV and the tail carries the logits") + func heldBackOneRoutesTailThroughMain() async throws { + let tokens: [Int32] = Array(0..<10) + let vocabSize = 3 + let recorder = Recorder() + + let result = try await runChunkedPrefill( + tokens: tokens[...], + chunkSize: 4, + heldBack: 1, + vocabSize: vocabSize + ) { chunk, isHeldBack in + recorder.calls.append((Array(chunk), isHeldBack)) + // Prefill-graph chunks produce no logits; only the held-back tail does. + return isHeldBack ? self.rows(for: chunk, vocabSize: vocabSize) : [] + } + + // Nine tokens prefilled as [4, 4, 1] (not held back), then token 9 held back. + #expect(recorder.calls.map(\.tokens) == [[0, 1, 2, 3], [4, 5, 6, 7], [8], [9]]) + #expect(recorder.calls.map(\.heldBack) == [false, false, false, true]) + + // The held-back token 9 supplies the logits. + #expect(result == Array(repeating: LogitsScalarType(9), count: vocabSize)) + } + + @Test("With one held back and a one-token prompt, nothing is prefilled") + func heldBackOneSingleToken() async throws { + let tokens: [Int32] = [7] + let recorder = Recorder() + + let result = try await runChunkedPrefill( + tokens: tokens[...], chunkSize: 4, heldBack: 1, vocabSize: 2 + ) { chunk, isHeldBack in + recorder.calls.append((Array(chunk), isHeldBack)) + return isHeldBack ? self.rows(for: chunk, vocabSize: 2) : [] + } + + // The lone token is the held-back tail; the prefill loop never runs. + #expect(recorder.calls.map(\.tokens) == [[7]]) + #expect(recorder.calls[0].heldBack == true) + #expect(result == [LogitsScalarType(7), LogitsScalarType(7)]) + } + + @Test("With one held back, an exact multiple leaves only the tail") + func heldBackOneExactMultiple() async throws { + // 9 tokens: 8 prefilled as two width-4 chunks, token 8 held back. + let tokens: [Int32] = Array(0..<9) + let recorder = Recorder() + + _ = try await runChunkedPrefill( + tokens: tokens[...], chunkSize: 4, heldBack: 1, vocabSize: 1 + ) { chunk, isHeldBack in + recorder.calls.append((Array(chunk), isHeldBack)) + return isHeldBack ? self.rows(for: chunk, vocabSize: 1) : [] + } + + #expect(recorder.calls.map(\.tokens) == [[0, 1, 2, 3], [4, 5, 6, 7], [8]]) + #expect(recorder.calls.map(\.heldBack) == [false, false, true]) + } + + // MARK: - Coverage invariants + + @Test("Chunks are contiguous and cover the whole prompt") + func chunksCoverPrompt() async throws { + for heldBack in [0, 1] { + for count in [1, 2, 7, 8, 9, 16, 17] { + let tokens: [Int32] = Array(0..