From 94314784cdff1b8f6f9c121e78cd054543bbd1cb Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Sun, 27 Sep 2026 13:50:59 -0700 Subject: [PATCH 1/6] Split inference includeLogits bool into orthogonal tokens + logits axes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace InferenceOptions.includeLogits: Bool with two orthogonal axes — tokens: TokenRequest {.none, .sample} and logits: LogitsRequest {.none, .lastPosition, .allPositions} — plus presets (.prefill / .extend / .eval / .guided / .sampleWithLogits). Default (.sample, .none) equals the old includeLogits: false, so decode is byte-identical. All call sites migrated to presets; engines derive returnsLogits = (logits != .none); the pipelined engine rejects logits != .none (was includeLogits). --- .../ConstrainedDecodingStrategy.swift | 2 +- .../ConstrainedGenerator.swift | 2 +- .../DecodingStrategies/DecodingStrategy.swift | 2 +- .../VanillaDecodingStrategy.swift | 2 +- .../CoreAIPipelinedEngine.swift | 2 +- .../CoreAISequentialEngine.swift | 2 +- .../CoreAISequentialVLMEngine.swift | 2 +- .../CoreAIStaticShapeEngine.swift | 2 +- .../InferenceEngines/InferenceEngine.swift | 60 ++++++++++++++++--- .../TextGeneration/TextGenerator.swift | 8 +-- .../VLM/CoreAIVisionLanguageModel.swift | 2 +- .../Tools/benchmark/BenchmarkMain.swift | 2 +- .../Tools/llm-runner/LLMRunnerMain.swift | 6 +- .../Tools/llm-server/CompletionHandler.swift | 3 +- .../LanguageModelsTests/TestUtilities.swift | 2 +- .../UnifiedGenerationAPITests.swift | 43 ++++++++----- 16 files changed, 97 insertions(+), 45 deletions(-) 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.logits != .none { 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..1cca904b 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.logits != .none) 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..2d8cffcb 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.logits != .none) 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 45d79e58..ea7a9c99 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift @@ -689,7 +689,7 @@ extension StaticShapeEngine.GenerationSequence { ) { self.engine = engine self.samplingConfiguration = samplingConfiguration - self.returnsLogits = inferenceOptions.includeLogits + self.returnsLogits = (inferenceOptions.logits != .none) 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..2b4b2785 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,28 +32,70 @@ 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]? 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 } } +/// Presets for the common (tokens, logits) pairs. +extension InferenceOptions { + /// Warm the KV cache; no token or logits generation. + 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 /// Configuration-specific errors with user-friendly messages @@ -96,7 +138,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.. Date: Wed, 7 Oct 2026 14:19:03 -0700 Subject: [PATCH 2/6] Extract returnsLogits computed property; document tokens x logits combinations Add an InferenceOptions.returnsLogits computed property (equivalent to logits != .none) and route all engines (sequential, static-shape, VLM, pipelined guard) and the MockEngine through it instead of open-coding the comparison. Document on the raw initializer that the (tokens, logits) space is intentionally permissive with no invalid combinations, that the presets cover the meaningful pairings, and that (tokens: .sample, logits: .allPositions) is the one representable pair without a dedicated preset. Pure readability refactor; behavior unchanged. --- .../InferenceEngines/CoreAIPipelinedEngine.swift | 2 +- .../CoreAISequentialEngine.swift | 2 +- .../CoreAISequentialVLMEngine.swift | 2 +- .../CoreAIStaticShapeEngine.swift | 2 +- .../InferenceEngines/InferenceEngine.swift | 16 ++++++++++++++++ .../LanguageModelsTests/TestUtilities.swift | 2 +- 6 files changed, 21 insertions(+), 5 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index d5ebc4e6..f3bf2798 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -117,7 +117,7 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions ) async throws -> GenerationSequence { - if inferenceOptions.logits != .none { + 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 1cca904b..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.logits != .none) + 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 2d8cffcb..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.logits != .none) + 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 9b14e310..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.logits != .none) + 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 2b4b2785..1c5f4af6 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -54,6 +54,16 @@ public struct InferenceOptions: Sendable { /// 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, tokens: TokenRequest = .sample, @@ -65,6 +75,12 @@ public struct InferenceOptions: Sendable { self.logits = logits self.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 } } /// Presets for the common (tokens, logits) pairs. diff --git a/swift/Tests/LanguageModelsTests/TestUtilities.swift b/swift/Tests/LanguageModelsTests/TestUtilities.swift index 6e60733a..b0933999 100644 --- a/swift/Tests/LanguageModelsTests/TestUtilities.swift +++ b/swift/Tests/LanguageModelsTests/TestUtilities.swift @@ -157,7 +157,7 @@ class MockEngine: InferenceEngine, @unchecked Sendable { generationToken: GenerationToken ) { self.engine = engine - self.returnsLogits = (inferenceOptions.logits != .none) + self.returnsLogits = inferenceOptions.returnsLogits self.forcedContinuation = inferenceOptions.forcedContinuation self.stopReasonStore = stopReasonStore self.generationToken = generationToken From 8cf894bab30d73155169bfb797a03e4e67648cd2 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Fri, 9 Oct 2026 11:37:49 -0700 Subject: [PATCH 3/6] Keep includeLogits as a deprecated compatibility shim MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reintroduce the old InferenceOptions.includeLogits boolean as a deprecated initializer parameter and computed property that bridge to the (tokens, logits) axes: true -> logits: .lastPosition, false -> .none, with tokens still sampled. This avoids a hard public-API break from splitting the bool into the orthogonal axes — external callers keep compiling, with a deprecation warning pointing to the new API. Addresses review feedback. --- .../InferenceEngines/InferenceEngine.swift | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 1c5f4af6..9bf5320b 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -76,11 +76,38 @@ public struct InferenceOptions: Sendable { 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. From 62a56b741dc3af0cc927863ea537cd0a726b2b4b Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Fri, 9 Oct 2026 12:20:48 -0700 Subject: [PATCH 4/6] Clarify .prefill doc as request intent, not a generation guarantee Reword the .prefill preset from "no token or logits generation" to describe what the request needs. The (tokens, logits) axes express intent, and engines do not yet gate sampling on tokens: .none, so the old wording over-claimed runtime behavior. Addresses review feedback. --- .../CoreAILanguageModels/InferenceEngines/InferenceEngine.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 9bf5320b..9781aa83 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -112,7 +112,7 @@ public struct InferenceOptions: Sendable { /// Presets for the common (tokens, logits) pairs. extension InferenceOptions { - /// Warm the KV cache; no token or logits generation. + /// 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) } From 5cd27704e21b31ab1d60da8c23a487251c5724dc Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Mon, 26 Oct 2026 21:12:28 -0700 Subject: [PATCH 5/6] Tighten InferenceOptions docs and test the includeLogits shim Doc comments state what each option and preset does, and note that engines currently sample under both TokenRequest cases and return one row per step for both LogitsRequest positions. Tests cover returnsLogits and the deprecated includeLogits initializer and accessor. --- .../InferenceEngines/InferenceEngine.swift | 48 +++++++++---------- .../UnifiedGenerationAPITests.swift | 36 ++++++++++++++ 2 files changed, 58 insertions(+), 26 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 2c0c9103..c20d052e 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -32,13 +32,19 @@ public struct InferenceOutput: Sendable { // MARK: - Inference Options -/// Whether to sample a token this step. +/// Whether the engine samples a token each step. +/// +/// Engines currently sample under both cases. `.none` marks requests whose caller does not +/// read `InferenceOutput.tokenId`. public enum TokenRequest: Sendable, Equatable { case none case sample } -/// Which positions' logits to return. +/// Which positions' logits the engine returns. +/// +/// Engines currently return one `[vocabSize]` row per output step for both `.lastPosition` and +/// `.allPositions`. The pipelined engine rejects any value other than `.none`. public enum LogitsRequest: Sendable, Equatable { case none case lastPosition @@ -49,21 +55,15 @@ public enum LogitsRequest: Sendable, Equatable { public struct InferenceOptions: Sendable { /// Max tokens to generate. Nil = until EOS or context limit. public var maxTokens: Int? + /// Whether to sample a token each step. Default `.sample`. public var tokens: TokenRequest + /// Which positions' logits to return. Default `.none`. 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. + /// Creates options from explicit `tokens` and `logits` values. Every pair is valid; prefer the + /// presets below. public init( maxTokens: Int? = nil, tokens: TokenRequest = .sample, @@ -76,9 +76,8 @@ public struct InferenceOptions: Sendable { 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. + /// Creates options from the legacy `includeLogits` flag: `true` maps to `logits: .lastPosition`, + /// `false` to `logits: .none`. `tokens` is `.sample`. @available( *, deprecated, message: "Use the tokens:/logits: initializer; includeLogits: true == logits: .lastPosition." @@ -95,14 +94,11 @@ public struct InferenceOptions: Sendable { 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. + /// Whether `logits` is anything other than `.none`. 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. + /// Legacy flag. Reads `returnsLogits`; writing `true` sets `logits` to `.lastPosition`, `false` + /// to `.none`. @available(*, deprecated, message: "Use `logits` / `returnsLogits` instead.") public var includeLogits: Bool { get { returnsLogits } @@ -112,11 +108,11 @@ public struct InferenceOptions: Sendable { /// Presets for the common (tokens, logits) pairs. extension InferenceOptions { - /// Warm the KV cache; does not need token or logits + /// No token, no logits. For prefill and warmup chunks. public static func prefill(maxTokens: Int? = nil) -> InferenceOptions { InferenceOptions(maxTokens: maxTokens, tokens: .none, logits: .none) } - /// Decode one token. + /// Sampled tokens, no logits. The default for decode. public static func extend(maxTokens: Int? = nil) -> InferenceOptions { InferenceOptions(maxTokens: maxTokens, tokens: .sample, logits: .none) } @@ -124,16 +120,16 @@ extension InferenceOptions { public static func `default`(maxTokens: Int? = nil) -> InferenceOptions { .extend(maxTokens: maxTokens) } - /// Logits at every position, for eval / PPL. + /// No token, logits at every position. For perplexity and scoring with `forcedContinuation`. 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. + /// No token, last-position logits. For constrained decoding, where the caller masks and samples. public static func guided(maxTokens: Int? = nil) -> InferenceOptions { InferenceOptions(maxTokens: maxTokens, tokens: .none, logits: .lastPosition) } - /// Decode, and also return the last-row logits. + /// Sampled tokens plus last-position logits. For inspecting logits during decode. public static func sampleWithLogits(maxTokens: Int? = nil) -> InferenceOptions { InferenceOptions(maxTokens: maxTokens, tokens: .sample, logits: .lastPosition) } diff --git a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift index 223f5848..2e2d8c52 100644 --- a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift +++ b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift @@ -64,6 +64,42 @@ struct InferenceOptionsTests { #expect(InferenceOptions().tokens == .sample) #expect(InferenceOptions().logits == .none) } + + @Test("returnsLogits is false only for logits .none") + func returnsLogits() { + #expect(InferenceOptions(logits: .none).returnsLogits == false) + #expect(InferenceOptions(logits: .lastPosition).returnsLogits == true) + #expect(InferenceOptions(logits: .allPositions).returnsLogits == true) + } + + @available(*, deprecated) + @Test("Legacy includeLogits initializer maps to sample + lastPosition or none") + func legacyInitializer() { + let forced: [Int32] = [4, 5] + let on = InferenceOptions(maxTokens: 3, includeLogits: true, forcedContinuation: forced) + #expect(on.maxTokens == 3) + #expect(on.tokens == .sample) + #expect(on.logits == .lastPosition) + #expect(on.forcedContinuation == forced) + + let off = InferenceOptions(includeLogits: false) + #expect(off.tokens == .sample) + #expect(off.logits == .none) + } + + @available(*, deprecated) + @Test("Legacy includeLogits accessor reads returnsLogits and writes the logits axis") + func legacyAccessor() { + var opts = InferenceOptions.eval() + #expect(opts.includeLogits == true) + + opts.includeLogits = false + #expect(opts.logits == .none) + #expect(opts.tokens == .none) + + opts.includeLogits = true + #expect(opts.logits == .lastPosition) + } } // MARK: - generate() Default Extension Tests From d3e8b83272c486d11250f4a4155323b19f6cbb4c Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Tue, 27 Oct 2026 18:14:52 -0700 Subject: [PATCH 6/6] Describe what the tokens: .none presets request --- .../InferenceEngines/InferenceEngine.swift | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index c20d052e..19ec4bbd 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -106,9 +106,11 @@ public struct InferenceOptions: Sendable { } } -/// Presets for the common (tokens, logits) pairs. +/// Presets for the common (tokens, logits) pairs. Engines currently sample a token under every +/// preset; see `TokenRequest`. extension InferenceOptions { - /// No token, no logits. For prefill and warmup chunks. + /// Requests no token and no logits, for prefill and warmup chunks. Pass `maxTokens`; nil runs + /// until EOS. public static func prefill(maxTokens: Int? = nil) -> InferenceOptions { InferenceOptions(maxTokens: maxTokens, tokens: .none, logits: .none) } @@ -120,12 +122,14 @@ extension InferenceOptions { public static func `default`(maxTokens: Int? = nil) -> InferenceOptions { .extend(maxTokens: maxTokens) } - /// No token, logits at every position. For perplexity and scoring with `forcedContinuation`. + /// Requests logits at every position and no token. For perplexity and scoring with + /// `forcedContinuation`. public static func eval(maxTokens: Int? = nil, forcedContinuation: [Int32]? = nil) -> InferenceOptions { InferenceOptions( maxTokens: maxTokens, tokens: .none, logits: .allPositions, forcedContinuation: forcedContinuation) } - /// No token, last-position logits. For constrained decoding, where the caller masks and samples. + /// Requests last-position logits and no token. For constrained decoding, where the caller masks + /// and samples. public static func guided(maxTokens: Int? = nil) -> InferenceOptions { InferenceOptions(maxTokens: maxTokens, tokens: .none, logits: .lastPosition) }