minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] MORI-EP support for EP4. (#19578)

Browse Source
This commit is contained in:
Duyi-Wang
2026-02-28 13:13:46 -08:00
committed by GitHub
parent 560b867cce
commit 8240a87306
2 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/rocm.Dockerfile
+1 -1
View File
@@ -98,7 +98,7 @@ ARG ENABLE_MORI=0
ARG NIC_BACKEND=none ARG NIC_BACKEND=none
ARG MORI_REPO="https://github.com/ROCm/mori.git" ARG MORI_REPO="https://github.com/ROCm/mori.git"
ARG MORI_COMMIT="20920706a9004018dbd87c7387f207d08d0e05af" ARG MORI_COMMIT="2f88d06aba75400262ca5c1ca5986cf1fdf4cd82"
# AMD AINIC apt repo settings # AMD AINIC apt repo settings
ARG AINIC_VERSION=1.117.5 ARG AINIC_VERSION=1.117.5
python/sglang/srt/layers/moe/token_dispatcher/moriep.py
+3
View File
@@ -189,6 +189,8 @@ def init_mori_op(
world_size = get_moe_expert_parallel_world_size() world_size = get_moe_expert_parallel_world_size()
rank = get_moe_expert_parallel_rank() rank = get_moe_expert_parallel_rank()
gpu_per_node = 8 if world_size >= 8 else world_size
cpu_group = group.cpu_group cpu_group = group.cpu_group
torch._C._distributed_c10d._register_process_group("mori", cpu_group) torch._C._distributed_c10d._register_process_group("mori", cpu_group)
mori.shmem.shmem_torch_process_group_init("mori") mori.shmem.shmem_torch_process_group_init("mori")
@@ -227,6 +229,7 @@ def init_mori_op(
warp_num_per_block=warp_num_per_block, warp_num_per_block=warp_num_per_block,
block_num=block_num, block_num=block_num,
kernel_type=kernel_type, kernel_type=kernel_type,
gpu_per_node=gpu_per_node,
rdma_block_num=rdma_block_num, rdma_block_num=rdma_block_num,
num_qp_per_pe=2, num_qp_per_pe=2,
) )