RoLaQ: Joint Rotation and Range Calibration for 2-Bit KV Caches in Multi-Head Latent Attention

1 Introduction

• Because the shared latent feeds distinct key and value projections, its quantization errors can alter both attention matching and value readout.Uneven K/V projection energy and the effect of reconstruction gain motivate learning rotation and range together.Source: §1 · pp. 1–2
• RoLaQ jointly refines a shared rotation within K/V-defined subspaces and learns per-layer quantization thresholds with model weights frozen.Prediction KL aligns the model with a BF16 teacher, while attention-output supervision guides individual layers.Source: §1 · p. 2
• RoLaQ achieves the highest average accuracy among compared 2-bit methods across three MLA models and six benchmarks, with and without token protection.RULER results at 8K and 16K show calibration gains beyond the 4K training context.Source: §1 · p. 2
• 3.96× the BF16 KV capacity at equal memory budgets and 2.51× its decode throughput with 256K input tokens are reported for DeepSeek-V2-Lite.The serving measurements use SGLang on one A100 80GB.Source: Figure 1 · p. 1, §1 · p. 2

2.1 One Cache, Two Roles

• The cached content latent feeds distinct key and value projections, so its quantization error affects both attention matching and value readout.The lower-dimensional RoPE key is retained in BF16 because of its quantization sensitivity.Source: §2.1 · p. 2

2.2 Projection-Dependent Distortion

• 62–75% of each K/V projection’s energy lies in the leading half of the joint spectrum across the three models, averaged over layers.This shared eigenbasis separates latent directions by their influence on both projections and guides the rotation search.Source: §2.2 · p. 3

2.3 Sensitivity to Reconstruction Scale

• 52% lower next-token KL to BF16 was reached for DeepSeek-V2-Lite on 8K-token documents by adjusting gain alone.The sweep fixed OSCAR’s rotation and threshold; the observations motivate joint rotation and range calibration.Source: §2.3 · p. 3

3 Method

• ROLAQ jointly calibrates the shared representation and quantizer to preserve attention matching and value readout.The model weights remain frozen during calibration.Source: §3 · p. 3, §3.3 · p. 4

3.1 K/V-Aware Shared Rotation

• Rotation refinement stays within two K/V-defined spectral subspaces, preserving their joint projection-energy distinction while adapting the quantization basis.The orthogonal parameterization starts from OSCAR’s covariance-aware rotation and uses half as many trainable parameters as a full generator.Source: §3.1 · p. 3

3.2 Response-Preserving Range Calibration

• Each layer learns a threshold that assigns standardized coordinates to four levels, while its linked gain preserves unit response of the centered signal.The threshold controls inner versus outer assignments, and the gain scales reconstruction around the group mean.Source: §3.2 · p. 4

3.3 Joint Calibration and Inference

• A frozen BF16 teacher guides joint learning with next-token prediction KL and attention-output loss; inference stores rotated latents as 2-bit codes and folds read-side rotation into absorbed projections.The output loss supplies a direct signal to each layer through attention matching and value readout; rotation, thresholds, and gains stay fixed during serving.Source: §3.3 · pp. 4–5

4.1 Setup

• The evaluation covers three MLA models and six downstream benchmarks, plus eight RULER retrieval tasks; main results average three calibration seeds.All methods use 2-bit content caches and BF16 RoPE keys; protection retains the first 4 sink tokens and most recent 16 tokens in BF16.Source: §4.1 · pp. 5–6

4.2 Main Results

• 46.7, 60.1, and 59.4 are ROLAQ’s unprotected six-task averages on DeepSeek-V2-Lite, Moonlight, and MiniCPM3, beating the strongest baseline by 3.8, 2.8, and 1.1 points.It has the highest six-task average on all three models under both protection settings.Source: §4.2 · p. 6
• Retaining 4 sink and 16 recent tokens in BF16 further improves ROLAQ, which keeps the highest average accuracy on all three models.On DeepSeek-V2-Lite, the unprotected ROLAQ average of 46.7 exceeds the strongest baseline’s protected average of 45.3.Source: §4.2 · p. 6

4.3 Long-Context Retrieval

