Vramkit: Measuring LLM Weight and KV-Cache Compression on Apple Silicon
A benchmarking toolkit for optimizing LLM inference memory usage via lossless weight and KV-cache compression.
1Vramkit Framework Overview
Vramkit provides a suite of tools for measuring LLM weight and KV-cache compression. Built entirely on Apple Silicon, it allows developers to benchmark memory-saving techniques without requiring datacenter-grade GPUs. The project focuses on two primary memory bottlenecks: static model weights and dynamic KV-cache, providing runnable benchmarks for every reported performance metric.
2Lossless Weight Compression
The tool implements a 32-bit range coder with a 2-D adaptive context model to compress quantized LLM weights. On Qwen2.5-0.5B models, this approach achieves a 33% reduction compared to raw 4-bit weights. It consistently outperforms standard compressors like zstd-19 and xz-9e, reaching the 2-D entropy floor on real-world quantized tensors.
3KV-Cache Quantization Strategy
The research highlights that keys and values require different quantization axes for optimal performance. Keys benefit significantly from per-channel quantization to handle persistent outliers in activation dimensions. Conversely, values remain well-behaved under per-token quantization. This distinction is critical for maintaining model accuracy while aggressively reducing the memory footprint of the KV-cache.
4End-Task Perplexity Validation
Unlike proxy metrics, Vramkit validates compression quality using real end-task perplexity on diverse prose. Testing on Qwen2.5-0.5B shows that per-channel int8 keys achieve a 2x reduction in cache size with only a 0.1% increase in perplexity. This near-lossless result confirms that per-channel scaling is superior to the common per-token approach.
5Addressing RoPE Positioning
The study clarifies the impact of Rotary Positional Embeddings (RoPE) on quantization. Quantizing keys before RoPE versus after yields different results, with int8 per-channel keys holding stability in both positions. This finding ensures that the 2x compression ratio is applicable to standard transformer caches, providing a reliable drop-in solution for existing inference pipelines.
6Device-Agnostic Implementation
Vramkit is built as a device-agnostic module using plain PyTorch and NumPy. This architecture ensures compatibility across various hardware backends, including AMD ROCm, NVIDIA CUDA, Apple MPS, and standard CPUs. By providing a unified interface, it allows developers to apply the same quantization logic regardless of the underlying hardware acceleration stack.
7Impact on ROCm Ecosystem
Long-context inference remains a significant challenge for AMD ROCm environments, where KV-cache memory is often the primary bottleneck. By enabling a near-lossless 2x to 2.67x reduction in cache size, Vramkit helps close the performance gap with NVIDIA hardware. This allows for higher concurrent request capacity or longer context windows per GPU.
8Future Development Roadmap
The current implementation serves as a foundation for further optimization. Future goals include developing a grouped per-channel key path for streaming decode, similar to KIVI, and creating specialized Metal or ROCm dequantization kernels. These additions will allow compressed forms to be used directly during attention computation rather than just for storage.