Efficient AI · Undergraduate thesis

Behavior-Guided
KV-Cache Quantization

Reducing long-context cache storage while preserving attention behavior.

Project period: Jan 2026 – May 2026
Independent thesis · SCUT · Supervised by Prof. Ziqian Zeng

73.4%Less measured KV-cache storage than FP16, across four configurations; not total GPU memory.
6 modelsInstruction models from Qwen2.5, LLaMA-3.1, and Mistral.
40%Lower TPOT than reference INT4 for Qwen2.5-14B at 32K on one H20, batch 1.

The framework

Calibrate once. Reuse at inference.

Preserve attention behavior, then decide where precision matters most.

01

Offline preparation

Model weights stay unchanged
  1. Input

    Establish a reference

    A fixed calibration set, FP16 attention behavior, and candidate quantization paths.

    Same samples · controlled comparison
  2. Calibrate

    Measure behavior change

    K
    Attention-distribution KL
    V
    Aggregation-output perturbation
    Search candidates · select robust settings
  3. Freeze & allocate

    Produce fixed artifacts

    θ* · Selected quantization parameters and K/V formats.

    S → b* · Layer sensitivity informs precision allocation under a budget.

Fixed artifactsQuantization parameters θ* + layer precision b*
02

Online inference

Reuse the chosen configuration
  1. Load θ* and b*Read-only configuration
  2. Write quantized KVPack and store the cache
  3. Decode next tokensUse the quantized cache

Evaluate separatelyHeld-out task quality · cache storage · decoding latency

Original technical framework
Behavior-guided KV-cache quantization framework connecting offline calibration and sensitivity-based budget allocation to read-only online inference
Original framework: calibration produces frozen cache parameters and sensitivity profiles; the allocator outputs the layer precision configuration used during inference.

My Contribution

01

Calibration & diagnosis

Designed the behavior-guided calibration protocol, K/V role diagnostics, and layer-sensitivity representation.

02

Quantized inference

Implemented custom cache formats, INT4-RoleAlign, the generation loop, and fused Triton decoding.

03

Evaluation & thesis

Independently built and ran the cross-model experiment, analysis, and visualization pipeline, and wrote the thesis.

Preserve what attention does

During autoregressive decoding, every generated token adds new Key and Value tensors to the KV cache. Its memory footprint therefore grows with sequence length, batch size, model depth, the number of KV heads, and storage precision. Lowering the cache from FP16 to INT8 or INT4 can release substantial memory, but the resulting numerical error is not equally important in every tensor, layer, model, or task.

Key perturbations can change attention logits and reorder the attention distribution after softmax. Value perturbations act later through weighted aggregation and change the representation returned by attention. This distinction motivated a calibration framework that observes how quantization changes the computation actually consumed by the model, rather than relying on tensor error alone.

K/V Role-Aware Low-Bit Quantization

The INT8 path provides a conservative reference for behavior-guided calibration. INT4 creates a more demanding setting: in the evaluated Qwen models, direct symmetric INT4 caused stepwise failures in retrieval-oriented tests. K-only and V-only ablations showed that lowering Key precision was the more direct source of instability in those cases, while LLaMA-3.1-8B exhibited a different response.

INT4-RoleAlign therefore treats the two cache roles separately. Key is quantized per channel so that feature directions involved in attention ranking retain independent endpoints. Value is quantized per token so that each time step receives its own parameters for the aggregation path. Offline percentile candidates are selected using attention-distribution KL for Key and an aggregation-output perturbation proxy for Value.

Across the three complete KIVI-style and RoleAlign comparisons, the RoleAlign perplexity difference relative to KIVI-style remained within 0.15, and both paths achieved 100% pass rates on the two Needle retrieval tests. Qwen2.5-14B also reached 100% on both tests under RoleAlign, although a paired KIVI-style result was not available for that model. The evidence therefore supports low-bit retrieval recovery and a reusable offline calibration interface, rather than a universal task-quality advantage over KIVI-style.

RoleAlign offline percentile calibration with separate Key distribution and Value output objectives
The calibrated percentile rules are saved for inference; this step does not update model weights. Click to enlarge.
INT4-RoleAlign diagram showing per-channel Key quantization for attention ranking and per-token Value quantization for content aggregation
INT4-RoleAlign quantizes Key per channel to preserve attention ranking and Value per token to preserve content aggregation.
Paired Key and Value reconstruction-error heatmaps and layer-wise error profiles for Qwen2.5-1.5B and LLaMA-3.1-8B
Paired K/V reconstruction-error profiles illustrate that low-bit sensitivity varies across tensor roles, layers, and model architectures.

Storage savings, with a conditional speed trade-off

Quantized cache runtime pipeline and separately measured cache storage and decode latency
Runtime integration and two separate measurement protocols. Fused INT4 improves over reference INT4 in this test, while FP16 remains faster. Click to enlarge.

I implemented custom FP16, INT8, symmetric INT4, mixed-precision, and role-aware cache paths within a controlled generation loop. The system includes quantization and dequantization routines, packed low-bit storage, calibration-artifact loading, GQA head mapping, and Triton decode kernels. The evaluation pipeline measures perplexity, Needle-in-a-Haystack retrieval, RULER, LongBench-style task scores, selected official LongBench tasks, cache memory, peak memory, prefill throughput, and decode latency.

The reported INT4 cache layout reduced measured KV-cache storage by approximately 73.4% relative to FP16 across four representative model configurations. Runtime gains were more conditional. For Qwen2.5-14B at 32K context, fused Triton INT4 decoding reduced TPOT from 190.23 ms to 113.16 ms, approximately 40%, relative to the reference INT4 path on a single NVIDIA H20 with batch size 1, 64 generated tokens, and 10 runs. The crossover depended on the number of KV heads, context length, and backend implementation, so this result is reported as a system-specific within-INT4 improvement rather than a universal speedup.

