diff --git a/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py b/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py new file mode 100644 index 000000000..c38126679 --- /dev/null +++ b/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py @@ -0,0 +1,274 @@ +from typing import Optional + +import torch + +from sglang.srt.utils.custom_op import register_custom_op + + +def _fake_fp8_block_scale_moe( + routing_logits: torch.Tensor, + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int = 0, + use_shuffled_weight: bool = False, + weight_layout: int = 0, + enable_pdl: Optional[bool] = None, + tune_max_num_tokens: int = 8192, + fp8_quantization_type: Optional[int] = None, +) -> torch.Tensor: + return torch.empty( + hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device + ) + + +@register_custom_op(fake_impl=_fake_fp8_block_scale_moe) +def trtllm_fp8_block_scale_moe_wrapper( + routing_logits: torch.Tensor, + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int = 0, + use_shuffled_weight: bool = False, + weight_layout: int = 0, + enable_pdl: Optional[bool] = None, + tune_max_num_tokens: int = 8192, + fp8_quantization_type: Optional[int] = None, +) -> torch.Tensor: + try: + from flashinfer.fused_moe import trtllm_fp8_block_scale_moe + except ImportError as e: + raise ImportError( + "Can't import trtllm_fp8_block_scale_moe from flashinfer. " + "Please check flashinfer version." + ) from e + kwargs = { + "routing_logits": routing_logits, + "routing_bias": routing_bias, + "hidden_states": hidden_states, + "hidden_states_scale": hidden_states_scale, + "gemm1_weights": gemm1_weights, + "gemm1_weights_scale": gemm1_weights_scale, + "gemm2_weights": gemm2_weights, + "gemm2_weights_scale": gemm2_weights_scale, + "num_experts": num_experts, + "top_k": top_k, + "n_group": n_group, + "topk_group": topk_group, + "intermediate_size": intermediate_size, + "local_expert_offset": local_expert_offset, + "local_num_experts": local_num_experts, + "routed_scaling_factor": routed_scaling_factor, + "routing_method_type": routing_method_type, + "use_shuffled_weight": use_shuffled_weight, + "weight_layout": weight_layout, + "enable_pdl": enable_pdl, + "tune_max_num_tokens": tune_max_num_tokens, + } + if fp8_quantization_type is not None: + from flashinfer.fused_moe import Fp8QuantizationType + + kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type) + + return trtllm_fp8_block_scale_moe(**kwargs) + + +def _fake_fp8_block_scale_routed_moe( + topk_ids: torch.Tensor, + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int = 0, + use_shuffled_weight: bool = False, + weight_layout: int = 0, + enable_pdl: Optional[bool] = None, + tune_max_num_tokens: int = 8192, + fp8_quantization_type: Optional[int] = None, +) -> torch.Tensor: + return torch.empty( + hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device + ) + + +@register_custom_op(fake_impl=_fake_fp8_block_scale_routed_moe) +def trtllm_fp8_block_scale_routed_moe_wrapper( + topk_ids: torch.Tensor, + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int = 0, + use_shuffled_weight: bool = False, + weight_layout: int = 0, + enable_pdl: Optional[bool] = None, + tune_max_num_tokens: int = 8192, + fp8_quantization_type: Optional[int] = None, +) -> torch.Tensor: + try: + from flashinfer.fused_moe import trtllm_fp8_block_scale_routed_moe + except ImportError as e: + raise ImportError( + "Can't import trtllm_fp8_block_scale_routed_moe from flashinfer. " + "Please check flashinfer version." + ) from e + kwargs = { + "topk_ids": topk_ids, + "routing_bias": routing_bias, + "hidden_states": hidden_states, + "hidden_states_scale": hidden_states_scale, + "gemm1_weights": gemm1_weights, + "gemm1_weights_scale": gemm1_weights_scale, + "gemm2_weights": gemm2_weights, + "gemm2_weights_scale": gemm2_weights_scale, + "num_experts": num_experts, + "top_k": top_k, + "n_group": n_group, + "topk_group": topk_group, + "intermediate_size": intermediate_size, + "local_expert_offset": local_expert_offset, + "local_num_experts": local_num_experts, + "routed_scaling_factor": routed_scaling_factor, + "routing_method_type": routing_method_type, + "use_shuffled_weight": use_shuffled_weight, + "weight_layout": weight_layout, + "enable_pdl": enable_pdl, + "tune_max_num_tokens": tune_max_num_tokens, + } + if fp8_quantization_type is not None: + from flashinfer.fused_moe import Fp8QuantizationType + + kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type) + + return trtllm_fp8_block_scale_routed_moe(**kwargs) + + +def _fake_fp8_per_tensor_scale_moe( + routing_logits: torch.Tensor, + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + gemm1_weights: torch.Tensor, + output1_scales_scalar: torch.Tensor, + output1_scales_gate_scalar: torch.Tensor, + gemm2_weights: torch.Tensor, + output2_scales_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + use_routing_scales_on_input: bool, + routing_method_type: int = 0, + enable_pdl: Optional[bool] = None, + tune_max_num_tokens: int = 8192, +) -> torch.Tensor: + return torch.empty( + hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device + ) + + +@register_custom_op(fake_impl=_fake_fp8_per_tensor_scale_moe) +def trtllm_fp8_per_tensor_scale_moe_wrapper( + routing_logits: torch.Tensor, + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + gemm1_weights: torch.Tensor, + output1_scales_scalar: torch.Tensor, + output1_scales_gate_scalar: torch.Tensor, + gemm2_weights: torch.Tensor, + output2_scales_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + use_routing_scales_on_input: bool, + routing_method_type: int = 0, + enable_pdl: Optional[bool] = None, + tune_max_num_tokens: int = 8192, +) -> torch.Tensor: + # lazy import + try: + from flashinfer.fused_moe import trtllm_fp8_per_tensor_scale_moe + except ImportError as e: + raise ImportError( + "Can't import trtllm_fp8_per_tensor_scale_moe from flashinfer. " + "Please check flashinfer version." + ) from e + + kwargs = { + "routing_logits": routing_logits, + "routing_bias": routing_bias, + "hidden_states": hidden_states, + "gemm1_weights": gemm1_weights, + "output1_scales_scalar": output1_scales_scalar, + "output1_scales_gate_scalar": output1_scales_gate_scalar, + "gemm2_weights": gemm2_weights, + "output2_scales_scalar": output2_scales_scalar, + "num_experts": num_experts, + "top_k": top_k, + "n_group": n_group, + "topk_group": topk_group, + "intermediate_size": intermediate_size, + "local_expert_offset": local_expert_offset, + "local_num_experts": local_num_experts, + "routed_scaling_factor": routed_scaling_factor, + "use_routing_scales_on_input": use_routing_scales_on_input, + "routing_method_type": routing_method_type, + "enable_pdl": enable_pdl, + "tune_max_num_tokens": tune_max_num_tokens, + } + + return trtllm_fp8_per_tensor_scale_moe(**kwargs) 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 a54e9ea7d..d60fb92c5 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -40,6 +40,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardDispatcher, + StandardDispatchOutput, ) from sglang.srt.layers.moe.topk import ( BypassedTopKOutput, @@ -967,10 +968,7 @@ class FusedMoE(torch.nn.Module): def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput): if is_in_piecewise_cuda_graph(): - if not TopKOutputChecker.format_is_standard(topk_output): - # Make sure there is torch lib op registration for the whole moe layer - return self.forward_impl(hidden_states, topk_output) - else: + if TopKOutputChecker.format_is_standard(topk_output): return moe_forward_piecewise_cuda_graph_impl( hidden_states, topk_output.topk_weights, @@ -978,6 +976,20 @@ class FusedMoE(torch.nn.Module): topk_output.router_logits, self.layer_id, ) + elif TopKOutputChecker.format_is_bypassed(topk_output): + return fused_moe_bypassed_piecewise_cuda_graph_impl( + hidden_states, + topk_output.router_logits, + topk_output.topk_config.top_k, + topk_output.topk_config.topk_group, + topk_output.topk_config.num_expert_group, + topk_output.topk_config.correction_bias, + topk_output.topk_config.renormalize, + self.layer_id, + ) + else: + # Make sure there is torch lib op registration for the whole moe layer + return self.forward_impl(hidden_states, topk_output) else: return self.forward_impl(hidden_states, topk_output) @@ -1134,6 +1146,116 @@ class FusedMoE(torch.nn.Module): self.meta_overlap_args = None +class FlashInferFusedMoE(FusedMoE): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput): + assert TopKOutputChecker.format_is_bypassed( + topk_output + ), "Only bypassed topk output is supported for flashinfer trtllm moe" + + if is_in_piecewise_cuda_graph(): + return flashinfer_bf16_moe_forward_piecewise_cuda_graph_impl( + hidden_states, + topk_output.router_logits, + topk_output.topk_config.top_k, + topk_output.topk_config.topk_group, + topk_output.topk_config.num_expert_group, + topk_output.topk_config.correction_bias, + topk_output.topk_config.renormalize, + self.layer_id, + ) + else: + return self.forward_impl(hidden_states, topk_output) + + def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput): + assert ( + self.moe_runner_config.activation == "silu" + ), "Only silu is supported for flashinfer trtllm moe" + assert self.quant_method is not None + assert ( + topk_output.topk_config.renormalize + ), "Renormalize is required for flashinfer trtllm moe" + assert ( + self.num_fused_shared_experts == 0 + ), "Fused shared experts are not supported for flashinfer trtllm moe" + assert ( + self.moe_runner_config.is_gated + ), "Only gated MoEs are supported for flashinfer trtllm moe" + + router_logits = topk_output.router_logits + topk_config = topk_output.topk_config + correction_bias = topk_config.correction_bias + routed_scaling_factor = self.moe_runner_config.routed_scaling_factor + + if isinstance(self.quant_method, UnquantizedFusedMoEMethod): + # lazy import + try: + from flashinfer.fused_moe import trtllm_bf16_moe + except ImportError as e: + raise ImportError( + "Can't import trtllm_bf16_moe from flashinfer. " + "Please check flashinfer version to use bf16 with flashinfer_trtllm backend." + ) from e + + # Allocate output inside symmetric memory context + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + # TODO: Now trtllm_bf16_moe doesn't support inplace output, + # we can move this out when it support that. + symm_output = torch.empty( + hidden_states.shape[0], + hidden_states.shape[1], + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # Move kernel call outside context manager to avoid graph breaks + # during torch.compile for piecewise cuda graph + moe_result = trtllm_bf16_moe( + routing_logits=router_logits, + routing_bias=correction_bias, + hidden_states=hidden_states, + gemm1_weights=self.w13_weight, + gemm2_weights=self.w2_weight, + num_experts=self.num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=self.moe_ep_rank * self.num_local_experts, + local_num_experts=self.num_local_experts, + routing_method_type=self.routing_method_type, + tune_max_num_tokens=next_power_of_2(hidden_states.shape[0]), + ) + # Copy result to symmetric memory output + symm_output.copy_(moe_result) + final_hidden_states = symm_output + + else: + + final_hidden_states = self.quant_method.apply( + layer=self, + dispatch_output=StandardDispatchOutput( + hidden_states=hidden_states, + hidden_states_scale=None, + topk_output=topk_output, + ), + ).hidden_states + + # NOTE for symmetric memory tagging: + # We do not create the context in this function. + # Instead, we create the context and tagging inside each FusedMoEMethodBase + # This can allow fine-grained tagging. + + if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states + + class FlashInferFP4MoE(FusedMoE): """FP4 TRTLLM MoE implementation using FlashInfer.""" @@ -1300,6 +1422,60 @@ def moe_forward_piecewise_cuda_graph_impl( return moe_layer.forward_impl(hidden_states, topk_output) +@register_custom_op(out_shape="hidden_states") +def fused_moe_bypassed_piecewise_cuda_graph_impl( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + top_k: int, + topk_group: Optional[int], + num_expert_group: Optional[int], + correction_bias: Optional[torch.Tensor], + renormalize: bool, + layer_id: int, +) -> torch.Tensor: + topk_output = BypassedTopKOutput( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=TopKConfig( + top_k=top_k, + topk_group=topk_group, + num_expert_group=num_expert_group, + correction_bias=correction_bias, + renormalize=renormalize, + ), + ) + forward_context = get_forward_context() + moe_layer = forward_context.moe_layers[layer_id] + return moe_layer.forward_impl(hidden_states, topk_output) + + +@register_custom_op(out_shape="hidden_states") +def flashinfer_bf16_moe_forward_piecewise_cuda_graph_impl( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + top_k: int, + topk_group: Optional[int], + num_expert_group: Optional[int], + correction_bias: Optional[torch.Tensor], + renormalize: bool, + layer_id: int, +) -> torch.Tensor: + topk_output = BypassedTopKOutput( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=TopKConfig( + top_k=top_k, + topk_group=topk_group, + num_expert_group=num_expert_group, + correction_bias=correction_bias, + renormalize=renormalize, + ), + ) + forward_context = get_forward_context() + moe_layer = forward_context.moe_layers[layer_id] + return moe_layer.forward_impl(hidden_states, topk_output) + + @register_custom_op(out_shape="hidden_states") def flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl( hidden_states: torch.Tensor, 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 9c3eac87c..1a3ef033b 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -7,6 +7,8 @@ import torch from torch.nn import Module from torch.nn.parameter import Parameter +# Import to register custom ops for torch.compile compatibility +import sglang.srt.layers.moe.flashinfer_trtllm_moe # noqa: F401 from sglang.srt.distributed import get_tp_group from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -298,12 +300,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( runner_config: MoeRunnerConfig, use_routed_topk: bool = False, ) -> StandardCombineInput: - from flashinfer.fused_moe import ( - Fp8QuantizationType, - trtllm_fp8_block_scale_moe, - trtllm_fp8_block_scale_routed_moe, - trtllm_fp8_per_tensor_scale_moe, - ) + from flashinfer.fused_moe import Fp8QuantizationType from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker @@ -354,92 +351,95 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( ) a_sf_t = a_sf.t().contiguous() + # Allocate output inside symmetric memory context with use_symmetric_memory( get_tp_group(), disabled=not is_allocation_symmetric() ): - if use_routed_topk: - assert ( - runner_config.top_k is not None - ), "runner_config.top_k is required for flashinfer_trtllm_routed." - assert TopKOutputChecker.format_is_standard(topk_output) - packed_topk_ids = _pack_topk_for_flashinfer_routed( - topk_ids=topk_output.topk_ids, - topk_weights=topk_output.topk_weights, - ) + symm_output = torch.empty( + hidden_states.shape[0], + hidden_states.shape[1], + dtype=torch.bfloat16, + device=hidden_states.device, + ) - output = trtllm_fp8_block_scale_routed_moe( - topk_ids=packed_topk_ids, - routing_bias=None, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=quant_info.w13_weight, - gemm1_weights_scale=quant_info.w13_weight_scale_inv, - gemm2_weights=quant_info.w2_weight, - gemm2_weights_scale=quant_info.w2_weight_scale_inv, - num_experts=quant_info.global_num_experts, - top_k=runner_config.top_k, - n_group=None, - topk_group=None, - intermediate_size=quant_info.intermediate_size, - local_expert_offset=quant_info.local_expert_offset, - local_num_experts=quant_info.local_num_experts, - routed_scaling_factor=( - runner_config.routed_scaling_factor - if runner_config.routed_scaling_factor is not None - else 1.0 - ), - routing_method_type=( - RoutingMethodType.TopK - if routing_method_type == RoutingMethodType.DeepSeekV3 - else routing_method_type - ), - use_shuffled_weight=use_shuffled_weight, - weight_layout=0, - tune_max_num_tokens=next_power_of_2(a_q.shape[0]), - fp8_quantization_type=fp8_quantization_type, - ) - else: - assert TopKOutputChecker.format_is_bypassed(topk_output) + # Move kernel call outside context manager to avoid graph breaks + # during torch.compile for piecewise cuda graph. + # Use custom op wrapper for torch.compile compatibility. + if use_routed_topk: + assert ( + runner_config.top_k is not None + ), "runner_config.top_k is required for flashinfer_trtllm_routed." + assert TopKOutputChecker.format_is_standard(topk_output) + packed_topk_ids = _pack_topk_for_flashinfer_routed( + topk_ids=topk_output.topk_ids, + topk_weights=topk_output.topk_weights, + ) - # FIXME: there is a bug in the trtllm_fp8_block_scale_moe. - # It ignored the `output` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325 - # so we put the whole function under the ``use_symmetric_memory`` context manager. - # If the bug is fixed, we can only put the output tensor allocation under the context manager. - output = trtllm_fp8_block_scale_moe( - routing_logits=( - router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits - ), - routing_bias=correction_bias, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=quant_info.w13_weight, - gemm1_weights_scale=quant_info.w13_weight_scale_inv, - gemm2_weights=quant_info.w2_weight, - gemm2_weights_scale=quant_info.w2_weight_scale_inv, - num_experts=quant_info.global_num_experts, - top_k=topk_config.top_k, - n_group=( - topk_config.num_expert_group - if topk_config.num_expert_group - else 0 - ), - topk_group=topk_config.topk_group if topk_config.topk_group else 0, - intermediate_size=quant_info.intermediate_size, - local_expert_offset=quant_info.local_expert_offset, - local_num_experts=quant_info.local_num_experts, - routed_scaling_factor=( - runner_config.routed_scaling_factor - if runner_config.routed_scaling_factor is not None - else 1.0 - ), - routing_method_type=routing_method_type, - use_shuffled_weight=use_shuffled_weight, - weight_layout=0, - tune_max_num_tokens=next_power_of_2(a_q.shape[0]), - fp8_quantization_type=fp8_quantization_type, - ) + output = torch.ops.sglang.trtllm_fp8_block_scale_routed_moe_wrapper( + topk_ids=packed_topk_ids, + routing_bias=None, + hidden_states=a_q, + hidden_states_scale=a_sf_t, + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale_inv, + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale_inv, + num_experts=quant_info.global_num_experts, + top_k=runner_config.top_k, + n_group=None, + topk_group=None, + intermediate_size=quant_info.intermediate_size, + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=( + runner_config.routed_scaling_factor + if runner_config.routed_scaling_factor is not None + else 1.0 + ), + routing_method_type=( + RoutingMethodType.TopK + if routing_method_type == RoutingMethodType.DeepSeekV3 + else routing_method_type + ), + use_shuffled_weight=use_shuffled_weight, + tune_max_num_tokens=next_power_of_2(a_q.shape[0]), + fp8_quantization_type=int(fp8_quantization_type), + ) + else: + assert TopKOutputChecker.format_is_bypassed(topk_output) + + output = torch.ops.sglang.trtllm_fp8_block_scale_moe_wrapper( + routing_logits=( + router_logits.to(torch.float32) + if routing_method_type == RoutingMethodType.DeepSeekV3 + else router_logits + ), + routing_bias=correction_bias, + hidden_states=a_q, + hidden_states_scale=a_sf_t, + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale_inv, + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale_inv, + num_experts=quant_info.global_num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=quant_info.intermediate_size, + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=( + runner_config.routed_scaling_factor + if runner_config.routed_scaling_factor is not None + else 1.0 + ), + routing_method_type=routing_method_type, + use_shuffled_weight=use_shuffled_weight, + tune_max_num_tokens=next_power_of_2(a_q.shape[0]), + fp8_quantization_type=int(fp8_quantization_type), + ) + symm_output.copy_(output) + output = symm_output else: assert quant_info.w13_input_scale is not None assert quant_info.output1_scales_scalar is not None @@ -451,37 +451,48 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( None if correction_bias is None else correction_bias.to(torch.bfloat16) ) + # Allocate output inside symmetric memory context with use_symmetric_memory( get_tp_group(), disabled=not is_allocation_symmetric() ): - output = trtllm_fp8_per_tensor_scale_moe( - routing_logits=router_logits.to(torch.bfloat16), - routing_bias=routing_bias_cast, - hidden_states=a_q, - gemm1_weights=quant_info.w13_weight, - output1_scales_scalar=quant_info.output1_scales_scalar, - output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar, - gemm2_weights=quant_info.w2_weight, - output2_scales_scalar=quant_info.output2_scales_scalar, - num_experts=quant_info.global_num_experts, - top_k=topk_config.top_k, - n_group=( - topk_config.num_expert_group if topk_config.num_expert_group else 0 - ), - topk_group=topk_config.topk_group if topk_config.topk_group else 0, - intermediate_size=quant_info.intermediate_size, - local_expert_offset=quant_info.local_expert_offset, - local_num_experts=quant_info.local_num_experts, - routed_scaling_factor=( - runner_config.routed_scaling_factor - if runner_config.routed_scaling_factor is not None - else 1.0 - ), - use_routing_scales_on_input=quant_info.use_routing_scales_on_input, - routing_method_type=routing_method_type, - tune_max_num_tokens=next_power_of_2(a_q.shape[0]), + symm_output = torch.empty( + hidden_states.shape[0], + hidden_states.shape[1], + dtype=torch.bfloat16, + device=hidden_states.device, ) + # Move kernel call outside context manager to avoid graph breaks + # during torch.compile for piecewise cuda graph. + # Use custom op wrapper for torch.compile compatibility. + output = torch.ops.sglang.trtllm_fp8_per_tensor_scale_moe( + routing_logits=router_logits.to(torch.bfloat16), + routing_bias=routing_bias_cast, + hidden_states=a_q, + gemm1_weights=quant_info.w13_weight, + output1_scales_scalar=quant_info.output1_scales_scalar, + output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar, + gemm2_weights=quant_info.w2_weight, + output2_scales_scalar=quant_info.output2_scales_scalar, + num_experts=quant_info.global_num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=int(quant_info.w2_weight.shape[2]), + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=( + runner_config.routed_scaling_factor + if runner_config.routed_scaling_factor is not None + else 1.0 + ), + use_routing_scales_on_input=False, + routing_method_type=routing_method_type, + tune_max_num_tokens=next_power_of_2(a_q.shape[0]), + ) + symm_output.copy_(output) + output = symm_output + return StandardCombineInput(hidden_states=output)