Dual-qk reduces memory use and speeds up long context language models
Dual-QK: Sharp Queries and Flat Keys for Prunable 2-bit KV Caches
Machine Learning
Summary
Handling very long texts in AI language models requires storing and accessing lots of data, which can slow things down and use a lot of memory. The authors present Dual-QK, a technique that changes how the model processes and stores this data by using smart math tricks to reduce storage size and speed up access. This method balances how data is compressed and searched, allowing the model to run faster and use less memory without losing much accuracy. Tests showed Dual-QK works well across different tasks and large input sizes.
What this means in practice
- •For ai language model developers: Reduce memory usage and speed up decoding for models handling very long text inputs using 2-bit compressed key-value caches with query pruning.
- •For cloud ai service operators: Lower infrastructure costs by reducing memory bandwidth and storage needs during long-context language model serving without major accuracy loss.
Authors
Sunjoo Whang, Jungjun Oh, Minsung Kim, Dongho Seo, Jisu Shin, Gregory Kielian, Hoi-Jun Yoo, Sangjin Kim
Abstract
Long inputs and extended generation increase the storage and access costs of the key-value (KV) cache. Low-bit quantization reduces storage and memory traffic, while query-channel pruning can further reduce key-cache reads. Rotation-based quantization redistributes the energy of key outliers across channels. To maintain computational invariance, the same orthogonal transform must be applied to queries, preserving query-key dot products. However, this rotation can disperse query energy, weakening the separation between a few large components to retain and many small ones to prune. We introduce Dual-QK, which uses paired non-orthogonal query and key transforms to address this conflict. Using calibrated query and key statistics, Dual-QK combines partial key whitening with a query-aligned basis to balance key scales for INT2 quantization and concentrate query energy for dynamic channel pruning. Channel-0 protection and bucket-relative RoPE support low-bit accuracy over long contexts. Experiments on four models across five generative benchmarks and long-context retrieval tasks show improved accuracy over OSCAR on most tasks at 40% query-channel sparsity. At a 128K context, Dual-QK provides $6.8\times$ KV-cache compression and an estimated $8.3\times$ reduction in KV read volume relative to unpruned BF16. Under the evaluated configurations, our SGLang implementation achieves up to $3.75\times$ the decoding throughput of unpruned BF16.