[Fix] compressed-tensors: Serve split-expert MXFP4 MoE via FP4 Marlin

python/sglang/kernels/jit/csrc/gemm/marlin/gptq_marlin.cuh

@@ -264,9 +264,13 @@ bool is_valid_config(
       q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS &&                  \
       thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS &&        \
       num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) {                                                      \
+    constexpr auto S_TYPE = W_TYPE == host::kFE2M1f                                                                    \
+                                ? (GROUP_BLOCKS == 1 ? host::kFE4M3fn : host::kFE8M0fnu)                               \
+                                : (std::is_same<scalar_t, half>::value ? host::kFloat16 : host::kBFloat16);            \
     kernel = Marlin<                                                                                                   \
         scalar_t,                                                                                                      \
         W_TYPE.id(),                                                                                                   \
+        S_TYPE.id(),                                                                                                   \
         NUM_THREADS,                                                                                                   \
         THREAD_M_BLOCKS,                                                                                               \
         THREAD_N_BLOCKS,                                                                                               \
@@ -282,7 +286,8 @@ bool is_valid_config(
 // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
 // FZP: cases for float-zero-point (is_zp_float = true)
 // ACT: cases for act order case (group_blocks == 0)
-// FP4: cases for nvfp4(e2m1) (group_blocks == 1)
+// NVFP4: cases for nvfp4(e2m1) (group_blocks == 1)
+// MXFP4: cases for mxfp4(e2m1) (group_blocks == 2)
 #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
   _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
   _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false)   \
@@ -339,22 +344,39 @@ bool is_valid_config(
   BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)  \
   BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128)
 
-#define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)        \
+#define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
   _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
   _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
 
-#define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+#define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
   _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
   _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
   _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
 
-#define FP4_GET_IF(W_TYPE)            \
-  FP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
-  FP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
-  FP4_GET_IF_M1(W_TYPE, 4, 8, 128)    \
-  FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
-  FP4_GET_IF_M234(W_TYPE, 8, 4, 128)  \
-  FP4_GET_IF_M234(W_TYPE, 4, 8, 128)
+#define NVFP4_GET_IF(W_TYPE)            \
+  NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+  NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+  NVFP4_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+  NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
+  NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+  NVFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
+
+#define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
+
+#define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
+
+#define MXFP4_GET_IF(W_TYPE)            \
+  MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+  MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+  MXFP4_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+  MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
+  MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+  MXFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
 
 // We currently have 4-bit models only with group_blocks == 4
 #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
@@ -413,7 +435,7 @@ MarlinFuncPtr get_marlin_kernel(
   COMMON_GET_IF(host::kU4B8)
   COMMON_GET_IF(host::kU8B128)
 
-  FP4_GET_IF(host::kFE2M1f)
+  NVFP4_GET_IF(host::kFE2M1f)
 
   BIGGROUP_GET_IF(host::kFE4M3fn)
 
@@ -426,6 +448,16 @@ MarlinFuncPtr get_marlin_kernel(
     FZP_GET_IF(host::kU4)
   }
 
+  // MXFP4's E8M0 scale decoder is only defined for bf16 (see dequant.h), so the
+  // half branch must not even be instantiated -- `if constexpr`, not a runtime
+  // `if`, which would still compile the body and fail to link half2/kFE8M0fnu.
+  if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+    if (false) {
+      // _GET_IF expands to "else if", so it needs a preceding if to chain onto.
+    }
+    MXFP4_GET_IF(host::kFE2M1f)
+  }
+
   return kernel;
 }
 
@@ -918,14 +950,6 @@ void gptq_marlin_gemm(
     }
   }
 
-  // Verify global_scale
-  int64_t global_scale_size = global_scale.size(0);
-  if (global_scale_size > 0) {
-    RuntimeCheck(b_q_type == kFE2M1f, "global_scale can only be used for float4_e2m1f.");
-  } else {
-    RuntimeCheck(!(b_q_type == kFE2M1f), "the global_scale parameter must be passed for float4_e2m1f.");
-  }
-
   // Derive group_size
   int group_size = -1;
   if (has_act_order) {
@@ -945,6 +969,29 @@ void gptq_marlin_gemm(
     }
   }
 
