[ROCm][Quant] mxfp8 moe/linear gfx950 tuning for MiniMax-M3 (#45725)

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

vllm/model_executor/kernels/linear/mxfp8/rocm_native.py

@@ -89,7 +89,13 @@ def _mxfp8_dot_scaled_linear(
     N = w.shape[0]
     x_q, x_scale = mxfp8_e4m3_quantize(x)
     out = torch.empty((M, N), dtype=x.dtype, device=x.device)
-    BLOCK_M, BLOCK_N, BLOCK_K = 64, 128, 128
+    # Regime-gated launch tiles for gfx950, tuned at MiniMax-M3 shapes:
+    # for example, 8k/1k, 1k/1k
+    if M >= 1024:
+        BLOCK_M, BLOCK_N, num_warps, num_stages = 128, 256, 8, 2
+    else:
+        BLOCK_M, BLOCK_N, num_warps, num_stages = 64, 64, 4, 2
+    BLOCK_K = 128
     grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
     _mxfp8_linear_kernel[grid](
         x_q,
@@ -113,7 +119,8 @@ def _mxfp8_dot_scaled_linear(
         BLOCK_M=BLOCK_M,
         BLOCK_N=BLOCK_N,
         BLOCK_K=BLOCK_K,
-        num_warps=8,
+        num_warps=num_warps,
+        num_stages=num_stages,
     )
     return out
 

vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py

@@ -140,6 +140,9 @@ def _grouped_gemm_mxfp8(
     a_div: int,
     mul_weight_by: torch.Tensor | None = None,
     expert_map: torch.Tensor | None = None,
+    block_n: int = 128,
+    num_warps: int = 8,
+    num_stages: int = 2,
 ) -> torch.Tensor:
     M_routed = num_valid_tokens
     E, N, K = w.shape
@@ -149,9 +152,8 @@ def _grouped_gemm_mxfp8(
     # written — zero them so the downstream reduction ignores their garbage.
     alloc = torch.zeros if expert_map is not None else torch.empty
     out = alloc((M_routed, N), dtype=out_dtype, device=a_q.device)
-    BLOCK_N = 128
     BLOCK_K = 128
-    grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, BLOCK_N))
+    grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, block_n))
     _mxfp8_grouped_gemm_kernel[grid](
         a_q,
         a_scale,
@@ -181,13 +183,29 @@ def _grouped_gemm_mxfp8(
         A_DIV=a_div,
         MUL_WEIGHT=mul_weight_by is not None,
         BLOCK_M=block_m,
-        BLOCK_N=BLOCK_N,
+        BLOCK_N=block_n,
         BLOCK_K=BLOCK_K,
-        num_warps=8,
+        num_warps=num_warps,
+        num_stages=num_stages,
     )
     return out
 
 
+# Tuned native-MXFP8 launch tiles for gfx950 (CDNA4) at MiniMax-M3 MoE shapes.
+# For example, 8k/1k, 1k/1k cases.
+
+_MXFP8_PREFILL_TILES = dict(block_m=128, block_n=256, num_warps=8, num_stages=2)
+_MXFP8_DECODE_TILES = dict(block_m=64, block_n=64, num_warps=4, num_stages=2)
+_MXFP8_PREFILL_MIN_TOKENS = 1024
+
+
+def _mxfp8_moe_tiles(num_tokens: int) -> dict:
+    """Pick grouped-GEMM launch tiles by regime (token count)."""
+    if num_tokens >= _MXFP8_PREFILL_MIN_TOKENS:
+        return _MXFP8_PREFILL_TILES
+    return _MXFP8_DECODE_TILES
+
+
 def fused_moe_mxfp8_native(
     hidden_states: torch.Tensor,  # [T, H] bf16
     w13: torch.Tensor,  # [E, 2I, H] fp8
@@ -207,7 +225,8 @@ def fused_moe_mxfp8_native(
     top_k = topk_ids.shape[1]
     M = T * top_k
 
-    block_m = 64
+    tiles = _mxfp8_moe_tiles(T)
+    block_m = tiles["block_m"]
     sorted_ids, expert_ids, num_post = moe_align_block_size(
         topk_ids,
         block_m,
@@ -232,6 +251,9 @@ def fused_moe_mxfp8_native(
         hidden_states.dtype,
         a_div=top_k,
         expert_map=expert_map,
+        block_n=tiles["block_n"],
+        num_warps=tiles["num_warps"],
+        num_stages=tiles["num_stages"],
     )  # [M, 2I]
 
     # SwiGLU-OAI (split layout: gate=g1[:, :I], up=g1[:, I:]) FUSED with the
@@ -260,6 +282,9 @@ def fused_moe_mxfp8_native(
         a_div=1,
         mul_weight_by=topk_weights.reshape(-1).to(torch.float32),
         expert_map=expert_map,
+        block_n=tiles["block_n"],
+        num_warps=tiles["num_warps"],
+        num_stages=tiles["num_stages"],
     )  # [M, H] == [T*top_k, H]
 
     return g2.view(T, top_k, H).sum(dim=1).to(hidden_states.dtype)