Skip to content

Commit a270998

Browse files
stikvesstikves
authored
Add InferenceStream with StopReason for structured generation termination (apple#16)
Introduces InferenceStream as a concrete return type from InferenceEngine.generate(), replacing the opaque `some AsyncSequence<InferenceOutput, Error>`. The nested StopReason enum (.maxTokens, .eos, .stopSequence, .cancelled, .error) gives callers structured visibility into why generation ended — previously they had to guess from local state. - InferenceStream: final class, AsyncSequence, Sendable - Engines set .maxTokens/.cancelled/.error; decoders set .eos - FM adaptor (respondVanilla) sets .eos on EOS token detection - Removes associatedtype OutputSequence from InferenceEngine protocol - All existing callers compile unchanged (for-await still works)
1 parent 9592167 commit a270998

8 files changed

Lines changed: 310 additions & 146 deletions

File tree

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
8787
with input: [TokenId],
8888
samplingConfiguration: SamplingConfiguration,
8989
inferenceOptions: InferenceOptions
90-
) throws -> some AsyncSequence<InferenceOutput, Error> {
90+
) throws -> InferenceStream {
9191
if inferenceOptions.includeLogits {
9292
throw InferenceRuntimeError.invalidArgument(
9393
"CoreAI pipelined engine does not support logits (GPU-side sampling). "
@@ -101,18 +101,14 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
101101
)
102102
}
103103
let maxTokens = inferenceOptions.maxTokens
104-
let (stream, outputContinuation) = AsyncThrowingStream<InferenceOutput, any Error>.makeStream()
104+
let (inferenceStream, outputContinuation) = InferenceStream.makeStream()
105105
Task {
106106
self.acquireEngine()
107107
defer { self.releaseEngine() }
108108
do {
109-
// Bridge: runCompletion yields TokenId via a TokenId continuation.
110-
// We wrap each token into InferenceOutput in the yield callback.
111109
let (tokenStream, tokenContinuation) =
112110
AsyncThrowingStream<InferenceEngine.TokenId, any Error>.makeStream()
113111

114-
// Forward tokens from tokenStream → outputContinuation as InferenceOutput.
115-
// This must run concurrently with runCompletion.
116112
async let forwarding: Void = {
117113
do {
118114
for try await token in tokenStream {
@@ -131,12 +127,17 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
131127
)
132128
tokenContinuation.finish()
133129
await forwarding
130+
inferenceStream.setStopReason(.maxTokens)
131+
outputContinuation.finish()
132+
} catch is CancellationError {
133+
inferenceStream.setStopReason(.cancelled)
134134
outputContinuation.finish()
135135
} catch {
136+
inferenceStream.setStopReason(.error)
136137
outputContinuation.finish(throwing: error)
137138
}
138139
}
139-
return stream
140+
return inferenceStream
140141
}
141142

142143
/// Wait for any in-flight generate() Task to return the engine.

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift‎

Lines changed: 62 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -338,72 +338,77 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
338338
with input: [TokenId],
339339
samplingConfiguration: SamplingConfiguration,
340340
inferenceOptions: InferenceOptions
341-
) throws -> some AsyncSequence<InferenceOutput, Error> {
342-
AsyncThrowingStream { continuation in
343-
Task {
344-
self.generating.withLock { $0 = true }
345-
defer { self.generating.withLock { $0 = false } }
346-
do {
347-
let maxTokens: Int
348-
if let forced = inferenceOptions.forcedContinuation {
349-
maxTokens = forced.count
350-
} else {
351-
maxTokens = min(
352-
inferenceOptions.maxTokens ?? Int.max,
353-
max(0, self.config.maxContextLength - input.count)
354-
)
355-
}
356-
let returnsLogits = inferenceOptions.includeLogits
357-
var inputTokens = input
358-
359-
for i in 0..<maxTokens {
360-
try Task.checkCancellation()
341+
) throws -> InferenceStream {
342+
let (stream, continuation) = InferenceStream.makeStream()
343+
Task {
344+
self.generating.withLock { $0 = true }
345+
defer { self.generating.withLock { $0 = false } }
346+
do {
347+
let maxTokens: Int
348+
if let forced = inferenceOptions.forcedContinuation {
349+
maxTokens = forced.count
350+
} else {
351+
maxTokens = min(
352+
inferenceOptions.maxTokens ?? Int.max,
353+
max(0, self.config.maxContextLength - input.count)
354+
)
355+
}
356+
let returnsLogits = inferenceOptions.includeLogits
357+
var inputTokens = input
361358

362-
guard self.processedTokenCount < inputTokens.count else {
363-
throw InferenceRuntimeError.invalidState("No new tokens to process")
364-
}
359+
for i in 0..<maxTokens {
360+
try Task.checkCancellation()
365361

366-
let newTokens = inputTokens[self.processedTokenCount...]
367-
let strategy = self.selectPrefillStrategy(newTokenCount: newTokens.count)
368-
369-
let logitBuffer: [LogitsScalarType]
370-
switch strategy {
371-
case .chunked(let chunkSize):
372-
logitBuffer = try await self.processChunkedPrompt(
373-
tokens: newTokens, chunkSize: chunkSize)
374-
case .wholeBatch:
375-
let allLogits = try await self.processTokenBatch(newTokens)
376-
logitBuffer = lastTokenLogits(
377-
from: allLogits, vocabSize: self.config.vocabSize)
378-
case .oneAtATime:
379-
var lastLogits: [LogitsScalarType] = []
380-
for j in newTokens.indices {
381-
lastLogits = try await self.processTokenBatch(newTokens[j...j])
382-
}
383-
logitBuffer = lastLogits
384-
}
362+
guard self.processedTokenCount < inputTokens.count else {
363+
throw InferenceRuntimeError.invalidState("No new tokens to process")
364+
}
385365

386-
let nextToken: Int32
387-
if let forced = inferenceOptions.forcedContinuation {
388-
nextToken = forced[i]
389-
} else {
390-
var mutableLogits = logitBuffer
391-
nextToken = samplingConfiguration.fallbackSampler(from: &mutableLogits)
366+
let newTokens = inputTokens[self.processedTokenCount...]
367+
let strategy = self.selectPrefillStrategy(newTokenCount: newTokens.count)
368+
369+
let logitBuffer: [LogitsScalarType]
370+
switch strategy {
371+
case .chunked(let chunkSize):
372+
logitBuffer = try await self.processChunkedPrompt(
373+
tokens: newTokens, chunkSize: chunkSize)
374+
case .wholeBatch:
375+
let allLogits = try await self.processTokenBatch(newTokens)
376+
logitBuffer = lastTokenLogits(
377+
from: allLogits, vocabSize: self.config.vocabSize)
378+
case .oneAtATime:
379+
var lastLogits: [LogitsScalarType] = []
380+
for j in newTokens.indices {
381+
lastLogits = try await self.processTokenBatch(newTokens[j...j])
392382
}
383+
logitBuffer = lastLogits
384+
}
393385

394-
continuation.yield(
395-
InferenceOutput(
396-
tokenId: nextToken,
397-
logits: returnsLogits ? logitBuffer : nil
398-
))
399-
inputTokens.append(nextToken)
386+
let nextToken: Int32
387+
if let forced = inferenceOptions.forcedContinuation {
388+
nextToken = forced[i]
389+
} else {
390+
var mutableLogits = logitBuffer
391+
nextToken = samplingConfiguration.fallbackSampler(from: &mutableLogits)
400392
}
401-
continuation.finish()
402-
} catch {
403-
continuation.finish(throwing: error)
393+
394+
continuation.yield(
395+
InferenceOutput(
396+
tokenId: nextToken,
397+
logits: returnsLogits ? logitBuffer : nil
398+
))
399+
inputTokens.append(nextToken)
404400
}
401+
stream.setStopReason(.maxTokens)
402+
continuation.finish()
403+
} catch is CancellationError {
404+
stream.setStopReason(.cancelled)
405+
continuation.finish()
406+
} catch {
407+
stream.setStopReason(.error)
408+
continuation.finish(throwing: error)
405409
}
406410
}
411+
return stream
407412
}
408413

409414
// MARK: - Lifecycle

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift‎

Lines changed: 42 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -316,48 +316,51 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
316316
with input: [TokenId],
317317
samplingConfiguration: SamplingConfiguration,
318318
inferenceOptions: InferenceOptions
319-
) throws -> some AsyncSequence<InferenceOutput, Error> {
320-
AsyncThrowingStream { continuation in
321-
Task {
322-
do {
323-
let forced = inferenceOptions.forcedContinuation
324-
let maxTokens: Int
325-
if let forced {
326-
maxTokens = forced.count
327-
} else {
328-
maxTokens = min(
329-
inferenceOptions.maxTokens ?? Int.max,
330-
max(0, self.config.maxContextLength - input.count)
331-
)
332-
}
333-
let returnsLogits = inferenceOptions.includeLogits
334-
var inputTokens = input
335-
336-
for i in 0..<maxTokens {
337-
try Task.checkCancellation()
338-
// When forced, we still need the forward pass (for logits + KV cache update)
339-
// but skip the sampler — the next token is predetermined.
340-
let (logits, sampledToken) = try await self.inference(
341-
inputTokens: inputTokens,
342-
samplingConfig: samplingConfiguration,
343-
returnsLogits: returnsLogits || forced != nil
344-
)
345-
346-
let nextToken = forced?[i] ?? sampledToken
347-
348-
continuation.yield(
349-
InferenceOutput(
350-
tokenId: nextToken,
351-
logits: returnsLogits ? logits : nil
352-
))
353-
inputTokens.append(nextToken)
354-
}
355-
continuation.finish()
356-
} catch {
357-
continuation.finish(throwing: error)
319+
) throws -> InferenceStream {
320+
let (stream, continuation) = InferenceStream.makeStream()
321+
Task {
322+
do {
323+
let forced = inferenceOptions.forcedContinuation
324+
let maxTokens: Int
325+
if let forced {
326+
maxTokens = forced.count
327+
} else {
328+
maxTokens = min(
329+
inferenceOptions.maxTokens ?? Int.max,
330+
max(0, self.config.maxContextLength - input.count)
331+
)
358332
}
333+
let returnsLogits = inferenceOptions.includeLogits
334+
var inputTokens = input
335+
336+
for i in 0..<maxTokens {
337+
try Task.checkCancellation()
338+
let (logits, sampledToken) = try await self.inference(
339+
inputTokens: inputTokens,
340+
samplingConfig: samplingConfiguration,
341+
returnsLogits: returnsLogits || forced != nil
342+
)
343+
344+
let nextToken = forced?[i] ?? sampledToken
345+
346+
continuation.yield(
347+
InferenceOutput(
348+
tokenId: nextToken,
349+
logits: returnsLogits ? logits : nil
350+
))
351+
inputTokens.append(nextToken)
352+
}
353+
stream.setStopReason(.maxTokens)
354+
continuation.finish()
355+
} catch is CancellationError {
356+
stream.setStopReason(.cancelled)
357+
continuation.finish()
358+
} catch {
359+
stream.setStopReason(.error)
360+
continuation.finish(throwing: error)
359361
}
360362
}
363+
return stream
361364
}
362365

363366
// MARK: - Inference

‎swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,6 @@ public enum ConfigurationError: Error, LocalizedError {
8585
///
8686
/// KV cache is preserved between `generate()` calls. Call `reset()` to clear.
8787
public protocol InferenceEngine: Sendable {
88-
associatedtype OutputSequence: AsyncSequence<InferenceOutput, Error>
8988
typealias TokenId = Int32
9089

9190
// MARK: - Primary API
@@ -96,12 +95,12 @@ public protocol InferenceEngine: Sendable {
9695
/// - input: Token IDs (prompt, context, or continuation).
9796
/// - sampling: Sampling configuration (temperature, topK, etc.).
9897
/// - generation: Inference options (maxTokens, includeLogits).
99-
/// - Returns: Async sequence of `InferenceOutput`.
98+
/// - Returns: `InferenceStream` — iterate for tokens, read `stopReason` after.
10099
func generate(
101100
with input: [TokenId],
102101
samplingConfiguration: SamplingConfiguration,
103102
inferenceOptions: InferenceOptions
104-
) throws -> OutputSequence
103+
) throws -> InferenceStream
105104

106105
// MARK: - Lifecycle
107106

‎swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceStream.swift‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
// Copyright 2026 Apple Inc.
2+
//
3+
// Use of this source code is governed by a BSD-3-clause license that can
4+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
5+
6+
import Synchronization
7+
8+
/// Concrete async sequence returned by `InferenceEngine.generate()`.
9+
///
10+
/// Wraps an `AsyncThrowingStream<InferenceOutput, any Error>` and tracks
11+
/// why generation ended. Reference type — the producer sets `stopReason`
12+
/// and the consumer reads it after iteration.
13+
public final class InferenceStream: AsyncSequence, Sendable {
14+
public typealias Element = InferenceOutput
15+
16+
// MARK: - StopReason
17+
18+
/// Why token generation terminated.
19+
public enum StopReason: Sendable, Equatable {
20+
/// The maximum token limit was reached.
21+
case maxTokens
22+
/// An end-of-sequence token was generated.
23+
case eos
24+
/// A stop sequence was matched in the output.
25+
case stopSequence(String)
26+
/// Generation was cancelled (Task cancellation or explicit cancel).
27+
case cancelled
28+
/// An unrecoverable error occurred during generation.
29+
case error
30+
}
31+
32+
// MARK: - Init
33+
34+
private let base: AsyncThrowingStream<InferenceOutput, any Error>
35+
private let _stopReason: Mutex<StopReason?>
36+
37+
init(base: AsyncThrowingStream<InferenceOutput, any Error>) {
38+
self.base = base
39+
self._stopReason = Mutex(nil)
40+
}
41+
42+
// MARK: - Public API
43+
44+
/// Why generation stopped. Nil while the stream is still active.
45+
/// Guaranteed non-nil after the `for try await` loop exits.
46+
public var stopReason: StopReason? {
47+
_stopReason.withLock { $0 }
48+
}
49+
50+
// MARK: - Package-internal
51+
52+
/// Engines and decoders call this when they know why generation ended.
53+
func setStopReason(_ reason: StopReason) {
54+
_stopReason.withLock { $0 = reason }
55+
}
56+
57+
// MARK: - Factory
58+
59+
/// Create a stream + continuation pair for engines to drive.
60+
static func makeStream() -> (
61+
stream: InferenceStream,
62+
continuation: AsyncThrowingStream<InferenceOutput, any Error>.Continuation
63+
) {
64+
let (base, continuation) = AsyncThrowingStream<InferenceOutput, any Error>.makeStream()
65+
return (InferenceStream(base: base), continuation)
66+
}
67+
68+
// MARK: - AsyncSequence
69+
70+
public struct AsyncIterator: AsyncIteratorProtocol {
71+
var base: AsyncThrowingStream<InferenceOutput, any Error>.AsyncIterator
72+
let stream: InferenceStream
73+
74+
public mutating func next() async throws -> InferenceOutput? {
75+
do {
76+
guard let element = try await base.next() else {
77+
return nil
78+
}
79+
return element
80+
} catch is CancellationError {
81+
stream.setStopReason(.cancelled)
82+
throw CancellationError()
83+
} catch {
84+
stream.setStopReason(.error)
85+
throw error
86+
}
87+
}
88+
}
89+
90+
public func makeAsyncIterator() -> AsyncIterator {
91+
AsyncIterator(base: base.makeAsyncIterator(), stream: self)
92+
}
93+
}

0 commit comments

Comments
 (0)