Skip to content

Mxfp8 mtp triton tile tuning #45875

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

Open
nehaprakriya wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Mxfp8 mtp triton tile tuning #45875

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
49ed428
triton GEMM tuning
nehaprakriya Jun 16, 2026
91b1ee4
flatten MXFP8 tile dispatch into (K, N) - keyed table
nehaprakriya Jun 16, 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
100 changes: 57 additions & 43 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 @@ -76,55 +76,70 @@ def _mxfp8_linear_kernel(

o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
tl.store(
o_ptrs, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :]
)
o_ptrs, acc.to(out_ptr.dtype.element_ty),
mask=m_mask[:, None] & n_mask[None, :])



# Tuned launch tiles for gfx950 (CDNA4) at MiniMax-M3 MTP decode shapes.
# Each table maps (M_upper_bound, (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages)).
# Decode path (MTP EAGLE3, num_spec=3, conc=32): verify M≈128, draft M≈32.
# K=6144 skinny-N projections (qkv/gate_up) favour BLOCK_N=32; wider shapes use
# BLOCK_N=64/128. Prefill (M≥1024) uses the wide 128×256 tile for efficiency.
# sentinel None in the last row catches all remaining M.
_TILES: dict[tuple[int, int], list[tuple[int | None, tuple[int, ...]]]] = {

@tjtanaa tjtanaa Jun 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what is the impact on the non-mtp use case?

@nehaprakriya nehaprakriya Jun 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tjtanaa, thanks for the review. There isn’t an MTP vs non‑MTP switch here: tile choice is purely shape‑based from the runtime GEMM (M, K, …).

For non‑MTP workloads that still isn’t special‑cased — it’s the same table:

When M is large (the prefill / long‑context regime, e.g. the M ≥ 1024 bucket in the comment), we intentionally select the wide baseline tile (same high‑throughput strategy as before for large M). So this isn’t “large M gets a new experimental tile”; it’s “large M stays on the known‑good wide path.”
For smaller M (typical decode / skinny projections), we pick per‑bucket tuned tiles instead of a single fixed tile — that’s where the measured win came from in the MTP decode benchmark, but any workload (MTP or not) that hits those M ranges gets the same selection.
The None sentinel row is just the catch‑all for shapes that don’t match an explicit bucket so we don’t leave gaps (not “non‑MTP fallback” specifically).
Net: non‑MTP doesn’t bypass the table; for large M it should land in the wide‑tile row (baseline‑aligned), and for small M it can pick a decode‑tuned row when applicable. I’ll add a short non‑MTP bench line in the PR if you want hard numbers for that path too.

(6144, 2304): [(48, (16, 32, 256, 2, 2)), # qkv_proj decode
(80, (64, 32, 256, 2, 2)),
(112, (32, 32, 512, 2, 2)),
(256, (64, 32, 512, 4, 2))],
(6144, 1536): [(48, (32, 32, 256, 4, 2)), # mlp gate_up decode
(80, (32, 32, 256, 2, 2)),
(112, (16, 32, 512, 2, 2)),
(256, (64, 32, 256, 2, 2))],
}
_GENERIC_TILES: list[tuple[int | None, tuple[int, ...]]] = [
(64, (32, 64, 256, 4, 2)),
(128, (64, 64, 256, 4, 2)),
(1023, (64, 128, 256, 8, 2)),
(None, (128, 256, 256, 8, 2)), # prefill: wide tile for long sequences
]


def _pick_tile(
table: list[tuple[int | None, tuple[int, ...]]], M: int
) -> tuple[int, ...]:
return next(cfg for ub, cfg in table if ub is None or M <= ub)


def _vllm_dot_scaled_kernel(x_q, x_scale, w, w_scale, out_dtype):
M, K = x_q.shape
N = w.shape[0]
out = torch.empty((M, N), dtype=out_dtype, device=x_q.device)

table = _TILES.get((K, N), _GENERIC_TILES)
BLOCK_M, BLOCK_N, BLOCK_K, NUM_WARPS, NUM_STAGES = _pick_tile(table, M)

