[MiniMax-M3] Fuse MXFP8 UE8M0 scale packing into ep_scatter (#21)

Remove the `use_deep_gemm_packed_mxfp8` flag and the special MXFP8
quantization branch in `moe_kernel_quantize_input`. MXFP8 activations are
now always quantized to the plain non-swizzled (M, K/32) uint8 UE8M0 scale
layout, and the pack into DeepGEMM's consumed scale layout (int32, MN-major,
TMA-aligned, 4 UE8M0 per int32) is fused directly into the expert-permute
scatter instead of relying on the GEMM's internal repack.

- utils.py: drop the flag + branch; mxfp8 always uses _mxfp8_e4m3_quantize
  (non-swizzled).
- config.py / no_dp_ep.py / oracle/fp8.py: remove the flag plumbing.
- deep_gemm_utils.py: add a PACK_UE8M0 path to _fwd_kernel_ep_scatter_2 that
  concatenates 4 UE8M0 bytes per int32 and stores MN-major; deepgemm_moe_permute
  allocates the TMA-aligned int32 buffer for the uint8 path. Float32 (FP8/FP4)
  scales keep the row-major path unchanged.

mm1 now feeds the grouped GEMM pre-packed int32 scales with recipe_a=(1,32),
matching what the mm2 activation-quant path already does (validate-only
transform).

AI assistance (Claude Code) was used for this change.

Tests run on GB200 (SM100):
- Numerical unit test: fused packed scatter matches a torch reference for
  data placement and byte-packing; output layout/stride is identical to
  per_token_group_quant_fp8_packed_for_deepgemm.
- gsm8k 5-shot, MiniMax-M3-preview:
    TP=4 monolithic : exact_match 0.9249 +/- 0.0073
    DP=4 + EP       : exact_match 0.9325 +/- 0.0069


Co-authored-by: Claude

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

vllm/model_executor/layers/fused_moe/config.py

@@ -256,11 +256,6 @@ class FusedMoEQuantConfig:
 
     mx_alignment: int = 0
 
