From a3bf25dc620f31fc672aeced1465d6fe7c81d28f Mon Sep 17 00:00:00 2001 From: inkcherry Date: Wed, 16 Sep 2026 17:13:59 +0800 Subject: [PATCH] [AMD] Clamp MORI intranode grid GPUs (#39763) Co-authored-by: Duyi-Wang --- python/sglang/srt/layers/moe/token_dispatcher/moriep.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index f97221278..37a85ee6a 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -319,6 +319,15 @@ def init_mori_op( block_num = 256 warp_num_per_block = 16 + if kernel_type == mori.ops.EpDispatchCombineKernelType.IntraNode: + # mori's intranode combine ends in a grid-wide barrier every block spins + # on. A 256-block grid deadlocks on a 192-CU device; clamping to the CU + # count is a conservative bound, not the measured residency limit. + cu_count = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + block_num = min(block_num, cu_count) + # Fp8 blockwise combine uses its own internal scale_dim driven which can be # overridden by env ``MORI_FP8_COMBINE_SCALE_DIM`` (default 56) # See https://github.com/ROCm/mori/blob/96ffa169710f214e76e07abe5008d686fe54522b/python/mori/ops/dispatch_combine.py#L81-L84