+  // Verify global_scale (Optional unwrap done in Python).
+  // Only NVFP4 (group_size 16) carries one; MXFP4 folds its scale into the E8M0
+  // group scales.
+  int64_t global_scale_size = global_scale.size(0);
+  if (global_scale_size > 0) {
+    RuntimeCheck(b_q_type == kFE2M1f && group_size == 16, "global_scale can only be used for nvfp4 format.");
+  } else {
+    RuntimeCheck(
+        !(b_q_type == kFE2M1f && group_size == 16), "the global_scale parameter must be passed for nvfp4 format.");
+  }
+
+  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);
+    // The E8M0 scale decoder is only instantiated for bf16 (see dequant.h), so a
+    // half call would otherwise dispatch to the no-op MarlinDefault kernel.
+    RuntimeCheck(
+        group_size != 32 || std::is_same<scalar_t, nv_bfloat16>::value,
+        "MXFP4 Marlin with E8M0 scales is only instantiated for bfloat16 activations.");
+  }
+
   // Verify workspace and get device info
   RuntimeCheck(
       size_n % device::marlin::min_thread_n == 0,

python/sglang/kernels/jit/csrc/gemm/marlin/kernel.h

@@ -16,6 +16,7 @@ namespace device::marlin {
 template <
     typename scalar_t,                   // compute dtype, half or nv_float16
     const host::ScalarTypeId w_type_id,  // weight ScalarType id
+    const host::ScalarTypeId s_type_id,  // weight scale ScalarType id
     const int threads,                   // number of threads in a threadblock
     const int thread_m_blocks,           // number of 16x16 blocks in the m
                                          // dimension (batchsize) of the

python/sglang/kernels/jit/csrc/gemm/marlin/marlin_template.h

@@ -269,6 +269,7 @@ __device__ inline void wait_negative_and_add(int* lock) {
 template <
     typename scalar_t,                   // compute dtype, half or nv_float16
     const host::ScalarTypeId w_type_id,  // weight ScalarType id
+    const host::ScalarTypeId s_type_id,  // weight scale ScalarType id
     const int threads,                   // number of threads in a threadblock
     const int thread_m_blocks,           // number of 16x16 blocks in the m
                                          // dimension (batchsize) of the
@@ -324,17 +325,29 @@ __global__ void Marlin(
   using FragZP = typename ScalarType<scalar_t>::FragZP;
 
   static constexpr auto w_type = host::ScalarType::from_id(w_type_id);
+  static constexpr auto s_type = host::ScalarType::from_id(s_type_id);
+  if constexpr (w_type == host::kFE2M1f) {
+    static_assert(s_type == host::kFE4M3fn && group_blocks == 1 || s_type == host::kFE8M0fnu && group_blocks == 2);
+  } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+    static_assert(s_type == host::kBFloat16);
+  } else if constexpr (std::is_same<scalar_t, half>::value) {
+    static_assert(s_type == host::kFloat16);
+  }
+
   constexpr bool has_zp = w_type == host::kU4 || w_type == host::kU8;
-  constexpr bool is_int_type =
-      w_type == host::kU4 || w_type == host::kU8 || w_type == host::kU4B8 || w_type == host::kU8B128;
-  // see comments of dequant.h for more details
-  constexpr bool dequant_skip_flop = !is_int_type ||
+  // see comments of dequant.h for more details.
+  // The first two clauses enumerate the float weight types that can skip the flop
+  // rather than accepting every non-int type: validity depends on the scale
+  // encoding, so MXFP4 (kFE2M1f with E8M0 scales) is excluded here. A new float
+  // weight type gets false until its dequant path is analyzed.
+  constexpr bool dequant_skip_flop = w_type == host::kFE4M3fn || w_type == host::kFE2M1f && s_type == host::kFE4M3fn ||
                                      has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
                                      has_zp && !is_zp_float && !(w_type == host::kU8);
 
   scalar_t2 global_scale;
 
-  if constexpr (w_type == host::kFE2M1f) {
+  // MXFP4 folds its scale into the E8M0 group scales; only NVFP4 has a global one.
+  if constexpr (w_type == host::kFE2M1f && s_type == host::kFE4M3fn) {
     uint16_t val = scale2_ptr[0];
     global_scale = Dtype::num2num2(*reinterpret_cast<scalar_t*>(&val));
   }
@@ -486,10 +499,10 @@ __global__ void Marlin(
   constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
 
   // Scale sizes/strides without act_order.
-  // NVFP4 packs FP8 (8-bit) scales into shared/global memory at twice the
-  // density of the half-precision scale path, so the strides scale with the
-  // element size.
-  constexpr bool is_8bit_scale = w_type == host::kFE2M1f;
+  // FP4 packs 8-bit scales (NVFP4 E4M3, MXFP4 E8M0) into shared/global memory at
+  // twice the density of the half-precision scale path, so the strides scale with
+  // the element size.
+  constexpr bool is_8bit_scale = s_type.size_bits() == 8;
   int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8);
   constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8);
   constexpr int s_tb_groups =
@@ -876,7 +889,9 @@ __global__ void Marlin(
 
           int4* sh_s_stage = sh_s + s_sh_stage * pipe;
 
-          if constexpr (w_type_id != host::kFE2M1f.id()) {
+          // 8-bit scales are half as wide, so they read int2 rather than int4. This
+          // is a property of s_type, not of the weight type that usually implies it.
+          if constexpr (!is_8bit_scale) {
             reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
           } else {
             reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
@@ -1080,8 +1095,8 @@ __global__ void Marlin(
       int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
       int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
 
-      dequant_fp8_scales<scalar_t2>(s_quant_0, reinterpret_cast<scalar_t2*>(&frag_s[k2]));
-      dequant_fp8_scales<scalar_t2>(s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
+      dequant_fp8_scales<scalar_t2, s_type_id>(s_quant_0, reinterpret_cast<scalar_t2*>(&frag_s[k2]));
+      dequant_fp8_scales<scalar_t2, s_type_id>(s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
     }
 
 // We have the m dimension as the inner loop in order to encourage overlapping
@@ -1363,7 +1378,7 @@ __global__ void Marlin(
         res = __hmul2(res, s[0]);
       }
 
-      if constexpr (w_type == host::kFE2M1f) {
+      if constexpr (w_type == host::kFE2M1f && s_type == host::kFE4M3fn) {
         res = __hmul2(res, global_scale);
       }
 

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

@@ -349,6 +349,7 @@ def __init__(
         params_dtype: Optional[torch.dtype] = None,
         reduce_results: bool = False,
         quant_config: Optional[QuantizationConfig] = None,
+        serves_fused_mxfp4: bool = True,
         prefix: str = "",
         activation: str = "silu",
         apply_router_weight_on_input: bool = False,
@@ -374,6 +375,10 @@ def __init__(
             params_dtype = torch.get_default_dtype()
 
         self.params_dtype = params_dtype
+        # False for a model whose loader keeps compressed-tensors' `weight_packed`
+        # expert names; Mxfp4MoEMethod's fused `w13_weight`/`w2_weight` cannot
+        # serve those, so such a model routes to a compressed-tensors MoE scheme.
+        self.serves_fused_mxfp4 = serves_fused_mxfp4
         self.layer_name = prefix
         self.layer_id = layer_id
         self.is_draft = get_flags().moe.in_speculative_scope

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

@@ -45,6 +45,8 @@
     CompressedTensorsW4A4Fp4,
     CompressedTensorsW4A4Nvfp4MoE,
     CompressedTensorsW4A16Fp4,
+    CompressedTensorsW4A16Mxfp4,
+    CompressedTensorsW4A16Mxfp4MoE,
     CompressedTensorsW4A16Nvfp4MoE,
     CompressedTensorsW4AFP8MoE,
     CompressedTensorsW8A8Fp8,
@@ -232,8 +234,14 @@ def get_quant_method(
             # handles all MoE backends, bypassing the scheme abstraction.
             # On NPU the dedicated Mxfp4MoEMethod does not apply, so fall
             # through to the scheme-based path and let get_moe_scheme select
-            # NPUCompressedTensorsW4A8mxfp4MoE.
-            if self._is_mxfp4_moe(layer_name=prefix) and not _is_npu:
+            # NPUCompressedTensorsW4A8mxfp4MoE. A model that keeps the
+            # `weight_packed` expert names falls through too: Mxfp4MoEMethod
+            # registers the fused `w13_weight`/`w2_weight` instead.
+            if (
+                self._is_mxfp4_moe(layer_name=prefix)
+                and layer.serves_fused_mxfp4
+                and not _is_npu
+            ):
                 from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
 
                 logger.info_once(
@@ -630,6 +638,25 @@ def _is_fp4a16_nvfp4(
             and weight_quant.symmetric
         )
 
+    def _is_mxfp4_linear(self, weight_quant: QuantizationArgs) -> bool:
+        """MXFP4 dense linear: fp4 weights with E8M0 group scales, group_size 32.
+
+        Selects the weight-only FP4 Marlin scheme, the only MXFP4 linear path in
+        tree; MoE is handled separately by Mxfp4MoEMethod via _is_mxfp4_moe.
+        Predicated on the weights alone so it covers both mxfp4a16 and the w4a4
+        variant, whose weights blocks are byte-identical.
+        """
+        if weight_quant is None:
+            return False
+
+        return (
+            weight_quant.strategy == QuantizationStrategy.GROUP.value
+            and weight_quant.type == QuantizationType.FLOAT
+            and weight_quant.num_bits == 4
+            and weight_quant.group_size == 32
+            and weight_quant.symmetric
+        )
+
     def _is_wNa16_group_channel(
         self, weight_quant: BaseModel, input_quant: BaseModel
     ) -> bool:
@@ -750,6 +777,24 @@ def _get_scheme_from_parts(
                     "Other method (CompressedTensorsW4A16Sparse24) is not supported now"
                 )
 
+        # Dispatched ahead of the fp4 branches below: a w4a4 MXFP4 checkpoint
+        # would otherwise reach the non-nvfp4 NotImplementedError, which reports
+        # an unsupported group_size rather than selecting the weight-only path.
+        if self._is_mxfp4_linear(weight_quant):
+            # FP4 Marlin is weight-only on every architecture: there is no native
+            # MXFP4 linear kernel, on Blackwell or elsewhere. A w4a4 checkpoint
+            # therefore runs with unquantized activations, and unlike NVFP4 it
+            # carries no input_global_scale to drop.
+            if input_quant is not None:
+                logger.warning_once(
+                    "MXFP4 w4a4 has no native kernel; serving this checkpoint "
+                    "weight-only via FP4 Marlin instead. Activation "
+                    "quantization is ignored."
+                )
+            return CompressedTensorsW4A16Mxfp4(
+                has_input_activations=input_quant is not None
+            )
+
         if is_activation_quantization_format(quant_format):
             if self._is_fp4a16_nvfp4(weight_quant, input_quant):
                 return CompressedTensorsW4A16Fp4()
@@ -956,6 +1001,18 @@ def get_moe_scheme(
                 ):
                     logger.info_once("Using NPUCompressedTensorsW4A16Int4DynamicMoE")
                     return NPUCompressedTensorsW4A16Int4DynamicMoE(self)
+        elif self._is_mxfp4_linear(weight_quant):
+            if input_quant is not None:
+                # MXFP4 has no FP4-activation MoE kernel in tree; the Marlin
+                # path is weight-only and the input global scales are dropped.
+                logger.warning_once(
+                    "MXFP4 w4a4 MoE is served weight-only via FP4 Marlin. "
+                    "Activation quantization is ignored."
+                )
+            logger.info_once("Using CompressedTensorsW4A16Mxfp4MoE")
+            return CompressedTensorsW4A16Mxfp4MoE(
+                has_input_global_scale=input_quant is not None
+            )
         elif self._is_fp4a16_nvfp4(weight_quant, input_quant):
             logger.info_once("Using CompressedTensorsW4A16Nvfp4MoE")
             return CompressedTensorsW4A16Nvfp4MoE()

python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py

@@ -10,6 +10,8 @@
 from .compressed_tensors_w4a8_fp8_moe import CompressedTensorsW4AFP8MoE
 from .compressed_tensors_w4a8_int8_moe import NPUCompressedTensorsW4A8Int8DynamicMoE
 from .compressed_tensors_w4a8_mxfp4_moe import NPUCompressedTensorsW4A8mxfp4MoE
+from .compressed_tensors_w4a16_mxfp4 import CompressedTensorsW4A16Mxfp4
+from .compressed_tensors_w4a16_mxfp4_moe import CompressedTensorsW4A16Mxfp4MoE
 from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
 from .compressed_tensors_w4a16_nvfp4_moe import CompressedTensorsW4A16Nvfp4MoE
 from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
@@ -47,6 +49,8 @@
     "NPUCompressedTensorsW4A16Int4DynamicMoE",
     "WNA16_SUPPORTED_BITS",
     "CompressedTensorsW4A16Fp4",
+    "CompressedTensorsW4A16Mxfp4",
+    "CompressedTensorsW4A16Mxfp4MoE",
     "CompressedTensorsW4A16Nvfp4MoE",
     "CompressedTensorsW4A4Fp4",
     "CompressedTensorsW4A4Nvfp4MoE",