Support piecewise cuda graph for fused marlin moe (#15100)

python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.utils import is_cuda
+from sglang.srt.utils import direct_register_custom_op, is_cuda
 
 _is_cuda = is_cuda()
 
@@ -41,7 +41,7 @@ def fused_marlin_moe(
     num_bits: int = 8,
     is_k_full: bool = True,
     inplace: bool = False,
-    routed_scaling_factor: float = None,
+    routed_scaling_factor: Optional[float] = None,
 ) -> torch.Tensor:
     """
     This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -225,15 +225,26 @@ def fused_marlin_moe_fake(
     gating_output: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
     g_idx1: Optional[torch.Tensor] = None,
     g_idx2: Optional[torch.Tensor] = None,
     sort_indices1: Optional[torch.Tensor] = None,
     sort_indices2: Optional[torch.Tensor] = None,
     w1_zeros: Optional[torch.Tensor] = None,
     w2_zeros: Optional[torch.Tensor] = None,
+    workspace: Optional[torch.Tensor] = None,
     num_bits: int = 8,
     is_k_full: bool = True,
     inplace: bool = False,
-    routed_scaling_factor: float = None,
+    routed_scaling_factor: Optional[float] = None,
 ) -> torch.Tensor:
     return torch.empty_like(hidden_states)
+
+
+direct_register_custom_op(
+    op_name="fused_marlin_moe",
+    op_func=fused_marlin_moe,
+    mutates_args=[],
+    fake_impl=fused_marlin_moe_fake,
+)

python/sglang/srt/layers/moe/moe_runner/marlin.py

@@ -80,7 +80,9 @@ def fused_experts_none_to_marlin(
     runner_config: MoeRunnerConfig,
 ) -> StandardCombineInput:
     global MARLIN_MOE_WORKSPACE
-    from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import fused_marlin_moe
+    from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (  # noqa
+        fused_marlin_moe,
+    )
     from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
     from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
 
@@ -97,7 +99,7 @@ def fused_experts_none_to_marlin(
             hidden_states.device, max_blocks_per_sm=4
         )
 
-    output = fused_marlin_moe(
+    output = torch.ops.sglang.fused_marlin_moe(
         hidden_states=hidden_states,
         w1=quant_info.w13_qweight,
         w2=quant_info.w2_qweight,

python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -943,7 +943,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
         layer: torch.nn.Module,
         dispatch_output: StandardDispatchOutput,
     ) -> CombineInput:
-        from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (
+        from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (  # noqa
             fused_marlin_moe,
         )
         from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
@@ -967,7 +967,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
             if expert_map is not None:
                 global_num_experts = self.moe_runner_config.num_experts
 
-        output = fused_marlin_moe(
+        output = torch.ops.sglang.fused_marlin_moe(
             x,
             layer.w13_weight_packed,
             layer.w2_weight_packed,

python/sglang/srt/layers/quantization/gptq.py

@@ -1099,32 +1099,3 @@ if _is_cuda:
     @register_fake_if_exists("sgl_kernel::gptq_shuffle")
     def _(q_weight, q_perm, bit):
         return
-
-    @register_fake_if_exists("sgl_kernel::moe_wna16_marlin_gemm")
-    def _(
-        a,
-        c,
-        b_q_weight,
-        b_scales,
-        b_zeros,
-        g_idx,
-        perm,
-        workspace,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        topk_weights,
-        moe_block_size,
-        top_k,
-        mul_topk_weights,
-        is_ep,
-        b_q_type_id,
-        size_m,
-        size_n,
-        size_k,
-        is_k_full,
-        use_atomic_add,
-        use_fp32_reduce,
-        is_zp_float,
-    ):
-        return c