49 · RadixAttention, eviction và cache churn
Nguồn: fanout.sh / inference-eng / 10-07 · ~6 phút
Ghi chú
Prompt thật chia sẻ theo từng lớp, nên các phần chung tạo thành một cái cây. RadixAttention (SGLang) cache cái cây đó trong một radix tree mà mỗi cạnh chứa một dãy token cùng K, V của chúng, so khớp tới từng token một; request đã xong, kể cả đầu ra, ở lại cho lượt kế tiếp. Cache là hữu hạn: SGLang đuổi lá không bị ghim và lâu nhất không được dùng (LRU) trước, nên thân cây dùng chung sống lâu nhất, và phục vụ prefix đã cache dài nhất trước sẽ duyệt cây theo chiều sâu để có hit rate tốt nhất (được báo tới 6,4 lần throughput). Khi working set lớn hơn cache, LRU có thể churn: 5 prefix xoay vòng qua một cache chứa được 4 cho 0 hit trên 10.
1. Chia sẻ theo từng lớp
Một radix tree chứa cả ba khuôn mẫu; các lượt chat được nén thành một cạnh duy nhất cho tới khi có thứ rẽ nhánh khỏi nó.
- Một system prompt, rồi vài ví dụ, rồi những câu hỏi khác nhau.
- Một cuộc chat lớn dần từng lượt một.
- Một agent rẽ thành nhiều lần thử từ cùng một điểm.
- Các phần chung tạo thành một cái cây, và cache chứa nó là hữu hạn: phải có thứ gì đó quyết định giữ cái gì.
2. RadixAttention: cache cái cây
◆ Dựng bằng một radix tree ~30 dòng trên token là từ; các con số là đầu ra thật của nó.
- Đến từ dự án SGLang. K và V đã cache nằm trong một radix tree: mỗi cạnh là một dãy token và giữ KV cache của các token đó.
- Request 1 (system prompt + Q1) thành một đường đi duy nhất từ gốc.
- Request 2 dùng chung system prompt nhưng hỏi Q2. Lượt tra cứu đi xuống từng token cho tới khi token khác nhau; tại đó cạnh tách ra thành một phần chung và hai nhánh. Phần chung được dùng lại; chỉ Q2 phải tính.
- Khi một request xong, cả đường đi của nó, kể cả câu trả lời đã sinh, ở lại trong cây, nên lượt kế tiếp thấy toàn bộ lịch sử đã được cache.
3. So với prefix caching của vLLM
Số token được dùng lại theo độ dài prefix chung.
| vLLM | SGLang RadixAttention | |
|---|---|---|
| Cấu trúc | chuỗi hash của các block | radix tree |
| Độ mịn khi so khớp | bậc 16 token | tới từng token |
| Chia sẻ | reference count | reference count |
| Eviction | lâu nhất không dùng đi trước | lâu nhất không dùng đi trước |
- Cùng mục tiêu: dùng lại mọi prefix chung, một cách tự động.
4. Eviction
◆ Các số thứ tự đến từ việc mô phỏng các quy tắc bên dưới trên cây này.
- Trong SGLang, cache và các request đang chạy dùng chung một pool bộ nhớ. Khi batch lớn hơn cần chỗ, cache co lại.
- Một số node không bao giờ bị đuổi: mỗi request đang chạy ghim đường đi của nó bằng reference count. Chỉ node có số đếm 0 mới là ứng viên.
- Chỉ lá mới bị đuổi trước, và trong số đó, lá lâu nhất không được dùng bị đuổi.
- Khi mọi con của một node đã đi hết, nó thành lá và có thể bị đuổi sau. Vì vậy thân cây dùng chung sống lâu nhất, đúng như điều bạn muốn.
5. Thứ tự quan trọng
◆ Mô phỏng cache LRU chứa được một prefix, cùng 9 request theo hai thứ tự.
- Các request đang chờ thường chia thành các nhóm dùng chung những prefix khác nhau. Phục vụ theo thứ tự đến, cache cứ phải đổi qua lại giữa các nhóm.
- SGLang có thể phục vụ prefix đã cache dài nhất trước. Cách đó duyệt cây theo chiều sâu, xong một nhánh rồi mới sang nhánh kế.
- Với một batch cố định, paper chứng minh thứ tự này cho hit rate tốt nhất có thể. ◆ Định lý giả định cache chứa được ít nhất request dài nhất.
- Rủi ro là bỏ đói (starvation): một request không có prefix chung có thể cứ bị đẩy lùi.
- SGLang (NeurIPS '24) báo throughput tới 6,4 lần so với các hệ thống khác cùng thời, trên các workload đầy cấu trúc chia sẻ.
6. Cache churn
Trái: ví dụ của bài, mô phỏng lại. Phải ◆: cùng chu kỳ đó lặp 200 lần ở các kích thước cache khác nhau.
- Cache chứa được 4 prefix; 5 prefix đến theo chu kỳ lặp: A, B, C, D, E.
- E không vừa, nên LRU đuổi A. Rồi A quay lại và đuổi B; B quay lại và đuổi C… 0 hit trên 10 request, dù prefix nào cũng lặp lại.
- Gom cùng lưu lượng đó theo prefix (A, A, B, B, …) thì một nửa trúng (5/10).
Cảnh báo
◆ Với truy cập theo chu kỳ, LRU luôn đuổi đúng prefix sắp cần tới. Vì vậy hit rate giữ ở 0% cho tới khi cache chứa được toàn bộ working set, rồi nhảy lên 100%. Thiếu một prefix là mất tất cả.
Triệu chứng và cách sửa, như liệt kê trong bài.
7. Tóm tắt
