From 99c7d499571fb3d5a0567e944fd4c92c002342d3 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Sat, 5 Sep 2026 13:16:55 -0700 Subject: [PATCH 1/3] Fix pipelined prefix reuse skipping the last sampled token After pipelined generation, history.count exceeds processedTokenCount by 1 because the last sampled token is yielded but never fed back through the model. On the next multi-turn call, resolve() matches that unprocessed token, placing new tokens at the wrong KV cache position. Clamp commonPrefix to processedTokenCount before the backup-by-1 step so the unprocessed trailing token is re-included in the next prefill. Fixes #234 --- .../CoreAIPipelinedEngine.swift | 6 +++ .../LanguageModelsTests/TestUtilities.swift | 10 ++++- .../UnifiedGenerationAPITests.swift | 37 +++++++++++++++++++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 23da5c2e..b11dec94 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -168,6 +168,12 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable // 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...] + } + // 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..4ed02125 100644 --- a/swift/Tests/LanguageModelsTests/TestUtilities.swift +++ b/swift/Tests/LanguageModelsTests/TestUtilities.swift @@ -76,9 +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 lastPrefixHitCount = commonPrefix + // 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...] + } + if commonPrefix < processedTokenCount { // Input diverged — rewind processedTokenCount = commonPrefix diff --git a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift index 81cc436f..fbfbc06d 100644 --- a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift +++ b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift @@ -715,6 +715,43 @@ struct PrefixCachingTests { ) {} #expect(engine.lastPrefixHitCount == 3) } + + @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). + // Without the clamp, lastPrefixHitCount would be 6 and token 30 would be skipped. + #expect(engine.lastPrefixHitCount <= 6) + // processedTokenCount should account for: 5 (carried) + 3 (token 30 + two user) + 2 (generated) = 10 + #expect(engine.processedTokenCount == 10) + } } // MARK: - Deterministic RNG for Tests From 669afcb84e6f989bb85fd9df72da28d78d4070af Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Fri, 11 Sep 2026 10:20:28 -0700 Subject: [PATCH 2/3] Report reused-token count after clamp and truncate stale history Addresses review on #237: - Move lastPrefixHitCount after the KV-range clamp so it reports tokens actually reused (was over-reporting by one -- the re-prefilled token). - Truncate history to the clamped prefix so the re-prefilled trailing token is not duplicated when new tokens are appended, avoiding spurious divergence on the third and later turns. - Add multi-turn and multi-token-gap (cancellation/early-EOS) regression tests. --- .../CoreAIPipelinedEngine.swift | 4 +- .../LanguageModelsTests/TestUtilities.swift | 4 +- .../UnifiedGenerationAPITests.swift | 83 ++++++++++++++++++- 3 files changed, 87 insertions(+), 4 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 679f4486..1aff70fb 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -163,7 +163,6 @@ 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 @@ -172,8 +171,11 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable 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 4ed02125..c807370b 100644 --- a/swift/Tests/LanguageModelsTests/TestUtilities.swift +++ b/swift/Tests/LanguageModelsTests/TestUtilities.swift @@ -78,15 +78,17 @@ class MockEngine: InferenceEngine, @unchecked Sendable { // Implicit prefix caching: resolve input against history let (rawCommonPrefix, _) = history.resolve(input: input) var commonPrefix = rawCommonPrefix - lastPrefixHitCount = commonPrefix // 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 { // Input diverged — rewind processedTokenCount = commonPrefix diff --git a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift index fbfbc06d..a4fc55a7 100644 --- a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift +++ b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift @@ -747,11 +747,90 @@ struct PrefixCachingTests { } // The clamp must cap the prefix to processedTokenCount (5), not history.count (6). - // Without the clamp, lastPrefixHitCount would be 6 and token 30 would be skipped. - #expect(engine.lastPrefixHitCount <= 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)" + ) + } + } + } + + @Test("prefix clamp handles a multi-token pipelined gap (cancellation / early-EOS)") + func multiTurnPrefixClampWithLargerPipelinedGap() async throws { + let engine = MockEngine(tokens: [10, 20, 30, 40, 50], maxContextLength: 200) + + // Turn 1: prompt [1, 2, 3], generate 5 tokens. + var context: [Int32] = [1, 2, 3] + for try await output in try await engine.generate( + with: context, + samplingConfiguration: .greedy, + inferenceOptions: InferenceOptions(maxTokens: 5) + ) { + context.append(output.tokenId) + } + // context = [1, 2, 3, 10, 20, 30, 40, 50], history.count = 8, processedTokenCount = 8 + #expect(context.count == 8) + + // A cancellation/early-EOS drain can leave more than one trailing token + // yielded but never fed back through the model. Simulate a gap of three. + engine.processedTokenCount -= 3 + // Now: history.count = 8, processedTokenCount = 5 (positions 5, 6, 7 lack KV) + + // 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 caps the prefix to processedTokenCount (5), re-prefilling + // 30, 40, 50, 77, 78 — regardless of how many tokens the gap spans. + #expect(engine.lastPrefixHitCount == 5) + // History must mirror context exactly — no duplicated trailing tokens. + #expect(engine.history.tokens.count == context.count) + // processedTokenCount: 5 (carried) + 5 (re-prefill) + 2 (generated) = 12 + #expect(engine.processedTokenCount == 12) + } } // MARK: - Deterministic RNG for Tests From e092b78fcb5bacafcdd96a33e5b81bf56f7c2e42 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Mon, 14 Sep 2026 09:41:28 -0700 Subject: [PATCH 3/3] Drop unreachable multi-token pipelined-gap test The pipelined gap is structurally always 1: runCompletion has no early-EOS break and increments processedTokenCount eagerly at launch, so no path yields a gap >1. Remove the synthetic gap-of-3 test and its cancellation/early-EOS comment, and note the remaining clamp tests validate the algorithm against a MockEngine copy of the clamp, not the real engine's KV interaction. --- .../UnifiedGenerationAPITests.swift | 44 ++----------------- 1 file changed, 4 insertions(+), 40 deletions(-) diff --git a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift index a4fc55a7..06955dd9 100644 --- a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift +++ b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift @@ -716,6 +716,10 @@ 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) @@ -791,46 +795,6 @@ struct PrefixCachingTests { } } } - - @Test("prefix clamp handles a multi-token pipelined gap (cancellation / early-EOS)") - func multiTurnPrefixClampWithLargerPipelinedGap() async throws { - let engine = MockEngine(tokens: [10, 20, 30, 40, 50], maxContextLength: 200) - - // Turn 1: prompt [1, 2, 3], generate 5 tokens. - var context: [Int32] = [1, 2, 3] - for try await output in try await engine.generate( - with: context, - samplingConfiguration: .greedy, - inferenceOptions: InferenceOptions(maxTokens: 5) - ) { - context.append(output.tokenId) - } - // context = [1, 2, 3, 10, 20, 30, 40, 50], history.count = 8, processedTokenCount = 8 - #expect(context.count == 8) - - // A cancellation/early-EOS drain can leave more than one trailing token - // yielded but never fed back through the model. Simulate a gap of three. - engine.processedTokenCount -= 3 - // Now: history.count = 8, processedTokenCount = 5 (positions 5, 6, 7 lack KV) - - // 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 caps the prefix to processedTokenCount (5), re-prefilling - // 30, 40, 50, 77, 78 — regardless of how many tokens the gap spans. - #expect(engine.lastPrefixHitCount == 5) - // History must mirror context exactly — no duplicated trailing tokens. - #expect(engine.history.tokens.count == context.count) - // processedTokenCount: 5 (carried) + 5 (re-prefill) + 2 (generated) = 12 - #expect(engine.processedTokenCount == 12) - } } // MARK: - Deterministic RNG for Tests