Add repetition penalty to the llm-runner

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,18 @@ public struct ConstrainedDecodingStrategy: DecodingStrategy {
         }
 
         var maskedLogits = logits
+        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(penalty)
+            )
+        }
         _ = session.applyMask(to: &maskedLogits)
 
         let bestToken = CompositeSampler.sample(from: &maskedLogits, config: samplingConfiguration)
@@ -296,6 +309,7 @@ extension ConstrainedDecodingStrategy.ConstrainedDecodedSequence {
                 do {
                     result = try await ConstrainedDecodingStrategy.generateOneToken(
                         inputTokens: inputTokens,
+                        generatedTokens: generatedTokens,
                         session: &session,
                         inferenceEngine: inferenceEngine,
                         samplingConfiguration: samplingConfiguration,

swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedGenerator.swift

@@ -204,6 +204,18 @@ public struct ConstrainedGenerator: DecodingStrategy {
             }
 
             var maskedLogits = logits
+            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(penalty)
+                )
+            }
             _ = session.applyMask(to: &maskedLogits)
 
             let bestToken = CompositeSampler.sample(from: &maskedLogits, config: samplingConfiguration)

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift

@@ -566,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
@@ -817,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,
@@ -977,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()
@@ -994,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,

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)

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)

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,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)
+        let nextToken = samplingConfig.fallbackSampler(
+            from: &logitBuffer, tokenHistory: inputTokens[generationStartOffset...])
         sampleSpan.end()
         CLILogger.log("Token: \(nextToken), processed: \(processedTokenCount)")
         return (logits: actualLogits, token: nextToken)
@@ -657,6 +659,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 +678,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 +715,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