KV cache phân trang, hay vì sao đặt trước theo độ dài tối đa là hoang phí
KV cache phân trang, hay vì sao đặt trước theo độ dài tối đa là hoang phí
Bài trước tính ra bộ đệm tốn bao nhiêu byte. Bài này hỏi câu tiếp theo, và là câu quyết định số người bạn phục vụ được: cấp phát đống byte đó ra sao khi nhiều yêu cầu chạy cùng lúc.
Ở bài KV cache bạn đã tính được cái giá: mỗi token trong bộ đệm ngốn một số byte cố định, và với hình dạng của Llama 3.1 8B thì con số đó là 128 KiB. Bài đó dừng ở một chuỗi. Máy chủ thật thì không: nó phục vụ hàng chục yêu cầu cùng lúc, và mỗi yêu cầu cần một vùng bộ đệm riêng.
Câu hỏi là vùng đó được cấp ra sao. Cách hiển nhiên nhất, và cũng là cách mọi bản cài đặt đầu tiên đều làm, là đặt trước một khối liền mạch theo độ dài tối đa. Bộ đệm cần liền mạch vì phép nhân chú ý muốn đọc K và V của cả quá khứ trong một dải bộ nhớ, và phải theo độ dài tối đa vì lúc yêu cầu mới đến, không ai biết nó sẽ sinh ra bao nhiêu chữ.
Nghe hợp lý. Nhưng hãy nhìn cái giá của nó.
| # | độ dài ✎ | số trang | phân trang giữ | trang cuối | thừa | ngây thơ giữ | ngây thơ phí |
|---|---|---|---|---|---|---|---|
| 1 | 26 | 416 | 3/16 | 13 | 512 | 109 | |
| 2 | 6 | 96 | 16/16 | 0 | 512 | 416 | |
| 3 | 16 | 256 | 1/16 | 15 | 512 | 271 | |
| 4 | 2 | 32 | 14/16 | 2 | 512 | 482 | |
| 5 | 11 | 176 | 16/16 | 0 | 512 | 336 | |
| tổng | 946 | 61 | 976 | 2 trọn | 30 | 2.560 | 1.614 |
Phân mảnh trong, và nó lớn cỡ nào
Bảng mở ra với năm yêu cầu dài 403, 96, 241, 30 và 176 token, tổng cộng 946 token dữ liệu thật. Mức đặt trước là 512 token mỗi yêu cầu, nên máy chủ giữ:
5 × 512 = 2.560 token
Trong đó chỉ 946 token có dữ liệu. 1.614 token, tức 63,0% bộ nhớ đã cấp, nằm không. Ở giá 128 KiB mỗi token thì đó là 201,8 MiB bị khoá lại mà không ai dùng, trong tổng số 320,0 MiB đã cấp phát.
Đây gọi là phân mảnh trong: bộ nhớ bị lãng phí nằm bên trong vùng đã cấp cho một yêu cầu, chứ không phải rơi rớt giữa các vùng. Bộ cấp phát không thể lấy lại nó cho ai khác, vì nó vẫn thuộc về yêu cầu kia, phòng khi yêu cầu đó còn sinh tiếp.
Nhìn lưới bên trái để thấy hình dạng của sự lãng phí: mỗi ô là một trang, mỗi khối 32 ô là vùng đặt trước của một yêu cầu, và 99 trong số 160 ô là ô nét đứt, tức đã đặt trước nhưng chưa có gì trong đó.
Bấm tải mẫu Một yêu cầu rất dài, còn lại rất ngắn để thấy trường hợp xấu nhất: một yêu cầu 1.990 token kéo mức đặt trước lên 2.048, năm yêu cầu còn lại chỉ vài chục token nhưng vẫn phải trả theo mức đó. Lãng phí lên 81,9%.
Phân trang: bỏ luôn ràng buộc liền mạch
Ý tưởng đến thẳng từ bộ nhớ ảo của hệ điều hành, và nó chỉ có đúng một câu: chia bể nhớ thành các trang cỡ cố định, rồi cho mỗi yêu cầu giữ một bảng ánh xạ trang thay vì một khối liền mạch. Các trang của cùng một yêu cầu không cần nằm cạnh nhau, vì bảng ánh xạ biết trang nào ở đâu.
Hệ quả là yêu cầu chỉ giữ đúng số trang nó cần ngay lúc này:
số trang = trần(độ dài / cỡ trang)
Với cỡ trang 16 token, năm yêu cầu mặc định cần 26, 6, 16, 2 và 11 trang, tổng 61 trang tức 976 token. So với 2.560 token của cách ngây thơ, bộ nhớ giữ lại tụt còn 122,0 MiB thay vì 320,0 MiB.
Phần lãng phí còn lại không biến mất hẳn, nó chỉ co lại thành cái đuôi của trang cuối. Yêu cầu 403 token cần 26 trang, tức 416 token, nên 13 token cuối trong trang thứ 26 vẫn trống. Cộng cả năm yêu cầu, phần thừa là 30 token: 3,1% thay vì 63,0%.
Bảng có một quy tắc kiểm tra rất dễ nhớ, và cổng kiểm số của bài canh đúng nó ở mọi cấu hình:
token dùng thật + phần thừa của mọi trang cuối = số trang × cỡ trang
946 + 30 = 61 × 16 = 976
Hai mốc đáng tự thử. Đặt cỡ trang bằng 1: phần thừa về đúng 0, vì không có trang nào có đuôi. Đặt cỡ trang bằng 512, đúng bằng mức đặt trước: mỗi yêu cầu vừa vặn một trang và phân trang trùng khớp cách ngây thơ, cả hai giữ 2.560 token và cùng lãng phí 63,0%. Phân trang không phải phép màu, nó chỉ là hạ đơn vị cấp phát xuống. Hạ càng sâu thì phân mảnh trong càng nhỏ.
Con số thật sự đáng quan tâm: phục vụ được bao nhiêu người
Tiết kiệm phần trăm nghe hay nhưng trừu tượng. Điều người vận hành hỏi là: cùng một cái thẻ, tôi chạy được bao nhiêu yêu cầu cùng lúc.
Bảng chia bể nhớ ra thành chỗ chứa token. Với 2 GiB dành cho bộ đệm và 128 KiB mỗi token:
2 × 1.073.741.824 / 131.072 = 16.384 chỗ
Cách ngây thơ tiêu 512 chỗ cho mỗi yêu cầu bất kể yêu cầu đó dài bao nhiêu, nên 16.384 / 512 = 32 yêu cầu. Cách phân trang tiêu đúng số trang mỗi yêu cầu cần, trung bình 195,2 chỗ, nên bể nhớ nhận được 83 yêu cầu.
32 so với 83, tức 2,6 lần. Cùng một phần cứng, cùng một mô hình, cùng một chỗ bộ nhớ. Khác nhau chỉ ở cách chia. Đây là lý do phân trang bộ đệm là thay đổi hạ tầng phục vụ đáng kể nhất mấy năm qua, và cũng là lý do bạn thấy thông lượng của các máy chủ suy luận nhảy vọt mà kiến trúc mô hình không đổi gì.
Cách bảng đếm con số đó cần nói rõ, vì nó là một quy ước chứ không phải chân lý: các yêu cầu được coi là đến lặp lại theo đúng thứ tự trong danh sách và không bao giờ kết thúc, rồi bảng nhận vào lần lượt cho tới khi hết chỗ. Vì thế thứ tự danh sách có ảnh hưởng tới lượt cuối cùng: ở trạng thái mặc định, sau 16 vòng trọn vẹn còn dư 768 chỗ, vừa đúng cho ba yêu cầu đầu (416 cộng 96 cộng 256 bằng đúng 768) và yêu cầu thứ tư cần 32 chỗ nữa thì không còn. Đổi thứ tự danh sách, con số cuối có thể lệch một vài đơn vị. Đó là tính chất thật của mô hình, không phải chỗ làm tròn giấu đi.
Chia sẻ trang: chỗ mà phân trang làm được việc cách kia không làm nổi
Có một món quà đi kèm mà cách khối liền mạch không thể có. Nếu nhiều yêu cầu cùng bắt đầu bằng một đoạn giống hệt nhau, ví dụ cùng một prompt hệ thống dài, cùng một tài liệu nền, cùng vài ví dụ mẫu, thì phần K và V của đoạn đó là y hệt nhau. Đã phân trang thì các bảng ánh xạ chỉ việc trỏ vào cùng những trang đó. Tính một lần, giữ một bản.
Nhưng chỉ chia sẻ được tới ranh giới trang, và đây là chi tiết hay bị bỏ qua:
số trang tiết kiệm = (số yêu cầu - 1) × sàn(độ dài tiền tố / cỡ trang)
Hai chỗ cần để ý trong công thức này. Thứ nhất là - 1: bản đầu tiên vẫn phải tồn tại, chỉ các bản sao mới được miễn. Thứ hai là hàm sàn: chỉ những trang trọn vẹn của tiền tố mới dùng chung được, vì trang cuối của tiền tố còn chứa cả token đầu tiên mà các yêu cầu khác nhau, mà đã khác một token thì K và V của cả trang đó không còn dùng chung được nữa.
Ở trạng thái mặc định, ba yêu cầu đầu dùng chung 64 token đầu với trang 16 token, nên mỗi bản chia sẻ được 4 trang trọn vẹn và tiết kiệm (3 - 1) × 4 = 8 trang, tức 128 token. Bộ nhớ giữ lại tụt từ 976 xuống 848 token.
Hãy tự thử mốc quan trọng nhất: đặt độ dài tiền tố thành 15 với cỡ trang 16. Phần tiết kiệm về đúng 0, không phải "gần 0". Tiền tố ngắn hơn một trang thì không chia sẻ được gì cả. Rồi tăng lên 16 và nó nhảy lên 2 trang ngay.
Cái bảng này không nói gì
Đây là chỗ phải nói thẳng, vì một hoá đơn tính đúng vẫn có thể bị đọc quá xa.
Nó đếm đúng phần kế toán bộ nhớ, và chỉ có thế. Với danh sách yêu cầu, cỡ trang và mức đặt trước bạn nhập, các con số về trang, phần thừa, tỉ lệ lãng phí và số chỗ trong bể là chính xác từng token. Cổng kiểm số của bài canh đúng những phép tính đó.
Nó bỏ qua chi phí của chính bảng ánh xạ trang. Mỗi yêu cầu phải giữ một danh sách số hiệu trang, và bộ cấp phát phải giữ danh sách trang rỗng. Cỡ trang càng nhỏ thì bảng càng dài, nên cỡ trang 1 token tuy cho lãng phí 0 trong bảng này lại là lựa chọn tệ trong thực tế. Không có byte nào của phần đó được tính ở đây.
Nó bỏ qua cái giá của việc đọc bộ nhớ rời rạc. Nhân chú ý chạy nhanh nhất khi đọc một dải liền mạch. Bộ đệm phân trang bắt nó đi qua bảng ánh xạ và nhảy giữa các trang, nên phải viết lại nhân cho tử tế thì mới không mất phần lớn cái vừa tiết kiệm được. Bảng này không đo được một nano giây nào.
Nó không có gì về lập lịch. Máy chủ thật còn phải quyết định nhận yêu cầu nào trước, có tạm dừng yêu cầu đang chạy để nhường chỗ không, có đẩy bộ đệm ra bộ nhớ chính rồi kéo về không. Con số phục vụ đồng thời ở đây là một phép chia bộ nhớ, không phải thông lượng. Đừng đọc nó thành "máy chủ của tôi sẽ nhanh gấp 2,6 lần".
Và phân trang không phải lúc nào cũng thắng đậm. Bấm tải mẫu Mọi yêu cầu gần bằng mức đặt trước: lãng phí của cách ngây thơ chỉ còn 2,1%, phân trang cho 1,5%, tức tiết kiệm được đúng 16 token. Khi độ dài thật sát với mức đặt trước thì phân mảnh trong vốn đã nhỏ sẵn, và mọi phức tạp thêm vào gần như không mua được gì. Cái quyết định không phải là phân trang hay không, mà là độ dài yêu cầu lệch nhau tới đâu.
Đặt trước theo độ dài tối đa nghĩa là mọi yêu cầu đều trả giá cho yêu cầu dài nhất, nên khi độ dài lệch nhau, phần lớn bộ nhớ đã cấp sẽ nằm không. Phân trang hạ đơn vị cấp phát từ "cả một chuỗi tối đa" xuống "một trang cỡ cố định", nên phần lãng phí co lại còn đúng cái đuôi của trang cuối, và bể nhớ cũ nhận được nhiều yêu cầu hơn hẳn. Đổi lại, bạn nhận thêm một bảng ánh xạ phải nuôi và một nhân chú ý phải viết lại. Còn hai công thức thì chỉ có bấy nhiêu: trần(độ dài / cỡ trang) cho phần cấp phát, và (số yêu cầu - 1) × sàn(độ dài tiền tố / cỡ trang) cho phần chia sẻ.
- 1Cỡ trang 16 token. Một yêu cầu dài 100 token thì được cấp mấy trang, và phần thừa của trang cuối là bao nhiêu?
- 2Bốn yêu cầu cùng bắt đầu bằng một prompt hệ thống dài 40 token, cỡ trang 16. Chia sẻ trang tiết kiệm được bao nhiêu trang?
- 3Ở trạng thái mặc định, bể nhớ 16.384 chỗ nhận được 32 yêu cầu theo cách ngây thơ và 83 yêu cầu theo cách phân trang. Vì sao lại chênh nhau như vậy?