diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index e63c8ff5..623030d0 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -165,11 +165,19 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable // Implicit prefix caching: resolve input against history var (commonPrefix, resolvedNewTokens) = self.history.resolve(input: input) - self.lastPrefixHitCount = commonPrefix // Detect TRUE divergence before backup (tokens actually differ) let isDivergence = commonPrefix < input.count && commonPrefix < self.history.count + // Pipelined decode yields the last token without processing it; cap to KV-valid range. + if commonPrefix > self.engine.processedTokenCount { + commonPrefix = self.engine.processedTokenCount + resolvedNewTokens = input[commonPrefix...] + self.history.truncate(to: commonPrefix) + } + + self.lastPrefixHitCount = commonPrefix + // Ensure at least 1 token for prefill (seeds the decode loop). // Back up by 1 if the entire input is cached. if resolvedNewTokens.isEmpty && commonPrefix > 0 { diff --git a/swift/Tests/LanguageModelsTests/TestUtilities.swift b/swift/Tests/LanguageModelsTests/TestUtilities.swift index 801cd65a..c807370b 100644 --- a/swift/Tests/LanguageModelsTests/TestUtilities.swift +++ b/swift/Tests/LanguageModelsTests/TestUtilities.swift @@ -76,7 +76,17 @@ class MockEngine: InferenceEngine, @unchecked Sendable { _activeToken.withLock { $0 = token } // Implicit prefix caching: resolve input against history - let (commonPrefix, resolvedNewTokens) = history.resolve(input: input) + let (rawCommonPrefix, _) = history.resolve(input: input) + var commonPrefix = rawCommonPrefix + + // Pipelined decode yields the last token without processing it; cap to KV-valid range. + var resolvedNewTokens = input[commonPrefix...] + if commonPrefix > processedTokenCount { + commonPrefix = processedTokenCount + resolvedNewTokens = input[commonPrefix...] + history.truncate(to: commonPrefix) + } + lastPrefixHitCount = commonPrefix if commonPrefix < processedTokenCount { diff --git a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift index 81cc436f..06955dd9 100644 --- a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift +++ b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift @@ -715,6 +715,86 @@ struct PrefixCachingTests { ) {} #expect(engine.lastPrefixHitCount == 3) } + + // MockEngine embeds the same clamp as CoreAIPipelinedEngine and advances history and + // processedTokenCount in lockstep, so the +1 pipelined gap is injected by hand below rather than + // produced organically. These tests cover the clamp arithmetic and history/count invariants; the + // real engine's KV interaction and the organic gap are exercised by the #234 GPU reproduction. + @Test("multi-turn prefix clamp when processedTokenCount trails history") + func multiTurnPrefixClampWithPipelinedGap() async throws { + let engine = MockEngine(tokens: [10, 20, 30, 40, 50], maxContextLength: 200) + + // Turn 1: prompt [1, 2, 3], generate 3 tokens + var context: [Int32] = [1, 2, 3] + for try await output in try await engine.generate( + with: context, + samplingConfiguration: .greedy, + inferenceOptions: InferenceOptions(maxTokens: 3) + ) { + context.append(output.tokenId) + } + // context = [1, 2, 3, 10, 20, 30], history.count = 6, processedTokenCount = 6 + #expect(context.count == 6) + + // Simulate pipelined engine gap: last sampled token was never processed. + engine.processedTokenCount -= 1 + // Now: history.count = 6, processedTokenCount = 5 (position 5 has no KV entry) + + // Turn 2: append new user tokens, generate again + context.append(contentsOf: [77, 78]) + for try await output in try await engine.generate( + with: context, + samplingConfiguration: .greedy, + inferenceOptions: InferenceOptions(maxTokens: 2) + ) { + context.append(output.tokenId) + } + + // The clamp must cap the prefix to processedTokenCount (5), not history.count (6). + #expect(engine.lastPrefixHitCount == 5) + // processedTokenCount should account for: 5 (carried) + 3 (token 30 + two user) + 2 (generated) = 10 + #expect(engine.processedTokenCount == 10) + } + + @Test("five-turn prefix caching with pipelined gap stays consistent") + func fiveTurnPrefixCachingWithPipelinedGap() async throws { + let engine = MockEngine(tokens: [10, 20, 30, 40, 50], maxContextLength: 2000) + + var context: [Int32] = [1, 2, 3] + + for turn in 1...5 { + // Append user tokens for each turn beyond the first + if turn > 1 { + context.append(contentsOf: [Int32(turn * 100 + 1), Int32(turn * 100 + 2)]) + } + + let preGenProcessed = engine.processedTokenCount + for try await output in try await engine.generate( + with: context, + samplingConfiguration: .greedy, + inferenceOptions: InferenceOptions(maxTokens: 2) + ) { + context.append(output.tokenId) + } + + // Simulate pipelined gap: last token yielded but not processed. + engine.processedTokenCount -= 1 + + // History must mirror context (no duplicates from missing truncation). + #expect( + engine.history.tokens.count == context.count, + "Turn \(turn): history (\(engine.history.tokens.count)) drifted from context (\(context.count))" + ) + + // Prefix hit should reuse all prior KV-valid tokens, not trigger divergence. + if turn > 1 { + #expect( + engine.lastPrefixHitCount == preGenProcessed, + "Turn \(turn): expected prefix hit \(preGenProcessed), got \(engine.lastPrefixHitCount)" + ) + } + } + } } // MARK: - Deterministic RNG for Tests