From 5ddc84e33e64c7891597ecb9a4346a13c94e94c0 Mon Sep 17 00:00:00 2001 From: Duyi-Wang Date: Mon, 16 Feb 2026 12:59:39 +0800 Subject: [PATCH] [AMD] MORI-EP inter kernel type switch (#18437) Co-authored-by: HAI --- docker/rocm.Dockerfile | 2 +- docker/rocm720.Dockerfile | 2 +- docs/references/environment_variables.md | 1 + .../srt/layers/moe/token_dispatcher/moriep.py | 18 +++++++++++++++--- 4 files changed, 18 insertions(+), 5 deletions(-) diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 98d854cd4..712646a1f 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -70,7 +70,7 @@ ARG ENABLE_MORI=0 ARG NIC_BACKEND=none ARG MORI_REPO="https://github.com/ROCm/mori.git" -ARG MORI_COMMIT="b0dce4beebeb1f26c784eee17d5fd9785ee9447f" +ARG MORI_COMMIT="20920706a9004018dbd87c7387f207d08d0e05af" # AMD AINIC apt repo settings ARG AINIC_VERSION=1.117.5 diff --git a/docker/rocm720.Dockerfile b/docker/rocm720.Dockerfile index 02845f428..421b3cacc 100644 --- a/docker/rocm720.Dockerfile +++ b/docker/rocm720.Dockerfile @@ -94,7 +94,7 @@ ARG ENABLE_MORI=0 ARG NIC_BACKEND=none ARG MORI_REPO="https://github.com/ROCm/mori.git" -ARG MORI_COMMIT="b0dce4beebeb1f26c784eee17d5fd9785ee9447f" +ARG MORI_COMMIT="20920706a9004018dbd87c7387f207d08d0e05af" # AMD AINIC apt repo settings ARG AINIC_VERSION=1.117.5 diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 9a42ae7b8..7812574ed 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -76,6 +76,7 @@ SGLang supports various environment variables that can be used to configure its | --- | --- | --- | | `SGLANG_MORI_FP8_DISP` | Use FP8 for dispatch | `"false"` | | `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | Maximum number of dispatch tokens per rank for MORI-EP buffer allocation | `4096` | +| `SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD` | Threshold for switching between `InterNodeV1` and `InterNodeV1LL` kernel types. `InterNodeV1LL` is used if `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` is less than or equal to this threshold; otherwise, `InterNodeV1` is used. | `256` | | `SGLANG_MORI_QP_PER_TRANSFER` | Number of RDMA Queue Pairs (QPs) used per transfer operation | `1` | | `SGLANG_MORI_POST_BATCH_SIZE` | Number of RDMA work requests posted in a single batch to each QP | `-1` | | `SGLANG_MORI_NUM_WORKERS` | Number of worker threads in the RDMA executor thread pool | `1` | diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 84d4b3c7c..6ee2443b4 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -85,9 +85,21 @@ class EpDispatchConfig: rdma_block_num: int -def get_ep_dispatch_configs(): +def get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank: int = 4096): import mori + # Selects the inter-node kernel. `InterNodeV1LL` is used if `num_max_dispatch_tokens_per_rank` + # is less than or equal to the threshold, otherwise `InterNodeV1` is used. The threshold defaults to 256. + inter_kernel_switch_threshold = get_int_env_var( + "SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD", 256 + ) + + inter_kernel_type = ( + mori.ops.EpDispatchCombineKernelType.InterNodeV1LL + if num_max_dispatch_tokens_per_rank <= inter_kernel_switch_threshold + else mori.ops.EpDispatchCombineKernelType.InterNodeV1 + ) + return { EpMode.INTRA_NODE: EpDispatchConfig( kernel_type=mori.ops.EpDispatchCombineKernelType.IntraNode, @@ -96,7 +108,7 @@ def get_ep_dispatch_configs(): rdma_block_num=0, ), EpMode.INTER_NODE: EpDispatchConfig( - kernel_type=mori.ops.EpDispatchCombineKernelType.InterNodeV1, + kernel_type=inter_kernel_type, warp_num_per_block=8, block_num=64, rdma_block_num=32, @@ -130,7 +142,7 @@ def init_mori_op( ) mode = EpMode.INTRA_NODE if world_size <= 8 else EpMode.INTER_NODE - cfg = get_ep_dispatch_configs()[mode] + cfg = get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank)[mode] kernel_type = cfg.kernel_type warp_num_per_block = cfg.warp_num_per_block