From 31afa4f1bb1118b7cd83ad8041777cd4dff92a34 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sun, 27 Sep 2026 19:04:58 -0500 Subject: [PATCH] [ROCm] Bump AITER to v0.1.23 (#58867) Signed-off-by: fai Co-authored-by: Cursor Co-authored-by: Andreas Karatzas --- docker/Dockerfile.rock_base | 2 +- docker/Dockerfile.rocm_base | 2 +- tests/kernels/moe/test_ocp_mx_moe.py | 15 +++------------ .../fused_moe/experts/aiter_mxfp4_w4a16_moe.py | 10 +++++++++- 4 files changed, 14 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base index 79896933e7..d3ba0bc842 100644 --- a/docker/Dockerfile.rock_base +++ b/docker/Dockerfile.rock_base @@ -8,7 +8,7 @@ ARG TRITON_VERSION=3.8.0+gitc0a142ff ARG TRITON_REPO="https://github.com/ROCm/triton.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.22.post1" +ARG AITER_BRANCH="v0.1.23" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="f95cd5e0ca7b730f92f9ffdfb358484d6965d9f7" # TheRock compatible v1.1.0 ARG MORI_REPO="https://github.com/ROCm/mori.git" diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 750d4ce4d9..1c59968005 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.22.post1" +ARG AITER_BRANCH="v0.1.23" 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/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index e4b11c9548..295a92482f 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -47,7 +47,6 @@ if ROCM_AVAILABLE: ROCM_GFX950 = on_gfx950() if ROCM_AITER_AVAILABLE: - from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp from aiter.ops.triton.quant import dynamic_mxfp4_quant if TRTLLM_GEN_MXFP4_AVAILABLE: @@ -191,11 +190,7 @@ def mxfp8_dequantize(x, scale): def mxfp4_quant_dequant(x: torch.Tensor) -> torch.Tensor: shape = x.shape quantized, scale = dynamic_mxfp4_quant(x.to(torch.bfloat16).flatten(0, -2)) - return ( - upcast_from_mxfp(quantized.view(torch.uint8), scale, torch.bfloat16, axis=-1) - .reshape(shape) - .float() - ) + return mxfp4_dequantize(quantized.view(torch.uint8), scale).reshape(shape) def reference_moe( @@ -1809,12 +1804,8 @@ def test_rocm_mxfp4_moe_oracle( assert out.abs().max() > 0.01, "Output is effectively zero" # Dequantize weights for reference computation - w13_dq = upcast_from_mxfp( - w13_quant_ref.view(torch.uint8), w13_scale_ref, torch.bfloat16, axis=-1 - ) - w2_dq = upcast_from_mxfp( - w2_quant_ref.view(torch.uint8), w2_scale_ref, torch.bfloat16, axis=-1 - ) + w13_dq = mxfp4_dequantize(w13_quant_ref.view(torch.uint8), w13_scale_ref) + w2_dq = mxfp4_dequantize(w2_quant_ref.view(torch.uint8), w2_scale_ref) # Determine activation type and layout # SWIGLUOAI uses interleaved layout (gate/up alternating) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index fb3256e1a4..13c34698c8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py @@ -139,7 +139,7 @@ def aiter_triton_kernel_w4a16_moe_forward( from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( should_use_cdna4_mx_scale_swizzle, ) - from vllm.platforms.rocm import on_gfx1250 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 try: from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 @@ -236,6 +236,12 @@ def aiter_triton_kernel_w4a16_moe_forward( # kernel indexes a swizzled buffer as if it were linear. swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None + # AITER's gfx950 Gluon a16w4 kernel (default backend="None"/auto-detect + # picks it) calls a Triton-Gluon AMD API that isn't available in vLLM's + # pinned Triton, so force the Triton backend on gfx950. gfx1250 already + # used Gluon before this and is unaffected, so leave it on auto-detect. + a16w4_backend = "triton" if on_gfx950() else None + intermediate = moe_gemm_a16w4( hidden_states, w1_data, @@ -254,6 +260,7 @@ def aiter_triton_kernel_w4a16_moe_forward( swiglu_add_residual=swiglu_add_residual, unpadded_N=unpadded_N_w1, unpadded_K=unpadded_K_w1, + backend=a16w4_backend, ) out = moe_gemm_a16w4( @@ -270,6 +277,7 @@ def aiter_triton_kernel_w4a16_moe_forward( swizzle_mx_scale=swz, unpadded_N=unpadded_N_w2, unpadded_K=unpadded_K_w2, + backend=a16w4_backend, ) return out