Tailoring the Quantization Space for 1-Bit KV Cache Compression

• At rates near 1 bit per channel, each codebook must represent larger channel groups with limited centroids, making effective use of its capacity challenging.Quantization errors differ in their effects on attention logits, and VQ can exploit channel dependencies only when those channels share a codebook.Source: §1 · p. 1
• TaSQ combines query-guided channel weighting, covariance-aware grouping, and cross-head normalization while retaining conventional VQ lookup with negligible runtime overhead.Its channel-wise scaling and permutation avoid dense rotation; transforms are mostly absorbed into projection weights and codebooks, with runtime reconstruction fused into serving kernels.Source: §1 · p. 2
• TaSQ substantially improves accuracy over existing VQ baselines across general, reasoning, and long-context retrieval benchmarks in the 1-bit regime.Its SGLang implementation supports up to 14× larger batches and achieves up to 1.87× higher throughput than the BF16 baseline.Source: §1 · p. 2

2.1 LLM GENERATION AND KV CACHE

• The KV cache stores preceding tokens’ keys and values for reuse during decoding, but grows linearly with sequence length and is repeatedly accessed.As contexts lengthen, this places increasing pressure on memory capacity and bandwidth.Source: §2.1 · p. 2

2.2 VECTOR QUANTIZATION

• Vector quantization replaces a multidimensional input with the index of its nearest codeword from a finite codebook.Codebooks are typically learned with k-means on representative samples, using cluster centroids as codewords.Source: §2.2 · p. 2
• Representing a g-dimensional vector with one of K codewords costs log2 K/g bits per channel, and joint channel representation can exploit distributional structure.This makes VQ particularly attractive for ultra-low-bit compression.Source: §2.2 · p. 2

3 MOTIVATION

• Key channels vary in query sensitivity, have non-uniform inter-channel correlations, and undergo substantial distributional changes after RoPE.Compared with keys, values show more uniform channel-wise loss sensitivity within each head and weaker average inter-channel correlations, so value-tailored VQ is left to future work.Source: §3 · p. 3

3.1 SENSITIVITY OF KEY CHANNELS TO QUERIES

• Query activations vary substantially across channels, so key-channel quantization errors affect attention scores unequally.Each key-channel error is scaled by its corresponding query activation in the query–key dot product.Source: §3.1 · p. 3

3.2 INTER-CHANNEL CORRELATION

• Some key-channel pairs are strongly correlated while others are nearly independent, making VQ’s ability to exploit dependencies depend on group membership.Source: §3.2 · p. 3

3.3 PRE-ROPE VS. POST-ROPE KEYS

• RoPE spreads key distributions, and pre-RoPE VQ has 35% lower total reconstruction error across all 32 layers, motivating quantization of pre-RoPE keys.Source: §3.3 · p. 3

4 METHOD

• Calibration applies query-guided weighting, normalization, and RoPE-pair-preserving covariance grouping to pre-RoPE keys before fitting a VQ codebook for each group.At inference, keys are encoded as codebook indices and a scale, then reconstructed for attention with RoPE applied on the fly.Source: §4 · p. 3, Figure 2 · p. 4, §4.4 · p. 6

4.1 QUERY-GUIDED CHANNEL WEIGHTING

• Query-guided weighting derives channel weights from query–key dot-product error, so Euclidean VQ in weighted space directly incorporates query sensitivity.The weighting consistently lowers QK⊤ error and attention KL divergence across layers.Source: §4.1 · pp. 3–4

4.2 CROSS-HEAD SHARED-SCALE NORMALIZATION

• A single RMS scale shared across all H KV heads handles token-level magnitude outliers at a factor-H lower scale cost than head-wise scaling.The shared scale is stored in FP16.Source: §4.2 · p. 5

4.3 COVARIANCE-AWARE CHANNEL GROUPING

• The covariance-determinant proxy motivates grouping, and hierarchical matching merges RoPE-pair groups until each group contains g channels.The partition keeps each RoPE pair together for efficient online RoPE application.Source: §4.3 · p. 5
• Covariance-aware grouping’s cost tracks VQ reconstruction MSE (r = 0.994), and optimizing it lowers both the grouping cost and reconstruction error versus contiguous grouping.The reported correlation is for a representative layer; correlations are consistently strong across layers.Source: §4.3 · p. 6

4.4 VECTOR QUANTIZATION IN THE TAILORED QUANTIZATION SPACE

