diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 5e07b18d..e09a2d5a 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -528,14 +528,12 @@ extension CoreAISequentialEngine.GenerationSequence { self.generationToken = generationToken self.inputTokens = input self.generationStartOffset = input.count - if let forced = inferenceOptions.forcedContinuation { - self.maxTokens = forced.count - } else { - self.maxTokens = Swift.min( - inferenceOptions.maxTokens ?? Int.max, - Swift.max(0, engine.config.maxContextLength - input.count) - ) - } + self.maxTokens = SequentialIterator.clampMaxTokens( + requested: inferenceOptions.maxTokens, + forcedCount: inferenceOptions.forcedContinuation?.count, + inputCount: input.count, + maxContextLength: engine.config.maxContextLength + ) } deinit { @@ -655,14 +653,13 @@ extension CoreAISequentialEngine.GenerationSequence { return nil } - let nextToken: Int32 - if let forced = forcedContinuation { - nextToken = forced[step] - } else { - var mutableLogits = logitBuffer - nextToken = samplingConfiguration.fallbackSampler( - from: &mutableLogits, tokenHistory: inputTokens[generationStartOffset...]) - } + let nextToken = SequentialIterator.nextToken( + fromLogits: logitBuffer, + forced: forcedContinuation, + step: step, + sampling: samplingConfiguration, + tokenHistory: inputTokens[generationStartOffset...] + ) inputTokens.append(nextToken) step += 1 diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift index 54526733..bf332fcd 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift @@ -3,7 +3,6 @@ // 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 -// TODO: Refactor to re-use common components with CoreAISequentialEngine // TODO: Add pipelined engine variant for higher throughput import CoreAI @@ -67,8 +66,9 @@ public struct VLMModelConfig: InferenceConfiguration, Codable, Sendable { /// 2. `generate(with: InputEmbeddings, tokens:, ...)` — embed tokens, scatter-merge with vision /// embeddings at placeholder positions, run LLM prefill, then standard autoregressive decode /// -/// KV cache is managed identically to `CoreAISequentialEngine`: starts small and grows -/// dynamically with 2x expansion. +/// Shares its decode-loop machinery with `CoreAISequentialEngine` (KV cache, `GenerationTokenBox`, +/// `runChunkedPrefill`, `SequentialIterator`). VLM-specific: the vision/embed pipeline, +/// scatter-merge, and the embeddings-input LLM contract. public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchecked Sendable { public typealias ConfigType = VLMModelConfig public typealias OutputSequence = GenerationSequence @@ -886,9 +886,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec if tokenIndex == 0 { tokenBox.cancelActive() let resetSpan = InstrumentsProfiler.beginReset(engine: "CoreAIVLM") - processedTokenCount = 0 - kvCache.reset() - additionalStates?.reset() + clearGenerationState() resetSpan.end() } else { processedTokenCount = tokenIndex @@ -910,6 +908,11 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec let dummyTokens: ArraySlice = [Int32(1)][...] _ = try await processTokenBatch(dummyTokens) // Reset state after warmup + clearGenerationState() + } + + /// Rewind to an empty context: clear the token cursor and zero the persistent KV state. + private func clearGenerationState() { processedTokenCount = 0 kvCache.reset() additionalStates?.reset() @@ -993,14 +996,12 @@ extension CoreAISequentialVLMEngine.GenerationSequence { self.inputTokens = input self.generationStartOffset = input.count self.embeddedInput = embeddedInput - if let forced = inferenceOptions.forcedContinuation { - self.maxTokens = forced.count - } else { - self.maxTokens = Swift.min( - inferenceOptions.maxTokens ?? Int.max, - Swift.max(0, engine.config.maxContextLength - input.count) - ) - } + self.maxTokens = SequentialIterator.clampMaxTokens( + requested: inferenceOptions.maxTokens, + forcedCount: inferenceOptions.forcedContinuation?.count, + inputCount: input.count, + maxContextLength: engine.config.maxContextLength + ) } deinit { @@ -1075,14 +1076,13 @@ extension CoreAISequentialVLMEngine.GenerationSequence { } // Sample next token - let nextToken: Int32 - if let forced = forcedContinuation { - nextToken = forced[step] - } else { - var mutableLogits = logitBuffer - nextToken = samplingConfiguration.fallbackSampler( - from: &mutableLogits, tokenHistory: inputTokens[generationStartOffset...]) - } + let nextToken = SequentialIterator.nextToken( + fromLogits: logitBuffer, + forced: forcedContinuation, + step: step, + sampling: samplingConfiguration, + tokenHistory: inputTokens[generationStartOffset...] + ) inputTokens.append(nextToken) step += 1 diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/SequentialIterator.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/SequentialIterator.swift new file mode 100644 index 00000000..64100c0a --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/SequentialIterator.swift @@ -0,0 +1,45 @@ +// 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 + +/// Token-count clamp and next-token selection shared by the sequential engines' iterators. +/// Each engine keeps its own `next()` control flow; only these identical pure-value helpers +/// live here. +enum SequentialIterator { + /// The generation-length cap for an iterator. + /// + /// A forced continuation replays exactly its own tokens; otherwise the requested budget + /// (`nil` meaning "unbounded") is clamped to the context left after the prompt. + static func clampMaxTokens( + requested: Int?, + forcedCount: Int?, + inputCount: Int, + maxContextLength: Int + ) -> Int { + if let forcedCount { + return forcedCount + } + return min(requested ?? Int.max, max(0, maxContextLength - inputCount)) + } + + /// Select the next token: the forced-continuation token when replaying, otherwise the + /// sampler's choice. + /// + /// `logits` is taken by value; the sampler mutates a copy-on-write copy so the caller's + /// buffer (which it may also return to the consumer) is left untouched — matching the + /// engines' previous `var mutableLogits = logitBuffer` behavior. + static func nextToken( + fromLogits logits: [LogitsScalarType], + forced: [Int32]?, + step: Int, + sampling: SamplingConfiguration, + tokenHistory: ArraySlice + ) -> Int32 { + if let forced { + return forced[step] + } + var mutableLogits = logits + return sampling.fallbackSampler(from: &mutableLogits, tokenHistory: tokenHistory) + } +} diff --git a/swift/Tests/LanguageModelsTests/SequentialIteratorTests.swift b/swift/Tests/LanguageModelsTests/SequentialIteratorTests.swift new file mode 100644 index 00000000..d7b7eaa7 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/SequentialIteratorTests.swift @@ -0,0 +1,44 @@ +// 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 Testing + +@testable import CoreAILanguageModels + +/// Covers the pure-value iterator helpers both sequential engines share: the generation-length +/// clamp and next-token selection. +@Suite("SequentialIterator") +struct SequentialIteratorTests { + @Test("clampMaxTokens: forced replays its own count, else clamps the request to remaining context") + func clampMaxTokens() { + // Forced continuation ignores the request and the context budget. + #expect( + SequentialIterator.clampMaxTokens( + requested: 999, forcedCount: 5, inputCount: 100, maxContextLength: 128) == 5) + // No request ("unbounded") clamps to the context left after the prompt. + #expect( + SequentialIterator.clampMaxTokens( + requested: nil, forcedCount: nil, inputCount: 100, maxContextLength: 128) == 28) + // A request smaller than the remaining context is honored as-is. + #expect( + SequentialIterator.clampMaxTokens( + requested: 10, forcedCount: nil, inputCount: 100, maxContextLength: 128) == 10) + } + + @Test("nextToken: replays the forced token, else returns the sampler's greedy argmax") + func nextToken() { + let logits: [LogitsScalarType] = [0.1, 0.2, 0.9, 0.3] + // Forced continuation replays the token at `step`, ignoring the logits. + #expect( + SequentialIterator.nextToken( + fromLogits: logits, forced: [7, 42], step: 1, + sampling: SamplingConfiguration(temperature: 0), tokenHistory: []) == 42) + // No forced tokens: greedy sampling (temperature 0) picks the argmax. + #expect( + SequentialIterator.nextToken( + fromLogits: logits, forced: nil, step: 0, + sampling: SamplingConfiguration(temperature: 0), tokenHistory: []) == 2) + } +}