diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift index f5873808..838f3186 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift @@ -260,7 +260,7 @@ extension ConstrainedDecodingStrategy.ConstrainedDecodedSequence { private let inferenceEngine: any InferenceEngine private let samplingConfiguration: SamplingConfiguration private let stopSequences: StopSequences - private let constrainedOptions = InferenceOptions(maxTokens: 1, includeLogits: true) + private let constrainedOptions = InferenceOptions.guided(maxTokens: 1) // Generation state, seeded eagerly from the prepared setup. private var session: ConstrainedGenerationSession? diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift index 9e621323..0bd23a9b 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift @@ -188,7 +188,7 @@ public struct ConstrainedGenerator: DecodingStrategy { for _ in 0.. GenerationSequence { - if inferenceOptions.includeLogits { + if inferenceOptions.returnsLogits { throw InferenceRuntimeError.invalidArgument( "CoreAI pipelined engine does not support logits (GPU-side sampling). " + "Use a sequential engine for constrained generation or evaluation." diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 37dca8c7..b0d9a196 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -603,7 +603,7 @@ extension CoreAISequentialEngine.GenerationSequence { self.engine = engine self.sessionState = sessionState self.samplingConfiguration = samplingConfiguration.normalized() - self.returnsLogits = inferenceOptions.includeLogits + self.returnsLogits = inferenceOptions.returnsLogits self.forcedContinuation = inferenceOptions.forcedContinuation self.stopReasonStore = stopReasonStore self.generationToken = generationToken diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift index 46353fda..19caaaad 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift @@ -983,7 +983,7 @@ extension CoreAISequentialVLMEngine.GenerationSequence { ) { self.engine = engine self.samplingConfiguration = samplingConfiguration.normalized() - self.returnsLogits = inferenceOptions.includeLogits + self.returnsLogits = inferenceOptions.returnsLogits self.forcedContinuation = inferenceOptions.forcedContinuation self.stopReasonStore = stopReasonStore self.generationToken = generationToken diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift index a6ae1365..3890860f 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift @@ -793,7 +793,7 @@ extension StaticShapeEngine.GenerationSequence { ) { self.engine = engine self.samplingConfiguration = samplingConfiguration - self.returnsLogits = inferenceOptions.includeLogits + self.returnsLogits = inferenceOptions.returnsLogits self.forcedContinuation = inferenceOptions.forcedContinuation self.stopReasonStore = stopReasonStore self.generationToken = generationToken diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 9fc24ddd..9781aa83 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -21,7 +21,7 @@ public typealias LogitsScalarType = Float public struct InferenceOutput: Sendable { public let tokenId: Int32 - /// Populated when `InferenceOptions.includeLogits` is true. Shape: [vocabSize]. + /// Populated when `InferenceOptions.logits` is not `.none`. Shape: [vocabSize]. public let logits: [LogitsScalarType]? public init(tokenId: Int32, logits: [LogitsScalarType]? = nil) { @@ -32,26 +32,111 @@ public struct InferenceOutput: Sendable { // MARK: - Inference Options +/// Whether to sample a token this step. +public enum TokenRequest: Sendable, Equatable { + case none + case sample +} + +/// Which positions' logits to return. +public enum LogitsRequest: Sendable, Equatable { + case none + case lastPosition + case allPositions +} + /// Controls what the engine produces and how much. -/// Struct-based for additive extensibility (future: embeddings, attention maps). public struct InferenceOptions: Sendable { /// Max tokens to generate. Nil = until EOS or context limit. public var maxTokens: Int? - /// Include raw logits in each `InferenceOutput`. May incur GPU→CPU copy cost. - public var includeLogits: Bool - /// When set, engines use these token IDs instead of sampling. - /// Used by MMLU-style evaluation to compute P(continuation|context). + public var tokens: TokenRequest + public var logits: LogitsRequest + /// Force these token IDs instead of sampling (MMLU-style scoring). public var forcedContinuation: [Int32]? + /// Raw initializer over the full `(tokens, logits)` space. + /// + /// The two axes are intentionally orthogonal and this initializer is + /// deliberately permissive: every `(tokens, logits)` pair is representable + /// and there are no invalid combinations to reject. The `prefill`, `extend`, + /// `eval`, `guided`, and `sampleWithLogits` presets below cover all the + /// meaningful pairings; prefer them to constructing options by hand. The + /// only representable pair without a dedicated preset is + /// `(tokens: .sample, logits: .allPositions)` — sampling a token while also + /// returning logits at every position — which is odd but harmless. public init( maxTokens: Int? = nil, - includeLogits: Bool = false, + tokens: TokenRequest = .sample, + logits: LogitsRequest = .none, forcedContinuation: [Int32]? = nil ) { self.maxTokens = maxTokens - self.includeLogits = includeLogits + self.tokens = tokens + self.logits = logits self.forcedContinuation = forcedContinuation } + + /// Legacy initializer bridging the old `includeLogits` boolean to the + /// `(tokens, logits)` axes: `true` maps to `logits: .lastPosition`, `false` to + /// `logits: .none`. Tokens are sampled, as they always were under this API. + @available( + *, deprecated, + message: "Use the tokens:/logits: initializer; includeLogits: true == logits: .lastPosition." + ) + public init( + maxTokens: Int? = nil, + includeLogits: Bool, + forcedContinuation: [Int32]? = nil + ) { + self.init( + maxTokens: maxTokens, + tokens: .sample, + logits: includeLogits ? .lastPosition : .none, + forcedContinuation: forcedContinuation) + } + + /// Whether this request asks the engine to return logits at any position. + /// + /// Equivalent to `logits != .none`; use this instead of open-coding the + /// comparison so the intent reads clearly at the call site. + public var returnsLogits: Bool { logits != .none } + + /// Legacy accessor for the old `includeLogits` boolean. Reads as `returnsLogits`; + /// setting `true` requests last-position logits, `false` requests none. + @available(*, deprecated, message: "Use `logits` / `returnsLogits` instead.") + public var includeLogits: Bool { + get { returnsLogits } + set { logits = newValue ? .lastPosition : .none } + } +} + +/// Presets for the common (tokens, logits) pairs. +extension InferenceOptions { + /// Warm the KV cache; does not need token or logits + public static func prefill(maxTokens: Int? = nil) -> InferenceOptions { + InferenceOptions(maxTokens: maxTokens, tokens: .none, logits: .none) + } + /// Decode one token. + public static func extend(maxTokens: Int? = nil) -> InferenceOptions { + InferenceOptions(maxTokens: maxTokens, tokens: .sample, logits: .none) + } + /// Same as `extend`. + public static func `default`(maxTokens: Int? = nil) -> InferenceOptions { + .extend(maxTokens: maxTokens) + } + /// Logits at every position, for eval / PPL. + public static func eval(maxTokens: Int? = nil, forcedContinuation: [Int32]? = nil) -> InferenceOptions { + InferenceOptions( + maxTokens: maxTokens, tokens: .none, logits: .allPositions, forcedContinuation: forcedContinuation) + } + /// Last-row logits for constrained decoding. + public static func guided(maxTokens: Int? = nil) -> InferenceOptions { + InferenceOptions(maxTokens: maxTokens, tokens: .none, logits: .lastPosition) + } + /// Decode, and also return the last-row logits. + public static func sampleWithLogits(maxTokens: Int? = nil) -> InferenceOptions { + InferenceOptions(maxTokens: maxTokens, tokens: .sample, logits: .lastPosition) + } } // MARK: - Configuration Data Structures @@ -96,7 +181,7 @@ public protocol InferenceEngine: Sendable { /// - Parameters: /// - input: Token IDs (prompt, context, or continuation). /// - sampling: Sampling configuration (temperature, topK, etc.). - /// - generation: Inference options (maxTokens, includeLogits). + /// - generation: Inference options (maxTokens, tokens, logits). /// - Returns: An `InferenceOutputSequence` — iterate for tokens, read /// `stopReason` after the loop to learn why generation ended. func generate( diff --git a/swift/Sources/CoreAILanguageModels/TextGeneration/TextGenerator.swift b/swift/Sources/CoreAILanguageModels/TextGeneration/TextGenerator.swift index 17807fd1..ae5736a7 100644 --- a/swift/Sources/CoreAILanguageModels/TextGeneration/TextGenerator.swift +++ b/swift/Sources/CoreAILanguageModels/TextGeneration/TextGenerator.swift @@ -83,7 +83,7 @@ public class TextGenerator { tokenizer: tokenizer, inferenceEngine: inferenceEngine, samplingConfiguration: samplingConfiguration, - options: InferenceOptions(maxTokens: maxTokens, includeLogits: true), + options: .sampleWithLogits(maxTokens: maxTokens), stopSequences: effectiveStopSequences ) @@ -134,9 +134,8 @@ public class TextGenerator { let continuationTokens = Array( encoding.tokens[encoding.continuationStartIndex..