Int8 KV cache quantization for 2× memory reduction

[msnabiel]

Summary

Implementation of int8 per-token KV cache quantization that halves cache memory with negligible accuracy impact, aligned with Apple's Foundation Language Models approach.

Background

Apple's Intelligence Foundation Language Models (2025) employ int8 quantization of the KV cache to reduce memory footprint during inference. This implementation brings the same technique to coreai-models export pipeline, extending it with:

• iOS StaticShapeEngine 4-state support
• macOS opt-in flag for runtime compatibility
• Per-token-per-head scale granularity (macOS) for improved accuracy

Implementation Highlights

iOS (always-on)

• Per-token symmetric quantization: scale = abs_max(k) / 127.0 per token
• Float16 scale caches: [n_layers, 1, 1, 1, max_seq_len] minimal overhead
• StaticShapeEngine: 4-state support (k_cache, v_cache, k_scale_cache, v_scale_cache)
• Memory savings: ~2× reduction (e.g., 16k context 7B: 4GB → 2GB cache)
• Transparent: Quantization in KVCacheHandler.update_and_fetch, SDPA sees float16

macOS (opt-in)

• Flag: ExportConfig.quantize_kv_cache (default False - runtime needs 4-state)
• Per-token-per-head scales: [n_layers, 1, n_kv_heads, max_seq_len, 1]
• Float16 fallback: Uses upstream's mutable_cache_update_and_fetch when scales absent

Pull Request

Branch: msnabiel:coreai-models:nabiel/int8-kv-cache
Compare: main...msnabiel:coreai-models:nabiel/int8-kv-cache

Files changed: 18 files (+402/-37 lines)
Commits:

1. Add int8 KV cache quantization for iOS (2× memory reduction)
2. Add opt-in int8 KV cache quantization for macOS export pipeline

Testing Status

• iOS 7B models at 16k context
• macOS sequential/pipelined engines with --quantize-kv-cache
• Float16 fallback path compatibility

References

Apple Research:

• Apple Intelligence Foundation Language Models (2025): "int8 quantization of the KV cache to reduce memory footprint"
• CommVQ (ICML 2025): 87.5% compression with minimal accuracy loss
• QuantSpec (2025): 4-bit hierarchical KV cache with ~2.5× speedup

Industry Standard:

• llama.cpp, MLX, vLLM/TensorRT-LLM all default to int8 KV cache for unified-memory systems
• Per-token asymmetric quantization well-established as best practice

Happy to address any review feedback or questions.

[stikves]

Thanks for the feedback

Yes, KV-Cache quantization would be a useful feature.