From 1df336c3bab8e2ac2eb3d325bd90e62b94656e9c Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Thu, 17 Sep 2026 18:05:48 +0300 Subject: [PATCH] [Perf][ROCm] Enable AITER QuickReduce + RMSNorm fusion (#48249) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_quick_all_reduce.py | 172 +++++++++++++++++++++ vllm/_aiter_ops.py | 107 ++++++++++++- 2 files changed, 278 insertions(+), 1 deletion(-) diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py index f9228fbd68..107d95fa61 100644 --- a/tests/distributed/test_quick_all_reduce.py +++ b/tests/distributed/test_quick_all_reduce.py @@ -3,12 +3,17 @@ import multiprocessing import random +import sys +import types +from types import SimpleNamespace import pytest import ray import torch import torch.distributed as dist +import vllm._aiter_ops as aiter_ops +import vllm.distributed as vllm_distributed import vllm.envs as envs from vllm import _custom_ops as ops from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa @@ -233,6 +238,173 @@ def test_quick_allreduce_passes_dynamic_quant_level( assert called_quant_level == QuickReduceRegime.FP.value +def test_rocm_aiter_fused_rmsnorm_uses_aiter_qr_rmsnorm_for_prefill( + monkeypatch: pytest.MonkeyPatch, +): + class FakeAiterAllReduce: + aiter_ca = SimpleNamespace(world_size=2, fully_connected=True) + + def use_1stage_fused_ar_rms(self, inp): + return False + + def custom_fused_ar_rms(self, *args, **kwargs): + raise AssertionError("expected fused QR+RMSNorm dispatch") + + class FakeQuickReduce: + disabled = False + + def __init__(self): + self.checked_input = None + + def should_quick_allreduce(self, inp): + self.checked_input = inp + return True + + def _get_qr_quant_level(self, inp): + return QuickReduceRegime.INT4.value + + qr_comm = FakeQuickReduce() + device_comm = SimpleNamespace(cpu_group=object(), qr_comm=qr_comm) + monkeypatch.setattr( + vllm_distributed, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_comm), + ) + monkeypatch.setattr( + aiter_ops.rocm_aiter_ops, + "get_aiter_allreduce", + lambda: FakeAiterAllReduce(), + ) + monkeypatch.setattr(aiter_ops.dist, "get_world_size", lambda group: 2) + monkeypatch.setattr(aiter_ops.dist, "get_rank", lambda group: 0) + monkeypatch.setattr( + aiter_ops.dist, + "all_gather_object", + lambda handles, handle, group: handles.__setitem__(0, handle), + ) + + calls = [] + fake_aiter = types.SimpleNamespace( + init_custom_qr=lambda rank, world_size, qr_max_size: 123, + qr_get_handle=lambda ptr: b"handle", + qr_open_handles=lambda ptr, handles: None, + ) + + def fake_qr_all_reduce_rmsnorm( + ptr, + inp, + residual, + residual_out, + out, + weight, + eps, + hidden_dim, + quant_level, + cast_bf2half, + ): + calls.append( + { + "ptr": ptr, + "inp": inp, + "residual": residual, + "residual_out": residual_out, + "out": out, + "weight": weight, + "eps": eps, + "hidden_dim": hidden_dim, + "quant_level": quant_level, + "cast_bf2half": cast_bf2half, + } + ) + out.copy_(inp) + residual_out.copy_(residual) + + fake_aiter.qr_all_reduce_rmsnorm = fake_qr_all_reduce_rmsnorm + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None) + monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", True) + + # token_num > 80 forces vLLM's AITER fused AR+RMS wrapper onto its + # non-1stage prefill path. hidden_dim=4096 bf16 gives an 8 KiB row, so four + # independent RMSNorm rows fit in one 32 KiB QR tile. + inp = torch.randn(81, 4096, dtype=torch.bfloat16) + residual = torch.randn_like(inp) + weight = torch.randn(4096, dtype=torch.bfloat16) + + out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( + inp, residual, weight, 1e-6 + ) + + assert qr_comm.checked_input is inp + assert len(calls) == 1 + assert calls[0]["ptr"] == 123 + assert calls[0]["inp"] is inp + assert calls[0]["residual"] is residual + assert calls[0]["out"] is out + assert calls[0]["residual_out"] is residual_out + assert calls[0]["weight"] is weight + assert calls[0]["hidden_dim"] == 4096 + assert calls[0]["quant_level"] == QuickReduceRegime.INT4.value + assert calls[0]["cast_bf2half"] is True + torch.testing.assert_close(out, inp) + torch.testing.assert_close(residual_out, residual) + + +def test_rocm_aiter_fused_rmsnorm_keeps_1stage_decode_path( + monkeypatch: pytest.MonkeyPatch, +): + calls = [] + + class FakeAiterCA: + world_size = 2 + fully_connected = True + + def custom_fused_ar_rms( + self, inp, residual, weight, epsilon, *, use_1stage, gemma_norm + ): + calls.append( + { + "inp": inp, + "residual": residual, + "weight": weight, + "eps": epsilon, + "use_1stage": use_1stage, + "gemma_norm": gemma_norm, + } + ) + return inp + 1, residual + 1 + + class FakeAiterAllReduce: + aiter_ca = FakeAiterCA() + + def use_1stage_fused_ar_rms(self, inp): + return True + + monkeypatch.setattr( + aiter_ops.rocm_aiter_ops, + "get_aiter_allreduce", + lambda: FakeAiterAllReduce(), + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + inp = torch.randn(4, 4096, dtype=torch.bfloat16) + residual = torch.randn_like(inp) + weight = torch.randn(4096, dtype=torch.bfloat16) + + out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( + inp, residual, weight, 1e-6 + ) + + assert len(calls) == 1 + assert calls[0]["inp"] is inp + assert calls[0]["residual"] is residual + assert calls[0]["weight"] is weight + assert calls[0]["use_1stage"] is True + assert calls[0]["gemma_norm"] is False + torch.testing.assert_close(out, inp + 1) + torch.testing.assert_close(residual_out, residual + 1) + + @ray.remote(num_gpus=1, max_calls=1) def graph_quickreduce( monkeypatch: pytest.MonkeyPatch, diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 101e1dd10a..4b9a145f48 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -4,9 +4,10 @@ import ctypes import functools import os from collections.abc import Callable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch +import torch.distributed as dist from torch._ops import OpOverload import vllm.envs as envs @@ -34,6 +35,63 @@ except ImportError: # which is a host op, so we cache it once here. FP8_DTYPE = current_platform.fp8_dtype() _HIPB_MM_INITIALIZED_DEVICES: set[int] = set() +KB = 1024 +MB = 1024 * KB + + +def _get_or_create_aiter_qr_rmsnorm_comm( + device_comm: Any, +) -> tuple[int, int, bool] | None: + """Create the AITER QR communicator used by fused QR+RMSNorm. + + vLLM's normal QuickReduce communicator owns vLLM C++ QR buffers. The fused + QR+RMSNorm kernel added in AITER needs an AITER DeviceComms instance, so we + lazily create one on the same TP CPU group and keep it on the vLLM device + communicator for reuse. + """ + existing = getattr(device_comm, "_aiter_qr_rmsnorm_comm", None) + if existing is False: + return None + if existing is not None: + return existing + + try: + import aiter + except Exception: + device_comm._aiter_qr_rmsnorm_comm = False + return None + + if not hasattr(aiter, "qr_all_reduce_rmsnorm"): + device_comm._aiter_qr_rmsnorm_comm = False + return None + + group = getattr(device_comm, "cpu_group", None) + if group is None: + device_comm._aiter_qr_rmsnorm_comm = False + return None + + world_size = dist.get_world_size(group=group) + rank = dist.get_rank(group=group) + if world_size not in (2, 4, 8): + device_comm._aiter_qr_rmsnorm_comm = False + return None + + qr_max_size_mb = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB + qr_max_size = None if qr_max_size_mb is None else qr_max_size_mb * MB + + try: + ptr = aiter.init_custom_qr(rank, world_size, qr_max_size) + handle = aiter.qr_get_handle(ptr) + handles = [None] * world_size + dist.all_gather_object(handles, handle, group=group) + aiter.qr_open_handles(ptr, handles) + except Exception: + device_comm._aiter_qr_rmsnorm_comm = False + return None + + comm = (ptr, world_size, envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16) + device_comm._aiter_qr_rmsnorm_comm = comm + return comm def _ensure_hipb_mm_extension_initialized() -> None: @@ -991,6 +1049,53 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( ca = aiter_ar.aiter_ca use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_) + if not use_1stage: + from vllm.distributed import get_tp_group + + device_comm = get_tp_group().device_communicator + qr_comm = getattr(device_comm, "qr_comm", None) + hidden_dim = input_.shape[-1] + row_size = hidden_dim * input_.element_size() + fused_qr_rmsnorm_ok = ( + qr_comm is not None + and not getattr(qr_comm, "disabled", True) + and hasattr(qr_comm, "should_quick_allreduce") + and qr_comm.should_quick_allreduce(input_) + and input_.shape == residual.shape + and input_.dtype in (torch.bfloat16, torch.float16) + and input_.dtype == residual.dtype + and input_.dtype == weight.dtype + and weight.dim() == 1 + and weight.numel() == hidden_dim + and input_.numel() % hidden_dim == 0 + and row_size > 0 + and row_size <= 32 * KB + and (32 * KB) % row_size == 0 + ) + if fused_qr_rmsnorm_ok: + assert qr_comm is not None + aiter_qr_comm = _get_or_create_aiter_qr_rmsnorm_comm(device_comm) + if aiter_qr_comm is not None: + import aiter + + ptr, _, cast_bf2half = aiter_qr_comm + quant_level = qr_comm._get_qr_quant_level(input_) + out = torch.empty_like(input_) + residual_out = torch.empty_like(residual) + aiter.qr_all_reduce_rmsnorm( + ptr, + input_, + residual, + residual_out, + out, + weight, + epsilon, + hidden_dim, + quant_level, + cast_bf2half, + ) + return out, residual_out + result = ca.custom_fused_ar_rms( input_, residual,