stikves · Oct 28, 2026 · Sep 27, 2026 · Oct 2, 2026 · Oct 7, 2026 · Oct 7, 2026
@@ -272,7 +272,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?

@@ -188,7 +188,7 @@ public struct ConstrainedGenerator: DecodingStrategy {
         for _ in 0..<maxTokens {
             if session.isTerminated { break }
 
-            let options = InferenceOptions(maxTokens: 1, includeLogits: true)
+            let options = InferenceOptions.guided(maxTokens: 1)
 
             var rawLogits: [LogitsScalarType]? = nil
             for try await output in try await inferenceEngine.generate(

@@ -171,7 +171,7 @@ public protocol DecodingStrategy: Sendable {
     ///   - tokenizer: Tokenizer for encoding/decoding
     ///   - inferenceEngine: Engine for model inference
     ///   - samplingConfiguration: Sampling parameters (temperature, topK, etc.)
-    ///   - options: Inference options (maxTokens, includeLogits)
+    ///   - options: Inference options (maxTokens, tokens, logits)
     ///   - stopSequences: Token sequences that halt generation
     /// - Returns: Stream of `GenerationResult` (text + optional logits)
     func decode(

@@ -24,7 +24,7 @@ public struct VanillaDecodingStrategy: DecodingStrategy {
     ///   - tokenizer: Tokenizer for encoding/decoding
     ///   - inferenceEngine: Engine for model inference
     ///   - samplingConfiguration: Sampling parameters (temperature, topK, etc.)
-    ///   - options: Inference options (maxTokens, includeLogits)
+    ///   - options: Inference options (maxTokens, tokens, logits)
     ///   - stopSequences: Token sequences that halt generation
     /// - Returns: Stream of `GenerationResult` (text + optional logits)
     public func decode(

@@ -117,7 +117,7 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable
         samplingConfiguration: SamplingConfiguration,
         inferenceOptions: InferenceOptions
     ) async throws -> 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."

@@ -992,7 +992,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

@@ -985,7 +985,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

@@ -925,7 +925,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

@@ -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 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 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
+    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).
+    /// 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]?
 
+    /// Creates options from explicit `tokens` and `logits` values. Every pair is valid; prefer the
+    /// presets below.
     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
     }
+
+    /// 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."
+    )
+    public init(
+        maxTokens: Int? = nil,
+        includeLogits: Bool,
+        forcedContinuation: [Int32]? = nil
+    ) {
+        self.init(
+            maxTokens: maxTokens,
+            tokens: .sample,
+            logits: includeLogits ? .lastPosition : .none,
+            forcedContinuation: forcedContinuation)
+    }
+
+    /// Whether `logits` is anything other than `.none`.
+    public var returnsLogits: Bool { logits != .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 }
+        set { logits = newValue ? .lastPosition : .none }
+    }
+}
+
+/// Presets for the common (tokens, logits) pairs. Engines currently sample a token under every
+/// preset; see `TokenRequest`.
+extension InferenceOptions {
+    /// 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)
+    }
+    /// Sampled tokens, no logits. The default for decode.
+    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)
+    }
+    /// 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)
+    }
+    /// 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)
+    }
+    /// 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)
+    }
 }
 
 // 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(

@@ -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..<encoding.tokens.count])
 
-        let options = InferenceOptions(
+        let options = InferenceOptions.eval(
             maxTokens: continuationTokens.count,
-            includeLogits: true,
             forcedContinuation: continuationTokens
         )
 
@@ -193,9 +192,8 @@ public class TextGenerator {
 
         try await inferenceEngine.reset()
 
-        let options = InferenceOptions(
+        let options = InferenceOptions.eval(
             maxTokens: continuationTokens.count,
-            includeLogits: true,
             forcedContinuation: continuationTokens
         )
 

@@ -161,7 +161,7 @@ public struct CoreAIVLMExecutor: LanguageModelExecutor {
             with: embeddedInput,
             tokens: promptTokens,
             samplingConfiguration: SamplingConfiguration(temperature: 1.0, topK: 1),
-            inferenceOptions: InferenceOptions(maxTokens: maxTokens, includeLogits: false)
+            inferenceOptions: .extend(maxTokens: maxTokens)
         )
 
         var generatedCount = 0

@@ -221,7 +221,7 @@ struct LLMBenchmark: AsyncParsableCommand {
         try? await Task.sleep(for: .milliseconds(50))
         try await engine.reset()
 
-        let options = InferenceOptions(maxTokens: generationTokens, includeLogits: false)
+        let options = InferenceOptions.extend(maxTokens: generationTokens)
         let start = SuspendingClock.now
         let stream = try await engine.generate(
             with: prompt, samplingConfiguration: sampling, inferenceOptions: options

@@ -1223,10 +1223,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
             with: embeddedInput,
             tokens: vlmTokens,
             samplingConfiguration: samplingConfiguration,
-            inferenceOptions: InferenceOptions(
-                maxTokens: maxTokens,
-                includeLogits: printLogits || saveLogits != nil
-            )
+            inferenceOptions: (printLogits || saveLogits != nil)
+                ? .sampleWithLogits(maxTokens: maxTokens) : .extend(maxTokens: maxTokens)
         )
 
         CLILogger.log("VLM generate started, maxTokens=\(maxTokens)", component: "VLM")

@@ -161,9 +161,8 @@ private func preparePrompt(_ prompt: Prompt, state: ServerState) throws -> Prepa
 private func scoreOnEngine(_ prompt: PreparedPrompt, topN: Int, state: ServerState) async throws
     -> LogProbabilities
 {
-    let options = InferenceOptions(
+    let options = InferenceOptions.eval(
         maxTokens: prompt.targets.count,
-        includeLogits: true,
         forcedContinuation: prompt.targets
     )
 

@@ -157,7 +157,7 @@ class MockEngine: InferenceEngine, @unchecked Sendable {
                 generationToken: GenerationToken
             ) {
                 self.engine = engine
-                self.returnsLogits = inferenceOptions.includeLogits
+                self.returnsLogits = inferenceOptions.returnsLogits
                 self.forcedContinuation = inferenceOptions.forcedContinuation
                 self.stopReasonStore = stopReasonStore
                 self.generationToken = generationToken