NexusQuant: Cách nén KV Cache của LLM lên tới 33 lần mà không cần huấn luyện

Nếu bạn từng cố gắng chạy một Mô hình Ngôn ngữ Lớn (LLM) với ngữ cảnh dài (long context), chắc hẳn bạn đã gặp phải vấn đề nan giải: hết bộ nhớ.

Kẻ thủ ác ở đây chính là KV cache — bảng lưu trữ các khóa (keys) và giá trị (values) mà mô hình duy trì cho mỗi token. Ở một mô hình hiện đại với ngữ cảnh 128K, bộ nhớ đệm này có thể dễ dàng chiếm tới 80 GB. Tức là một GPU A100 entier chỉ để lưu trữ attention trung gian.

NexusQuant loại bỏ nút thắt cổ chai này. Không cần huấn luyện (training), không cần hiệu chuẩn (calibration), chỉ cần một dòng code duy nhất.

Cài đặt nhanh

Bạn có thể cài đặt NexusQuant qua pip như sau:

pip install nexusquant-kv
pip install "nexusquant-kv[hf]"  # với HuggingFace transformers

Quickstart

Việc tích hợp vào quy trình làm việc của bạn vô cùng đơn giản:

from nexusquant import nexusquant_evict

with nexusquant_evict(model, quality="balanced"):
    output = model.generate(input_ids, max_new_tokens=512)

Về cơ bản, đó là tất cả những gì bạn cần làm. Mô hình không bị sửa đổi — các "hooks" sẽ được tự động cài đặt và gỡ bỏ bởi context manager.

Những con số ấn tượng

Được đo trên Mistral-7B, GPU A100, FP16. Tất cả các tỷ số dưới đây đều đã bao gồm cả overhead (các thang đo, chỉ mục, metadata).

Preset	Tỷ lệ nén	Sự giảm sút PPL	Ngữ cảnh trên 80 GB
high	10x	+0.4%	~1.3M tokens
balanced	17x	+1.3%	~2.2M tokens
max	33x	+2.6%	~4.2M tokens

Cụ thể hơn: Từ 128K tokens ngữ cảnh → có thể lên tới 4.2M tokens với chế độ max trên cùng một GPU.

Cơ chế hoạt động (6 bước)

NexusQuant hoạt động dựa trên quy trình tối ưu hóa 6 bước:

1. Chấm điểm quan trọng (Importance scoring): Xếp hạng các token dựa trên tích phẩm nội (dot product) chéo giữa các đầu attention.
2. Loại bỏ token (Eviction): Loại bỏ các token ít quan trọng nhất; luôn bảo toàn token BOS và một cửa sổ trượt (sliding window) gần đây.
3. Gỡ bỏ RoPE (RoPE removal): Hoàn tác các embedding xoay trong keys để chúng chia sẻ một không gian con chung, giúp giảm lỗi lượng tử hóa ~0.7 pp.
4. Xoay Hadamard (Hadamard rotation): Phân phối năng lượng đồng đều giữa các chiều để không có giá trị ngoại lệ nào chi phối thang đo lượng tử hóa.
5. Lượng tử hóa mạng E8 (E8 lattice quantization): Lượng tử hóa các nhóm 8 số thực trong mạng gốc E8 (cách đóng gói hình cầu dày đặc nhất trong 8D), đạt 2 bits/dim.
6. Delta coding + zstd: Các token liên tiếp tạo ra các chỉ mục mạng tương tự; việc lưu trữ các delta và nén với zstd mang lại thêm 2-3x dung lượng cho luồng chỉ mục.

Quá trình Eviction giúp giảm số lượng (~2.5x với 60% eviction). E8 giúp giảm độ chính xác (~7x sau entropy coding). Kết hợp lại: 17x.

So sánh với các đối thủ

Phương pháp	Tỷ lệ nén	Sự giảm sút PPL	Cần huấn luyện
NexusQuant	10-33x	+0.4-2.6%	Không
TurboQuant (Google)	~5-6x	~0%	Không
KVTC (NVIDIA)	upto 20x	<1%	Có (calibration)
CommVQ (Apple)	~8x	~0%	Có (full training)
Palu	11x	~25% rel	Có (calibration)

NexusQuant là phương pháp không cần huấn luyện (training-free) có tỷ lệ nén cao nhất hiện nay. Các con số của đối thủ được lấy từ các bài báo đã xuất bản của họ, không phải chạy lại trên phần cứng của người viết.

Các mô hình được hỗ trợ

Bất kỳ LM nhân quả HuggingFace nào sử dụng split-half RoPE (tiêu chuẩn kể từ Llama-2):

• Gia đình Llama (Llama-2, Llama-3, Llama-3.1)
• Mistral / Mixtral
• Qwen
• Phi
• Gemma

Chưa hỗ trợ: Các mô hình với RoPE xen kẽ (GPT-NeoX, GPT-J).

Một số hạn chế

Sự trung thực là ưu tiên hàng đầu, dưới đây là những lưu ý quan trọng:

• Chất lượng phụ thuộc vào văn bản: Văn bản sáng tạo/kể chuyện bị giảm chất lượng nhiều hơn so với văn bản có cấu trúc/kỹ thuật ở cùng một tỷ lệ. Hãy kiểm tra trên khối lượng công việc thực tế của bạn trước khi triển khai.
• Prefix ngắn bị ảnh hưởng: Các tiền tố dưới 500 token sẽ thấy mức độ giảm sút lớn hơn so với các con số trên (được đo ở mức 1600-3500 token).
• E8 bị giới hạn bởi CPU: Việc triển khai production cần các kernel Triton/CUDA cho bước lượng tử hóa.
• Eviction là vĩnh viễn: Các token bị loại bỏ sẽ biến mất. Nếu tác vụ của bạn yêu cầu khả năng ghi nhớ chính xác một token cụ thể, hãy đo lường độ nhạy cảm với eviction trước.

Thử nghiệm ngay trên Colab

Tác giả cung cấp một notebook có thể chạy trong dưới 2 phút trên gói miễn phí của Colab, sử dụng TinyLlama trên CPU:

github.com/jagmarques/nexusquant/blob/main/examples/nexusquant_demo.ipynb