• At evaluated contexts beyond the 4K calibration length, ROLAQ reaches 73.7% on DeepSeek-V2-Lite retrieval, 6.3 points above the strongest baseline.On Moonlight it reaches 94.8%, versus BF16’s 96.2%; DeepSeek-V2-Lite MK2 accuracy rises from OSCAR’s 7.0 to 55.0.Source: §4.3 · p. 6, Table 2 · p. 6

5.1 Calibration Ablations

• Joint calibration gives the highest accuracy on both models; learning the range increases rotation calibration’s MBPP gain from 2.2 to 4.0 points on DeepSeek-V2-Lite.The results support adapting the quantizer jointly with the latent representation.Source: §5.1 · pp. 6–7

5.2 Rotation Search Space

• HK + HV uses half the trainable rotation parameters of full exponential and Cayley rotations while matching or improving six-task accuracy.It also reduces validation KL by 6.3% and 9.1% versus full exponential rotation on DeepSeek-V2-Lite and Moonlight, respectively.Source: §5.2 · p. 7, Table 4 · p. 7

5.3 Quantizer Response

• κ = 1 gives the lowest prediction KL across contexts from 2K to 32K, while amplification becomes more costly as context grows.The comparison varies inference gain while holding the calibrated rotation and thresholds fixed.Source: §5.3 · p. 7, Figure 4 · p. 7

5.4 Token Protection

• ROLAQ maintains the highest six-task mean across all three token-protection budgets, and with 4 sink plus 8 recent tokens exceeds the strongest baseline using 4 sink plus 16 recent tokens on each model.The same calibrated checkpoints are reused across protection budgets.Source: §5.4 · pp. 7–8, Figure 5 · p. 7

5.5 Serving Efficiency

• 3.96× BF16 KV capacity at the same memory budget and 2.51× BF16 decode throughput with 256K input tokens; at 256 concurrent requests, P99 TTFT falls 37.8%, from 588.6 s to 366.0 s.Measurements use DeepSeek-V2-Lite on one A100 80GB.Source: §5.5 · p. 8, Figure 6 · p. 8

• KV-cache methods include statistics-aware quantization, outlier handling, residual compensation, low-rank or sparse error correction, and higher precision for salient tokens.Some approaches retain recent tokens or attention sinks in full precision.Source: §6 · p. 8
• Rotation and calibration methods redistribute outliers or learn transformations and quantization ranges, while attention-aware approaches optimize cache representations for attention computations.Examples include OSCAR’s attention-aware covariance rotations and OptR’s attention-output reconstruction objective.Source: §6 · p. 8
• ROLAQ uses MLA’s K/V projection structure to guide joint calibration of the shared latent rotation and quantization range.Source: §6 · p. 8

• ROLAQ combines a compact K/V-geometry rotation search space with coupled thresholds and gains to control reconstruction response.Source: §7 · p. 9
• Experiments show higher average accuracy across three MLA models, better long-context retrieval, and further gains from token protection.Source: §7 · p. 9
• The SGLang implementation increases KV capacity and throughput while reducing latency at high concurrency.Source: §7 · p. 9

• Content-latent quantization leaves RoPE keys in BF16 and positional keys and group metadata as a substantial share of the cache; short-context overheads can offset bandwidth savings.Source: Limitations · p. 9

A.1 Models and Evaluation

• The evaluation spans two 16B mixture-of-experts models and a dense 4B model, using answer-selection, generation, and long-context retrieval benchmarks.The models are DeepSeek-V2-Lite, Moonlight, and MiniCPM3; RULER evaluates eight retrieval tasks as context grows.Source: §A.1 · p. 11

A.2 Baseline Implementations

• RTN, QuaRot, KVarN, and OSCAR are compared on the shared content latent, with BF16 RoPE keys and specified quantization groupings.RTN, QuaRot, and OSCAR group latent coordinates within each token, whereas KVarN groups tokens for channel-wise quantization.Source: §A.2 · p. 11

A.3 Calibration and Storage

• Calibration updates only rotation and thresholds with the model and teacher frozen, using unprotected tokens and reusing checkpoints across protection budgets.Checkpoints are selected by held-out prediction fidelity across short, 2K, and 4K contexts.Source: §A.3 · p. 12
• An unprotected token uses 288 bytes in DeepSeek-V2-Lite and Moonlight and 144 bytes in MiniCPM3 under the stated cache allocation.The accounting includes packed content codes, group statistics, and the BF16 RoPE key; protected tokens also store BF16 content.Source: §A.3 · p. 12

