iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][ROCm] Enable AITER QuickReduce + RMSNorm fusion (#48249)

Browse Source 
Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Mikko Tukiainen
2026-09-17 15:05:48 +00:00
committed by GitHub
co-authored by Cursor mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 4f9a41ac94
commit 1df336c3ba
2 changed files with 278 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/distributed/test_quick_all_reduce.py
+172
View File
@@ -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,
vllm/_aiter_ops.py
+106 -1
View File
@@ -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,