• TaSQ fits a separate K-centroid codebook for each layer, KV head, and channel group, using a Fisher-weighted k-means objective.At inference, each group is encoded by its nearest codeword.Source: §4.4 · p. 6
• Offline projection-weight and codebook transformations let the attention kernel fuse lookup, scale restoration, pairwise RoPE, and query–key dot products without extra shuffling, dense transforms, or kernel launches.The permutation preserves complete RoPE pairs, allowing standard block-diagonal RoPE structure with reordered frequencies.Source: §4.4 · p. 6
• Values retain standard contiguous-group VQ because their channel-wise characteristics are less pronounced than keys’.Further value-side analysis and alternatives are provided in Appendix C.Source: §4.4 · p. 6

5.1 EXPERIMENTAL SETUP

• Evaluation covers five models, with general, reasoning, and retrieval tasks assigned to distinct model groups.General evaluation uses Llama-3.1-8B-Instruct and Qwen3-4B; reasoning uses three other models, and retrieval uses three models.Source: §5.1 · p. 7
• Benchmarks span math, coding, general knowledge, reasoning, and needle-in-a-haystack retrieval, with retrieval tested from 4k to 64k context.Reasoning uses three sampling seeds; retrieval uses 200 examples and three seeds per setting.Source: §5.1 · p. 7
• CQ, NovaKV, and TaSQ share 64 calibration windows, while comparisons include pre-RoPE CQ and post-RoPE NSNQuant and NovaKV.The implementation uses SGLang with Triton attention kernels; efficiency experiments use RTX 6000 Ada GPUs.Source: §5.1 · pp. 7–8

5.2 EXPERIMENT RESULTS

• TaSQ leads existing VQ methods on general tasks, preserves more BF16 performance on long-CoT reasoning, and degrades less as retrieval context grows.The retrieval evaluation reports mean and standard deviation over three seeds at increasing context lengths.Source: §5.2 · p. 8, Table 3 · p. 8
• Figure 5 frames GSM8K and MBPP rate–accuracy comparisons for Llama-3.1-8B-Instruct and Qwen3-4B, normalized to each model’s BF16 score.The caption notes that values remain in BF16.Source: Figure 5 · p. 9

6.1 EFFICIENCY ANALYSIS

• 1.87× higher peak throughput than BF16, with peak throughput rising from 220.8 to 412.6 tokens/s on one RTX 6000 Ada GPU.TaSQ expands the KV cache pool from 235,257 to 2,935,184 tokens and maximum batch size from 6 to 84; TTFT rises 10% to 14% for 8k–32k prompts, amortized over long generations.Source: §6.1 · p. 9
• TaSQ keeps output lengths and cap-hit rates close to BF16, preserving reasoning stability while retaining low-bit cache benefits.Avoiding cap hits improves accuracy and efficiency because unterminated generations mostly fail while consuming the full decoding budget.Source: §6.1 · p. 9

6.2 ABLATION STUDY

• Removing grouping causes the largest perplexity increase among individual component removals, suggesting grouping contributes most to quantization quality.The ablation uses WikiText-2 with Llama-3.1-8B at context length 2,048, without retaining any KV cache in full precision.Source: §6.2 · p. 9
• TaSQ achieves the strongest rate–accuracy trade-off across both models, with a larger advantage at lower key rates.Values remain in BF16 for all methods while only key rates vary, because NovaKV cannot sweep value rates below one bit.Source: §6.2 · p. 9

7 CONCLUSION

• TaSQ combines query-guided weighting, cross-head normalization, and covariance-aware grouping for 1-bit KV-cache VQ while retaining conventional VQ serving efficiency.It outperforms existing VQ baselines across general, reasoning, and long-context retrieval tasks while preserving generation stability.Source: §7 · p. 10
• TaSQ supports up to 14× larger batches and up to 1.87× higher throughput than BF16 in the SGLang implementation.Source: §7 · p. 10

A RELATED WORK

• Scalar quantization treats weights independently, while VQ jointly encodes weight blocks and offers more favorable ultra-low-bit rate–distortion tradeoffs (Frantar et al., 2022; Lin et al., 2024; Xiao et al., 2023; Van Baalen et al., 2024; Egiazarian et al., 2024; Tseng et al., 2024).Source: A RELATED WORK · p. 13
• KV-cache quantization includes scalar methods (Liu et al., 2024; Hooper et al., 2024; Duanmu et al., 2024) and VQ methods (Zhang et al., 2024; Son et al., 2026); NovaKV uses a hybrid design (Fern´andez-Mendui˜na et al., 2026).Source: A RELATED WORK · p. 13
• TaSQ belongs to the KV-cache VQ line and studies the native-basis pre-RoPE target space.Source: A RELATED WORK · p. 13

