From 9d679af48400d243f77e423a086ebad4d578cdc8 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Sat, 26 Sep 2026 18:21:28 -0700 Subject: [PATCH 1/3] Add frequency_penalty, presence_penalty, and logit_bias sampling options Adds the three chat-completions sampling controls to /v1/chat/completions and wires them through every sampling path. - ServerAPITypes: parse frequency_penalty, presence_penalty, logit_bias (token-id-string -> bias map); validated in ChatHandler (penalties in [-2, 2], bias in [-100, 100], token ids bounds-checked; invalid input returns 400). - SamplingConfiguration + AdditivePenaltyProcessor: additive, count-scaled frequency penalty, once-per-token presence penalty, and per-token logit bias, applied on the CPU sampler path (sequential/static/VLM engines). Distinct from the existing multiplicative repetition penalty. - Pipelined GPU engine: a fused additive-delta stage in both the composite and argmax MPSGraph samplers (AdditivePenaltyGPUState), so the params take effect on the default engine. Greedy + penalties/bias is supported. The additive delta carries the same N-1 pipeline-delay handling as the repetition penalty. - Guided generation: penalties/bias applied before the grammar mask so token bans stay authoritative. - Fixes a latent cross-request issue: the pipelined engine's cached sampler is now invalidated when the penalty/bias configuration changes across requests that share a prefix (PenaltySignature). Unit tests cover the CPU processor, the GPU delta state, and request decoding. --- .../CoreAILMCommon/ServerAPITypes.swift | 15 + .../ConstrainedDecodingStrategy.swift | 13 + .../ConstrainedGenerator.swift | 13 + .../CoreAIPipelinedEngine.swift | 115 ++++- .../Samplers/AdditivePenaltyGPUState.swift | 166 +++++++ .../Samplers/AdditivePenaltyProcessor.swift | 73 +++ .../Samplers/MPSGraphSamplers.swift | 434 +++++++++++++++++- .../Samplers/SamplingConfiguration.swift | 92 +++- .../Tools/llm-server/ChatHandler.swift | 71 ++- .../Tools/llm-server/ServerState.swift | 23 +- .../ServerAPITypesTests.swift | 31 ++ .../AdditivePenaltyGPUStateTests.swift | 83 ++++ .../AdditivePenaltyProcessorTests.swift | 95 ++++ 13 files changed, 1175 insertions(+), 49 deletions(-) create mode 100644 swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift create mode 100644 swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift create mode 100644 swift/Tests/LanguageModelsTests/AdditivePenaltyGPUStateTests.swift create mode 100644 swift/Tests/LanguageModelsTests/AdditivePenaltyProcessorTests.swift 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..6cbe1f07 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,7 @@ private struct EngineImpl: ~Copyable { self.logits = logitsRef self.cachedSampler = nil self.cachedSamplerTemperature = nil + self.cachedSamplerPenaltySignature = nil CLILogger.log("CoreAI pipelined engine initialized — Vocab: \(config.vocabSize)") } @@ -880,7 +883,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 @@ -901,6 +913,23 @@ private struct EngineImpl: ~Copyable { } } + // 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, + windowSize: config.repetitionPenaltyWindow + ) + } + } + let newSampler = try MPSGraphSamplerFactory.makeSampler( device: device, vocabSize: self.config.vocabSize, @@ -908,9 +937,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 +1110,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 +1131,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 +1172,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 +1861,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..2433617a --- /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, + presencePenalty: Double, + logitBias: [Int32: Float]?, + windowSize: Int? + ) 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..0de9e4f9 --- /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 || presencePenalty != 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..089fc588 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,38 @@ 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 + } + let ptr = neutralBuf.contents().assumingMemoryBound(to: Float16.self) + for i in 0.. [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 +436,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 +470,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 +499,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 +526,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 +541,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 +666,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 +814,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 +898,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 +925,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 +954,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 +980,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 +992,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 +1031,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 +1086,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 +1131,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 +1186,23 @@ 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 + } + let ptr = neutralBuf.contents().assumingMemoryBound(to: Float16.self) + for i in 0.. [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 +1277,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 +1318,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 +1326,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 +1340,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 +1372,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 +1405,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 +1421,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 +1481,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 +1509,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 +1684,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 200a034e..b7005c3b 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) } @@ -369,6 +442,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) + } +} From add27a9a0d555fc4fd87d93f0001b6ca38e05d90 Mon Sep 17 00:00:00 2001 From: Sukru Date: Thu, 8 Oct 2026 18:56:35 -0700 Subject: [PATCH 2/3] Apply suggestion from @kevchengcodes Co-authored-by: kevchengcodes <59463423+kevchengcodes@users.noreply.github.com> --- .../Samplers/AdditivePenaltyProcessor.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift index 0de9e4f9..dcc3220d 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyProcessor.swift @@ -35,7 +35,7 @@ public struct AdditivePenaltyProcessor { ) { let vocabSize = logits.count - if frequencyPenalty != 0 || presencePenalty != 0 { + 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 } From 8dd9f19d6f98fd1869968960bb5c96528d62050a Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 8 Oct 2026 19:11:22 -0700 Subject: [PATCH 3/3] Address review: unify penalty window across engines, cleanups - Fix cross-engine divergence (carinapeng): a nil repetitionPenaltyWindow means "all history" per the SamplingConfiguration contract. The CPU path honored it, but the GPU state silently capped at 256, so the same request could diverge by engine. Size the GPU ring to maxContextLength when the window is nil, for both the additive (frequency/presence) and repetition penalties. - Reset penaltyState / additivePenaltyState in the pipelined engine init for symmetry with the other cached-sampler fields. - Default AdditivePenaltyGPUState's frequency/presence/logitBias/window params so callers can omit them. - Replace the neutral additive buffer's per-element Float16 fill with a single memset (0.0 is all-zero bits). --- .../InferenceEngines/CoreAIPipelinedEngine.swift | 12 ++++++++++-- .../Samplers/AdditivePenaltyGPUState.swift | 8 ++++---- .../Samplers/MPSGraphSamplers.swift | 12 ++++-------- 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 6cbe1f07..5005e194 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -859,6 +859,8 @@ private struct EngineImpl: ~Copyable { self.cachedSampler = nil self.cachedSamplerTemperature = nil self.cachedSamplerPenaltySignature = nil + self.penaltyState = nil + self.additivePenaltyState = nil CLILogger.log("CoreAI pipelined engine initialized — Vocab: \(config.vocabSize)") } @@ -908,7 +910,10 @@ 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 ) } } @@ -925,7 +930,10 @@ private struct EngineImpl: ~Copyable { frequencyPenalty: config.frequencyPenalty ?? 0, presencePenalty: config.presencePenalty ?? 0, logitBias: config.logitBias, - 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 ) } } diff --git a/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift index 2433617a..4beb1f4c 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/AdditivePenaltyGPUState.swift @@ -46,10 +46,10 @@ final class AdditivePenaltyGPUState: @unchecked Sendable { device: MTLDevice, vocabSize: Int, pipelineDepth: Int, - frequencyPenalty: Double, - presencePenalty: Double, - logitBias: [Int32: Float]?, - windowSize: Int? + frequencyPenalty: Double = 0, + presencePenalty: Double = 0, + logitBias: [Int32: Float]? = nil, + windowSize: Int? = nil ) throws { self.vocabSize = vocabSize self.pipelineDepth = pipelineDepth diff --git a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift index 089fc588..b09c9bfc 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift @@ -369,10 +369,8 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable { guard let neutralBuf = device.makeBuffer(length: neutralByteCount, options: .storageModeShared) else { throw MPSGraphSamplerError.bufferAllocationFailed } - let ptr = neutralBuf.contents().assumingMemoryBound(to: Float16.self) - for i in 0..