新智元報(bào)道

編輯：LRST

【新智元導(dǎo)讀】通過(guò)完全啟用并發(fā)多塊執(zhí)行，支持任意專家數(shù)量（MAX_EXPERT_NUMBER==256），并積極利用共享內(nèi)存（5kB LDS）和寄存器（52 VGPRs，48 SGPRs），MoE Align & Sort邏輯被精心設(shè)計(jì)，實(shí)現(xiàn)了顯著的性能提升：A100提升3倍，H200提升3倍，MI100提升10倍，MI300X/MI300A提升7倍...

MoE（Mixture of Experts）模型模仿了人腦的低功耗運(yùn)作模式：功能被劃分為多個(gè)獨(dú)立的部分，在思考時(shí)通過(guò)自適應(yīng)路由部分激活，從而提高計(jì)算效率。

牛津大學(xué)研究論文中的人腦皮層示意圖，來(lái)源于互聯(lián)網(wǎng)

首個(gè)可在CUDA真正可行的版本是Switch Transformer[1]，隨后通過(guò)循環(huán)利用(Up Cycling)稠密模型Mistral[2]進(jìn)一步優(yōu)化了該設(shè)計(jì)。

SwitchTransformer-MoE

隨后，DeepSeek V2/V3/R1[3][4][5]通過(guò)引入共享專家[3]和門(mén)控偏差（gating bias）[4][5]進(jìn)一步改進(jìn)了MoE，最終實(shí)現(xiàn)了無(wú)輔助損失（auxiliary loss free）的MoE模型 [4][5]。這一優(yōu)化本質(zhì)上歸因于一個(gè)關(guān)鍵事實(shí)：當(dāng)使用共享專家（DeepSeek團(tuán)隊(duì)選擇的值為1）時(shí)，可以通過(guò)在較大的專家池（256個(gè)上施加偏差分?jǐn)?shù)的懲罰，從而緩解專家路由的不均衡問(wèn)題[11]。

MoE層本質(zhì)上是由多個(gè)專家前饋網(wǎng)絡(luò)（FFN）組成的層，其中包含門(mén)控函數(shù)（gating functions），用于根據(jù)Top-K門(mén)控分?jǐn)?shù)（DeepSeek V3/R1中引入偏差）進(jìn)行激活路由，并在所選的FFN層上通過(guò)Group GEMM計(jì)算logits。

該功能在很大程度上依賴于基數(shù)排序（radix sort）邏輯。借助MoE Align & Sort，機(jī)器學(xué)習(xí)研究人員和實(shí)踐者可以按照專家ID對(duì)tokens進(jìn)行排序。

在某些應(yīng)用中，例如TransformerEngine[6][7]，該操作最初是通過(guò)已廢棄的cub::DeviceRadixSort實(shí)現(xiàn)的，而增加的permute操作用于記錄源（左）到目標(biāo)（右）的映射，其梯度操作為unpermute。

MoE Permute示例

盡管cub::DeviceRadixSort大量使用共享內(nèi)存，相比于基于__shfl_xor_sync（僅使用線程本地內(nèi)存）的實(shí)現(xiàn)略慢，但它不支持對(duì)齊排序（alignment sorting）。

對(duì)齊排序?qū)τ贕roup GEMM的效率至關(guān)重要，因?yàn)樗试S專家以塊（block 為單位處理tokens。

SGLang 中的MoE Align & Sort算法采用了對(duì)齊排序，但在支持多達(dá)256個(gè)專家的大規(guī)模prefill操作時(shí)效率并不理想。該問(wèn)題已在issue#2732中被確認(rèn)。

目前的實(shí)現(xiàn)將MoE Align & Sort拆分為兩個(gè)kernel啟動(dòng)（kernel launches）：

  • 對(duì)齊（alignment）：在單個(gè)block內(nèi)執(zhí)行傳統(tǒng)基數(shù)排序算法對(duì)齊后的偏移計(jì)算（alignment-based offsets computation）;

  • 放置（placement）：根據(jù)在多個(gè)block并行計(jì)算出的偏移量，并行放置tokens;

研究人員提出并編寫(xiě)了AMD友好的CUDA設(shè)備代碼，采用了該設(shè)計(jì)的MoE Align & Sort算法。因此，在AMD平臺(tái)上的性能分析和優(yōu)化將被充分考慮。

