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

2026-05-13 16:33:46 -07:00
parent 371cb2ade2
commit 37f18438c5
7 changed files with 146 additions and 36 deletions
@@ -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");
@@ -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, (
-            "MXFP4 Marlin expects w1_scale to be torch.float8_e8m0fnu, "
-            f"got {w1_scale.dtype}"
-        )
-        assert w2_scale.dtype == torch.float8_e8m0fnu, (
-            "MXFP4 Marlin expects w2_scale to be torch.float8_e8m0fnu, "
-            f"got {w2_scale.dtype}"
+        assert hidden_states.dtype == torch.bfloat16, (
+            "MXFP4 Marlin with E8M0 scales is only instantiated for bfloat16 "
+            f"activations, got {hidden_states.dtype}"
        )
    else:
        assert (
@@ -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
-    w2_scale = layer.w2_weight_scale_inv.data
-    w13_bias = getattr(layer, "w13_bias", None)
-    w2_bias = getattr(layer, "w2_bias", None)
+    w13_scale = _get_optional_param(layer, "w13_weight_scale", "w13_weight_scale_inv")
+    w2_scale = _get_optional_param(layer, "w2_weight_scale", "w2_weight_scale_inv")
+    w13_bias = _get_optional_param(layer, "w13_weight_bias", "w13_bias")
+    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)
-    w2_scale_marlin = _permute_scales(w2_scale, False)
+    w13_scale_marlin = _permute_scales(w13_scale_data, True)
+    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(
-        w13_scale_marlin, requires_grad=False
-    )
-    layer.w2_weight_scale_inv = torch.nn.Parameter(w2_scale_marlin, requires_grad=False)
+    layer.w13_weight_scale = torch.nn.Parameter(w13_scale_marlin, requires_grad=False)
+    layer.w2_weight_scale = torch.nn.Parameter(w2_scale_marlin, requires_grad=False)

-    if w13_bias is not None:
-        layer.w13_bias = torch.nn.Parameter(
-            _permute_bias(w13_bias), requires_grad=False
+    if w13_bias_data is not None:
+        layer.w13_weight_bias = torch.nn.Parameter(
+            _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)
@@ -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:
@@ -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 —
-        # the raw weight shapes are the same; only post-loading processing differs.
-        self._fp8.create_weights(
-            layer,
-            num_experts,
-            hidden_size,
-            intermediate_size_per_partition,
-            params_dtype,
-            **extra_weight_attrs,
+        from sglang.srt.layers.moe.fused_moe_triton import (
+            FusedMoeWeightScaleSupported,
        )

+        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,
-            w2_scales=layer.w2_weight_scale_inv,
+            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,