Add frequency_penalty, presence_penalty, and logit_bias sampling options

swift/Sources/CoreAILMCommon/ServerAPITypes.swift

@@ -16,6 +16,15 @@ public struct ChatCompletionRequest: Decodable, Sendable {
     public let topP: Double?
     public let topK: Int?
     public let seed: UInt64?
+    /// Additive penalty scaled by how often a token has already appeared (range [-2, 2]).
+    /// Applied as `logit -= frequencyPenalty * count`.
+    public let frequencyPenalty: Double?
+    /// Additive penalty applied once if a token has appeared at all (range [-2, 2]).
+    /// Applied as `logit -= presencePenalty` for any seen token.
+    public let presencePenalty: Double?
+    /// Per-token logit bias keyed by token-ID string → bias value (range [-100, 100]).
+    /// Applied as `logit += bias` before sampling.
+    public let logitBias: [String: Double]?
     public let stream: Bool?
     public let stop: [String]?
     public let responseFormat: ResponseFormat?
@@ -33,6 +42,9 @@ public struct ChatCompletionRequest: Decodable, Sendable {
         case maxCompletionTokens = "max_completion_tokens"
         case topP = "top_p"
         case topK = "top_k"
+        case frequencyPenalty = "frequency_penalty"
+        case presencePenalty = "presence_penalty"
+        case logitBias = "logit_bias"
         case responseFormat = "response_format"
         case toolChoice = "tool_choice"
         case parallelToolCalls = "parallel_tool_calls"
@@ -49,6 +61,9 @@ public struct ChatCompletionRequest: Decodable, Sendable {
         topP = try container.decodeIfPresent(Double.self, forKey: .topP)
         topK = try container.decodeIfPresent(Int.self, forKey: .topK)
         seed = try container.decodeIfPresent(UInt64.self, forKey: .seed)
+        frequencyPenalty = try container.decodeIfPresent(Double.self, forKey: .frequencyPenalty)
+        presencePenalty = try container.decodeIfPresent(Double.self, forKey: .presencePenalty)
+        logitBias = try container.decodeIfPresent([String: Double].self, forKey: .logitBias)
         stream = try container.decodeIfPresent(Bool.self, forKey: .stream)
         responseFormat = try container.decodeIfPresent(ResponseFormat.self, forKey: .responseFormat)
         tools = try container.decodeIfPresent([ToolDefinition].self, forKey: .tools)

swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift

@@ -148,6 +148,19 @@ public struct ConstrainedDecodingStrategy: DecodingStrategy {
                 penalty: Float(penalty)
             )
         }
+        if samplingConfiguration.needsAdditiveLogitProcessing {
+            let window =
+                samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) }
+                ?? generatedTokens.count
+            AdditivePenaltyProcessor.apply(
+                to: &maskedLogits,
+                recentTokenIds: generatedTokens.suffix(window),
+                frequencyPenalty: Float(samplingConfiguration.frequencyPenalty ?? 0),
+                presencePenalty: Float(samplingConfiguration.presencePenalty ?? 0),
+                logitBias: samplingConfiguration.logitBias
+            )
+        }
+        // Apply the grammar mask last so its token bans stay authoritative over any bias.
         _ = session.applyMask(to: &maskedLogits)
 
         let bestToken = samplingConfiguration.sampleToken(from: &maskedLogits, step: generatedTokens.count)

swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift

@@ -216,6 +216,19 @@ public struct ConstrainedGenerator: DecodingStrategy {
                     penalty: Float(penalty)
                 )
             }
+            if samplingConfiguration.needsAdditiveLogitProcessing {
+                let window =
+                    samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) }
+                    ?? generatedTokens.count
+                AdditivePenaltyProcessor.apply(
+                    to: &maskedLogits,
+                    recentTokenIds: generatedTokens.suffix(window),
+                    frequencyPenalty: Float(samplingConfiguration.frequencyPenalty ?? 0),
+                    presencePenalty: Float(samplingConfiguration.presencePenalty ?? 0),
+                    logitBias: samplingConfiguration.logitBias
+                )
+            }
+            // Apply the grammar mask last so its token bans stay authoritative over any bias.
             _ = session.applyMask(to: &maskedLogits)
 
             let bestToken = samplingConfiguration.sampleToken(from: &maskedLogits, step: generatedTokens.count)

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift

