diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index 312ccb1ff..a9ea71e33 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/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, +) diff --git a/python/sglang/srt/layers/moe/moe_runner/marlin.py b/python/sglang/srt/layers/moe/moe_runner/marlin.py index 45104dd27..ea35b875f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/marlin.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py index b5e3964c8..0e03c2057 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/gptq.py b/python/sglang/srt/layers/quantization/gptq.py index 74f64174c..ab9262c9f 100644 --- a/python/sglang/srt/layers/quantization/gptq.py +++ b/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 diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py index b8f3511b8..4fc5934d6 100644 --- a/test/srt/test_piecewise_cuda_graph.py +++ b/test/srt/test_piecewise_cuda_graph.py @@ -214,6 +214,41 @@ class TestPiecewiseCudaGraphAWQ(CustomTestCase): self.assertGreaterEqual(metrics["score"], 0.65) +class TestPiecewiseCudaGraphGPTQ(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen3-30B-A3B-GPTQ-Int4" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--enable-piecewise-cuda-graph"], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_mgsm_accuracy(self): + num_examples = 1319 + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mgsm_en", + num_examples=num_examples, + num_threads=min(num_examples, 1024), + ) + + metrics = run_eval(args) + print(f"MGSM Accuracy: {metrics['score']:.3f}") + + # Expected accuracy: 0.948, allow some variance + self.assertGreaterEqual(metrics["score"], 0.92) + + class TestPiecewiseCudaGraphFP8(CustomTestCase): """Test piecewise CUDA graph with FP8 quantized model"""