KV Cache trong LLM: Hiểu và triển khai từ đầu

KV cache lưu trữ các vector key (K) và value (V) từ các phép tính attention trước đó để tái sử dụng khi tạo ra các token tiếp theo, giúp tăng tốc suy luận bằng cách tránh tính toán lại. Nếu không có KV cache, mỗi bước tạo token sẽ tính toán lại key và value cho tất cả các token trước đó; còn với cache, chỉ có vector của token mới được tính và nối thêm vào cache. Bài viết cung cấp một bản cài đặt mã dựa trên mô hình kiểu GPT từ cuốn sách của tác giả. Các thay đổi chính là: thêm bộ đệm cache (cache_k và cache_v) trong lớp MultiHeadAttention, sửa đổi phương thức forward để sử dụng cache có điều kiện, thêm phương thức reset, và truyền cờ use_cache qua toàn bộ mô hình. Trong quá trình tạo token, khi use_cache là True, mô hình chỉ xử lý token mới sau prompt ban đầu, trong khi không có cache, nó xử lý toàn bộ chuỗi ở mỗi bước. Một so sánh hiệu suất đơn giản cho thấy KV cache làm tăng tốc độ tạo token lên gần gấp đôi trong ví dụ được thử nghiệm.