[ROCm][MiniMax-M2] Dispatch fused QK-norm + AllReduce via AITER (#44849)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Co-authored-by: Pawel Kowalski <pawel.kowalski@amd.com>
akii96 committed
Jul 14, 2026 at 06:10 UTC
b50ef9c6ed978206bd8cfca0e8f9fb1202813f88
1 file changed
+12
vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py
+12
@@ -6,6 +6,7 @@ from functools import partial
6
import torch
7
from torch import nn
8
9
+from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
11
from vllm.distributed.parallel_state import (
12
get_tensor_model_parallel_rank,
@@ -235,6 +236,17 @@ def _minimax_qk_norm_fusion(
236
tp_world,
237
eps,
238
)
239
+ # ROCm AITER fused QK-norm + AllReduce; skips unaligned shapes.
240
+ if tp_world > 1 and num_tokens <= MINIMAX_QK_NORM_MAX_TOKEN_NUM:
241
+ pack_size = 16 // qkv.element_size()
242
+ warp_work_size = 32 * pack_size
243
+ if q_size % warp_work_size == 0 and kv_size % warp_work_size == 0:
244
+ aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
245
+ if aiter_ar is not None:
246
+ q_out, k_out, _ = aiter_ar.aiter_ca.custom_fused_qknorm_ar(
247
+ qkv, q_weight, k_weight, eps
248
+ )
249
+ return q_out, k_out
250
return _minimax_qk_norm_tp_fallback(
251
qkv, q_weight, k_weight, q_size, kv_size, tp_rank, tp_world, eps
252
)