diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index abc98df80ab..364608a806a 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index 33851fdc862..fa6e902396f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/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)