From 24a25ffa20649b97f66ce91673cd501f3c08b161 Mon Sep 17 00:00:00 2001
From: b8zhong <b8zhong@uwaterloo.ca>
Date: Sat, 15 Nov 2025 19:03:19 -0800
Subject: [PATCH] [Piecewise CUDA Graph] Support ModelOpt FP4 (#13101)

---
 .../srt/layers/quantization/modelopt_quant.py | 48 ++++++++++++++++++-
 test/srt/test_piecewise_cuda_graph.py         | 41 +++++++++++++++-
 2 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py
index f3f7b85ea..62e4677db 100755
--- a/python/sglang/srt/layers/quantization/modelopt_quant.py
+++ b/python/sglang/srt/layers/quantization/modelopt_quant.py
@@ -92,6 +92,50 @@ except ImportError:
 # Initialize logger for the module
 logger = logging.getLogger(__name__)
 
+
+@torch.library.custom_op("sglang::fp4_gemm", mutates_args=())
+def _sglang_fp4_gemm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    alpha: torch.Tensor,
+    out_dtype: torch.dtype,
+    out_features: int,
+) -> torch.Tensor:
+    backend = FLASHINFER_FP4_GEMM_BACKEND if FLASHINFER_FP4_GEMM_BACKEND else "cutlass"
+    if enable_flashinfer_fp4_gemm:
+        return fp4_gemm(
+            input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend
+        )
+    else:
+        return fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype)
+
+
+@torch.library.register_fake("sglang::fp4_gemm")
+def _sglang_fp4_gemm_fake(
+    input,
+    weight,
+    input_sf,
+    weight_sf,
+    alpha,
+    out_dtype,
+    out_features: int,
+):
+    M = input.shape[-2]
+    N = int(out_features)
+    return input.new_empty((M, N), dtype=out_dtype)
+
+
+if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None):
+
+    @torch.library.register_fake("sgl_kernel::scaled_fp4_quant")
+    def _sgl_kernel_scaled_fp4_quant_fake(
+        output, input, output_scale, input_global_scale
+    ):
+        return
+
+
 CUTEDSL_MOE_SCALAR_INPUT_SCALE = get_bool_env_var(
     "SGLANG_CUTEDSL_MOE_SCALAR_INPUT_SCALE", "true"
 )
@@ -1078,14 +1122,14 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         backend = (
             FLASHINFER_FP4_GEMM_BACKEND if FLASHINFER_FP4_GEMM_BACKEND else "cutlass"
         )
-        out = fp4_gemm(
+        out = _sglang_fp4_gemm(
             x_fp4,
             w,
             x_scale_interleaved,
             w_scale_interleaved,
             layer.alpha,
             output_dtype,
-            **(dict(backend=backend)),
+            w_n,
         )
         if bias is not None:
             out = out + bias
diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py
index 173e8b56f..19723e5e6 100644
--- a/test/srt/test_piecewise_cuda_graph.py
+++ b/test/srt/test_piecewise_cuda_graph.py
@@ -1,6 +1,6 @@
 import unittest
 
-from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils import get_device_sm, kill_process_tree
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -54,6 +54,45 @@ class TestPiecewiseCudaGraphBenchmark(CustomTestCase):
         self.assertLess(prefill_latency, 0.015)
 
 
+@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
+class TestPiecewiseCudaGraphLlama31FP4(CustomTestCase):
+    """MGSM test: piecewise CUDA graph with NVFP4 Llama3.1 8B on Blackwell."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "nvidia/Llama-3.1-8B-Instruct-FP4"
+        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",
+                "--quantization",
+                "modelopt_fp4",
+                "--mem-fraction-static",
+                "0.8",
+            ],
+        )
+
+    @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}")
+        self.assertGreaterEqual(metrics["score"], 0.78)
+
+
 class TestPiecewiseCudaGraphQwen3MoE(CustomTestCase):
     """Test piecewise CUDA graph with Qwen3-Coder-30B-A3B-Instruct MoE model"""