Skip to content

perf(cache): write the decode KV row in place instead of copying the cache #1959

Closed

Description

@inureyes

Problem / Background

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:

  1. 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.
  2. Fuse the sliding-layer k RoPE into the same kernel.
  3. 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.
  • Benchmark record in docs/benchmark_results/.

Activity

  1. added
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    platform:macosmacOS (Apple Silicon) specific
    and removed on Sep 23, 2026
  2. added a commit that references this issue on Sep 23, 2026
    820788d
  3. added a commit that references this issue on Sep 29, 2026
    d3f6ee7
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:macosmacOS (Apple Silicon) specificpriority:highHigh prioritystatus:doneCompletedtype:performancePerformance improvements

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions