Skip to content

Commit d260f79

Browse files
authored
[FEAT] [ROCm] Upgrade AITER Fused MoE kernels. (#18271)
Signed-off-by: vllmellm <[email protected]>
1 parent b50602d commit d260f79

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)