Harden FlashInfer FP4 imports in standard dispatcher (#21776)

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