Tăng tốc SGLang HiCache với bộ xử lý bộ nhớ Netpreme X-Mem™

Netpreme X-Mem™ Memory Processing Unit (MPU) giúp SGLang HiCache trở nên nhanh hơn và có khả năng mở rộng tốt hơn bằng cách bổ sung tầng lưu trữ Host DRAM chậm hơn bằng một tầng bộ nhớ KV băng thông cao được thiết kế chuyên biệt.

Tóm tắt (TL;DR)

Tại sao Prefix Caching cần một tầng bộ nhớ tốc độ cao

Prefix caching được thiết kế để giảm thiểu các tính toán prefill tốn kém cho các khối lượng công việc mà các yêu cầu liên tục chia sẻ các prompt dài, giống hệt nhau về token và chỉ khác biệt ở phần cuối ngắn. SGLang đã xây dựng một nền tảng phần mềm vững chắc cho việc tái sử dụng prefix với RadixAttention và HiCache. RadixAttention cho phép thực hiện prefix caching hiệu quả trong bộ nhớ GPU, trong khi HiCache mở rộng khả năng này thành một hệ thống phân cấp vượt ra ngoài GPU HBM.

Tuy nhiên, khi GPU ngày càng nhanh hơn, tốc độ tải KV cache vào HBM trở nên quan trọng. Để ước tính yêu cầu băng thông cho việc offload KV-cache, chúng tôi mô phỏng một khối lượng công việc coding-agent với ngữ cảnh dài, gồm các prompt 64K-token và độ đồng thời prefill là 8. Khi tỷ lệ hit của prefix tăng lên, các tính toán dư thừa giảm nhanh chóng, nhưng lượng dữ liệu KV được lấy từ bộ nhớ thứ cấp lại tăng lên. Khi tỷ lệ hit vượt quá 95%, việc mở rộng băng thông của tầng KV-cache có thể giảm đáng kể TTFT!

Nhưng liệu tỷ lệ hit 95% của prefix cache có thực sự khả thi? Hãy cùng xem xét việc tái sử dụng KV-cache trong các coding agent — thị trường lớn nhất cho các ứng dụng AI hiện nay.

Để định lượng điều này, chúng tôi đo lường tỷ lệ hit của prefix-cache bằng cách sử dụng các dấu vết (traces) thu thập được từ Claude Code agent khi chạy các quy trình lập trình đa lượt trên SWE-bench. Chúng tôi quan sát thấy rằng tỷ lệ cache-hit liên tục vượt quá 95% trong các vòng tiếp theo cho cùng một phiên làm việc, với tỷ lệ hit trung bình khoảng 98%. Trong điều kiện này, việc mở rộng băng thông KV-caching vượt xa những gì có thể đạt được với Host CPU offloading sẽ mang lại hiệu suất tăng gấp 3 lần!

Đây chính là lúc Netpreme X-Mem™ phát huy tác dụng!

Netpreme X-Mem™: Một tầng KV chuyên dụng cho SGLang HiCache

Netpreme X-Mem™ MPU là một giải pháp được xây dựng chuyên biệt để mở rộng bộ nhớ GPU thêm hàng chục TB thông qua các cấu trúc mạng tốc độ cao. Về mặt khái niệm, X-Mem™ là một node bộ nhớ chuyên dụng trong tủ rack AI, hoạt động ngang hàng (peer-to-peer) với tất cả các GPU khác. Một node bộ nhớ đơn lẻ đi kèm với dung lượng lên tới 24 TB, có thể được truy cập bởi tất cả các GPU trong cùng một tủ rack với tốc độ 4 TB/s. Không có giới hạn về số lượng node bộ nhớ có thể được cung cấp trong tủ rack, và tổng băng thông sẽ được nhân lên nếu sử dụng nhiều node MPU. Về phía phần mềm, không gian địa chỉ của X-Mem™ được hiển thị như một phần của bộ nhớ ảo thống nhất, và về mặt ngữ nghĩa, nó không khác gì việc truy cập vào HBM cục bộ hoặc bộ nhớ của GPU từ xa.

Khi tích hợp với SGLang HiCache, Netpreme X-Mem™ đóng vai trò là một tầng bộ nhớ chuyên dụng để offload KV cache. Thay vì sử dụng Host DRAM hoặc RDMA làm tầng L2 KV cache chính, MPU cung cấp một tầng bộ nhớ chuyên biệt, hướng tới bộ tăng tốc (accelerator-oriented) được tối ưu hóa cho việc di chuyển dữ liệu KV.

Điều này làm cho X-Mem™ đặc biệt hữu ích cho các khối lượng công việc với:

Netpreme X-Mem™ tích hợp với SGLang thông qua các API tương thích với CUDA và PyTorch. Các API này cho phép bất kỳ ứng dụng ML nào khai thác tầng bộ nhớ băng thông cao một cách minh bạch và chỉ cần sửa đổi tối thiểu mã nguồn ứng dụng.

