diff --git a/swift/Sources/CoreAILMCommon/ServerAPITypes.swift b/swift/Sources/CoreAILMCommon/ServerAPITypes.swift index 09360dce..e5d58d89 100644 --- a/swift/Sources/CoreAILMCommon/ServerAPITypes.swift +++ b/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) diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift index f5873808..637ee06e 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift +++ b/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) diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift index 9e621323..2a113170 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift +++ b/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) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 8acb75ca..5005e194 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/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() diff --git a/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift new file mode 100644 index 00000000..4beb1f4c --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift @@ -0,0 +1,166 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Foundation +import Metal + +/// Manages per-pipeline-depth additive delta buffers for GPU frequency/presence penalties and +/// logit bias. +/// +/// Produces an ADDITIVE delta buffer `[1, vocabSize]` (f16) where, for each token id: +/// +/// delta[id] = logitBias[id] - frequencyPenalty * count[id] - (count[id] > 0 ? presence : 0) +/// +/// The GPU sampler simply adds this delta to the logits. This mirrors the CPU +/// `AdditivePenaltyProcessor` exactly. +/// +/// Uses the same split design as ``RepetitionPenaltyGPUState`` to avoid races between GPU reads +/// and CPU writes: +/// - `recordToken()`: updates only the CPU-side ring buffer and queues per-slot delta increments +/// - `buffer(forStep:)`: applies the queued increments to one buffer slot at encode time, when the +/// pipeline gate guarantees that slot is not being read by the GPU +/// +/// Unlike the multiplicative repetition penalty (whose neutral value is 1.0), the additive delta's +/// neutral value is 0.0. Each buffer starts at the static `logitBias` baseline and accumulates the +/// negative frequency/presence increments as tokens enter and leave the window. +final class AdditivePenaltyGPUState: @unchecked Sendable { + let deltaBuffers: [MTLBuffer] + let vocabSize: Int + let pipelineDepth: Int + let frequencyPenalty: Float + let presencePenalty: Float + let windowSize: Int + + /// Static per-token bias baseline (applied to every buffer at init and on reset). + private let logitBias: [Int32: Float] + + private var ring: [Int32] + private var writeIndex: Int = 0 + private var count: Int = 0 + private var refCounts: [Int32: Int] = [:] + private var dirtyIncrements: [[(token: Int32, change: Float)]] + + init( + device: MTLDevice, + vocabSize: Int, + pipelineDepth: Int, + frequencyPenalty: Double = 0, + presencePenalty: Double = 0, + logitBias: [Int32: Float]? = nil, + windowSize: Int? = nil + ) throws { + self.vocabSize = vocabSize + self.pipelineDepth = pipelineDepth + self.frequencyPenalty = Float(frequencyPenalty) + self.presencePenalty = Float(presencePenalty) + self.windowSize = windowSize ?? 256 + self.logitBias = logitBias?.filter { $0.key >= 0 && Int($0.key) < vocabSize } ?? [:] + + let bufferSize = vocabSize * MemoryLayout.size + var buffers: [MTLBuffer] = [] + for _ in 0.. MTLBuffer { + let slot = step % pipelineDepth + let buf = deltaBuffers[slot] + let ptr = buf.contents().assumingMemoryBound(to: Float16.self) + + let applied = dirtyIncrements[slot] + for (token, change) in applied { + let idx = Int(token) + ptr[idx] = Float16(Float(ptr[idx]) + change) + } + dirtyIncrements[slot].removeAll(keepingCapacity: true) + + return buf + } + + /// Record a newly generated token (CPU-side bookkeeping only). + /// + /// Called from the completion callback. Does NOT write to MTLBuffers directly. Instead, queues + /// per-slot delta increments to be applied at the next `buffer(forStep:)` call. Increments are + /// computed from `refCounts` at event time so the frequency/presence math is exact. + func recordToken(_ token: Int32) { + guard token >= 0 && Int(token) < vocabSize else { return } + + // Evict the oldest token if the window is full. + if count == windowSize { + let candidate = ring[writeIndex] + if candidate >= 0 { + let oldCount = refCounts[candidate, default: 0] + let newCount = oldCount - 1 + if newCount <= 0 { + refCounts.removeValue(forKey: candidate) + } else { + refCounts[candidate] = newCount + } + // Removing one occurrence relaxes the frequency penalty; if the token leaves the + // window entirely, the presence penalty is lifted too. + var change = frequencyPenalty + if newCount <= 0 { change += presencePenalty } + enqueue(token: candidate, change: change) + } + } else { + count += 1 + } + + // Add the new token. + let oldCount = refCounts[token, default: 0] + var change = -frequencyPenalty + if oldCount == 0 { change -= presencePenalty } + enqueue(token: token, change: change) + + ring[writeIndex] = token + writeIndex = (writeIndex + 1) % windowSize + refCounts[token, default: 0] += 1 + } + + private func enqueue(token: Int32, change: Float) { + guard change != 0 else { return } + for i in 0...size + for buf in deltaBuffers { + memset(buf.contents(), 0, bufferSize) + let ptr = buf.contents().assumingMemoryBound(to: Float16.self) + for (tokenId, bias) in logitBias { + ptr[Int(tokenId)] = Float16(bias) + } + } + ring = [Int32](repeating: -1, count: windowSize) + writeIndex = 0 + count = 0 + refCounts.removeAll(keepingCapacity: true) + dirtyIncrements = Array(repeating: [], count: pipelineDepth) + } +} diff --git a/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift new file mode 100644 index 00000000..dcc3220d --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift @@ -0,0 +1,73 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import CoreAIShared + +/// Applies additive logit adjustments to the logits before sampling: count-scaled frequency +/// penalty, once-per-token presence penalty, and a fixed per-token logit bias. +/// +/// These are additive and differ from ``RepetitionPenaltyProcessor`` (which is multiplicative and +/// count-agnostic): +/// - **frequency penalty**: `logit -= frequencyPenalty * count` for each token in the history +/// - **presence penalty**: `logit -= presencePenalty` once for any token that appeared at all +/// - **logit bias**: `logit += bias` for each biased token id, regardless of history +/// +/// The frequency/presence penalties are computed from the recent token history; the logit bias is +/// a static, request-level vector. +public struct AdditivePenaltyProcessor { + /// Apply the additive frequency/presence penalties and the logit bias to logits in-place. + /// + /// - Parameters: + /// - logits: Mutable logits array (vocab-sized). Modified in-place. + /// - recentTokenIds: Token IDs from recent generation history, used for the count-scaled + /// frequency penalty and the presence penalty. + /// - frequencyPenalty: Additive penalty scaled by occurrence count (0 = disabled). + /// - presencePenalty: Additive penalty applied once per seen token (0 = disabled). + /// - logitBias: Optional per-token-id additive bias applied after the penalties. + public static func apply>( + to logits: inout [LogitsScalarType], + recentTokenIds: C, + frequencyPenalty: Float, + presencePenalty: Float, + logitBias: [Int32: Float]? + ) { + let vocabSize = logits.count + + if frequencyPenalty != 0.0 || presencePenalty != 0.0 { + var counts = [Int32: Int](minimumCapacity: min(recentTokenIds.count, 512)) + for tokenId in recentTokenIds { + guard tokenId >= 0 && Int(tokenId) < vocabSize else { continue } + counts[tokenId, default: 0] += 1 + } + for (tokenId, count) in counts { + let idx = Int(tokenId) + let logit = Float(logits[idx]) + let delta = frequencyPenalty * Float(count) + presencePenalty + logits[idx] = LogitsScalarType(logit - delta) + } + } + + if let logitBias { + applyLogitBias(to: &logits, logitBias: logitBias) + } + } + + /// Apply only the per-token logit bias in-place (no history required). + /// + /// - Parameters: + /// - logits: Mutable logits array (vocab-sized). Modified in-place. + /// - logitBias: Per-token-id additive bias (`logit += bias`). + public static func applyLogitBias( + to logits: inout [LogitsScalarType], + logitBias: [Int32: Float] + ) { + let vocabSize = logits.count + for (tokenId, bias) in logitBias { + guard tokenId >= 0 && Int(tokenId) < vocabSize else { continue } + let idx = Int(tokenId) + logits[idx] = LogitsScalarType(Float(logits[idx]) + bias) + } + } +} diff --git a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift index 99e15060..b09c9bfc 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift @@ -114,7 +114,9 @@ enum MPSGraphSamplerFactory { config: SamplingConfiguration ) throws -> any MPSGraphSampler { if config.temperature == 0 { - return try MPSGraphArgmaxSampler(device: device, vocabSize: vocabSize) + return try MPSGraphArgmaxSampler( + device: device, vocabSize: vocabSize, + additiveEnabled: config.needsAdditiveLogitProcessing) } // Determine effective K for the topK operation: @@ -137,7 +139,8 @@ enum MPSGraphSamplerFactory { temperature: Float(config.temperature), topP: config.topP.map { Float($0) } ?? 1.0, minP: config.minP.map { Float($0) } ?? 0.0, - penaltyEnabled: config.needsRepetitionPenalty + penaltyEnabled: config.needsRepetitionPenalty, + additiveEnabled: config.needsAdditiveLogitProcessing ) } @@ -203,6 +206,21 @@ private func buildBitmaskExpansionGraph( return graph.select(predicate: predicate, trueTensor: logits, falseTensor: negInf, name: "masked_logits") } +/// Builds the additive-delta stage: `logits + delta`, where `delta` carries the frequency/presence +/// penalty and the logit bias. The delta placeholder is f16 with a neutral value of 0.0. +/// +/// - Parameters: +/// - graph: The graph to add the operation to. +/// - logits: The current logits tensor (float32). +/// - additiveTensor: The f16 additive-delta placeholder `[1, vocabSize]`. +/// - Returns: `logits + cast(additiveTensor, float32)`. +private func applyAdditiveStage( + graph: MPSGraph, logits: MPSGraphTensor, additiveTensor: MPSGraphTensor, name: String +) -> MPSGraphTensor { + let additiveF32 = graph.cast(additiveTensor, to: .float32, name: "\(name)_f32") + return graph.addition(logits, additiveF32, name: name) +} + // MARK: - MPSGraph Argmax Sampler /// MPSGraph-based argmax sampler using Apple's optimized reductionArgMaximum. @@ -230,12 +248,16 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { private let mpsDevice: MPSGraphDevice private let graph: MPSGraph private let inputPlaceholder: MPSGraphTensor + private let additivePlaceholder: MPSGraphTensor? private let outputTensor: MPSGraphTensor private let executable: MPSGraphExecutable /// The vocabulary size this sampler was compiled for let vocabSize: Int + /// Whether an additive delta (frequency/presence penalty + logit bias) is compiled into the graph. + let additiveEnabled: Bool + // Pre-allocated objects reused every step to avoid ~70µs of CPU object creation. // MPSGraphTensorData wraps MTLBuffer references — safe to reuse when buffers match. private var cachedInputData: MPSGraphTensorData? @@ -243,10 +265,15 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { private var cachedInputBuffer: MTLBuffer? private var cachedOutputBuffer: MTLBuffer? + // Pre-allocated neutral additive buffer (all 0.0 in f16) used by non-additive encode paths + // when the executable was compiled with additive support. + private let neutralAdditiveData: MPSGraphTensorData? + // Constrained sampling — compiled lazily on first applyBitmask: true call. private var constrainedExecutable: MPSGraphExecutable? private var constrainedBitmaskBuffer: MTLBuffer? private var constrainedBitmaskData: MPSGraphTensorData? + private var constrainedAdditivePlaceholder: MPSGraphTensor? /// Number of Int32 words in the bitmask. let bitmaskSize: Int @@ -255,13 +282,16 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { /// - Parameters: /// - device: Metal device /// - vocabSize: Vocabulary size (fixed for compilation) - init(device: MTLDevice, vocabSize: Int) throws { + /// - additiveEnabled: Whether to compile in an additive-delta stage (frequency/presence + /// penalty + logit bias) applied before argmax. + init(device: MTLDevice, vocabSize: Int, additiveEnabled: Bool = false) throws { guard vocabSize > 0 else { throw MPSGraphSamplerError.graphCompilationFailed } self.device = device self.mpsDevice = MPSGraphDevice(mtlDevice: device) self.vocabSize = vocabSize + self.additiveEnabled = additiveEnabled self.bitmaskSize = (vocabSize + 31) / 32 // Build the argmax graph @@ -276,10 +306,25 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { ) self.inputPlaceholder = inputPlaceholder + // Optional additive delta [1, vocabSize] f16 (frequency/presence penalty + logit bias) + let argmaxInput: MPSGraphTensor + if additiveEnabled { + let additive = graph.placeholder( + shape: [1, vocabSize as NSNumber], dataType: .float16, name: "additive") + self.additivePlaceholder = additive + let logitsF32 = graph.cast(inputPlaceholder, to: .float32, name: "logits_f32") + let biased = applyAdditiveStage( + graph: graph, logits: logitsF32, additiveTensor: additive, name: "additive") + argmaxInput = biased + } else { + self.additivePlaceholder = nil + argmaxInput = inputPlaceholder + } + // Argmax along axis 1 (vocab dimension) - returns Int64 // No reshape needed! Just reduce along the vocab dimension. let argmaxInt64 = graph.reductionArgMaximum( - with: inputPlaceholder, + with: argmaxInput, axis: 1, // axis 1 = vocab dimension in [1, vocabSize] name: "argmax" ) @@ -293,12 +338,15 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { self.outputTensor = outputTensor // Compile to executable - let feeds: [MPSGraphTensor: MPSGraphShapedType] = [ + var feeds: [MPSGraphTensor: MPSGraphShapedType] = [ inputPlaceholder: MPSGraphShapedType( shape: [1, vocabSize as NSNumber], dataType: .float16 ) ] + if let additive = additivePlaceholder { + feeds[additive] = MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16) + } let targetTensors = [outputTensor] @@ -313,6 +361,36 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { targetOperations: nil, compilationDescriptor: compilationDescriptor ) + + // Pre-allocate a neutral additive buffer (all 0.0 in f16) so non-additive encode paths can + // feed the additive-enabled executable without a caller-provided delta buffer. + if additiveEnabled { + let neutralByteCount = vocabSize * MemoryLayout.size + guard let neutralBuf = device.makeBuffer(length: neutralByteCount, options: .storageModeShared) else { + throw MPSGraphSamplerError.bufferAllocationFailed + } + // Float16 0.0 is all-zero bits, so a single memset is the additive identity. + memset(neutralBuf.contents(), 0, neutralByteCount) + self.neutralAdditiveData = MPSGraphTensorData( + neutralBuf, shape: [1, vocabSize as NSNumber], dataType: .float16) + } else { + self.neutralAdditiveData = nil + } + } + + /// Build the argmax executable inputs in feed order, injecting the additive delta when enabled. + private func argmaxInputs( + logitsData: MPSGraphTensorData, additiveData: MPSGraphTensorData? + ) -> [MPSGraphTensorData] { + let effectiveAdditive = additiveData ?? neutralAdditiveData + return executable.feedTensors!.map { tensor -> MPSGraphTensorData in + switch tensor.operation.name { + case "logits": return logitsData + case "additive": return effectiveAdditive! + default: + fatalError("MPSGraphArgmaxSampler: unknown feed tensor '\(tensor.operation.name)'") + } + } } // MARK: - Constrained Sampling (Lazy) @@ -356,13 +434,29 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { graph: cGraph, logits: cLogits, bitmaskPlaceholder: cBitmask, vocabSize: vocabSize, bitmaskSize: bitmaskSize) - let cArgmax = cGraph.reductionArgMaximum(with: maskedLogits, axis: 1, name: "argmax") + // Apply the additive delta after masking so grammar bans stay authoritative. + let argmaxInput: MPSGraphTensor + if additiveEnabled { + let cAdditive = cGraph.placeholder( + shape: [1, vocabSize as NSNumber], dataType: .float16, name: "additive") + self.constrainedAdditivePlaceholder = cAdditive + let maskedF32 = cGraph.cast(maskedLogits, to: .float32, name: "logits_f32") + argmaxInput = applyAdditiveStage( + graph: cGraph, logits: maskedF32, additiveTensor: cAdditive, name: "additive") + } else { + argmaxInput = maskedLogits + } + + let cArgmax = cGraph.reductionArgMaximum(with: argmaxInput, axis: 1, name: "argmax") let cOutput = cGraph.cast(cArgmax, to: .int32, name: "token_id") - let cFeeds: [MPSGraphTensor: MPSGraphShapedType] = [ + var cFeeds: [MPSGraphTensor: MPSGraphShapedType] = [ cLogits: MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16), cBitmask: MPSGraphShapedType(shape: [bitmaskSize as NSNumber], dataType: .int32), ] + if let cAdditive = constrainedAdditivePlaceholder { + cFeeds[cAdditive] = MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16) + } let desc = MPSGraphCompilationDescriptor() desc.optimizationLevel = .level0 @@ -374,6 +468,22 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { return (exec, constrainedBitmaskData!) } + /// Build the constrained argmax inputs in feed order, injecting bitmask and additive delta. + private func constrainedInputs( + logitsData: MPSGraphTensorData, bitmaskData: MPSGraphTensorData, additiveData: MPSGraphTensorData? + ) -> [MPSGraphTensorData] { + let effectiveAdditive = additiveData ?? neutralAdditiveData + return constrainedExecutable!.feedTensors!.map { tensor -> MPSGraphTensorData in + switch tensor.operation.name { + case "logits": return logitsData + case "bitmask": return bitmaskData + case "additive": return effectiveAdditive! + default: + fatalError("MPSGraphArgmaxSampler: unknown constrained feed tensor '\(tensor.operation.name)'") + } + } + } + /// Encode argmax sampling with optional bitmask constraint. /// /// When `applyBitmask` is true, the bitmask in `bitmaskBuffer` is applied to @@ -387,6 +497,23 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { outputOffset: Int, applyBitmask: Bool, completion: @escaping (Int32, Error?) -> Void + ) throws { + try encode( + to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset, + outputBuffer: outputBuffer, outputOffset: outputOffset, + applyBitmask: applyBitmask, additiveBuffer: nil, completion: completion) + } + + /// Encode argmax sampling with optional bitmask constraint and an optional additive delta. + func encode( + to queue: MTLCommandQueue, + logitsBuffer: MTLBuffer, + logitsOffset: Int, + outputBuffer: MTLBuffer, + outputOffset: Int, + applyBitmask: Bool, + additiveBuffer: MTLBuffer?, + completion: @escaping (Int32, Error?) -> Void ) throws { if applyBitmask { guard let (constrained, bitmaskTensorData) = try? ensureConstrainedResources() else { @@ -397,6 +524,9 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { logitsBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) let outputData = MPSGraphTensorData( outputBuffer, shape: [1 as NSNumber], dataType: .int32) + let additiveData = additiveBuffer.map { + MPSGraphTensorData($0, shape: [1, vocabSize as NSNumber], dataType: .float16) + } let execDesc = MPSGraphExecutableExecutionDescriptor() execDesc.completionHandler = { [outputBuffer, outputOffset] (_, error) in if let error = error { @@ -409,11 +539,14 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { completion(result, nil) } constrained.runAsync( - with: queue, inputs: [inputData, bitmaskTensorData], + with: queue, + inputs: constrainedInputs( + logitsData: inputData, bitmaskData: bitmaskTensorData, additiveData: additiveData), results: [outputData], executionDescriptor: execDesc) } else { encode( to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset, + additiveBuffer: additiveBuffer, outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) } @@ -531,7 +664,51 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { executable.runAsync( with: queue, - inputs: [inputData], + inputs: argmaxInputs(logitsData: inputData, additiveData: nil), + results: [outputData], + executionDescriptor: execDescriptor + ) + } + + /// Encode argmax sampling with an additive delta (frequency/presence penalty + logit bias). + func encode( + to queue: MTLCommandQueue, + logitsBuffer: MTLBuffer, + logitsOffset: Int, + additiveBuffer: MTLBuffer?, + outputBuffer: MTLBuffer, + outputOffset: Int, + completion: @escaping (Int32, Error?) -> Void + ) { + guard additiveEnabled, let additiveBuffer else { + encode( + to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + return + } + let inputData = MPSGraphTensorData( + logitsBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) + let additiveData = MPSGraphTensorData( + additiveBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) + let outputData = MPSGraphTensorData( + outputBuffer, shape: [1 as NSNumber], dataType: .int32) + + let execDescriptor = MPSGraphExecutableExecutionDescriptor() + execDescriptor.completionHandler = { [outputBuffer, outputOffset] (_, error) in + if error != nil { + completion(0, error) + return + } + let result = outputBuffer.contents() + .advanced(by: outputOffset) + .assumingMemoryBound(to: Int32.self) + .pointee + completion(result, nil) + } + + executable.runAsync( + with: queue, + inputs: argmaxInputs(logitsData: inputData, additiveData: additiveData), results: [outputData], executionDescriptor: execDescriptor ) @@ -635,11 +812,58 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { // Run async - GPU naturally orders this after the blit due to queue ordering executable.runAsync( with: queue, - inputs: [inputData], + inputs: argmaxInputs(logitsData: inputData, additiveData: nil), results: [outputData], executionDescriptor: execDescriptor ) } + + /// Encode argmax sampling with slice support and an additive delta. + /// + /// For prefill (`queryLength > 1`) the last token's logits are blitted to a temporary buffer + /// and sampled with the additive delta applied. + func encodeWithSlice( + to queue: MTLCommandQueue, + logitsBuffer: MTLBuffer, + queryLength: Int, + additiveBuffer: MTLBuffer?, + outputBuffer: MTLBuffer, + outputOffset: Int, + completion: @escaping (Int32, Error?) -> Void + ) { + guard additiveEnabled, let additiveBuffer else { + encodeWithSlice( + to: queue, logitsBuffer: logitsBuffer, queryLength: queryLength, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + return + } + if queryLength == 1 { + encode( + to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0, + additiveBuffer: additiveBuffer, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + return + } + let logitsOffset = (queryLength - 1) * vocabSize * MemoryLayout.size + let sliceSize = vocabSize * MemoryLayout.size + guard let tempBuffer = device.makeBuffer(length: sliceSize, options: .storageModeShared), + let blitCmdBuffer = queue.makeCommandBuffer(), + let blitEncoder = blitCmdBuffer.makeBlitCommandEncoder() + else { + completion(0, MPSGraphSamplerError.bufferAllocationFailed) + return + } + blitEncoder.copy( + from: logitsBuffer, sourceOffset: logitsOffset, + to: tempBuffer, destinationOffset: 0, size: sliceSize) + blitEncoder.endEncoding() + blitCmdBuffer.commit() + + encode( + to: queue, logitsBuffer: tempBuffer, logitsOffset: 0, + additiveBuffer: additiveBuffer, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + } } // Conformance to MPSGraphSampler protocol @@ -672,6 +896,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { // Graph tensors private let logitsPlaceholder: MPSGraphTensor private let penaltyPlaceholder: MPSGraphTensor? + private let additivePlaceholder: MPSGraphTensor? private let temperaturePlaceholder: MPSGraphTensor private let randomPlaceholder: MPSGraphTensor private let topPPlaceholder: MPSGraphTensor @@ -698,6 +923,9 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { /// Whether repetition penalty is compiled into this sampler's graph let penaltyEnabled: Bool + /// Whether an additive delta (frequency/presence penalty + logit bias) is compiled into the graph. + let additiveEnabled: Bool + /// Pre-allocated buffer for random value private let randomBuffer: MTLBuffer @@ -724,10 +952,15 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { // encode paths when the executable was compiled with penalty support. private let neutralPenaltyData: MPSGraphTensorData? + // Pre-allocated neutral additive buffer (all 0.0 in f16) used by non-additive + // encode paths when the executable was compiled with additive support. + private let neutralAdditiveData: MPSGraphTensorData? + // Constrained sampling — compiled lazily on first applyBitmask: true call. private var constrainedExecutable: MPSGraphExecutable? private var constrainedBitmaskBuffer: MTLBuffer? private var constrainedBitmaskData: MPSGraphTensorData? + private var constrainedAdditivePlaceholder: MPSGraphTensor? /// Number of Int32 words in the bitmask. let bitmaskSize: Int @@ -745,7 +978,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { /// - minP: Minimum probability threshold (0.0 = disabled) init( device: MTLDevice, vocabSize: Int, k: Int = 40, temperature: Float = 1.0, topP: Float = 1.0, minP: Float = 0.0, - penaltyEnabled: Bool = false + penaltyEnabled: Bool = false, additiveEnabled: Bool = false ) throws { @@ -757,6 +990,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { self.topP = topP self.minP = minP self.penaltyEnabled = penaltyEnabled + self.additiveEnabled = additiveEnabled self.bitmaskSize = (vocabSize + 31) / 32 // Pre-allocate buffers @@ -795,6 +1029,17 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { self.penaltyPlaceholder = nil } + if additiveEnabled { + let ap = graph.placeholder( + shape: [1, vocabSize as NSNumber], + dataType: .float16, + name: "additive" + ) + self.additivePlaceholder = ap + } else { + self.additivePlaceholder = nil + } + // Temperature scalar [1] let temperaturePlaceholder = graph.placeholder( shape: [1 as NSNumber], @@ -839,8 +1084,18 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { penalizedLogits = logitsFloat32 } + // Apply the additive delta (frequency/presence penalty + logit bias) after the + // multiplicative repetition penalty, matching the CPU ordering. + let adjustedLogits: MPSGraphTensor + if additiveEnabled { + adjustedLogits = applyAdditiveStage( + graph: graph, logits: penalizedLogits, additiveTensor: additivePlaceholder!, name: "additive") + } else { + adjustedLogits = penalizedLogits + } + let (topKValues, topKIndices) = Self.topKStage( - graph: graph, logits: penalizedLogits, k: k, name: "topk") + graph: graph, logits: adjustedLogits, k: k, name: "topk") let scaledValues = Self.temperatureStage( graph: graph, values: topKValues, temperature: temperaturePlaceholder, name: "temp") @@ -874,6 +1129,9 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { if let pp = penaltyPlaceholder { feeds[pp] = MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16) } + if let ap = additivePlaceholder { + feeds[ap] = MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16) + } let compilationDescriptor = MPSGraphCompilationDescriptor() compilationDescriptor.optimizationLevel = .level0 @@ -926,6 +1184,21 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { } else { self.neutralPenaltyData = nil } + + // Pre-allocate a neutral additive buffer (all 0.0 in f16) so that non-additive encode + // paths can feed the additive-enabled executable without a caller-provided delta buffer. + if additiveEnabled { + let neutralByteCount = vocabSize * MemoryLayout.size + guard let neutralBuf = device.makeBuffer(length: neutralByteCount, options: .storageModeShared) else { + throw MPSGraphSamplerError.bufferAllocationFailed + } + // Float16 0.0 is all-zero bits, so a single memset is the additive identity. + memset(neutralBuf.contents(), 0, neutralByteCount) + self.neutralAdditiveData = MPSGraphTensorData( + neutralBuf, shape: [1, vocabSize as NSNumber], dataType: .float16) + } else { + self.neutralAdditiveData = nil + } } // MARK: - Feed Tensor Ordering @@ -938,13 +1211,16 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { /// (e.g. f32 scalar data landing where an f16 penalty tensor is expected). private func buildInputs( logitsData: MPSGraphTensorData, - penaltyData: MPSGraphTensorData? = nil + penaltyData: MPSGraphTensorData? = nil, + additiveData: MPSGraphTensorData? = nil ) -> [MPSGraphTensorData] { let effectivePenalty = penaltyData ?? neutralPenaltyData + let effectiveAdditive = additiveData ?? neutralAdditiveData return executable.feedTensors!.map { tensor -> MPSGraphTensorData in switch tensor.operation.name { case "logits": return logitsData case "penalty": return effectivePenalty! + case "additive": return effectiveAdditive! case "temperature": return temperatureData case "random": return randomData case "topP": return topPData @@ -997,7 +1273,20 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { vocabSize: vocabSize, bitmaskSize: bitmaskSize) let cLogitsF32 = cGraph.cast(cMaskedLogits, to: .float32, name: "logits_f32") - let cTopK = cGraph.topK(cLogitsF32, k: k, name: "topk") + + // Apply the additive delta after masking so grammar bans stay authoritative. + let cAdjusted: MPSGraphTensor + if additiveEnabled { + let cAdditive = cGraph.placeholder( + shape: [1, vocabSize as NSNumber], dataType: .float16, name: "additive") + self.constrainedAdditivePlaceholder = cAdditive + cAdjusted = applyAdditiveStage( + graph: cGraph, logits: cLogitsF32, additiveTensor: cAdditive, name: "additive") + } else { + cAdjusted = cLogitsF32 + } + + let cTopK = cGraph.topK(cAdjusted, k: k, name: "topk") let cScaled = cGraph.division(cTopK[0], cTemp, name: "scaled") let cProbs = cGraph.softMax(with: cScaled, axis: 1, name: "probs") @@ -1025,7 +1314,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { let cSelFlat = cGraph.reshape(cSelI32, shape: [1 as NSNumber], name: "selected_flat") let cOutput = cGraph.gatherAlongAxis(0, updates: cIndFlat, indices: cSelFlat, name: "token_id") - let cFeeds: [MPSGraphTensor: MPSGraphShapedType] = [ + var cFeeds: [MPSGraphTensor: MPSGraphShapedType] = [ cLogits: MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16), cTemp: MPSGraphShapedType(shape: [1 as NSNumber], dataType: .float32), cRandom: MPSGraphShapedType(shape: [1 as NSNumber], dataType: .float32), @@ -1033,6 +1322,9 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { cMinP: MPSGraphShapedType(shape: [1 as NSNumber], dataType: .float32), cBitmask: MPSGraphShapedType(shape: [bitmaskSize as NSNumber], dataType: .int32), ] + if let cAdditive = constrainedAdditivePlaceholder { + cFeeds[cAdditive] = MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16) + } let desc = MPSGraphCompilationDescriptor() desc.optimizationLevel = .level0 @@ -1044,6 +1336,26 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { return (exec, constrainedBitmaskData!) } + /// Build the constrained composite inputs in feed order, injecting the additive delta. + private func constrainedInputs( + logitsData: MPSGraphTensorData, bitmaskData: MPSGraphTensorData, additiveData: MPSGraphTensorData? + ) -> [MPSGraphTensorData] { + let effectiveAdditive = additiveData ?? neutralAdditiveData + return constrainedExecutable!.feedTensors!.map { tensor -> MPSGraphTensorData in + switch tensor.operation.name { + case "logits": return logitsData + case "bitmask": return bitmaskData + case "additive": return effectiveAdditive! + case "temperature": return temperatureData + case "random": return randomData + case "topP": return topPData + case "minP": return minPData + default: + fatalError("MPSGraphCompositeSampler: unknown constrained feed tensor '\(tensor.operation.name)'") + } + } + } + /// Encode composite sampling with optional bitmask constraint. /// /// When `applyBitmask` is true, the bitmask in `bitmaskBuffer` masks logits @@ -1056,6 +1368,23 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { outputOffset: Int, applyBitmask: Bool, completion: @escaping (Int32, Error?) -> Void + ) throws { + try encode( + to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset, + outputBuffer: outputBuffer, outputOffset: outputOffset, + applyBitmask: applyBitmask, additiveBuffer: nil, completion: completion) + } + + /// Encode composite sampling with optional bitmask constraint and an optional additive delta. + func encode( + to queue: MTLCommandQueue, + logitsBuffer: MTLBuffer, + logitsOffset: Int, + outputBuffer: MTLBuffer, + outputOffset: Int, + applyBitmask: Bool, + additiveBuffer: MTLBuffer?, + completion: @escaping (Int32, Error?) -> Void ) throws { if applyBitmask { guard let (constrained, bitmaskTensorData) = try? ensureConstrainedResources() else { @@ -1072,6 +1401,9 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { logitsBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) let outputData = MPSGraphTensorData( outputBuffer, shape: [1 as NSNumber], dataType: .int32) + let additiveData = additiveBuffer.map { + MPSGraphTensorData($0, shape: [1, vocabSize as NSNumber], dataType: .float16) + } let execDesc = MPSGraphExecutableExecutionDescriptor() execDesc.completionHandler = { [outputBuffer, outputOffset] (_, error) in if let error = error { @@ -1085,11 +1417,13 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { } constrained.runAsync( with: queue, - inputs: [logitsData, temperatureData, randomData, topPData, minPData, bitmaskTensorData], + inputs: constrainedInputs( + logitsData: logitsData, bitmaskData: bitmaskTensorData, additiveData: additiveData), results: [outputData], executionDescriptor: execDesc) } else { encode( to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset, + penaltyBuffer: nil, additiveBuffer: additiveBuffer, outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) } @@ -1143,16 +1477,20 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { /// Encode sampling with repetition penalty buffer. /// The penalty buffer must be Float16[vocabSize] with 1.0 for unpenalized tokens. + /// The additive buffer must be Float16[vocabSize] with 0.0 for unbiased tokens. func encode( to queue: MTLCommandQueue, logitsBuffer: MTLBuffer, logitsOffset: Int, - penaltyBuffer: MTLBuffer, + penaltyBuffer: MTLBuffer?, + additiveBuffer: MTLBuffer? = nil, outputBuffer: MTLBuffer, outputOffset: Int, completion: @escaping (Int32, Error?) -> Void ) { - guard penaltyEnabled else { + // If the executable carries neither penalty nor additive placeholders, or the caller + // supplied no buffers, fall back to the plain path (which feeds neutral values). + if (!penaltyEnabled && !additiveEnabled) || (penaltyBuffer == nil && additiveBuffer == nil) { encode( to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset, outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) @@ -1167,12 +1505,16 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { let logitsData = MPSGraphTensorData( logitsBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) - let penaltyData = MPSGraphTensorData( - penaltyBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) + let penaltyData = penaltyBuffer.map { + MPSGraphTensorData($0, shape: [1, vocabSize as NSNumber], dataType: .float16) + } + let additiveData = additiveBuffer.map { + MPSGraphTensorData($0, shape: [1, vocabSize as NSNumber], dataType: .float16) + } let outputData = MPSGraphTensorData( outputBuffer, shape: [1 as NSNumber], dataType: .int32) - let inputs = buildInputs(logitsData: logitsData, penaltyData: penaltyData) + let inputs = buildInputs(logitsData: logitsData, penaltyData: penaltyData, additiveData: additiveData) let execDesc = MPSGraphExecutableExecutionDescriptor() execDesc.completionHandler = { [outputBuffer, outputOffset] (_, error) in @@ -1338,6 +1680,52 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { ) } + /// Encode composite sampling with slice support and an additive delta (frequency/presence + /// penalty + logit bias). Used for the first sampled token after prefill. + func encodeWithSlice( + to queue: MTLCommandQueue, + logitsBuffer: MTLBuffer, + queryLength: Int, + additiveBuffer: MTLBuffer?, + outputBuffer: MTLBuffer, + outputOffset: Int, + completion: @escaping (Int32, Error?) -> Void + ) { + guard additiveEnabled, let additiveBuffer else { + encodeWithSlice( + to: queue, logitsBuffer: logitsBuffer, queryLength: queryLength, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + return + } + if queryLength == 1 { + encode( + to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0, + penaltyBuffer: nil, additiveBuffer: additiveBuffer, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + return + } + + let logitsOffset = (queryLength - 1) * vocabSize * MemoryLayout.size + let sliceSize = vocabSize * MemoryLayout.size + guard let tempBuffer = device.makeBuffer(length: sliceSize, options: .storageModeShared), + let blitCmdBuffer = queue.makeCommandBuffer(), + let blitEncoder = blitCmdBuffer.makeBlitCommandEncoder() + else { + completion(0, MPSGraphSamplerError.bufferAllocationFailed) + return + } + blitEncoder.copy( + from: logitsBuffer, sourceOffset: logitsOffset, + to: tempBuffer, destinationOffset: 0, size: sliceSize) + blitEncoder.endEncoding() + blitCmdBuffer.commit() + + encode( + to: queue, logitsBuffer: tempBuffer, logitsOffset: 0, + penaltyBuffer: nil, additiveBuffer: additiveBuffer, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + } + // MARK: - Graph Stage Helpers /// Apply repetition penalty: where(logits > 0, logits / penalty, logits * penalty) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift b/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift index d094bd53..5d971bd2 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift @@ -91,6 +91,8 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { /// Unlike TopP, it does not require sorting — it operates as a simple threshold in logit space. public let minP: Double? + // MARK: Penalties and logit bias + /// Repetition penalty factor applied to tokens that appear in the generation history. /// /// - **nil** or **1.0**: No penalty (disabled) @@ -111,6 +113,36 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { /// - **256**: Moderate window public let repetitionPenaltyWindow: Int? + /// Additive penalty scaled by how many times a token has appeared in the recent history. + /// + /// - **nil** or **0**: No penalty (disabled) + /// - **positive**: Discourages repetition; applied as `logit -= frequencyPenalty * count` + /// - **negative**: Encourages repetition + /// + /// Valid range is [-2, 2]. Unlike `repetitionPenalty` (multiplicative, count-agnostic), + /// this is additive and scales with occurrence count. + public let frequencyPenalty: Double? + + /// Additive penalty applied once to any token that has appeared in the recent history. + /// + /// - **nil** or **0**: No penalty (disabled) + /// - **positive**: Discourages reusing tokens already seen; applied as `logit -= presencePenalty` + /// - **negative**: Encourages reuse + /// + /// Valid range is [-2, 2]. Independent of how many times the token appeared. + public let presencePenalty: Double? + + /// Fixed additive bias per token id, applied to the logits before sampling. + /// + /// - **nil**: No bias + /// - Keyed by token id; each entry adds `bias` to that token's logit (`logit += bias`) + /// + /// Typical values are in [-100, 100]; large magnitudes effectively force or ban a token. + /// This is a static, request-level vector — it does not depend on generation history. + public let logitBias: [Int32: Float]? + + // MARK: Execution options + /// A boolean flag that requests the sampling operation be combined /// with logit inference. /// @@ -148,6 +180,9 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { /// - repetitionPenaltyWindow: Optional window size. Must be > 0 if set. /// - combined: Whether to combine sampling with logit inference. Defaults to true. /// - seed: Optional seed for reproducible sampling. Defaults to nil (system generator). + /// - frequencyPenalty: Optional additive count-scaled penalty. Must be in [-2, 2] if set. + /// - presencePenalty: Optional additive presence penalty. Must be in [-2, 2] if set. + /// - logitBias: Optional per-token-id additive bias applied before sampling. public init( temperature: Double, topK: Int? = nil, @@ -156,7 +191,10 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { repetitionPenalty: Double? = nil, repetitionPenaltyWindow: Int? = nil, combined: Bool = true, - seed: UInt64? = nil + seed: UInt64? = nil, + frequencyPenalty: Double? = nil, + presencePenalty: Double? = nil, + logitBias: [Int32: Float]? = nil ) { precondition(temperature >= 0, "Temperature must be non-negative.") precondition(topK == nil || topK! > 0, "TopK must be positive if set.") @@ -168,6 +206,12 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { precondition( repetitionPenaltyWindow == nil || repetitionPenaltyWindow! > 0, "Repetition penalty window must be > 0 if set.") + precondition( + frequencyPenalty == nil || (frequencyPenalty! >= -2 && frequencyPenalty! <= 2), + "Frequency penalty must be in [-2, 2] if set.") + precondition( + presencePenalty == nil || (presencePenalty! >= -2 && presencePenalty! <= 2), + "Presence penalty must be in [-2, 2] if set.") self.temperature = temperature self.topK = topK @@ -177,6 +221,9 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { self.repetitionPenaltyWindow = repetitionPenaltyWindow self.combined = combined self.seed = seed + self.frequencyPenalty = frequencyPenalty + self.presencePenalty = presencePenalty + self.logitBias = logitBias } /// A predefined configuration for deterministic, greedy token generation. @@ -215,6 +262,22 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { return penalty > 1.0 } + /// Whether an additive frequency or presence penalty is active. + public var needsAdditivePenalty: Bool { + (frequencyPenalty ?? 0) != 0 || (presencePenalty ?? 0) != 0 + } + + /// Whether a per-token logit bias is active. + public var needsLogitBias: Bool { + guard let bias = logitBias else { return false } + return !bias.isEmpty + } + + /// Whether any logit modification beyond the multiplicative repetition penalty is active. + public var needsAdditiveLogitProcessing: Bool { + needsAdditivePenalty || needsLogitBias + } + /// Validates the configuration and returns warnings for potentially suboptimal settings. /// /// This method checks for: @@ -287,8 +350,10 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { /// /// - Returns: A new configuration with redundant settings removed. public func normalized() -> SamplingConfiguration { - // topK=1 with temperature>0 is equivalent to greedy — promote it - if let k = topK, k == 1, temperature > 0 { + // topK=1 with temperature>0 is equivalent to greedy — promote it, unless additive + // logit processing (frequency/presence penalty or logit bias) is configured, since the + // bare .greedy configuration would drop those fields. + if let k = topK, k == 1, temperature > 0, !needsAdditiveLogitProcessing { CLILogger.log( "⚠️ SamplingConfiguration: topK=1 normalized to greedy (temperature=0)", component: "Sampling") @@ -319,7 +384,10 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { repetitionPenalty: repetitionPenalty, repetitionPenaltyWindow: repetitionPenaltyWindow, combined: combined, - seed: seed + seed: seed, + frequencyPenalty: frequencyPenalty, + presencePenalty: presencePenalty, + logitBias: logitBias ) } } @@ -341,6 +409,11 @@ extension SamplingConfiguration { !needsRepetitionPenalty, "Use fallbackSampler(from:tokenHistory:) when repetition penalty is configured" ) + // Frequency/presence penalties require token history; the logit bias does not, so it is + // still honored on this history-free path. + if needsLogitBias { + AdditivePenaltyProcessor.applyLogitBias(to: &logits, logitBias: logitBias!) + } return sampleToken(from: &logits, step: step) } @@ -375,6 +448,17 @@ extension SamplingConfiguration { penalty: Float(repetitionPenalty!) ) } + if needsAdditiveLogitProcessing { + let window = repetitionPenaltyWindow.map { min($0, tokenHistory.count) } ?? tokenHistory.count + let recentTokens = tokenHistory.suffix(window) + AdditivePenaltyProcessor.apply( + to: &logits, + recentTokenIds: recentTokens, + frequencyPenalty: Float(frequencyPenalty ?? 0), + presencePenalty: Float(presencePenalty ?? 0), + logitBias: logitBias + ) + } return sampleToken(from: &logits, step: step) } diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index 42aeba95..a98a6271 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -191,13 +191,7 @@ func runChatCompletion(chatRequest: ChatCompletionRequest, state: ServerState, s let requestID = RequestID.next() let created = Int(Date().timeIntervalSince1970) - let samplingConfig = state.makeSamplingConfig( - temperature: chatRequest.temperature, - topP: chatRequest.topP, - topK: chatRequest.topK, - minP: nil, - seed: chatRequest.seed - ) + let samplingConfig = try makeSamplingConfig(for: chatRequest, state: state) let reasoningEffort = ReasoningEffort.resolve( request: chatRequest.reasoningEffort, default: state.config.defaultReasoningEffort) @@ -406,13 +400,7 @@ func prepareStreaming(chatRequest: ChatCompletionRequest, state: ServerState, se } let requestID = RequestID.next() - let samplingConfig = state.makeSamplingConfig( - temperature: chatRequest.temperature, - topP: chatRequest.topP, - topK: chatRequest.topK, - minP: nil, - seed: chatRequest.seed - ) + let samplingConfig = try makeSamplingConfig(for: chatRequest, state: state) let reasoningEffort = ReasoningEffort.resolve( request: chatRequest.reasoningEffort, default: state.config.defaultReasoningEffort) @@ -743,6 +731,61 @@ private func buildStopSequences(from request: ChatCompletionRequest, state: Serv ) } +// MARK: - Sampling options + +/// Builds the sampling configuration from a chat request, validating the additive-penalty and +/// logit-bias parameters and mapping logit-bias string keys to token ids. +private func makeSamplingConfig( + for chatRequest: ChatCompletionRequest, state: ServerState +) throws -> SamplingConfiguration { + try validatePenalty(chatRequest.frequencyPenalty, name: "frequency_penalty") + try validatePenalty(chatRequest.presencePenalty, name: "presence_penalty") + let logitBias = try parseLogitBias(chatRequest.logitBias, vocabSize: state.config.vocabSize) + return state.makeSamplingConfig( + temperature: chatRequest.temperature, + topP: chatRequest.topP, + topK: chatRequest.topK, + minP: nil, + seed: chatRequest.seed, + frequencyPenalty: chatRequest.frequencyPenalty, + presencePenalty: chatRequest.presencePenalty, + logitBias: logitBias + ) +} + +/// Validates that a frequency/presence penalty is within the accepted [-2, 2] range. +private func validatePenalty(_ value: Double?, name: String) throws { + guard let value else { return } + guard value >= -2 && value <= 2 else { + throw ServerError.badRequest("\(name) must be in [-2, 2], got \(value)") + } +} + +/// Parses the request's logit-bias object (token-id string → bias) into `[Int32: Float]`, +/// validating the token ids against the vocabulary and the bias values against [-100, 100]. +private func parseLogitBias(_ raw: [String: Double]?, vocabSize: Int?) throws -> [Int32: Float]? { + guard let raw, !raw.isEmpty else { return nil } + var result = [Int32: Float](minimumCapacity: raw.count) + for (key, value) in raw { + guard let tokenID = Int32(key) else { + throw ServerError.badRequest("logit_bias keys must be integer token ids, got '\(key)'") + } + guard tokenID >= 0 else { + throw ServerError.badRequest("logit_bias token id must be non-negative, got \(tokenID)") + } + if let vocabSize, Int(tokenID) >= vocabSize { + throw ServerError.badRequest( + "logit_bias token id \(tokenID) out of range (vocab size \(vocabSize))") + } + guard value >= -100 && value <= 100 else { + throw ServerError.badRequest( + "logit_bias values must be in [-100, 100], got \(value) for token \(tokenID)") + } + result[tokenID] = Float(value) + } + return result +} + // Per-request diagnostics go to stderr so a `--replay` stdout redirect captures clean JSONL. private func logSummary(_ line: String) { FileHandle.standardError.write(Data((line + "\n").utf8)) diff --git a/swift/Sources/Tools/llm-server/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index 0283c90b..30d588cc 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -286,18 +286,35 @@ final class ServerState: @unchecked Sendable { topP: Double?, topK: Int?, minP: Double?, - seed: UInt64? = nil + seed: UInt64? = nil, + frequencyPenalty: Double? = nil, + presencePenalty: Double? = nil, + logitBias: [Int32: Float]? = nil ) -> SamplingConfiguration { let temp = temperature ?? config.defaultTemperature + let hasAdditiveProcessing = + (frequencyPenalty ?? 0) != 0 || (presencePenalty ?? 0) != 0 || !(logitBias?.isEmpty ?? true) if temp == 0 { - return .greedy + // Greedy: additive penalties and the logit bias still shift the argmax, so build a + // full temperature-0 config carrying them rather than dropping to bare `.greedy`. + guard hasAdditiveProcessing else { return .greedy } + return SamplingConfiguration( + temperature: 0, + seed: seed, + frequencyPenalty: frequencyPenalty, + presencePenalty: presencePenalty, + logitBias: logitBias + ) } return SamplingConfiguration( temperature: temp, topK: topK ?? config.defaultTopK, topP: topP ?? config.defaultTopP, minP: minP ?? config.defaultMinP, - seed: seed + seed: seed, + frequencyPenalty: frequencyPenalty, + presencePenalty: presencePenalty, + logitBias: logitBias ) } } diff --git a/swift/Tests/CoreAILMCommonTests/ServerAPITypesTests.swift b/swift/Tests/CoreAILMCommonTests/ServerAPITypesTests.swift index ef69d084..60230a6d 100644 --- a/swift/Tests/CoreAILMCommonTests/ServerAPITypesTests.swift +++ b/swift/Tests/CoreAILMCommonTests/ServerAPITypesTests.swift @@ -53,6 +53,37 @@ struct ServerAPITypesTests { #expect(request.responseFormat?.type == "json_schema") } + @Test("frequency_penalty and presence_penalty decode") + func penaltiesDecoding() throws { + let json = """ + {"messages":[{"role":"user","content":"Hi"}],"frequency_penalty":1.5,"presence_penalty":-0.5} + """.data(using: .utf8)! + let request = try JSONDecoder().decode(ChatCompletionRequest.self, from: json) + #expect(request.frequencyPenalty == 1.5) + #expect(request.presencePenalty == -0.5) + } + + @Test("logit_bias object decodes to token-id keyed map") + func logitBiasDecoding() throws { + let json = """ + {"messages":[{"role":"user","content":"Hi"}],"logit_bias":{"50256":-100,"1917":12.5}} + """.data(using: .utf8)! + let request = try JSONDecoder().decode(ChatCompletionRequest.self, from: json) + #expect(request.logitBias?["50256"] == -100) + #expect(request.logitBias?["1917"] == 12.5) + } + + @Test("Absent penalty and logit_bias fields default to nil") + func penaltyFieldsDefaultNil() throws { + let json = """ + {"messages":[{"role":"user","content":"Hi"}]} + """.data(using: .utf8)! + let request = try JSONDecoder().decode(ChatCompletionRequest.self, from: json) + #expect(request.frequencyPenalty == nil) + #expect(request.presencePenalty == nil) + #expect(request.logitBias == nil) + } + @Test("Multimodal message content with image_url") func multimodalContent() throws { let json = """ diff --git a/swift/Tests/LanguageModelsTests/AdditivePenaltyGPUStateTests.swift b/swift/Tests/LanguageModelsTests/AdditivePenaltyGPUStateTests.swift new file mode 100644 index 00000000..751615c1 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/AdditivePenaltyGPUStateTests.swift @@ -0,0 +1,83 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Metal +import Testing + +@testable import CoreAILanguageModels + +/// Isolation tests for the GPU additive-delta state: verify the buffer contents it produces match +/// the CPU `AdditivePenaltyProcessor` math, independent of the pipelined engine's threading. +@Suite("AdditivePenaltyGPUState") +struct AdditivePenaltyGPUStateTests { + private func value(_ buf: MTLBuffer, _ idx: Int) -> Float { + Float(buf.contents().assumingMemoryBound(to: Float16.self)[idx]) + } + + @Test("logit bias is the static baseline") + func logitBiasBaseline() throws { + guard let device = MTLCreateSystemDefaultDevice() else { return } + let state = try AdditivePenaltyGPUState( + device: device, vocabSize: 8, pipelineDepth: 1, + frequencyPenalty: 0, presencePenalty: 0, logitBias: [2: 5.0, 5: -3.0], windowSize: 256) + let buf = state.buffer(forStep: 0) + #expect(abs(value(buf, 2) - 5.0) < 1e-2) + #expect(abs(value(buf, 5) - (-3.0)) < 1e-2) + #expect(value(buf, 0) == 0) + } + + @Test("frequency + presence accumulate per occurrence (pipelineDepth 1)") + func accumulateDepth1() throws { + guard let device = MTLCreateSystemDefaultDevice() else { return } + let state = try AdditivePenaltyGPUState( + device: device, vocabSize: 8, pipelineDepth: 1, + frequencyPenalty: 1.0, presencePenalty: 0.5, logitBias: nil, windowSize: 256) + _ = state.buffer(forStep: 0) // slot 0, no history yet + state.recordToken(3) + let b1 = state.buffer(forStep: 1) + // First occurrence: -freq*1 - presence = -1.5 + #expect(abs(value(b1, 3) - (-1.5)) < 1e-2) + state.recordToken(3) + let b2 = state.buffer(forStep: 2) + // Second occurrence: additional -freq = -2.5 total + #expect(abs(value(b2, 3) - (-2.5)) < 1e-2) + #expect(value(b2, 0) == 0) + } + + @Test("all pipeline slots converge to the same delta") + func slotsConverge() throws { + guard let device = MTLCreateSystemDefaultDevice() else { return } + let depth = 3 + let state = try AdditivePenaltyGPUState( + device: device, vocabSize: 8, pipelineDepth: depth, + frequencyPenalty: 2.0, presencePenalty: 0, logitBias: nil, windowSize: 256) + // Record 4 occurrences of token 4, cycling buffer() across slots as the engine would. + for step in 0..<8 { + _ = state.buffer(forStep: step) + if step < 4 { state.recordToken(4) } + } + // After all 4 occurrences have propagated, every slot must read -freq*4 = -8.0. + for slot in 0.. evict oldest (1), then add 1: window [2,1] + let b = state.buffer(forStep: 1) + // token 1: evicted once (+1), re-added (-1 -1 presence since it had left)... net over events. + // Simpler invariant: token 2 present once -> -freq-presence = -2.0 + #expect(abs(value(b, 2) - (-2.0)) < 1e-2) + } +} diff --git a/swift/Tests/LanguageModelsTests/AdditivePenaltyProcessorTests.swift b/swift/Tests/LanguageModelsTests/AdditivePenaltyProcessorTests.swift new file mode 100644 index 00000000..e14a07ab --- /dev/null +++ b/swift/Tests/LanguageModelsTests/AdditivePenaltyProcessorTests.swift @@ -0,0 +1,95 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import CoreAIShared +import Testing + +@testable import CoreAILanguageModels + +@Suite("AdditivePenaltyProcessor") +struct AdditivePenaltyProcessorTests { + @Test("Frequency penalty subtracts count times penalty") + func frequencyScalesWithCount() { + var logits: [LogitsScalarType] = [LogitsScalarType(5.0), LogitsScalarType(5.0), 0] + // Token 0 appears 3 times, token 1 appears once, penalty 0.5. + AdditivePenaltyProcessor.apply( + to: &logits, recentTokenIds: [0, 0, 0, 1] as [Int32], + frequencyPenalty: 0.5, presencePenalty: 0, logitBias: nil) + #expect(abs(Float(logits[0]) - (5.0 - 1.5)) < 1e-2) // 5 - 0.5*3 + #expect(abs(Float(logits[1]) - (5.0 - 0.5)) < 1e-2) // 5 - 0.5*1 + #expect(Float(logits[2]) == 0) + } + + @Test("Presence penalty subtracts once regardless of count") + func presenceAppliedOnce() { + var logits: [LogitsScalarType] = [LogitsScalarType(5.0), LogitsScalarType(5.0), 0] + AdditivePenaltyProcessor.apply( + to: &logits, recentTokenIds: [0, 0, 0, 1] as [Int32], + frequencyPenalty: 0, presencePenalty: 1.0, logitBias: nil) + #expect(abs(Float(logits[0]) - 4.0) < 1e-2) // 5 - 1.0 (once, despite 3 occurrences) + #expect(abs(Float(logits[1]) - 4.0) < 1e-2) // 5 - 1.0 + #expect(Float(logits[2]) == 0) // not seen + } + + @Test("Frequency and presence combine additively") + func frequencyAndPresenceCombine() { + var logits: [LogitsScalarType] = [LogitsScalarType(10.0), 0] + AdditivePenaltyProcessor.apply( + to: &logits, recentTokenIds: [0, 0] as [Int32], + frequencyPenalty: 0.5, presencePenalty: 1.0, logitBias: nil) + // 10 - (0.5*2 + 1.0) = 8.0 + #expect(abs(Float(logits[0]) - 8.0) < 1e-2) + } + + @Test("Logit bias adds to specified tokens") + func logitBiasAdds() { + var logits: [LogitsScalarType] = [LogitsScalarType(1.0), LogitsScalarType(2.0), LogitsScalarType(3.0)] + AdditivePenaltyProcessor.apply( + to: &logits, recentTokenIds: [] as [Int32], + frequencyPenalty: 0, presencePenalty: 0, logitBias: [0: 10.0, 2: -5.0]) + #expect(abs(Float(logits[0]) - 11.0) < 1e-2) + #expect(Float(logits[1]) == 2.0) + #expect(abs(Float(logits[2]) - (-2.0)) < 1e-2) + } + + @Test("Out-of-range token ids are ignored") + func outOfRangeIgnored() { + var logits: [LogitsScalarType] = [LogitsScalarType(1.0), LogitsScalarType(2.0)] + AdditivePenaltyProcessor.apply( + to: &logits, recentTokenIds: [-1, 5, 100] as [Int32], + frequencyPenalty: 1.0, presencePenalty: 1.0, logitBias: [-1: 50.0, 9: 50.0]) + #expect(Float(logits[0]) == 1.0) + #expect(Float(logits[1]) == 2.0) + } + + @Test("Zero penalties with nil bias is a no-op") + func zeroIsNoop() { + var logits: [LogitsScalarType] = [LogitsScalarType(3.0), LogitsScalarType(-1.0)] + let original = logits + AdditivePenaltyProcessor.apply( + to: &logits, recentTokenIds: [0, 1] as [Int32], + frequencyPenalty: 0, presencePenalty: 0, logitBias: nil) + #expect(logits == original) + } + + @Test("Greedy fallbackSampler honors logit bias to ban a token") + func greedyBiasBansToken() { + // Token 0 has the highest logit; a large negative bias should push token 1 to win. + let config = SamplingConfiguration( + temperature: 0, logitBias: [0: -100.0]) + var logits: [LogitsScalarType] = [LogitsScalarType(5.0), LogitsScalarType(4.0), LogitsScalarType(1.0)] + let token = config.fallbackSampler(from: &logits, tokenHistory: [] as [Int32]) + #expect(token == 1) + } + + @Test("Greedy fallbackSampler honors frequency penalty via history") + func greedyFrequencyShiftsArgmax() { + // Token 0 leads by 1.0 but appeared twice; frequency 1.0 * 2 = 2.0 penalty flips it. + let config = SamplingConfiguration(temperature: 0, frequencyPenalty: 1.0) + var logits: [LogitsScalarType] = [LogitsScalarType(5.0), LogitsScalarType(4.0), LogitsScalarType(1.0)] + let token = config.fallbackSampler(from: &logits, tokenHistory: [0, 0] as [Int32]) + #expect(token == 1) + } +}