Skip to content

[ROCm][Perf][MiniMax-M3.1] Tune MXFP8 linear tiles for 16 < M <= 256 #62509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jiacao-amd wants to merge 3 commits into from
+137 −15
Closed

[ROCm][Perf][MiniMax-M3.1] Tune MXFP8 linear tiles for 16 < M <= 256 #62509

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
99ec093
[ROCm][Perf][MiniMax-M3.1] Tune MXFP8 linear tiles for 16 < M <= 256
jiacao-amd Nov 2, 2026
4bfce4e
[ROCm] Address MiniMax-M3.1 MXFP8 tile review
jiacao-amd Nov 3, 2026
753f28e
Merge branch 'main' into minimax-m3-rocm-fp8-attn-proj
vllm-agent Nov 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 33 additions & 9 deletions tests/kernels/quantization/test_rocm_mxfp8_linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,11 @@
("shared_gate_up", 1536, 6144),
("shared_down", 6144, 768),
]
# Wider local shapes at TP2 (fused qkv, dense-layer gate_up).
M31_TP2_SHAPES = [
("tp2_fused_qkv", 4992, 6144),
("tp2_dense_gate_up", 12288, 6144),
]


def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:
Expand All @@ -72,23 +77,38 @@ def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:


@pytest.mark.parametrize(
"n,k,selected,fallback",
"n,k,selected,mid_m",
[
(2560, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)),
(1536, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)),
(6144, 768, (16, 16, 256, 2, 3), (16, 32, 256, 4, 3)),
(2560, 6144, (16, 16, 1024, 2, 3), (32, 16, 1024, 2, 3)),
(1536, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 3)),
(6144, 768, (16, 16, 256, 2, 3), (32, 16, 256, 2, 3)),
],
)
def test_rocm_mxfp8_linear_selects_m31_tp4_tiles(n, k, selected, fallback):
"""The gfx950 small-M tuning is shape-scoped and ends after M=16."""
def test_rocm_mxfp8_linear_selects_m31_tp4_tiles(n, k, selected, mid_m):
"""The gfx950 small-M tuning is shape-scoped through M=16; from M=17 the
mid-M table takes over."""
from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
from vllm.platforms.rocm import on_gfx950

if not on_gfx950():
pytest.skip("MiniMax-M3.1 small-M tuning requires gfx950")
for num_tokens in (1, 8, 16):
assert rocm_native._select_cfg(num_tokens, n, k) == selected
assert rocm_native._select_cfg(17, n, k) == fallback
assert rocm_native._select_cfg(17, n, k) == mid_m


@pytest.mark.parametrize("k", [768, 1280, 1536, 2048, 3072, 4096, 6144])
@pytest.mark.parametrize("n", [1536, 2560, 3072, 6144, 12288])
def test_gfx950_mid_m_tiles_divide_k(n, k):
"""The mid-M table's BLOCK_K always divides K (the K-loop is unmasked)."""
from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
_gfx950_mid_m_cfg,
)

for m in range(17, 257):
block_m, block_n, block_k, _, _ = _gfx950_mid_m_cfg(m, n, k)
assert k % block_k == 0 and block_k >= 128, (m, n, k, block_k)
assert block_m in (16, 32, 64) and block_n in (16, 32, 64)


