From be7a03ea65fa24d072798c6d388f3d62fdfe4698 Mon Sep 17 00:00:00 2001 From: akii96 Date: Fri, 15 May 2026 20:44:38 +0300 Subject: [PATCH] [ROCm] Widen AITER fused AR RMSNorm 1-stage gate (#42409) Signed-off-by: Aakif Nawaz --- vllm/_aiter_ops.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 7ac356a927b..a144d05eb44 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -744,7 +744,11 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] token_num = input_.shape[0] - hidden_ok = hidden_dim in (512, 1024, 2048, 4096, 7168) + if input_.dtype in (torch.bfloat16, torch.float16): + pack_size = 16 // input_.element_size() + hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 + else: + hidden_ok = False token_ok = token_num <= 80 world_size = aiter_ar.world_size full_nvlink = aiter_ar.fully_connected