A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 14:04 UTC
[ROCm] Widen AITER fused AR RMSNorm 1-stage gate (#42409)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
May 15, 2026 at 20:44 UTC
be7a03ea65fa24d072798c6d388f3d62fdfe4698
1 file changed
+5
-1
vllm/_aiter_ops.py
+5
-1
@@ -744,7 +744,11 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
744
total_bytes = input_.numel() * input_.element_size()
745
hidden_dim = input_.shape[-1]
746
token_num = input_.shape[0]
747
- hidden_ok = hidden_dim in (512, 1024, 2048, 4096, 7168)
747
+ if input_.dtype in (torch.bfloat16, torch.float16):
748
+ pack_size = 16 // input_.element_size()
749
+ hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
750
+ else:
751
+ hidden_ok = False
752
token_ok = token_num <= 80
753
world_size = aiter_ar.world_size
754
full_nvlink = aiter_ar.fully_connected