打開(kāi)網(wǎng)易新聞 查看精彩圖片

文章地址：https://shorturl.at/C23JF

通過(guò)在不同的工作負(fù)載下使用RocProfiler-Compute進(jìn)行分析，可以清楚地看到，即使不計(jì)入多次設(shè)備函數(shù)啟動(dòng)的額外開(kāi)銷，第一個(gè)kernel仍然消耗了33W個(gè)周期，第二個(gè)kernel消耗了8W個(gè)周期，總計(jì)41W周期：

the moe align kernel 1

the moe align kernel 2

在ROCm SDK 6.3.0 中，omniperf已更名為rocprof-compute。盡管MI300X/MI300A已得到積極支持，但該工具默認(rèn)未隨ROCm SDK 6.3.0一同發(fā)布。不過(guò)，在Tools-dockerhub中的展示一樣，ROCm計(jì)算分析工具的設(shè)置僅需簡(jiǎn)單三步。

現(xiàn)在，在PR#3613（https://shorturl.at/OPbiI）中應(yīng)用優(yōu)化方案后，片上計(jì)算開(kāi)銷將從之前的41W個(gè)周期立即降低至20W個(gè)周期。

在SGLang中實(shí)現(xiàn)高效的多塊（multi-blocks）MoE-Align

通過(guò)完全地多塊（multiple blocks）并發(fā)執(zhí)行，并支持任意專家數(shù)量（MAX_EXPERT_NUMBER==256），結(jié)合激進(jìn)使用共享內(nèi)存（5kB LDS）和寄存器（52 VGPRs，48 SGPRs），MoE Align & Sort邏輯被優(yōu)化，實(shí)現(xiàn)了以下性能提升3x in A100,3x in H200, 10x in MI100, and 7x in MI300X/Mi300A:

借助Rocprof-Compute，可以輕松收集捕獲設(shè)備代碼的一些關(guān)鍵指標(biāo)，并在遠(yuǎn)程GUI服務(wù)器上進(jìn)行可視化展示：

服務(wù)端開(kāi)啟Rocprof-Compute

總而言之，在AMD MI300X/MI300A上，所提出的高效多塊（multi-blocks）MoE Align & Sort算法充分利用了每個(gè)wave的向量寄存器（52個(gè)），且無(wú)寄存器溢出（我已將初始線程塊大小調(diào)整至最佳值）；同時(shí)，每個(gè)CU使用5kB LDS，且僅有6.8%的存儲(chǔ)銀行沖突率。

研究人員還分析了MoE Sort & Align的Roofline模型。該模型顯示，設(shè)備代碼的性能在受限于內(nèi)存帶寬的區(qū)域有所下降。

在AMD Compute Profile部分，研究人員詳細(xì)介紹了在ROCm平臺(tái)上算法設(shè)計(jì)的影響與性能數(shù)據(jù)。

本質(zhì)上，MI300X/MI300A是全球首款基于多芯片（multi-die）設(shè)計(jì)的高性能AI加速器架構(gòu)。因此，在該芯片上進(jìn)行算子優(yōu)化的方式將與NVIDIA平臺(tái)略有不同。

基本規(guī)則是，XCDs（加速計(jì)算芯片）之間的同步代價(jià)較高，因此最好充分利用XCDs，并利用L2緩存的局部性親和性來(lái)提高性能。

此外，研究人員應(yīng)避免昂貴的同步開(kāi)銷，具體方法包括：

  • 當(dāng)網(wǎng)格大小小于每顆芯片上的XCD數(shù)量（MI300X為8，MI300A為6）時(shí)，優(yōu)先使用最低速計(jì)算單元（MI300X使用XCD7，MI300A使用XCD5）。

  • 當(dāng)網(wǎng)格大小大于每顆芯片上的XCD數(shù)量時(shí)，將其調(diào)整為XCD數(shù)量的整數(shù)倍。

使用hipCooperativeLaunch啟動(dòng)協(xié)作設(shè)備代碼可能會(huì)增加L2緩存壓力（與紋理尋址器停滯率和忙碌率相關(guān)），特別是在數(shù)據(jù)交換（尤其是Die-Die交換增多的情況下。

