diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 9e22d8092..65d0fbc44 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -34,7 +34,6 @@ from sglang.srt.utils.common import ( get_bool_env_var, get_device, is_hip, - is_sm120_supported, ) _is_hip = is_hip() @@ -45,14 +44,13 @@ if TYPE_CHECKING: try: - if is_sm120_supported(): - from flashinfer import fp4_quantize - else: - from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize - from flashinfer import fp4_quantize as fp4_quantize_flashinfer + from flashinfer import ( + nvfp4_block_scale_interleave as nvfp4_block_scale_interleave_flashinfer, + ) except ImportError: - fp4_quantize = None + fp4_quantize_flashinfer = None + nvfp4_block_scale_interleave_flashinfer = None class StandardDispatchOutput(NamedTuple): @@ -114,8 +112,15 @@ class StandardDispatcher(BaseDispatcher): if should_use_flashinfer_cutlass_moe_fp4_allgather(): # all-gather fp4 hidden states - from flashinfer import nvfp4_block_scale_interleave - + if ( + fp4_quantize_flashinfer is None + or nvfp4_block_scale_interleave_flashinfer is None + ): + raise RuntimeError( + "FlashInfer fp4_quantize and nvfp4_block_scale_interleave " + "are required for the flashinfer_cutlass FP4 all-gather " + "path." + ) global_scale = self.quant_config.get("input_global_scale", None) assert global_scale is not None, "input_global_scale is not set" topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids @@ -140,7 +145,7 @@ class StandardDispatcher(BaseDispatcher): [topk_weights, topk_ids, x, x_sf], sizes=get_dp_global_num_tokens() ) # TODO: fuse into cutlass moe - x_sf = nvfp4_block_scale_interleave(x_sf) + x_sf = nvfp4_block_scale_interleave_flashinfer(x_sf) hidden_states = x hidden_states_scale = x_sf