def _mxfp8_dot_scaled_linear(
x: torch.Tensor, # [M, K] bf16/fp16
w: torch.Tensor, # [N, K] fp8 e4m3
w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0)
) -> torch.Tensor:
M, K = x.shape
N = w.shape[0]
x_q, x_scale = mxfp8_e4m3_quantize(x)
out = torch.empty((M, N), dtype=x.dtype, device=x.device)
# 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,
x_scale,
w,
w_scale,
out,
M,
N,
K,
x_q.stride(0),
x_q.stride(1),
x_scale.stride(0),
x_scale.stride(1),
w.stride(0),
w.stride(1),
w_scale.stride(0),
w_scale.stride(1),
out.stride(0),
out.stride(1),
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
BLOCK_K=BLOCK_K,
num_warps=num_warps,
num_stages=num_stages,
x_q, x_scale, w, w_scale, out, M, N, K,
x_q.stride(0), x_q.stride(1), x_scale.stride(0), x_scale.stride(1),
w.stride(0), w.stride(1), w_scale.stride(0), w_scale.stride(1),
out.stride(0), out.stride(1),
BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
num_warps=NUM_WARPS, num_stages=NUM_STAGES,
)
return out


def _mxfp8_dot_scaled_linear(
x: torch.Tensor,
w: torch.Tensor,
w_scale: torch.Tensor,
) -> torch.Tensor:
x_q, x_scale = mxfp8_e4m3_quantize(x)
return _vllm_dot_scaled_kernel(x_q, x_scale, w, w_scale, x.dtype)