在此示例中，之前main分支的實(shí)現(xiàn)使用了39個(gè)活躍CU，這已經(jīng)接近最佳，因?yàn)楸举|(zhì)上使用了兩個(gè)Die。

該實(shí)現(xiàn)在多塊（multi-blocks）執(zhí)行中使用了66個(gè)活躍CU，跨越兩個(gè)Die，并且塊級(jí)歸約（block-wise reduction）過(guò)程中Die-Die數(shù)據(jù)交換是不可避免的，將在本季度晚些時(shí)候向SGLang提交進(jìn)一步的V4優(yōu)化。

具體細(xì)節(jié)將在性能分析（profiling）部分進(jìn)一步討論。

SGLang中Fused MoE的回顧

SGLang團(tuán)隊(duì)采用Triton First方法實(shí)現(xiàn)了相關(guān)邏輯，并在2024年12月成功實(shí)現(xiàn)DeepSeek V3的Day-0支持。

SGLang的MoE調(diào)用了使用Triton實(shí)現(xiàn)的Fused MoE 設(shè)備代碼。

在設(shè)備代碼啟動(dòng)之前，會(huì)應(yīng)用MoE Align & Sort算法。MoE Align & Sort的Triton設(shè)備代碼被拆分為四個(gè)階段，其中直接訪問(wèn)DRAM，而不使用共享內(nèi)存，這與向量化 Triton版本形成對(duì)比。

與單塊（single block wise）CUDA實(shí)現(xiàn)相比，Triton版本的多次設(shè)備代碼觸發(fā)以及對(duì)LDS、本地緩存和寄存器（例如VGPR）的低效利用，導(dǎo)致了在小規(guī)模工作負(fù)載上的單次測(cè)試執(zhí)行效率較低。

隨后，CUDA實(shí)現(xiàn)最終被拆分為兩個(gè)階段，其中僅第二階段的執(zhí)行在多塊（multiple blocks）上進(jìn)行了加速。

MoE Align & Sort CUDA算法在其他開(kāi)源平臺(tái)的實(shí)現(xiàn)

FasterTransfomer

在Mistral[2]和DeepSeek V2[3]之前，開(kāi)放式稠密模型（open dense models）在推理場(chǎng)景中更為流行。這也是FasterTransformer[8]誕生的時(shí)期。

在FasterTransformer[8]項(xiàng)目中（由NVIDIA發(fā)起），MoE模型的支持主要依賴于cub::DeviceRadixSort，以及諸如moe_softmax（本質(zhì)上是cub::BlockReduce實(shí)現(xiàn)的softmax）、moe_top_k及其融合版本topk_gating_softmax、用于排列潛在向量logits的permute，最終執(zhí)行g(shù)roup gemm。

因此，融合優(yōu)化主要（按計(jì)算開(kāi)銷計(jì)算）限制在topk gating softmax和biased topk gating softmax，后續(xù)這些優(yōu)化被整合進(jìn)SGLang。

Megatron

在該算法發(fā)表之前，Megatron在FP16/BF16計(jì)算中主要采用FasterTransformer方法，但額外添加了permute的梯度操作unpermute，以支持訓(xùn)練任務(wù)。

這意味著MoE仍然沒(méi)有得到高效融合。

vLLM

SGLang使用了許多vLLM設(shè)備代碼，但vLLM的Fused MoE最初是由SGLang團(tuán)隊(duì)貢獻(xiàn)的。因此，它們采用了相同的方法進(jìn)行部署。

CK

首個(gè)AMD友好的Fused MoE版本于2024年11月26日在CK#1634（https://tinyurl.com/3fuj7yws）中提出。隨后，MoE Align & Sort被添加到CK#1771（https://tinyurl.com/5h4e8jat）和CK#1840（https://tinyurl.com/3wm8pdc3）中。

核心思路是將MoE 排序與Group GEMM進(jìn)行融合。此外，CK中的MoE & Sorting在很大程度上采用了SGLang團(tuán)隊(duì)的方法，但在CK pipeline及partitioner方面有所不同。

CK融合MoE思路[9]

