karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[ROCm] Widen AITER fused AR RMSNorm 1-stage gate (#42409)

Browse Source 
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
This commit is contained in:
akii96
2026-05-15 17:44:38 +00:00
committed by GitHub
parent 6147c70224
commit be7a03ea65
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/_aiter_ops.py
+5 -1
View File
@@ -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