Cloudflare tối ưu hóa Kimi và GLM trên Workers AI: Kỹ thuật nén trọng số và quản lý bộ nhớ đệm

Workers AI thực hiện suy luận cho một số mô hình mã nguồn mở tốt nhất thế giới trên các GPU tại các trung tâm dữ liệu của Cloudflare đặt gần người dùng của bạn. Hai trong số những mô hình mạnh mẽ và đòi hỏi khắt khe nhất là dòng Kimi K của Moonshot và GLM của Z.ai. Đây là các mô hình mixture-of-experts (hỗn hợp chuyên gia) có ngữ cảnh dài và kích thước lớn, mang lại trải nghiệm sử dụng tuyệt vời. Tuy nhiên, chúng cũng rất khó để phục vụ hiệu quả do các hạn chế về bộ nhớ.

Chúng tôi đã từng viết về cách chúng tôi phục vụ các mô hình lớn trên Workers AI cũng như việc tách biệt giai đoạn prefill (tiền nạp) và decode (giải mã) trong quá trình suy luận để tận dụng tối đa mỗi GPU. Bài viết này xem xét ba kỹ thuật mà chúng tôi áp dụng trên nền tảng đó để đưa các mô hình này vào bộ nhớ và duy trì tốc độ xử lý nhanh: lượng tử hóa KV cache, nén trọng số mô hình, và bảo vệ bộ nhớ đệm mà các yêu cầu chia sẻ, vì cả hai kỹ thuật trước đó đều giúp dồn nhiều yêu cầu hơn vào phần cứng dùng chung. Những tối ưu hóa này cho phép chúng tôi hỗ trợ nhiều khách hàng hơn với chi phí thấp hơn mà không làm thay đổi độ chính xác của mô hình.

Tất cả các thử nghiệm và lưu lượng truy cập thực tế của chúng tôi đều đang chạy và được đo kiểm (benchmark) bằng SGLang, một framework phục vụ suy luận mã nguồn mở. Chúng tôi nhận thấy SGLang mang lại hiệu suất tốt nhất trên thị trường và chúng tôi hợp tác chặt chẽ với đội ngũ SGLang để đóng góp các bản vá và tính năng mới nhằm đưa công việc của chúng tôi đến với cộng đồng mã nguồn mở.

Lượng tử hóa KV cache

Khi một mô hình tạo văn bản, nó lưu trữ các khóa (K) và giá trị (V) chú ý (attention) cho mọi token mà nó đã xử lý trong một cấu trúc gọi là KV cache. Bộ nhớ đệm này là thứ cho phép mô hình kéo dài một cuộc hội thoại dài mà không cần đọc lại toàn bộ ngữ cảnh cho mỗi token mới. Đối với một mô hình có ngữ cảnh dài, bộ nhớ này tăng lên rất nhanh, và thường thì chính KV cache, chứ không phải trọng số của mô hình, là thứ làm đầy bộ nhớ GPU trước tiên.

Theo mặc định, bộ nhớ đệm được lưu trữ ở độ chính xác 16-bit (BF16). Thay vào đó, chúng tôi lưu trữ nó ở dạng dấu phẩy động 8-bit (FP8, e4m3), giúp giảm một nửa kích thước của nó. Trên Kimi K2.6, điều đó làm tăng lượng ngữ cảnh mà chúng tôi có thể giữ trong bộ nhớ từ khoảng 686.000 token lên khoảng 1,37 triệu, gấp đôi so với trước.

Cần phải làm rõ lợi ích đến từ đâu, vì đó không phải là tốc độ thuần túy. Lượng tử hóa bộ nhớ đệm làm tăng thêm một chút khối lượng công việc cho mỗi token, vì nhân (kernel) chú ý FP8 phải chuyển đổi các giá trị khi nó đọc chúng. Điều nó thay đổi là số lượng yêu cầu mà chúng tôi có thể giữ thường trú cùng lúc. Các phép đo sau đây dành cho quá trình giải mã Kimi K2.6 trên một hệ thống triển khai H200 phân tách, so sánh trực tiếp các nhân chú ý:

Số yêu cầu đồng thời

KV cache BF16 (tok/s)

KV cache FP8 (tok/s)

137

125

731

689

16

1.106

1.028

32

1.558

1.489

64

Hết bộ nhớ (Out of memory)

2.192

Ở bất kỳ mức độ đồng thời nào, BF16 nhanh hơn vài phần trăm trên mỗi token. Nhưng BF16 bị hết bộ nhớ đệm ở mức 32 yêu cầu đồng thời và không thể tiếp nhận yêu cầu thứ 33, trong khi FP8 vẫn tiếp tục hoạt động đến mức 64 và đạt 2.192 token mỗi giây, cao hơn khoảng 41% so với đỉnh của BF16, với chi phí trên mỗi token thấp hơn khoảng 30%. Vì chúng tôi chạy prefill và decode như các nhóm riêng biệt, chúng tôi có thể áp dụng điều này ở nơi nó mang lại hiệu quả cao nhất: prefill bị giới hạn bởi tính toán thay vì bộ nhớ, vì vậy ở đó chúng tôi giữ bộ nhớ đệm ở định dạng BF16 để duy trì thông lượng cao hơn một chút.

Tất cả những điều này sẽ không có ý nghĩa nếu nó làm thay đổi câu trả lời của mô hình, vì vậy chúng tôi đã kiểm tra. Trong bộ đánh giá của chúng tôi, bộ nhớ đệm FP8 và BF16 là không thể phân biệt được:

Benchmark

KV BF16

KV FP8

GSM8K

94,24

94,09

ARC-Easy

89,06

89,14

ARC-Challenge

66,72

67,49

MMLU

89,11