Add piecewise cuda graph for Qwen3-Next FP8 flashinfer_trtllm moe backend (#18184)

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)

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):
+            if 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:
                 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,

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 (
+    from flashinfer.fused_moe import Fp8QuantizationType
-        Fp8QuantizationType,
-        trtllm_fp8_block_scale_moe,
-        trtllm_fp8_block_scale_routed_moe,
-        trtllm_fp8_per_tensor_scale_moe,
-    )
 
     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:
+            symm_output = torch.empty(
-                assert (
+                hidden_states.shape[0],
-                    runner_config.top_k is not None
+                hidden_states.shape[1],
-                ), "runner_config.top_k is required for flashinfer_trtllm_routed."
+                dtype=torch.bfloat16,
-                assert TopKOutputChecker.format_is_standard(topk_output)
+                device=hidden_states.device,
-                packed_topk_ids = _pack_topk_for_flashinfer_routed(
+            )
-                    topk_ids=topk_output.topk_ids,
-                    topk_weights=topk_output.topk_weights,
-                )
 
-                output = trtllm_fp8_block_scale_routed_moe(
+        # Move kernel call outside context manager to avoid graph breaks
-                    topk_ids=packed_topk_ids,
+        # during torch.compile for piecewise cuda graph.
-                    routing_bias=None,
+        # Use custom op wrapper for torch.compile compatibility.
-                    hidden_states=a_q,
+        if use_routed_topk:
-                    hidden_states_scale=a_sf_t,
+            assert (
-                    gemm1_weights=quant_info.w13_weight,
+                runner_config.top_k is not None
-                    gemm1_weights_scale=quant_info.w13_weight_scale_inv,
+            ), "runner_config.top_k is required for flashinfer_trtllm_routed."
-                    gemm2_weights=quant_info.w2_weight,
+            assert TopKOutputChecker.format_is_standard(topk_output)
-                    gemm2_weights_scale=quant_info.w2_weight_scale_inv,
+            packed_topk_ids = _pack_topk_for_flashinfer_routed(
-                    num_experts=quant_info.global_num_experts,
+                topk_ids=topk_output.topk_ids,
-                    top_k=runner_config.top_k,
+                topk_weights=topk_output.topk_weights,
-                    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)
 
-                # FIXME: there is a bug in the trtllm_fp8_block_scale_moe.
+            output = torch.ops.sglang.trtllm_fp8_block_scale_routed_moe_wrapper(
-                # It ignored the `output` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325
+                topk_ids=packed_topk_ids,
-                # so we put the whole function under the ``use_symmetric_memory`` context manager.
+                routing_bias=None,
-                # If the bug is fixed, we can only put the output tensor allocation under the context manager.
+                hidden_states=a_q,
-                output = trtllm_fp8_block_scale_moe(
+                hidden_states_scale=a_sf_t,
-                    routing_logits=(
+                gemm1_weights=quant_info.w13_weight,
-                        router_logits.to(torch.float32)
+                gemm1_weights_scale=quant_info.w13_weight_scale_inv,
-                        if routing_method_type == RoutingMethodType.DeepSeekV3
+                gemm2_weights=quant_info.w2_weight,
-                        else router_logits
+                gemm2_weights_scale=quant_info.w2_weight_scale_inv,
-                    ),
+                num_experts=quant_info.global_num_experts,
-                    routing_bias=correction_bias,
+                top_k=runner_config.top_k,
-                    hidden_states=a_q,
+                n_group=None,
-                    hidden_states_scale=a_sf_t,
+                topk_group=None,
-                    gemm1_weights=quant_info.w13_weight,
+                intermediate_size=quant_info.intermediate_size,
-                    gemm1_weights_scale=quant_info.w13_weight_scale_inv,
+                local_expert_offset=quant_info.local_expert_offset,
-                    gemm2_weights=quant_info.w2_weight,
+                local_num_experts=quant_info.local_num_experts,
-                    gemm2_weights_scale=quant_info.w2_weight_scale_inv,
+                routed_scaling_factor=(
-                    num_experts=quant_info.global_num_experts,
+                    runner_config.routed_scaling_factor
-                    top_k=topk_config.top_k,
+                    if runner_config.routed_scaling_factor is not None
-                    n_group=(
+                    else 1.0
-                        topk_config.num_expert_group
+                ),
-                        if topk_config.num_expert_group
+                routing_method_type=(
-                        else 0
+                    RoutingMethodType.TopK
-                    ),
+                    if routing_method_type == RoutingMethodType.DeepSeekV3
-                    topk_group=topk_config.topk_group if topk_config.topk_group else 0,
+                    else routing_method_type
-                    intermediate_size=quant_info.intermediate_size,
+                ),
-                    local_expert_offset=quant_info.local_expert_offset,
+                use_shuffled_weight=use_shuffled_weight,
-                    local_num_experts=quant_info.local_num_experts,
+                tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
-                    routed_scaling_factor=(
+                fp8_quantization_type=int(fp8_quantization_type),
-                        runner_config.routed_scaling_factor
+            )
-                        if runner_config.routed_scaling_factor is not None
+        else:
-                        else 1.0
+            assert TopKOutputChecker.format_is_bypassed(topk_output)
-                    ),
+
-                    routing_method_type=routing_method_type,
+            output = torch.ops.sglang.trtllm_fp8_block_scale_moe_wrapper(
-                    use_shuffled_weight=use_shuffled_weight,
+                routing_logits=(
-                    weight_layout=0,
+                    router_logits.to(torch.float32)
-                    tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+                    if routing_method_type == RoutingMethodType.DeepSeekV3
-                    fp8_quantization_type=fp8_quantization_type,
+                    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(
+            symm_output = torch.empty(
-                routing_logits=router_logits.to(torch.bfloat16),
+                hidden_states.shape[0],
-                routing_bias=routing_bias_cast,
+                hidden_states.shape[1],
-                hidden_states=a_q,
+                dtype=torch.bfloat16,
-                gemm1_weights=quant_info.w13_weight,
+                device=hidden_states.device,
-                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]),
             )
 
+        # 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)