Skip to content

Add KV cache quantization to mlx_lm.server (#1043)#1476

Open
katlun-lgtm wants to merge 1 commit into
ml-explore:mainfrom
katlun-lgtm:feat/server-kv-cache-quant
Open

Add KV cache quantization to mlx_lm.server (#1043)#1476
katlun-lgtm wants to merge 1 commit into
ml-explore:mainfrom
katlun-lgtm:feat/server-kv-cache-quant

Conversation

@katlun-lgtm

Copy link
Copy Markdown

Add KV cache quantization to mlx_lm.server (closes #1043)

What

Exposes the KV-cache quantization already supported by generate/stream_generate
as server CLI flags:

  • --kv-bits (default None → off)
  • --kv-group-size (default 64)
  • --quantized-kv-start (default 0)

When set, they are threaded into the server's sequential stream_generate path.

Why the batched path is handled by fallback, not plumbing

The server uses continuous batching (BatchGenerator) whenever a request
_is_batchable — which requires mergeable caches
(all(hasattr(c, "merge") ...)). QuantizedKVCache intentionally does not
implement merge(), so a quantized cache cannot participate in continuous
batching without new batch-membership support for quantized caches (a larger,
separate feature).

Rather than silently ignore --kv-bits in the default (batched) path, this PR
makes _is_batchable return False when --kv-bits is set, so those requests
use sequential generation, which applies the quantization. This keeps the
existing mergeable-cache invariant intact and makes the flag behave correctly.

Trade-off worth a maintainer call: with --kv-bits, requests no longer batch.
If you'd prefer, I can instead scope a follow-up that teaches the batched path
to carry quantized caches. Happy to go either way.

Changes

  • mlx_lm/server.py: 3 argparse flags; pass to stream_generate; _is_batchable
    returns False when kv_bits is set.

Test plan

Verified on Apple Silicon (M3 Max, mlx 0.31.2, mlx-lm 0.31.3) with
mlx-community/Qwen2.5-0.5B-Instruct-4bit:

  • --kv-bits flags parse with correct defaults (None / 64 / 0).
  • stream_generate(..., kv_bits=4) produces a QuantizedKVCache and coherent
    output; kv_bits=None stays a plain KVCache (regression check).
kv_bits=None: cache=KVCache          quantized=False coherent=True
kv_bits=4:    cache=QuantizedKVCache quantized=True  coherent=True
  • End-to-end HTTP smoke against a running mlx_lm.server --kv-bits 4 (recommended
    before merge; core generation path is exercised above).

Notes

Mirrors the existing generate.py flag names/help text verbatim for consistency.

…ize/--quantized-kv-start)

Wires the existing KV-cache quantization (already supported by generate/
stream_generate) into mlx_lm.server, closing ml-explore#1043.

Quantized KV caches do not implement merge(), so they cannot participate in
continuous batching; when --kv-bits is set, requests fall back to sequential
generation, which applies the quantization.
@katlun-lgtm

Copy link
Copy Markdown
Author

Verified end-to-end on M3 Max: mlx_lm.server --kv-bits 4 boots, /v1/chat/completions returns HTTP 200 and generates through the sequential (quantized-cache) path. The core stream_generate path is also confirmed to produce a QuantizedKVCache with kv_bits=4 (plain KVCache without). Output quality in any short snippet here reflects a 0.5B/4-bit test model, not the change itself.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add KV cache quantization support to server

1 participant