feat: min_quantize_tokens threshold + trim oversized KV buffers

[janhilgard]

Summary

Addresses both requests from PR #62 comment:

• Threshold-based quantization: Added kv_min_quantize_tokens (default 256) — sequences shorter than this skip quantization since the overhead exceeds memory savings for short entries
• Trim oversized KV buffers: Added _trim_to_offset() which trims pre-allocated KV arrays to their actual used size before storage. Applied unconditionally (both FP16 and quantized paths) since oversized buffers waste memory regardless of quantization

Changes

• vllm_mlx/memory_cache.py — new _trim_to_offset() helper, kv_min_quantize_tokens config field, updated store() logic
• vllm_mlx/scheduler.py — kv_cache_min_quantize_tokens field in SchedulerConfig, threaded into MemoryCacheConfig
• vllm_mlx/cli.py — --kv-cache-min-quantize-tokens CLI argument for both serve and bench commands
• tests/test_kv_cache_quantization.py — 7 new tests covering trim and threshold behavior

Test plan

• python3 -m pytest tests/test_kv_cache_quantization.py -v — all 23 tests pass (16 existing + 7 new)
• ruff check — no new lint errors introduced
• Manual test with server: --kv-cache-quantization --kv-cache-min-quantize-tokens 256

🤖 Generated with Claude Code

[waybarrios]

Pushed a commit with fixes for the _trim_to_offset function and the store() method.

The duplicate entry check in store() was happening after trim and quantize, so repeated tokens were doing expensive work for nothing. Moved it before those operations.

Rewrote _trim_to_offset to validate that offset is positive before slicing, which prevents silent corruption if offset ever goes negative. Replaced KVCache.new() with KVCache() so future mlx-lm versions that add attributes in init will not break silently. Added mx.eval on the sliced arrays so the original large buffer actually gets freed and memory tracking stays accurate instead of holding a lazy reference to the full allocation. Also added an early exit when no layer needs trimming to avoid unnecessary object creation on every store call.

Added validation for kv_min_quantize_tokens in MemoryCacheConfig so negative values get rejected at init time instead of silently forcing quantization on all sequences. Documented the field in the class docstring and added Args and Returns to the _trim_to_offset docstring.

[waybarrios]

@janhilgard I’m handing over the quantized KV work to you. It would be great if you could refactor the code, clean up any inconsistencies, and organize everything into a clearer, more maintainable structure.