-    # When True, MXFP8 (1x32) activations are quantized into DeepGEMM's packed
-    # UE8M0 scale layout in the prepare phase. Set only for the DeepGEMM mxfp8
-    # backend; other backends (FlashInfer/Marlin) want the unpacked layout.
-    use_deep_gemm_packed_mxfp8: bool = False
-
     def __post_init__(self):
         assert not self.per_act_token_quant or self.block_shape is None, (
             "illegal quantization"
@@ -511,7 +506,6 @@ class FusedMoEQuantConfig:
         gemm1_alpha: float | None = None,
         gemm1_beta: float | None = None,
         gemm1_clamp_limit: float | None = None,
-        use_deep_gemm_packed_mxfp8: bool = False,
     ) -> "FusedMoEQuantConfig":
         """
         General builder function for a FusedMoEQuantConfig.
@@ -583,7 +577,6 @@ class FusedMoEQuantConfig:
             gemm1_alpha=gemm1_alpha,
             gemm1_beta=gemm1_beta,
             gemm1_clamp_limit=gemm1_clamp_limit,
-            use_deep_gemm_packed_mxfp8=use_deep_gemm_packed_mxfp8,
         )
         assert quant_config.per_act_token_quant == per_act_token_quant
         assert quant_config.per_out_ch_quant == per_out_ch_quant

vllm/model_executor/layers/fused_moe/deep_gemm_utils.py

@@ -130,6 +130,9 @@ def _fwd_kernel_ep_scatter_2(
     HIDDEN_SIZE_PAD: tl.constexpr,
     SCALE_HIDDEN_SIZE: tl.constexpr,
     SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
+    PACK_UE8M0: tl.constexpr,
+    SCALE_PACKED_SIZE: tl.constexpr,
+    SCALE_PACKED_SIZE_PAD: tl.constexpr,
 ):
     start_token_id = tl.program_id(0)
     grid_num = tl.num_programs(0)
@@ -137,16 +140,47 @@ def _fwd_kernel_ep_scatter_2(
     offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
     mask = offset_in < HIDDEN_SIZE
 
-    offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
-    mask_s = offset_in_s < SCALE_HIDDEN_SIZE
-
     output_tensor_stride0 = output_tensor_stride0.to(tl.int64)
 
+    if PACK_UE8M0:
+        # One int32 per 4 consecutive 32-wide UE8M0 groups, stored MN-major.
+        offs_pk = tl.arange(0, SCALE_PACKED_SIZE_PAD)
+        mask_pk = offs_pk < SCALE_PACKED_SIZE
+    else:
+        offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
+        mask_s = offset_in_s < SCALE_HIDDEN_SIZE
+
     for token_id in range(start_token_id, total_token_num, grid_num):
         to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
-        to_copy_s = tl.load(
-            recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s
-        )
+
+        if PACK_UE8M0:
+            # Pack 4 UE8M0 bytes into one int32 (byte j = group 4*pk+j).
+            base_s = recv_x_scale + token_id * recv_x_scale_stride0
+            g0, g1 = offs_pk * 4, offs_pk * 4 + 1
+            g2, g3 = offs_pk * 4 + 2, offs_pk * 4 + 3
+            b0 = tl.load(
+                base_s + g0 * recv_x_scale_stride1, mask=g0 < SCALE_HIDDEN_SIZE
+            )
+            b1 = tl.load(
+                base_s + g1 * recv_x_scale_stride1, mask=g1 < SCALE_HIDDEN_SIZE
+            )
+            b2 = tl.load(
+                base_s + g2 * recv_x_scale_stride1, mask=g2 < SCALE_HIDDEN_SIZE
+            )
+            b3 = tl.load(
+                base_s + g3 * recv_x_scale_stride1, mask=g3 < SCALE_HIDDEN_SIZE
+            )
+            packed_s = (
+                b0.to(tl.int32)
+                | (b1.to(tl.int32) << 8)
+                | (b2.to(tl.int32) << 16)
+                | (b3.to(tl.int32) << 24)
+            )
+        else:
+            to_copy_s = tl.load(
+                recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s,
+                mask=mask_s,
+            )
 
         for topk_index in tl.range(0, topk_num, 1, num_stages=4):
             expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
@@ -164,11 +198,21 @@ def _fwd_kernel_ep_scatter_2(
                 output_tensor_ptr = (
                     output_tensor + dest_token_index_i64 * output_tensor_stride0
                 )
+                tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
+
                 output_tensor_scale_ptr = (
                     output_tensor_scale + dest_token_index * output_tensor_scale_stride0
                 )
-                tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
-                tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s)
+                if PACK_UE8M0:
+                    tl.store(
+                        output_tensor_scale_ptr + offs_pk * output_tensor_scale_stride1,
+                        packed_s,
+                        mask=mask_pk,
+                    )
+                else:
+                    tl.store(
+                        output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s
+                    )
 
 
 @torch.no_grad()
@@ -184,6 +228,7 @@ def ep_scatter(
     m_indices: torch.Tensor,
     output_index: torch.Tensor,
     block_size: int = 128,
+    pack_ue8m0: bool = False,
 ):
     BLOCK_E = 128  # token num of per expert is aligned to 128
     BLOCK_D = block_size  # block size of activation-scale quantization
@@ -196,6 +241,10 @@ def ep_scatter(
     assert m_indices.shape[0] % BLOCK_E == 0
     assert expert_start_loc.shape[0] == num_experts
 
+    # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
+    scale_hidden_size = hidden_size // BLOCK_D
+    scale_packed_size = (scale_hidden_size + 3) // 4 if pack_ue8m0 else 1
+
     _fwd_kernel_ep_scatter_1[(grid,)](
         num_recv_tokens_per_expert,
         expert_start_loc,
@@ -235,8 +284,11 @@ def ep_scatter(
         num_warps=num_warps,
         HIDDEN_SIZE=hidden_size,
         HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
-        SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D,
-        SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D),
+        SCALE_HIDDEN_SIZE=scale_hidden_size,
+        SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size),
+        PACK_UE8M0=pack_ue8m0,
+        SCALE_PACKED_SIZE=scale_packed_size,
+        SCALE_PACKED_SIZE_PAD=triton.next_power_of_2(scale_packed_size),
     )
     return
 
@@ -382,9 +434,21 @@ def deepgemm_moe_permute(
     if aq_out is None:
         aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype)
 
-    aq_scale_out = torch.empty(
-        (M_sum, H // block_k), device=device, dtype=torch.float32
-    )
+    # uint8 UE8M0 (MXFP8) -> scatter packs into DeepGEMM's int32 MN-major
+    # TMA-aligned layout; float32 (FP8/FP4) scattered row-major as-is.
+    pack_ue8m0 = aq_scale.dtype == torch.uint8
+    sf_k = H // block_k
+    if pack_ue8m0:
+        packed_sf_k = (sf_k + 3) // 4
+        tma_aligned_mn = round_up(M_sum, 4)
+        aq_scale_out = torch.empty_strided(
+            (M_sum, packed_sf_k),
+            (1, tma_aligned_mn),
+            device=device,
+            dtype=torch.int32,
+        )
+    else:
+        aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32)
 
     # DeepGEMM uses negative values in m_indices (here expert_ids) to mark
     # completely invalid / padded blocks that should be skipped. We always
@@ -419,6 +483,7 @@ def deepgemm_moe_permute(
         m_indices=expert_ids,
         output_index=inv_perm,
         block_size=block_k,
+        pack_ue8m0=pack_ue8m0,
     )
 
     return aq_out, aq_scale_out, expert_ids, inv_perm

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -560,11 +560,9 @@ def make_fp8_moe_quant_config(
             g2_alphas=(w2_scale * a2_scale).squeeze(),
             gemm1_clamp_limit=swiglu_limit,
         )
-    # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to the
-    # mxfp8 activation quant rather than standard FP8 block quantization.
-    # Non-swizzled layout is required since the TRTLLM kernel expects scales in
-    # (num_tokens, hidden_dim // 32) format. DeepGEMM instead needs its packed
-    # UE8M0 scale layout, selected via use_deep_gemm_packed_mxfp8.
+    # MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are
+    # the non-swizzled (num_tokens, hidden_dim // 32) uint8 UE8M0 layout for all
+    # backends; the DeepGEMM expert permute repacks them for the grouped GEMM.
     if block_shape == [1, 32]:
         return FusedMoEQuantConfig.make(
             "mxfp8",
@@ -579,7 +577,6 @@ def make_fp8_moe_quant_config(
             gemm1_alpha=gemm1_alpha,
             gemm1_beta=gemm1_beta,
             gemm1_clamp_limit=swiglu_limit,
-            use_deep_gemm_packed_mxfp8=(fp8_backend == Fp8MoeBackend.DEEPGEMM),
         )
 
     # All other backends use normal config.

vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py

@@ -32,7 +32,6 @@ def _quantize_input(
         block_shape=quant_config.block_shape,
         is_scale_swizzled=quant_config.is_scale_swizzled,
         mx_alignment=quant_config.mx_alignment,
-        use_deep_gemm_packed_mxfp8=quant_config.use_deep_gemm_packed_mxfp8,
     )
 
     return a1q, a1q_scale

vllm/model_executor/layers/fused_moe/utils.py

@@ -258,7 +258,6 @@ def moe_kernel_quantize_input(
     ocp_mx_scheme: str | None = None,
     quantization_emulation: bool = False,
     mx_alignment: int = 0,
-    use_deep_gemm_packed_mxfp8: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation
     if ocp_mx_scheme is not None:
@@ -314,13 +313,8 @@ def moe_kernel_quantize_input(
                 "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
                 "quantization emulation. Please open an issue."
             )
-        # DeepGEMM consumes per-(1, group) float32 scales (cast to UE8M0
-        # internally), exactly like the FP8 128-block path but with group=32.
-        # The expert permute (deepgemm_moe_permute) is given the same group so
-        # the scale width matches.
-        if use_deep_gemm_packed_mxfp8:
-            assert block_shape is not None
-            return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
+        # Non-swizzled (M, K/32) uint8 UE8M0 scales; deepgemm_moe_permute packs
+        # them for DeepGEMM, TRTLLM takes them as-is.
         return _mxfp8_e4m3_quantize(
             A,
             A_scale,