From 66ac385f520ca5abd8fff7790739f1972d559da4 Mon Sep 17 00:00:00 2001 From: Chengze Fan Date: Tue, 16 Jun 2026 21:02:06 -0700 Subject: [PATCH] fix(moe): MoRI EP init_mori_op missing BF16 dispatch branch (#28469) Signed-off-by: Chengze Fan --- python/sglang/srt/layers/moe/token_dispatcher/moriep.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index c43c95042..1e21ac701 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -257,7 +257,10 @@ def init_mori_op( data_type = fp8_dtype scale_type_size = torch.float32.itemsize - if dispatch_dtype == DispatchDtype.fp8: + if dispatch_dtype == DispatchDtype.bf16: + data_type = params_dtype + scale_dim = 0 + elif dispatch_dtype == DispatchDtype.fp8: scale_dim = hidden_size // FP8_BLOCK_SIZE elif dispatch_dtype == DispatchDtype.fp4: # FP4 kernel still takes the original hidden size and do quantization