class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
"""Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""

Expand Down Expand Up @@ -162,8 +177,7 @@ def apply_weights(
if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
raise ValueError(
f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got "
f"{layer.weight_scale.dtype}."
)
f"{layer.weight_scale.dtype}.")
out_shape = (*x.shape[:-1], layer.weight.shape[0])
x2d = x.reshape(-1, x.shape[-1])
if x2d.shape[-1] % 128 == 0:
Expand Down
104 changes: 49 additions & 55 deletions vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@
directly (no dequant-to-BF16), and activations are MXFP8-quantized per token.
On CDNA4 ``dot_scaled`` maps to the native MX matrix-core ops; on other archs
Triton upcasts to BF16 (so this stays correct, just not faster) — but the
oracle only selects this path on gfx950 and routes everything else to the
BF16 ``Mxfp8EmulationTritonExperts`` fallback.
kernel selector only picks this path on gfx950 and routes everything else to
the BF16 ``Mxfp8EmulationTritonExperts`` fallback.

Structure mirrors vLLM's ``fused_moe_kernel``: tokens are sorted by expert
(``moe_align_block_size``); each program computes a ``[BLOCK_M, BLOCK_N]`` tile
Expand Down Expand Up @@ -126,10 +126,10 @@ def _mxfp8_grouped_gemm_kernel(


def _grouped_gemm_mxfp8(
a_q: torch.Tensor, # [M, K] fp8 e4m3
a_scale: torch.Tensor, # [M, K//32] uint8 (E8M0)
w: torch.Tensor, # [E, N, K] fp8 e4m3
w_scale: torch.Tensor, # [E, N, K//32] uint8 (E8M0)
a_q: torch.Tensor, # [M, K] fp8 e4m3
a_scale: torch.Tensor, # [M, K//32] uint8 (E8M0)
w: torch.Tensor, # [E, N, K] fp8 e4m3
w_scale: torch.Tensor, # [E, N, K//32] uint8 (E8M0)
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor,
Expand All @@ -140,9 +140,6 @@ 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
Expand All @@ -152,8 +149,34 @@ 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_K = 128
grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, block_n))

# Tuned launch tiles for gfx950 at MiniMax-M3 MTP decode shapes (conc=32,
# T=128 tokens, E=128 experts, top-4 routing, block_m=16).
# GEMM1 (gate_up, K=6144->N=1536): BLOCK_N=32 + BLOCK_K=512 + waves=4 wins.
# GEMM2 (down, K=768->N=6144): BLOCK_N=64 + BLOCK_K=256 wins.
# Prefill (block_m=64) keeps the wider 128x256 tile.
WAVES = 2
if block_m <= 16 and K == 6144 and N == 1536: # GEMM1 gate_up (decode)
BLOCK_N, BLOCK_K, NUM_WARPS, WAVES = 32, 512, 4, 4
elif block_m <= 16 and K == 768 and N == 6144: # GEMM2 down (decode)
BLOCK_N, BLOCK_K, NUM_WARPS, WAVES = 64, 256, 4, 0
else:
BLOCK_N = 128
BLOCK_K = 256
NUM_WARPS = 4 if block_m <= 16 else 8

grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, BLOCK_N))
_launch_kw: dict = dict(
A_DIV=a_div,
MUL_WEIGHT=mul_weight_by is not None,
BLOCK_M=block_m,
BLOCK_N=BLOCK_N,
BLOCK_K=BLOCK_K,
num_warps=NUM_WARPS,
num_stages=2,
)
if WAVES > 0:
_launch_kw["waves_per_eu"] = WAVES
_mxfp8_grouped_gemm_kernel[grid](
a_q,
a_scale,
Expand All @@ -180,40 +203,19 @@ def _grouped_gemm_mxfp8(
w_scale.stride(2),
out.stride(0),
out.stride(1),
A_DIV=a_div,
MUL_WEIGHT=mul_weight_by is not None,
BLOCK_M=block_m,
BLOCK_N=block_n,
BLOCK_K=BLOCK_K,
num_warps=num_warps,
num_stages=num_stages,
**_launch_kw,
)
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
w13_scale: torch.Tensor, # [E, 2I, H//32] uint8
w2: torch.Tensor, # [E, H, I] fp8
w2_scale: torch.Tensor, # [E, H, I//32] uint8
topk_weights: torch.Tensor, # [T, top_k]
topk_ids: torch.Tensor, # [T, top_k] (global expert ids)
w13: torch.Tensor, # [E, 2I, H] fp8
w13_scale: torch.Tensor, # [E, 2I, H//32] uint8
w2: torch.Tensor, # [E, H, I] fp8
w2_scale: torch.Tensor, # [E, H, I//32] uint8
topk_weights: torch.Tensor, # [T, top_k]
topk_ids: torch.Tensor, # [T, top_k] (global expert ids)
*,
alpha: float,
beta: float,
Expand All @@ -225,8 +227,11 @@ def fused_moe_mxfp8_native(
top_k = topk_ids.shape[1]
M = T * top_k

tiles = _mxfp8_moe_tiles(T)
block_m = tiles["block_m"]
# Adaptive block_m: at decode (small M), tokens_per_expert ≈ 1-2 and
# block_m=64 wastes >95% of tile rows. block_m=16 cuts the sorted-id /
# grid overhead significantly. Prefill keeps 64 for compute efficiency.
tokens_per_expert = max(1, M // max(1, global_num_experts))
block_m = 16 if tokens_per_expert <= 16 else 64
sorted_ids, expert_ids, num_post = moe_align_block_size(
topk_ids,
block_m,
Expand All @@ -235,7 +240,7 @@ def fused_moe_mxfp8_native(
ignore_invalid_experts=expert_map is not None,
)

# GEMM1: x (mxfp8) @ w13^T -> [M, 2I]
# GEMM1: hidden_states (mxfp8) @ w13^T -> [M, 2I]
a_q, a_s = mxfp8_e4m3_quantize(hidden_states)
g1 = _grouped_gemm_mxfp8(
a_q,
Expand All @@ -251,18 +256,10 @@ 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
# GEMM2 MXFP8 activation-quant in one fp32 Triton pass — no bf16 ``act``
# round-trip to HBM. Bit-exact vs the unfused swiglu+quant chain on measured
# MoE shapes, and ~1.2-1.9x faster on that step in isolation. (Not the #22
# ``silu_and_mul_with_clamp`` op: it rounds intermediates to bf16, rel ~3e-3.)
# Lazy import: the amd.ops package pulls in the minimax_m3 platform dispatch,
# only resolvable after the model module finishes loading.
# SwiGLU-OAI activation (split layout: gate=g1[:, :I], up=g1[:, I:])
# followed by MXFP8 quantization before GEMM2.
from vllm.models.minimax_m3.amd.ops import swiglu_oai_quantize_mxfp8

# GEMM2: act (mxfp8) @ w2^T -> [M, H], weighted by topk_weights, then reduce.
Expand All @@ -282,9 +279,6 @@ 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)
Expand Down
Loading