TurboQuant của Google: Giải pháp tối ưu giúp KV Cache không còn "ngốn" VRAM

Nếu bạn đã từng làm việc với các mô hình Transformer, bạn chắc chắn biết rằng cơ chế Attention là "bộ não" của toàn bộ quá trình vận hành. Nó giúp mô hình xác định các token đang tương tác với nhau, và chính cơ chế này tạo nên khả năng ấn tượng của các Mô hình Ngôn ngữ Lớn (LLM).

Tuy nhiên, sự ưu việt này đi kèm cái giá phải trả. Trong quá trình suy luận (inference), mỗi khi một token mới được dự đoán, các ma trận Key (K) và Value (V) phải được tính toán lại cho tất cả các token trước đó. Để giải quyết sự lãng phí tính toán này, KV Cache ra đời: thay vì tính toán lại, ta chỉ cần lưu trữ K và V vào VRAM và tái sử dụng chúng.

Minh họa cơ chế hoạt động của TurboQuant

Mặc dù KV Cache giúp giảm độ trễ đáng kể, nó lại tạo ra một thách thức lớn về bộ nhớ. Với các mô hình LLM khổng lồ, KV Cache có thể chiếm tới 20-30% VRAM và con số này tiếp tục tăng theo độ dài ngữ cảnh. Các giải pháp trước đây như Grouped-Query Attention (GQA) hay PagedAttention thường phải đánh đổi độ chính xác để tiết kiệm bộ nhớ.

Đây là lúc TurboQuant của Google xuất hiện như một giải pháp đột phá, quản lý vừa nén dữ liệu vừa giữ nguyên độ chính xác ở mức tối ưu về mặt lý thuyết.

Cơ chế hoạt động của TurboQuant

TurboQuant hoạt động theo quy trình hai giai đoạn: PolarQuant và Phần dư (Residual Correction). Thay vì cố gắng tái tạo lại vector một cách hoàn hảo như các phương pháp lượng tử hóa truyền thống, TurboQuant tập trung vào việc tái tạo những gì cơ chế Attention thực sự cần.

Giai đoạn 1: PolarQuant

Vấn đề lớn nhất của lượng tử hóa truyền thống là việc xử lý kém các giá trị ngoại lai (outliers). PolarQuant giải quyết vấn đề này bằng cách xoay vector trong không gian nhiều chiều.

Hãy tưởng tượng một vector có một giá trị cực lớn so với các giá trị còn lại. Nếu lượng tử hóa trực tiếp, bộ lượng tử sẽ bị "giãn" để bao phủ giá trị ngoại lai đó, làm mất mát thông tin của các giá trị nhỏ. Bằng cách xoay vector, năng lượng của giá trị ngoại lai được phân phối đều sang các tọa độ khác, tạo ra một phân phối đồng nhất (isotropic).

Xoay vector giúp phân phối năng lượng đồng đều

Sau khi xoay, TurboQuant sử dụng thuật toán Lloyd-Max để đặt các mức lượng tử hóa (centroids) vào vị trí tối ưu nhất nhằm giảm thiểu sai số bình phương trung bình (MSE). Kết quả của giai đoạn này là các chỉ mục (indexes) được lưu trữ thay vì các giá trị thực.

Giai đoạn 2: Phần dư và QJL

Đây là điểm độc đáo của TurboQuant. Các phương pháp truyền thống thường bỏ qua phần sai số (residual) sau lượng tử hóa. TurboQuant thì không; nó trích xuất các đặc tính cốt lõi của phần sai số này thay vì lưu trữ toàn bộ.

Giai đoạn này sử dụng biến đổi Quantized Johnson-Lindenstrauss (QJL). Nó lưu trữ dấu (+1 hoặc -1) của phép chiếu ngẫu nhiên lên vector phần dư. Cùng với đó là chuẩn L2 (độ lớn) của vector phần dư. Kết hợp hai yếu tố này cho phép khôi phục lại thông tin bị mất mà không tốn nhiều bộ nhớ.

So sánh hiệu quả lưu trữ giữa các phương pháp

Kết quả và Tương lai

Theo báo cáo chính thức, TurboQuant đạt được mức nén KV Cache hơn 4.5–5 lần (tương đương 2.5–3.5 bit mỗi kênh) với độ mất mát chính xác gần như bằng không trong thực tế.

Điều này có ý nghĩa rất lớn cho tương lai của AI. Khi chúng ta ngày càng cần các cửa sổ ngữ cảnh (context windows) dài hơn, nút thắt cổ chai về bộ nhớ của KV Cache sẽ không còn là rào cản nếu áp dụng TurboQuant. Chúng ta không nhất thiết phải nâng cấp phần cứng đắt đỏ, mà chỉ cần xử lý dữ liệu thông minh hơn.

Với sự ra mắt của TurboQuant, liệu chương trình quản lý bộ nhớ KV Cache đã khép lại, hay đây chỉ là nền tảng cho những công cụ mạnh mẽ hơn trong tương lai?