[Perf][ROCm] Enable AITER QuickReduce + RMSNorm fusion (#48249)
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>
Mikko Tukiainen committed
Sep 17, 2026 at 18:05 UTC
1df336c3bab8e2ac2eb3d325bd90e62b94656e9c
2 files changed
+278
-1
tests/distributed/test_quick_all_reduce.py
+172
@@ -3,12 +3,17 @@
3
4
import multiprocessing
5
import random
6
+import sys
7
+import types
8
+from types import SimpleNamespace
9
10
import pytest
11
import ray
12
import torch
13
import torch.distributed as dist
14
15
+import vllm._aiter_ops as aiter_ops
16
+import vllm.distributed as vllm_distributed
17
import vllm.envs as envs
18
from vllm import _custom_ops as ops
19
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
@@ -233,6 +238,173 @@ def test_quick_allreduce_passes_dynamic_quant_level(
238
assert called_quant_level == QuickReduceRegime.FP.value
239
240
241
+def test_rocm_aiter_fused_rmsnorm_uses_aiter_qr_rmsnorm_for_prefill(
242
+ monkeypatch: pytest.MonkeyPatch,
243
+):
244
+ class FakeAiterAllReduce:
245
+ aiter_ca = SimpleNamespace(world_size=2, fully_connected=True)
246
+
247
+ def use_1stage_fused_ar_rms(self, inp):
248
+ return False
249
+
250
+ def custom_fused_ar_rms(self, *args, **kwargs):
251
+ raise AssertionError("expected fused QR+RMSNorm dispatch")
252
+
253
+ class FakeQuickReduce:
254
+ disabled = False
255
+
256
+ def __init__(self):
257
+ self.checked_input = None
258
+
259
+ def should_quick_allreduce(self, inp):
260
+ self.checked_input = inp
261
+ return True
262
+
263
+ def _get_qr_quant_level(self, inp):
264
+ return QuickReduceRegime.INT4.value
265
+
266
+ qr_comm = FakeQuickReduce()
267
+ device_comm = SimpleNamespace(cpu_group=object(), qr_comm=qr_comm)
268
+ monkeypatch.setattr(
269
+ vllm_distributed,
270
+ "get_tp_group",
271
+ lambda: SimpleNamespace(device_communicator=device_comm),
272
+ )
273
+ monkeypatch.setattr(
274
+ aiter_ops.rocm_aiter_ops,
275
+ "get_aiter_allreduce",
276
+ lambda: FakeAiterAllReduce(),
277
+ )
278
+ monkeypatch.setattr(aiter_ops.dist, "get_world_size", lambda group: 2)
279
+ monkeypatch.setattr(aiter_ops.dist, "get_rank", lambda group: 0)
280
+ monkeypatch.setattr(
281
+ aiter_ops.dist,
282
+ "all_gather_object",
283
+ lambda handles, handle, group: handles.__setitem__(0, handle),
284
+ )
285
+
286
+ calls = []
287
+ fake_aiter = types.SimpleNamespace(
288
+ init_custom_qr=lambda rank, world_size, qr_max_size: 123,
289
+ qr_get_handle=lambda ptr: b"handle",
290
+ qr_open_handles=lambda ptr, handles: None,
291
+ )
292
+
293
+ def fake_qr_all_reduce_rmsnorm(
294
+ ptr,
295
+ inp,
296
+ residual,
297
+ residual_out,
298
+ out,
299
+ weight,
300
+ eps,
301
+ hidden_dim,
302
+ quant_level,
303
+ cast_bf2half,
304
+ ):
305
+ calls.append(
306
+ {
307
+ "ptr": ptr,
308
+ "inp": inp,
309
+ "residual": residual,
310
+ "residual_out": residual_out,
311
+ "out": out,
312
+ "weight": weight,
313
+ "eps": eps,
314
+ "hidden_dim": hidden_dim,
315
+ "quant_level": quant_level,
316
+ "cast_bf2half": cast_bf2half,
317
+ }
318
+ )
319
+ out.copy_(inp)
320
+ residual_out.copy_(residual)
321
+
322
+ fake_aiter.qr_all_reduce_rmsnorm = fake_qr_all_reduce_rmsnorm
323
+ monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
324
+ monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None)
325
+ monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", True)
326
+
327
+ # token_num > 80 forces vLLM's AITER fused AR+RMS wrapper onto its
328
+ # non-1stage prefill path. hidden_dim=4096 bf16 gives an 8 KiB row, so four
329
+ # independent RMSNorm rows fit in one 32 KiB QR tile.
330
+ inp = torch.randn(81, 4096, dtype=torch.bfloat16)
331
+ residual = torch.randn_like(inp)
332
+ weight = torch.randn(4096, dtype=torch.bfloat16)
333
+
334
+ out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl(
335
+ inp, residual, weight, 1e-6
336
+ )
337
+
338
+ assert qr_comm.checked_input is inp
339
+ assert len(calls) == 1
340
+ assert calls[0]["ptr"] == 123
341
+ assert calls[0]["inp"] is inp
342
+ assert calls[0]["residual"] is residual
343
+ assert calls[0]["out"] is out
344
+ assert calls[0]["residual_out"] is residual_out
345
+ assert calls[0]["weight"] is weight
346
+ assert calls[0]["hidden_dim"] == 4096
347
+ assert calls[0]["quant_level"] == QuickReduceRegime.INT4.value
348
+ assert calls[0]["cast_bf2half"] is True
349
+ torch.testing.assert_close(out, inp)
350
+ torch.testing.assert_close(residual_out, residual)
351
+
352
+
353
+def test_rocm_aiter_fused_rmsnorm_keeps_1stage_decode_path(
354
+ monkeypatch: pytest.MonkeyPatch,
355
+):
356
+ calls = []
357
+
358
+ class FakeAiterCA:
359
+ world_size = 2
360
+ fully_connected = True
361
+
362
+ def custom_fused_ar_rms(
363
+ self, inp, residual, weight, epsilon, *, use_1stage, gemma_norm
364
+ ):
365
+ calls.append(
366
+ {
367
+ "inp": inp,
368
+ "residual": residual,
369
+ "weight": weight,
370
+ "eps": epsilon,
371
+ "use_1stage": use_1stage,
372
+ "gemma_norm": gemma_norm,
373
+ }
374
+ )
375
+ return inp + 1, residual + 1
376
+
377
+ class FakeAiterAllReduce:
378
+ aiter_ca = FakeAiterCA()
379
+
380
+ def use_1stage_fused_ar_rms(self, inp):
381
+ return True
382
+
383
+ monkeypatch.setattr(
384
+ aiter_ops.rocm_aiter_ops,
385
+ "get_aiter_allreduce",
386
+ lambda: FakeAiterAllReduce(),
387
+ )
388
+ monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
389
+
390
+ inp = torch.randn(4, 4096, dtype=torch.bfloat16)
391
+ residual = torch.randn_like(inp)
392
+ weight = torch.randn(4096, dtype=torch.bfloat16)
393
+
394
+ out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl(
395
+ inp, residual, weight, 1e-6
396
+ )
397
+
398
+ assert len(calls) == 1
399
+ assert calls[0]["inp"] is inp
400
+ assert calls[0]["residual"] is residual
401
+ assert calls[0]["weight"] is weight
402
+ assert calls[0]["use_1stage"] is True
403
+ assert calls[0]["gemma_norm"] is False
404
+ torch.testing.assert_close(out, inp + 1)
405
+ torch.testing.assert_close(residual_out, residual + 1)
406
+
407
+
408
@ray.remote(num_gpus=1, max_calls=1)
409
def graph_quickreduce(
410
monkeypatch: pytest.MonkeyPatch,
vllm/_aiter_ops.py
+106
-1
@@ -4,9 +4,10 @@ import ctypes
4
import functools
5
import os
6
from collections.abc import Callable
7
-from typing import TYPE_CHECKING
7
+from typing import TYPE_CHECKING, Any
8
9
import torch
10
+import torch.distributed as dist
11
from torch._ops import OpOverload
12
13
import vllm.envs as envs
@@ -34,6 +35,63 @@ except ImportError:
35
# which is a host op, so we cache it once here.
36
FP8_DTYPE = current_platform.fp8_dtype()
37
_HIPB_MM_INITIALIZED_DEVICES: set[int] = set()
38
+KB = 1024
39
+MB = 1024 * KB
40
+
41
+
42
+def _get_or_create_aiter_qr_rmsnorm_comm(
43
+ device_comm: Any,
44
+) -> tuple[int, int, bool] | None:
45
+ """Create the AITER QR communicator used by fused QR+RMSNorm.
46
+
47
+ vLLM's normal QuickReduce communicator owns vLLM C++ QR buffers. The fused
48
+ QR+RMSNorm kernel added in AITER needs an AITER DeviceComms instance, so we
49
+ lazily create one on the same TP CPU group and keep it on the vLLM device
50
+ communicator for reuse.
51
+ """
52
+ existing = getattr(device_comm, "_aiter_qr_rmsnorm_comm", None)
53
+ if existing is False:
54
+ return None
55
+ if existing is not None:
56
+ return existing
57
+
58
+ try:
59
+ import aiter
60
+ except Exception:
61
+ device_comm._aiter_qr_rmsnorm_comm = False
62
+ return None
63
+
64
+ if not hasattr(aiter, "qr_all_reduce_rmsnorm"):
65
+ device_comm._aiter_qr_rmsnorm_comm = False
66
+ return None
67
+
68
+ group = getattr(device_comm, "cpu_group", None)
69
+ if group is None:
70
+ device_comm._aiter_qr_rmsnorm_comm = False
71
+ return None
72
+
73
+ world_size = dist.get_world_size(group=group)
74
+ rank = dist.get_rank(group=group)
75
+ if world_size not in (2, 4, 8):
76
+ device_comm._aiter_qr_rmsnorm_comm = False
77
+ return None
78
+
79
+ qr_max_size_mb = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
80
+ qr_max_size = None if qr_max_size_mb is None else qr_max_size_mb * MB
81
+
82
+ try:
83
+ ptr = aiter.init_custom_qr(rank, world_size, qr_max_size)
84
+ handle = aiter.qr_get_handle(ptr)
85
+ handles = [None] * world_size
86
+ dist.all_gather_object(handles, handle, group=group)
87
+ aiter.qr_open_handles(ptr, handles)
88
+ except Exception:
89
+ device_comm._aiter_qr_rmsnorm_comm = False
90
+ return None
91
+
92
+ comm = (ptr, world_size, envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16)
93
+ device_comm._aiter_qr_rmsnorm_comm = comm
94
+ return comm
95
96
97
def _ensure_hipb_mm_extension_initialized() -> None:
@@ -991,6 +1049,53 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
1049
ca = aiter_ar.aiter_ca
1050
use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
1051
1052
+ if not use_1stage:
1053
+ from vllm.distributed import get_tp_group
1054
+
1055
+ device_comm = get_tp_group().device_communicator
1056
+ qr_comm = getattr(device_comm, "qr_comm", None)
1057
+ hidden_dim = input_.shape[-1]
1058
+ row_size = hidden_dim * input_.element_size()
1059
+ fused_qr_rmsnorm_ok = (
1060
+ qr_comm is not None
1061
+ and not getattr(qr_comm, "disabled", True)
1062
+ and hasattr(qr_comm, "should_quick_allreduce")
1063
+ and qr_comm.should_quick_allreduce(input_)
1064
+ and input_.shape == residual.shape
1065
+ and input_.dtype in (torch.bfloat16, torch.float16)
1066
+ and input_.dtype == residual.dtype
1067
+ and input_.dtype == weight.dtype
1068
+ and weight.dim() == 1
1069
+ and weight.numel() == hidden_dim
1070
+ and input_.numel() % hidden_dim == 0
1071
+ and row_size > 0
1072
+ and row_size <= 32 * KB
1073
+ and (32 * KB) % row_size == 0
1074
+ )
1075
+ if fused_qr_rmsnorm_ok:
1076
+ assert qr_comm is not None
1077
+ aiter_qr_comm = _get_or_create_aiter_qr_rmsnorm_comm(device_comm)
1078
+ if aiter_qr_comm is not None:
1079
+ import aiter
1080
+
1081
+ ptr, _, cast_bf2half = aiter_qr_comm
1082
+ quant_level = qr_comm._get_qr_quant_level(input_)
1083
+ out = torch.empty_like(input_)
1084
+ residual_out = torch.empty_like(residual)
1085
+ aiter.qr_all_reduce_rmsnorm(
1086
+ ptr,
1087
+ input_,
1088
+ residual,
1089
+ residual_out,
1090
+ out,
1091
+ weight,
1092
+ epsilon,
1093
+ hidden_dim,
1094
+ quant_level,
1095
+ cast_bf2half,
1096
+ )
1097
+ return out, residual_out
1098
+
1099
result = ca.custom_fused_ar_rms(
1100
input_,
1101
residual,