Tối ưu hóa suy luận Step3: Decode throughput trên AMD Instinct MI308X cao hơn 1,30 lần so với NVIDIA H20

Bối cảnh

Step3 của StepFun là một mô hình đa phương thức Mixture-of-Experts (MoE) với 321B tham số, kích hoạt 38B tham số cho mỗi token. Mô hình gồm 61 lớp với 56 lớp MoE sử dụng lựa chọn chuyên gia 3-in-48, và giới thiệu Multi-Matrix Factorization Attention (MFA) giúp giảm nhu cầu KV-cache xuống còn khoảng 22% chi phí attention trên mỗi token của DeepSeek V3.

StepFun đang triển khai Step3 trên GPU NVIDIA H20 và muốn đánh giá AMD Instinct MI308X như một phương án thay thế. Moreh được mời tối ưu hóa suy luận cho một mô hình riêng có cùng kiến trúc với Step3 trên MI308X, trước khi Step3 được phát hành công khai dưới dạng mã nguồn mở. Đây là một ví dụ về dịch vụ tối ưu hóa mô hình tùy chỉnh của Moreh, nơi chúng tôi điều chỉnh Moreh vLLM cho các kiến trúc mô hình độc quyền.

Tại sao chọn MI308X cho Decode

AMD Instinct MI308X là một biến thể của MI300X dành cho thị trường Trung Quốc. Nó có 1/4 số lõi tính toán của MI300X nhưng vẫn giữ nguyên dung lượng và băng thông bộ nhớ HBM3e. Điều này khiến MI308X đặc biệt phù hợp cho giai đoạn decode của suy luận LLM, vốn bị giới hạn bởi băng thông bộ nhớ chứ không phải năng lực tính toán: các token được tạo ra từng cái một theo cách tự hồi quy, và nút thắt cổ chai nằm ở việc tải trọng số mô hình và KV-cache từ bộ nhớ — chứ không phải thực hiện phép nhân ma trận.

Các kỹ thuật tối ưu hóa

• Kernel HIP attention tùy chỉnh: Kernel Triton attention mặc định của vLLM là nút thắt cổ chai lớn nhất, chiếm khoảng 50% thời gian GPU. Chúng tôi đã phát triển các kernel HIP attention tùy chỉnh được tối ưu hóa cho cấu hình MFA của Step3 (64 query heads, 1 KV head, head dimension 256) với data parallelism. Các kernel của chúng tôi giảm độ trễ attention 72% cho decode batch và 37% cho mixed prefill/decode batch.
• CUDA graph: Sau khi độ trễ GPU kernel được giảm đáng kể, chi phí phía CPU trở thành nút thắt cổ chai tiếp theo cho các bước decode. Chúng tôi đã bật tính năng CUDA graph capture hoàn chỉnh cho mô hình Step3 với DP8-EP8 parallelism, cải thiện decode throughput từ khoảng 2,900 lên 4,100 tok/s.
• Lượng tử hóa hỗn hợp BF16–FP8 blockscale: Tinh chỉnh GEMM toàn diện cho cả tính toán BF16 và FP8 blockscale để đạt được sự cân bằng tối ưu giữa độ chính xác và hiệu suất.
• Kernel MoE one-stage được tối ưu: Tối ưu hóa kernel tùy chỉnh cho lớp MoE của Step3, tập trung vào tham số inter_dim.
• Shared-expert MLP fusion: Tích hợp shared-expert MLP vào trong lớp MoE để giảm tính toán dư thừa và cải thiện độ trễ suy luận.
• Tích hợp MoRI EP: Tích hợp thư viện MoRI để thực hiện giao tiếp expert-parallel all-to-all hiệu quả trên AMD GPU.

Kết quả hiệu năng

Chúng tôi đã benchmark Moreh vLLM trên 8× MI308X so với số liệu được StepFun báo cáo trên 8× NVIDIA H20, sử dụng cùng cấu hình kiểm thử: ISL=4096, OSL=256, Concurrency=256, với DP8-EP8 parallelism (8-way data parallelism cho attention, 8-way expert parallelism cho MoE).

Kết quả cho thấy sự phân tách rõ ràng giữa hai giai đoạn:

• Decode: Moreh vLLM trên MI308X đạt 4,082 tok/s — throughput cao hơn 1.30× và latency thấp hơn 1.30× (63 ms vs. 82 ms) so với baseline H20 của StepFun.
• Prefill: H20 vẫn giữ ưu thế trong giai đoạn prefill phụ thuộc vào tính toán (13,780 vs. 9,601 tok/s), điều này là dự kiến do hệ thống cache on-chip mạnh hơn của nó.

Trong môi trường production với kiến trúc tách biệt prefill–decode, giai đoạn decode là nơi phân bổ nhiều GPU nhất. Hiệu năng decode mạnh mẽ của MI308X chuyển hóa trực tiếp thành khả năng triển khai tiết kiệm chi phí ở quy mô lớn.

Tổng kết

Dự án này chứng minh rằng AMD Instinct MI308X, kết hợp với các tối ưu hóa chuyên biệt cho từng mô hình của Moreh vLLM, có thể mang lại decode throughput cao hơn NVIDIA H20 cho các mô hình MoE lớn. Băng thông bộ nhớ cao của MI308X so với năng lực tính toán khiến nó trở thành lựa chọn hiệu quả về chi phí cho giai đoạn decode, vốn chiếm phần lớn phân bổ GPU trong các triển khai LLM serving production.

Moreh cung cấp dịch vụ tối ưu hóa vLLM tùy chỉnh cho các mô hình độc quyền và fine-tuned. Nếu bạn đang đánh giá AMD GPU cho mô hình của mình, hãy liên hệ với chúng tôi để thảo luận về cách chúng tôi có thể hỗ trợ.