@pytest.mark.parametrize(
Expand All @@ -108,9 +128,13 @@ def test_rocm_mxfp8_linear_preserves_non_gfx950_tiles(monkeypatch, n, k, fallbac


@pytest.mark.parametrize(
"name,n,k", V41_ATTN_SHAPES + M31_TP4_SHAPES, ids=lambda v: str(v)
"name,n,k",
V41_ATTN_SHAPES + M31_TP4_SHAPES + M31_TP2_SHAPES,
ids=lambda v: str(v),
)
@pytest.mark.parametrize(
"num_tokens", [1, 2, 7, 8, 16, 17, 31, 32, 48, 64, 96, 128, 192, 256, 1024]
)
@pytest.mark.parametrize("num_tokens", [1, 2, 7, 8, 16, 31, 32, 128, 1024])
@torch.inference_mode()
def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens):
"""Output is finite and matches an exact MXFP8 GEMM of the same inputs.
Expand Down
103 changes: 103 additions & 0 deletions vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,107 @@ def _mxfp8_dot_scaled_gemm(
_BK256_STAGES = 2


# gfx950 tiles for 16 < M <= 256 (decode batches carrying draft tokens, small
# prefill chunks) with local N > 1280: (BLOCK_M, BLOCK_N, preferred BLOCK_K,
# num_warps, num_stages) per M bucket and local-N bucket (<= 1536, <= 1792, <= 2560,
# <= 3072, <= 6144, wider), for K >= 1024. Graph-tuned on MiniMax-M3.1-Flash
# TP2/TP4 linears on MI355X. The 1792-wide bucket also covers V4.1 wqa/wkv
Comment thread
jiacao-amd marked this conversation as resolved.
Resolved
# (N=1792, K=5120), where a narrow N tile avoids a regression at M=64.
_GFX950_MID_M_BUCKETS = (32, 48, 64, 96, 128, 192, 256)
Comment thread
jiacao-amd marked this conversation as resolved.
Resolved
_GFX950_MID_N_BUCKETS = (1536, 1792, 2560, 3072, 6144)
_GFX950_MID_M_TILES = {
32: (
(16, 16, 1024, 2, 3),
(32, 16, 1024, 2, 3),
(32, 16, 1024, 2, 3),
(32, 16, 1024, 2, 3),
(32, 32, 512, 2, 3),
(32, 16, 512, 2, 2),
),
48: (
(32, 16, 1024, 2, 3),
(32, 16, 1024, 2, 2),
# M3.1 TP4 qkv (N=2560, K=6144) and V4.1 wo_a (2048, 4096)
# retain the narrow tile through M=48.
(16, 16, 1024, 2, 2),
(32, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 64, 512, 4, 2),
),
64: (
(32, 16, 1024, 2, 3),
(32, 16, 1024, 2, 2),
(32, 32, 512, 2, 3),
(32, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 64, 512, 4, 2),
),
96: (
(32, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 64, 512, 4, 2),
(64, 64, 256, 4, 2),
),
128: (
(32, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 32, 512, 2, 3),
(64, 64, 512, 4, 2),
(64, 64, 256, 4, 2),
),
192: (
(64, 32, 512, 2, 3),
(64, 64, 512, 4, 2),
(64, 64, 512, 4, 2),
(64, 64, 512, 4, 2),
(64, 64, 256, 4, 2),
(64, 64, 256, 4, 2),
),
256: (
(64, 32, 512, 2, 3),
(64, 64, 512, 4, 2),
(64, 64, 512, 4, 2),
(64, 64, 512, 4, 2),
(64, 64, 256, 4, 2),
(64, 64, 512, 4, 2),
),
}
# The same for short K (< 1024, e.g. a shared-expert down projection).
_GFX950_MID_M_SHORT_K_TILES = {
32: (32, 16, 256, 2, 3),
48: (64, 32, 256, 4, 3),
64: (64, 32, 256, 4, 3),
96: (64, 64, 256, 4, 3),
128: (64, 64, 256, 4, 3),
192: (64, 32, 256, 2, 2),
256: (64, 32, 256, 2, 2),
}


def _gfx950_mid_m_cfg(M, N, K):
m_bucket = next(b for b in _GFX950_MID_M_BUCKETS if b >= M)
if K < 1024:
block_m, block_n, block_k, num_warps, num_stages = _GFX950_MID_M_SHORT_K_TILES[
m_bucket
]
else:
n_idx = next(
(i for i, b in enumerate(_GFX950_MID_N_BUCKETS) if b >= N),
len(_GFX950_MID_N_BUCKETS),
)
block_m, block_n, block_k, num_warps, num_stages = _GFX950_MID_M_TILES[
m_bucket
][n_idx]
# The K-loop is unmasked: shrink BLOCK_K until it divides K (the entry
# requires K % 128 == 0).
while block_k > 128 and K % block_k:
block_k //= 2
return block_m, block_n, block_k, num_warps, num_stages


def _select_cfg(M, N, K):
"""(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950.

Expand All @@ -282,6 +383,8 @@ def _select_cfg(M, N, K):
return 16, 16, 1024, 2, 3
if K == 768 and N == 6144:
return 16, 16, 256, 2, 3
if _ON_GFX950 and 16 < M <= 256 and N > 1280:
return _gfx950_mid_m_cfg(M, N, K)
if M <= 64:
# decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The
# lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in
Expand Down
Loading