You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
perf(cache): write the decode KV row in place instead of copying the cache #1959
With pipelined decode, the previous step's command buffer holds the KV cache buffer, so the slice_update calls in KVCache::update_fp16 (src/lib/mlxcel-core/src/cache.rs:939-1010) cannot donate it and copy both the K and V caches of every layer on every token. At context 2048 on command-r7b that is about 268 MB read and written per token.
Reproduced in plain MLX: 64 updates on [1,8,CTX+256,128] fp16 take 0.52 ms (CTX 512) and 0.56 ms (CTX 2048) with no outstanding view, and 1.00 ms and 1.39 ms when the previous step's view is alive. mlx-lm has the same structure.
Proposed Solution
Add an MLX primitive that writes only the new rows into the cache buffer in place (the output shares the input buffer). This is safe while the previous step is in flight, because that step reads rows < pos and the write targets row pos. Buffers shared with another owner (prompt-cache snapshots, detached or cloned caches, forks) must keep the current copying slice_update, so KVCache has to track exclusive ownership of its buffer.
Plan, in order:
The in-place write for command-r (cohere2) only, measured through scripts/bench_decode.sh against MLXCEL_DIAG_SKIP_DECODE_KV_WRITE as the upper bound.
Fuse the sliding-layer k RoPE into the same kernel.
Ownership tracking in KVCache and rollout to every model.
Also record: the existing quantized KV modes do not help decode at short context (pp512/tg128: int8 91.7 tok/s and fp16+turbo4 30.7 tok/s against fp16 116 tok/s), since KV is about 1.5% of the bytes read per token there.
Acceptance Criteria
Greedy output is byte-identical to the copying path.
Decode gain measured at contexts 16, 512, and 2048.
Shared-buffer cases (prompt cache, detach, speculative rollback) fall back to the copy and are covered by tests.
Activity
perf(cache): write the decode KV row in place instead of copying the …
perf(cache): write rotating-cache warmup rows in place (#1963)
perf(cache): write the decode KV row in place instead of copying the …