[Perf][MiniMax-M3] Sequence parallelism and GEMM-RS for TP + MegaMoE

benchmarks/kernels/benchmark_gemm_rs_ar.py

@@ -18,6 +18,11 @@
 
     torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
         --model deepseek_v41 --backend flashinfer_cutedsl
+
+or the MiniMax-M3 MXFP8 attention O-proj and dense down-proj at TP4 with:
+
+    torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
+        --model minimax_m3 --backend flashinfer_cutedsl
 """
 
 import argparse
@@ -45,10 +50,12 @@
 
 # Per model: (hidden size N, global input widths K). Kimi-K3 lists the
 # shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the
-# ``wo_b`` output projection (o_groups * o_lora_rank).
+# ``wo_b`` output projection (o_groups * o_lora_rank). MiniMax-M3 lists the
+# attention O-proj and the dense-layer down-proj.
 _MODEL_PROJECTIONS = {
     "kimi_k3": (7168, (6144, 12288)),
     "deepseek_v41": (5120, (8192,)),
+    "minimax_m3": (6144, (8192, 12288)),
 }
 
 

tests/kernels/test_gemm_rs_ar.py

@@ -165,6 +165,31 @@ def _run_mode(
         torch.accelerator.synchronize(device)
         _assert_valid_rows_close(actual, expected, M, rank, all_reduce)
 
+        if quant_config is not None:
+            # Producers that already emit MXFP8 (e.g. MiniMax-M3's fused dense
+            # GEMM1) hand in a QuantizedActivation, which skips re-quantization.
+            from vllm.model_executor.layers.fusion.quant_activation import (
+                QuantizedActivation,
+            )
+            from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+                mxfp8_e4m3_quantize,
+            )
+            from vllm.model_executor.layers.quantization.utils.quant_utils import (
+                kMxfp8Dynamic,
+            )
+
+            x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+            x_qa = QuantizedActivation(
+                data=x_q,
+                scale=x_sf,
+                orig_dtype=x.dtype,
+                orig_shape=x.shape,
+                quant_key=kMxfp8Dynamic,
+            )
+            actual = gemm_rs_ar.apply(x_qa, projections[K])
+            torch.accelerator.synchronize(device)
+            _assert_valid_rows_close(actual, expected, M, rank, all_reduce)
+
     if all_reduce:
         # AR outputs must remain valid after the symmetric workspace is reused.
         lifetime_M, lifetime_K = 257, 768

vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py

@@ -41,7 +41,12 @@
 from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem
 from vllm.distributed import get_tp_group
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    as_quantized_activation,
+)
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
 from vllm.platforms import current_platform
 
 if TYPE_CHECKING:
@@ -984,7 +989,9 @@ def should_run(self, x: torch.Tensor) -> bool:
         # supported but faster on the existing LL path.
         return x.shape[0] >= 128
 
-    def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
+    def apply(
+        self, x: torch.Tensor | QuantizedActivation, linear: LinearBase
+    ) -> torch.Tensor:
         from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
             FlashInferCutedslMxfp8LinearKernel,
         )
@@ -998,16 +1005,23 @@ def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
             w = w.t()
         w_sf = getattr(linear, "weight_scale", None)
 
+        # A pre-quantized activation carries F8_128x4-swizzled MXFP8 scales,
+        # the layout the in-kernel quantization below produces.
+        qa = as_quantized_activation(x, kMxfp8Dynamic)
+        if qa is not None:
+            assert qa.orig_dtype == torch.bfloat16
+            x, x_sf = qa.data, qa.scale
+        else:
+            assert isinstance(x, torch.Tensor) and x.dtype == torch.bfloat16
+            x_sf = None
         assert x.ndim == 2
         M, K = x.shape
         assert 0 < M <= self.max_M
-        assert x.dtype == torch.bfloat16
         assert x.device == self.device
         assert x.is_contiguous()
         assert w.device == self.device
         N = self.N
         dtype = "mxfp8" if w.dtype == torch.float8_e4m3fn else "bf16"
-        x_sf = None
         if dtype == "mxfp8":
             from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
                 mxfp8_e4m3_quantize,
@@ -1019,13 +1033,14 @@ def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
             assert w_sf.numel() == N * (K // 32)
             assert w.shape == (N, K) and K % 128 == 0
             assert w.is_contiguous()
-            x, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+            if x_sf is None:
+                x, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
             # The kernel declares the scales as E8M0; the stored/quantized
             # tensors carry the same bits as uint8.
             x_sf = x_sf.view(torch.float8_e8m0fnu)
             w_sf = w_sf.view(torch.float8_e8m0fnu)
         else:
-            assert w_sf is None
+            assert x_sf is None and w_sf is None
             assert w.shape == (N, K) and K % 64 == 0
             assert w.dtype == torch.bfloat16
             assert w.is_contiguous()