From 37f18438c593809abfd6edc816f77b5257bb6797 Mon Sep 17 00:00:00 2001 From: shiyu7 <65223714+shiyu7@users.noreply.github.com> Date: Thu, 14 May 2026 07:33:46 +0800 Subject: [PATCH] [rebase]Deepseek_v4 support w4(mxfp4)a16 on hopper (#24986) --- .../csrc/gemm/marlin_moe/moe_wna16_marlin.cuh | 10 +++ .../moe/fused_moe_triton/fused_marlin_moe.py | 10 +-- .../layers/quantization/marlin_utils_fp4.py | 48 ++++++++----- .../sglang/srt/layers/quantization/mxfp4.py | 41 ++++++++++- .../layers/quantization/mxfp4_marlin_moe.py | 69 +++++++++++++++---- .../dsv4/test_deepseek_v4_flash_fp4_h200.py | 2 + .../dsv4/test_deepseek_v4_flash_fp8_h200.py | 2 + 7 files changed, 146 insertions(+), 36 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh b/python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh index 81c021dc8..a3b563151 100644 --- a/python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh +++ b/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::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"); 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 a2f3f845e..f62fe7278 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 @@ -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 ( diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py index 11a664c88..4ebca81d9 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py +++ b/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 - 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) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 4dedb85f5..70425a73c 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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: diff --git a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py index 90a3de66f..f7fc76dfc 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py +++ b/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 — - # 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, diff --git a/test/registered/dsv4/test_deepseek_v4_flash_fp4_h200.py b/test/registered/dsv4/test_deepseek_v4_flash_fp4_h200.py index 4c9423638..c5c7dbc17 100644 --- a/test/registered/dsv4/test_deepseek_v4_flash_fp4_h200.py +++ b/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", ], ) diff --git a/test/registered/dsv4/test_deepseek_v4_flash_fp8_h200.py b/test/registered/dsv4/test_deepseek_v4_flash_fp8_h200.py index 81893bc82..790985e08 100644 --- a/test/registered/dsv4/test_deepseek_v4_flash_fp8_h200.py +++ b/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",