Eagle 3.1: Bước tiến đột phá trong giải mã suy đoán LLM nhờ hợp tác giữa vLLM và TorchSpec

Chuỗi thuật toán EAGLE — bao gồm EAGLE 1, 2 và 3 — đã trở thành một trong những họp thuật toán giải mã suy đoán (speculative decoding) được áp dụng rộng rãi nhất trong cả hệ thống nghiên cứu và môi trường sản xuất. Hôm nay, các đội ngũ phát triển từ EAGLE, vLLM và TorchSpec chính thức giới thiệu EAGLE 3.1, một bước tiến lớn nhằm nâng cao độ ổn định, hiệu quả và khả năng triển khai của kỹ thuật này.

Thách thức về độ ổn định trong giải mã suy đoán

Mặc dù giải mã suy đoán hoạt động hiệu quả trong các môi trường được kiểm soát, hiệu suất của nó thường bị suy giảm khi đối mặt với các mẫu chat (chat templates) khác nhau, đầu vào ngữ cảnh dài (long-context) hoặc các câu lệnh hệ thống nằm ngoài phân phối dự kiến (out-of-distribution system prompts).

Đội ngũ EAGLE đã truy nguyên sự mong manh này đến một hiện tượng được gọi là "attention drift" (sự trôi dạt của sự chú ý). Khi độ sâu của sự suy đoán tăng lên, mô hình dự thảo (drafter) dần chuyển sự chú ý của nó khỏi các token đích (sink tokens) và hướng sang các token do chính nó tạo ra. Điều này dẫn đến việc đầu vào kết hợp trở nên mất cân bằng và độ lớn của trạng thái ẩn (hidden-state magnitude) tăng lên theo các bước suy đoán do đường dẫn dư (residual path) chưa được chuẩn hóa. Tất cả các yếu tố này khiến mô hình dự thảo ngày càng kém ổn định hơn ở các độ sâu suy đoán lớn hơn.

So sánh kiến trúc EAGLE 3 và EAGLE 3.1

Các cải tiến kiến trúc trong EAGLE 3.1

Để giải quyết vấn đề này, EAGLE 3.1 giới thiệu hai cải tiến kiến trúc chính:

• Thực hiện chuẩn hóa FC (FC normalization) sau mỗi trạng thái ẩn mục tiêu và trước lớp FC.
• Nạp các trạng thái ẩn đã chuẩn hóa (post-norm hidden states) vào bước giải mã tiếp theo.

Về mặt trực quan, thiết kế post-norm này giúp phương pháp hoạt động giống như việc gọi đệ quy mô hình dự thảo qua các bước giải mã, thay vì chỉ đơn giản là nối thêm các lớp vào mô hình mục tiêu. Những thay đổi này đã cải thiện đáng kể độ ổn định trong các kịch bản triển khai thực tế.

So với EAGLE 3, phiên bản 3.1 thể hiện khả năng ngoại suy tốt hơn từ thời điểm huấn luyện đến thời điểm suy luận, độ bền cao hơn với ngữ cảnh dài, khả năng phục hồi tốt hơn trước sự biến đổi của mẫu chat và câu lệnh hệ thống, cũng như độ dài chấp nhận ổn định hơn trong nhiều môi trường phục vụ khác nhau. Đặc biệt, trong các khối lượng công việc ngữ cảnh dài, EAGLE 3.1 đạt được độ dài chấp nhận dài hơn gấp 2 lần so với EAGLE 3.

Huấn luyện EAGLE 3.1 với TorchSpec

TorchSpec hiện cung cấp hỗ trợ huấn luyện hiệu quả cho EAGLE 3.1 và các thuật toán giải mã suy đoán trong tương lai. Bằng cách giảm chi phí huấn luyện và đơn giản hóa quy trình thử nghiệm, TorchSpec giúp tăng tốc vòng lặp phát triển và khám phá cho việc nghiên cứu cũng như triển khai giải mã suy đoán thế hệ tiếp theo.

