from __future__ import annotations import logging from typing import TYPE_CHECKING import torch 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, round_up, set_weight_attrs from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput logger = logging.getLogger(__name__) class Mxfp4MarlinMoEMethod: """MXFP4 (E8M0 scales) MoE quantization method using the Marlin backend.""" def __init__(self, fp8_method, prefix: str): self._fp8 = fp8_method self.prefix = prefix def create_moe_runner(self, layer, moe_runner_config): from sglang.srt.layers.moe.moe_runner import MoeRunner self.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config) def create_weights( self, layer: Module, num_experts: int, hidden_size: int, intermediate_size_per_partition: int, params_dtype: torch.dtype, **extra_weight_attrs, ): from sglang.srt.layers.moe.fused_moe_triton import ( FusedMoeWeightScaleSupported, ) layer._dsv4_mxfp4_backend = None # set in process_weights_after_loading fp4_block_k = 32 intermediate_size_per_partition = round_up(intermediate_size_per_partition, 128) hidden_size = round_up(hidden_size, 256) self.hidden_pad = hidden_size - layer.hidden_size 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, ) from sglang.srt.layers.quantization.marlin_utils_fp4 import ( deinterleave_moe_mxfp4_w13_for_marlin, prepare_moe_mxfp4_layer_for_marlin, ) # Let the FP8 base method handle ROCm normalization, etc. self._fp8.process_weights_after_loading(layer) if getattr(layer, "_mega_moe_weights_built", False): return if not is_sm90_supported() and not is_sm120_supported(): raise RuntimeError("MXFP4 Marlin requires SM90 or SM120.") if not check_moe_marlin_supports_layer(layer, 32, allow_tile_padding=True): raise RuntimeError( "Current MXFP4 MoE layer does not satisfy Marlin constraints." ) # NOTE: the Marlin MoE runner consumes w13 in the checkpoint's # native ``[w1; w3]`` order -- see ``silu_and_mul`` in # fused_marlin_moe.py which expects ``gate = intermediate[:, :N]`` # (first half) and ``up = intermediate[:, N:]`` (second half). # Unlike the flashinfer trtllm_fp4 kernel (which wants [w3, w1]), # we must *not* call ``reorder_w1w3_to_w3w1`` here. log_info_on_rank0( logger, f"Preparing MXFP4 experts for Marlin backend " f"(layer: {self.prefix})...", ) if self.runner.config.gemm1_alpha is not None: deinterleave_moe_mxfp4_w13_for_marlin(layer) prepare_moe_mxfp4_layer_for_marlin(layer) layer._dsv4_mxfp4_backend = "marlin" def apply( self, layer: Module, dispatch_output: DispatchOutput, ) -> CombineInput: from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker topk_output = dispatch_output.topk_output if not TopKOutputChecker.format_is_standard(topk_output): raise ValueError(f"Unsupported topk output format: {topk_output.format}") hidden_states = dispatch_output.hidden_states target_hidden_size = layer.w13_weight.shape[1] * 16 if hidden_states.shape[-1] == target_hidden_size: hidden_states_padded = hidden_states else: hidden_states_padded = torch.nn.functional.pad( hidden_states, (0, target_hidden_size - hidden_states.shape[-1]), mode="constant", value=0.0, ) 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, w13_bias=getattr(layer, "w13_weight_bias", None), w2_bias=getattr(layer, "w2_weight_bias", None), ) runner_output = self.runner.run( dispatch_output._replace(hidden_states=hidden_states_padded), quant_info=quant_info, ) return StandardCombineInput(hidden_states=runner_output.hidden_states)