Fix pipelined prefix reuse skipping the last sampled token - #237
Conversation
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 apple#234
|
Confirmed with the harness from #234: this fix passes on all three fixtures. Tested
One observation: Mac Studio M4 Max, macOS 27.0 |
Addresses review on apple#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.
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.
Thanks, updated to post process. However we can take the reported engine stats as "best effort", as correctness and performance are higher priorities. |
4 participants
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