B Supplementary Method Analysis

• The supplementary analysis separates rotation-space choices from quantizer-response control.It frames these as distinct design dimensions for the method analysis.Source: B Supplementary Method Analysis · p. 12

B.1 Motivating Observations

• Projection energy and reconstruction gain are distinct factors: latent directions affect K/V projections unequally, while gain changes reconstruction strength at fixed rotation and assignments.The gain intervention isolates scale from direction.Source: §B.1 · p. 12

B.2 Rotation Search Spaces

• Rotation variants differ in which directions can mix and how the orthogonal update is parameterized; the random control keeps the block sizes but removes K/V-based grouping.The random control retains the spectral variants’ trainable-parameter count.Source: §B.2 · p. 13

B.3 Response of the Four-Level Quantizer

• Coupling gain to threshold anchors the reconstructed component aligned with the input at unit response, although minimizing Gaussian reconstruction error alone can favor a weaker signal.At the official threshold, the stated MSE-optimal response is approximately 0.88, whereas the coupled setting sets κ = 1.Source: §B.3 · p. 14

C.1 Token Protection and Calibration Seeds

• ROLAQ retains the highest mean among the compared 2-bit methods on all three models at the intermediate budget of 4 sink and 8 recent BF16 tokens.The reported comparison covers all six tasks.Source: §C.1 · p. 14

C.2 Range Calibration on MiniCPM3

• Coupling gain to threshold lowers MiniCPM3’s mean validation KL ratio from 0.606 to 0.527 and reduces variation across calibration seeds.This comparison holds the HK + HV search space fixed.Source: §C.2 · p. 14

C.3 Learned Thresholds

• Across rotations and seeds, most learned thresholds fall below the initial Lloyd–Max threshold, clustering around 0.91, 0.89, and 0.85 for the three models.The narrower central assignment region adapts quantization assignments while the coupled gain preserves unit response.Source: §C.3 · p. 14

C.4 Case Study

• Across all 1,319 GSM8K problems, ROLAQ solves 277 that OSCAR misses and misses 69 that OSCAR solves; examples show ROLAQ counts both bags and subtracts the 6 feet already available.Source: §C.4 · p. 14
• In the GSM8K examples, ROLAQ counts both bags to answer 22 chips per person and subtracts 6 available feet to answer 594.QuaRot and OSCAR miss the second bag in one example and return the unadjusted 600-foot requirement in the other.Source: §C.4 · p. 14, Table 13 · p. 17
• In the 16K RULER examples, ROLAQ retrieves the correct single-key value or complete four-key set, unlike QuaRot and OSCAR.The comparison uses the same protection budget and calibration seed as the GSM8K examples.Source: §C.4 · p. 14, Table 14 · p. 17

D Serving Benchmarks

• Serving tests use one A100 80GB GPU, a 0.8 static memory fraction, and the final DeepSeek-V2-Lite checkpoint with 2-bit content and selected BF16 tokens.BF16 and ROLAQ use Triton prefill and decode kernels under the same workload and memory budget.Source: D Serving Benchmarks · p. 18
• OSCAR’s Figure 6 timing curves reuse ROLAQ’s timing values because both use the same packed-cache layout and kernels.The allocation runs give the same KV capacity for both methods.Source: D Serving Benchmarks · p. 18
• Long-input tests sweep 1K–256K input tokens with 128 output tokens, while long-output tests generate 256K tokens from 128-token inputs; both use four concurrent requests.Decode throughput is measured after prefill for the input sweep and in successive output-token intervals for the output sweep.Source: D Serving Benchmarks · p. 18, Table 16 · p. 18
• The concurrency test measures P99 TTFT as simultaneous arrivals rise from 1 to 256, with each request using 16K input and 1K output tokens.The metric includes queueing and prefill; the benchmarks report one timed run per workload after warmup.Source: D Serving Benchmarks · p. 18, Table 17 · p. 18

E AI Assistance

• The authors report using AI assistants for code development, manuscript organization, and language editing, and take responsibility for the final work.Source: E AI Assistance · p. 18