MultiAlloc: Budget-Aware Precision Allocation for KV Cache Quantization

What it claims

MultiAlloc treats KV cache quantization as one budget split between two competing uses of memory: raising the bit-width of whole layers, and keeping the top-k tokens per group at high precision. All decisions are made offline and frozen into a static policy π = (q, k). Here q is the per-layer bit-width and k is the number of high-precision tokens kept per group in each layer. Reported evaluations cover GSM8K, GPQA and LongBench-v2 on Llama-3-8B-Instruct, Mistral-7B-Instruct-v0.3, DeepSeek-R1-Distill-Qwen-7B/14B and QwQ-32B. Code is at https://github.com/RL-MIND/MultiAlloc.

How it works

• Layer sensitivity (offline). For each layer l and each candidate bit-width b (the figure shows 2, 4 and 8), the method takes the quantization error ΔK, ΔV on calibration samples. It weights that error by the loss gradients G^K, G^V, which is a first-order Taylor estimate, giving S_l(b) = E[‖ΔK⊙G^K‖₁ + ‖ΔV⊙G^V‖₁]. This needs backward passes through the model on calibration data.
• Token importance. A token's score is the L1 norm of its key vector, s_{l,t} = ‖K_{l,t}‖₁. Values are not used. Tokens are grouped into groups of size G, and the top-k_l tokens in each group stay at high precision (FP16 in the figure). The rest use the layer's assigned bit-width. The token "loss" for a given k_l is the importance mass left unprotected.
• Budget-coupled allocation. The method enumerates candidate token-capacity configurations o_i. For each one: 1. Subtract its token memory cost from the total budget B. 2. Solve for the best layer bit-widths within the remaining budget. 3. Score the result as J_i = L_layer + L_tok.

The candidate with the lowest J_i is chosen.

flowchart LR
  C[Calibration set] --> LS[Layer sensitivity S_l b]
  C --> TI[Token key-norm profile]
  TI --> O[Token-capacity candidates o_i]
  O --> R[Remaining budget B - M_tok]
  R --> Q[Best layer bits q_i]
  LS --> Q
  Q --> J[Score J_i]
  O --> J
  J -->|next candidate| O
  J -->|argmin| P[Static policy q*, k*]

What is new relative to the baselines the paper names: KVTuner and PM-KVQ assign mixed precision per layer but treat tokens uniformly within a layer. ZipCache and OTT protect tokens but ignore differences between layers. KVmix adapts at both levels, but its token side selects context rather than allocating precision.

Evidence

The fetched text stops partway through the method section, before any results tables. The only numbers are in the motivation figure (Figure 1), at a fixed KV budget:

Neither the model nor the benchmark for these figures is stated in the available text, so read them as illustrative. They are not a verified head-to-head result.

Limits

• The page is truncated. Results tables, ablations, latency and throughput figures, and any stated limitations are not in the fetched text.
• The available text does not give the calibration set size, the group size G, the candidate set sizes or the profiling cost.
• J_i adds a gradient-weighted error (layers) to an L1 key-norm mass (tokens). The text does not say how these two quantities are put on the same scale.
• The policy is static per model and per budget. A new budget or a new model means re-profiling.
• The method needs gradient access during profiling. This could be a constraint for very large models.
• The available text does not describe whether runtime top-k selection is supported by kernels or serving systems such as PagedAttention.