Repository navigation
perf(cache): write rotating-cache warmup rows in place - #1963
Merged
Merged
Conversation
RotatingKVCache kept the copying slice_update after #1959 moved the dense FP16 cache to in-place decode writes, so sliding-window layers still copied their whole buffer every decode step while the previous pipelined step held it. While the cache is still filling its window, each decode token lands on a slot that was never valid. Those writes now go through inplace_slice_write under the same proof as the dense cache: the cache's own last write produced the current buffers and the row lands at that write's length. After the window is full a step overwrites the oldest slot, which a prompt-cache snapshot may still hold and whose sharing cannot be detected, so those writes keep copying. MLXCEL_KV_INPLACE_WRITE=0 still restores the copying path everywhere. M1 Ultra, pp512/tg128 against main: Gemma 3 4B and Gemma 4 12B (window 1024) about +3% decode; Gemma 4 E4B (512) and gpt-oss 20B (128) unchanged, as they decode past their windows. Greedy output identical on five models. A new test checks that warmup writes leave a shared snapshot intact and a wrapped overwrite copies; letting the steady state write in place fails it. Closes #1962
inureyes
added
status:review
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #1961:
RotatingKVCachenow writes decode rows in place while it is still filling its sliding window. Every such row lands on a slot that was never valid, so the dense cache's ownership proof applies unchanged: the cache's own last write produced the current buffers and the row lands at that write's length. After the window is full a step overwrites the oldest slot, which a prompt-cache snapshot may still hold, so those writes keep the copyingslice_update.Measurement
M1 Ultra,
mlxcel-bench-decodewithbench_decode.sh's arguments (pp512, tg128), two runs per arm against main (rotating path off), load 6 to 7.5:The 1024-window models decode inside the warmup phase at this condition; the 512- and 128-window models fill their windows during the prompt and decode in the steady state, which is unchanged by design. Record:
docs/benchmark_results/rotating-kv-inplace-write-m1ultra-2026-09-24.md.Validation
MLXCEL_KV_INPLACE_WRITEon and off for Gemma 3 4B, Gemma 4 12B, Gemma 4 E4B, gpt-oss 20B and EXAONE 4.0 1.2B (two prompts, 300 tokens).make verify-test: 11431 passed, 0 failed. Workspace clippy and fmt pass.Closes #1962