Skip to content

[FEAT] [ROCm] Upgrade AITER Fused MoE kernels. #18271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
May 27, 2025
Merged

[FEAT] [ROCm] Upgrade AITER Fused MoE kernels. #18271

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
a989ce0
use single AITER fmoe module
vllmellm May 16, 2025
da0622c
bugfixes: remove the weight scale expension and set correct layout size
vllmellm May 16, 2025
350bf88
clean code in aiter fmoe module
vllmellm May 16, 2025
2b40c72
udpate docker file for new AITER package
vllmellm May 16, 2025
85a7151
update shuffle weights documentation.
vllmellm May 19, 2025
fbaa890
Merge remote-tracking branch 'origin/main' into update-aiter-fmoe
vllmellm May 21, 2025
33e36d6
only enable per tensor quantization for fp8 w8a8
vllmellm May 22, 2025
aa9e31d
fix precommit error
vllmellm May 22, 2025
ee9506f
fix precommit error
vllmellm May 22, 2025
55e11f7
avoid unnecessary multiplication of hidden_states and topk_weights in…
vllmellm May 22, 2025
730fb77
remove unnecessary layout argument
vllmellm May 23, 2025
11fb95d
Merge remote-tracking branch 'origin/main' into update-aiter-fmoe
vllmellm May 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docker/Dockerfile.rocm_base
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ ARG PYTORCH_REPO="https://github.com/pytorch/pytorch.git"
ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
ARG FA_BRANCH="1a7f4dfa"
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG AITER_BRANCH="5a77249"
ARG AITER_BRANCH="c1debd8"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"

FROM ${BASE_IMAGE} AS base
Expand Down
6 changes: 2 additions & 4 deletions vllm/model_executor/layers/fused_moe/layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -375,10 +375,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
shuffle_weights)

if self.rocm_aiter_moe_enabled:
# use 2stage ck moe layout
shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,
layer.w2_weight.data,
layout=(32, 32))
shuffled_w13, shuffled_w2 = shuffle_weights(
layer.w13_weight.data, layer.w2_weight.data)

layer.w13_weight.data = shuffled_w13
layer.w2_weight.data = shuffled_w2
Expand Down
Loading