Skip to content

Commit 6383108

Browse files
vllmellmamitm02
authored andcommitted
[FEAT] [ROCm] Upgrade AITER Fused MoE kernels. (vllm-project#18271)
Signed-off-by: vllmellm <[email protected]> Signed-off-by: amit <[email protected]>
1 parent 607a4fc commit 6383108

File tree

4 files changed

+133
-317
lines changed

4 files changed

+133
-317
lines changed

vllm/model_executor/layers/fused_moe/layer.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -419,10 +419,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
419419
shuffle_weights)
420420

421421
if self.rocm_aiter_moe_enabled:
422-
# use 2stage ck moe layout
423-
shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,
424-
layer.w2_weight.data,
425-
layout=(32, 32))
422+
shuffled_w13, shuffled_w2 = shuffle_weights(
423+
layer.w13_weight.data, layer.w2_weight.data)
426424

427425
layer.w13_weight.data = shuffled_w13
428426
layer.w2_weight.data = shuffled_w2

0 commit comments

Comments
 (0)