From bd55261705f5224967ab5cb02c91c61fb35a2861 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Fri, 14 Aug 2026 21:27:34 -0700 Subject: [PATCH 1/7] Add repetition penalty support for CPU-based engines Penalizes tokens that appear in recent generation history, discouraging repetitive output. Applied as a separate logit modification step before the existing sampling pipeline (temperature/topK/topP/minP). - Add repetitionPenalty and repetitionPenaltyWindow to SamplingConfiguration - Add RepetitionPenaltyProcessor (deduplicates, sign-aware divide/multiply) - Integrate into Sequential, StaticShape, VLM, and Constrained engines - Only penalize generated tokens (not prompt) via generationStartOffset - Pipelined engine: hard fail with clear error (GPU path in follow-up) - CLI: --repetition-penalty and --repetition-penalty-window flags --- .../ConstrainedDecodingStrategy.swift | 11 +++ .../ConstrainedGenerator.swift | 9 ++ .../CoreAIPipelinedEngine.swift | 6 ++ .../CoreAISequentialEngine.swift | 5 +- .../CoreAISequentialVLMEngine.swift | 5 +- .../CoreAIStaticShapeEngine.swift | 10 ++- .../Samplers/RepetitionPenaltyProcessor.swift | 46 ++++++++++ .../Samplers/SamplingConfiguration.swift | 84 +++++++++++++++++-- .../Tools/llm-runner/LLMRunnerMain.swift | 12 +++ .../RepetitionPenaltyProcessorTests.swift | 79 +++++++++++++++++ 10 files changed, 253 insertions(+), 14 deletions(-) create mode 100644 swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyProcessor.swift create mode 100644 swift/Tests/LanguageModelsTests/RepetitionPenaltyProcessorTests.swift diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift index 0d23f246..e04dde00 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift @@ -116,6 +116,7 @@ public struct ConstrainedDecodingStrategy: DecodingStrategy { /// Returns `(nil, nil)` if generation should stop. fileprivate static func generateOneToken( inputTokens: [Int32], + generatedTokens: [Int32], session: inout ConstrainedGenerationSession, inferenceEngine: any InferenceEngine, samplingConfiguration: SamplingConfiguration, @@ -135,6 +136,15 @@ public struct ConstrainedDecodingStrategy: DecodingStrategy { } var maskedLogits = logits + if samplingConfiguration.needsRepetitionPenalty { + let window = samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } + ?? generatedTokens.count + RepetitionPenaltyProcessor.apply( + to: &maskedLogits, + recentTokenIds: generatedTokens.suffix(window), + penalty: Float(samplingConfiguration.repetitionPenalty!) + ) + } _ = session.applyMask(to: &maskedLogits) let bestToken = CompositeSampler.sample(from: &maskedLogits, config: samplingConfiguration) @@ -296,6 +306,7 @@ extension ConstrainedDecodingStrategy.ConstrainedDecodedSequence { do { result = try await ConstrainedDecodingStrategy.generateOneToken( inputTokens: inputTokens, + generatedTokens: generatedTokens, session: &session, inferenceEngine: inferenceEngine, samplingConfiguration: samplingConfiguration, diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift index fc014a0e..35de12b2 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift @@ -204,6 +204,15 @@ public struct ConstrainedGenerator: DecodingStrategy { } var maskedLogits = logits + if samplingConfiguration.needsRepetitionPenalty { + let window = samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } + ?? generatedTokens.count + RepetitionPenaltyProcessor.apply( + to: &maskedLogits, + recentTokenIds: generatedTokens.suffix(window), + penalty: Float(samplingConfiguration.repetitionPenalty!) + ) + } _ = session.applyMask(to: &maskedLogits) let bestToken = CompositeSampler.sample(from: &maskedLogits, config: samplingConfiguration) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 000a4416..22ab3827 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -127,6 +127,12 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable + "Use a sequential engine for evaluation." ) } + if samplingConfiguration.needsRepetitionPenalty { + throw InferenceRuntimeError.invalidArgument( + "CoreAI pipelined engine does not yet support repetition penalty (GPU-side sampling). " + + "Use a sequential engine (--inference-engine-variant coreai-sequential)." + ) + } // Serialize: if a prior generation is still winding down (GPU drain), // cancel it and wait for the engine slot to be released. diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index ceb7167c..eed77bed 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -480,6 +480,7 @@ extension CoreAISequentialEngine.GenerationSequence { private let generationToken: GenerationToken private var inputTokens: [CoreAISequentialEngine.TokenId] + private let generationStartOffset: Int private var step: Int = 0 private var finished: Bool = false @@ -498,6 +499,7 @@ extension CoreAISequentialEngine.GenerationSequence { self.stopReasonStore = stopReasonStore self.generationToken = generationToken self.inputTokens = input + self.generationStartOffset = input.count if let forced = inferenceOptions.forcedContinuation { self.maxTokens = forced.count } else { @@ -570,7 +572,8 @@ extension CoreAISequentialEngine.GenerationSequence { nextToken = forced[step] } else { var mutableLogits = logitBuffer - nextToken = samplingConfiguration.fallbackSampler(from: &mutableLogits) + nextToken = samplingConfiguration.fallbackSampler( + from: &mutableLogits, tokenHistory: inputTokens[generationStartOffset...]) } inputTokens.append(nextToken) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift index 1aa411ec..62c9c45f 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift @@ -1071,6 +1071,7 @@ extension CoreAISequentialVLMEngine.GenerationSequence { private let generationToken: GenerationToken private var inputTokens: [CoreAISequentialVLMEngine.TokenId] + private let generationStartOffset: Int private var embeddedInput: InputEmbeddings? private var step: Int = 0 private var finished: Bool = false @@ -1092,6 +1093,7 @@ extension CoreAISequentialVLMEngine.GenerationSequence { self.stopReasonStore = stopReasonStore self.generationToken = generationToken self.inputTokens = input + self.generationStartOffset = input.count self.embeddedInput = embeddedInput if let forced = inferenceOptions.forcedContinuation { self.maxTokens = forced.count @@ -1180,7 +1182,8 @@ extension CoreAISequentialVLMEngine.GenerationSequence { nextToken = forced[step] } else { var mutableLogits = logitBuffer - nextToken = samplingConfiguration.fallbackSampler(from: &mutableLogits) + nextToken = samplingConfiguration.fallbackSampler( + from: &mutableLogits, tokenHistory: inputTokens[generationStartOffset...]) } inputTokens.append(nextToken) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift index e0ab4594..2a62fe3a 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift @@ -368,7 +368,8 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable { // MARK: - Inference public func inference( - inputTokens: [Int32], samplingConfig: SamplingConfiguration, returnsLogits: Bool + inputTokens: [Int32], samplingConfig: SamplingConfiguration, returnsLogits: Bool, + generationStartOffset: Int = 0 ) async throws -> (logits: [LogitsScalarType]?, token: Int32) { CLILogger.log("Inference: \(inputTokens.count) tokens, processed: \(processedTokenCount)") @@ -453,7 +454,7 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable { let actualLogits = returnsLogits ? logitBuffer : nil let sampleSpan = InstrumentsProfiler.beginSample(strategy: "cpu-fallback") - let nextToken = samplingConfig.fallbackSampler(from: &logitBuffer) + let nextToken = samplingConfig.fallbackSampler(from: &logitBuffer, tokenHistory: inputTokens[generationStartOffset...]) sampleSpan.end() CLILogger.log("Token: \(nextToken), processed: \(processedTokenCount)") return (logits: actualLogits, token: nextToken) @@ -657,6 +658,7 @@ extension StaticShapeEngine.GenerationSequence { private let generationToken: GenerationToken private var inputTokens: [StaticShapeEngine.TokenId] + private let generationStartOffset: Int private var step: Int = 0 private var finished: Bool = false @@ -675,6 +677,7 @@ extension StaticShapeEngine.GenerationSequence { self.stopReasonStore = stopReasonStore self.generationToken = generationToken self.inputTokens = input + self.generationStartOffset = input.count if let forced = inferenceOptions.forcedContinuation { self.maxTokens = forced.count } else { @@ -711,7 +714,8 @@ extension StaticShapeEngine.GenerationSequence { let (logits, sampledToken) = try await engine.inference( inputTokens: inputTokens, samplingConfig: samplingConfiguration, - returnsLogits: returnsLogits || forcedContinuation != nil + returnsLogits: returnsLogits || forcedContinuation != nil, + generationStartOffset: generationStartOffset ) // Update history with newly processed tokens diff --git a/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyProcessor.swift b/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyProcessor.swift new file mode 100644 index 00000000..c7fc784a --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyProcessor.swift @@ -0,0 +1,46 @@ +// 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 repetition penalty to logits based on token generation history. +/// +/// For each unique token ID in the recent history: +/// - If logit > 0: divide by penalty factor +/// - If logit < 0: multiply by penalty factor +/// +/// This discourages the model from re-emitting recently generated tokens. +public struct RepetitionPenaltyProcessor { + /// Apply repetition penalty to logits in-place. + /// + /// - Parameters: + /// - logits: Mutable logits array (vocab-sized). Modified in-place. + /// - recentTokenIds: Token IDs from recent generation history. + /// - penalty: The penalty factor (> 1.0 penalizes, 1.0 = no-op). + public static func apply>( + to logits: inout [LogitsScalarType], + recentTokenIds: C, + penalty: Float + ) { + guard penalty > 1.0 else { return } + guard !recentTokenIds.isEmpty else { return } + + let vocabSize = logits.count + var seen = Set(minimumCapacity: min(recentTokenIds.count, 512)) + + for tokenId in recentTokenIds { + guard tokenId >= 0 && Int(tokenId) < vocabSize else { continue } + guard seen.insert(tokenId).inserted else { continue } + + let idx = Int(tokenId) + let logit = Float(logits[idx]) + if logit > 0 { + logits[idx] = LogitsScalarType(logit / penalty) + } else if logit < 0 { + logits[idx] = LogitsScalarType(logit * penalty) + } + } + } +} diff --git a/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift b/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift index 15352be5..f6b19746 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift @@ -15,11 +15,12 @@ import CoreAIShared /// /// ## Sampling Algorithm Order /// When multiple parameters are set, they are applied in this order: -/// 1. Temperature scaling (logits / temperature) -/// 2. MinP filtering (relative probability threshold) -/// 3. TopP filtering (cumulative probability cutoff) -/// 4. TopK filtering (hard limit on vocabulary) -/// 5. Softmax and multinomial sampling +/// 1. Repetition penalty (logits modified based on token history) +/// 2. Temperature scaling (logits / temperature) +/// 3. MinP filtering (relative probability threshold) +/// 4. TopP filtering (cumulative probability cutoff) +/// 5. TopK filtering (hard limit on vocabulary) +/// 6. Softmax and multinomial sampling /// /// ## Usage Example /// ```swift @@ -90,6 +91,26 @@ 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? + /// Repetition penalty factor applied to tokens that appear in the generation history. + /// + /// - **nil** or **1.0**: No penalty (disabled) + /// - **1.1–1.3**: Common range for reducing repetition + /// - **>1.5**: Aggressive penalty, may hurt coherence + /// + /// For each token in recent history: + /// - If logit > 0: divide by penalty + /// - If logit < 0: multiply by penalty + /// + /// Applied before all other sampling steps (temperature, topK, topP, minP). + public let repetitionPenalty: Double? + + /// How many recent tokens to consider for repetition penalty. + /// + /// - **nil**: All tokens in generation history + /// - **64**: Only penalize tokens from the last 64 steps + /// - **256**: Moderate window + public let repetitionPenaltyWindow: Int? + /// A boolean flag that requests the sampling operation be combined /// with logit inference. /// @@ -107,20 +128,35 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { /// - topK: Optional top-K limit. Must be > 0 if set. /// - topP: Optional top-P threshold. Must be in (0, 1] if set. /// - minP: Optional min-P threshold. Must be in (0, 1] if set. + /// - repetitionPenalty: Optional repetition penalty factor. Must be >= 1.0 if set. + /// - repetitionPenaltyWindow: Optional window size. Must be > 0 if set. /// - combined: Whether to combine sampling with logit inference. Defaults to true. - /// - /// - Note: Call `validate()` to check for potentially suboptimal configurations. - public init(temperature: Double, topK: Int? = nil, topP: Double? = nil, minP: Double? = nil, combined: Bool = true) - { + public init( + temperature: Double, + topK: Int? = nil, + topP: Double? = nil, + minP: Double? = nil, + repetitionPenalty: Double? = nil, + repetitionPenaltyWindow: Int? = nil, + combined: Bool = true + ) { precondition(temperature >= 0, "Temperature must be non-negative.") precondition(topK == nil || topK! > 0, "TopK must be positive if set.") precondition(topP == nil || (topP! > 0 && topP! <= 1), "TopP must be in (0, 1] if set.") precondition(minP == nil || (minP! > 0 && minP! <= 1), "MinP must be in (0, 1] if set.") + precondition( + repetitionPenalty == nil || repetitionPenalty! >= 1.0, + "Repetition penalty must be >= 1.0 if set.") + precondition( + repetitionPenaltyWindow == nil || repetitionPenaltyWindow! > 0, + "Repetition penalty window must be > 0 if set.") self.temperature = temperature self.topK = topK self.topP = topP self.minP = minP + self.repetitionPenalty = repetitionPenalty + self.repetitionPenaltyWindow = repetitionPenaltyWindow self.combined = combined } @@ -154,6 +190,12 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { temperature > 0 && (topK != nil || topP != nil || minP != nil) } + /// Whether repetition penalty is active. + public var needsRepetitionPenalty: Bool { + guard let penalty = repetitionPenalty else { return false } + return penalty > 1.0 + } + /// Validates the configuration and returns warnings for potentially suboptimal settings. /// /// This method checks for: @@ -255,6 +297,8 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable { topK: effectiveTopK, topP: effectiveTopP, minP: effectiveMinP, + repetitionPenalty: repetitionPenalty, + repetitionPenaltyWindow: repetitionPenaltyWindow, combined: combined ) } @@ -272,4 +316,26 @@ extension SamplingConfiguration { public func fallbackSampler(from logits: inout [LogitsScalarType]) -> Int32 { return CompositeSampler.sample(from: &logits, config: self) } + + /// Samples the next token with repetition penalty applied first. + /// + /// Applies repetition penalty (if configured) to the logits based on token history, + /// then delegates to the standard sampler pipeline. + /// + /// - Parameters: + /// - logits: Mutable array of Float16 logits. May be modified during sampling. + /// - tokenHistory: Recent token IDs for repetition penalty. + /// - Returns: The sampled token ID. + public func fallbackSampler(from logits: inout [LogitsScalarType], tokenHistory: some Collection) -> Int32 { + if needsRepetitionPenalty { + let window = repetitionPenaltyWindow.map { min($0, tokenHistory.count) } ?? tokenHistory.count + let recentTokens = tokenHistory.suffix(window) + RepetitionPenaltyProcessor.apply( + to: &logits, + recentTokenIds: recentTokens, + penalty: Float(repetitionPenalty!) + ) + } + return CompositeSampler.sample(from: &logits, config: self) + } } diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index ea45fecf..b00bd7e9 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -108,6 +108,16 @@ struct LLMRunner: AsyncParsableCommand, Sendable { help: "Min-P sampling: keep tokens with probability >= minP × max probability (e.g., 0.05)") var minP: Double? + @Option( + name: .customLong("repetition-penalty"), + help: "Repetition penalty factor (>= 1.0). Penalizes tokens that appeared in recent generation (e.g., 1.2)") + var repetitionPenalty: Double? + + @Option( + name: .customLong("repetition-penalty-window"), + help: "Number of recent tokens to consider for repetition penalty (default: all)") + var repetitionPenaltyWindow: Int? + @Option(help: "Sampling strategy. Options: 'temperature' (default), 'greedy'") var samplingStrategy: String = "temperature" @@ -850,6 +860,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable { topK: topK, topP: topP, minP: minP, + repetitionPenalty: repetitionPenalty, + repetitionPenaltyWindow: repetitionPenaltyWindow, combined: !synchronousSampling ) case "greedy": diff --git a/swift/Tests/LanguageModelsTests/RepetitionPenaltyProcessorTests.swift b/swift/Tests/LanguageModelsTests/RepetitionPenaltyProcessorTests.swift new file mode 100644 index 00000000..f8f2b58a --- /dev/null +++ b/swift/Tests/LanguageModelsTests/RepetitionPenaltyProcessorTests.swift @@ -0,0 +1,79 @@ +// 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("RepetitionPenaltyProcessor") +struct RepetitionPenaltyProcessorTests { + @Test("Positive logits are divided by penalty") + func positiveLogitsDivided() { + var logits: [LogitsScalarType] = [0, 0, LogitsScalarType(2.0), 0, 0] + RepetitionPenaltyProcessor.apply(to: &logits, recentTokenIds: [2], penalty: 2.0) + #expect(abs(Float(logits[2]) - 1.0) < 1e-3) + } + + @Test("Negative logits are multiplied by penalty") + func negativeLogitsMultiplied() { + var logits: [LogitsScalarType] = [0, 0, LogitsScalarType(-2.0), 0, 0] + RepetitionPenaltyProcessor.apply(to: &logits, recentTokenIds: [2], penalty: 2.0) + #expect(abs(Float(logits[2]) - (-4.0)) < 1e-2) + } + + @Test("Zero logits unchanged") + func zeroLogitsUnchanged() { + var logits: [LogitsScalarType] = [0, 0, 0, 0, 0] + RepetitionPenaltyProcessor.apply(to: &logits, recentTokenIds: [0, 1, 2, 3, 4], penalty: 1.5) + for l in logits { + #expect(l == 0) + } + } + + @Test("Penalty of 1.0 is a no-op") + func penaltyOneIsNoop() { + var logits: [LogitsScalarType] = [LogitsScalarType(3.0), LogitsScalarType(-1.0)] + let original = logits + RepetitionPenaltyProcessor.apply(to: &logits, recentTokenIds: [0, 1], penalty: 1.0) + #expect(logits == original) + } + + @Test("Duplicate token IDs penalized only once") + func deduplication() { + var logits1: [LogitsScalarType] = [LogitsScalarType(4.0), 0, 0] + var logits2: [LogitsScalarType] = [LogitsScalarType(4.0), 0, 0] + RepetitionPenaltyProcessor.apply(to: &logits1, recentTokenIds: [0], penalty: 2.0) + RepetitionPenaltyProcessor.apply(to: &logits2, recentTokenIds: [0, 0, 0], penalty: 2.0) + #expect(logits1[0] == logits2[0]) + } + + @Test("Out-of-range token IDs are ignored") + func outOfRangeIgnored() { + var logits: [LogitsScalarType] = [LogitsScalarType(1.0), LogitsScalarType(2.0)] + RepetitionPenaltyProcessor.apply(to: &logits, recentTokenIds: [-1, 5, 100], penalty: 2.0) + #expect(Float(logits[0]) == 1.0) + #expect(Float(logits[1]) == 2.0) + } + + @Test("fallbackSampler with tokenHistory applies penalty") + func fallbackSamplerWithHistory() { + let config = SamplingConfiguration(temperature: 0, repetitionPenalty: 2.0) + // Token 0 has highest logit but is penalized; token 1 should win + var logits: [LogitsScalarType] = [LogitsScalarType(3.0), LogitsScalarType(2.0), LogitsScalarType(1.0)] + let token = config.fallbackSampler(from: &logits, tokenHistory: [0] as [Int32]) + #expect(token == 1) + } + + @Test("Window limits which tokens are penalized") + func windowLimitsScope() { + let config = SamplingConfiguration(temperature: 0, repetitionPenalty: 2.0, repetitionPenaltyWindow: 1) + // History: [0, 1] but window=1, so only token 1 is penalized + var logits: [LogitsScalarType] = [LogitsScalarType(2.0), LogitsScalarType(3.0), LogitsScalarType(1.0)] + let token = config.fallbackSampler(from: &logits, tokenHistory: [0, 1] as [Int32]) + // Token 1 (3.0/2.0=1.5) penalized, token 0 (2.0) not penalized → token 0 wins + #expect(token == 0) + } +} From acca06573f026d4713ae9b0d0eeb458824d07d42 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Sat, 15 Aug 2026 09:18:54 -0700 Subject: [PATCH 2/7] Add GPU repetition penalty for pipelined engine Extend MPSGraphCompositeSampler with an optional penalty stage (penaltyEnabled flag at init). When active, the compiled graph applies sign-aware penalty (divide positive logits, multiply negative) before topK. Refactor the monolithic graph-building init into composable static stage helpers (applyPenaltyStage, topKStage, temperatureStage, softmaxStage, minPStage, topPStage, maskAndNormalizeStage, multinomialStage, gatherTokenStage) that can be unit-tested independently. RepetitionPenaltyGPUState manages per-pipeline-depth rotating penalty buffers with dirty-tracking: recordToken() updates only CPU-side ring state, and buffer(forStep:) applies pending writes at encode time when the gate guarantees no in-flight GPU read on that slot. Inherent 2-token staleness from pipelineDepth=3 is acceptable for practical window sizes. Greedy + penalty on pipelined is rejected at entry (use sequential). --- .../ConstrainedDecodingStrategy.swift | 3 +- .../ConstrainedGenerator.swift | 3 +- .../CoreAIPipelinedEngine.swift | 45 +++- .../CoreAIStaticShapeEngine.swift | 3 +- .../Samplers/MPSGraphSamplers.swift | 250 ++++++++++++++---- .../Samplers/RepetitionPenaltyGPUState.swift | 125 +++++++++ 6 files changed, 361 insertions(+), 68 deletions(-) create mode 100644 swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift index e04dde00..bc5d4142 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift @@ -137,7 +137,8 @@ public struct ConstrainedDecodingStrategy: DecodingStrategy { var maskedLogits = logits if samplingConfiguration.needsRepetitionPenalty { - let window = samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } + let window = + samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } ?? generatedTokens.count RepetitionPenaltyProcessor.apply( to: &maskedLogits, diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift index 35de12b2..2f64b535 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift @@ -205,7 +205,8 @@ public struct ConstrainedGenerator: DecodingStrategy { var maskedLogits = logits if samplingConfiguration.needsRepetitionPenalty { - let window = samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } + let window = + samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } ?? generatedTokens.count RepetitionPenaltyProcessor.apply( to: &maskedLogits, diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 22ab3827..b3648693 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -127,12 +127,6 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable + "Use a sequential engine for evaluation." ) } - if samplingConfiguration.needsRepetitionPenalty { - throw InferenceRuntimeError.invalidArgument( - "CoreAI pipelined engine does not yet support repetition penalty (GPU-side sampling). " - + "Use a sequential engine (--inference-engine-variant coreai-sequential)." - ) - } // Serialize: if a prior generation is still winding down (GPU drain), // cancel it and wait for the engine slot to be released. @@ -572,6 +566,7 @@ private struct EngineImpl: ~Copyable { // GPU sampler — reuses MPSGraphSampler from MPSGraphSamplers.swift var cachedSampler: (any MPSGraphSampler)? var cachedSamplerTemperature: Double? + var penaltyState: RepetitionPenaltyGPUState? // State var processedTokenCount: Int = 0 @@ -823,6 +818,24 @@ private struct EngineImpl: ~Copyable { return existingSampler } + // Create penalized sampler if repetition penalty is configured + if config.needsRepetitionPenalty { + if config.temperature == 0 { + throw InferenceRuntimeError.invalidArgument( + "Repetition penalty with greedy sampling is not supported on pipelined engine. " + + "Use temperature > 0, or use a sequential engine.") + } + if penaltyState == nil { + penaltyState = try RepetitionPenaltyGPUState( + device: device, + vocabSize: self.config.vocabSize, + pipelineDepth: pipelineDepth, + penalty: config.repetitionPenalty!, + windowSize: config.repetitionPenaltyWindow + ) + } + } + let newSampler = try MPSGraphSamplerFactory.makeSampler( device: device, vocabSize: self.config.vocabSize, @@ -983,7 +996,10 @@ private struct EngineImpl: ~Copyable { let queue = pipelineQueue let localInFlightGate = inFlightGate + let localPenaltyState = penaltyState let completionCallback: (Int32, Error?) -> Void = { nextToken, error in + // Update penalty state BEFORE releasing the gate. + localPenaltyState?.recordToken(nextToken) // Release the pipeline slot acquired before encode. Happens on // Metal's callback thread — PipelineGate.release() is thread-safe. localInFlightGate.release() @@ -1000,7 +1016,22 @@ private struct EngineImpl: ~Copyable { } do { - if queryLength == 1 { + // Use penalty-aware path for decode steps when penalty is active. + if queryLength == 1, let state = penaltyState, + let compositeSampler = localGPUSampler as? MPSGraphCompositeSampler, + compositeSampler.penaltyEnabled + { + let penaltyBuf = state.buffer(forStep: currentStep) + compositeSampler.encode( + to: queue, + logitsBuffer: samplerLogitsBuffer, + logitsOffset: logitsOffset, + penaltyBuffer: penaltyBuf, + outputBuffer: outputBuffer, + outputOffset: 0, + completion: completionCallback + ) + } else if queryLength == 1 { try localGPUSampler.encode( to: queue, logitsBuffer: samplerLogitsBuffer, diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift index 2a62fe3a..457a988f 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift @@ -454,7 +454,8 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable { let actualLogits = returnsLogits ? logitBuffer : nil let sampleSpan = InstrumentsProfiler.beginSample(strategy: "cpu-fallback") - let nextToken = samplingConfig.fallbackSampler(from: &logitBuffer, tokenHistory: inputTokens[generationStartOffset...]) + let nextToken = samplingConfig.fallbackSampler( + from: &logitBuffer, tokenHistory: inputTokens[generationStartOffset...]) sampleSpan.end() CLILogger.log("Token: \(nextToken), processed: \(processedTokenCount)") return (logits: actualLogits, token: nextToken) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift index 51b01090..d8b0baa7 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift @@ -136,7 +136,8 @@ enum MPSGraphSamplerFactory { k: effectiveK, temperature: Float(config.temperature), topP: config.topP.map { Float($0) } ?? 1.0, - minP: config.minP.map { Float($0) } ?? 0.0 + minP: config.minP.map { Float($0) } ?? 0.0, + penaltyEnabled: config.needsRepetitionPenalty ) } @@ -670,6 +671,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { // Graph tensors private let logitsPlaceholder: MPSGraphTensor + private let penaltyPlaceholder: MPSGraphTensor? private let temperaturePlaceholder: MPSGraphTensor private let randomPlaceholder: MPSGraphTensor private let topPPlaceholder: MPSGraphTensor @@ -693,6 +695,9 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { /// The minP value (0.0 = disabled) let minP: Float + /// Whether repetition penalty is compiled into this sampler's graph + let penaltyEnabled: Bool + /// Pre-allocated buffer for random value private let randomBuffer: MTLBuffer @@ -734,7 +739,10 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { /// - temperature: Sampling temperature /// - topP: Nucleus sampling threshold (1.0 = disabled) /// - 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) + init( + device: MTLDevice, vocabSize: Int, k: Int = 40, temperature: Float = 1.0, topP: Float = 1.0, minP: Float = 0.0, + penaltyEnabled: Bool = false + ) throws { self.device = device @@ -744,6 +752,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { self.temperature = temperature self.topP = topP self.minP = minP + self.penaltyEnabled = penaltyEnabled self.bitmaskSize = (vocabSize + 31) / 32 // Pre-allocate buffers @@ -771,6 +780,17 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { ) self.logitsPlaceholder = logitsPlaceholder + if penaltyEnabled { + let pp = graph.placeholder( + shape: [1, vocabSize as NSNumber], + dataType: .float16, + name: "penalty" + ) + self.penaltyPlaceholder = pp + } else { + self.penaltyPlaceholder = nil + } + // Temperature scalar [1] let temperaturePlaceholder = graph.placeholder( shape: [1 as NSNumber], @@ -806,71 +826,50 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { // Cast logits to Float32 for numerical stability let logitsFloat32 = graph.cast(logitsPlaceholder, to: .float32, name: "logits_f32") - // Step 1: Get Top-K values and indices - let topKResult = graph.topK(logitsFloat32, k: k, name: "topk") - let topKValues = topKResult[0] // [1, k] sorted descending - let topKIndices = topKResult[1] // [1, k] as Int32 - - // Step 2: Apply temperature: values / temperature - let scaledValues = graph.division(topKValues, temperaturePlaceholder, name: "scaled") - - // Step 3: Softmax over the K dimension (axis 1) - let probabilities = graph.softMax(with: scaledValues, axis: 1, name: "probs") - - // Step 4: MinP filtering - // max_prob is the first element (topK returns sorted descending) - let maxProb = graph.sliceTensor(probabilities, dimension: 1, start: 0, length: 1, name: "max_prob") - // threshold = minP * max_prob - let minPThreshold = graph.multiplication(minPPlaceholder, maxProb, name: "minp_threshold") - // mask: probs >= threshold (broadcasts [1,1] to [1,k]) - let minPMask = graph.greaterThanOrEqualTo(probabilities, minPThreshold, name: "minp_mask") - - // Step 5: TopP filtering via exclusive cumulative sum - // exclusive_cumsum[i] = sum of probs[0..i-1], so position 0 always has value 0 - let exclusiveCumsum = graph.cumulativeSum( - probabilities, axis: 1, exclusive: true, reverse: false, name: "excl_cumsum") - // mask: exclusive_cumsum < topP (includes all tokens before cumsum reaches topP) - let topPMask = graph.lessThan(exclusiveCumsum, topPPlaceholder, name: "topp_mask") - - // Step 6: Combined mask = minP AND topP - let combinedMask = graph.logicalAND(minPMask, topPMask, name: "combined_mask") - let maskFloat = graph.cast(combinedMask, to: .float32, name: "mask_float") - - // Step 7: Apply mask and re-normalize - let maskedProbs = graph.multiplication(probabilities, maskFloat, name: "masked_probs") - let sumMasked = graph.reductionSum(with: maskedProbs, axis: 1, name: "sum_masked") - // Avoid division by zero: use max(sum, epsilon) - let epsilon = graph.constant(1e-10, dataType: .float32) - let safeDenominator = graph.maximum(sumMasked, epsilon, name: "safe_denom") - let normalizedProbs = graph.division(maskedProbs, safeDenominator, name: "normalized_probs") - - // Step 8: Multinomial sampling via cumulative sum + random comparison - let cumsum = graph.cumulativeSum(normalizedProbs, axis: 1, exclusive: false, reverse: false, name: "cumsum") - let selectionMask = graph.greaterThanOrEqualTo(cumsum, randomPlaceholder, name: "selection_mask") - let selectionMaskFloat = graph.cast(selectionMask, to: .float32, name: "selection_mask_float") - let selectedIdx = graph.reductionArgMaximum(with: selectionMaskFloat, axis: 1, name: "selected_idx") - - // Step 9: Gather the token index from topKIndices - let selectedIdxInt32 = graph.cast(selectedIdx, to: .int32, name: "selected_idx_i32") - let indicesFlat = graph.reshape(topKIndices, shape: [k as NSNumber], name: "indices_flat") - let selectedIdxFlat = graph.reshape(selectedIdxInt32, shape: [1 as NSNumber], name: "selected_flat") - - let outputTensor = graph.gatherAlongAxis( - 0, - updates: indicesFlat, - indices: selectedIdxFlat, - name: "token_id" - ) + // Build sampling pipeline using composable stage helpers + let penalizedLogits: MPSGraphTensor + if penaltyEnabled { + penalizedLogits = Self.applyPenaltyStage( + graph: graph, logits: logitsFloat32, penaltyTensor: penaltyPlaceholder!, name: "penalty") + } else { + penalizedLogits = logitsFloat32 + } + + let (topKValues, topKIndices) = Self.topKStage( + graph: graph, logits: penalizedLogits, k: k, name: "topk") + + let scaledValues = Self.temperatureStage( + graph: graph, values: topKValues, temperature: temperaturePlaceholder, name: "temp") + + let probabilities = Self.softmaxStage(graph: graph, values: scaledValues, name: "sm") + + let minPMask = Self.minPStage( + graph: graph, probs: probabilities, minP: minPPlaceholder, name: "minp") + + let topPMask = Self.topPStage( + graph: graph, probs: probabilities, topP: topPPlaceholder, name: "topp") + + let normalizedProbs = Self.maskAndNormalizeStage( + graph: graph, probs: probabilities, masks: [minPMask, topPMask], name: "norm") + + let selectedIdx = Self.multinomialStage( + graph: graph, probs: normalizedProbs, random: randomPlaceholder, name: "sample") + + let outputTensor = Self.gatherTokenStage( + graph: graph, topKIndices: topKIndices, selectedIdx: selectedIdx, k: k, name: "gather") self.outputTensor = outputTensor // Compile to executable - let feeds: [MPSGraphTensor: MPSGraphShapedType] = [ + var feeds: [MPSGraphTensor: MPSGraphShapedType] = [ logitsPlaceholder: MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16), temperaturePlaceholder: MPSGraphShapedType(shape: [1 as NSNumber], dataType: .float32), randomPlaceholder: MPSGraphShapedType(shape: [1 as NSNumber], dataType: .float32), topPPlaceholder: MPSGraphShapedType(shape: [1 as NSNumber], dataType: .float32), minPPlaceholder: MPSGraphShapedType(shape: [1 as NSNumber], dataType: .float32), ] + if let pp = penaltyPlaceholder { + feeds[pp] = MPSGraphShapedType(shape: [1, vocabSize as NSNumber], dataType: .float16) + } let compilationDescriptor = MPSGraphCompilationDescriptor() compilationDescriptor.optimizationLevel = .level0 @@ -1084,6 +1083,54 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { applyBitmask: applyBitmask, completion: completion) } + /// Encode sampling with repetition penalty buffer. + /// The penalty buffer must be Float16[vocabSize] with 1.0 for unpenalized tokens. + func encode( + to queue: MTLCommandQueue, + logitsBuffer: MTLBuffer, + logitsOffset: Int, + penaltyBuffer: MTLBuffer, + outputBuffer: MTLBuffer, + outputOffset: Int, + completion: @escaping (Int32, Error?) -> Void + ) { + guard penaltyEnabled else { + encode( + to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset, + outputBuffer: outputBuffer, outputOffset: outputOffset, completion: completion) + return + } + + temperatureBuffer.contents().assumingMemoryBound(to: Float.self).pointee = max(temperature, 0.01) + topPBuffer.contents().assumingMemoryBound(to: Float.self).pointee = topP + minPBuffer.contents().assumingMemoryBound(to: Float.self).pointee = minP + let randomValue = testingOnlyRandomOverride ?? Float.random(in: 0..<1) + randomBuffer.contents().assumingMemoryBound(to: Float.self).pointee = randomValue + + let logitsData = MPSGraphTensorData( + logitsBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) + let penaltyData = MPSGraphTensorData( + penaltyBuffer, shape: [1, vocabSize as NSNumber], dataType: .float16) + let outputData = MPSGraphTensorData( + outputBuffer, shape: [1 as NSNumber], dataType: .int32) + + let execDesc = MPSGraphExecutableExecutionDescriptor() + execDesc.completionHandler = { [outputBuffer, outputOffset] (_, error) in + if let error = error { + completion(0, error) + return + } + let result = outputBuffer.contents() + .advanced(by: outputOffset) + .assumingMemoryBound(to: Int32.self).pointee + completion(result, nil) + } + executable.runAsync( + with: queue, + inputs: [logitsData, penaltyData, temperatureData, randomData, topPData, minPData], + results: [outputData], executionDescriptor: execDesc) + } + /// Encode composite sampling asynchronously (protocol conformance). func encode( to queue: MTLCommandQueue, @@ -1230,6 +1277,93 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { executionDescriptor: prefillExecDescriptor ) } + + // MARK: - Graph Stage Helpers + + /// Apply repetition penalty: where(logits > 0, logits / penalty, logits * penalty) + static func applyPenaltyStage( + graph: MPSGraph, logits: MPSGraphTensor, penaltyTensor: MPSGraphTensor, name: String + ) -> MPSGraphTensor { + let penaltyF32 = graph.cast(penaltyTensor, to: .float32, name: "\(name)_f32") + let zero = graph.constant(0.0, dataType: .float32) + let positive = graph.greaterThan(logits, zero, name: "\(name)_pos") + let divided = graph.division(logits, penaltyF32, name: "\(name)_div") + let multiplied = graph.multiplication(logits, penaltyF32, name: "\(name)_mul") + return graph.select(predicate: positive, trueTensor: divided, falseTensor: multiplied, name: name) + } + + /// Extract top-K values and indices from logits. + static func topKStage( + graph: MPSGraph, logits: MPSGraphTensor, k: Int, name: String + ) -> (values: MPSGraphTensor, indices: MPSGraphTensor) { + let result = graph.topK(logits, k: k, name: name) + return (result[0], result[1]) + } + + /// Scale values by temperature: values / temperature. + static func temperatureStage( + graph: MPSGraph, values: MPSGraphTensor, temperature: MPSGraphTensor, name: String + ) -> MPSGraphTensor { + graph.division(values, temperature, name: name) + } + + /// Softmax over the K dimension (axis 1). + static func softmaxStage(graph: MPSGraph, values: MPSGraphTensor, name: String) -> MPSGraphTensor { + graph.softMax(with: values, axis: 1, name: name) + } + + /// MinP mask: probs >= minP * max_prob. + static func minPStage( + graph: MPSGraph, probs: MPSGraphTensor, minP: MPSGraphTensor, name: String + ) -> MPSGraphTensor { + let maxProb = graph.sliceTensor(probs, dimension: 1, start: 0, length: 1, name: "\(name)_max") + let threshold = graph.multiplication(minP, maxProb, name: "\(name)_thr") + return graph.greaterThanOrEqualTo(probs, threshold, name: "\(name)_mask") + } + + /// TopP mask: exclusive_cumsum < topP. + static func topPStage( + graph: MPSGraph, probs: MPSGraphTensor, topP: MPSGraphTensor, name: String + ) -> MPSGraphTensor { + let cumsum = graph.cumulativeSum(probs, axis: 1, exclusive: true, reverse: false, name: "\(name)_cs") + return graph.lessThan(cumsum, topP, name: "\(name)_mask") + } + + /// Combine boolean masks, apply to probs, and re-normalize. + static func maskAndNormalizeStage( + graph: MPSGraph, probs: MPSGraphTensor, masks: [MPSGraphTensor], name: String + ) -> MPSGraphTensor { + var combined = masks[0] + for i in 1.. MPSGraphTensor { + let cumsum = graph.cumulativeSum(probs, axis: 1, exclusive: false, reverse: false, name: "\(name)_cs") + let mask = graph.greaterThanOrEqualTo(cumsum, random, name: "\(name)_sel") + let maskFloat = graph.cast(mask, to: .float32, name: "\(name)_sf") + return graph.reductionArgMaximum(with: maskFloat, axis: 1, name: name) + } + + /// Gather the final token ID from topK indices using the selected position. + static func gatherTokenStage( + graph: MPSGraph, topKIndices: MPSGraphTensor, selectedIdx: MPSGraphTensor, k: Int, name: String + ) -> MPSGraphTensor { + let idxI32 = graph.cast(selectedIdx, to: .int32, name: "\(name)_i32") + let flat = graph.reshape(topKIndices, shape: [k as NSNumber], name: "\(name)_flat") + let idxFlat = graph.reshape(idxI32, shape: [1 as NSNumber], name: "\(name)_idx") + return graph.gatherAlongAxis(0, updates: flat, indices: idxFlat, name: name) + } } // Conformance to MPSGraphSampler protocol diff --git a/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift b/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift new file mode 100644 index 00000000..4c2a57ef --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift @@ -0,0 +1,125 @@ +// 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 penalty buffers for GPU repetition penalty. +/// +/// Uses a split design to avoid races between GPU reads and CPU writes: +/// - `recordToken()`: updates only the CPU-side ring buffer (no MTLBuffer writes) +/// - `syncBuffer(forStep:)`: writes the full penalty state to a specific buffer +/// slot, called at encode time when the gate guarantees that slot is not in use +/// +/// Thread safety: `recordToken` is called from Metal completion callbacks. +/// `syncBuffer` is called from the encode thread. The gate serializes them: +/// `syncBuffer(N)` is called only after the gate releases (meaning the previous +/// user of slot N%depth completed), and `recordToken` fires from that completion. +final class RepetitionPenaltyGPUState: @unchecked Sendable { + let penaltyBuffers: [MTLBuffer] + let vocabSize: Int + let pipelineDepth: Int + let penalty: Float16 + let windowSize: Int + + private var ring: [Int32] + private var writeIndex: Int = 0 + private var count: Int = 0 + private var refCounts: [Int32: Int] = [:] + private var dirtyTokens: [(added: [Int32], evicted: [Int32])] + + init(device: MTLDevice, vocabSize: Int, pipelineDepth: Int, penalty: Double, windowSize: Int?) throws { + self.vocabSize = vocabSize + self.pipelineDepth = pipelineDepth + self.penalty = Float16(penalty) + self.windowSize = windowSize ?? 256 + + let bufferSize = vocabSize * MemoryLayout.size + var buffers: [MTLBuffer] = [] + for _ in 0.. MTLBuffer { + let slot = step % pipelineDepth + let buf = penaltyBuffers[slot] + let ptr = buf.contents().assumingMemoryBound(to: Float16.self) + + let dirty = dirtyTokens[slot] + for tokenId in dirty.evicted { + ptr[Int(tokenId)] = Float16(1.0) + } + for tokenId in dirty.added { + ptr[Int(tokenId)] = penalty + } + dirtyTokens[slot] = (added: [], evicted: []) + + return buf + } + + /// Record a newly generated token (CPU-side bookkeeping only). + /// + /// Called from the completion callback. Does NOT write to MTLBuffers directly. + /// Instead, queues changes to be applied per-slot at the next `buffer(forStep:)` call. + func recordToken(_ token: Int32) { + guard token >= 0 && Int(token) < vocabSize else { return } + + var evictedToken: Int32 = -1 + if count == windowSize { + let evictSlot = writeIndex + let candidate = ring[evictSlot] + if candidate >= 0 { + refCounts[candidate, default: 0] -= 1 + if refCounts[candidate, default: 0] <= 0 { + refCounts.removeValue(forKey: candidate) + evictedToken = candidate + } + } + } else { + count += 1 + } + + ring[writeIndex] = token + writeIndex = (writeIndex + 1) % windowSize + refCounts[token, default: 0] += 1 + + for i in 0..= 0 { + dirtyTokens[i].evicted.append(evictedToken) + } + dirtyTokens[i].added.append(token) + } + } + + /// Reset all state (called on engine reset). + func reset() { + for buf in penaltyBuffers { + let ptr = buf.contents().assumingMemoryBound(to: Float16.self) + for i in 0.. Date: Wed, 19 Aug 2026 10:04:59 -0700 Subject: [PATCH 3/7] Fix validation, silent drops, and force-unwraps in repetition penalty - validate(): reject penalty < 1.0, orphan window, and penalty + json-schema (constrained generation does not support penalty on the pipelined engine) - Greedy strategy: forward repetition penalty to SamplingConfiguration (was silently constructing config without it) - Force-unwraps: replace repetitionPenalty! with guard-let in ConstrainedDecodingStrategy and ConstrainedGenerator - fallbackSampler(from:): precondition catches wrong overload usage --- .../ConstrainedDecodingStrategy.swift | 6 ++++-- .../ConstrainedGenerator.swift | 6 ++++-- .../Samplers/SamplingConfiguration.swift | 4 ++++ .../Tools/llm-runner/LLMRunnerMain.swift | 18 ++++++++++++++++-- 4 files changed, 28 insertions(+), 6 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift index bc5d4142..dc96430a 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift @@ -136,14 +136,16 @@ public struct ConstrainedDecodingStrategy: DecodingStrategy { } var maskedLogits = logits - if samplingConfiguration.needsRepetitionPenalty { + if samplingConfiguration.needsRepetitionPenalty, + let penalty = samplingConfiguration.repetitionPenalty + { let window = samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } ?? generatedTokens.count RepetitionPenaltyProcessor.apply( to: &maskedLogits, recentTokenIds: generatedTokens.suffix(window), - penalty: Float(samplingConfiguration.repetitionPenalty!) + penalty: Float(penalty) ) } _ = session.applyMask(to: &maskedLogits) diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift index 2f64b535..dfaf6c23 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift @@ -204,14 +204,16 @@ public struct ConstrainedGenerator: DecodingStrategy { } var maskedLogits = logits - if samplingConfiguration.needsRepetitionPenalty { + if samplingConfiguration.needsRepetitionPenalty, + let penalty = samplingConfiguration.repetitionPenalty + { let window = samplingConfiguration.repetitionPenaltyWindow.map { min($0, generatedTokens.count) } ?? generatedTokens.count RepetitionPenaltyProcessor.apply( to: &maskedLogits, recentTokenIds: generatedTokens.suffix(window), - penalty: Float(samplingConfiguration.repetitionPenalty!) + penalty: Float(penalty) ) } _ = session.applyMask(to: &maskedLogits) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift b/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift index f6b19746..ad9faae0 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift @@ -314,6 +314,10 @@ extension SamplingConfiguration { /// - Parameter logits: Mutable array of Float16 logits. May be modified during sampling. /// - Returns: The sampled token ID. public func fallbackSampler(from logits: inout [LogitsScalarType]) -> Int32 { + precondition( + !needsRepetitionPenalty, + "Use fallbackSampler(from:tokenHistory:) when repetition penalty is configured" + ) return CompositeSampler.sample(from: &logits, config: self) } diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index b00bd7e9..3cd3adb6 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -262,6 +262,16 @@ struct LLMRunner: AsyncParsableCommand, Sendable { if videoFrames < 1 { throw ValidationError("--video-frames must be >= 1") } + if let p = repetitionPenalty, p < 1.0 { + throw ValidationError("--repetition-penalty must be >= 1.0") + } + if repetitionPenaltyWindow != nil && repetitionPenalty == nil { + throw ValidationError("--repetition-penalty-window requires --repetition-penalty") + } + if repetitionPenalty != nil && jsonSchema != nil { + throw ValidationError( + "--repetition-penalty cannot be used with --json-schema (constrained generation does not support penalty on the pipelined engine)") + } } func run() async throws { @@ -865,13 +875,17 @@ struct LLMRunner: AsyncParsableCommand, Sendable { combined: !synchronousSampling ) case "greedy": - // Fatal error if topK/topP/minP set with greedy if topK != nil || topP != nil || minP != nil { print("Error: --top-k, --top-p, and --min-p cannot be used with --sampling-strategy greedy") print("Use --sampling-strategy temperature with --top-k/--top-p/--min-p, or remove them for greedy") throw ExitCode.failure } - config = SamplingConfiguration(temperature: 0, combined: !synchronousSampling) + config = SamplingConfiguration( + temperature: 0, + repetitionPenalty: repetitionPenalty, + repetitionPenaltyWindow: repetitionPenaltyWindow, + combined: !synchronousSampling + ) default: print("Error: Unknown sampling strategy '\(samplingStrategy)'") print("Valid options: 'temperature', 'greedy'") From 1f61e2547b0da753d335011efced69789f5a722e Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Wed, 19 Aug 2026 10:31:34 -0700 Subject: [PATCH 4/7] Add sentinel test for MPSGraph completion ordering assumption RepetitionPenaltyGPUState relies on completions firing in submission order (no additional synchronization). This is observed behavior on a single MTLCommandQueue but not documented by Apple. The test validates the assumption and will break if the dispatch model changes. --- .../Samplers/RepetitionPenaltyGPUState.swift | 10 +-- .../CoreAIPipelinedTests.swift | 64 +++++++++++++++++++ 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift b/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift index 4c2a57ef..762cebf8 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/RepetitionPenaltyGPUState.swift @@ -10,13 +10,13 @@ import Metal /// /// Uses a split design to avoid races between GPU reads and CPU writes: /// - `recordToken()`: updates only the CPU-side ring buffer (no MTLBuffer writes) -/// - `syncBuffer(forStep:)`: writes the full penalty state to a specific buffer +/// - `buffer(forStep:)`: writes the full penalty state to a specific buffer /// slot, called at encode time when the gate guarantees that slot is not in use /// -/// Thread safety: `recordToken` is called from Metal completion callbacks. -/// `syncBuffer` is called from the encode thread. The gate serializes them: -/// `syncBuffer(N)` is called only after the gate releases (meaning the previous -/// user of slot N%depth completed), and `recordToken` fires from that completion. +/// Thread safety: relies on MPSGraph runAsync completions being dispatched in +/// submission order on a single MTLCommandQueue (observed behavior, validated by +/// `MPSGraphCompletionOrderingTests`). The gate further ensures that +/// `buffer(forStep:)` does not overlap with `recordToken` for the same slot. final class RepetitionPenaltyGPUState: @unchecked Sendable { let penaltyBuffers: [MTLBuffer] let vocabSize: Int diff --git a/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift b/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift index 049366e3..decbefc3 100644 --- a/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift +++ b/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift @@ -6,6 +6,7 @@ import CoreAI import Foundation import Metal +import Synchronization import TestUtilities import Testing @@ -653,3 +654,66 @@ struct GPUSamplerContinuationSyncTests { #expect(received.count == tokenCount) } } + +// MARK: - MPSGraph Completion Ordering Sentinel + +/// Validates that MPSGraphExecutable.runAsync completionHandler calls are dispatched +/// in submission order when using a single MTLCommandQueue. This is not documented by +/// Apple but is relied upon by RepetitionPenaltyGPUState (recordToken is called from +/// these completions without additional synchronization). +/// +/// If this test fails, RepetitionPenaltyGPUState needs a lock. +@Suite("MPSGraph completion ordering", .enabled(if: !CIEnvironment.isVM)) +struct MPSGraphCompletionOrderingTests { + static let device: MTLDevice? = MTLCreateSystemDefaultDevice() + static let vocabSize = 512 + + @Test("completions fire in submission order on a single command queue") + func completionsAreSerial() async throws { + let device = try #require(Self.device) + let queue = try #require(device.makeCommandQueue()) + let sampler = try MPSGraphArgmaxSampler(device: device, vocabSize: Self.vocabSize) + + let stepCount = 16 + let orderRecord = Mutex<[Int]>([]) + + // Submit stepCount encode calls back-to-back. Each completion records its index. + for i in 0...size, + options: .storageModeShared)) + let ptr = logitsBuffer.contents().assumingMemoryBound(to: Float16.self) + for v in 0...size, options: .storageModeShared)) + + sampler.encode( + to: queue, + logitsBuffer: logitsBuffer, + logitsOffset: 0, + outputBuffer: outputBuffer, + outputOffset: 0, + completion: { _, _ in + orderRecord.withLock { $0.append(i) } + } + ) + } + + // Wait for all completions via a sentinel command buffer. + await withCheckedContinuation { (cont: CheckedContinuation) in + guard let cmdBuf = queue.makeCommandBuffer() else { + cont.resume() + return + } + cmdBuf.addCompletedHandler { _ in cont.resume() } + cmdBuf.commit() + } + + let observed = orderRecord.withLock { $0 } + #expect(observed == Array(0.. Date: Wed, 19 Aug 2026 10:34:37 -0700 Subject: [PATCH 5/7] Propagate encode errors via completion instead of swallowing with try? Also note GPU window cap (256) in --help text. --- .../Samplers/MPSGraphSamplers.swift | 24 ++++++++++++------- .../Tools/llm-runner/LLMRunnerMain.swift | 2 +- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift index d8b0baa7..751995f6 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift @@ -1056,10 +1056,14 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { completion: @escaping (Int32, Error?) -> Void ) { if queryLength == 1 { - try? encode( - to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0, - outputBuffer: outputBuffer, outputOffset: outputOffset, - applyBitmask: applyBitmask, completion: completion) + do { + try encode( + to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0, + outputBuffer: outputBuffer, outputOffset: outputOffset, + applyBitmask: applyBitmask, completion: completion) + } catch { + completion(0, error) + } return } let logitsOffset = (queryLength - 1) * vocabSize * MemoryLayout.size @@ -1077,10 +1081,14 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { blitEncoder.endEncoding() blitCmdBuffer.commit() - try? encode( - to: queue, logitsBuffer: tempBuffer, logitsOffset: 0, - outputBuffer: outputBuffer, outputOffset: outputOffset, - applyBitmask: applyBitmask, completion: completion) + do { + try encode( + to: queue, logitsBuffer: tempBuffer, logitsOffset: 0, + outputBuffer: outputBuffer, outputOffset: outputOffset, + applyBitmask: applyBitmask, completion: completion) + } catch { + completion(0, error) + } } /// Encode sampling with repetition penalty buffer. diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index 3cd3adb6..476cdd5a 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -115,7 +115,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable { @Option( name: .customLong("repetition-penalty-window"), - help: "Number of recent tokens to consider for repetition penalty (default: all)") + help: "Number of recent tokens to consider for repetition penalty (default: all; GPU engine caps at 256)") var repetitionPenaltyWindow: Int? @Option(help: "Sampling strategy. Options: 'temperature' (default), 'greedy'") From 4c3a355e939cba7bd33f716d7c15e3ab62ebc1fe Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Wed, 19 Aug 2026 10:40:09 -0700 Subject: [PATCH 6/7] Fix swift-format lint warnings --- swift/Sources/Tools/llm-runner/LLMRunnerMain.swift | 3 ++- swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index 476cdd5a..4a2dc14c 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -270,7 +270,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable { } if repetitionPenalty != nil && jsonSchema != nil { throw ValidationError( - "--repetition-penalty cannot be used with --json-schema (constrained generation does not support penalty on the pipelined engine)") + "--repetition-penalty cannot be used with --json-schema" + + " (constrained generation does not support penalty on the pipelined engine)") } } diff --git a/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift b/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift index decbefc3..7e132e15 100644 --- a/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift +++ b/swift/Tests/LanguageModelsTests/CoreAIPipelinedTests.swift @@ -713,7 +713,8 @@ struct MPSGraphCompletionOrderingTests { } let observed = orderRecord.withLock { $0 } - #expect(observed == Array(0.. Date: Wed, 19 Aug 2026 11:05:33 -0700 Subject: [PATCH 7/7] Use feedTensors to order runAsync inputs for penalty encode The feeds dictionary used at compile time has no guaranteed order. Using executable.feedTensors ensures the inputs array matches the order the compiled graph expects. --- .../Samplers/MPSGraphSamplers.swift | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift index 751995f6..f140c22f 100644 --- a/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift +++ b/swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift @@ -1122,6 +1122,16 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { let outputData = MPSGraphTensorData( outputBuffer, shape: [1 as NSNumber], dataType: .int32) + let tensorDataMap: [MPSGraphTensor: MPSGraphTensorData] = [ + logitsPlaceholder: logitsData, + penaltyPlaceholder!: penaltyData, + temperaturePlaceholder: temperatureData, + randomPlaceholder: randomData, + topPPlaceholder: topPData, + minPPlaceholder: minPData, + ] + let inputs = executable.feedTensors!.map { tensorDataMap[$0]! } + let execDesc = MPSGraphExecutableExecutionDescriptor() execDesc.completionHandler = { [outputBuffer, outputOffset] (_, error) in if let error = error { @@ -1135,7 +1145,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable { } executable.runAsync( with: queue, - inputs: [logitsData, penaltyData, temperatureData, randomData, topPData, minPData], + inputs: inputs, results: [outputData], executionDescriptor: execDesc) }