[ROCm][CI] add Aiter ops tests (#52208)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Aug 17, 2026 at 19:33 UTC
49fb2ee3481246ce88224220442c6ef30f88b209
3 files changed
+672
.buildkite/test-amd.yaml
+18
@@ -3404,6 +3404,24 @@ steps:
3404
3405
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
3406
3407
+- label: Kernels Core Operation Test %N # TBD
3408
+ timeout_in_minutes: 180
3409
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3410
+ dind: false
3411
+ agent_pool: mi355_1
3412
+ parallelism: 3
3413
+ working_dir: "/vllm-workspace/tests"
3414
+ source_file_dependencies:
3415
+ - csrc/
3416
+ - tests/kernels/core
3417
+ - tests/kernels/test_top_k_per_row.py
3418
+ - tests/kernels/test_concat_mla_q.py
3419
+ - vllm/model_executor/layers/rotary_embedding/
3420
+ - vllm/_aiter_ops.py
3421
+ - vllm/platforms/rocm.py
3422
+ commands:
3423
+ - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3424
+
3425
- label: Kernels (B200-MI355) # TBD
3426
timeout_in_minutes: 180
3427
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
tests/kernels/core/test_rocm_aiter_ops.py
new
+625
@@ -0,0 +1,625 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm AITER helper-op tests for CI stability.
4
+
5
+These tests directly validate AITER kernel correctness against PyTorch/reference
6
+implementations. They provide clear failure signals when AITER library updates
7
+introduce regressions, independent of vLLM's fusion/compilation passes.
8
+
9
+Tested ops:
10
+- ``rms_norm`` and ``rms_norm2d_with_add`` vs PyTorch reference
11
+- ``triton_rotary_embedding`` vs manual NeoX RoPE reference (xfail - known issue)
12
+- ``act_mul_and_fp8_group_quant`` (SiGLU + FP8 group quant)
13
+- fused RMSNorm + quantization ops vs sequential composition
14
+
15
+Related coverage:
16
+- per-token/per-tensor quant roundtrips: ``tests/rocm/aiter/test_quant_op_schema.py``
17
+- RMSNorm determinism: ``tests/rocm/aiter/test_quant_op_schema.py``
18
+- group_fp8_quant: ``tests/kernels/quantization/test_rocm_aiter_grouped_quant.py``
19
+"""
20
+
21
+import warnings
22
+
23
+import pytest
24
+import torch
25
+
26
+import vllm._aiter_ops # noqa: F401 - ensure ops are registered
27
+from tests.kernels.utils import bf16_ulp_distance
28
+from vllm.platforms import current_platform
29
+from vllm.utils.torch_utils import set_random_seed
30
+
31
+pytestmark = pytest.mark.skipif(
32
+ not current_platform.is_rocm(), reason="ROCm-specific tests"
33
+)
34
+
35
+
36
+@pytest.fixture(autouse=True)
37
+def setup_cuda_device():
38
+ """Set default device to CUDA and seed RNG for all tests."""
39
+ torch.set_default_device("cuda")
40
+ set_random_seed(0)
41
+
42
+
43
+def require_aiter():
44
+ from vllm._aiter_ops import is_aiter_found_and_supported
45
+
46
+ if not is_aiter_found_and_supported():
47
+ pytest.skip("aiter is not found or not supported on this hardware")
48
+
49
+
50
+def require_fp8():
51
+ if not current_platform.supports_fp8():
52
+ pytest.skip("FP8 is not supported on this hardware")
53
+
54
+
55
+def _format_observed_rate(count: int, total: int) -> str:
56
+ return f"{count / total:.4%} ({count}/{total})"
57
+
58
+
59
+def _format_allowed_rate(rate: float, total: int) -> str:
60
+ allowed_count = int(rate * total)
61
+ return f"{rate:.4%} (<= {allowed_count}/{total})"
62
+
63
+
64
+def _quantile(values: torch.Tensor, q: float) -> float:
65
+ if values.numel() == 0:
66
+ return 0.0
67
+ return torch.quantile(values, q).item()
68
+
69
+
70
+def _dequantize_grouped(
71
+ quantized: torch.Tensor, scales: torch.Tensor, group_size: int
72
+) -> torch.Tensor:
73
+ """Expand per-group scales and dequantize to float32."""
74
+ M, N = quantized.shape
75
+ scales_expanded = scales.repeat_interleave(group_size, dim=1)[:, :N]
76
+ return quantized.float() * scales_expanded
77
+
78
+
79
+def _rms_norm_reference(
80
+ x: torch.Tensor, weight: torch.Tensor, eps: float
81
+) -> torch.Tensor:
82
+ """RMSNorm reference implementation in float32."""
83
+ rms = x.float().pow(2).mean(-1, keepdim=True).add(eps).sqrt()
84
+ return x.float() / rms * weight.float()
85
+
86
+
87
+def _assert_quant_shapes(
88
+ x_quant: torch.Tensor,
89
+ scale: torch.Tensor,
90
+ M: int,
91
+ N: int,
92
+ fp8_dtype: torch.dtype,
93
+ num_groups: int | None = None,
94
+) -> None:
95
+ """Verify quantized output shapes and dtypes."""
96
+ assert x_quant.shape == (M, N)
97
+ assert x_quant.dtype == fp8_dtype
98
+ if num_groups is not None:
99
+ assert scale.shape == (M, num_groups)
100
+ else:
101
+ assert scale.shape == (M, 1)
102
+
103
+
104
+def _assert_close_budget(
105
+ actual: torch.Tensor,
106
+ expected: torch.Tensor,
107
+ *,
108
+ label: str,
109
+ atol: float,
110
+ rtol: float = 0.0,
111
+ pass_rate: float = 0.99999,
112
+ max_violation_factor: float = 3.0,
113
+) -> None:
114
+ actual_f = actual.detach().float().flatten()
115
+ expected_f = expected.detach().float().flatten()
116
+ abs_diff = (actual_f - expected_f).abs()
117
+ allowed = atol + rtol * expected_f.abs()
118
+
119
+ total = abs_diff.numel()
120
+ passed = int((abs_diff <= allowed).sum().item())
121
+ failed = total - passed
122
+ allowed_fail_rate = 1.0 - pass_rate
123
+ max_abs = abs_diff.max().item()
124
+ mean_abs = abs_diff.mean().item()
125
+ p99_abs = _quantile(abs_diff, 0.99)
126
+ p999_abs = _quantile(abs_diff, 0.999)
127
+ worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item()
128
+ max_atol = max_violation_factor * atol
129
+ above_max_count = int((abs_diff > max_atol).sum().item())
130
+
131
+ msg = (
132
+ "[rocm_aiter_ops] "
133
+ f"{label}: "
134
+ f"pass={passed / total:.4%} ({passed}/{total}) "
135
+ f"fail={_format_observed_rate(failed, total)} "
136
+ f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
137
+ f"atol={atol:g} "
138
+ f"rtol={rtol:g} "
139
+ f"abs>{max_atol:g}={_format_observed_rate(above_max_count, total)} "
140
+ f"allowed_above_max={_format_allowed_rate(0.0, total)} "
141
+ f"max_abs={max_abs:.6g} "
142
+ f"mean_abs={mean_abs:.6g} "
143
+ f"p99_abs={p99_abs:.6g} "
144
+ f"p999_abs={p999_abs:.6g} "
145
+ f"worst_ratio={worst_ratio:.6g}"
146
+ )
147
+ print(msg)
148
+ if failed > 0:
149
+ warnings.warn(msg, stacklevel=2)
150
+
151
+ assert passed / total >= pass_rate, msg
152
+ assert max_abs <= max_atol, msg
153
+ assert mean_abs <= atol * 0.25, msg
154
+
155
+
156
+def _assert_rel_error_quality(
157
+ actual: torch.Tensor,
158
+ expected: torch.Tensor,
159
+ *,
160
+ label: str,
161
+ mean_limit: float,
162
+ preferred_rel: float,
163
+ max_rel: float,
164
+ max_fail_rate: float,
165
+) -> None:
166
+ rel = (
167
+ (actual.float() - expected.float()).abs()
168
+ / expected.float().abs().clamp_min(1e-5)
169
+ ).flatten()
170
+ total = rel.numel()
171
+ within_preferred_count = int((rel <= preferred_rel).sum().item())
172
+ fail_count = total - within_preferred_count
173
+ above_max_count = int((rel > max_rel).sum().item())
174
+ mean_rel = rel.mean().item()
175
+ max_rel_err = rel.max().item()
176
+ p99 = _quantile(rel, 0.99)
177
+ p999 = _quantile(rel, 0.999)
178
+
179
+ msg = (
180
+ "[rocm_aiter_ops] "
181
+ f"{label}: "
182
+ f"rel<={preferred_rel:g} pass={within_preferred_count / total:.4%} "
183
+ f"({within_preferred_count}/{total}) "
184
+ f"fail={_format_observed_rate(fail_count, total)} "
185
+ f"mean_limit={mean_limit:.4%} "
186
+ f"rel>{max_rel:g}={_format_observed_rate(above_max_count, total)} "
187
+ f"allowed_above_max={_format_allowed_rate(max_fail_rate, total)} "
188
+ f"mean_rel={mean_rel:.6g} "
189
+ f"max_rel={max_rel_err:.6g} "
190
+ f"p99={p99:.6g} "
191
+ f"p999={p999:.6g}"
192
+ )
193
+ print(msg)
194
+ if fail_count > 0:
195
+ warnings.warn(msg, stacklevel=2)
196
+
197
+ assert mean_rel < mean_limit, msg
198
+ assert above_max_count / total <= max_fail_rate, msg
199
+
200
+
201
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
202
+def test_rocm_aiter_rms_norm_vs_torch(dtype):
203
+ """rocm_aiter_rms_norm matches PyTorch manual RMSNorm for float16 and bfloat16."""
204
+ require_aiter()
205
+ from vllm._aiter_ops import rocm_aiter_ops
206
+
207
+ M, N = 8, 128
208
+ x = torch.randn(M, N, dtype=dtype)
209
+ weight = torch.ones(N, dtype=dtype)
210
+ eps = 1e-5
211
+
212
+ # Reference: float32 RMSNorm for precision
213
+ ref = _rms_norm_reference(x, weight, eps).to(dtype)
214
+
215
+ out = rocm_aiter_ops.rms_norm(x, weight, eps)
216
+
217
+ # Allow max 1 ULP difference due to different accumulation order.
218
+ ulp = bf16_ulp_distance(out, ref)
219
+ max_ulp = int(ulp.max().item())
220
+ assert max_ulp <= 1, f"rms_norm dtype={dtype}: max ULP distance {max_ulp} > 1"
221
+
222
+
223
+# -- Numerical accuracy tests for AITER custom ops -------------------------
224
+
225
+
226
+def test_rocm_aiter_rmsnorm_with_add_vs_torch():
227
+ """rocm_aiter_rmsnorm2d_fwd_with_add matches manual residual+RMSNorm reference."""
228
+ require_aiter()
229
+ from vllm._aiter_ops import rocm_aiter_ops
230
+
231
+ M, N = 16, 256
232
+ x = torch.randn(M, N, dtype=torch.bfloat16)
233
+ residual = torch.randn(M, N, dtype=torch.bfloat16)
234
+ weight = torch.ones(N, dtype=torch.bfloat16)
235
+ eps = 1e-5
236
+
237
+ # Reference: add residual, then RMSNorm
238
+ h = x.float() + residual.float()
239
+ rms = h.pow(2).mean(-1, keepdim=True).add(eps).sqrt()
240
+ ref_normed = (h / rms * weight.float()).to(torch.bfloat16)
241
+ ref_residual = h.to(torch.bfloat16)
242
+
243
+ out, res_out = rocm_aiter_ops.rms_norm2d_with_add(x, residual, weight, eps)
244
+
245
+ # Allow max 1 ULP difference due to different accumulation order.
246
+ ulp_normed = bf16_ulp_distance(out, ref_normed)
247
+ max_ulp_normed = int(ulp_normed.max().item())
248
+ assert max_ulp_normed <= 1, (
249
+ f"rmsnorm_with_add normed: max ULP distance {max_ulp_normed} > 1"
250
+ )
251
+
252
+ ulp_residual = bf16_ulp_distance(res_out, ref_residual)
253
+ max_ulp_residual = int(ulp_residual.max().item())
254
+ assert max_ulp_residual <= 1, (
255
+ f"rmsnorm_with_add residual: max ULP distance {max_ulp_residual} > 1"
256
+ )
257
+
258
+
259
+@pytest.mark.xfail(
260
+ strict=True,
261
+ raises=AssertionError,
262
+ reason="AITER Triton RoPE precision issue: https://github.com/ROCm/aiter/issues/4765",
263
+)
264
+def test_rocm_aiter_triton_rotary_embedding_vs_torch():
265
+ """rocm_aiter_triton_rotary_embedding matches manual NeoX-style RoPE reference.
266
+
267
+ The AITER kernel calls rope_cached_thd_positions_offsets_2c_fwd_inplace with
268
+ reuse_freqs_front_part=True. This means it reads only the first head_size//2
269
+ entries of cos and sin from the cache and applies them to both halves of the
270
+ head (pairwise NeoX rotation). The cos_sin_cache must be built with the second
271
+ half mirroring the first so the reference and kernel agree.
272
+ """
273
+ require_aiter()
274
+
275
+ num_tokens = 8
276
+ num_heads = 4
277
+ head_size = 64
278
+ half_dim = head_size // 2
279
+ max_pos = 32
280
+
281
+ positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.long)
282
+ query = torch.randn(num_tokens, num_heads * head_size, dtype=torch.bfloat16)
283
+ key = torch.randn(num_tokens, num_heads * head_size, dtype=torch.bfloat16)
284
+
285
+ # Build cos/sin cache with second half mirroring first half.
286
+ # The kernel uses reuse_freqs_front_part=True: it reads only cos/sin[:, :half_dim]
287
+ # and applies those same frequencies to both the first and second
288
+ # halves of the head.
289
+ cos_half = torch.randn(max_pos, half_dim, dtype=torch.bfloat16)
290
+ sin_half = torch.randn(max_pos, half_dim, dtype=torch.bfloat16)
291
+ cos_cache = torch.cat([cos_half, cos_half], dim=-1) # [max_pos, head_size]
292
+ sin_cache = torch.cat([sin_half, sin_half], dim=-1) # [max_pos, head_size]
293
+ cos_sin_cache = torch.cat([cos_cache, sin_cache], dim=-1) # [max_pos, 2*head_size]
294
+
295
+ # Reference: NeoX-style pairwise rotation with front-half frequencies only.
296
+ # rotate_style=0 (NeoX): [x1*c - x2*s, x2*c + x1*s]
297
+ cos_pos = cos_half[positions] # [num_tokens, half_dim]
298
+ sin_pos = sin_half[positions] # [num_tokens, half_dim]
299
+
300
+ def apply_rope_ref(t: torch.Tensor) -> torch.Tensor:
301
+ t_r = t.float().view(num_tokens, num_heads, head_size)
302
+ c = cos_pos.float().unsqueeze(1) # [num_tokens, 1, half_dim]
303
+ s = sin_pos.float().unsqueeze(1)
304
+ x1, x2 = t_r[..., :half_dim], t_r[..., half_dim:]
305
+ rotated = torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)
306
+ return rotated.to(t.dtype).view(num_tokens, num_heads * head_size)
307
+
308
+ ref_q = apply_rope_ref(query)
309
+ ref_k = apply_rope_ref(key)
310
+
311
+ # AITER in-place RoPE (modifies query/key in-place)
312
+ q_aiter = query.clone()
313
+ k_aiter = key.clone()
314
+ torch.ops.vllm.rocm_aiter_triton_rotary_embedding(
315
+ positions,
316
+ q_aiter,
317
+ k_aiter,
318
+ head_size,
319
+ cos_sin_cache,
320
+ True, # is_neox style -> rotate_style=0
321
+ )
322
+
323
+ _assert_close_budget(
324
+ q_aiter.float(),
325
+ ref_q.float(),
326
+ label="triton_rope query",
327
+ atol=1e-3,
328
+ rtol=1.6e-2,
329
+ )
330
+ _assert_close_budget(
331
+ k_aiter.float(),
332
+ ref_k.float(),
333
+ label="triton_rope key",
334
+ atol=1e-3,
335
+ rtol=1.6e-2,
336
+ )
337
+
338
+
339
+def test_rocm_aiter_act_mul_fp8_group_quant_roundtrip():
340
+ """act_mul_and_fp8_group_quant: dequantized output matches SiLU gate reference."""
341
+ require_aiter()
342
+ require_fp8()
343
+
344
+ M, N = 32, 512 # N even: N//2 gate, N//2 up
345
+ group_size = 128
346
+ x = torch.randn(M, N, dtype=torch.bfloat16)
347
+
348
+ x_quant, scale = torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant(
349
+ x, group_size
350
+ )
351
+
352
+ N_half = N // 2
353
+ # Reference: SiLU(gate) * up
354
+ gate = x.float()[:, :N_half]
355
+ up = x.float()[:, N_half:]
356
+ ref = torch.sigmoid(gate) * gate * up # SiGLU
357
+
358
+ x_dequant = _dequantize_grouped(x_quant, scale, group_size)
359
+
360
+ _assert_rel_error_quality(
361
+ x_dequant,
362
+ ref,
363
+ label="act_mul_fp8_group_quant",
364
+ mean_limit=0.1,
365
+ preferred_rel=0.1,
366
+ max_rel=1.0,
367
+ max_fail_rate=0.01,
368
+ )
369
+
370
+
371
+# -- Fused RMSNorm + quantization accuracy tests ---------------------------
372
+
373
+
374
+def test_rocm_aiter_rmsnorm_fused_dynamic_quant_vs_sequential():
375
+ """Fused RMSNorm+per-token-FP8-quant matches sequential rms_norm->per_token_quant.
376
+
377
+ Tests that the fused kernel produces the same result as the two-step
378
+ sequential composition: rms_norm(x) followed by per_token_quant.
379
+ The fused path is used in production for inference throughput.
380
+ """
381
+ require_aiter()
382
+ require_fp8()
383
+ from vllm._aiter_ops import rocm_aiter_ops
384
+
385
+ M, N = 32, 512
386
+ x = torch.randn(M, N, dtype=torch.bfloat16)
387
+ weight = torch.ones(N, dtype=torch.bfloat16)
388
+ eps = 1e-5
389
+ fp8_dtype = current_platform.fp8_dtype()
390
+
391
+ # Sequential reference
392
+ normed = rocm_aiter_ops.rms_norm(x, weight, eps)
393
+ ref_q, ref_scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype)
394
+ ref_dequant = ref_q.float() * ref_scale.float()
395
+
396
+ # Fused op
397
+ fused_q, fused_scale = torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant(
398
+ x, weight, eps, fp8_dtype
399
+ )
400
+ _assert_quant_shapes(fused_q, fused_scale, M, N, fp8_dtype)
401
+ fused_dequant = fused_q.float() * fused_scale.float()
402
+
403
+ # Fused vs sequential: both should recover RMSNorm output within FP8 error
404
+ _assert_rel_error_quality(
405
+ fused_dequant,
406
+ ref_dequant,
407
+ label="rmsnorm_fused_dynamic_quant",
408
+ mean_limit=0.05,
409
+ preferred_rel=0.05,
410
+ max_rel=0.5,
411
+ max_fail_rate=0.01,
412
+ )
413
+
414
+
415
+def test_rocm_aiter_rmsnorm_fused_add_dynamic_quant_vs_reference():
416
+ """Fused (residual-add + RMSNorm + per-token-FP8-quant)
417
+ matches sequential reference.
418
+
419
+ Production path: input + residual -> RMSNorm -> FP8 quantize, returning
420
+ both the quantized norm output and the residual sum for the next layer.
421
+ """
422
+ require_aiter()
423
+ require_fp8()
424
+ from vllm._aiter_ops import rocm_aiter_ops
425
+
426
+ M, N = 16, 256
427
+ x = torch.randn(M, N, dtype=torch.bfloat16)
428
+ residual = torch.randn(M, N, dtype=torch.bfloat16)
429
+ weight = torch.ones(N, dtype=torch.bfloat16)
430
+ eps = 1e-5
431
+ fp8_dtype = current_platform.fp8_dtype()
432
+
433
+ # Sequential reference: add residual -> rms_norm -> per_token_quant
434
+ h = (x.float() + residual.float()).to(torch.bfloat16)
435
+ normed = rocm_aiter_ops.rms_norm(h, weight, eps)
436
+ ref_q, ref_scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype)
437
+ ref_residual_out = h
438
+
439
+ # Fused op: returns (x_quant, residual_out, scale)
440
+ fused_q, fused_res_out, fused_scale = (
441
+ torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant(
442
+ x, residual, weight, eps, fp8_dtype
443
+ )
444
+ )
445
+ _assert_quant_shapes(fused_q, fused_scale, M, N, fp8_dtype)
446
+ assert fused_res_out.shape == (M, N)
447
+
448
+ torch.testing.assert_close(fused_res_out, ref_residual_out, atol=1e-2, rtol=1.6e-2)
449
+ # Dequantized output matches sequential path
450
+ fused_dequant = fused_q.float() * fused_scale.float()
451
+ ref_dequant = ref_q.float() * ref_scale.float()
452
+ _assert_rel_error_quality(
453
+ fused_dequant,
454
+ ref_dequant,
455
+ label="rmsnorm_fused_add_dynamic_quant",
456
+ mean_limit=0.05,
457
+ preferred_rel=0.05,
458
+ max_rel=0.5,
459
+ max_fail_rate=0.01,
460
+ )
461
+
462
+
463
+def test_rocm_aiter_rmsnorm_fp8_group_quant_vs_sequential():
464
+ """Fused RMSNorm+FP8-group-quant matches sequential rms_norm->group_fp8_quant.
465
+
466
+ Tests both output shapes and dequantized accuracy against the two-step
467
+ sequential composition.
468
+ """
469
+ require_aiter()
470
+ require_fp8()
471
+ from vllm._aiter_ops import rocm_aiter_ops
472
+
473
+ M, N = 32, 512
474
+ group_size = 128
475
+ x = torch.randn(M, N, dtype=torch.bfloat16)
476
+ weight = torch.ones(N, dtype=torch.bfloat16)
477
+ eps = 1e-5
478
+ fp8_dtype = current_platform.fp8_dtype()
479
+ expected_groups = (N + group_size - 1) // group_size
480
+
481
+ # Fused op: (x_quant, scales)
482
+ fused_q, fused_scales = torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant(
483
+ x, weight, eps, group_size
484
+ )
485
+ _assert_quant_shapes(fused_q, fused_scales, M, N, fp8_dtype, expected_groups)
486
+
487
+ # Dequantize and compare to reference: rms_norm -> group quant -> dequant
488
+ normed = rocm_aiter_ops.rms_norm(x, weight, eps)
489
+ ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(normed, group_size)
490
+ ref_dequant = _dequantize_grouped(ref_q, ref_scales, group_size)
491
+ fused_dequant = _dequantize_grouped(fused_q, fused_scales, group_size)
492
+
493
+ _assert_rel_error_quality(
494
+ fused_dequant,
495
+ ref_dequant,
496
+ label="rmsnorm_fp8_group_quant",
497
+ mean_limit=0.05,
498
+ preferred_rel=0.05,
499
+ max_rel=0.5,
500
+ max_fail_rate=0.01,
501
+ )
502
+
503
+
504
+def test_rocm_aiter_rmsnorm_with_add_fp8_group_quant_residual_accuracy():
505
+ """Fused rmsnorm_with_add_fp8_group_quant residual output matches x + residual."""
506
+ require_aiter()
507
+ require_fp8()
508
+ from vllm._aiter_ops import rocm_aiter_ops
509
+
510
+ M, N = 16, 256
511
+ group_size = 128
512
+ x = torch.randn(M, N, dtype=torch.bfloat16)
513
+ residual = torch.randn(M, N, dtype=torch.bfloat16)
514
+ weight = torch.ones(N, dtype=torch.bfloat16)
515
+ eps = 1e-5
516
+
517
+ fused_q, fused_res, fused_scales = (
518
+ torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant(
519
+ x, residual, weight, eps, group_size
520
+ )
521
+ )
522
+
523
+ # Residual output must equal x + residual
524
+ ref_residual = (x.float() + residual.float()).to(torch.bfloat16)
525
+ _assert_close_budget(
526
+ fused_res.float(),
527
+ ref_residual.float(),
528
+ label="rmsnorm_with_add_fp8_group_quant residual",
529
+ atol=1e-2,
530
+ rtol=1e-2,
531
+ )
532
+
533
+ # Dequantized quant output must match rms_norm(x + residual)
534
+ h = ref_residual
535
+ ref_normed = _rms_norm_reference(h, weight, eps).to(torch.bfloat16)
536
+ ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(ref_normed, group_size)
537
+ ref_dequant = _dequantize_grouped(ref_q, ref_scales, group_size)
538
+ fused_dequant = _dequantize_grouped(fused_q, fused_scales, group_size)
539
+ _assert_rel_error_quality(
540
+ fused_dequant,
541
+ ref_dequant,
542
+ label="rmsnorm_with_add_fp8_group_quant",
543
+ mean_limit=0.05,
544
+ preferred_rel=0.05,
545
+ max_rel=0.5,
546
+ max_fail_rate=0.01,
547
+ )
548
+
549
+
550
+# -- End-to-end inference chain test ---------------------------------------
551
+
552
+
553
+def test_rocm_aiter_rms_norm_then_per_token_quant_e2e():
554
+ """End-to-end: BF16 RMSNorm -> per-token FP8 quantization -> dequantize.
555
+
556
+ Simulates the inference path through a transformer layer norm before a
557
+ linear projection: verifies the full chain produces accurate output
558
+ compared to a float32 reference.
559
+ """
560
+ require_aiter()
561
+ require_fp8()
562
+ from vllm._aiter_ops import rocm_aiter_ops
563
+
564
+ # Llama-style hidden dim
565
+ M, N = 32, 4096
566
+ x = torch.randn(M, N, dtype=torch.bfloat16)
567
+ weight = torch.randn(N, dtype=torch.bfloat16) # learned scale
568
+ eps = 1e-5
569
+ fp8_dtype = current_platform.fp8_dtype()
570
+
571
+ # Float32 reference for the full chain
572
+ ref_normed_f32 = _rms_norm_reference(x, weight, eps)
573
+
574
+ # AITER chain: RMSNorm -> per-token FP8 quant -> dequant
575
+ normed = rocm_aiter_ops.rms_norm(x, weight, eps)
576
+ x_q, scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype)
577
+ x_dequant = x_q.float() * scale.float() # scale: [M, 1]
578
+
579
+ _assert_quant_shapes(x_q, scale, M, N, fp8_dtype)
580
+ # Dequantized result should match the float32 reference within FP8 quant error
581
+ _assert_rel_error_quality(
582
+ x_dequant,
583
+ ref_normed_f32,
584
+ label="rms_norm_then_per_token_quant_e2e",
585
+ mean_limit=0.05,
586
+ preferred_rel=0.05,
587
+ max_rel=0.5,
588
+ max_fail_rate=0.01,
589
+ )
590
+
591
+
592
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
593
+def test_rocm_aiter_rms_norm_then_group_fp8_quant_e2e(dtype):
594
+ """End-to-end: RMSNorm (fp16/bf16) -> FP8 group quantization -> dequantize.
595
+
596
+ Covers both float16 and bfloat16 inputs to verify dtype-agnostic
597
+ behavior of the RMSNorm+group-quant pipeline.
598
+ """
599
+ require_aiter()
600
+ require_fp8()
601
+ from vllm._aiter_ops import rocm_aiter_ops
602
+
603
+ M, N = 16, 512
604
+ group_size = 128
605
+ x = torch.randn(M, N, dtype=dtype)
606
+ weight = torch.randn(N, dtype=dtype)
607
+ eps = 1e-5
608
+
609
+ # Float32 reference
610
+ ref_normed_f32 = _rms_norm_reference(x, weight, eps)
611
+
612
+ # AITER chain: RMSNorm -> group FP8 quant -> dequant
613
+ normed = rocm_aiter_ops.rms_norm(x, weight, eps)
614
+ x_fp8, scales = rocm_aiter_ops.group_fp8_quant(normed.bfloat16(), group_size)
615
+ x_dequant = _dequantize_grouped(x_fp8, scales, group_size)
616
+
617
+ _assert_rel_error_quality(
618
+ x_dequant,
619
+ ref_normed_f32,
620
+ label=f"rms_norm_then_group_fp8_quant_e2e dtype={dtype}",
621
+ mean_limit=0.1,
622
+ preferred_rel=0.1,
623
+ max_rel=1.0,
624
+ max_fail_rate=0.01,
625
+ )
vllm/_aiter_ops.py
+29
@@ -2340,6 +2340,35 @@ class rocm_aiter_ops:
2340
def get_fused_mla_dual_rms_norm_per_token_quant_op() -> OpOverload:
2341
return torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default
2342
2343
+ @staticmethod
2344
+ def rms_norm(
2345
+ x: torch.Tensor,
2346
+ weight: torch.Tensor,
2347
+ epsilon: float,
2348
+ ) -> torch.Tensor:
2349
+ """RMSNorm via AITER kernel."""
2350
+ import aiter
2351
+
2352
+ return aiter.rms_norm(x, weight, epsilon)
2353
+
2354
+ @staticmethod
2355
+ def rms_norm2d_with_add(
2356
+ x: torch.Tensor,
2357
+ residual: torch.Tensor,
2358
+ weight: torch.Tensor,
2359
+ epsilon: float,
2360
+ ) -> tuple[torch.Tensor, torch.Tensor]:
2361
+ """Fused residual-add + RMSNorm via AITER kernel.
2362
+
2363
+ Returns (normalized_output, residual_sum).
2364
+ """
2365
+ import aiter
2366
+
2367
+ out = torch.empty_like(x)
2368
+ residual_out = torch.empty_like(x)
2369
+ aiter.rmsnorm2d_fwd_with_add(out, x, residual, residual_out, weight, epsilon, 0)
2370
+ return out, residual_out
2371
+
2372
@staticmethod
2373
def w8a8_gemm(
2374
A: torch.Tensor,