Bộ nhớ đệm KV phân tầng cho các mô hình ngôn ngữ lớn trên Amazon SageMaker HyperPod với Curvine

Vận hành suy luận mô hình ngôn ngữ lớn (large language model – LLM) ở quy mô lớn thường buộc phải đánh đổi với bộ nhớ đệm KV: hoặc trả tiền cho các phiên bản bộ xử lý đồ họa (graphics processing unit – GPU) quá khổ để chứa bộ nhớ đệm KV ngày càng lớn, hoặc chấp nhận thởi gian tạo token đầu tiên (time-to-first-token – TTFT) chậm vì các prompt giống hệt nhau bị tính toán lại ở mỗi yêu cầu. Với các đội ngũ triển khai một danh mục đa dạng các mô hình nền tảng (foundation model – FM) công khai như Qwen, Llama, DeepSeek và nhiều mô hình khác, trên các endpoint theo từng mảng kinh doanh, các pipeline tạo sinh tăng cường truy xuất (Retrieval Augmented Generation – RAG), hay các ứng dụng hội thoại nhiều lượt, sự đánh đổi này trực tiếp dẫn đến chi phí hạ tầng cao hơn và trải nghiệm người dùng bị suy giảm. Nguyên nhân gốc rễ là trong quá trình sinh, vLLM lưu trữ các khóa (key) và giá trị (value) của cơ chế chú ý cho mọi token trong bộ nhớ đệm KV, và prefix caching tái sử dụng bộ nhớ đệm đó giữa các yêu cầu có chung phần token đầu. Trên các phiên bản tiết kiệm chi phí như ml.g6e.4xlarge (48 GB mỗi GPU), bộ nhớ còn lại cho prefix caching rất hạn chế, tỷ lệ trúng cache giảm khi prompt dài, các system prompt giống hệt nhau bị prefill lại ở mỗi yêu cầu, và các bản sao vLLM mở rộng theo chiều ngang đều duy trì bộ nhớ đệm riêng biệt. Giải pháp là xây dựng một hệ thống phân cấp bộ nhớ đệm ba tầng: L0 (bộ nhớ đệm tiền tố trên GPU), L1 (offload sang bộ nhớ của bộ xử lý trung tâm – CPU), và L2 (pool NVMe phân tán dùng chung thông qua Curvine). L0 là tầng paged-attention gốc của vLLM, L1 dùng LMCache để offload các khối GPU bị loại bỏ sang bộ nhớ truy cập ngẫu nhiên động (dynamic random-access memory – DRAM) của máy chủ, còn L2 gom các ổ NVMe (Non-Volatile Memory Express) cục bộ thành một không gian tên dùng chung qua Curvine, được gắn (mount) dưới dạng PVC ReadWriteMany (PersistentVolumeClaim) vào mọi Pod suy luận. Intelligent Routing của HyperPod điều hướng yêu cầu tới các bản sao có khả năng cao nhất đang chứa các khối KV liên quan, với các chiến lược gồm prefix-aware, kv-aware và round-robin. Hiệu quả tổng thể: chỉ khi trượt hoàn toàn (miss), hệ thống mới phải prefill lại từ đầu, nhờ đó giảm đáng kể TTFT cho các khối lượng công việc có mức trùng lặp prompt từ trung bình đến cao. Trong một lần triển khai thử nghiệm, giải pháp đạt tỷ lệ trúng cache liên Pod lên tới 100%, cải thiện TTFT tới 2,7 lần, và độ trễ đọc L2 liên nút khoảng 56 ms cho một prompt khoảng 1.900 token. Với kiến trúc này, các khối lượng công việc trước đây đòi hỏi phiên bản P5 nay có thể chạy trên các phiên bản G6e chi phí thấp hơn, qua đó giảm chi phí trên mỗi endpoint. Việc triển khai yêu cầu SageMaker HyperPod đã bật Tiered Storage, một cụm EKS, và cài đặt Inference Operator cùng Curvine. Bài viết trình bày năm giai đoạn: bật Tiered Storage, cài đặt Inference Operator và các thành phần phụ thuộc, cài đặt Curvine, vá (patch) Inference Operator để hỗ trợ L2 dựa trên hệ thống tệp, và chạy benchmark.