融合per_group_token_quant（用于在線FP8量化）、MoE排序和Group GEMM可以通過(guò)將Radix Sort計(jì)算邏輯納入Group GEMM pipeline輕松解決：即統(tǒng)計(jì)出現(xiàn)次數(shù)以計(jì)算偏移量，隨后進(jìn)行并行放置。

其中最關(guān)鍵的問(wèn)題之一是如何平衡Radix Sorting和Group GEMM這兩種計(jì)算負(fù)載。

在AMD數(shù)據(jù)中心芯片中，Group GEMM片段更可能均勻分布在XCD內(nèi)的所有可用計(jì)算單元。然而，當(dāng)涉及多個(gè)XCD時(shí)，不同CU之間的數(shù)據(jù)交換主要通過(guò)低速L2 Cache及其互聯(lián)結(jié)構(gòu)（L2 Cache fabric）進(jìn)行。

編寫(xiě)CK設(shè)備代碼需要先編寫(xiě)主機(jī)端CK解決方案啟動(dòng)器：

// Here is the entry of fused MoE :
//   https://github.com/ROCm/composable_kernel/blob/1342ecf7fbf64f43d8621cf6665c583fdc49b2c6/example/ck_tile/15_fused_moe/instances/fused_moegemm_api_internal.hpp
using f_pipeline    = ck_tile::FusedMoeGemmPipeline_FlatmmUk ;
using f_partitioner = ck_tile::FusedMoeGemmTilePartitioner_Linear ;
using f_kernel      = ck_tile::FusedMoeGemmKernel void >;
const dim3 grids                       = f_kernel::GridSize(a);
constexpr dim3 blocks                  = f_kernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = 1;
static int printed = 0;
auto kargs = f_kernel::MakeKargs(a);
if(s.log_level_ > 0 && printed == 0)
{
std::cout << ", " << f_kernel::GetName() << std::flush;
printed = 1;
}
return ck_tile::launch_kernel(
s, ck_tile::make_kernel (f_kernel{}, grids, blocks,  0, kargs));

AMD CK分區(qū)器和階段流水線（stages pipeliner）在Fused MoE的最終匯編過(guò)程中扮演了重要角色，確實(shí)值得深入研究，但已超出本文討論范圍。

但需要記住，MoE Align & Sort是生產(chǎn)者代碼的一部分：

// https://github.com/ROCm/composable_kernel/blame/fdaff5603ebae7f8eddd070fcc02941d84f20538/include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp#L438
CK_TILE_DEVICE void moe_align_block_size_kernel(...)
{
const index_t tid       = static_cast
             
 (threadIdx.x); const index_t start_idx = tid * tokens_per_thread; ... #if 1 if(tid < num_experts){ // each thread reduce a column segment of tokens_cnts with # blockDim.x elements ... } #else ... #endif __syncthreads(); // do cumsum to compute offsets based on condition ... // do parallel placement based on the offsets computed ... }

因此，在AMD CK方案中，MoE Align & Sort的實(shí)現(xiàn)幾乎與SGLang主實(shí)現(xiàn)保持一致，僅在分區(qū)器（partitioner）和流水線（pipeliner）方面有所不同。

需要注意的是，該實(shí)現(xiàn)并不總是能在AMD平臺(tái)上提供最佳性能（請(qǐng)參考AITER中的asm MoE）。

由于AMD CDNA3架構(gòu)并不支持類似Graphcore的片上（on-chip）洗牌操作（在2023年已經(jīng)將PopART[12] & PopRT的Remapping操作進(jìn)行抽象與泛化），而這一特性已在NVIDIA H100/H200/B200中得到了支持，并通過(guò)高效的SM<->SM片上通信實(shí)現(xiàn)。

因此，在AMD 開(kāi)源解決方案中，如何以低開(kāi)銷方式在塊（block）之間優(yōu)化數(shù)據(jù)布局將是一個(gè)非常有趣的研究方向。

從哲學(xué)上講，這兩類不同工作負(fù)載的基于 Tiling 的融合代碼可能并不總是比非融合版本更優(yōu)。相關(guān)研究的詳細(xì)內(nèi)容將在V4 版本發(fā)布時(shí)進(jìn)一步探討。

AITER

AI Tensor Engine[10]

AITER在今年早些時(shí)候被引入，以便整合在不同項(xiàng)目中使用的LLM設(shè)備代碼。它通過(guò)ck moe、asm 版本的 MoE 通過(guò) hipModule和triton fused moe支持MoE融合。

因此，AITER是部分開(kāi)源的，因?yàn)椴煌该鞯膮R編代碼和開(kāi)發(fā)計(jì)劃是針對(duì)MI300X開(kāi)發(fā)者的。

