Skip to content

Commit e1d2d6c

Browse files
[AMD] Fix MORI EP crash when mori package lacks AsyncLL kernel type
PR sgl-project#17953 added EpDispatchCombineKernelType.AsyncLL but the CI Docker image's mori package doesn't have it yet, causing an AttributeError on every mori EP test. Guard the LOW_LATENCY config behind a hasattr check and fall back to INTRA_NODE/INTER_NODE mode when AsyncLL is unavailable.
1 parent aaa07b8 commit e1d2d6c

File tree

1 file changed

+20
-6
lines changed
  • python/sglang/srt/layers/moe/token_dispatcher

1 file changed

+20
-6
lines changed

python/sglang/srt/layers/moe/token_dispatcher/moriep.py

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ def get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank: int = 4096):
146146
else mori.ops.EpDispatchCombineKernelType.InterNodeV1
147147
)
148148

149-
return {
149+
configs = {
150150
# TODO(billishyahao): need to tune different configs for intra node async
151151
# Also could be tuned for different AMD platform
152152
EpMode.INTRA_NODE: EpDispatchConfig(
@@ -161,13 +161,17 @@ def get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank: int = 4096):
161161
block_num=64,
162162
rdma_block_num=32,
163163
),
164-
EpMode.LOW_LATENCY: EpDispatchConfig(
164+
}
165+
166+
if hasattr(mori.ops.EpDispatchCombineKernelType, "AsyncLL"):
167+
configs[EpMode.LOW_LATENCY] = EpDispatchConfig(
165168
kernel_type=mori.ops.EpDispatchCombineKernelType.AsyncLL,
166169
warp_num_per_block=8,
167170
block_num=64,
168171
rdma_block_num=32,
169-
),
170-
}
172+
)
173+
174+
return configs
171175

172176

173177
# init_mori_op only needs do once in model initial stage
@@ -195,14 +199,24 @@ def init_mori_op(
195199

196200
mode = EpMode.INTRA_NODE if world_size <= 8 else EpMode.INTER_NODE
197201
async_mode = deepep_mode.enable_low_latency()
202+
203+
dispatch_configs = get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank)
204+
198205
if async_mode:
199-
mode = EpMode.LOW_LATENCY
206+
if EpMode.LOW_LATENCY in dispatch_configs:
207+
mode = EpMode.LOW_LATENCY
208+
else:
209+
logger.warning(
210+
"MORI LOW_LATENCY mode requested but not supported by the installed "
211+
"mori package (missing AsyncLL kernel). Falling back to %s mode.",
212+
mode.value,
213+
)
200214

201215
logger.info(
202216
f"[MORI init] {world_size=} {rank=} {hidden_size=} {params_dtype=} {num_max_dispatch_tokens_per_rank=} {num_local_experts=} {router_topk=} {mode=}"
203217
)
204218

205-
cfg = get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank)[mode]
219+
cfg = dispatch_configs[mode]
206220

207221
kernel_type = cfg.kernel_type
208222
warp_num_per_block = cfg.warp_num_per_block

0 commit comments

Comments
 (0)