Trong bài viết này
Từ được đánh dấu có giải thích. Di chuột hoặc chạm để xem, hay mở từ điển AI.
Một máy có thể nạp được trọng số mô hình nhưng hết bộ nhớ khi nhiều người cùng trò chuyện dài. Phần tăng thêm thường nằm ở trạng thái suy luận, trong đó KV cache là một thành phần lớn. Đếm dung lượng tệp trọng số chưa đủ để lập kế hoạch phục vụ.
Trước khi đọc, hãy nắm self-attention, causal mask và sự khác nhau giữa huấn luyện với suy luận. Các phép tính dưới đây dùng cấu hình giả định và đơn vị nhị phân; chúng không phải số đo của một dịch vụ cụ thể.
Vì sao key và value cũ có thể dùng lại?
Trong Transformer decoder có causal attention, biểu diễn tại token quá khứ không phụ thuộc token tương lai. Vì vậy, sau khi đã tính key và value ở mỗi lớp cho phần tiền tố, ta có thể giữ chúng lại. Khi thêm token mới, chỉ tính phần mới rồi cho query mới attention vào toàn bộ key và value hợp lệ.
Query quá khứ thường không cần giữ cho bước sinh kế tiếp: đầu ra mới chỉ cần query của vị trí mới. Đây là lý do tên cache gồm K và V. Tài liệu How caching works của Hugging Face giải thích cơ chế tái sử dụng này và vai trò của cache khi sinh tự hồi quy.
Cache không phải bộ nhớ kiến thức cập nhật trọng số. Xóa cache khiến hệ thống phải tính lại tiền tố; nó không khiến mô hình quên những gì đã học trong huấn luyện. Hai loại “nhớ” này cần được tách khi giải thích chi phí và quyền riêng tư.
Hai giai đoạn có hình dạng công việc khác nhau
Prefill xử lý các token của prompt, tạo trạng thái ban đầu và KV cache. Nhiều vị trí trong prompt có thể được xử lý song song dưới causal mask. Decode sinh từng token tiếp theo, thường đọc lại cache ngày càng dài cho mỗi bước.
Vì thế thời gian đến token đầu tiên và tốc độ các token sau đó là hai phép đo khác nhau. Prompt dài có thể làm prefill chậm; câu trả lời dài kéo dài nhiều vòng decode. Một tối ưu cho giai đoạn này chưa chắc cải thiện giai đoạn kia.
Mô hình tự hồi quy tính xác suất token tiếp theo từ ngữ cảnh hiện có, chọn một token rồi nối vào chuỗi. Chuỗi mới trở thành ngữ cảnh cho lượt tiếp theo.
↔ Vuốt ngang để xem toàn bộ sơ đồ.
01 / Đọc ngữ cảnh
Mô hình xử lý các token của yêu cầu và phần trả lời đã có.
02 / Tính xác suất
Từ ngữ cảnh đó, mô hình tính phân bố xác suất cho những token có thể xuất hiện tiếp theo.
03 / Chọn token
Cách giải mã quyết định token được chọn; không phải lúc nào cũng lấy token có xác suất cao nhất.
04 / Nối và lặp
Token được nối vào chuỗi để tạo ngữ cảnh mới. Quá trình lặp đến điều kiện dừng.
Các ô dùng từ nguyên vẹn để dễ quan sát; đây không phải kết quả của một tokenizer cụ thể. Token thực tế có thể chỉ là một phần của từ.
Sơ đồ minh họa vòng sinh tiếp từng token. Nó không hiển thị vùng nhớ KV hay phân chia prefill/decode; hãy coi cache là trạng thái bổ sung phục vụ vòng lặp đó, không phải một token đặc biệt trong câu.
Tính lượng bộ nhớ theo từng thành phần
Một ước lượng cơ bản cho cache đầy đủ là: byte = 2 × L × B × T × Hkv × Dhead × S. L là số lớp, B là số chuỗi, T là số token giữ lại, Hkv là số đầu key/value, Dhead là chiều mỗi đầu, S là byte mỗi phần tử. Hệ số hai dành cho cả K và V.
Lấy L = 32, B = 1, T = 4096, Hkv = 8, Dhead = 128 và S = 2. Kết quả là 536.870.912 byte, bằng 512 MiB. Bốn chuỗi dài tương đương cần khoảng 2 GiB riêng cho cache theo giả định này.
Đây chưa phải tổng bộ nhớ chạy. Còn trọng số, activation tạm, workspace kernel, bộ đệm và chi phí cấp phát. Grouped-query attention có thể có nhiều query head hơn KV head, nên dùng số query head thay Hkv sẽ tính sai. Các cơ chế sliding window hoặc cache lượng tử hóa cũng cần công thức phù hợp hơn.
Chia sẻ tiền tố không phải phép sao chép tùy ý
Nếu nhiều yêu cầu dùng cùng tiền tố, hệ thống có thể tái sử dụng phần cache tương ứng khi triển khai hỗ trợ. Tiền tố phải khớp token, vị trí và điều kiện mô hình. Cùng văn bản nhìn thấy chưa chắc tạo cùng token nếu chat template hoặc tokenizer khác nhau.
Cũng cần kiểm tra cách hệ thống cô lập các phiên và hủy trạng thái hết hạn. Cache tối ưu thời gian tính không phải lý do để đưa dữ liệu của người này vào ngữ cảnh người khác. Quy tắc truy cập nguồn dữ liệu vẫn thuộc ứng dụng.
Ngân sách cache đặt ra đánh đổi giữa độ dài ngữ cảnh và số yêu cầu đồng thời. Thu ngắn prompt, giới hạn đầu ra, thay chiến lược batch hoặc thay biểu diễn cache đều có tác động khác nhau; hãy đo trên phân bố chiều dài thật trước khi chọn.
Bài tập lập giới hạn phục vụ
Với cấu hình trên, tăng T lên 8192 và B lên 4. Cache phải bằng 4 GiB. Sau đó đổi Hkv từ 8 xuống 4, giữ các giả định khác: kết quả giảm còn 2 GiB. Nếu bạn nhận mức giảm theo bình phương T, bạn đã nhầm kích thước cache với ma trận điểm attention.
Lập bảng riêng cho trọng số, cache và phần dự phòng chưa đo. Tiêu chí hoàn thành là không gọi tổng ước lượng ấy là số đo thực nghiệm, đồng thời chỉ rõ phép đo peak memory cần chạy. Đọc FlashAttention và IO để hiểu một nguồn tiết kiệm bộ nhớ khác.