Tối ưu hóa giải mã suy đoán (Speculative Decoding) trên GPU AMD với vLLM

TL;DR: Speculative decoding cho phép vLLM xác thực nhiều token dự đoán trong một lần chạy mô hình mục tiêu (target-model pass). Trong các thử nghiệm của chúng tôi, hiệu quả của nó đối với thông lượng token đầu ra thay đổi tùy theo phương pháp dự đoán và độ dài đề xuất, đồng thời cũng phụ thuộc vào dòng mô hình, checkpoint dự đoán, khối lượng công việc và hành vi chấp nhận.

Giới thiệu

Các mô hình ngôn ngữ lớn (LLM) hỗ trợ nhiều ứng dụng đa dạng, nhưng việc phục vụ chúng ở quy mô lớn đòi hỏi sự tối ưu hóa cẩn thận. Giải mã tự hồi quy (autoregressive decoding) tiêu chuẩn là nền tảng được hầu hết các hệ thống phục vụ LLM sử dụng: mô hình tạo ra một token, thêm nó vào chuỗi, sau đó sử dụng chuỗi đã cập nhật để tạo token tiếp theo. Quá trình này đơn giản và đáng tin cậy, nhưng vòng lặp phục vụ vẫn chỉ tiến thêm một token đã cam kết tại mỗi thời điểm vì các token đầu ra phải được tạo theo thứ tự nghiêm ngặt từ trái sang phải.

Speculative decoding [1] xây dựng dựa trên nền tảng này thông qua cơ chế dự đoán và xác thực (draft-and-verify). Một thành phần dự đoán nhẹ sẽ đề xuất các token tương lai tiềm năng, và mô hình mục tiêu sẽ xác thực các ứng viên đó trước khi chúng được cam kết. Khi một vài token dự đoán được chấp nhận, hệ thống có thể cam kết nhiều token đầu ra từ một bước xác thực duy nhất của mô hình mục tiêu trong khi vẫn bảo toàn hành vi đầu ra của mô hình đó.

Bài viết này khám phá cách thức hoạt động của speculative decoding trong vLLM và chia sẻ các phép đo từ môi trường thử nghiệm của chúng tôi. Trước tiên, chúng tôi xem xét lại nền tảng giải mã tự hồi quy và quy trình dự đoán và xác thực. Sau đó, chúng tôi kiểm tra năm phương pháp dự đoán suy luận: native MTP, Gemma 4 MTP, EAGLE-3, DFlash và DSpark. Các phương pháp này khác nhau ở cách thành phần dự đoán nhận thông tin từ mô hình mục tiêu và liệu các token ứng viên được tạo ra theo trình tự, tự hồi quy, song song hay thông qua phương pháp lai. Cuối cùng, chúng tôi chỉ ra cách kích hoạt các phương pháp đã thử nghiệm trong môi trường của mình, báo cáo các phép đo từ các thử nghiệm trên GPU AMD Instinct™ MI300X và MI355X sử dụng nền tảng phần mềm mở ROCm™, đồng thời thảo luận về các cân nhắc thực tế trong việc tinh chỉnh và khả năng quan sát.

Nền tảng giải mã tự hồi quy

Trong giải mã tự hồi quy tiêu chuẩn, mỗi bước giải mã tạo ra và cam kết một token mới. Ví dụ, việc tạo bốn token đầu ra yêu cầu bốn bước giải mã tuần tự:

Bước 1: context → model → T1

Bước 2: context + T1 → model → T2

Bước 3: context + T1 T2 → model → T3

Bước 4: context + T1 T2 T3 → model → T4

Sau mỗi bước, token được tạo ra sẽ được thêm vào chuỗi và trở thành một phần của đầu vào cho bước tiếp theo. Điều này làm cho vòng lặp giải mã trở nên đơn giản, nhưng nó cũng yêu cầu một bước giải mã mô hình cho mỗi token đầu ra. Trong quá trình tạo văn bản dài, vòng lặp từng token này có thể chiếm ưu thế về độ trễ và giới hạn thông lượng phục vụ.

Do đó, câu hỏi then chốt đằng sau speculative decoding là:

Liệu chúng ta có thể bảo toàn hành vi đầu ra của mô hình gốc trong khi giảm tần suất thế hệ tiến thêm chỉ một token tại mỗi thời điểm?

Speculative decoding giải quyết vấn đề này bằng cách tách biệt việc đề xuất và xác thực. Một thành phần dự đoán trước tiên đề xuất một vài token tương lai tiềm năng. Mô hình gốc, đóng vai trò là mô hình mục tiêu, sau đó xác thực các ứng viên đó trước khi chúng được cam kết.

Ý tưởng cốt lõi của speculative decoding

Speculative decoding không thay thế mô hình gốc. Thay vào đó, nó giữ mô hình gốc làm mô hình mục tiêu, vốn vẫn chịu trách nhiệm cho đầu ra cuối cùng, và thêm một giai đoạn đề xuất nhanh hơn ở phía trước.

Quy trình này có hai phần:

Trong mỗi vòng speculative decoding, như minh họa trong Hình 1, một thành phần dự đoán nhẹ sẽ đề xuất một hoặc nhiều token tương lai. Các token này chỉ là ứng viên và không được cam kết ngay lập tức. Sau đó, mô hình mục tiêu đánh giá chuỗi token ứng viên trong một lần xác thực.

Quá trình xác thực diễn ra từ trái sang phải. Mỗi token dự đoán được kiểm tra bằng kết quả của mô hình mục tiêu tại vị trí tương ứng. Các token được chấp nhận sẽ được cam kết vào chuỗi đầu ra. Khi một token dự đoán bị từ chối, các ứng viên sau đó từ cùng một đề xuất sẽ không còn được chấp nhận nữa.

Nếu một token dự đoán bị từ chối, mô hình mục tiêu sẽ cung cấp token tiếp theo. Các token dự đoán còn lại sẽ bị loại bỏ và quá trình tạo tiếp tục từ chuỗi đã cập nhật.

Về mặt khái niệm, giải mã tự hồi quy tiêu chuẩn tiến triển như sau:

target model → T1

target model → T2

target model → T3

target model → T4

Ngược lại, speculative decoding cho phép đánh giá nhiều vị trí ứng viên cùng nhau:

draft proposes T1 T2 T3 T4

model verifies ✓ ✓ ✗ stop

commit T1 T2 replacement token -

Điều này có thể giảm số vòng giải mã của mô hình mục tiêu khi nhiều ứng viên được chấp nhận. Khi thành phần dự đoán tạo ra các token mà mô hình mục tiêu chấp nhận, nhiều token đầu ra có thể được cam kết từ một bước xác thực của mô hình mục tiêu. Khi một đề xuất bị từ chối, kết quả từ phía mô hình mục tiêu sẽ quyết định cách thức tạo tiếp tục.

Một ví dụ đơn giản về chấp nhận/từ chối

Hình 2 đưa ra ví dụ về một vòng speculative decoding. Các ô màu xanh lá cây là các token dự đoán vượt qua xác thực, ô màu đỏ đánh dấu token dự đoán bị từ chối đầu tiên và ô màu xám là token dự đoán sau đó bị loại bỏ. Token màu xanh dương trong đầu ra đến từ mô hình mục tiêu, không phải từ đề xuất dự đoán.

Giả sử prompt hiện tại là:

The weather today is

Thành phần dự đoán đề xuất một vài token tương lai:

sunny and warm outside

Mô hình mục tiêu xác thực các token dự đoán từ trái sang phải:

draft proposes sunny and warm outside