From 88ecb3d58139dacf01b963f55a4dfc7487dfc252 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Fri, 17 Jul 2026 08:48:14 -0700 Subject: [PATCH] Stop pipelined generation when consumer drops the stream When the consumer breaks out of the token stream (e.g. at EOS), the forwarding loop now detects the terminated yield result and finishes the inner token continuation, which signals the GPU producer to stop. Additionally, outputContinuation.onTermination finishes the inner stream so cancellation propagates even if the forwarding loop hasn't started iterating yet. Also cancel any prior active generation at the start of generate() on all engine types (sequential, static-shape, VLM). This ensures the cancel-and-replace contract is upheld regardless of which engine is in use, preventing state overlap between turns. --- .../CoreAIPipelinedEngine.swift | 17 +++++++++++++++-- .../CoreAISequentialEngine.swift | 7 ++++++- .../CoreAISequentialVLMEngine.swift | 8 ++++++++ .../CoreAIStaticShapeEngine.swift | 6 ++++++ 4 files changed, 35 insertions(+), 3 deletions(-) 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)