| 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, fp8_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 | def test_rocm_aiter_triton_rotary_embedding_vs_torch(): |
| 260 | """rocm_aiter_triton_rotary_embedding matches manual NeoX-style RoPE reference. |
| 261 | |
| 262 | The AITER kernel calls rope_cached_thd_positions_offsets_2c_fwd_inplace with |
| 263 | reuse_freqs_front_part=True. This means it reads only the first head_size//2 |
| 264 | entries of cos and sin from the cache and applies them to both halves of the |
| 265 | head (pairwise NeoX rotation). The cos_sin_cache must be built with the second |
| 266 | half mirroring the first so the reference and kernel agree. |
| 267 | """ |
| 268 | require_aiter() |
| 269 | |
| 270 | num_tokens = 8 |
| 271 | num_heads = 4 |
| 272 | head_size = 64 |
| 273 | half_dim = head_size // 2 |
| 274 | max_pos = 32 |
| 275 | |
| 276 | positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.long) |
| 277 | query = torch.randn(num_tokens, num_heads * head_size, dtype=torch.bfloat16) |
| 278 | key = torch.randn(num_tokens, num_heads * head_size, dtype=torch.bfloat16) |
| 279 | |
| 280 | # Build cos/sin cache with second half mirroring first half. |
| 281 | # The kernel uses reuse_freqs_front_part=True: it reads only cos/sin[:, :half_dim] |
| 282 | # and applies those same frequencies to both the first and second |
| 283 | # halves of the head. |
| 284 | cos_half = torch.randn(max_pos, half_dim, dtype=torch.bfloat16) |
| 285 | sin_half = torch.randn(max_pos, half_dim, dtype=torch.bfloat16) |
| 286 | cos_cache = torch.cat([cos_half, cos_half], dim=-1) # [max_pos, head_size] |
| 287 | sin_cache = torch.cat([sin_half, sin_half], dim=-1) # [max_pos, head_size] |
| 288 | cos_sin_cache = torch.cat([cos_cache, sin_cache], dim=-1) # [max_pos, 2*head_size] |
| 289 | |
| 290 | # Reference: NeoX-style pairwise rotation with front-half frequencies only. |
| 291 | # rotate_style=0 (NeoX): [x1*c - x2*s, x2*c + x1*s] |
| 292 | cos_pos = cos_half[positions] # [num_tokens, half_dim] |
| 293 | sin_pos = sin_half[positions] # [num_tokens, half_dim] |
| 294 | |
| 295 | def apply_rope_ref(t: torch.Tensor) -> torch.Tensor: |
| 296 | # AITER triton keeps native dtype (no fp32 upcast), so reference must match. |
| 297 | t_r = t.view(num_tokens, num_heads, head_size) |
| 298 | c = cos_pos.unsqueeze(1) # [num_tokens, 1, half_dim] |
| 299 | s = sin_pos.unsqueeze(1) |
| 300 | x1, x2 = t_r[..., :half_dim], t_r[..., half_dim:] |
| 301 | rotated = torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1) |
| 302 | return rotated.view(num_tokens, num_heads * head_size) |
| 303 | |
| 304 | ref_q = apply_rope_ref(query) |
| 305 | ref_k = apply_rope_ref(key) |
| 306 | |
| 307 | # AITER in-place RoPE (modifies query/key in-place) |
| 308 | q_aiter = query.clone() |
| 309 | k_aiter = key.clone() |
| 310 | torch.ops.vllm.rocm_aiter_triton_rotary_embedding( |
| 311 | positions, |
| 312 | q_aiter, |
| 313 | k_aiter, |
| 314 | head_size, |
| 315 | cos_sin_cache, |
| 316 | True, # is_neox style -> rotate_style=0 |
| 317 | ) |
| 318 | |
| 319 | torch.testing.assert_close(q_aiter, ref_q) |
| 320 | torch.testing.assert_close(k_aiter, ref_k) |
| 321 | |
| 322 | |
| 323 | def test_rocm_aiter_act_mul_fp8_group_quant_roundtrip(): |
| 324 | """act_mul_and_fp8_group_quant: dequantized output matches SiLU gate reference.""" |
| 325 | require_aiter() |
| 326 | require_fp8() |
| 327 | |
| 328 | M, N = 32, 512 # N even: N//2 gate, N//2 up |
| 329 | group_size = 128 |
| 330 | x = torch.randn(M, N, dtype=torch.bfloat16) |
| 331 | |
| 332 | x_quant, scale = torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant( |
| 333 | x, group_size |
| 334 | ) |
| 335 | |
| 336 | N_half = N // 2 |
| 337 | # Reference: SiLU(gate) * up |
| 338 | gate = x.float()[:, :N_half] |
| 339 | up = x.float()[:, N_half:] |
| 340 | ref = torch.sigmoid(gate) * gate * up # SiGLU |
| 341 | |
| 342 | x_dequant = _dequantize_grouped(x_quant, scale, group_size) |
| 343 | |
| 344 | _assert_rel_error_quality( |
| 345 | x_dequant, |
| 346 | ref, |
| 347 | label="act_mul_fp8_group_quant", |
| 348 | mean_limit=0.1, |
| 349 | preferred_rel=0.1, |
| 350 | max_rel=1.0, |
| 351 | max_fail_rate=0.01, |
| 352 | ) |
| 353 | |
| 354 | |
| 355 | # -- Fused RMSNorm + quantization accuracy tests --------------------------- |
| 356 | |
| 357 | |
| 358 | def test_rocm_aiter_rmsnorm_fused_dynamic_quant_vs_sequential(): |
| 359 | """Fused RMSNorm+per-token-FP8-quant matches sequential rms_norm->per_token_quant. |
| 360 | |
| 361 | Tests that the fused kernel produces the same result as the two-step |
| 362 | sequential composition: rms_norm(x) followed by per_token_quant. |
| 363 | The fused path is used in production for inference throughput. |
| 364 | """ |
| 365 | require_aiter() |
| 366 | require_fp8() |
| 367 | from vllm._aiter_ops import rocm_aiter_ops |
| 368 | |
| 369 | M, N = 32, 512 |
| 370 | x = torch.randn(M, N, dtype=torch.bfloat16) |
| 371 | weight = torch.ones(N, dtype=torch.bfloat16) |
| 372 | eps = 1e-5 |
| 373 | fp8_dtype = current_platform.fp8_dtype() |
| 374 | |
| 375 | # Sequential reference |
| 376 | normed = rocm_aiter_ops.rms_norm(x, weight, eps) |
| 377 | ref_q, ref_scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype) |
| 378 | ref_dequant = ref_q.float() * ref_scale.float() |
| 379 | |
| 380 | # Fused op |
| 381 | fused_q, fused_scale = torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant( |
| 382 | x, weight, eps, fp8_dtype |
| 383 | ) |
| 384 | _assert_quant_shapes(fused_q, fused_scale, M, N, fp8_dtype) |
| 385 | fused_dequant = fused_q.float() * fused_scale.float() |
| 386 | |
| 387 | # Fused vs sequential: both should recover RMSNorm output within FP8 error |
| 388 | _assert_rel_error_quality( |
| 389 | fused_dequant, |
| 390 | ref_dequant, |
| 391 | label="rmsnorm_fused_dynamic_quant", |
| 392 | mean_limit=0.05, |
| 393 | preferred_rel=0.05, |
| 394 | max_rel=0.5, |
| 395 | max_fail_rate=0.01, |
| 396 | ) |
| 397 | |
| 398 | |
| 399 | def test_rocm_aiter_rmsnorm_fused_add_dynamic_quant_vs_reference(): |
| 400 | """Fused (residual-add + RMSNorm + per-token-FP8-quant) |
| 401 | matches sequential reference. |
| 402 | |
| 403 | Production path: input + residual -> RMSNorm -> FP8 quantize, returning |
| 404 | both the quantized norm output and the residual sum for the next layer. |
| 405 | """ |
| 406 | require_aiter() |
| 407 | require_fp8() |
| 408 | from vllm._aiter_ops import rocm_aiter_ops |
| 409 | |
| 410 | M, N = 16, 256 |
| 411 | x = torch.randn(M, N, dtype=torch.bfloat16) |
| 412 | residual = torch.randn(M, N, dtype=torch.bfloat16) |
| 413 | weight = torch.ones(N, dtype=torch.bfloat16) |
| 414 | eps = 1e-5 |
| 415 | fp8_dtype = current_platform.fp8_dtype() |
| 416 | |
| 417 | # Sequential reference: add residual -> rms_norm -> per_token_quant |
| 418 | h = (x.float() + residual.float()).to(torch.bfloat16) |
| 419 | normed = rocm_aiter_ops.rms_norm(h, weight, eps) |
| 420 | ref_q, ref_scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype) |
| 421 | ref_residual_out = h |
| 422 | |
| 423 | # Fused op: returns (x_quant, residual_out, scale) |
| 424 | fused_q, fused_res_out, fused_scale = ( |
| 425 | torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant( |
| 426 | x, residual, weight, eps, fp8_dtype |
| 427 | ) |
| 428 | ) |
| 429 | _assert_quant_shapes(fused_q, fused_scale, M, N, fp8_dtype) |
| 430 | assert fused_res_out.shape == (M, N) |
| 431 | |
| 432 | torch.testing.assert_close(fused_res_out, ref_residual_out, atol=1e-2, rtol=1.6e-2) |
| 433 | # Dequantized output matches sequential path |
| 434 | fused_dequant = fused_q.float() * fused_scale.float() |
| 435 | ref_dequant = ref_q.float() * ref_scale.float() |
| 436 | _assert_rel_error_quality( |
| 437 | fused_dequant, |
| 438 | ref_dequant, |
| 439 | label="rmsnorm_fused_add_dynamic_quant", |
| 440 | mean_limit=0.05, |
| 441 | preferred_rel=0.05, |
| 442 | max_rel=0.5, |
| 443 | max_fail_rate=0.01, |
| 444 | ) |
| 445 | |
| 446 | |
| 447 | def test_rocm_aiter_rmsnorm_fp8_group_quant_vs_sequential(): |
| 448 | """Fused RMSNorm+FP8-group-quant matches sequential rms_norm->group_fp8_quant. |
| 449 | |
| 450 | Tests both output shapes and dequantized accuracy against the two-step |
| 451 | sequential composition. |
| 452 | """ |
| 453 | require_aiter() |
| 454 | require_fp8() |
| 455 | from vllm._aiter_ops import rocm_aiter_ops |
| 456 | |
| 457 | M, N = 32, 512 |
| 458 | group_size = 128 |
| 459 | x = torch.randn(M, N, dtype=torch.bfloat16) |
| 460 | weight = torch.ones(N, dtype=torch.bfloat16) |
| 461 | eps = 1e-5 |
| 462 | fp8_dtype = current_platform.fp8_dtype() |
| 463 | expected_groups = (N + group_size - 1) // group_size |
| 464 | |
| 465 | # Fused op: (x_quant, scales) |
| 466 | fused_q, fused_scales = torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant( |
| 467 | x, weight, eps, group_size |
| 468 | ) |
| 469 | _assert_quant_shapes(fused_q, fused_scales, M, N, fp8_dtype, expected_groups) |
| 470 | |
| 471 | # Dequantize and compare to reference: rms_norm -> group quant -> dequant |
| 472 | normed = rocm_aiter_ops.rms_norm(x, weight, eps) |
| 473 | ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(normed, group_size) |
| 474 | ref_dequant = _dequantize_grouped(ref_q, ref_scales, group_size) |
| 475 | fused_dequant = _dequantize_grouped(fused_q, fused_scales, group_size) |
| 476 | |
| 477 | _assert_rel_error_quality( |
| 478 | fused_dequant, |
| 479 | ref_dequant, |
| 480 | label="rmsnorm_fp8_group_quant", |
| 481 | mean_limit=0.05, |
| 482 | preferred_rel=0.05, |
| 483 | max_rel=0.5, |
| 484 | max_fail_rate=0.01, |
| 485 | ) |
| 486 | |
| 487 | |
| 488 | def test_rocm_aiter_rmsnorm_with_add_fp8_group_quant_residual_accuracy(): |
| 489 | """Fused rmsnorm_with_add_fp8_group_quant residual output matches x + residual.""" |
| 490 | require_aiter() |
| 491 | require_fp8() |
| 492 | from vllm._aiter_ops import rocm_aiter_ops |
| 493 | |
| 494 | M, N = 16, 256 |
| 495 | group_size = 128 |
| 496 | x = torch.randn(M, N, dtype=torch.bfloat16) |
| 497 | residual = torch.randn(M, N, dtype=torch.bfloat16) |
| 498 | weight = torch.ones(N, dtype=torch.bfloat16) |
| 499 | eps = 1e-5 |
| 500 | |
| 501 | fused_q, fused_res, fused_scales = ( |
| 502 | torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant( |
| 503 | x, residual, weight, eps, group_size |
| 504 | ) |
| 505 | ) |
| 506 | |
| 507 | # Residual output must equal x + residual |
| 508 | ref_residual = (x.float() + residual.float()).to(torch.bfloat16) |
| 509 | _assert_close_budget( |
| 510 | fused_res.float(), |
| 511 | ref_residual.float(), |
| 512 | label="rmsnorm_with_add_fp8_group_quant residual", |
| 513 | atol=1e-2, |
| 514 | rtol=1e-2, |
| 515 | ) |
| 516 | |
| 517 | # Dequantized quant output must match rms_norm(x + residual) |
| 518 | h = ref_residual |
| 519 | ref_normed = _rms_norm_reference(h, weight, eps).to(torch.bfloat16) |
| 520 | ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(ref_normed, group_size) |
| 521 | ref_dequant = _dequantize_grouped(ref_q, ref_scales, group_size) |
| 522 | fused_dequant = _dequantize_grouped(fused_q, fused_scales, group_size) |
| 523 | _assert_rel_error_quality( |
| 524 | fused_dequant, |
| 525 | ref_dequant, |
| 526 | label="rmsnorm_with_add_fp8_group_quant", |
| 527 | mean_limit=0.05, |
| 528 | preferred_rel=0.05, |
| 529 | max_rel=0.5, |
| 530 | max_fail_rate=0.01, |
| 531 | ) |
| 532 | |
| 533 | |
| 534 | @pytest.mark.parametrize("m", [1, 32]) |
| 535 | def test_rocm_aiter_mla_dual_rmsnorm_group_quant_vs_sequential(m): |
| 536 | """Fused q/kv RMSNorm + q group-quant matches rms_norm -> group_fp8_quant. |
| 537 | |
| 538 | M=1 covers the single-token decode case. M=0 is guarded by the caller |
| 539 | (the ROCm DeepSeek-V4 attention checks num_tokens > 0) and the op has |
| 540 | no zero-row guard itself, so it is intentionally not tested here. |
| 541 | """ |
| 542 | require_aiter() |
| 543 | require_fp8() |
| 544 | from vllm._aiter_ops import rocm_aiter_ops |
| 545 | |
| 546 | nq, nk, group_size = 1536, 576, 128 |
| 547 | q = torch.randn(m, nq, dtype=torch.bfloat16) |
| 548 | kv = torch.randn(m, nk, dtype=torch.bfloat16) |
| 549 | q_weight = torch.ones(nq, dtype=torch.bfloat16) |
| 550 | kv_weight = torch.ones(nk, dtype=torch.bfloat16) |
| 551 | eps = 1e-5 |
| 552 | fp8_dtype = current_platform.fp8_dtype() |
| 553 | |
| 554 | fused_q, fused_scales, fused_kv = rocm_aiter_ops.fused_qk_rmsnorm_group_quant( |
| 555 | q, q_weight, eps, kv, kv_weight, eps, group_size=group_size |
| 556 | ) |
| 557 | _assert_quant_shapes(fused_q, fused_scales, m, nq, fp8_dtype, nq // group_size) |
| 558 | assert fused_kv.shape == (m, nk) |
| 559 | assert fused_kv.dtype == kv.dtype |
| 560 | |
| 561 | # q: dequantized comparison against the sequential composition. |
| 562 | normed_q = rocm_aiter_ops.rms_norm(q, q_weight, eps) |
| 563 | ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(normed_q, group_size) |
| 564 | _assert_rel_error_quality( |
| 565 | _dequantize_grouped(fused_q, fused_scales, group_size), |
| 566 | _dequantize_grouped(ref_q, ref_scales, group_size), |
| 567 | label="mla_dual_rmsnorm_group_quant", |
| 568 | mean_limit=0.05, |
| 569 | preferred_rel=0.05, |
| 570 | max_rel=0.5, |
| 571 | max_fail_rate=0.01, |
| 572 | ) |
| 573 | |
| 574 | # kv: both sides are aiter RMSNorm kernels; allow one bf16 ULP. |
| 575 | ref_kv = rocm_aiter_ops.rms_norm(kv, kv_weight, eps) |
| 576 | max_ulp = int(bf16_ulp_distance(fused_kv, ref_kv).max().item()) |
| 577 | assert max_ulp <= 1, f"fused kv: max ULP {max_ulp} > 1" |
| 578 | |
| 579 | |
| 580 | def test_rocm_aiter_mla_dual_rmsnorm_group_quant_strided_split_view(): |
| 581 | """Split-view inputs, as produced by the DeepSeek-V4 q/kv split, match |
| 582 | contiguous inputs.""" |
| 583 | require_aiter() |
| 584 | require_fp8() |
| 585 | from vllm._aiter_ops import rocm_aiter_ops |
| 586 | |
| 587 | m, nq, nk, group_size = 16, 1536, 576, 128 |
| 588 | qr_kv = torch.randn(m, nq + nk, dtype=torch.bfloat16) |
| 589 | q_view, kv_view = qr_kv.split([nq, nk], dim=-1) |
| 590 | assert not q_view.is_contiguous() |
| 591 | q_weight, kv_weight = torch.ones(nq + nk, dtype=torch.bfloat16).split([nq, nk]) |
| 592 | eps = 1e-5 |
| 593 | |
| 594 | def run(q, kv): |
| 595 | return rocm_aiter_ops.fused_qk_rmsnorm_group_quant( |
| 596 | q, q_weight, eps, kv, kv_weight, eps, group_size=group_size |
| 597 | ) |
| 598 | |
| 599 | out_strided = run(q_view, kv_view) |
| 600 | out_contig = run(q_view.contiguous(), kv_view.contiguous()) |
| 601 | |
| 602 | q_ulp = int(fp8_ulp_distance(out_strided[0], out_contig[0]).max().item()) |
| 603 | kv_ulp = int(bf16_ulp_distance(out_strided[2], out_contig[2]).max().item()) |
| 604 | assert q_ulp <= 1, f"strided q: max fp8 ULP {q_ulp} > 1" |
| 605 | assert kv_ulp <= 1, f"strided kv: max bf16 ULP {kv_ulp} > 1" |
| 606 | torch.testing.assert_close(out_strided[1], out_contig[1]) |
| 607 | |
| 608 | |
| 609 | # -- End-to-end inference chain test --------------------------------------- |
| 610 | |
| 611 | |
| 612 | def test_rocm_aiter_rms_norm_then_per_token_quant_e2e(): |
| 613 | """End-to-end: BF16 RMSNorm -> per-token FP8 quantization -> dequantize. |
| 614 | |
| 615 | Simulates the inference path through a transformer layer norm before a |
| 616 | linear projection: verifies the full chain produces accurate output |
| 617 | compared to a float32 reference. |
| 618 | """ |
| 619 | require_aiter() |
| 620 | require_fp8() |
| 621 | from vllm._aiter_ops import rocm_aiter_ops |
| 622 | |
| 623 | # Llama-style hidden dim |
| 624 | M, N = 32, 4096 |
| 625 | x = torch.randn(M, N, dtype=torch.bfloat16) |
| 626 | weight = torch.randn(N, dtype=torch.bfloat16) # learned scale |
| 627 | eps = 1e-5 |
| 628 | fp8_dtype = current_platform.fp8_dtype() |
| 629 | |
| 630 | # Float32 reference for the full chain |
| 631 | ref_normed_f32 = _rms_norm_reference(x, weight, eps) |
| 632 | |
| 633 | # AITER chain: RMSNorm -> per-token FP8 quant -> dequant |
| 634 | normed = rocm_aiter_ops.rms_norm(x, weight, eps) |
| 635 | x_q, scale = rocm_aiter_ops.per_token_quant(normed, fp8_dtype) |
| 636 | x_dequant = x_q.float() * scale.float() # scale: [M, 1] |
| 637 | |
| 638 | _assert_quant_shapes(x_q, scale, M, N, fp8_dtype) |
| 639 | # Dequantized result should match the float32 reference within FP8 quant error |
| 640 | _assert_rel_error_quality( |
| 641 | x_dequant, |
| 642 | ref_normed_f32, |
| 643 | label="rms_norm_then_per_token_quant_e2e", |
| 644 | mean_limit=0.05, |
| 645 | preferred_rel=0.05, |
| 646 | max_rel=0.5, |
| 647 | max_fail_rate=0.01, |
| 648 | ) |
| 649 | |
| 650 | |
| 651 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 652 | def test_rocm_aiter_rms_norm_then_group_fp8_quant_e2e(dtype): |
| 653 | """End-to-end: RMSNorm (fp16/bf16) -> FP8 group quantization -> dequantize. |
| 654 | |
| 655 | Covers both float16 and bfloat16 inputs to verify dtype-agnostic |
| 656 | behavior of the RMSNorm+group-quant pipeline. |
| 657 | """ |
| 658 | require_aiter() |
| 659 | require_fp8() |
| 660 | from vllm._aiter_ops import rocm_aiter_ops |
| 661 | |
| 662 | M, N = 16, 512 |
| 663 | group_size = 128 |
| 664 | x = torch.randn(M, N, dtype=dtype) |
| 665 | weight = torch.randn(N, dtype=dtype) |
| 666 | eps = 1e-5 |
| 667 | |
| 668 | # Float32 reference |
| 669 | ref_normed_f32 = _rms_norm_reference(x, weight, eps) |
| 670 | |
| 671 | # AITER chain: RMSNorm -> group FP8 quant -> dequant |
| 672 | normed = rocm_aiter_ops.rms_norm(x, weight, eps) |
| 673 | x_fp8, scales = rocm_aiter_ops.group_fp8_quant(normed.bfloat16(), group_size) |
| 674 | x_dequant = _dequantize_grouped(x_fp8, scales, group_size) |
| 675 | |
| 676 | _assert_rel_error_quality( |
| 677 | x_dequant, |
| 678 | ref_normed_f32, |
| 679 | label=f"rms_norm_then_group_fp8_quant_e2e dtype={dtype}", |
| 680 | mean_limit=0.1, |
| 681 | preferred_rel=0.1, |
| 682 | max_rel=1.0, |
| 683 | max_fail_rate=0.01, |
| 684 | ) |