KV Cache Quantization — Nén cache INT4/INT3, mở khóa ultra-long context

Bạn muốn chạy LLaMA 70B với context 1 triệu token? Vấn đề không phải ở GPU compute — H100 có đủ FLOPs để xử lý — mà ở chỗ bạn cần hơn 200GB VRAM chỉ để lưu KV cache. Trong khi đó, weights của model chỉ chiếm ~40GB. Đây là lý do tại sao các hệ thống như vLLM, SGLang và TensorRT-LLM đang chuyển sang nén KV cache xuống INT4, thậm chí INT3, mở đường cho "ultra-long context" mà không cần cluster GPU khổng lồ.

Vấn đề

Trong inference autoregressive, mỗi token mới sinh ra cần truy cập Key (K) và Value (V) của tất cả token trước đó để tính attention. Điều này tạo ra một bài toán memory bậc nhất với độ phức tạp O(L) — chi phí memory tăng tuyến tính theo độ dài sequence.

Với Llama-2-7B ở FP16:

• 10k tokens: ~5GB cache (chấp nhận được)
• 100k tokens: >50GB cache (vượt quá VRAM của hầu hết GPU consumer)
• 1M tokens: >500GB (chỉ có thể chạy trên cluster đắt đỏ)

Grouped Query Attention (GQA) đã giảm 4x bằng cách chia sẻ K/V heads, nhưng với context thực sự dài (100k+), điều đó vẫn chưa đủ. Bạn cần nén sâu hơn — nhưng làm sao để nén từ 16-bit xuống 4-bit mà không biến model thành "ngu ngốc"?

Ý tưởng cốt lõi

Attention không phải là bộ nhớ thanh ghi (register file), mà là bộ nhớ liên kết (associative memory) dựa trên "soft matching".

Hãy tưởng tượng bạn đang tìm kiếm trong thư viện 32.000 cuốn sách (context length). Thay vì phải nhớ chính xác từng trang sách (FP16), bạn chỉ cần nhớ gần đúng nội dung và vị trí. Khi attention mechanism tính softmax(Q·K^T), nó đang làm một việc: tính trung bình có trọng số của Value vectors.

Đây là "aha moment": Nếu bạn thêm nhiễu nhỏ (quantization error) vào từng vector K và V, khi lấy trung bình qua 32k tokens, các lỗi này triệt tiêu lẫn nhau theo luật số lớn (Law of Large Numbers). Attention là "fuzzy matcher" — nó không cần byte-precise, chỉ cần biết "token nào liên quan" với độ tin cậy tương đối.

Tại sao Key và Value lại khác nhau?

• Keys quyết định "nhìn đâu" (direction in high-D space). Dot product Q·K chỉ quan tâm góc giữa vectors, không phải magnitude. Quantization làm thay đổi góc nhỏ, nhưng cosine similarity chịu đựng tốt perturbation này — miễn là relative ordering của similarity scores được bảo toàn. Do đó, Keys dùng per-channel quantization (quan trọng vì RoPE rotation).

• Values quyết định "lấy gì" (content). Đây là weighted sum, nên magnitude error accumulate tuyến tính. Do đó Values cần per-token quantization (mỗi token có variance khác nhau).

Vấn đề Outlier — "Kẻ phá bĩnh":

Softmax rất nhạy cảm với outliers. Một giá trị K cực lớn sẽ dominate cả attention pattern (chiếm 99% weight). Nếu bạn nén outlier này sai, bạn phá vỡ hoàn toàn khả năng "nhớ" một token cụ thể. Đây là lý do naive INT4 thất bại, còn các phương pháp hiện đại như KVQuant/KIVI thành công: họ dùng Hadamard transform để "spread" outlier energy đều ra các chiều khác trước khi nén.

RoPE — Cái bẫy tinh vi:

Keys được xoay bởi Rotary Position Embedding. Nếu bạn quantize sau khi rotate, error sẽ bị amplify bởi rotation matrix. Insight của KVQuant: quantize Keys ở trạng thái "gốc" (pre-RoPE), lưu ở low-bit, rồi rotate on-the-fly khi inference. Điều này giữ lại thông tin position chính xác hơn.

Mixed Precision — "Nóng và Lạnh":