@@ -604,7 +604,9 @@ private struct EngineImpl: ~Copyable {
     // GPU sampler — reuses MPSGraphSampler from MPSGraphSamplers.swift
     var cachedSampler: (any MPSGraphSampler)?
     var cachedSamplerTemperature: Double?
+    var cachedSamplerPenaltySignature: PenaltySignature?
     var penaltyState: RepetitionPenaltyGPUState?
+    var additivePenaltyState: AdditivePenaltyGPUState?
 
     // State
     var processedTokenCount: Int = 0
@@ -856,6 +858,9 @@ private struct EngineImpl: ~Copyable {
         self.logits = logitsRef
         self.cachedSampler = nil
         self.cachedSamplerTemperature = nil
+        self.cachedSamplerPenaltySignature = nil
+        self.penaltyState = nil
+        self.additivePenaltyState = nil
 
         CLILogger.log("CoreAI pipelined engine initialized — Vocab: \(config.vocabSize)")
     }
@@ -880,7 +885,16 @@ private struct EngineImpl: ~Copyable {
                 throw InferenceRuntimeError.genericError(
                     "Temperature changed mid-generation (\(existingTemp) -> \(temperature)). Call reset() first.")
             }
-            return existingSampler
+            // Penalty/bias config is part of the sampler's identity, not just temperature. A
+            // prefix-shared request that changed it (reset() skipped) rebuilds instead of reusing.
+            if cachedSamplerPenaltySignature == PenaltySignature(config) {
+                return existingSampler
+            }
+            cachedSampler = nil
+            cachedSamplerTemperature = nil
+            cachedSamplerPenaltySignature = nil
+            penaltyState = nil
+            additivePenaltyState = nil
         }
 
         // Create penalized sampler if repetition penalty is configured
@@ -896,7 +910,30 @@ private struct EngineImpl: ~Copyable {
                     vocabSize: self.config.vocabSize,
                     pipelineDepth: pipelineDepth,
                     penalty: config.repetitionPenalty!,
-                    windowSize: config.repetitionPenaltyWindow
+                    // nil window means "all history" (SamplingConfiguration contract); the GPU
+                    // ring is fixed-size, so size it to the max possible history (context length)
+                    // to match the CPU path instead of silently capping at the state's default.
+                    windowSize: config.repetitionPenaltyWindow ?? self.config.maxContextLength
+                )
+            }
+        }
+
+        // Create the additive-delta state if a frequency/presence penalty or logit bias is
+        // configured. Unlike the multiplicative repetition penalty, this is supported with greedy
+        // sampling — the additive stage is compiled into the argmax sampler as well.
+        if config.needsAdditiveLogitProcessing {
+            if additivePenaltyState == nil {
+                additivePenaltyState = try AdditivePenaltyGPUState(
+                    device: device,
+                    vocabSize: self.config.vocabSize,
+                    pipelineDepth: pipelineDepth,
+                    frequencyPenalty: config.frequencyPenalty ?? 0,
+                    presencePenalty: config.presencePenalty ?? 0,
+                    logitBias: config.logitBias,
+                    // nil window means "all history" (SamplingConfiguration contract); the GPU
+                    // ring is fixed-size, so size it to the max possible history (context length)
+                    // to match the CPU path instead of silently capping at the state's default.
+                    windowSize: config.repetitionPenaltyWindow ?? self.config.maxContextLength
                 )
             }
         }
@@ -908,9 +945,28 @@ private struct EngineImpl: ~Copyable {
         )
         cachedSampler = newSampler
         cachedSamplerTemperature = temperature
+        cachedSamplerPenaltySignature = PenaltySignature(config)
         return newSampler
     }
 
+    /// Identifies the penalty/logit-bias configuration a cached sampler was built for, so the cache
+    /// can be invalidated when a new request changes it (see `getOrCreateSampler`).
+    struct PenaltySignature: Equatable {
+        let repetitionPenalty: Double?
+        let repetitionPenaltyWindow: Int?
+        let frequencyPenalty: Double
+        let presencePenalty: Double
+        let logitBias: [Int32: Float]
+
+        init(_ config: SamplingConfiguration) {
+            self.repetitionPenalty = config.needsRepetitionPenalty ? config.repetitionPenalty : nil
+            self.repetitionPenaltyWindow = config.repetitionPenaltyWindow
+            self.frequencyPenalty = config.frequencyPenalty ?? 0
+            self.presencePenalty = config.presencePenalty ?? 0
+            self.logitBias = config.logitBias ?? [:]
+        }
+    }
+
     // MARK: - Core Encode Step
 
     /// Encodes inference + GPU sampling for one step.