Baseline configurations and bit accounting

• Effective cache rates include codebook-index and per-token metadata costs, but exclude static codebooks and shared full-precision prefix and recent windows.All evaluated models use head dimension 128; most have 8 KV heads, with Phi4-14B-Reasoning-Plus using 10.Source: B BASELINE CONFIGURATIONS AND BIT… · p. 13, Table 5 · p. 13
• For the 1-bit comparison, CQ and NSNQuant use their published configurations, while NovaKV retains its original designs with 1,024 key centroids.The 1,024-centroid key codebook matches CQ and TaSQ.Source: B BASELINE CONFIGURATIONS AND BIT… · p. 13
• The rate–accuracy sweep varies key codebook sizes by method, keeps values in BF16, and includes method-specific metadata in reported key rates.CQ, NovaKV, and TaSQ use K = 64–1,024; NSNQuant uses K = 32–512 to cover a comparable effective-rate range.Source: B BASELINE CONFIGURATIONS AND BIT… · p. 13
• Figure 6 summarizes channel sensitivity variation and inter-channel correlation across layers, with both metrics averaged over the eight KV heads.Lower coefficient of variation indicates flatter sensitivity; the correlation metric excludes the diagonal.Source: Figure 6 · p. 14

Value quantization analysis

• Values have lower average channel-sensitivity CV than keys, 0.176 versus 0.841, indicating more uniform sensitivity across channels.The paper argues this makes the original value space better aligned with Euclidean VQ’s uniform error metric.Source: C VALUE QUANTIZATION ANALYSIS · p. 14
• Values also have weaker mean absolute inter-channel correlation than keys, 0.067 versus 0.136, motivating the paper’s focus on tailoring key quantization.The analysis reports this pattern across all layers and leaves value-space tailoring to future work.Source: C VALUE QUANTIZATION ANALYSIS · p. 14
• Value-side sensitivity weights use diag(W_O,h^T W_O,h), since attention weights scale value channels equally and the output projection determines their relative error importance.The value reconstruction error contributes to the attention output through the head’s output-projection block.Source: C VALUE QUANTIZATION ANALYSIS · p. 14
• Combining value weighting, cross-head normalization, and covariance-aware grouping reduces WikiText-2 perplexity from 8.45594 to 8.42058, a 0.42% gain, but the final method keeps plain contiguous VQ for values.The key path is identical in both Table 6 rows; the paper cites the modest gain and added scale bits as reasons not to retain these value transforms.Source: C VALUE QUANTIZATION ANALYSIS · pp. 14–15, Table 6 · p. 15

Covariance-aware grouping and shared scale

• Cross-head normalization lowers scale metadata from bs/d to bs/(Hd), allowing the saved bit budget to be reassigned to codebook indices and more centroids.The method approximately solves its equal-size, RoPE-pair-preserving partition problem through hierarchical minimum-weight matching, which does not guarantee a globally optimal final partition.Source: D METHOD DETAILS AND ANALYSIS · p. 15

Cross-head shared scale

• Pooling saves 0.109 bits per channel at equal K with relative SSE within ±4%; pooled K = 2048 spends 0.016 bits more than per-head K = 1024 and cuts SSE at least 20% in all 68 layers.The relative SSE comparison uses separately fitted codebooks across layers.Source: Table 7 · p. 15, D METHOD DETAILS AND ANALYSIS · pp. 15–16

Generalization beyond calibration length

• Beyond the 2,048-token calibration length, TaSQ maintains the lowest attention-logit NMSE across positions; NovaKV’s error rises sharply outside the calibration range.The evaluation uses eight 16K WikiText-2 sequences per model; the trend matches the reported long-context retrieval comparison.Source: Figure 7 · p. 16, D METHOD DETAILS AND ANALYSIS · p. 16

Diagonal approximation and restoration cost

• Dense key restoration overhead grows from 36% of attention time at 1k tokens to 103% at 16k tokens.The table measures standalone BF16 GEMM restoration over reconstructed keys, excluding VQ decoding, on one RTX 3090 at batch size 1.Source: Table 8 · p. 16, D METHOD DETAILS AND ANALYSIS · p. 16

Diagonal approximation and weighting accuracy

• Dense weighting slightly lowers calibration-corpus error but yields about 4× higher error on GSM8K; the diagonal approximation has lower GSM8K error.Both variants use the same calibration data and K = 1024 codebooks; the paper reports the dense variant’s off-diagonal couplings do not transfer in this experiment.Source: Table 9 · p. 17, D METHOD DETAILS AND ANALYSIS · p. 17