AITER中fused MoE的三倍加速[10]已由Bruce Xu[13]驗(yàn)證，并且這一加速主要來(lái)自于在不同形狀的Group GEMM中觀察到的加速：一個(gè)GEMM操作，其中每個(gè)專家的FFN權(quán)重與一塊隱藏狀態(tài)的token進(jìn)行相乘。

這一證明可以在PR#199（https://shorturl.at/F8y0F）中找到，asm gemm幾乎帶來(lái)了三倍的性能提升。

ASM版本扁平矩陣乘

值得注意的是，仍然有一些情況下，選擇了來(lái)自SGLang社區(qū)的triton設(shè)備代碼。為了在MI300X/MI300A上高效運(yùn)行triton設(shè)備代碼，它們采用了基于多芯片架構(gòu)的特定邏輯，將線程塊映射到不同的計(jì)算單元（dies）上：

# https://github.com/ROCm/triton/blob/f669d3038f4c03ee7a60835e875937c65b5cec35/python/perf-kernels/gemm.py#L115
...
## pid remapping on xcds
# Number of pids per XCD in the new arrangement
pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
# When GRID_MN cannot divide NUM_XCDS, some xcds will have
# pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
# We calculate the number of xcds that have pids_per_xcd pids as
# tall_xcds
tall_xcds = GRID_MN % NUM_XCDS
tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
# Compute current XCD and local pid within the XCD
xcd = pid % NUM_XCDS
local_pid = pid // NUM_XCDS
# Calculate new pid based on the new grouping
# Note that we need to consider the following two cases:
# 1. the current pid is on a tall xcd
# 2. the current pid is on a short xcd
if xcd < tall_xcds:
pid = xcd * pids_per_xcd + local_pid
else:
pid = tall_xcds * pids_per_xcd + (xcd - tall_xcds) * (pids_per_xcd - 1) + local_pid
if GROUP_SIZE_M == 1:
pid_m = pid // num_pid_n
pid_n = pid % num_pid_n
else:
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
...

此外，在CK fused MoE中使用了多種AMD芯片內(nèi)建函數(shù)（intrinsics），例如：

  • __builtin_nontemporal_load,

  • __builtin_amdgcn_ds_swizzle,

  • __builtin_amdgcn_ds_permute/__builtin_amdgcn_ds_bpermute,

  • _builtin_amdgcn_mov_dpp

等等。這些內(nèi)建函數(shù)可能最終影響fused MoE的匯編實(shí)現(xiàn)和性能。

例如，使用__builtin_nontemporal_load可以跳過(guò)L2緩存，從而為預(yù)測(cè)將被重復(fù)使用的數(shù)據(jù)留出更多L2緩存行空間。

Cutlass v3.8

Fused MoE尚未在NVIDIA Cutlass 3.8.0中公開(kāi)支持。因此，當(dāng)前該倉(cāng)庫(kù)中沒(méi)有提供MoE Align & Sort功能。

TRT-LLM

在v0.16.0之前，TRT-LLM基本上遵循了FasterTransformer的方法。自v0.17.0版本起，MoE部分開(kāi)始公開(kāi)。

編寫(xiě)對(duì)AMD設(shè)備友好的CUDA實(shí)現(xiàn)，并帶來(lái)超過(guò)3x~7x加速

該算法采用了多塊執(zhí)行方案，并由三個(gè)不同的部分（D-C-P）組成：

  • 分布式并發(fā)計(jì)數(shù)

  • 計(jì)算累積和（cumsum）

    • 并行非對(duì)齊本地累積和

    • 合并非對(duì)齊累積和

    • 對(duì)齊全局累積和

    • 存儲(chǔ)全局累積和

  • 并行放置

高效MoE Align& Sort算法

并行非對(duì)齊本地累積和

并行非對(duì)齊本地累積和

該算法首次由在PR#2970（https://shorturl.at/CuBs5）中提出并實(shí)現(xiàn)。

