Unified Radix Cache: Giải pháp cây đơn nhất hóa bộ nhớ đệm cho mô hình hỗn hợp

Giới thiệu

Prefix caching (bộ nhớ đệm tiền tố) tái sử dụng KV khi các yêu cầu chia sẻ cùng một tiền tố token. Trong cơ chế full attention, một khi KV cho một tiền tố dùng chung đã được tính toán, nó vẫn giữ nguyên hiệu lực khi có thêm các token mới được nối vào. Một yêu cầu sau đó có cùng tiền tố có thể tái sử dụng các mục KV đó thay vì phải tính toán lại trong quá trình prefill. SGLang theo dõi ánh xạ này bằng một radix tree (cây cơ số) với khóa là các chuỗi token. Trước khi prefill, cây này tìm kiếm tiền tố có thể tái sử dụng dài nhất và trả về các vị trí KV của nó cho bộ lập lịch.

Các mô hình lai (hybrid models) phá vỡ quy tắc tái sử dụng đơn lẻ này. Một yêu cầu có thể kết hợp KV của full attention, KV của sliding window attention và các trạng thái tái phát (recurrent states), mỗi loại có một khái niệm tái sử dụng khác nhau. KV của full attention có thể tái sử dụng trên toàn bộ tiền tố khớp. KV của sliding window attention chỉ bao phủ cửa sổ phía sau. Một trạng thái tái phát chỉ có hiệu lực tại một điểm kiểm tra (checkpoint) tiền tố chính xác. Chúng chia sẻ cùng một tiền tố token nhưng không có cùng ranh giới tái sử dụng. Việc ép buộc một ranh giới duy nhất cho tất cả chúng sẽ dẫn đến việc loại bỏ khả năng tái sử dụng hợp lệ hoặc cho phép tái sử dụng không hợp lệ.

Các quy tắc tái sử dụng này xuất hiện dưới nhiều kết hợp khác nhau giữa các dòng mô hình. Việc mã hóa mỗi kết hợp thành một lớp cache chuyên biệt tạo ra một ma trận lớp tổ hợp, đặc biệt là khi các khả năng trực giao như HiCache được thêm vào. Các triển khai trước đây tuân theo mô hình này, dẫn đến việc trùng lặp logic khớp, chèn, khóa và giải phóng (eviction) giữa các biến thể cache.

Unified Radix Cache giải quyết thiết kế tổ hợp này bằng cách tách biệt danh tính tiền tố dùng chung khỏi tính hợp lệ của việc tái sử dụng theo thành phần. Một cấu trúc liên kết radix duy nhất dựa trên khóa token cung cấp tọa độ chuẩn cho mỗi tiền tố, trong khi KV của full attention, KV của sliding window attention và các checkpoint của Mamba được gắn vào như các thành phần. HiCache là thành phần gốc trong cùng vòng đời, mở rộng danh tính thành phần trên GPU L1, Host L2 và tầng L3 bên ngoài. Các nhóm phụ trợ (auxiliary pools) có thể đi kèm với các thành phần như các sidecar mà không cần xác định ranh giới tái sử dụng mới. Các dòng mô hình mới có thể kết hợp các khả năng này mà không cần giới thiệu thêm một cây cache khác.

Điểm nổi bật

Một cây, các thành phần có thể kết hợp

Các mô hình lai kết hợp các giá trị có thể lưu vào cache, chia sẻ cùng một tiền tố token nhưng tuân theo các quy tắc tái sử dụng khác nhau. Unified Radix Cache ánh xạ tiền tố dùng chung vào một cấu trúc liên kết radix và mỗi quy tắc tái sử dụng vào một TreeComponent. UnifiedTreeCore thực hiện các cơ chế khớp, tách, chèn, khóa và giải phóng chung. UnifiedRadixCache điều phối các hoạt động của nhóm, trong khi mỗi thành phần chỉ xác định các ngữ nghĩa thay đổi.