Measured KV-cache and peak-memory scaling for FP16, INT8, KIVI-style, and symmetric INT4 cache paths across batch sizes Decode throughput, per-sequence throughput, and prefill throughput across batch sizes for full-precision and quantized cache paths
Measured memory and throughput behavior for the evaluated cache paths. Runtime results are tied to the reported NVIDIA H20, batch, sequence-length, and backend settings.

The 40% reduction is against the reference INT4 path. In the same table, fused INT4 at 113.16 ms remained slower than the 43.13 ms FP16 latency baseline.

Calibration protocol

The calibration pipeline evaluates full-precision reference configurations and candidate quantized configurations on a fixed calibration set. It records two complementary behavioral signals: the KL divergence between reference and quantized attention distributions, and the perturbation of the attention-aggregation output. Downstream task scores remain a separate audit layer rather than being folded into the calibration objective.

After candidate parameters are evaluated, the selected scales, clipping thresholds, quantization-axis metadata, and layer-wise sensitivity readings are frozen before downstream evaluation. This separation keeps calibration decisions independent from the test results and makes the same offline artifacts reusable by the cache implementation and allocation experiments.

Sensitivity-Driven Budget Allocation

The per-layer calibration readings are summarized into behavior-sensitivity profiles. Under a fixed budget, a behavior-guided allocator protects the most sensitive layers with a higher-precision path and compares that choice with uniform quantization and position-based heuristics.

AutoK extends this idea by selecting the smallest number of protected layers that reaches a specified fraction of cumulative sensitivity. It is a model-level candidate generator, not a universally optimal policy. The experiments instead identify model-dependent regimes: early-layer protection for Qwen2.5-3B, a clustered middle-layer regime for LLaMA-3.1-8B, a strong AutoK example on Mistral-7B, and several closely performing configurations on Qwen2.5-14B.

Scope and Limitations
  • The official LongBench comparison was limited to Qwen2.5-1.5B, three tasks, at most 50 samples per task, and one seed; it is an audit of direction rather than a general real-world guarantee.
  • The 100% Needle figures are exact-match pass rates on two controlled retrieval tests; they do not establish general long-context question-answering performance.
  • “Three models with complete KIVI-style/RoleAlign comparisons” refers to models with results for both paths. Qwen2.5-14B had RoleAlign coverage but no paired KIVI-style result.
  • The 73.4% figure measures KV-cache storage reduction across four representative configurations, not total model memory or end-to-end peak GPU memory.
  • The approximately 40% TPOT reduction was measured against the reference INT4 path for Qwen2.5-14B at 32K on a single NVIDIA H20, with batch size 1, 64 generated tokens, 10 runs, greedy decoding, and the reported backend. This is a within-INT4 backend comparison: in the same table, fused INT4 at 113.16 ms remained slower than the 43.13 ms FP16 latency baseline. Other hardware and serving systems may move the crossover point.
  • The allocation experiments compare methods within similar INT4 budget bands rather than establishing a strict matched-budget winner against every baseline.
  • The evaluated models are decoder-only architectures dominated by grouped-query attention. MHA, MQA, mixture-of-experts, sliding-window attention, and other architectures require separate validation.
  • AutoK is an offline, model-level budget proposer. Prompt-adaptive routing and production serving integration remain future work.
Future Directions
  • Evaluate role-aware allocation under strictly matched memory budgets.
  • Study prompt-conditioned policy selection after the static sensitivity profiles are better established.
  • Extend the evaluation to additional attention architectures, hardware platforms, and serving workloads.
  • Separate quantization-format effects, calibration effects, and architectural effects through more controlled ablations.
Project Information
  • Duration: January 2026 – May 2026
  • Type: Sole-authored undergraduate thesis
  • Institution: South China University of Technology
  • Supervisor: Prof. Ziqian Zeng
  • Models: Qwen2.5-1.5B, Qwen2.5-3B, Qwen2.5-7B, Qwen2.5-14B, LLaMA-3.1-8B, and Mistral-7B
  • Tools: Python, PyTorch, Hugging Face Transformers, Triton, FlashInfer, NumPy, Pandas, and Matplotlib
Reported experimental outcomes

Key Outcomes

  • Built a reproducible KV-cache quantization and evaluation pipeline covering six instruction models from the Qwen2.5, LLaMA-3.1, and Mistral families.
  • Established a behavior-guided INT8 reference path and evaluated its fidelity against FP16 under a controlled task protocol.
  • Symmetric INT4 caused retrieval collapse on Qwen2.5-1.5B and Qwen2.5-7B under the controlled Needle protocol, while the role-aware per-channel-Key and per-token-Value format restored a usable low-bit retrieval path.
  • INT4-RoleAlign achieved 100% pass rates on both Needle-single-retrieval and MK-NIAH-2 for each of the three models with complete KIVI-style/RoleAlign comparisons under the reported five-seed protocol: Qwen2.5-1.5B, Qwen2.5-7B, and LLaMA-3.1-8B.
  • Across Qwen2.5-1.5B, Qwen2.5-7B, LLaMA-3.1-8B, and Qwen2.5-14B, the reported INT4 cache layout reduced measured KV-cache storage by approximately 73.4% relative to FP16.
  • For Qwen2.5-14B at 32K context, fused Triton INT4 decoding reduced TPOT from 190.23 ms to 113.16 ms, approximately 40%, relative to the reference INT4 path on a single NVIDIA H20 with batch size 1.
  • Converted layer-wise behavior sensitivity into fixed-budget allocation profiles and an AutoK mechanism for proposing model-level protection budgets.