[ROCm][Perf] MXFP8 dense-linear + grouped-MoE GEMM optimizations for MiniMax-M3 (#46117)
Signed-off-by: amd-ethany <amd-ethany@users.noreply.github.com> Co-authored-by: amd-ethany <amd-ethany@users.noreply.github.com>
Ethan Yang committed
Jul 8, 2026 at 12:03 UTC
d9e57ea82e3e873a74d0ed1c3439fccea2164706
3 files changed
+226
-42
tests/kernels/test_minimax_m3_amd_ops.py
+87
@@ -284,6 +284,93 @@ def test_mxfp8_native_moe(T, H, inter, E, top_k):
284
assert _relerr(got, ref) < 5e-2
285
286
287
+# --------------------------------------------------------------------------- #
288
+# Native MXFP8 grouped GEMM (dot_scaled) vs pure-PyTorch grouped matmul
289
+# --------------------------------------------------------------------------- #
290
+def _ref_grouped_gemm(a_deq, w_deq, topk_ids, a_div, num_valid, mul_weight=None):
291
+ """Pure-PyTorch reference for ``_grouped_gemm_mxfp8``.
292
+
293
+ For each routed (expanded) token ``tid in [0, num_valid)`` the kernel writes
294
+ ``out[tid] = a[tid // a_div] @ w[expert(tid)].T`` (fp32 accumulate), optionally
295
+ scaled by ``mul_weight[tid]``. The expert for ``tid`` is ``topk_ids.flatten()
296
+ [tid]`` (row-major expansion: ``tid = token*top_k + slot``). This is computed
297
+ here with plain ``torch.matmul`` on the dequantized operands — independent of
298
+ the Triton ``dot_scaled`` path and of the (separate) aiter backend.
299
+ """
300
+ eids = topk_ids.reshape(-1)
301
+ n = w_deq.shape[1]
302
+ out = torch.empty(num_valid, n, dtype=torch.float32, device=a_deq.device)
303
+ for tid in range(num_valid):
304
+ e = int(eids[tid].item())
305
+ out[tid] = a_deq[tid // a_div].float() @ w_deq[e].float().T
306
+ if mul_weight is not None:
307
+ out[tid] *= float(mul_weight[tid].item())
308
+ return out
309
+
310
+
311
+@requires_gfx950
312
+@pytest.mark.parametrize("T,N,K,E,top_k", [(8, 256, 128, 8, 2), (5, 512, 256, 16, 4)])
313
+@pytest.mark.parametrize("weighted", [False, True])
314
+@torch.inference_mode()
315
+def test_mxfp8_grouped_gemm_native(T, N, K, E, top_k, weighted):
316
+ """Directly exercise ``_grouped_gemm_mxfp8`` against a non-Triton reference.
317
+
318
+ Covers both call modes used by ``fused_moe_mxfp8_native``:
319
+ * ``weighted=False`` -> g1: ``a_div=top_k`` (a-row shared across the top_k
320
+ expansions of a token), no per-token weight.
321
+ * ``weighted=True`` -> g2: ``a_div=1`` (one a-row per expansion), output
322
+ scaled by ``topk_weights``.
323
+ """
324
+ from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
325
+ _grouped_gemm_mxfp8,
326
+ )
327
+ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
328
+ moe_align_block_size,
329
+ )
330
+
331
+ torch.manual_seed(0)
332
+ block_m = 64
333
+ a_div = 1 if weighted else top_k
334
+ m_routed = T * top_k
335
+ # a-rows: g1 reads one row per token (a_div=top_k); g2 one per expansion.
336
+ a_rows = m_routed if weighted else T
337
+ a_bf16 = torch.randn(a_rows, K, device=DEVICE, dtype=torch.bfloat16) * 0.5
338
+ w_bf16 = torch.randn(E, N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1
339
+ a_fp8, a_scale = _mxfp8_e4m3_quantize_torch(a_bf16, is_sf_swizzled_layout=False)
340
+ w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
341
+
342
+ logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
343
+ topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
344
+ topk_weights = topk_weights.to(torch.float32)
345
+ topk_ids = topk_ids.to(torch.int32)
346
+ mul = topk_weights.reshape(-1) if weighted else None
347
+
348
+ sorted_ids, expert_ids, num_post = moe_align_block_size(
349
+ topk_ids, block_m, E, None, ignore_invalid_experts=False
350
+ )
351
+ got = _grouped_gemm_mxfp8(
352
+ a_fp8,
353
+ a_scale,
354
+ w_fp8,
355
+ w_scale,
356
+ sorted_ids,
357
+ expert_ids,
358
+ num_post,
359
+ m_routed,
360
+ top_k,
361
+ block_m,
362
+ torch.bfloat16,
363
+ a_div=a_div,
364
+ mul_weight_by=mul,
365
+ )
366
+ # Reference: dequant the SAME bits the kernel reads, plain torch matmul.
367
+ a_deq = dequant_mxfp8_to_bf16(a_fp8, a_scale)
368
+ w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
369
+ ref = _ref_grouped_gemm(a_deq, w_deq, topk_ids, a_div, m_routed, mul)
370
+ assert got.shape == (m_routed, N)
371
+ assert _relerr(got, ref) < 5e-2
372
+
373
+
374
# --------------------------------------------------------------------------- #
375
# MXFP8 linear emulation: BF16-at-load (default) vs per-step dequant + switch
376
# --------------------------------------------------------------------------- #
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+76
-7
@@ -89,13 +89,7 @@ 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
- # 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
92
+ BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K)
93
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
94
_mxfp8_linear_kernel[grid](
95
x_q,
@@ -125,6 +119,81 @@ def _mxfp8_dot_scaled_linear(
119
return out
120
121
122
+def _select_cfg(M, N, K):
123
+ """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950.
124
+
125
+ The M-bucketed, shape-adaptive tile selection here is the speedup over the
126
+ upstream 2-bucket launcher. Tiles are pipelined (num_stages>=2, larger BLOCK_K)
127
+ and occupancy- and shape-aware: keyed on the LOCAL (M, N, K), so it adapts to the
128
+ TP-sharded shapes (e.g. MiniMax-M3 TP=4 vs TP=8, where local N and K differ) —
129
+ large-K prefill uses BLOCK_K=256; short-K (K=768) widens N. BLOCK_K must divide K
130
+ (the K-loop is unmasked), so every BLOCK_K below is guarded to be K-divisible
131
+ (served K: 384/768/1024/2048/6144).
132
+ """
133
+ if M <= 64:
134
+ # decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The
135
+ # lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in
136
+ # parallel) + LARGE BLOCK_K (fewer K-iters, bigger coalesced weight loads).
137
+ # Tuned by CUDA-graph replay latency. Optimal at both TP=4 and TP=8.
138
+ if K % 1024 == 0: # K=2048, 6144 -> graph-best 16x16x1024 (all M)
139
+ return 16, 16, 1024, 2, 2
140
+ if K % 512 == 0:
141
+ return 16, 16, 512, 2, 3
142
+ if K % 256 == 0: # K=768 (shared_down) -> graph-best 16x32x256
143
+ return 16, 32, 256, 4, 3
144
+ return 16, 32, 128, 4, 3
145
+ # mid-M (65..256) on SMALL local-N: still occupancy-bound (a 64x64 tile makes too
146
+ # few N-tiles), so the narrow-BLOCK_N decode-style tile fills the CUs better.
147
+ # N<=1536 covers the real fused-qkv local N at TP=8: q heads shard but the GQA KV
148
+ # (4) + sparse-indexer (4) heads are < TP=8, so vLLM replicates them to 1/rank ->
149
+ # N = 1024 + 4*128 = 1536 (not 2560/2=1280). For the wider 1280<N<=1536 band the
150
+ # narrow tile only wins up to M=128 (at M=256 the 64x64 tile is better), so cap it
151
+ # there; N<=1280 keeps the narrow tile through M=256. TP=4 qkv N=2560 is unchanged.
152
+ if (M <= 256 and N <= 1280) or (M <= 128 and N <= 1536):
153
+ if K % 1024 == 0:
154
+ return 16, 16, 1024, 2, 2
155
+ if K % 512 == 0:
156
+ return 16, 16, 512, 2, 3
157
+ if K % 256 == 0:
158
+ return 16, 16, 256, 2, 3
159
+ return 16, 16, 128, 2, 3
160
+ # right-sized launch grid (host-side), used to gate the tall 256-BLOCK_M tile.
161
+ occ = triton.cdiv(M, 256) * triton.cdiv(N, 128)
162
+ if K <= 1024: # short-K (shared_down K=384/768; TP=8 o_proj K=1024)
163
+ if M <= 256:
164
+ return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2)
165
+ # large prefill. (The former 256x128x256 tall tile was faster only on triton
166
+ # 3.6; on triton 3.7 its large BLOCK_M register/LDS footprint spills or hits
167
+ # "out of resources", so use 128x128x256 -- within the known-good footprint.)
168
+ if M >= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256:
169
+ return 128, 128, 256, 8, 3
170
+ return 128, 128, 128, 8, 3
171
+ # large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked):
172
+ # served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since
173
+ # the entry requires K%128==0) for any other K to stay correct.
174
+ if M <= 256: # conc~128 decode + small prefill chunk: occupancy tile
175
+ if K % 512 == 0:
176
+ return 64, 64, 512, 8, 2
177
+ return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2)
178
+ if M <= 1024: # medium prefill chunk: BN=64 keeps small-N occupied
179
+ return (128, 64, 256, 8, 3) if K % 256 == 0 else (128, 64, 128, 8, 3)
180
+ # large prefill (M > 1024). The previously graph-tuned tall 256x128x256 and deep
181
+ # 128x128x512 tiles won only on triton 3.6; on triton 3.7 their larger BLOCK_M /
182
+ # BLOCK_K register+LDS footprint spills (or hits "out of resources" on stricter
183
+ # ROCm/triton builds). The 128x128x256 tile is equal-or-faster on triton 3.7 (the
184
+ # M=4096,N=2048,K=6144 shape: ~104us vs the 256x128x256 tile's ~184us), within ~5%
185
+ # on 3.6, and inside the footprint of the tiles used elsewhere in this selector.
186
+ # Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K
187
+ # / very-large-M shapes.
188
+ if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))):
189
+ return 128, 128, 256, 8, 3
190
+ # small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the
191
+ # N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there).
192
+ if N <= 1024 and K % 256 == 0:
193
+ return 128, 64, 256, 8, 3
194
+ return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3)
195
+
196
+
197
class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
198
"""Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
199
vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py
+63
-35
@@ -19,7 +19,6 @@ and the top-k weighted reduction run in PyTorch between/after the two GEMMs.
19
import torch
20
21
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
22
-from vllm.logger import init_logger
22
from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
23
Mxfp8TritonExpertsBase,
24
)
@@ -32,7 +31,26 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
31
from vllm.platforms import current_platform
32
from vllm.triton_utils import tl, triton
33
35
-logger = init_logger(__name__)
34
+
35
+def _select_cfg(M, N, K, block_m):
36
+ """Pick the launch config from host constants only (M=num_valid_tokens, N, K,
37
+ block_m) — graph-capture safe (no GPU-scalar branch)."""
38
+ # Per-regime winners (measured, isolated cuda-event A/B on gfx950, GPU 3):
39
+ # BLOCK_K=256 (fewer K-iters + bigger MX scale-load coalesced with the dot),
40
+ # num_stages=2 (software-pipeline overlaps the E8M0 scale-load with the scaled
41
+ # MFMA), GROUP_SIZE_M=4 (XCD-friendly swizzle that keeps each touched expert's
42
+ # A rows + B column-tiles L2/MALL-resident across its N-tiles -> kills the
43
+ # redundant A re-read). num_warps stays 8 (wave64 here did NOT spill: VGPR fits
44
+ # and 4 warps was neutral/worse in measurement). BLOCK_K must be a power of two
45
+ # and divide K (K=6144 and K=768 are both multiples of 256).
46
+ BLOCK_K = 256 if K % 256 == 0 else 128
47
+ return {
48
+ "BLOCK_N": 128,
49
+ "BLOCK_K": BLOCK_K,
50
+ "GROUP_SIZE_M": 4,
51
+ "num_warps": 8,
52
+ "num_stages": 2,
53
+ }
54
55
56
@triton.jit
@@ -46,6 +64,7 @@ def _mxfp8_grouped_gemm_kernel(
64
sorted_token_ids_ptr,
65
expert_ids_ptr,
66
num_tokens_post_padded_ptr,
67
+ EM,
68
N,
69
K,
70
num_valid_tokens,
@@ -67,9 +86,33 @@ def _mxfp8_grouped_gemm_kernel(
86
BLOCK_M: tl.constexpr,
87
BLOCK_N: tl.constexpr,
88
BLOCK_K: tl.constexpr,
89
+ GROUP_SIZE_M: tl.constexpr,
90
):
71
- pid_m = tl.program_id(0)
72
- pid_n = tl.program_id(1)
91
+ # ---- Grid-swizzle (super-grouping over M) so consecutive program-ids cover a
92
+ # GROUP_SIZE_M x grid_n super-block. This keeps a touched expert's A rows + its
93
+ # B-weight column-tiles L2/MALL-resident across the N-tiles they share, killing
94
+ # most of the redundant A re-read (the single largest redundant-traffic source:
95
+ # dot_scaled keeps operands in registers, so A is otherwise re-fetched per N-tile).
96
+ # ``num_pid_m`` uses EM (= sorted_token_ids.shape[0]), the SAME bound the grid
97
+ # is sized from, NOT the runtime ``num_tokens_post_padded`` (<= EM). This keeps
98
+ # the swizzle consistent with the grid: ``group_size_m`` is always >= 1 for every
99
+ # launched program, so the ``% group_size_m`` / ``// group_size_m`` below can never
100
+ # hit modulo/division-by-zero. ``num_tokens_post_padded`` is loaded afterwards and
101
+ # only gates the early-return. Mirrors the reference ``fused_moe_kernel``.
102
+ pid = tl.program_id(0)
103
+ num_pid_m = tl.cdiv(EM, BLOCK_M)
104
+ num_pid_n = tl.cdiv(N, BLOCK_N)
105
+ if GROUP_SIZE_M == 1:
106
+ pid_m = pid % num_pid_m
107
+ pid_n = pid // num_pid_m
108
+ else:
109
+ num_pid_in_group = GROUP_SIZE_M * num_pid_n
110
+ group_id = pid // num_pid_in_group
111
+ first_pid_m = group_id * GROUP_SIZE_M
112
+ group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
113
+ pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
114
+ pid_n = (pid % num_pid_in_group) // group_size_m
115
+
116
num_post = tl.load(num_tokens_post_padded_ptr)
117
if pid_m * BLOCK_M >= num_post:
118
return
@@ -140,9 +183,6 @@ def _grouped_gemm_mxfp8(
183
a_div: int,
184
mul_weight_by: torch.Tensor | None = None,
185
expert_map: torch.Tensor | None = None,
143
- block_n: int = 128,
144
- num_warps: int = 8,
145
- num_stages: int = 2,
186
) -> torch.Tensor:
187
M_routed = num_valid_tokens
188
E, N, K = w.shape
@@ -152,8 +192,16 @@ def _grouped_gemm_mxfp8(
192
# written — zero them so the downstream reduction ignores their garbage.
193
alloc = torch.zeros if expert_map is not None else torch.empty
194
out = alloc((M_routed, N), dtype=out_dtype, device=a_q.device)
155
- BLOCK_K = 128
156
- grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, block_n))
195
+
196
+ cfg = _select_cfg(M_routed, N, K, block_m)
197
+ BLOCK_N = cfg["BLOCK_N"]
198
+ BLOCK_K = cfg["BLOCK_K"]
199
+ GROUP_SIZE_M = cfg["GROUP_SIZE_M"]
200
+
201
+ n_pid_m = triton.cdiv(sorted_token_ids.shape[0], block_m)
202
+ n_pid_n = triton.cdiv(N, BLOCK_N)
203
+ grid = (n_pid_m * n_pid_n,)
204
+
205
_mxfp8_grouped_gemm_kernel[grid](
206
a_q,
207
a_scale,
@@ -164,6 +212,7 @@ def _grouped_gemm_mxfp8(
212
sorted_token_ids,
213
expert_ids,
214
num_tokens_post_padded,
215
+ sorted_token_ids.shape[0], # EM: sizes both the grid and the swizzle
216
N,
217
K,
218
num_valid_tokens,
@@ -183,29 +232,15 @@ def _grouped_gemm_mxfp8(
232
A_DIV=a_div,
233
MUL_WEIGHT=mul_weight_by is not None,
234
BLOCK_M=block_m,
186
- BLOCK_N=block_n,
235
+ BLOCK_N=BLOCK_N,
236
BLOCK_K=BLOCK_K,
188
- num_warps=num_warps,
189
- num_stages=num_stages,
237
+ GROUP_SIZE_M=GROUP_SIZE_M,
238
+ num_warps=cfg["num_warps"],
239
+ num_stages=cfg["num_stages"],
240
)
241
return out
242
243
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
-
244
def fused_moe_mxfp8_native(
245
hidden_states: torch.Tensor, # [T, H] bf16
246
w13: torch.Tensor, # [E, 2I, H] fp8
@@ -225,8 +260,7 @@ def fused_moe_mxfp8_native(
260
top_k = topk_ids.shape[1]
261
M = T * top_k
262
228
- tiles = _mxfp8_moe_tiles(T)
229
- block_m = tiles["block_m"]
263
+ block_m = 64
264
# Bin by the actual number of expert weight rows. With fused shared experts
265
# the weight tensor has more rows than ``global_num_experts`` (the routed
266
# count), and their ids fall outside [0, global_num_experts); binning by the
@@ -257,9 +291,6 @@ def fused_moe_mxfp8_native(
291
hidden_states.dtype,
292
a_div=top_k,
293
expert_map=expert_map,
260
- block_n=tiles["block_n"],
261
- num_warps=tiles["num_warps"],
262
- num_stages=tiles["num_stages"],
294
) # [M, 2I]
295
296
# SwiGLU-OAI (split layout: gate=g1[:, :I], up=g1[:, I:]) FUSED with the
@@ -288,9 +319,6 @@ def fused_moe_mxfp8_native(
319
a_div=1,
320
mul_weight_by=topk_weights.reshape(-1).to(torch.float32),
321
expert_map=expert_map,
291
- block_n=tiles["block_n"],
292
- num_warps=tiles["num_warps"],
293
- num_stages=tiles["num_stages"],
322
) # [M, H] == [T*top_k, H]
323
324
return g2.view(T, top_k, H).sum(dim=1).to(hidden_states.dtype)