Mingjie Lin

Author directory

2026

Key–value (KV) caching accelerates autoregressive decoding but makes long-context serving memory-intensive. Most KV quantizers minimize local reconstruction distortion, although this proxy ignores how perturbations propagate through downstream computation. We introduce Beyond, a loss-aware non-uniform quantizer that learns representation levels and decision thresholds from next-token negative log-likelihood while freezing the backbone. Training and inference use the same hard quantizer, with straight-through and boundary-localized pseudo-gradients updating only its parameters. With 4-bit codes and group size 32, Beyond uses 5.0 effective bits per KV scalar, providing 3.2× compression relative to a 16-bit cache. Across three instruction-tuned LLMs, it lowers held-out Pile perplexity below the full-precision baseline by 2.2–7.0% and attains the highest LongBench Avg21 among full-precision Native, KV-FP8, and TurboQuant for each model. On an NVIDIA B300, its fused packed-cache kernel achieves a peak 1.50× speedup over FlashAttention-4, while the resulting vLLM path improves end-to-end throughput by up to 1.17×. Our code is available at https://github.com/wsxhjnb1/Beyond-KV-Cache.