Thành phần FULL luôn hiện diện. SGLang thêm thành phần SWA cho sliding window attention lai và thành phần MAMBA cho các lớp tái phát lai. Ví dụ, DeepSeek-V4 kết hợp FULL và SWA, Kimi-K3 kết hợp FULL và MAMBA cho trạng thái tái phát KDA của nó, và Inkling kết hợp cả ba thành phần trên cùng một cây. Một dòng mô hình mới có thể tái sử dụng một kết hợp thành phần hiện có. Nếu nó giới thiệu một quy tắc tái sử dụng mà tập hợp hiện tại không thể biểu diễn, SGLang có thể thêm một TreeComponent mới mà không cần tạo ra một triển khai cây khác.

FULL cung cấp khả năng tái sử dụng đường dẫn. Nó giữ KV cho mọi token trong tiền tố khớp và bảo vệ đường dẫn tổ tiên tương ứng. SWA cung cấp khả năng tái sử dụng cửa sổ. Nó yêu cầu một cửa sổ phía sau liên tục, trong khi các khe SWA cũ hơn có thể là các "tombstone" (đánh dấu xóa) trống mà các nút radix của chúng vẫn nằm trong cấu trúc liên kết dùng chung. MAMBA cung cấp khả năng tái sử dụng checkpoint. Nó yêu cầu một checkpoint tái phát tại biên tái sử dụng và sao chép trạng thái dùng chung vào một khe yêu cầu riêng trước khi thay đổi. Các thành phần này áp dụng các quy tắc khác nhau cho cùng một ranh giới ứng viên.

Tìm kiếm ranh giới tái sử dụng an toàn

Trong quá trình khớp tiền tố, UnifiedTreeCore đi theo đường dẫn FULL chuẩn và coi mỗi nút đã truy cập là một ranh giới ứng viên. Chỉ riêng khớp FULL là chưa đủ. Mỗi thành phần hoạt động tạo ra một trình xác thực (validator), và ranh giới tái sử dụng chỉ tiến lên khi tất cả các trình xác thực chấp nhận ứng viên đó. Việc từ chối không dừng quá trình duyệt vì một thành phần có thể chấp nhận một nút sau đó. Trong Hình 2, n1 và n2 vượt qua mọi trình xác thực, trong khi n3 và n4 thất bại ở ít nhất một kiểm tra thành phần. Quá trình duyệt đạt đến n4 nhưng giữ lại n2 là kết quả an toàn sâu nhất.

Sau khi duyệt, lõi xây dựng một MatchResult. Các trình hoàn thiện thành phần (component finalizers) sau đó chuẩn bị các giá trị đã chọn để tái sử dụng, bao gồm cả việc sao chép cần thiết khi một checkpoint MAMBA dùng chung trở thành riêng tư cho một yêu cầu.

Các hook thành phần trong suốt vòng đời của cây

Hợp đồng thành phần tương tự bao gồm phần còn lại của vòng đời cây:

Hợp đồng này giữ cho lõi cây ở dạng tổng quát trong khi cho phép các thành phần duy trì các quy tắc đúng đắn khác nhau. Việc xóa một payload thành phần không phải lúc nào cũng xóa nút radix. Cấu trúc liên kết còn lại vẫn có thể neo giữ các thành phần khác, và khe thành phần trống có thể tồn tại như một tombstone cho đến khi thành phần đó được khôi phục hoặc nút trở nên không cần thiết. Vì ngữ nghĩa thành phần vẫn gắn liền với một danh tính tiền tố, HiCache có thể mở rộng payload thành phần qua các tầng bộ nhớ mà không cần giới thiệu thêm một cây khác.

HiCache gốc trên các tầng bộ nhớ

Các thành phần xác định những gì có thể được tái sử dụng. HiCache xác định nơi payload có thể tái sử dụng cư trú. Unified Radix Cache mang cùng danh tính thành phần qua GPU L1, Host L2 và tầng L3 bên ngoài, vì vậy việc di chuyển dữ liệu giữa các tầng không làm thay đổi danh tính tiền tố hoặc quy tắc tái sử dụng của nó. Các thành phần mô tả các chuyển đổi cần thiết, và HybridCacheController thực thi I/O vật lý.

Các thành phần, neo (anchor) và sidecar

