From 579aef4e8dfa2c9d498a702dd04c5ec5db333266 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Thu, 3 Sep 2026 18:28:40 -0500 Subject: [PATCH] [ROCm] Bump AITER to 0.1.21.post1 (#52826) Signed-off-by: Rohan138 Signed-off-by: Rohan Potdar Signed-off-by: fai Co-authored-by: Claude Opus 4.8 Co-authored-by: Andreas Karatzas Co-authored-by: fai --- docker/Dockerfile.rocm_base | 2 +- tests/kernels/quantization/test_rocm_mxfp4.py | 11 +- .../layers/fused_moe/fused_flydsl_moe.py | 257 ++++++++---------- 3 files changed, 129 insertions(+), 141 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 9191e5382e..1634c1f13b 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.11.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.19" +ARG AITER_BRANCH="v0.1.21.post1" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index 2e497c8cde..e60765eeaf 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -50,7 +50,16 @@ SKINNY_GEMM_PASS_RATES = { } PRESHUFFLED_SHAPES = [ (64, 4096, 8192), - (32, 8192, 8192), + # aiter 0.1.20 Triton preshuffled fp4 GEMM OOBs (GPU memory access fault -> + # SIGABRT, crashing the process) for (M=32, N=8192, K=8192) on gfx950, despite + # advertising the (N, K) as tuned. A hard GPU fault can't be xfail'd, so skip it + # until the aiter kernel is fixed (ROCm/aiter#4867). + pytest.param( + (32, 8192, 8192), + marks=pytest.mark.skip( + reason="aiter 0.1.20 preshuffled fp4 GEMM OOB (M32,N8192,K8192, gfx950)" + ), + ), ] diff --git a/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py index 27c07815cb..2533c3a895 100644 --- a/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py @@ -6,13 +6,11 @@ import functools import json import os -import flydsl.compiler as flyc -import flydsl.expr as fx import torch from aiter.fused_moe import moe_sorting as aiter_moe_sorting -from aiter.ops.flydsl.kernels.moe_gemm_2stage import ( - compile_moe_gemm1, - compile_moe_gemm2, +from aiter.ops.flydsl.kernels.moe_2stage_a16wmix import ( + flydsl_a16w4_gemm1, + flydsl_a16w4_gemm2, ) from vllm.logger import init_logger @@ -21,33 +19,58 @@ from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) -_FLYDSL_MOE_GEMM1_CACHE: dict = {} -_FLYDSL_MOE_GEMM2_CACHE: dict = {} - +# Valid tiles for the aiter a16w-mix int4 kernels (aiter#4646): stage1 requires +# tile_n in {64,128} and tile_k in {128,256} (k_wave=1); stage2 requires tile_n2==128 +# and tile_k2 in {128,256}. The pre-#4646 defaults (tile_k=512/64, tile_n2=256) are no +# longer registered and make the port emit an empty-fragment kernel. _FLYDSL_MOE_DEFAULT_CONFIG = { - 1: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256}, - 2: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128}, - 4: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128}, - 8: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256}, - 16: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 256}, - 24: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, - 32: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, - 48: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 1: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 2: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 4: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 8: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 16: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 24: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 32: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 48: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, 64: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, - 128: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, - 256: {"tile_m": 16, "tile_n": 128, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, - 512: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, - 1024: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, - 2048: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64}, - 4096: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, - 8192: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64}, + 128: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 256: {"tile_m": 16, "tile_n": 128, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 512: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 1024: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 2048: {"tile_m": 64, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 4096: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 8192: {"tile_m": 64, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, } +# aiter#4646's a16w-mix int4 kernels only register a fixed tile set. Tiles from before +# the bump (default config or tuned JSONs) carry tile_k=512/64 and tile_n2=256, which +# emit an empty-fragment kernel (IndexError at emit / wrong output). Snap any resolved +# tiles to the nearest registered value so a stale tuned config can't crash the kernel. +_FLYDSL_INT4_TILE_M = (16, 32, 64, 128) +_FLYDSL_INT4_TILE_N1 = (64, 128) +_FLYDSL_INT4_TILE_K = (128, 256) -def _as_flydsl_ptr(tensor: torch.Tensor | None): - """Adapt a torch tensor to the pointer ABI expected by AITER FlyDSL.""" - data_ptr = None if tensor is None else tensor.data_ptr() - return flyc.from_c_void_p(fx.Uint8, data_ptr) + +def _snap_flydsl_int4_tiles(tile_m, tile_n, tile_k, tile_n2, tile_k2): + def _nearest(v, allowed): + return min(allowed, key=lambda a: (abs(a - v), a)) + + snapped = ( + _nearest(tile_m, _FLYDSL_INT4_TILE_M), + _nearest(tile_n, _FLYDSL_INT4_TILE_N1), + _nearest(tile_k, _FLYDSL_INT4_TILE_K), + 128, # stage2 N tile is only registered at 128 + _nearest(tile_k2, _FLYDSL_INT4_TILE_K), + ) + if snapped != (tile_m, tile_n, tile_k, tile_n2, tile_k2): + logger.warning_once( + "FlyDSL int4 MoE tiles %s are not registered post-aiter#4646; " + "using %s. Retune the config for the a16w-mix kernel.", + (tile_m, tile_n, tile_k, tile_n2, tile_k2), + snapped, + scope="local", + ) + return snapped def moe_sorting( @@ -185,6 +208,9 @@ def fused_flydsl_moe_impl( tile_k = tuned_config["tile_k"] tile_n2 = tuned_config["tile_n2"] tile_k2 = tuned_config["tile_k2"] + tile_m, tile_n, tile_k, tile_n2, tile_k2 = _snap_flydsl_int4_tiles( + tile_m, tile_n, tile_k, tile_n2, tile_k2 + ) routing = build_routing_buffers( topk_ids=topk_ids, @@ -202,123 +228,76 @@ def fused_flydsl_moe_impl( blocks, ) = routing - scale_x_1d = torch.empty((0,), device=device, dtype=torch.float32) + # aiter#4646 (v0.1.20) removed compile_moe_gemm{1,2} + the flyc.compile pointer ABI + # and replaced them with the a16w-mix wrappers, which compile+launch internally + # (functools-cached) and take torch tensors directly. w_dtype="int4" is the a16wi4 + # (W4A16) path. The stage-1 intermediate now lives in sorted + # [sorted_size, inter_dim] layout, and topk weighting is applied in stage-2 via + # sorted_weights (the port has no stage-1 doweight). sorted_weights_1d = sorted_weights.view(-1).contiguous() - out_stage1 = torch.empty( - (tokens, topk, inter_dim), device=device, dtype=out_torch_dtype + empty_scale = torch.empty((0,), device=device, dtype=torch.uint8) + cumsum_tensor = num_valid_ids.to(torch.int32).contiguous() + m_indices = sorted_token_ids.to(torch.int32).contiguous() + + # stage1: bf16 A x int4 W1 + SiLU -> bf16 sorted intermediate. + inter_sorted = torch.empty( + (sorted_size, inter_dim), device=device, dtype=torch.bfloat16 + ) + w1_scale_u8 = ( + w1_scale.view(torch.uint8).contiguous().view(-1) + if w1_scale is not None + else empty_scale + ) + flydsl_a16w4_gemm1( + a_bf16=hidden_states.to(torch.bfloat16).contiguous(), + w1_u8=w1.view(torch.uint8).contiguous(), + w1_scale_u8=w1_scale_u8, + sorted_expert_ids=sorted_expert_ids, + cumsum_tensor=cumsum_tensor, + m_indices=m_indices, + inter_sorted_bf16=inter_sorted, + n_tokens=tokens, + NE=num_experts, + D_HIDDEN=model_dim, + D_INTER=inter_dim, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + waves_per_eu=None, + act="silu", + w_dtype="int4", + w_layout="standard", ) - stream = torch.cuda.current_stream() - - key1 = ( - model_dim, - inter_dim, - num_experts, - topk, - in_dtype, - out_dtype, - group_size, - tile_m, - tile_n, - tile_k, - bool(doweight_stage1), - False, + # stage2: bf16 inter x int4 W2 -> weighted reduce into [tokens, model_dim]. + out_stage2 = torch.zeros((tokens, model_dim), device=device, dtype=out_torch_dtype) + w2_scale_u8 = ( + w2_scale.view(torch.uint8).contiguous().view(-1) + if w2_scale is not None + else empty_scale ) - - compiled_exe1 = _FLYDSL_MOE_GEMM1_CACHE.get(key1) - args1 = ( - _as_flydsl_ptr(out_stage1), - _as_flydsl_ptr(hidden_states), - _as_flydsl_ptr(w1), - _as_flydsl_ptr(scale_x_1d), - _as_flydsl_ptr(w1_scale), - _as_flydsl_ptr(sorted_token_ids), - _as_flydsl_ptr(sorted_expert_ids), - _as_flydsl_ptr(sorted_weights_1d), - _as_flydsl_ptr(num_valid_ids), - tokens, - inter_dim, - model_dim, - int(blocks), - stream, + flydsl_a16w4_gemm2( + inter_sorted_bf16=inter_sorted, + w2_u8=w2.view(torch.uint8).contiguous(), + w2_scale_u8=w2_scale_u8, + sorted_expert_ids=sorted_expert_ids, + cumsum_tensor=cumsum_tensor, + sorted_token_ids=sorted_token_ids, + sorted_weights=sorted_weights_1d, + flat_out=out_stage2.view(-1), + M_logical=tokens, + max_sorted=sorted_size, + NE=num_experts, + D_HIDDEN=model_dim, + D_INTER=inter_dim, + topk=topk, + tile_m=tile_m, + tile_n=tile_n2, + tile_k=tile_k2, + waves_per_eu=None, + w_dtype="int4", ) - if compiled_exe1 is None: - exe1 = compile_moe_gemm1( - model_dim=model_dim, - inter_dim=inter_dim, - experts=num_experts, - topk=topk, - in_dtype=in_dtype, - out_dtype=out_dtype, - group_size=group_size, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - doweight_stage1=bool(doweight_stage1), - use_cshuffle_epilog=False, - scale_is_bf16=scale_is_bf16, - ) - compiled_exe1 = flyc.compile(exe1, *args1) - _FLYDSL_MOE_GEMM1_CACHE[key1] = compiled_exe1 - else: - compiled_exe1(*args1) - - a2_1d = out_stage1.view(-1).contiguous() - a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32) - out_stage2 = torch.empty((tokens, model_dim), device=device, dtype=out_torch_dtype) - doweight_stage2 = not bool(doweight_stage1) - - key2 = ( - model_dim, - inter_dim, - num_experts, - topk, - in_dtype, - out_dtype, - group_size, - tile_m, - tile_n2, - tile_k2, - bool(doweight_stage2), - ) - - compiled_exe2 = _FLYDSL_MOE_GEMM2_CACHE.get(key2) - args2 = ( - _as_flydsl_ptr(out_stage2), - _as_flydsl_ptr(a2_1d), - _as_flydsl_ptr(w2), - _as_flydsl_ptr(a2_scale_1d), - _as_flydsl_ptr(w2_scale), - _as_flydsl_ptr(sorted_token_ids), - _as_flydsl_ptr(sorted_expert_ids), - _as_flydsl_ptr(sorted_weights_1d), - _as_flydsl_ptr(num_valid_ids), - tokens, - model_dim, - inter_dim, - int(blocks), - stream, - ) - out_stage2.zero_() - if compiled_exe2 is None: - exe2 = compile_moe_gemm2( - model_dim=model_dim, - inter_dim=inter_dim, - experts=num_experts, - topk=topk, - in_dtype=in_dtype, - out_dtype=out_dtype, - group_size=group_size, - tile_m=tile_m, - tile_n=tile_n2, - tile_k=tile_k2, - doweight_stage2=bool(doweight_stage2), - scale_is_bf16=scale_is_bf16, - ) - compiled_exe2 = flyc.compile(exe2, *args2) - _FLYDSL_MOE_GEMM2_CACHE[key2] = compiled_exe2 - else: - compiled_exe2(*args2) return out_stage2