研究人員將每個(gè)塊中的累積和執(zhí)行進(jìn)行了負(fù)載均衡，分配給kElementsPerThr(16)個(gè)線程，每個(gè)線程需要處理kElementsPerThr+kElementsPerThr+threadIdx.x次加法操作。

因此，與當(dāng)前倉(cāng)庫(kù)中的單線程版本相比，波前（wavefront）更快地到達(dá)，可以觀察到此版本實(shí)現(xiàn)的性能提升了30%。

合并非對(duì)齊累積和

一旦獲得了每個(gè)塊中的本地非對(duì)齊累積和（Unaligned Cumsum），就可以在預(yù)分配的HBM緩沖區(qū)中進(jìn)行塊級(jí)別的累積和歸約。

研究人員選擇了FRAG_SIZE_M(16)xFRAG_SIZE_N(16)xFRAGS_PER_BLOCK(4)的SRAM塊進(jìn)行塊級(jí)歸約，其中FRAGS_PER_BLOCK是可調(diào)的：

打開(kāi)網(wǎng)易新聞 查看精彩圖片

塊級(jí)歸約

在AMD平臺(tái)上，計(jì)算是基于「1 warp加載/1warp計(jì)算」的方式進(jìn)行的，而在NVIDIA平臺(tái)上則是「2warps加載和1warp計(jì)算」。

該設(shè)計(jì)充分利用了AMD CDNA3架構(gòu)中64個(gè)SIMD通道的優(yōu)勢(shì)。并且，在這種多芯片架構(gòu)中，塊的數(shù)量始終是XCD數(shù)量的倍數(shù)。

FRAGS_PER_BLOCK被設(shè)置為4，以便在多輪中復(fù)用SMEM。

對(duì)齊全局累積和和存儲(chǔ)全局累積和

研究人員改進(jìn)了向量化代碼，并處理了如果輸入數(shù)據(jù)大小與kElementsPerAccess常量不對(duì)齊時(shí)的循環(huán)尾部情況。

基準(zhǔn)測(cè)試顯示，合并率有所提高，但仍然限制在30%左右。研究人員將在V4版本中繼續(xù)優(yōu)化此問(wèn)題。

編寫(xiě)AMD友好的CUDA代碼

編寫(xiě)PyTorch擴(kuò)展可以自動(dòng)將CUDA設(shè)備代碼轉(zhuǎn)換為HIP設(shè)備代碼，配合ROCm SDK進(jìn)行使用。

但是，有些情況下HIP設(shè)備代碼與CUDA設(shè)備代碼表現(xiàn)不同：

  • Warp大小是一個(gè)與架構(gòu)相關(guān)的全局變量，并在ROCm SDK中定義為warpSize；在CDNA3架構(gòu)中，warpSize定義為64。

  • 設(shè)備函數(shù)簽名可能與CUDA不完全對(duì)齊，因此需要條件編譯來(lái)支持這些符號(hào)。

  • 需要特別關(guān)注多芯片架構(gòu)中的L2緩存優(yōu)化。

基準(zhǔn)測(cè)試

在沒(méi)有CUDA圖捕獲的情況下，研究人員針對(duì)DeepSeek V3模型的大規(guī)模工作負(fù)載進(jìn)行了廣泛測(cè)試。因此，專家數(shù)量設(shè)置為256。當(dāng)前的算法不支持在CUDA圖捕獲下運(yùn)行，將在V4版本中解決此問(wèn)題。

由于GPU虛擬化和測(cè)試節(jié)點(diǎn)上分配的CPU數(shù)量，性能可能會(huì)與裸機(jī)測(cè)試時(shí)有所不同。

因此，研究人員使用Triton實(shí)現(xiàn)作為基準(zhǔn)，展示MoE Align & Sort算法在加速倍數(shù)和效率上的表現(xiàn)。

每個(gè)測(cè)試首先進(jìn)行了驗(yàn)證，之后才開(kāi)始基準(zhǔn)測(cè)試。在基準(zhǔn)測(cè)試中，可以觀察到，在AMD平臺(tái)上，Triton的運(yùn)行時(shí)間顯著長(zhǎng)于在NVIDIA平臺(tái)上的運(yùn)行時(shí)間，因此建議進(jìn)一步優(yōu)化Triton的MLIR，以獲得比NVIDIA Triton更高效的降級(jí)過(guò)程。

