KVarN: Khi quantize KV-Cache không còn làm hỏng reasoning của LLM

Nếu bạn đã từng chạy inference một model LLM lớn và cố optimize memory bằng cách quantize KV-Cache xuống 4-bit hay 8-bit, chắc bạn cũng đã gặp cái cảm giác khó chịu đó: model bắt đầu nói bậy, chain-of-thought bị đứt đoạn, hoặc tệ hơn là đưa ra kết quả sai một cách rất tự tin. Đây không phải bug của bạn đây là vấn đề cơ bản của KV-Cache quantization khi gặp các tác vụ reasoning dài.

Paper KVarN vừa drop trên HuggingFace tuần này đánh thẳng vào vấn đề này, và cách tiếp cận của họ khá thú vị.

Tại sao quantize KV-Cache lại đặc biệt nguy hiểm với reasoning?

Để hiểu KVarN giải quyết cái gì, mình cần nói một chút về bản chất của vấn đề trước.

Khi model generate token, nó cần attend đến tất cả các token trước đó thông qua KV-Cache. Với các tác vụ reasoning toán học, lập luận logic, code generation chain-of-thought có thể dài hàng nghìn token. Mỗi bước suy luận phụ thuộc vào các bước trước, tạo ra một chuỗi phụ thuộc rất dài.

Vấn đề là khi bạn quantize KV-Cache (nén từ float16/bfloat16 xuống int4/int8), bạn đưa vào một lượng quantization error nhất định. Với text generation thông thường, cái error này không đáng kể model vẫn generate ra câu văn đọc được. Nhưng với reasoning tasks, error này tích lũy qua từng bước suy luận. Một sai số nhỏ ở bước 3 có thể kéo theo sai lệch lớn hơn ở bước 10, rồi đến bước 20 thì kết quả đã lệch hoàn toàn so với ground truth.

Đây là lý do tại sao bạn thấy các model quantized thường vẫn chat tốt nhưng fail ở math benchmarks hay code reasoning.

Variance Normalization ý tưởng core của KVarN

KVarN tiếp cận bài toán từ góc độ: tại sao quantization error lại tích lũy?

Câu trả lời của họ là do sự phân bố không đồng đều của các giá trị trong KV-Cache. Một số channel có variance rất cao (giá trị trải dài từ -10 đến +10), trong khi channel khác lại rất nhỏ (từ -0.1 đến +0.1). Khi bạn quantize cả hai với cùng một scheme, channel có variance cao sẽ bị quantize thô hơn nhiều, tạo ra error lớn hơn. Và những channel này thường là những channel quan trọng nhất với attention mechanism.

Giải pháp của KVarN: normalize theo variance trước khi quantize. Về cơ bản, họ scale từng channel về cùng một phân phối chuẩn hơn trước khi nén, sau đó lưu lại scaling factor để dequantize sau. Điều này giúp error phân bố đều hơn thay vì tập trung vào các channel quan trọng.

Theo kinh nghiệm của mình khi làm với các kỹ thuật quantization khác nhau, cái insight này không mới hoàn toàn SmoothQuant và AWQ cũng dùng tư tưởng tương tự cho weight quantization. Điểm hay của KVarN là họ áp dụng nó online cho KV-Cache, tức là không cần calibration dataset hay post-training step phức tạp.

So sánh với các approach hiện tại

Cái mình thấy hay nhất ở KVarN là không cần calibration data. Trong thực tế deploy production, việc phải chuẩn bị calibration dataset rồi chạy calibration pass trước là một friction point khá lớn, đặc biệt khi bạn muốn swap model thường xuyên hoặc fine-tune cho domain mới.

Cái gì thực sự bị fix?

Error accumulation trong reasoning tasks xảy ra theo một pattern khá đặc trưng. Hãy tưởng tượng model đang giải một bài toán nhiều bước:

1Bước 1: Tính X = A + B  → quantization error nhỏ, kết quả ~đúng
2Bước 2: Tính Y = X * C  → error từ bước 1 được nhân lên
3Bước 3: Tính Z = Y / D  → error tiếp tục lan truyền
4...
5Bước N: Kết quả cuối    → sai hoàn toàn

Với variance normalization, error ở mỗi bước nhỏ hơn và phân bố đều hơn, nên khi tích lũy qua N bước, tổng error vẫn trong ngưỡng chấp nhận được.

Anh em lưu ý một điểm quan trọng: KVarN không loại bỏ hoàn toàn quantization error không có kỹ thuật nào làm được điều đó mà không tăng memory footprint. Cái họ làm là kiểm soát cách error phân bố để nó ít gây hại hơn cho các tác vụ đòi hỏi reasoning dài.

Practical implications cho người làm inference optimization

Nếu bạn đang chạy LLM inference và đang cân nhắc KV-Cache quantization, đây là cách mình nhìn nhận KVarN trong bức tranh tổng thể:

Với các use case như chatbot thông thường, summarization, hay translation quantization naive (INT8) vẫn ổn và đơn giản hơn nhiều. Overhead của variance normalization không đáng để thêm vào.

Nhưng nếu bạn đang serve các model cho coding assistant, math tutoring, hay bất kỳ tác vụ nào đòi hỏi multi-step reasoning KVarN trở nên rất relevant. Đây chính xác là những use case mà dân dev đang đổ tiền vào nhất hiện tại, từ Cursor đến các coding copilot nội bộ.

Mình cũng nghĩ đây sẽ là một technique quan trọng khi các model reasoning như DeepSeek-R1, QwQ hay các variant của o1 ngày càng phổ biến hơn. Những model này generate chain-of-thought cực dài đôi khi hàng chục nghìn token nên KV-Cache memory là bottleneck thực sự, và bất kỳ kỹ thuật quantization nào cũng cần đảm bảo không phá vỡ reasoning chain.

Còn gì cần theo dõi?

Paper mới ra nên chưa có nhiều benchmark độc lập. Mình muốn thấy kết quả trên các model reasoning mới nhất (R1, QwQ, Gemini thinking modes) thay vì chỉ các model standard. Ngoài ra, overhead của việc lưu thêm scaling factors cho từng channel cần được đánh giá kỹ hơn trong điều kiện batch inference thực tế lý thuyết là thấp nhưng thực tế production đôi khi có surprises.

Tác giả Philippe Bich có mention sẵn sàng thảo luận về paper trên HuggingFace, nên nếu bạn có câu hỏi cụ thể về implementation thì đây là cơ hội tốt để hỏi thẳng. Cộng đồng ML research khá open trong việc discuss technical details, đặc biệt khi paper vừa mới ra.

Nhìn chung, KVarN là một contribution thực tế và có giá trị ứng dụng rõ ràng không phải loại paper chỉ improve benchmark 0.1% rồi thôi. Nếu team bạn đang làm inference optimization cho reasoning models, đây là thứ đáng đọc kỹ.