A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 10:00 UTC
[ROCm] Enable MXFP4 MoE weight pre-shuffling on gfx950 and update aiter (#34192)
Signed-off-by: Doug Lehr <douglehr@amd.com> Co-authored-by: Doug Lehr <douglehr@amd.com> Co-authored-by: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com> Co-authored-by: tjtanaavllm <tunjian.tan@amd.com>
Douglas Lehr committed
Feb 12, 2026 at 07:06 UTC
8a798be929d62a6467fd079c03c83632f8231b11
2 files changed
+11
-3
docker/Dockerfile.rocm_base
+3
-3
@@ -1,5 +1,5 @@
1
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
2
-ARG TRITON_BRANCH="f332c492"
2
+ARG TRITON_BRANCH="57c693b6"
3
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
4
ARG PYTORCH_BRANCH="89075173"
5
ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
9
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
10
ARG FA_BRANCH="0e60e394"
11
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12
-ARG AITER_BRANCH="6af8b687"
12
+ARG AITER_BRANCH="v0.1.10.post2"
13
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
ARG MORI_BRANCH="2d02c6a9"
15
ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -239,7 +239,7 @@ RUN pip install pyyaml && cd aiter \
239
export HIP_CLANG_PATH=/opt/sccache-wrappers \
240
&& sccache --show-stats; \
241
fi \
242
- && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
242
+ && GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
243
&& if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
244
&& ls /app/aiter/dist/*.whl
245
RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
vllm/model_executor/layers/quantization/quark/quark_moe.py
+8
@@ -933,7 +933,15 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
933
layer.w2_weight.view(self.fp4_dtype),
934
requires_grad=layer.w2_weight.requires_grad,
935
)
936
+ # Pre-shuffle weight
937
+ shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
938
+ layer.w13_weight.data, layer.w2_weight.data
939
+ )
940
941
+ layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)
942
+ layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
943
+ layer.w13_weight.is_shuffled = True
944
+ layer.w2_weight.is_shuffled = True
945
torch.cuda.empty_cache()
946
947
def get_fused_moe_quant_config(