對(duì)于AMD Triton，可以觀察到MI300X的速度比MI100快1.5倍，因此MI300X的性能提升幅度不像MI100那么顯著。此外，盡管普遍認(rèn)為MI300X比MI100更快，但在測(cè)試中，MI100上的算法性能要優(yōu)于MI300X。

這部分歸因于內(nèi)存瓶頸操作，在多芯片之間的通信降低了執(zhí)行速度。

在兩個(gè)平臺(tái)上，都觀察到了應(yīng)用該算法后顯著的性能改進(jìn)，其中現(xiàn)有的CUDA實(shí)現(xiàn)幾乎與Triton消耗相同的時(shí)間。

AMD系統(tǒng)準(zhǔn)備

為了最大化使用AMD異構(gòu)系統(tǒng)，建議進(jìn)行以下檢查。

  • NVIDIA Grace CPU和AMD EPYC 9004系統(tǒng)通常建議禁用NUMA自動(dòng)平衡，以便與GPU協(xié)同工作；然而，在某些情況下，可能不建議禁用 NUMA自動(dòng)平衡。

  • 啟用虛擬化時(shí)，建議啟用IOMMU直通模式，以消除DMA翻譯，從而帶來(lái)性能提升。


MI100基準(zhǔn)測(cè)試

git clone 
https://github.com/yiakwy-xpu-ml-framework-team/AMD-sglang-benchmark-fork.git
 -b optimize_moe_align_v3 && cd sgl-kernel && python setup_rocm.py install

可以驗(yàn)證不同輸入令牌和專家數(shù)量組合的可行性 :

cd ../benchmark/kernels/fused_moe_trition && python benchmark_deepseekv3_moe_align_blocks.py --verify

A100 性能測(cè)試

H200 性能測(cè)試

MI300X 性能測(cè)試

AMD Compute Profile


設(shè)置

在ROCm 6.3.3版本中，設(shè)置rocprof-compute只需三步即可完成，詳細(xì)的設(shè)置步驟可以在這里找到：Tools-dockerhub中的rocprof-compute設(shè)置。

向量L1緩存的分析結(jié)果

在分析中，工作負(fù)載為16384個(gè)tokens x（從256個(gè)專家中選擇8個(gè)），除非另有說(shuō)明。

研究人員在算法中最大化了VGPRs的使用，但減少了SGPRs的總使用量。數(shù)據(jù)也表明，VGPRs/SGPRs的溢出為零，這表明寄存器的使用是健康的，并且此設(shè)備代碼沒(méi)有性能損失。

向量L1緩存（vL1D）是每個(gè)CU的本地單元，命中率記錄了從L2緩存請(qǐng)求到CU時(shí)的緩存行命中率。30%的L2緩存請(qǐng)求通過(guò)vL1D的紋理尋址器合并，達(dá)到了61%的命中率，如果需要，稍后可以進(jìn)一步提升。

當(dāng)數(shù)據(jù)從CU請(qǐng)求到vL1D的尋址處理單元（紋理尋址器）時(shí)，復(fù)雜的決策邏輯決定是否接受數(shù)據(jù)請(qǐng)求或回滾數(shù)據(jù)請(qǐng)求。以下是四種狀態(tài)：

  • Busy（忙碌）：紋理尋址器正在處理地址。

  • Address Stall（地址停頓）：紋理尋址器無(wú)法發(fā)送地址到vL1D。

  • Data Sending Stall（數(shù)據(jù)發(fā)送停頓）：紋理尋址器無(wú)法發(fā)送數(shù)據(jù)到vL1D。

  • Data Waiting Stall（數(shù)據(jù)等待停頓）：紋理尋址器等待發(fā)送數(shù)據(jù)到vL1D的數(shù)據(jù)處理單元。

有關(guān)這種微架構(gòu)行為的詳細(xì)信息，可以在AMD CDNA3的ISA文檔以及rocProfiler-compute文檔中找到。

vL1D 尋址器停頓

研究人員在該算法設(shè)計(jì)中觀察到了18.61%的數(shù)據(jù)等待停頓率來(lái)自于向量L1緩存。

