diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 80735f3da..43ad055a2 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -76,21 +76,11 @@ from sglang.srt.utils.custom_op import register_custom_op if is_flashinfer_available(): from flashinfer import fp4_quantize -# Try to import FP4 TRTLLM function if flashinfer is available -trtllm_fp4_block_scale_moe = None -if get_moe_runner_backend().is_flashinfer_trtllm(): - try: - from flashinfer.fused_moe import trtllm_fp4_block_scale_moe - except ImportError: - trtllm_fp4_block_scale_moe = None - _is_hip = is_hip() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip -logger = logging.getLogger(__name__) - def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: a2a_backend = get_moe_a2a_backend() @@ -1329,6 +1319,8 @@ class FlashInferFP4MoE(FusedMoE): hidden_states: Input tensor topk_output: TopKOutput object with Bypassed format """ + from flashinfer.fused_moe import trtllm_fp4_block_scale_moe + assert isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod) assert (