Ba tối ưu hóa giúp suy luận MoE nhanh hơn trên NVIDIA DGX Spark

Các kỹ sư tối ưu hóa mô hình ngôn ngữ lớn thường gặp rào cản khi việc bổ sung thêm tài nguyên tính toán không còn giúp tăng tốc độ tạo sinh. Một đánh giá kỹ thuật gần đây về mô hình Ling-3.0-flash trên phần cứng NVIDIA DGX Spark đã chỉ ra rằng hiệu suất giải mã bị giới hạn bởi băng thông bộ nhớ chứ không phải sức mạnh tính toán thô. Bằng cách nhắm mục tiêu vào các điểm nghẽn cụ thể trong việc đọc trọng số và xác minh token, các nhà nghiên cứu đã đạt được những cải thiện đáng kể về độ trễ mà không làm giảm độ chính xác.

Chuyện gì đã xảy ra

Đánh giá tập trung vào Ling-3.0-flash, một mô hình hỗn hợp chuyên gia (MoE) của Inclusion AI với tổng cộng 124 tỷ tham số nhưng chỉ có 5,1 tỷ tham số hoạt động cho mỗi token. Khi chạy với kích thước lô bằng 1, các công cụ phân tích cho thấy thời gian dành cho việc tạo token chủ yếu bị chi phối bởi hai thao tác: bản thân phép tính MoE và phép chiếu lớp đầu ra, hay còn gọi là LM head. Các bước này đòi hỏi phải đọc một lượng lớn dữ liệu trọng số từ bộ nhớ, tạo ra một điểm nghẽn mà các kỹ thuật tối ưu hóa tiêu chuẩn không thể giải quyết triệt để.

Để khắc phục điều này, nhóm đã thử nghiệm ba phương pháp tối ưu hóa riêng biệt. Thứ nhất, họ thay thế backend MoE FlashInfer MXFP4 mặc định bằng thư viện kernel Humming. Thứ hai, họ áp dụng lượng tử hóa FP8 trực tuyến (online quantization) cụ thể cho LM head để giảm dung lượng bộ nhớ chiếm dụng. Thứ ba, họ triển khai giải mã suy đoán (speculative decoding) với dự đoán đa token sử dụng lớp NEXTN tích hợp sẵn của mô hình. Mỗi thay đổi đều nhắm vào một khía cạnh khác nhau của quy trình suy luận, từ hiệu quả thực thi kernel đến khối lượng dữ liệu và số bước xử lý.

Kết quả được đo lường trên cả các tải công việc tổng hợp và toàn bộ tập kiểm tra GSM8K. Trong khi các bài kiểm tra tổng hợp cho thấy mức tăng trưởng khổng lồ do các mẫu đầu ra lặp đi lặp lại, thì đánh giá thực tế trên 1.319 câu hỏi GSM8K đã chứng minh mức tăng 42,5% về thông lượng giải mã và giảm 24,5% về độ trễ trung bình. Quan trọng hơn, những cải tiến này đạt được mà không có sự thay đổi nào về độ chính xác của mô hình có ý nghĩa thống kê, duy trì điểm số ở mức 96,82% so với mức cơ sở 96,89%.

Cách thức hoạt động

Phương pháp tối ưu hóa đầu tiên tận dụng Humming, một thư viện kernel độc lập sử dụng biên dịch thời gian chạy NVIDIA CUDA Runtime Compilation để tạo ra mã được tùy chỉnh theo các chiều cụ thể của mô hình. Khác với đường dẫn FlashInfer cơ sở, vốn biên dịch nhiều mục tiêu chung chung, Humming tạo ra các binary chuyên biệt cho các hình dạng ma trận thực tế mà Ling-3.0-flash sử dụng. Sự chuyên biệt hóa này giảm thời gian cho mỗi token đầu ra bằng cách đảm bảo GPU chỉ thực hiện các lệnh cần thiết cho kích thước trung gian hẹp của mô hình.

Kỹ thuật thứ hai giải quyết vấn đề của LM head, nơi chiếu các trạng thái ẩn sang một từ vựng gồm hơn 157.000 token. Trong cấu hình cơ sở, lớp này sử dụng độ chính xác BF16, đòi hỏi băng thông bộ nhớ đáng kể cho mỗi token được tạo ra. Bằng cách chuyển sang lượng tử hóa FP8 trực tuyến chỉ cho lớp này, hệ thống giảm một nửa khối lượng dữ liệu trọng số được đọc trong mỗi bước. Thay đổi này loại trừ lẫn nhau với tùy chọn LM head FP32 cơ sở và yêu cầu cấu hình cẩn thận để tránh các vấn đề tương thích với các phần khác của mô hình.