Không phải token nào cũng bình đẳng. Token gần nhất (hot) cần độ chính xác cao để tính toán chính xác; token cũ (cold) có thể chịu đựng nén mạnh hơn. Chiến lược MixKVQ giữ 4-bit cho recent tokens, 2-bit cho old tokens.

That's it. Bạn không cần lưu FP16 cho toàn bộ lịch sử — chỉ cần "gợi ý đủ tốt" về quá khứ để attention mechanism hoạt động.

Tại sao nó hoạt động

Mathematically, quantization error trong attention hoạt động như sau:

1. Asymmetric Quantization:

   • Keys: Per-channel (dọc theo chiều head_dim) để bảo vệ cấu trúc RoPE
   • Values: Per-token (mỗi position có scale riêng) vì variance thay đổi theo vị trí

2. Hadamard Rotation: Trước khi quantize, nhân vector với Hadamard matrix để phân tán outlier energy:

   # Pseudo-code cho TurboQuant-style approach
   rotated_k = hadamard_transform(k)  # Làm phẳng phân bố
   quantized_k = int4_quantize(rotated_k)  # Nén

3. Pre-RoPE Storage: Lưu Keys trước khi áp dụng rotation, tránh error amplification:

   # Thay vì: rotate -> quantize -> store -> dequantize -> use
   # Ta làm: quantize (pre-RoPE) -> store -> rotate on-the-fly -> use

4. Softmax Insensitivity: Với N tokens, lỗi quantization σ trên mỗi Value tạo ra lỗi tổng cộng ~σ/√N sau khi weighted average (theo Central Limit Theorem). Với N=100k, lỗi giảm 300 lần so với lỗi đơn lẻ.

Ý nghĩa thực tế

Benchmarks thực tế:

• KVQuant: Đạt 10 triệu token context với < 1% perplexity drop so với FP16 baseline
• NVFP4: 3x latency reduction và 20% cache hit rate improvement so với FP8
• Kitty (2-bit): Throughput tăng 2.35x với accuracy loss < 2% trên LongBench

Trade-offs:

• Tốt cho: Summarization dài, RAG với nhiều documents, multi-turn chat history
• Không tốt cho: Coding (cần nhớ chính xác syntax token-by-token), math reasoning (precision accumulation quan trọng), hay bất kỳ task nào cần "exact memorization" của token cụ thể ở xa
• GQA Interaction: Trên model hiện đại (Llama-3, Qwen) đã dùng GQA (giảm 4x KV cache), gain từ quantization thấp hơn — nhưng vẫn cần thiết cho ultra-long context (>200k)

Ai đang dùng:

• vLLM/SGLang: Hỗ trợ KV cache quantization với cache_quant_bit=4
• NVIDIA TensorRT-LLM: NVFP4 format cho production serving
• llama.cpp: Q4_0/Q8_0 KV cho local inference
• MLX (Apple Silicon): Tối ưu cho unified memory architecture

Đào sâu hơn

Paper gốc:

• KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (NeurIPS 2024) — Pre-RoPE quantization và outlier handling
• KIVI: A Tuning-Free Asymmetric 2-bit Quantization for KV Cache (2024) — Asymmetric per-channel/per-token approach
• Kitty: Accurate and Efficient 2-bit KV Cache Quantization (2025) — Kernel optimization cho 2-bit

Cùng cụm (Inference Frontier):

• EAGLE — Draft head gắn vào internal layer, 3x speedup — Kết hợp với KV quantization để tăng tốc inference
• TurboQuant — Random rotation trước khi quantize, 3-bit KV cache không mất accuracy — Nâng cao của KV quantization với Hadamard transform
• Prompt Lookup Decoding — Speculative decoding không cần draft model — Kỹ thuật tăng tốc khác cho long-context RAG

Đọc tiếp:

• KV Cache — Đừng tính lại cái đã tính, cache K và V — Foundation của KV cache (Level 1)
• Quantization INT8/INT4 — Ít bit hơn, nhỏ hơn 8x, vẫn chính xác — Nguyên lý quantization nói chung (Level 1)
• Native Long-Context Training — Train dài từ đầu, reason tốt hơn extend — Alternative approach cho long context (Level 2)