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

[AMD] Clamp MORI intranode grid GPUs (#39763)

Browse Source 
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>
This commit is contained in:
inkcherry
2026-09-16 02:13:59 -07:00
committed by GitHub
co-authored by Duyi-Wang
parent 7d5696b3a1
commit a3bf25dc62
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/moriep.py
+9
View File
@@ -319,6 +319,15 @@ def init_mori_op(
block_num = 256 block_num = 256
warp_num_per_block = 16 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 # Fp8 blockwise combine uses its own internal scale_dim driven which can be
# overridden by env ``MORI_FP8_COMBINE_SCALE_DIM`` (default 56) # 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 # See https://github.com/ROCm/mori/blob/96ffa169710f214e76e07abe5008d686fe54522b/python/mori/ops/dispatch_combine.py#L81-L84