Consolidate VLM sequential engine: shared iterator helpers, tidy (B3+B4)

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

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> = [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

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>
+    ) -> Int32 {
+        if let forced {
+            return forced[step]
+        }
+        var mutableLogits = logits
+        return sampling.fallbackSampler(from: &mutableLogits, tokenHistory: tokenHistory)
+    }
+}

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