Không phải mọi nhóm vật lý đều cần thành phần riêng của nó. Một neo xác định ngữ nghĩa tái sử dụng hoặc cung cấp các chỉ mục trang mà các nhóm khác tuân theo. Một sidecar lưu trữ một payload riêng biệt nhưng tái sử dụng các chỉ mục của nhóm nguồn đã khai báo. Nó di chuyển cùng với nguồn đó mà không cần bỏ phiếu về ranh giới tái sử dụng hoặc thêm một khe khác vào cấu trúc liên kết radix.

DeepSeek-V4 làm rõ sự khác biệt này. FULL bao phủ tiền tố logic, trong khi SWA chỉ bao phủ cửa sổ phía sau của nó, vì vậy cả hai đều là thành phần. Chúng cũng sử dụng các không gian chỉ mục thiết bị độc lập. Trong ví dụ sáu trang được chuẩn hóa ở Hình 3, bộ cấp phát ánh xạ các khe đuôi FULL F4, F5 tới các khe SWA S0, S1 tại thời điểm chạy. Các nhóm KV nén C4 và C128, bộ đệm chỉ mục và trạng thái nén không xác định ranh giới tái sử dụng mới. Chúng đăng ký dưới dạng sidecar, với ba nhóm theo sau FULL và hai nhóm theo sau SWA.

Kết quả kiểm chuẩn HiCache đa lượt (Multi-Turn)

Các khối lượng công việc đa lượt phát triển một tiền tố hội thoại có thể tái sử dụng sau mỗi vòng. Nếu các tầng thấp hơn bảo tồn tiền tố đó sau khi dung lượng GPU cạn kiệt, tỷ lệ trúng cache sẽ vẫn cao và TTFT sẽ tăng chậm hơn.

Chúng tôi so sánh ba cấu hình cache trên hai mô hình lai: chỉ GPU L1, GPU L1 với Host L2, và GPU L1 với Host L2 cộng với tầng bộ nhớ phân tán Mooncake Store 500 GiB làm L3. DeepSeek-V4-Flash sử dụng FULL và SWA trên bốn GPU H200 với TP4, 48 client, 60 vòng, và 4.096 token đầu vào cộng 16 token đầu ra mỗi lượt. Inkling-Small sử dụng FULL, SWA và MAMBA trên tám GPU H200 với TP8, 64 client, 30 vòng, và 1.216 token đầu vào cộng 64 token đầu ra mỗi lượt.

Các lệnh dưới đây chứa các trình giữ chỗ cho đường dẫn mô hình và cấu hình Mooncake client. Chúng ghi lại thời gian chạy và các cờ khối lượng công việc được sử dụng ở đây, nhưng không phải là một môi trường tái tạo hoàn chỉnh.

Chạy độc lập từng cấu hình cache và dừng server của nó trước khi bắt đầu cấu hình tiếp theo.

DeepSeek-V4

Inkling

Hình 4 báo cáo TTFT trung bình và tỷ lệ trúng cache prompt-token theo vòng. Đối với mỗi vòng, tỷ lệ trúng là tổng các token tiền tố được lưu trong cache trên các yêu cầu chia cho tổng độ dài prompt hoàn chỉnh của chúng. Trong cả hai khối lượng công việc, L1 mất các tiền tố có thể tái sử dụng trước, L2 trì hoãn giới hạn dung lượng, và L3 vẫn ở mức cao sau khi khởi động và kết thúc trên 96%.

Trên DeepSeek-V4-Flash, L3 giữ tỷ lệ trúng gần 98%, giữ TTFT trung bình dưới 9 giây, và đạt 145,5K token đầu vào hiệu dụng/giây, so với 9,4K cho L1 và 14,3K cho L1 cộng L2. Trên Inkling-Small, L3 kết thúc với tỷ lệ trúng 96,8% và TTFT 1,23 giây trong khi đạt 67,1K token đầu vào hiệu dụng/giây, so với 15,5K cho L1 và 21,1K cho L1 cộng L2.

