Suy luận DeepSeek 21K output token mỗi giây trên GPU AMD Instinct MI300X với Expert Parallelism

Tổng quan

Một trong những bước đột phá lớn của mô hình ngôn ngữ lớn (LLM) là việc áp dụng kiến trúc Mixture-of-Experts (MoE). Thay vì một lớp feed-forward network (FFN) lớn duy nhất, mô hình chứa nhiều lớp FFN nhỏ hơn, mỗi lớp được gọi là một expert. Đối với mỗi token, một mạng gating nhẹ sẽ tự động chọn chỉ một tập con nhỏ các expert để xử lý ở mỗi bước. Cơ chế tính toán có điều kiện này cho phép mô hình MoE mở rộng đến số lượng tham số cực lớn trong khi vẫn duy trì hiệu quả tính toán. Ví dụ, mô hình DeepSeek-R1 671B gây chấn động khi phát hành vào tháng 1 năm 2025 và mô hình mã nguồn mở phổ biến GPT-OSS 120B của OpenAI phát hành vào tháng 8 năm 2025 đều sử dụng kiến trúc MoE.

Tuy nhiên, do kích thước lớn và thiết kế thưa, các mô hình như DeepSeek-R1 cần các kỹ thuật tối ưu hóa nâng cao để phục vụ hiệu quả ở quy mô lớn. Đặc biệt, chúng ta có thể đạt được thông lượng cao (tức là tổng số output token mỗi giây) bằng cách áp dụng Expert Parallelism (EP). Điều này là do (1) chỉ một tập con expert được lưu trong bộ nhớ GPU, cho phép kích thước batch lớn hơn, và (2) các tham số của từng expert có thể được tái sử dụng nhiều hơn sau khi được tải, từ đó cải thiện FLOPS per byte. Tuy nhiên, thách thức nảy sinh trong việc triển khai hiệu quả mẫu truyền thông chi tiết được gọi là dispatch/combine — chuyển các giá trị kích hoạt đến các GPU chịu trách nhiệm cho các expert được kích hoạt — và trong việc giải quyết sự mất cân bằng giữa các expert.

Moreh, đối tác phần mềm của AMD, gần đây đã chứng minh rằng suy luận DeepSeek-R1 có thể được thực thi với thông lượng cao bằng cách triển khai EP trên ngăn xếp phần mềm ROCm. Để thực sự đạt được hiệu suất cao như vậy trong suy luận LLM, cần có phần mềm được tối ưu hóa cẩn thận trên nhiều tầng của ngăn xếp. Đội ngũ Moreh áp dụng nhiều phương pháp tối ưu hóa ở cả cấp độ thư viện và triển khai mô hình. Kết quả là, họ đạt được thông lượng giải mã >21,000 tokens/sec trên máy chủ trang bị 8x AMD Instinct MI300X GPU.

Phương pháp song song hóa

Trong mô hình DeepSeek R1, mỗi khối decoder chứa 256 expert, trong đó 8 expert được kích hoạt cho mỗi token. Ngoài ra, có một shared expert luôn được thực thi bất kể kết quả định tuyến, và một lớp Multi-Head Latent Attention (MLA), cơ chế attention mới do DeepSeek giới thiệu. EP thường được triển khai kết hợp với data parallelism (DP). Mỗi GPU giữ và thực thi một tập con expert khác nhau. Trong khi đó, đối với các thành phần luôn được thực thi — như shared expert và MLA — mỗi GPU xử lý một batch yêu cầu riêng biệt song song. Theo đó, chúng tôi áp dụng cấu hình song song hóa DP=8 và EP=8 cho mỗi node. Nghĩa là, mỗi GPU chịu trách nhiệm cho 32 expert.

Hơn nữa, việc áp dụng prefill-decode disaggregation — tức là thực thi giai đoạn prefill và decode trên các node riêng biệt — có thể tăng thêm thông lượng tổng thể của hệ thống cluster. Giai đoạn prefill xử lý toàn bộ chuỗi đầu vào cùng một lúc và thiên về tính toán (compute-bound), trong khi giai đoạn decode tạo các output token từng cái một theo phương thức tự hồi quy và thiên về bộ nhớ (memory-bound). Bằng cách không chỉ tách biệt việc thực thi mà còn điều chỉnh cấu hình song song hóa và tối ưu hóa cho các đặc tính tính toán riêng biệt của prefill và decode, hiệu quả GPU có thể được cải thiện thêm. Trong bài viết này, chúng tôi tập trung đánh giá chủ yếu vào hiệu suất của giai đoạn decode, vì trong các ứng dụng thông thường, phần lớn GPU được dành cho giải mã.

Tối ưu hóa hiệu suất

Chỉ đơn thuần triển khai một phương pháp song song hóa mới không tự động mang lại hiệu suất tốt nhất. Để khai thác hết tiềm năng phần cứng, toàn bộ ngăn xếp phần mềm từ thư viện đến triển khai mô hình phải được tối ưu hóa để phù hợp với các mẫu tính toán và truyền thông mới. Đội ngũ Moreh đã triển khai nhiều phương pháp tối ưu hóa dựa trên vLLM, được mô tả dưới đây.