Đánh giá SGLang + Netpreme X-Mem™

Chúng tôi thực hiện hai thử nghiệm, một là micro-benchmark đo lường TTFT của một yêu cầu đơn lẻ và thử nghiệm còn lại là thử nghiệm suy luận LLM toàn diện (end-to-end) với các khối lượng công việc dạng agent. Trong suốt quá trình đánh giá, chúng tôi so sánh ba cấu hình:

X-Mem™ làm phẳng TTFT của một yêu cầu đơn lẻ

Kết quả cho thấy MPU của Netpreme giảm TTFT của yêu cầu đơn lẻ lên tới ~6,7 lần so với Host DRAM, làm cho nó gần như phẳng ngay cả trong ngữ cảnh rất dài.

X-Mem™ tăng tính tương tác và công suất hệ thống của engine phục vụ LLM

Việc giảm TTFT nêu trên chuyển đổi như thế nào thành những cải tiến trong suy luận LLM toàn diện? Để trả lời câu hỏi này, chúng tôi chạy benchmark phục vụ LLM toàn diện, NVIDIA AIPerf. Chúng tôi sử dụng các khối lượng công việc đại diện cho suy luận AI dạng agent, bao gồm 1K token cho system prompt và 20K token cho ngữ cảnh mỗi người dùng. Mỗi người dùng nhập 26 token trong một lượt, và số lượt trung bình là 20. Các điểm đặc (solid points) biểu thị các điểm tối ưu Pareto; các điểm mờ (shaded points) biểu thị các điểm không tối ưu Pareto.

Kết quả chứng minh rằng X-Mem™ của Netpreme cung cấp TPS/người dùng (tính tương tác) cao hơn 33% trong trường hợp tải trung bình và tính tương tác cao hơn 50% cùng TPS cao hơn 30% (công suất hệ thống) trong trường hợp tải cao. Điều này là do GPU được sử dụng hiệu quả hơn, dành ít thời gian hơn để chờ dữ liệu được sao chép khi có prefix cache hit.

Hướng tới tương lai

Khi các coding agent, trợ lý ngữ cảnh dài và hệ thống gợi ý phát triển, các hệ thống phục vụ sẽ cần phải lưu giữ và di chuyển các tập dữ liệu KV làm việc lớn hơn nhiều. Để đạt được mục tiêu này, Netpreme sẽ tạo điều kiện cho:

Tầm nhìn dài hạn của Netpreme là cung cấp một giải pháp bộ nhớ toàn diện (full-stack) cho các khối lượng công việc dạng agent. Điều này không chỉ bao gồm KV cache mà còn cả trọng số mô hình (model weights), các giá trị kích hoạt (activations), embeddings và các cấu trúc dữ liệu khác. Bằng cách cung cấp một tầng bộ nhớ hiệu năng cao được tối ưu hóa cho các khối lượng công việc ML, Netpreme cho phép các kiến trúc phần mềm và các tối ưu hóa mới mà các giải pháp Host DRAM và GPU HBM truyền thống không thể thực hiện được.

Kết luận

SGLang đã cung cấp một nền tảng phần mềm vững chắc cho prefix caching. RadixAttention và HiCache giúp việc tái sử dụng KV cache trở nên khả thi trên các khối lượng công việc có ngữ cảnh dài, đa lượt và sử dụng nhiều prefix. Nhưng khi prefix caching trở nên hiệu quả hơn, nút thắt cổ chai sẽ dịch chuyển.

Thách thức không còn chỉ là tìm kiếm các prefix có thể tái sử dụng. Đó là việc di chuyển KV đã cache trở lại bộ nhớ GPU đủ nhanh để việc tái sử dụng mang lại sự cải thiện thực sự về độ trễ và thông lượng.

Việc offload dựa trên Host DRAM giúp tăng dung lượng, nhưng băng thông ở mức PCIe của nó có thể trở thành yếu tố hạn chế đối với suy luận đòi hỏi nhiều KV, đặc biệt là khi TTFT SLO bị thắt chặt. Netpreme X-Mem™ giải quyết nút thắt này bằng cách cung cấp cho SGLang HiCache một tầng bộ nhớ KV băng thông cao được xây dựng chuyên biệt.

Cùng với nhau, SGLang HiCache và Netpreme X-Mem™ giúp prefix caching nhanh hơn và có khả năng mở rộng tốt hơn cho thế hệ khối lượng công việc LLM tiếp theo, bao gồm coding agent, trợ lý ngữ cảnh dài và các hệ thống gợi ý.

Lời cảm ơn

Chúng tôi cảm ơn đội ngũ SGLang vì những thảo luận kỹ thuật mang tính xây dựng và sự hướng dẫn của họ trong suốt quá trình tích hợp này.

Tài liệu tham khảo

Phụ lục

Thiết lập đánh giá