diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 5f95bf20f..3ccdfd66f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -361,7 +361,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( symm_output = torch.empty( hidden_states.shape[0], hidden_states.shape[1], - dtype=torch.bfloat16, + dtype=hidden_states.dtype, device=hidden_states.device, ) @@ -441,9 +441,11 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( tune_max_num_tokens=next_power_of_2(a_q.shape[0]), fp8_quantization_type=int(fp8_quantization_type), ) + # TODO: Once https://github.com/flashinfer-ai/flashinfer/issues/2703 is fixed, pass output to moe kernel and remove this copy. symm_output.copy_(output) output = symm_output else: + assert TopKOutputChecker.format_is_bypassed(topk_output) assert quant_info.w13_input_scale is not None assert quant_info.output1_scales_scalar is not None assert quant_info.output1_scales_gate_scalar is not None diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 217b3fce7..9d0d15716 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1369,7 +1369,10 @@ class Fp8MoEMethod(FusedMoEMethodBase): self.process_weights_hip_scale_padding(layer) # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled - if get_moe_runner_backend().is_flashinfer_trtllm(): + if ( + get_moe_runner_backend().is_flashinfer_trtllm() + or get_moe_runner_backend().is_flashinfer_trtllm_routed() + ): from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( align_fp8_moe_weights_for_flashinfer_trtllm, ) @@ -1619,7 +1622,8 @@ class Fp8MoEMethod(FusedMoEMethodBase): local_num_experts=num_local_experts, intermediate_size=layer.w2_weight.shape[2], routing_method_type=int( - getattr(layer, "routing_method_type", RoutingMethodType.DeepSeekV3) + getattr(layer, "routing_method_type", None) + or RoutingMethodType.DeepSeekV3 ), block_quant=self.block_quant, use_mxfp8=getattr(self.quant_config, "use_mxfp8", False), diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 80b0e959b..297e4bbd7 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2046,6 +2046,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): if backend_str not in [ "flashinfer_trtllm", + # TODO: Enable for flashinfer_trtllm_routed once https://github.com/flashinfer-ai/flashinfer/issues/2749 is fixed. + # "flashinfer_trtllm_routed", "flashinfer_mxfp4", # TODO: flashinfer_cutlass will cause some flashinfer compilation errors. To be fixed. # "flashinfer_cutlass",