Cân bằng tải giữa các GPU

Một trong những thách thức lớn khi triển khai EP là khối lượng công việc không được phân bổ đều giữa các GPU và thay đổi tùy thuộc vào kết quả định tuyến. Mặc dù mô hình MoE được huấn luyện để giảm thiểu sự mất cân bằng này càng nhiều càng tốt, nhưng vẫn tồn tại những hạn chế cố hữu. DeepSeek cũng thừa nhận vấn đề này và báo cáo đã sử dụng phương pháp cân bằng tải riêng gọi là EPLB để giải quyết. Khi không áp dụng chiến lược cân bằng tải rõ ràng nào và chỉ đơn giản chia expert thành tám khối liên tiếp — tức là gán expert 0-31 cho GPU đầu tiên, 32-63 cho GPU thứ hai, …, và 224-255 cho GPU cuối cùng — chúng tôi quan sát thấy sự mất cân bằng khối lượng công việc lên đến 2 lần giữa các GPU. Chúng tôi đã thiết kế một thuật toán đo tần suất kích hoạt của mỗi expert, và nhóm 256 expert thành tám bộ gồm 32 expert sao cho tổng tần suất kích hoạt của mỗi bộ cân bằng nhất có thể. (Điều này có thể được xem như một biến thể của bài toán phân hoạch số cân bằng.) Thuật toán sau đó sắp xếp lại các expert sao cho mỗi bộ được đặt liền kề và áp dụng EP trên đó. Vì hoán vị expert đã thay đổi, mạng gating cũng phải tạo ra kết quả định tuyến phản ánh hoán vị expert mới. Kết quả là, chúng tôi đã giảm thành công sự mất cân bằng khối lượng công việc giữa các GPU xuống dưới 5%. Theo hiểu biết của chúng tôi, đây là triển khai cân bằng tải EP đầu tiên trên AMD GPU. Trong các kịch bản phục vụ thực tế, cân bằng tải liên tục có thể đạt được bằng cách định kỳ đo lại tần suất kích hoạt và tạo các hoán vị expert mới tương ứng.

Chồng lấp tính toán-truyền thông

Một thách thức khác khi áp dụng EP là giảm thiểu chi phí do truyền thông dispatch/combine all-to-all. Việc chia batch thành hai microbatch và thực thi đồng thời cho phép truyền thông của microbatch này chồng lấp với tính toán của microbatch kia. Khi kích thước batch đủ lớn — tức là trong các kịch bản thông lượng cao — phương pháp này có thể che giấu hiệu quả độ trễ truyền thông mà không hy sinh hiệu suất tính toán. Có nhiều biến thể triển khai của phương pháp này, trong đó chúng tôi chọn hệ thống DBO (Dual Batch Overlap) của vLLM. Chúng tôi đã sửa đổi hệ thống DBO để sử dụng MORI-EP, và đặc biệt, chúng tôi đã nâng cấp thư viện MORI-EP để các thao tác truyền thông của nó có thể được gọi từ hai luồng worker của DBO và được đồng bộ hóa đúng cách với các tính toán khác.

Tối ưu hóa thư viện

Đội ngũ Moreh cũng đã áp dụng nhiều phương pháp tối ưu hóa thư viện để tối đa hóa hiệu quả GPU và cải thiện cả thông lượng (tokens/sec) lẫn độ trễ (thời gian đến token đầu tiên và độ trễ giữa các token). Dưới đây là một số ví dụ tối ưu hóa tiêu biểu:

• Lựa chọn GEMM và Attention Kernel tối ưu: tự động chọn các kernel GEMM và Attention tối ưu mà không cần profiling trực tuyến hay điều chỉnh thủ công.
• Tối ưu hóa Fused MoE Kernel: triển khai kernel fused MoE được tối ưu hóa cao, mang lại hiệu suất tốt hơn thư viện AITER của AMD, đặc biệt với kích thước batch nhỏ.
• Hỗ trợ FP8 KV Cache: triển khai các kernel Multi-head Latent Attention (MLA) cho phép KV cache được lưu trữ và tải ở định dạng FP8, cải thiện hiệu suất đặc biệt trong các kịch bản ngữ cảnh dài.
• Kernel Fusion dọc và ngang: sử dụng cả fusion dọc (ví dụ: kernel fused RoPE) và fusion ngang (ví dụ: hợp nhất nhiều GEMM trong shared expert) để giảm chi phí kernel và cải thiện hiệu quả tính toán.

Để biết chi tiết về các tối ưu hóa cấp thư viện này và cải thiện hiệu suất độc lập với EP, vui lòng tham khảo báo cáo kỹ thuật của Moreh.

Tận dụng HIP Graphs

