A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][Quant] mxfp8 moe/linear gfx950 tuning for MiniMax-M3 (#45725)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
Hongxia Yang committed
Jun 16, 2026 at 14:50 UTC
f2beaa80c8d672dc45c3313e0c8b2cec5a9b3ee3
2 files changed
+39
-7
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+9
-2
@@ -89,7 +89,13 @@ def _mxfp8_dot_scaled_linear(
89
N = w.shape[0]
90
x_q, x_scale = mxfp8_e4m3_quantize(x)
91
out = torch.empty((M, N), dtype=x.dtype, device=x.device)
92
- BLOCK_M, BLOCK_N, BLOCK_K = 64, 128, 128
92
+ # Regime-gated launch tiles for gfx950, tuned at MiniMax-M3 shapes:
93
+ # for example, 8k/1k, 1k/1k
94
+ if M >= 1024:
95
+ BLOCK_M, BLOCK_N, num_warps, num_stages = 128, 256, 8, 2
96
+ else:
97
+ BLOCK_M, BLOCK_N, num_warps, num_stages = 64, 64, 4, 2
98
+ BLOCK_K = 128
99
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
100
_mxfp8_linear_kernel[grid](
101
x_q,
@@ -113,7 +119,8 @@ def _mxfp8_dot_scaled_linear(
119
BLOCK_M=BLOCK_M,
120
BLOCK_N=BLOCK_N,
121
BLOCK_K=BLOCK_K,
116
- num_warps=8,
122
+ num_warps=num_warps,
123
+ num_stages=num_stages,
124
)
125
return out
126
vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py
+30
-5
@@ -140,6 +140,9 @@ def _grouped_gemm_mxfp8(
140
a_div: int,
141
mul_weight_by: torch.Tensor | None = None,
142
expert_map: torch.Tensor | None = None,
143
+ block_n: int = 128,
144
+ num_warps: int = 8,
145
+ num_stages: int = 2,
146
) -> torch.Tensor:
147
M_routed = num_valid_tokens
148
E, N, K = w.shape
@@ -149,9 +152,8 @@ def _grouped_gemm_mxfp8(
152
# written — zero them so the downstream reduction ignores their garbage.
153
alloc = torch.zeros if expert_map is not None else torch.empty
154
out = alloc((M_routed, N), dtype=out_dtype, device=a_q.device)
152
- BLOCK_N = 128
155
BLOCK_K = 128
154
- grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, BLOCK_N))
156
+ grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, block_n))
157
_mxfp8_grouped_gemm_kernel[grid](
158
a_q,
159
a_scale,
@@ -181,13 +183,29 @@ def _grouped_gemm_mxfp8(
183
A_DIV=a_div,
184
MUL_WEIGHT=mul_weight_by is not None,
185
BLOCK_M=block_m,
184
- BLOCK_N=BLOCK_N,
186
+ BLOCK_N=block_n,
187
BLOCK_K=BLOCK_K,
186
- num_warps=8,
188
+ num_warps=num_warps,
189
+ num_stages=num_stages,
190
)
191
return out
192
193
194
+# Tuned native-MXFP8 launch tiles for gfx950 (CDNA4) at MiniMax-M3 MoE shapes.
195
+# For example, 8k/1k, 1k/1k cases.
196
+
197
+_MXFP8_PREFILL_TILES = dict(block_m=128, block_n=256, num_warps=8, num_stages=2)
198
+_MXFP8_DECODE_TILES = dict(block_m=64, block_n=64, num_warps=4, num_stages=2)
199
+_MXFP8_PREFILL_MIN_TOKENS = 1024
200
+
201
+
202
+def _mxfp8_moe_tiles(num_tokens: int) -> dict:
203
+ """Pick grouped-GEMM launch tiles by regime (token count)."""
204
+ if num_tokens >= _MXFP8_PREFILL_MIN_TOKENS:
205
+ return _MXFP8_PREFILL_TILES
206
+ return _MXFP8_DECODE_TILES
207
+
208
+
209
def fused_moe_mxfp8_native(
210
hidden_states: torch.Tensor, # [T, H] bf16
211
w13: torch.Tensor, # [E, 2I, H] fp8
@@ -207,7 +225,8 @@ def fused_moe_mxfp8_native(
225
top_k = topk_ids.shape[1]
226
M = T * top_k
227
210
- block_m = 64
228
+ tiles = _mxfp8_moe_tiles(T)
229
+ block_m = tiles["block_m"]
230
sorted_ids, expert_ids, num_post = moe_align_block_size(
231
topk_ids,
232
block_m,
@@ -232,6 +251,9 @@ def fused_moe_mxfp8_native(
251
hidden_states.dtype,
252
a_div=top_k,
253
expert_map=expert_map,
254
+ block_n=tiles["block_n"],
255
+ num_warps=tiles["num_warps"],
256
+ num_stages=tiles["num_stages"],
257
) # [M, 2I]
258
259
# SwiGLU-OAI (split layout: gate=g1[:, :I], up=g1[:, I:]) FUSED with the
@@ -260,6 +282,9 @@ def fused_moe_mxfp8_native(
282
a_div=1,
283
mul_weight_by=topk_weights.reshape(-1).to(torch.float32),
284
expert_map=expert_map,
285
+ block_n=tiles["block_n"],
286
+ num_warps=tiles["num_warps"],
287
+ num_stages=tiles["num_stages"],
288
) # [M, H] == [T*top_k, H]
289
290
return g2.view(T, top_k, H).sum(dim=1).to(hidden_states.dtype)