Thông lượng token đầu vào hiệu dụng tuân theo bench_multiturn.py: tổng độ dài prompt hoàn chỉnh chia cho thời gian thực tế. Nó ghi nhận các token tiền tố trúng cache, vì vậy nó đo lường tiến độ phục vụ dưới sự tái sử dụng tiền tố thay vì thông lượng tính toán prefill thô. Trong các lần chạy này, lợi ích của L3 chủ yếu đến từ việc giữ các tiền tố có thể tái sử dụng khả dụng sau khi các tầng nhỏ hơn đạt đến dung lượng.

Việc giải phóng (eviction) nhận biết phiên (session-aware) được triển khai trực tiếp trong UnifiedRadixCache. Cơ chế này cung cấp một tín hiệu tái sử dụng mà LRU thông thường không có. LRU ghi lại các mục cache nào được truy cập gần đây, nhưng không ghi lại tiền tố nào thuộc về các phiên đang hoạt động và có khả năng được tái sử dụng trong lượt tiếp theo của chúng. Dưới áp lực bộ nhớ, nó có thể giải phóng GPU KV của một phiên đang hoạt động trong khi giữ lại các mục không liên quan.

Các ứng dụng gắn một session_id ổn định vào mọi yêu cầu. Sau khi một yêu cầu kết thúc thành công, Unified Radix Cache đăng ký vùng có thể tái sử dụng cho phiên đó. FULL theo dõi đường dẫn tiền tố của nó, SWA theo dõi cửa sổ phía sau của nó, và MAMBA theo dõi biên tái sử dụng của nó. Tất cả các phiên vẫn chia sẻ một cấu trúc liên kết radix, và mọi lượt vẫn cung cấp prompt hoàn chỉnh của nó.

Các tham chiếu này thay đổi thứ tự giải phóng thay vì ghim bộ nhớ. FULL sắp xếp các ứng viên theo việc chúng có được tham chiếu hay không, số lượng tham chiếu phiên của chúng và ưu tiên giải phóng cơ sở đã cấu hình. SWA và MAMBA trước tiên quét các mục không được tham chiếu trong các vùng có thể tái sử dụng của riêng chúng, sau đó quay lại các mục được tham chiếu khi cần thêm không gian. Chính sách hiện tại bao gồm GPU L1 và Host L2, không bao gồm các tầng L3 bên ngoài.

Khi một ứng dụng gọi /close_session, Unified Radix Cache xóa các tham chiếu của phiên đó mà không xóa ngay lập tức các mục cache của nó. Các thế hệ phiên và các tombstone phiên đóng bị giới hạn ngăn chặn các yêu cầu cũ kết thúc sau khi đóng hoặc mở lại khôi phục các tham chiếu đã giải phóng.

HiCache nhận biết phiên trên khối lượng công việc SWE-bench

Chúng tôi đánh giá DeepSeek-V4-Pro và Qwen3.5-397B-A17B với TP8 và HiCache trên các quỹ đạo tác nhân SWE-bench. Cơ sở so sánh sử dụng HiRadixCache thông thường với LRU. Phép so sánh kích hoạt Unified Radix Cache và --enable-session-radix-cache. Vì điều này thay đổi cả triển khai cache và chính sách giải phóng, các khác biệt quan sát được không nên được hiểu là một sự cắt bỏ cô lập của nhận thức phiên. Bản ghi kiểm chuẩn cung cấp các cờ server và cấu hình sandbox.

Hàng trên cùng của Hình 6 xếp chồng các tỷ lệ trúng cache thiết bị và host. Tại batch size 128, tỷ lệ trúng thiết bị của DeepSeek-V4-Pro tăng từ khoảng 42% lên 51%. Tại batch size 32, Qwen3.5-397B-A17B tăng từ khoảng 5% lên 34%. Tại batch size 64, tổng tỷ lệ trúng thiết bị cộng host của Qwen tăng từ khoảng 58% lên 67%.

Hàng dưới cùng báo cáo TTFT tương ứng. So với cơ sở HiRadixCache thông thường, cấu hình Unified Radix Cache nhận biết phiên quan sát thấy TTFT thấp hơn 11,0% và 2,9% cho DeepSeek-V4-Pro ở batch size 128 và 256. Qwen3.5-397B-A17B quan sát thấy TTFT thấp hơn 13,5% và 16,6% ở batch size 32 và 64.