數(shù)據(jù)的讀寫(xiě)負(fù)載平衡大大減少，從8kB的讀取操作和27B的寫(xiě)入操作，轉(zhuǎn)變?yōu)?strong>109B的讀取操作，468B的寫(xiě)入操作和202B的原子操作的組合。

L2緩存的分析結(jié)果

在CDNA3架構(gòu)中，L2緩存是所有計(jì)算單元（CU）共享的，且是線程塊之間共享數(shù)據(jù)的主要通道，這些線程塊分布在不同的CUs上。

通過(guò)多通道和地址交錯(cuò)設(shè)計(jì)，向L2緩存的請(qǐng)求可以大大并行處理。

此外，使用AMD特有的內(nèi)置函數(shù)如__builtin_nontemporal_load，可以繞過(guò)L2緩存來(lái)處理那些不需要再次訪問(wèn)的數(shù)據(jù)。

更多L2緩存研究細(xì)節(jié)將在V4版本中揭示。

結(jié)論

新的算法通過(guò)最大化使用LDS和向量寄存器，顯著加速了CUDA和ROCm平臺(tái)上的MoE Align & Sort，提升幅度高達(dá)3x~7x。

還可以觀察到，相較于單個(gè)芯片，內(nèi)存密集型操作在多晶粒異構(gòu)集成架構(gòu)下可能表現(xiàn)更差，這表明在多芯片如MI300X/MI300A和B200/B300設(shè)備上編程時(shí)，可能需要新的微調(diào)方向。

然而，該算法的細(xì)節(jié)仍有進(jìn)一步優(yōu)化空間，以提高緩存命中率和主內(nèi)存合并率。

致謝

特別感謝來(lái)自NUS團(tuán)隊(duì)的覃含章教授，王昀鴻博士在MI100/MI250性能驗(yàn)證中的合作，Zev Rekhter在MI300X性能驗(yàn)證中的合作，范舒宜在H200驗(yàn)證中的合作，以及BBuf在SGLang解決方案的討論和審閱。

請(qǐng)注意，這是SGLang社區(qū)的獨(dú)立工作。

作者介紹

本文作者王翼之前在Graphcore擔(dān)任機(jī)器學(xué)習(xí)專家，后加入美國(guó)知名半導(dǎo)體公司擔(dān)任AI架構(gòu)師（SMTS主任工程師）。

參與貢獻(xiàn)諸多機(jī)器學(xué)習(xí)社區(qū)開(kāi)源軟件，主要研究興趣在LLM訓(xùn)練、推理的軟件棧優(yōu)化，并應(yīng)用計(jì)算機(jī)體系結(jié)構(gòu)知識(shí)協(xié)同設(shè)計(jì)軟硬件解決方案。

參考資料：

[1]W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. CoRR, abs/2101.03961, 2021. URL https://arxiv.org/abs/2101.03961.

[2]A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. d. l. Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.

[3]DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. CoRR, abs/2405.04434, 2024c. URL https://doi.org/10.48550/arXiv.2405.04434.

[4]DeepSeek V3 : https://arxiv.org/abs/2412.19437; Retrieved on 2025-03-18

[5]DeepSeek R1 : https://arxiv.org/pdf/2501.12948; Retrieved on 2025-03-18

[6]TransformerEngine : https://github.com/NVIDIA/TransformerEngine; Retrieved on 2025-03-18

[7]NV Group GEMM : https://github.com/yiakwy-xpu-ml-framework-team/NV_grouped_gemm; Retrieved on 2025-03-18

[8]FasterTransformer : https://github.com/NVIDIA/FasterTransformer; Retrieved on 2025-03-18

[9]CK Fused MoE V1 : ROCm/composable_kernel#1634

[10]AMD 3X MOE : https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html

[11]Lean Wang and Huazuo Gao and Chenggang Zhao and Xu Sun and Damai Dai Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts, 2024. URL https://arxiv.org/abs/2408.15664.

[12]PopART on chip TensorRemap : https://github.com/graphcore/popart/tree/sdk-release-3.4

[13] DeepSeek V3 Optimization based on AITER backend : sgl-project/sglang#4344

原文地址：

Github（https://shorturl.at/bEEn9），Hugging Face（https://shorturl.at/PMH3F）