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