OScaR: La navaja de Occam para la cuantización extrema de caché KV en LLMs y más allá

Resumen

El rápido avance hacia el razonamiento de contexto largo y la inteligencia multimodal ha convertido la huella de memoria de la caché de Clave-Valor (KV) en un cuello de botella dominante para el despliegue eficiente. Si bien la cuantización por canal establecida se adapta eficazmente a los valores atípicos intrínsecos por canal en los tensores Clave, su eficacia disminuye bajo una compresión extrema. En este trabajo, revisitamos las limitaciones inherentes del paradigma de cuantización por canal desde perspectivas tanto empíricas como teóricas. Nuestro análisis identifica el Desequilibrio de Norma de Tokens (TNI) como el principal cuello de botella para la fidelidad de la cuantización. Demostramos que TNI amplifica sistemáticamente los errores cuando se requieren parámetros de cuantización compartidos para abarcar grupos de tokens que exhiben disparidades sustanciales de norma. En lugar de depender de complejos pipelines de cuantización (p. ej., TurboQuant), proponemos OScaR (Rotación Canalizada a Escala Omni), un framework ligero y preciso de compresión de caché KV para X-LLMs (es decir, LLMs solo de texto, multimodales y omni-modales). Avanzando el paradigma por canal, OScaR emplea Rotación Canalizada seguida de Escalado de Tokens Omni para mitigar la varianza dimensional de secuencia inducida por TNI tanto de forma efectiva como eficiente, respaldado además por nuestro diseño de sistema optimizado y núcleos CUDA. Evaluaciones exhaustivas en X-LLMs muestran que OScaR supera consistentemente a los métodos existentes y logra un rendimiento casi sin pérdidas bajo cuantización INT2, estableciéndose como un framework robusto, de baja complejidad y universal que define un nuevo frente de Pareto. En comparación con la línea base BF16 FlashDecoding-v2, nuestra implementación de OScaR logra una notable aceleración de hasta 3.0x en decodificación, reduce la huella de memoria en 5.3x y aumenta el rendimiento en 4.1x. El código de OScaR está disponible públicamente en https://github.com/ZunhaiSu/OScaR-KV-Quant.

English

The rapid advancement toward long-context reasoning and multi-modal intelligence has made the memory footprint of the Key-Value (KV) cache a dominant memory bottleneck for efficient deployment. While the established per-channel quantization effectively accommodates intrinsic channel-wise outliers in Key tensors, its efficacy diminishes under extreme compression. In this work, we revisit the inherent limitations of the per-channel quantization paradigm from both empirical and theoretical perspectives. Our analysis identifies Token Norm Imbalance (TNI) as the primary bottleneck to quantization fidelity. We demonstrate that TNI systematically amplifies errors when shared quantization parameters are required to span token groups exhibiting substantial norm disparities. Instead of relying on intricate quantization pipelines (e.g., TurboQuant), we propose OScaR (Omni-Scaled Canalized Rotation), an accurate and lightweight KV cache compression framework for X-LLMs (i.e., text-only, multi-modal, and omni-modal LLMs). Advancing the per-channel paradigm, OScaR employs Canalized Rotation followed by Omni-Token Scaling to mitigate TNI-induced sequence-dimensional variance both effectively and efficiently, further supported by our optimized system design and CUDA kernels. Extensive evaluations across X-LLMs show that OScaR consistently outperforms existing methods and achieves near-lossless performance under INT2 quantization, establishing it as a robust, low-complexity, and universal framework that defines a new Pareto front. Compared with the BF16 FlashDecoding-v2 baseline, our OScaR implementation achieves a notable up to 3.0x speedup in decoding, reduces memory footprint by 5.3x, and increases throughput by 4.1x. The code for OScaR is publicly available at https://github.com/ZunhaiSu/OScaR-KV-Quant.