[ROCm] Bump AITER to v0.1.23 (#58867)

Signed-off-by: fai <fangzhouai@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

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"

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"

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)

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