diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 1ee9839..2c28784 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -45,6 +45,12 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable { typealias ConfigType = ModelConfig nonisolated(unsafe) private var engine: EngineImpl + + /// Token history for implicit prefix caching. Marked nonisolated(unsafe) because + /// mutations are serialized by the generation lifecycle: generate() awaits any prior + /// Task before starting, and the forwarding `async let` only appends tokens while + /// runCompletion holds the engine lock. No concurrent writes are possible when the + /// cancel-and-await contract is upheld. nonisolated(unsafe) private var history = TokenHistory() nonisolated(unsafe) private(set) var lastPrefixHitCount: Int = 0 private let engineInUse = Atomic(false) @@ -147,6 +153,10 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable { let (tokenStream, tokenContinuation) = AsyncThrowingStream.makeStream() + outputContinuation.onTermination = { @Sendable _ in + tokenContinuation.finish() + } + // Implicit prefix caching: resolve input against history var (commonPrefix, resolvedNewTokens) = self.history.resolve(input: input) self.lastPrefixHitCount = commonPrefix @@ -181,9 +191,12 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable { async let forwarding: Void = { do { for try await token in tokenStream { - // Track generated tokens in history self.history.append(token) - outputContinuation.yield(InferenceOutput(tokenId: token)) + let result = outputContinuation.yield(InferenceOutput(tokenId: token)) + if case .terminated = result { + tokenContinuation.finish() + break + } } } catch { outputContinuation.finish(throwing: error) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 620b9a8..1a768dc 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -349,7 +349,12 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions ) async throws -> GenerationSequence { - // Implicit prefix caching: resolve before creating Iterator. + // Cancel any prior generation so its Iterator stops on next poll. + _activeToken.withLock { + $0?.cancel() + $0 = nil + } + // Implicit prefix caching: resolve input against history. if history.count > 0 { let (commonPrefix, _) = history.resolve(input: input) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift index 2241017..f49ad8c 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift @@ -775,6 +775,10 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions ) async throws -> GenerationSequence { + _activeToken.withLock { + $0?.cancel() + $0 = nil + } let token = GenerationToken() _activeToken.withLock { $0 = token } return GenerationSequence( @@ -800,6 +804,10 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions ) async throws -> GenerationSequence { + _activeToken.withLock { + $0?.cancel() + $0 = nil + } let token = GenerationToken() _activeToken.withLock { $0 = token } return GenerationSequence( diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift index 6eebff9..5fee411 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift @@ -332,6 +332,12 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable { samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions ) async throws -> GenerationSequence { + // Cancel any prior generation so its Iterator stops on next poll. + _activeToken.withLock { + $0?.cancel() + $0 = nil + } + // Implicit prefix caching: resolve input against history. if history.count > 0 { let (commonPrefix, _) = history.resolve(input: input)