Dựa trên TorchSpec và vLLM, các đội ngũ cũng đã huấn luyện và mã nguồn mở một mô hình dự thảo EAGLE 3.1 cho Kimi K2.6 tại đường dẫn: https://huggingface.co/lightseekorg/kimi-k2.6-eagle3.1-mla. Mô hình này đóng vai trò là ví dụ điển hình cho việc triển khai EAGLE 3.1 với sự hỗ trợ huấn luyện từ TorchSpec và khả năng phục vụ từ vLLM trên một mô hình phục vụ thực tế.

Tích hợp sâu với vLLM

EAGLE 3.1 đã được tích hợp vào vLLM dưới dạng một tiện ích mở rộng dựa trên cấu hình của bản triển khai EAGLE 3 hiện có. Quá trình tích hợp bao gồm hỗ trợ chuẩn hóa FC, phản hồi trạng thái ẩn post-norm và loại bỏ các giả định mã hóa cứng liên quan đến trạng thái ẩn mục tiêu.

Đáng chú ý, tính tương thích ngược với các điểm kiểm tra (checkpoints) EAGLE 3 hiện có được bảo toàn hoàn toàn. Kết quả là, các mô hình dự thảo EAGLE 3.1 có thể được cắm trực tiếp thông qua cùng một đường dẫn mã giải mã suy đoán.

Ví dụ, người dùng có thể chạy lệnh sau để triển khai:

vllm serve nvidia/Kimi-K2.6-NVFP4 \
--trust-remote-code \
--tensor-parallel-size 4 \
--tool-call-parser kimi_k2 \
--enable-auto-tool-choice \
--reasoning-parser kimi_k2 \
--attention-backend tokenspeed_mla \
--speculative-config '{"model":"lightseekorg/kimi-k2.6-eagle3.1-mla","method":"eagle3","num_speculative_tokens":3}' \
--language-model-only

Điều này giúp việc nâng cấp mô hình dự thảo trong hệ thống phục vụ vLLM sản xuất trở nên mượt mà và dễ dàng. Hỗ trợ này đã được hợp nhất vào nhánh chính hiện tại của vLLM và sẽ có sẵn trong bản phát hành nightly cũng như phiên bản v0.22.0 sắp tới.

Hiệu năng và Benchmark

Dựa trên dữ liệu ban đầu, đội ngũ đã benchmark mô hình dự thảo Kimi K2.6 EAGLE 3.1 trên Kimi-K2.6-NVFP4 với vLLM (TP=4, GB200, non-disagg) sử dụng bộ dữ liệu lập trình SPEED-Bench. Kết quả cho thấy EAGLE 3.1 mang lại thông lượng đầu ra trên mỗi người dùng cao hơn 2,03 lần ở mức đồng thời (concurrency) 1. Tăng tốc này vẫn giữ ý nghĩa khi mức đồng thời tăng lên (1,71 lần ở C=4 và 1,66 lần ở C=16).

Biểu đồ thông lượng đầu ra trên mỗi người dùng

Sự hợp tác mã nguồn mở trong hệ sinh thái

Sự hợp tác giữa đội ngũ EAGLE, vLLM và TorchSpec là một ví dụ mạnh mẽ về sự cộng tác mã nguồn mở xuyên suốt nghiên cứu thuật toán, tối ưu hóa hệ thống và cơ sở hạ tầng huấn luyện. Đội ngũ EAGLE tiếp tục thúc đẩy các thuật toán giải mã suy đoán, vLLM giúp đưa những đổi mới này vào các hệ thống suy luận sản xuất quy mô lớn, trong khi TorchSpec cho phép huấn luyện hiệu quả và thử nghiệm nhanh chóng cho các thuật toán trong tương lai.

Cùng nhau, họ hy vọng sẽ tiếp tục nâng cao mức cơ sở chung cho giải mã suy đoán và thúc đẩy những cải tiến hơn nữa về hiệu suất token trong toàn bộ hệ sinh thái LLM rộng lớn hơn.