Phương pháp thứ ba sử dụng giải mã suy đoán để tạo ra nhiều token trong một vòng xác minh duy nhất. Ling-3.0-flash bao gồm một lớp NEXTN gốc đóng vai trò như một mô hình nháp (draft model), dự đoán các token tiếp theo mà mô hình chính sau đó sẽ xác minh. Thay vì tạo từng token một, hệ thống soạn thảo ba token và xác minh đồng thời bốn vị trí. Cách tiếp cận này giảm tổng số bước giải mã cần thiết, mặc dù nó làm tăng độ phức tạp của mỗi bước xác minh do số lượng chuyên gia được kích hoạt cao hơn.

Chi tiết quan trọng

• Ling-3.0-flash có tổng cộng 124 tỷ tham số với 5,1 tỷ tham số hoạt động cho mỗi token qua 42 lớp.
• Việc thay thế backend FlashInfer bằng Humming đã giảm thời gian cho mỗi token đầu ra từ 27,71 ms xuống còn 25,35 ms, tức là lợi ích 9,3%.
• Lượng tử hóa FP8 trực tuyến của LM head đã giảm một nửa khối lượng đọc trọng số cho lớp đó, góp phần vào mức cải thiện thông lượng tích lũy 21,7%.
• Dự đoán đa token với NEXTN đạt 3,90 token được chấp nhận mỗi vòng trên dữ liệu tổng hợp, tăng thông lượng lên 185,1% trong bài kiểm tra cụ thể đó.
• Trên benchmark GSM8K, các tối ưu hóa kết hợp đã tăng thông lượng giải mã lên 42,5% và giảm độ trễ trung bình xuống 24,5%.
• Độ chính xác vẫn ổn định, chỉ giảm 0,08 điểm phần trăm từ 96,89% xuống 96,82%, mức giảm này không có ý nghĩa thống kê.

Tại sao điều này quan trọng

Đối với các nhà phát triển xây dựng quy trình làm việc dựa trên agent hoặc ứng dụng thời gian thực, độ trễ thường là ràng buộc chính. Phân tích này xác nhận rằng băng thông bộ nhớ, chứ không chỉ là sức mạnh tính toán, là yếu tố giới hạn đối với nhiều mô hình MoE ở kích thước lô thấp. Bằng cách tập trung vào việc giảm lưu lượng trọng số và tăng số token mỗi bước, các kỹ sư có thể đạt được những cải thiện hiệu suất đáng kể mà không cần đến các cụm phần cứng lớn hơn. Điều này đặc biệt liên quan đến các triển khai tại edge hoặc môi trường đám mây nhạy cảm về chi phí, nơi mỗi mili giây đều quý giá.

Tuy nhiên, nghiên cứu cũng nêu bật rủi ro khi kết hợp các tối ưu hóa mà thiếu kiến thức sâu về tích hợp. Các nhà nghiên cứu phát hiện ra rằng việc bật lượng tử hóa FP8 LM head và dự đoán đa token cùng lúc đòi hỏi mô hình nháp phải chia sẻ toàn bộ module LM-head. Nếu không có giao diện cụ thể này, hệ thống âm thầm quay trở lại cấu hình chậm hơn, diễn giải các trọng số FP8 thành BF16. Điều này nhấn mạnh tầm quan trọng của việc xác minh mọi tối ưu hóa khi khởi động thay vì giả định khả năng tương thích dựa trên các bài kiểm tra thành phần riêng lẻ.

Bạn có thể làm gì

• Phân tích quy trình suy luận MoE hiện tại của bạn để xác định xem việc giải mã bị giới hạn bởi bộ nhớ hay bởi tính toán trước khi áp dụng các tối ưu hóa.
• Thử nghiệm thư viện kernel Humming như một lựa chọn thay thế trực tiếp cho FlashInfer nếu mô hình của bạn có các chiều ma trận chuyên biệt.
• Bật lượng tử hóa FP8 trực tuyến cho LM head bằng biến môi trường SGLANG_ENABLE_FP8_LM_HEAD để giảm mức sử dụng băng thông bộ nhớ.
• Triển khai dự đoán đa token sử dụng các lớp nháp gốc như NEXTN để giảm số bước giải mã tuần tự.
• Xác minh rằng các mô hình nháp và mô hình đích chia sẻ đúng các module khi kết hợp lượng tử hóa với giải mã suy đoán để tránh suy giảm hiệu suất âm thầm.
• Đánh giá benchmark các thay đổi trên các tải công việc thực tế như GSM8K thay vì chỉ dựa vào các bài kiểm tra tổng hợp, vốn có thể phóng đại lợi ích của giải mã suy đoán.