[rebase]Deepseek_v4 support w4(mxfp4)a16 on hopper (#24986)

python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh

@@ -1006,6 +1006,16 @@ void moe_wna16_marlin_gemm(
         "points.");
   }
 
+  if (b_q_type == kFE2M1f) {
+    RuntimeCheck(
+        group_size == 16 || group_size == 32,
+        "float4_e2m1f only supports group_size == 16 (NVFP4) or group_size == 32 (MXFP4). Got group_size = ",
+        group_size);
+    RuntimeCheck(
+        group_size != 32 || std::is_same<scalar_t, nv_bfloat16>::value,
+        "MXFP4 Marlin with E8M0 scales is only instantiated for bfloat16 activations.");
+  }
+
   // Verify b_zeros
   if (has_zp) {
     RuntimeCheck(b_zeros.dim() == 3, "b_zeros rank = ", b_zeros.dim(), " is not 3");

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

@@ -119,13 +119,9 @@ def fused_marlin_moe(
         and w2_scale.dtype == torch.float8_e8m0fnu
     )
     if is_mxfp4_marlin:
-        assert w1_scale.dtype == torch.float8_e8m0fnu, (
+        assert hidden_states.dtype == torch.bfloat16, (
-            "MXFP4 Marlin expects w1_scale to be torch.float8_e8m0fnu, "
+            "MXFP4 Marlin with E8M0 scales is only instantiated for bfloat16 "
-            f"got {w1_scale.dtype}"
+            f"activations, got {hidden_states.dtype}"
-        )
-        assert w2_scale.dtype == torch.float8_e8m0fnu, (
-            "MXFP4 Marlin expects w2_scale to be torch.float8_e8m0fnu, "
-            f"got {w2_scale.dtype}"
         )
     else:
         assert (

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

@@ -52,14 +52,30 @@ def _normalize_scale_tensor(
     raise TypeError(f"Unsupported MXFP4 scale dtype for Marlin: {scales.dtype}")
 
 
+def _get_optional_param(layer: torch.nn.Module, *names: str) -> torch.Tensor | None:
+    for name in names:
+        value = getattr(layer, name, None)
+        if value is not None:
+            return value
+    return None
+
+
 def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
     group_size = 32
     w13 = layer.w13_weight.data
     w2 = layer.w2_weight.data
-    w13_scale = layer.w13_weight_scale_inv.data
+    w13_scale = _get_optional_param(layer, "w13_weight_scale", "w13_weight_scale_inv")
-    w2_scale = layer.w2_weight_scale_inv.data
+    w2_scale = _get_optional_param(layer, "w2_weight_scale", "w2_weight_scale_inv")
-    w13_bias = getattr(layer, "w13_bias", None)
+    w13_bias = _get_optional_param(layer, "w13_weight_bias", "w13_bias")
-    w2_bias = getattr(layer, "w2_bias", None)
+    w2_bias = _get_optional_param(layer, "w2_weight_bias", "w2_bias")
+
+    if w13_scale is None or w2_scale is None:
+        raise ValueError("MXFP4 Marlin requires w13/w2 weight scales.")
+
+    w13_scale_data = w13_scale.data if hasattr(w13_scale, "data") else w13_scale
+    w2_scale_data = w2_scale.data if hasattr(w2_scale, "data") else w2_scale
+    w13_bias_data = w13_bias.data if hasattr(w13_bias, "data") else w13_bias
+    w2_bias_data = w2_bias.data if hasattr(w2_bias, "data") else w2_bias
 
     num_experts = w13.shape[0]
     intermediate_size = w13.shape[1] // 2
@@ -67,7 +83,7 @@ def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
     param_dtype = getattr(
         layer,
         "orig_dtype",
-        w13_bias.dtype if w13_bias is not None else torch.bfloat16,
+        w13_bias_data.dtype if w13_bias_data is not None else torch.bfloat16,
     )
 
     device = w13.device
@@ -129,19 +145,19 @@ def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
 
     w13_marlin = _repack_weight(w13, True)
     w2_marlin = _repack_weight(w2, False)
-    w13_scale_marlin = _permute_scales(w13_scale, True)
+    w13_scale_marlin = _permute_scales(w13_scale_data, True)
-    w2_scale_marlin = _permute_scales(w2_scale, False)
+    w2_scale_marlin = _permute_scales(w2_scale_data, False)
 
     layer.w13_weight = torch.nn.Parameter(w13_marlin, requires_grad=False)
     layer.w2_weight = torch.nn.Parameter(w2_marlin, requires_grad=False)
-    layer.w13_weight_scale_inv = torch.nn.Parameter(
+    layer.w13_weight_scale = torch.nn.Parameter(w13_scale_marlin, requires_grad=False)
-        w13_scale_marlin, requires_grad=False
+    layer.w2_weight_scale = torch.nn.Parameter(w2_scale_marlin, requires_grad=False)
-    )
-    layer.w2_weight_scale_inv = torch.nn.Parameter(w2_scale_marlin, requires_grad=False)
 
-    if w13_bias is not None:
+    if w13_bias_data is not None:
-        layer.w13_bias = torch.nn.Parameter(
+        layer.w13_weight_bias = torch.nn.Parameter(
-            _permute_bias(w13_bias), requires_grad=False
+            _permute_bias(w13_bias_data), requires_grad=False
+        )
+    if w2_bias_data is not None:
+        layer.w2_weight_bias = torch.nn.Parameter(
+            _permute_bias(w2_bias_data), requires_grad=False
         )
-    if w2_bias is not None:
-        layer.w2_bias = torch.nn.Parameter(_permute_bias(w2_bias), requires_grad=False)

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

@@ -35,6 +35,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 )
 from sglang.srt.layers.dp_attention import is_allocation_symmetric
 from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
+from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
 from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend
 from sglang.srt.layers.quantization.base_config import (
@@ -342,6 +343,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         self.use_triton_kernels = get_moe_runner_backend().is_triton_kernels()
         self.with_bias = False
         self.use_flashinfer = get_moe_runner_backend().is_flashinfer_mxfp4()
+        self.use_marlin = get_moe_runner_backend().is_marlin()
         self.flashinfer_mxfp4_moe_precision = (
             get_global_server_args().flashinfer_mxfp4_moe_precision
         )
@@ -507,6 +509,25 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         set_weight_attrs(w2_weight_bias, extra_weight_attrs)
 
     def process_weights_after_loading(self, layer):
+        if self.use_marlin:
+            from sglang.srt.layers.quantization.marlin_utils import (
+                check_moe_marlin_supports_layer,
+            )
+            from sglang.srt.layers.quantization.marlin_utils_fp4 import (
+                prepare_moe_mxfp4_layer_for_marlin,
+            )
+
+            if not is_sm90_supported():
+                raise RuntimeError("MXFP4 Marlin requires Hopper/SM90 or above.")
+            if not check_moe_marlin_supports_layer(layer, 32):
+                raise RuntimeError(
+                    "Current MXFP4 MoE layer is not supported by Marlin."
+                )
+
+            prepare_moe_mxfp4_layer_for_marlin(layer)
+            layer._mxfp4_backend = "marlin"
+            return
+
         if self._fi_kernel == "cutlass_sm90":
             self._process_weights_for_sm90_cutlass(layer)
             return
@@ -954,7 +975,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             self.runner = MoeRunner(
                 moe_runner_backend, replace(moe_runner_config, activation="swiglu")
             )
-        elif moe_runner_backend.is_triton_kernels() or moe_runner_backend.is_triton():
+        elif (
+            moe_runner_backend.is_triton_kernels()
+            or moe_runner_backend.is_triton()
+            or moe_runner_backend.is_marlin()
+        ):
             self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
         else:
             # TODO(cwan): refactor other backends
@@ -1040,6 +1065,20 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         x = dispatch_output.hidden_states
         topk_output = dispatch_output.topk_output
 
+        if self.use_marlin:
+            assert TopKOutputChecker.format_is_standard(topk_output)
+            quant_info = MarlinMoeQuantInfo(
+                w13_qweight=layer.w13_weight,
+                w2_qweight=layer.w2_weight,
+                w13_scales=layer.w13_weight_scale,
+                w2_scales=layer.w2_weight_scale,
+                w13_g_idx_sort_indices=None,
+                w2_g_idx_sort_indices=None,
+                weight_bits=4,
+                is_k_full=True,
+            )
+            return self.runner.run(dispatch_output, quant_info)
+
         if self._fi_kernel == "cutlass_sm90":
             return self._apply_sm90_cutlass(layer, x, topk_output)
         if self.use_flashinfer:

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

@@ -8,7 +8,7 @@ from torch.nn import Module
 
 from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
 from sglang.srt.layers.moe.utils import MoeRunnerBackend
-from sglang.srt.utils import log_info_on_rank0
+from sglang.srt.utils import log_info_on_rank0, set_weight_attrs
 from sglang.srt.utils.common import is_sm90_supported
 
 if TYPE_CHECKING:
@@ -38,17 +38,62 @@ class Mxfp4MarlinMoEMethod:
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        # Delegate to the underlying FP8 method for weight creation —
+        from sglang.srt.layers.moe.fused_moe_triton import (
-        # the raw weight shapes are the same; only post-loading processing differs.
+            FusedMoeWeightScaleSupported,
-        self._fp8.create_weights(
-            layer,
-            num_experts,
-            hidden_size,
-            intermediate_size_per_partition,
-            params_dtype,
-            **extra_weight_attrs,
         )
 
+        fp4_block_k = 32
+
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size // 2,
+                dtype=torch.int8,
+            ),
+            requires_grad=False,
+        )
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition // 2,
+                dtype=torch.int8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        w13_weight_scale = torch.nn.Parameter(
+            torch.ones(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size // fp4_block_k,
+                dtype=torch.float32,
+            ),
+            requires_grad=False,
+        )
+        w2_weight_scale = torch.nn.Parameter(
+            torch.ones(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition // fp4_block_k,
+                dtype=torch.float32,
+            ),
+            requires_grad=False,
+        )
+        w13_weight_scale.format_ue8m0 = False
+        w2_weight_scale.format_ue8m0 = False
+        scale_attrs = dict(extra_weight_attrs)
+        scale_attrs["quant_method"] = FusedMoeWeightScaleSupported.BLOCK.value
+        layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
+        set_weight_attrs(w13_weight_scale, scale_attrs)
+        layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
+        set_weight_attrs(w2_weight_scale, scale_attrs)
+
     def process_weights_after_loading(self, layer: Module) -> None:
         from sglang.srt.layers.quantization.marlin_utils import (
             check_moe_marlin_supports_layer,
@@ -102,8 +147,8 @@ class Mxfp4MarlinMoEMethod:
         quant_info = MarlinMoeQuantInfo(
             w13_qweight=layer.w13_weight,
             w2_qweight=layer.w2_weight,
-            w13_scales=layer.w13_weight_scale_inv,
+            w13_scales=layer.w13_weight_scale,
-            w2_scales=layer.w2_weight_scale_inv,
+            w2_scales=layer.w2_weight_scale,
             w13_g_idx_sort_indices=None,
             w2_g_idx_sort_indices=None,
             weight_bits=4,

test/registered/dsv4/test_deepseek_v4_flash_fp4_h200.py

@@ -66,6 +66,8 @@ class TestDSV4FlashFP4H200(ServerSanityMixin, CustomTestCase):
                 "1",
                 "--speculative-num-draft-tokens",
                 "4",
+                "--watchdog-timeout",
+                "900",
             ],
         )
 

test/registered/dsv4/test_deepseek_v4_flash_fp8_h200.py

@@ -63,6 +63,8 @@ class TestDSV4FlashFP8H200(ServerSanityMixin, CustomTestCase):
                 "128",
                 "--deepep-config",
                 DEEPEP_CONFIG,
+                "--watchdog-timeout",
+                "900",
             ],
             env={
                 "SGLANG_DSV4_FP4_EXPERTS": "0",