@@ -1062,9 +1118,11 @@ private struct EngineImpl: ~Copyable {
         let queue = pipelineQueue
         let localInFlightGate = inFlightGate
         let localPenaltyState = penaltyState
+        let localAdditiveState = additivePenaltyState
         let completionCallback: (Int32, Error?) -> Void = { nextToken, error in
-            // Update penalty state BEFORE releasing the gate.
+            // Update penalty/additive state BEFORE releasing the gate.
             localPenaltyState?.recordToken(nextToken)
+            localAdditiveState?.recordToken(nextToken)
             // Release the pipeline slot acquired before encode. Happens on
             // Metal's callback thread — PipelineGate.release() is thread-safe.
             localInFlightGate.release()
@@ -1081,17 +1139,34 @@ private struct EngineImpl: ~Copyable {
         }
 
         do {
-            // Use penalty-aware path for decode steps when penalty is active.
-            if queryLength == 1, let state = penaltyState,
-                let compositeSampler = localGPUSampler as? MPSGraphCompositeSampler,
-                compositeSampler.penaltyEnabled
+            if queryLength == 1, let compositeSampler = localGPUSampler as? MPSGraphCompositeSampler,
+                compositeSampler.penaltyEnabled || compositeSampler.additiveEnabled
             {
-                let penaltyBuf = state.buffer(forStep: currentStep)
+                // Composite decode with repetition penalty and/or additive delta. Each buffer is
+                // fetched only when used, so its slot's queued increments are consumed exactly once.
+                let penaltyBuf = compositeSampler.penaltyEnabled ? penaltyState?.buffer(forStep: currentStep) : nil
+                let additiveBuf =
+                    compositeSampler.additiveEnabled ? additivePenaltyState?.buffer(forStep: currentStep) : nil
                 compositeSampler.encode(
                     to: queue,
                     logitsBuffer: samplerLogitsBuffer,
                     logitsOffset: logitsOffset,
                     penaltyBuffer: penaltyBuf,
+                    additiveBuffer: additiveBuf,
+                    outputBuffer: outputBuffer,
+                    outputOffset: 0,
+                    completion: completionCallback
+                )
+            } else if queryLength == 1, let argmaxSampler = localGPUSampler as? MPSGraphArgmaxSampler,
+                argmaxSampler.additiveEnabled
+            {
+                // Greedy decode with an additive delta (frequency/presence penalty + logit bias).
+                let additiveBuf = additivePenaltyState?.buffer(forStep: currentStep)
+                argmaxSampler.encode(
+                    to: queue,
+                    logitsBuffer: samplerLogitsBuffer,
+                    logitsOffset: logitsOffset,
+                    additiveBuffer: additiveBuf,
                     outputBuffer: outputBuffer,
                     outputOffset: 0,
                     completion: completionCallback
@@ -1105,6 +1180,36 @@ private struct EngineImpl: ~Copyable {
                     outputOffset: 0,
                     completion: completionCallback
                 )
+            } else if let compositeSampler = localGPUSampler as? MPSGraphCompositeSampler,
+                compositeSampler.additiveEnabled
+            {
+                // Prefill's first token with an additive delta (composite path). The multiplicative
+                // repetition penalty is intentionally not applied during prefill, matching the
+                // existing behavior; the additive delta at this point is just the static bias.
+                let additiveBuf = additivePenaltyState?.buffer(forStep: currentStep)
+                compositeSampler.encodeWithSlice(
+                    to: queue,
+                    logitsBuffer: samplerLogitsBuffer,
+                    queryLength: actualTokenCount,
+                    additiveBuffer: additiveBuf,
+                    outputBuffer: outputBuffer,
+                    outputOffset: 0,
+                    completion: completionCallback
+                )
+            } else if let argmaxSampler = localGPUSampler as? MPSGraphArgmaxSampler,
+                argmaxSampler.additiveEnabled
+            {
+                // Prefill's first token with an additive delta (greedy path).
+                let additiveBuf = additivePenaltyState?.buffer(forStep: currentStep)
+                argmaxSampler.encodeWithSlice(
+                    to: queue,
+                    logitsBuffer: samplerLogitsBuffer,
+                    queryLength: actualTokenCount,
+                    additiveBuffer: additiveBuf,
+                    outputBuffer: outputBuffer,
+                    outputOffset: 0,
+                    completion: completionCallback
+                )
             } else {
                 localGPUSampler.encodeWithSlice(
                     to: queue,
@@ -1764,6 +1869,10 @@ private struct EngineImpl: ~Copyable {
         step = 0
         cachedSampler = nil
         cachedSamplerTemperature = nil
+        cachedSamplerPenaltySignature = nil
+        // Drop the additive-delta state so the next request rebuilds it with its own logit bias and
+        // penalties (the delta baseline is request-specific, unlike the reusable repetition state).
+        additivePenaltyState = nil
         // Zero SSM states so the next conversation starts from a clean slate.
         additionalStates?.reset()
         span.end()