Source-linked AI summary
PN-Mix: Pair-Norm Guided Mixed-Precision Key Cache Quantization for Efficient Long-Output LLM Inference
Gil Jae Song, Sungjoo Yoo
TL;DR
Long-output inference faces KV-cache memory and bandwidth bottlenecks, while low-bit quantization can degrade accuracy and prior mixed-precision methods estimate saliency from keys alone. PN-Mix uses attention-logit perturbation analysis to select salient RoPE channel pairs during prefill and reuse them during decoding. Across reasoning, code-generation, and AIME benchmarks, it achieves comparable or better average accuracy than prior mixed-precision baselines at lower average key-cache precision.
Problem
Prior mixed-precision key-cache methods estimate channel saliency from key statistics alone, despite the KV-cache bottleneck in long-output inference.
Source: §1 · pp. 1–2Method
PN-Mix uses attention-logit perturbation analysis to identify salient RoPE channel pairs from query–key pair-norm products and selects them during prefill for reuse in decoding.
Source: §1 · p. 2Results
Across reasoning, code-generation, and AIME benchmarks, PN-Mix achieves comparable or better average accuracy than prior mixed-precision baselines with lower average key-cache precision.
Source: §7 · p. 9Takeaways & Limitations
PN-Mix remains robust in the sub-2.25-bit regime and expands the feasible serving region for long-output inference.
Source: §7 · p. 9PN-Mix is most directly applicable to RoPE-based decoder-only LLMs; extending it to models without RoPE-style channel pairs may require adaptation.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
Long-output inference in large language models suffers from severe KV cache memory and bandwidth bottlenecks. While KV cache quantization alleviates this issue, low-bit quantization (e.g., 2-bit) severely degrades model accuracy. Recent mixed-precision key-cache quantization methods mitigate this degradation by allocating higher precision to important key channels, but estimate channel saliency using only key statistics. Through attention logit perturbation analysis at the RoPE channel-pair level, we show that quantization errors in channel pairs with large query--key pair-norm products contribute significantly to logit variation. Based on this insight, we propose PN-Mix, a mixed-precision key-cache quantization method for RoPE-based decoder-only LLMs, which allocates higher precision to salient channel pairs identified once during the prefill stage and reused throughout decoding. Across standard reasoning, code-generation, and AIME benchmarks, PN-Mix with 2.26-bit average precision achieves comparable or better average accuracy than prior mixed-precision methods using 2.32--2.39 bits. Furthermore, in our Pareto sweep, PN-Mix degrades gracefully down to 1.5 bits, whereas the compared baselines exhibit sharper performance drops below 2.25 bits.
1 Introduction
Long-output inference is constrained by KV-cache growth, while aggressive low-bit quantization can harm reasoning accuracy. PN-Mix guides key-cache precision using query-aware RoPE channel-pair saliency and selects salient pairs during prefill for reuse throughout decoding.
1 Introduction
- KV-cache growth increases memory pressure, reducing feasible serving batch size and limiting maximum generation length.The cache grows linearly with sequence length and batch size.Source: §1 · p. 1
- Aggressive low-bit KV-cache quantization, such as 2-bit quantization, can substantially degrade reasoning accuracy over long chains of thought.The limited bit budget therefore needs to be allocated to critical components.Source: §1 · p. 1
- Prior mixed-precision methods estimate saliency from key statistics alone, overlooking query–key interactions and RoPE channel-pair structure.PN-Mix addresses both omissions with a query-aware, pair-level saliency criterion.Source: §1 · pp. 1–2
- PN-Mix assigns higher precision to RoPE channel pairs with large query–key pair-norm products, selected once during prefill for reuse throughout decoding.The criterion comes from attention-logit perturbation analysis and avoids repeated decoding-time saliency recomputation.Source: §1 · p. 2
2 Background
The background describes attention with RoPE and the storage of past keys and values in the KV cache. It then outlines grouped quantization and the residual-buffer workflow used during online decoding.
2 Background
- Grouped Query Attention projects hidden states into query heads and key/value heads, while RoPE applies position-dependent rotations over adjacent channel pairs.Causal attention at each position is computed over preceding positions.Source: §2 · p. 2
- During autoregressive decoding, previously computed post-RoPE keys and values are stored and reused as the KV cache, whose memory scales as O(BLHkvTd).Here B is batch size, L is the number of layers, Hkv is the number of KV heads, T is sequence length, and d is head dimension.Source: §2 · p. 2
- The setup uses asymmetric group quantization, which reconstructs grouped values from quantized values, a scale, and a zero-point.The passage specifies b-bit quantization for a group of size G.Source: §2 · p. 2
- Keys are grouped per channel and values per token; newly generated entries remain in an FP16 residual buffer until it fills, then are quantized and appended.Attention combines the dequantized cache with the FP16 residual buffer.Source: §2 · p. 2
3 Related Work
Related work spans several KV-cache reduction strategies, including quantization. The review contrasts uniform-precision approaches with mixed-precision allocation and identifies how PN-Mix differs from prior channel-level methods.
3 Related Work
- KV-cache quantization reduces memory by lowering cached-tensor precision while retaining cache structure, alongside token eviction, low-rank approximation, merging, and pruning as other reduction directions.Quantization methods are grouped broadly into uniform-precision and mixed-precision approaches.Source: §3 · p. 2
3.1 Uniform-Precision Quantization
- Uniform-precision methods use one bit-width for all quantized elements and reduce error through grouping, distribution-aware schemes, residual correction, preservation, rotations, or specialized quantization.Because precision is not adjusted to element sensitivity, important-component errors may accumulate over long reasoning chains.Source: §3.1 · p. 3
3.2 Mixed-Precision Quantization
- Mixed-precision methods allocate bits by estimated importance; recent channel-level methods use key-channel saliency, with ChanMix relying on offline calibration and Kitty re-estimating average key magnitude by group.PN-Mix instead uses query–key interactions at the RoPE channel-pair level for saliency and precision allocation.Source: §3.2 · p. 3
4 Preliminary Analysis
Section 4 examines cache sensitivity, attention-aware importance scoring, and efficient score estimation, using these analyses to guide PN-Mix’s design.
4.1 Sensitivity of Key and Value Caches to Quantization
- V-only quantization causes less degradation than K-only quantization, and preserving the first 32 tokens in FP16 improves both.V-only quantization approaches the FP16 baseline, while K-only quantization still shows non-negligible degradation; the analysis motivates preserving sink tokens and refining key quantization more carefully.Source: §4.1 · p. 3
4.2 Designing an Attention-Aware Importance Score
- On both models, the proposed QK pair norm achieves the lowest D_A KL among mixed-precision criteria at the same effective ~2.25-bit budget.Its improvement over the K-pair norm indicates additional benefit from query-side information.Source: Figure 2 · p. 4, §4.2 · p. 5
- At the RoPE channel-pair level, query–key pair-norm products scale how quantization perturbations contribute to attention-logit error.The unquantized pair-norm product is used as a pre-quantization proxy because exact perturbation terms depend on the reconstructed key.Source: §4.2 · p. 4
- The separable score largely preserves the salient-pair ordering of the causal-joint score while avoiding explicit computation over causal query–key pairs.It summarizes query and key pair norms independently for channel-pair ranking.Source: §4.2 · p. 4
4.3 Efficient Importance Estimation for Long-Output Inference
- Prefill-based selection yields D_A KL close to the full-sequence counterpart on both models, the highest average Jaccard overlap with decoding-time salient pairs, and no systematic overlap decline during decoding.The results support computing the score once during prefill and reusing selected pairs throughout decoding.Source: §4.3 · pp. 5–6
- Approximately 97–99% of total importance is concentrated in the top 50% of channel pairs, using medians across layers, KV heads, and prompts.This selectivity supports allocating higher precision to a subset of pairs.Source: §4.3 · p. 5
5 Methodology
PN-Mix assigns mixed precision to key-cache channel pairs using their importance scores, while retaining sink tokens and the residual buffer in FP16.
5 Methodology
- PN-Mix assigns key-cache channel-pair precision according to σh,i while keeping sink tokens and the residual buffer in FP16.The scheme integrates the importance score from §4.2 into the online KV-cache quantization setup.Source: §5 · p. 6
5.1 Prefill-Stage Salient Channel-Pair Prediction
- During prefill, PN-Mix computes σh,i from prefill Q and K states for each layer, KV head, and channel pair, averaging query-side scores across query heads that share a KV head under GQA.For each KV head, it selects the top-r% pairs and stores the selection as a binary mask.Source: §5.1 · p. 6
5.2 Online KV-Cache Quantization Scheme
- During decoding, sink tokens and the residual buffer stay in FP16; when the buffer fills, selected key pairs use bH bits and remaining pairs use bL bits, while values use bV-bit per-token quantization.Keys use per-channel asymmetric groups; 1-bit key groups receive a midpoint correction to reduce extreme min–max quantization bias.Source: §5.2 · p. 6
6 Experiments
The experiments compare PN-Mix with uniform- and mixed-precision baselines across reasoning, code-generation, and AIME tasks, then examine design choices, low-bit behavior, and inference efficiency. PN-Mix maintains competitive accuracy at low key-cache precision and expands feasible serving capacity for long outputs.
6.1 Experimental Setup
- Evaluation spans standard reasoning and code-generation benchmarks and AIME, comparing KIVI, PolarQuant, ChanMix, and Kitty; PN-Mix tests 2.26- and 2.29-bit configurations.The baselines are set to 2.25–2.4 average key-cache bits for accuracy comparisons.Source: §6.1 · p. 6, §6.2 · p. 7
6.2 Downstream Task Accuracy
- At 2.26 bits, PN-Mix (1) stays within one FP16 Avg. point on all four standard models and is comparable to or better than mixed-precision baselines; it leads in five of six AIME settings despite fewer key bits.The AIME pattern also holds when the value cache is quantized to 2 bits.Source: Table 2 · p. 7, §6.2 · pp. 7–8, Table 3 · p. 8
6.3 Ablation Study
- Removing sink-token preservation or the 1-bit correction reduces accuracy, with removing the correction causing the larger drop.The ablations use PN-Mix (2) with FP16 values on Qwen3-8B.Source: Table 4 · p. 8, §6.3 · p. 8
- Replacing the query–key pair norm with a key-only criterion lowers mean accuracy, especially on AIME24; online reselection underperforms and requires group-specific pair-index storage.The online variant accumulates pair-norm statistics during decoding and reselects salient pairs for each quantization group.Source: §6.3 · p. 8
- Below 2.25 bits, PN-Mix degrades more gracefully than Kitty in the AIME24 sweep, extending its advantage beyond the 2.25-bit operating point.The sweep covers 1.5–3.0 bits on Qwen3-8B and R1-Llama-8B.Source: §6.3 · pp. 8–9, Figure 3 · p. 9
6.4 Inference Efficiency
- 32K generation at batch size 32 and batch size 110 at 8K fit PN-Mix’s 40GB memory budget, respectively 8× longer and 5.5× larger than FP16.The larger feasible batch range raises maximum per-GPU throughput to 2.37× FP16’s maximum.Source: §6.4 · p. 9
- For long prompts, PN-Mix’s TTFT differs from Kitty by at most 12 ms; decode time is 4.0–6.5% lower across the three batch-size-one settings.On the long-CoT trajectory, marginal per-step latency growth is 8.8% lower than Kitty’s.Source: §6.4 · p. 9
- At batch size 32 with a 141-token prompt and 32K generation, only PN-Mix fits in memory, using 37.12 GB.The 141-token prompt corresponds to a single AIME24 problem.Source: §6.4 · p. 9, Table 5 · p. 9
7 Conclusion
PN-Mix guides mixed-precision key-cache quantization by RoPE channel-pair query–key pair norms, selecting salient pairs once during prefill. Across the reported benchmarks, it delivers competitive accuracy at lower average key-cache precision and remains robust below 2.25 bits.
7 Conclusion
- PN-Mix assigns higher precision to RoPE channel pairs with large query–key pair-norm products, identifying salient pairs during prefill without repeated decoding-time saliency recomputation.The method is designed for RoPE-based decoder-only LLMs.Source: §7 · p. 9
- Across standard reasoning, code-generation, and AIME benchmarks, PN-Mix achieves comparable or better average accuracy than prior mixed-precision baselines with lower average key-cache precision.The reported 2.26-bit average-precision setting is the headline comparison in the paper context.Source: §7 · p. 9
- PN-Mix remains robust below 2.25 bits and expands the feasible serving region for long-output inference.Source: §7 · p. 9
Limitations
The analysis finds that low-bit key-cache quantization is more harmful than value-cache quantization, while prefill-based query–key pair-norm scores closely match full-sequence scores. PN-Mix is scoped to RoPE-based decoder-only LLMs, and measured speedups may vary across deployments.
Limitations
- PN-Mix is most directly applicable to RoPE-based decoder-only LLMs, and extending it to models without RoPE-style channel pairs may require adaptation.Efficiency results use a single A100 GPU and a Triton-based implementation; exact speedups may vary with hardware, kernel optimization, batching, and serving infrastructure.Source: Limitations · pp. 9–10
A Analysis on Key- and Value-Cache Quantization Sensitivity
- Value-only quantization is less harmful than key-only quantization at the same average bit-width across both models and AIME benchmarks.The difference is especially pronounced when the first 32 tokens are preserved in FP16.Source: A Analysis on Key- and Value-Cache… · p. 12
B Key-Cache Quantization Simulation Details
- The query–key pair-norm criterion yields the lowest attention KL divergence among mixed-precision criteria, and its prefill-only version closely matches full-sequence scoring.The simulation compares uniform baselines and mixed-precision settings that differ by importance score and whether they use full-sequence or prefill-only importance.Source: B Key-Cache Quantization Simulation… · p. 12
C Salient Channel-Pair Prediction Accuracy
Query–key pair-norm scores best match decoding-time salient channel pairs, with prefill-only predictions nearly matching full-sequence predictions. Pair-norm importance is concentrated in a small fraction of channel pairs, including at larger model scales.
C Salient Channel-Pair Prediction Accuracy
- Across both models and selection ratios, query–key pair norm has the highest Jaccard similarity to decoding-time salient-pair ground truth, followed by key pair norm, key magnitude, and key range.Its advantage is larger at r = 10%, where query-side information is particularly useful for identifying the most salient pairs.Source: C Salient Channel-Pair Prediction… · p. 13
- For query–key pair norm, prefill-only predictions yield nearly the same Jaccard similarity as full-sequence predictions, suggesting decoding states add little prediction benefit.Prediction accuracy is evaluated against decoding-time salient sets using Jaccard similarity.Source: C Salient Channel-Pair Prediction… · p. 13
D Concentration of Pair-Norm Importance Scores
- On Qwen3-14B, the top 10%, 30%, and 50% channel pairs capture 89.3%, 97.0%, and 98.7% of total importance; on Qwen3-32B, they capture 86.0%, 94.8%, and 97.4%.Figure 5 reports median cumulative ratios computed from prefill-based scores across layers, KV heads, and prompts.Source: D Concentration of Pair-Norm… · pp. 13–14, Figure 5 · p. 13
E Memory Scaling and Throughput Analysis
On Qwen3-8B, PN-Mix reduces memory use enough to support longer generations and larger batches than FP16, while its throughput advantage appears when memory limits FP16. It also reduces decode-latency growth relative to Kitty, with Triton kernels handling quantization and attention operations.
E Memory Scaling and Throughput Analysis
- At batch size 32, PN-Mix remains within 40GB through 32K generation, versus FP16 exceeding the budget beyond 4K; at 8K generation, PN-Mix supports batch size 110 versus FP16’s 20.Both quantized methods reduce memory growth, with PN-Mix using the least memory in the compared settings.Source: E Memory Scaling and Throughput Analysis · p. 14
- At batch size 20, FP16 reaches 469.8 tok/s versus PN-Mix’s 420.8 tok/s, but PN-Mix reaches 1114.0 tok/s at batch size 110, or 2.37× FP16’s maximum feasible throughput.FP16 runs out of memory at batch size 21; PN-Mix exceeds its maximum feasible throughput at batch size 23 and reaches 38.21 GB at batch size 110.Source: E Memory Scaling and Throughput Analysis · pp. 14–15
F Per-Step Decoding Latency Analysis
- Over decoding steps 11K–31K, PN-Mix’s latency-growth slope is 2.17 versus Kitty’s 2.38 µs/step, an 8.8% reduction; at 32K steps it is 6.8 ms faster.Both methods stay near a 45–48 ms plateau through approximately 10K steps on the B = 1, P = 141, Lgen = 32K trajectory.Source: F Per-Step Decoding Latency Analysis · p. 15, Figure 7 · p. 15
G Kernel Implementation
- The Triton implementation packs low bits for all key channels and upper bits only for salient pairs, then dequantizes cache groups on-chip immediately before dot-product computation.It uses separate decoding kernels for query–key computation and attention-weighted value computation.Source: G Kernel Implementation · p. 15
H Average Key Bit-Width Calculation and Method Configurations
The bit-width accounting includes quantized key payload and method-specific metadata, so comparisons use a narrow low-bit range rather than exact bit matches or each method’s recommended operating point. The section also specifies metadata, configuration, and profiling choices for the evaluated methods and Pareto sweep.
H Average Key Bit-Width Calculation and Method Configurations
- The reported comparison targets accuracy preservation in a 2.25–2.4-bit regime, not exact bit matching or each method’s best-performing or recommended operating point.Table 10’s configurations span approximately 2.25–2.39 average key bits.Source: H Average Key Bit-Width Calculation… · pp. 15, 17
- Average key-cache bit-width includes payload, applicable sink-token and scale/zero-point storage, and saliency metadata, while excluding the FP16 residual buffer from storage and element normalization.Uniform methods account for payload and scale/zero-point overhead; mixed-precision methods also account for their precision allocation.Source: H Average Key Bit-Width Calculation… · pp. 15–16
- Kitty stores channel-selection indices per quantization page, whereas PN-Mix selects RoPE pairs once during prefill and shares a mask across quantization groups, making its mask overhead negligible for long sequences.PN-Mix’s mask records one high-precision assignment bit per RoPE pair per KV head.Source: H Average Key Bit-Width Calculation… · p. 16
- ChanMix’s average precision accounts for promoted retrieval heads, scale and zero-point storage, and shared bit-width and channel-order metadata.Its effective promoted-head ratio uses the rounded number of promoted heads.Source: H Average Key Bit-Width Calculation… · pp. 16–17
- The 2-bit radius/angle PolarQuant setting is used for the low-bit comparison and may amplify error; its results are not comparisons against its original 3-bit operating point.The stated comparison evaluates accuracy preservation under comparable low key-cache budgets.Source: H Average Key Bit-Width Calculation… · p. 17
- For models without official ChanMix profiles, the authors reproduce its profiling procedure and report 0.98–0.99 Pearson correlations and 90–91% overlap with official profiles on three validation models.Results using reproduced profiles are marked with a dagger in Tables 2 and 3.Source: H Average Key Bit-Width Calculation… · p. 17
I Configurations for the Bit-Rate Pareto Frontier
- The Pareto sweep selects configurations nearest target budgets of 3.0, 2.5, 2.25, 2.0, and 1.5 bits; below 2.25 bits, Kitty also sweeps high and low precisions because its original setting cannot reach those budgets.For both Kitty and PN-Mix, the first 32 sink tokens remain full precision and the value cache stays at 16 bits.Source: I Configurations for the Bit-Rate… · p. 17
J Additional Evaluation on Long-Input and Multilingual Tasks
PN-Mix performs strongly across long-input and multilingual evaluations, compares favorably with PM-KVQ at lower average KV-bit budgets, and retains stable prefill-to-decoding saliency overlap in tested stress settings.
J Additional Evaluation on Long-Input and Multilingual Tasks
- PN-Mix achieved the highest average accuracy among compressed methods on RULER at every tested context length, remaining close to FP16 even at 64K.On the challenging MK3 task, its accuracy stayed nearly unchanged as context grew while competing quantized methods degraded substantially.Source: J Additional Evaluation on Long-Input… · p. 17
- 33.98 vs 34.15: PN-Mix (2) achieved the highest compressed-method average score on InfiniteBench, close to FP16 across diverse long-context workloads.The evaluation covered retrieval, dialogue QA, long-document QA, and summarization.Source: J Additional Evaluation on Long-Input… · pp. 17–18
- PN-Mix (2) achieved the best compressed-model accuracy on MGSM in Chinese, French, and Japanese, and remained competitive across all evaluated languages.Source: J Additional Evaluation on Long-Input… · p. 18
K Additional Comparison with PM-KVQ
- PN-Mix matched PM-KVQ’s accuracy on R1-Llama-8B and had higher mean AIME24 and AIME25 accuracy on R1-Qwen-7B while using fewer average KV bits.The compared PM-KVQ configurations used approximately 2.38 and 2.32 average KV bits, versus 2.26–2.28 for PN-Mix.Source: K Additional Comparison with PM-KVQ · p. 18
L Step-Wise Stability of Prefill-Based Saliency Selection
- Across 32,672-token generation and a 60-question multi-topic test, prefill-selected pairs’ overlap stayed stable; pair-norm scores were approximately 0.72 and 0.75, respectively.The long-generation setting used 96 prefill tokens, and the multi-topic questions spanned 57 MMLU subjects.Source: L Step-Wise Stability of Prefill-Based… · p. 19
- The pair-norm criterion had the highest overlap among compared selection criteria in both stress tests, with similar stability at r = 10% and r = 50%.Source: L Step-Wise Stability of Prefill-Based… · p. 19