Tối ưu suy luận MiMo-V2.5: Biến Hybrid SWA hoạt động thực tế trong sản xuất

Nhóm MiMo của Xiaomi đã công bố một bài phân tích sâu về các tối ưu suy luận đằng sau các mô hình MiMo-V2.5 và V2.5-Pro của họ. Điểm nhấn: Hybrid Sliding Window Attention (SWA) hứa hẹn giảm 7× lưu trữ và tính toán KVCache, nhưng để thu hẹp khoảng cách từ lý thuyết đến sản xuất đòi hỏi phải tái cấu trúc toàn bộ stack caching và lập lịch.

Dòng V2.5 sử dụng kiến trúc 70 lớp — 10 lớp Full Attention và 60 lớp SWA với kích thước cửa sổ 128. Về lý thuyết, điều này cắt giảm tính toán attention và KVCache xuống còn khoảng 1/7 so với mô hình full-attention. Trong thực tế, như nhóm lưu ý, các triển khai ngây thơ trong các framework như SGLang v0.5.5 lưu trữ toàn bộ KVCache cho các lớp SWA, vô hiệu hóa lợi thế.

Thiết kế KVCache hai pool

Sửa lỗi cốt lõi là chia KVCache thành hai pool độc lập: một cho các lớp Full Attention (lưu trữ O(N)) và một cho các lớp SWA (lưu trữ O(W)). Ở lớp vật lý, pool SWA chỉ được định kích thước cho cửa sổ và hỗ trợ loại bỏ độc lập. Ở lớp logic, một chế độ xem chuỗi thống nhất được hiển thị cho cây tiền tố và bộ lập lịch, với ánh xạ Full → SWA được duy trì cho lưu trữ phân tầng minh bạch. Riêng điều này đã khôi phục hiệu quả dung lượng 7×.

Cây bộ nhớ đệm tiền tố nhận biết SWA

RadixAttention truyền thống giả định các chuỗi token bằng nhau tạo ra KV có thể tái sử dụng. Dưới SWA, điều đó bị phá vỡ — một nút cây tiền tố vẫn có thể biểu diễn logic một chuỗi token, nhưng KV SWA của nó có thể đã bị loại bỏ ra ngoài cửa sổ. Nhóm đã sửa đổi các quy tắc khớp thành "độ dài an toàn cửa sổ": độ dài khớp được cắt bớt để W token cuối phải có slot hợp lệ trong pool SWA. Việc loại bỏ được gắn với vòng đời yêu cầu (hoàn thành chunk, kết thúc, mỗi N token giải mã), giữ cho mức sử dụng pool SWA không đổi ở kích thước cửa sổ thay vì tăng theo độ dài chuỗi. Các nút hiện mang hai chỉ mục — chỉ mục phân đoạn Full Attention và ánh xạ phân đoạn SWA — cho phép loại bỏ độc lập các phân đoạn SWA ngoài cửa sổ trong khi vẫn giữ các phân đoạn Full Attention để tái sử dụng.

Bộ nhớ đệm phân tầng và nhất quán phân tán

Hệ thống HiCache ba tầng (thiết bị, máy chủ, backend lưu trữ) đã được tái cấu trúc để duy trì các chỉ mục SWA hợp lệ trên mỗi tầng. Nhóm phát hiện ra rằng các lần trúng Full Attention Cache với các lần trượt SWA Cache gây ra sự cắt giảm độ dài khớp nghiêm trọng, buộc phải tính toán lại. Họ đã tối ưu hóa tính nhất quán phân tán giữa các triển khai, độ dài tiền tố được chia sẻ và di chuyển dữ liệu không đồng bộ để giữ cho hai bên đồng bộ — rất quan trọng để duy trì tỷ lệ trúng dưới các quy tắc khớp mới.

MoE và các nút thắt đa phương thức

Ngoài SWA, bài viết đề cập đến lập lịch MoE và cân bằng tải, đồng thời lưu ý rằng các bộ mã hóa đa phương thức vẫn là nút thắt thông lượng đối với đầu vào hình ảnh lớn và video dài. Nhóm đã tối ưu hóa pipeline thực thi Prefill/Decode và chiến lược lập lịch để xử lý những vấn đề này, nhưng bài viết chủ yếu tập trung vào công việc KVCache.

Kết luận: Hiệu quả lý thuyết của Hybrid SWA là có thật, nhưng chỉ sau khi bạn xây dựng lại hệ thống caching từ đầu. Cách tiếp cận của nhóm MiMo — hai pool, cây tiền tố an toàn cửa sổ, nhất quán phân tầng — hiện đang được các framework suy luận chính thống áp dụng, điều này nói lên tầm quan trọng của công việc này.