HIP Graphs là kỹ thuật giảm chi phí CPU runtime bằng cách capture nhiều thao tác GPU thành một đồ thị tĩnh và phát hành tất cả cùng một lúc. Điều này đặc biệt cần thiết trong giai đoạn decode, nơi các thao tác riêng lẻ tương đối ngắn, khiến việc tận dụng tối đa GPU trở nên quan trọng. Tuy nhiên, việc xây dựng đồ thị tĩnh yêu cầu kích thước tensor được truyền giữa các thao tác phải cố định. Trong EP, kích thước đầu vào cho mỗi expert vốn dĩ là động và được xác định bởi kết quả định tuyến, khiến việc biểu diễn dưới dạng đồ thị tĩnh cực kỳ khó khăn. Chúng tôi đã sửa đổi triển khai mô hình DeepSeek trong vLLM để làm cho kích thước tensor tĩnh nhất có thể trong khi vẫn hỗ trợ EP. Điều này cho phép chúng tôi tận dụng tối đa HIP Graph, từ đó giảm thiểu chi phí CPU. Việc tích hợp các thao tác thư viện MoRI vào đồ thị gặp nhiều thách thức kỹ thuật, nhưng cuối cùng chúng tôi đã thành công trong việc đưa các thao tác dispatch/combine của MORI-EP vào như một phần của việc thực thi đồ thị.

Phương pháp thực nghiệm

Các thí nghiệm của chúng tôi được thực hiện trên máy chủ MI300X với các thông số kỹ thuật sau.

• CPU: 2x AMD EPYC 9474F 48-core 3.6 GHz
• Main memory: 2,304 GB
• GPU: 8x AMD Instinct MI300X OAM GPU 192 GB
• Server: Gigabyte G593-ZX1-AAX1
• Operating system: Ubuntu 22.04.4 LTS
• ROCm version: 6.4.1

Chúng tôi thiết kế phương pháp đánh giá bằng cách tham khảo bài đăng blog của LMSys. Trong bài đăng đó, các chỉ số chính để đánh giá hiệu suất EP — giả định prefill/decode disaggregation — là tổng thông lượng (tokens/sec) mỗi node decode và độ trễ giữa các token. Trong khi LMSys triển khai prefill/decode disaggregation trên 12 node, chúng tôi hướng đến việc đo hiệu suất decode ngay cả trong môi trường đơn node. Để đạt được điều này, chúng tôi đã thực hiện một số sửa đổi đối với vLLM như sau:

• Ở phía máy chủ, chúng tôi loại trừ thời gian prefill và chỉ tích lũy thời gian decode và số lượng token để đo thông lượng decode riêng biệt. Độ trễ giữa các token vẫn có thể được đo ở phía client bằng phương pháp benchmark tiêu chuẩn.
• Bộ lập lịch của engine vLLM V1 ban đầu thực thi các yêu cầu prefill và decode trong một batch duy nhất (gọi là mixed prefill), khiến việc loại trừ thời gian prefill là không thể. Chúng tôi đã sửa đổi để các yêu cầu prefill và decode được gửi vào các hàng đợi riêng biệt và thực thi độc lập.

Có thể sử dụng lệnh sau để gửi yêu cầu đến máy chủ vLLM (đã sửa đổi) và đo hiệu suất.

vllm bench serve \
  --backend vllm \
  --model "deepseek-ai/DeepSeek-R1" \
  --metric-percentiles "10,25,50,75,90" \
  --percentile-metrics "itl,tps,ttft,e2el" \
  --port 8001 \
  --num-prompts 2048 \
  --max-concurrency 2048 \
  --ignore-eos \
  --dataset-name sharegpt \
  --dataset-path /app/dataset/ShareGPT_V3_unfiltered_cleaned_split.json \
  --sharegpt-input-len 2000 \
  --sharegpt-output-len 100

Kết quả thực nghiệm

Chúng tôi đo thông lượng (mỗi node) và độ trễ giữa các token cho bốn mức đồng thời khác nhau (1024, 2048, 3072 và 4096) để quan sát sự đánh đổi giữa độ trễ và thông lượng. Kết quả như sau.

Chúng tôi đạt được thông lượng 21,224.6 tokens/sec ở mức đồng thời cao nhất. Con số này gần tương đương với 22,282 tokens/sec được báo cáo bởi đội ngũ SGLang trên máy chủ 8x NVIDIA H100 GPU, xác nhận rằng triển khai của chúng tôi mang lại hiệu suất EP tiên tiến nhất. Trong khi đó, ngay cả trong cấu hình tối ưu cho độ trễ tối thiểu (với độ trễ giữa các token trung vị là 91.48 ms), thông lượng chỉ giảm dưới 15%, cho thấy hệ thống có thể linh hoạt điều chỉnh mức đồng thời tối đa theo các SLO (Service Level Objectives) mục tiêu.

Kết luận

Moreh đã triển khai thành công suy luận thông lượng cao với EP trên GPU AMD Instinct MI300 series bằng cách áp dụng và tối ưu hóa nhiều kỹ thuật cho mô hình MoE. Chúng tôi cũng cung cấp MoAI Inference Framework, tự động áp dụng các kỹ thuật suy luận phân tán đa dạng bao gồm cả triển khai EP hiệu quả được trình bày trong bài viết này, trên các cluster GPU AMD Instinct. Để biết thêm thông tin, vui lòng truy cập trang web Moreh.