feat: KV cache compression pipeline (Attention Matching + KVTC)

[buftar]

Summary

Adds a two-stage opt-in KV cache compression pipeline for long-context
inference on Apple Silicon. No changes to the public inference API —
existing behaviour is unchanged when the new flags are not set.

• Attention Matching (AM): compacts the KV cache by selecting the most important tokens per head using NNLS beta-fitting and OLS value-fitting. Sink tokens are always preserved. Configurable ratio (default 4x, range 1–8x).
• KVTC (KV Cache Transform Coding): PCA-based byte compression of key tensors using a calibration bundle generated offline.
• Calibration CLI: omlx calibrate-kv <model> generates per-layer PCA bundles and per-head entropy curves.
• Cache integration: CompressedPagedSSDCacheManager wraps the existing paged SSD cache, applying the pipeline transparently on each eviction cycle.
• Admin dashboard: compression status (ratio, latency, ok/fail counts) integrated into the Runtime Cache Observability card; AM ratio slider with live update via POST /admin/api/compression/config.
• Benchmark CLI: omlx benchmark-compression for cosine similarity and throughput validation.

Usage

# 1. Calibrate
omlx calibrate-kv Qwen/Qwen2.5-7B

# 2. Serve with compression
omlx serve \
  --compression-bundle ~/.omlx/calibration/kv_pca_calibration.npz \
  --compression-am-ratio 4.0 \
  --paged-cache-dir /path/to/ssd

See docs/KV_Cache_Compression.md for calibration, configuration, and troubleshooting.

Test plan

Compression test suite run locally — 135 passed, 0 failed:

pytest tests/test_am.py tests/test_kvtc.py tests/test_linalg_utils.py \
       tests/test_calibrator.py tests/test_pipeline.py \
       tests/test_cache_integration.py tests/test_compression_benchmark.py \
       tests/test_observability.py -v

Full fast suite (pytest -m "not slow"): 3120 passed. The 30 remaining failures are pre-existing on main and unrelated to this PR (embedding model loading, hardware detection, integration server tests).

References

• Fast KV Compaction via Attention Matching
• KV Cache Transform Coding

[jundot]

hey @buftar, this is a really interesting approach. The two-stage AM + KVTC design based on the attention matching and transform coding papers is a solid idea, and i can see the potential here for long-context workloads on Apple Silicon.

After going through the code though, i think there's a fair bit of distance between the current state and something production-ready. Merging as-is would introduce a lot of technical debt that would be hard to unwind later. I'd suggest closing this one and re-submitting the core algorithm (AM compaction + KVTC compression) as a smaller, focused PR once the issues below are addressed. That way it's easier to review and validate properly.

Here's what i found:

• _dp_allocate_bits() budget includes n_tokens multiplier, making it so large that every component always gets max_bits=8. Variable-rate quantization never kicks in. All spike measurements were taken with this bug, so compression ratios and quality need re-validation.
• Doc claims KVTC ~16x, combined ~64x, <1% on GSM8K/MMLU. Spike measured 6.8x, 16x, and 0.72 cosine. No GSM8K/MMLU benchmarks were actually run.
• Lazy _pipeline init has no lock. Concurrent calls can duplicate the full pipeline construction.
• No validation that a calibration bundle matches the model's n_layers/n_heads/head_dim.
• Admin endpoint accepts zero or negative am_ratio, causing ZeroDivisionError.

Thanks for putting this together. The underlying idea is worth pursuing and i'd be happy to review a tighter version.