Zoom-in: KV Cache

Khi bạn trò chuyện với một LLM, các chữ cái thường hiện ra trên màn hình một cách liên tục và mượt mà. Tuy nhiên, ít ai biết rằng để sinh ra từ tiếp theo, mô hình phải thực hiện lại hàng tỷ phép tính nhân ma trận từ đầu.

Nếu không có sự hỗ trợ của KV Cache (Key-Value Cache), tốc độ phản hồi của mô hình sẽ chậm dần đến mức đóng băng khi cuộc hội thoại kéo dài.

Hãy phóng to vào kiến trúc phần cứng và giải thuật suy luận (inference).

---

Layer 1 — Vấn đề: Thuộc tính Auto-regressive và sự lãng phí tính toán

Các LLM hoạt động theo cơ chế tự hồi quy (auto-regressive): Để tạo ra token thứ $N$, mô hình cần phải đọc lại toàn bộ dữ liệu đầu vào cộng với $N-1$ token mà nó vừa tự sinh ra trước đó.

Ví dụ, để viết câu "Học máy thú vị":

• Bước 1: Input Học $\to$ Output máy
• Bước 2: Input Học máy $\to$ Output thú
• Bước 3: Input Học máy thú $\to$ Output vị

Trong kiến trúc Transformer, trái tim của mô hình là cơ chế Self-Attention (tự chú ý). Nó giúp mô hình hiểu mối liên hệ giữa các từ bằng cách nhân ma trận ba vectơ: Query (Q), Key (K), và Value (V) cho từng token.

Mỗi token mới đi vào đều cần:
- Vectơ Query (Q) để truy vấn.
- Vectơ Key (K) để đối chiếu thông tin.
- Vectơ Value (V) để chứa nội dung.

Ở bước 2, khi xử lý từ máy, mô hình phải tính toán lại toàn bộ các vectơ Query, Key và Value cho từ Học mặc dù từ này đã được tính toán ở bước 1. Khi cuộc trò chuyện kéo dài đến hàng nghìn từ, sự lãng phí này sẽ tăng theo bậc hai $O(N^2)$, khiến GPU tiêu tốn năng lực xử lý chỉ để tính đi tính lại những thứ đã biết.

---

Layer 2 — Giải pháp: KV Cache – Tính một lần, dùng nhiều lần

KV Cache được ra đời để phá vỡ vòng lặp lãng phí đó. Nguyên lý của nó rất trực quan: Những gì đã tính toán một lần thì lưu vào bộ nhớ đệm, không bao giờ tính lại.

sequenceDiagram
    participant LLM as Mô hình ngôn ngữ
    participant Cache as KV Cache (VRAM)
    
    Note over LLM, Cache: Xử lý token 1: 'Học'
    LLM->>Cache: Lưu vectơ Key và Value của 'Học'
    Note over LLM, Cache: Sinh token 2: 'máy'
    LLM->>LLM: Chỉ tính Q, K, V cho riêng token 'máy'
    LLM->>Cache: Lấy Key/Value của 'Học' + Lưu Key/Value của 'máy'
    LLM->>LLM: Tính toán Attention và sinh từ tiếp theo

Nhờ cơ chế này:

• Mô hình chỉ tính toán đầy đủ bộ ba vectơ Q, K, V cho duy nhất token mới xuất hiện tại bước hiện tại.
• Đối với tất cả các token trước đó, mô hình chỉ việc truy xuất vectơ Key (K) và Value (V) từ KV Cache.

Sự tối ưu này loại bỏ việc tính toán lại các vectơ Key và Value cho các token cũ (giảm chi phí tạo projection về mức hằng số $O(1)$ cho mỗi token đã qua). Dù vậy, mô hình vẫn cần thực hiện phép tính attention giữa token mới và toàn bộ KV Cache cũ (độ phức tạp thực tế của bước attention vẫn tăng tuyến tính theo chiều dài chuỗi $O(N)$). Đây là lý do giúp AI có thể phản hồi nhanh hơn rất nhiều so với việc tính toán lại từ đầu.

---

Layer 3 — Sự đánh đổi: Khủng hoảng VRAM

Mặc dù giải phóng năng lực tính toán của GPU, KV Cache lại tạo ra một thách thức rất lớn khác: Nó cực kỳ ngốn VRAM.

Dung lượng của KV Cache tăng tuyến tính theo:

1. Batch size: Số lượng request (hoặc người dùng) đang xử lý đồng thời.
2. Context length: Chiều dài tối đa của cuộc trò chuyện.
3. Cấu trúc mô hình: Số lượng layer (lớp mạng) và số lượng attention head (đầu chú ý).

[!WARNING] Với một mô hình cỡ trung bình như Llama-3-8B chạy ở context length 8,000 với batch size là 16, riêng KV Cache đã chiếm khoảng 8GB VRAM – tương đương với dung lượng để tải toàn bộ các tham số của chính mô hình đó.

Nếu dung lượng cache này vượt quá giới hạn vật lý của phần cứng, hệ thống sẽ gặp lỗi OOM (Out of Memory - Cạn kiệt bộ nhớ) và crash hoàn toàn. Đây chính là lý do vì sao các kỹ thuật tối ưu hóa cache như Grouped-Query Attention (GQA) hay PagedAttention (quản lý bộ nhớ đệm KV giống như cơ chế phân trang bộ nhớ ảo của hệ điều hành) đang là những chủ đề nóng bỏng nhất trong giới kỹ sư tối ưu hóa LLM hiện nay.

---

Full picture

sequenceDiagram
    participant Tokenizer as Tokenizer
    participant GPU as GPU Processing
    participant VRAM as KV Cache (VRAM)
    
    Note over Tokenizer, VRAM: Lần chạy đầu tiên (Prefill Phase)
    Tokenizer->>GPU: Nhập prompt (N tokens)
    GPU->>GPU: Tính toán Q, K, V cho toàn bộ N tokens
    GPU->>VRAM: Lưu tất cả vectơ Key & Value của N tokens vào cache
    GPU-->>Tokenizer: Sinh ra token tiếp theo (N+1)
    
    Note over Tokenizer, VRAM: Các bước sinh từ sau (Decode Phase)
    Tokenizer->>GPU: Nhập token mới nhất (N+1)
    GPU->>GPU: Chỉ tính Q, K, V cho riêng token N+1
    VRAM-->>GPU: Lấy Key & Value của các token cũ từ 1 đến N
    GPU->>GPU: Tính attention nhanh gọn
    GPU->>VRAM: Lưu Key & Value của token N+1 vào cache
    GPU-->>Tokenizer: Sinh tiếp token N+2

---

Takeaway

KV Cache là kỹ thuật tối ưu hóa tốc độ cốt lõi giúp LLM tránh việc tính toán lặp lại các vectơ Key và Value của các token cũ trong quá trình sinh từ tự hồi quy. Bằng cách lưu trữ các ma trận này vào bộ nhớ VRAM, chi phí tính projection cho các token cũ được giảm về mức hằng số $O(1)$. Mặc dù phép tính attention ở mỗi bước sinh từ vẫn tăng tuyến tính $O(N)$ theo chiều dài ngữ cảnh, sự tối ưu này đã loại bỏ đáng kể lượng phép tính nhân ma trận khổng lồ. Đổi lại, hệ thống phải gánh dung lượng bộ nhớ đồ họa (VRAM) rất lớn, thúc đẩy sự phát triển của các giải pháp tối ưu cấu trúc bộ nhớ cache như GQA và PagedAttention.