| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm AITER FP4 and MXFP4 tests. |
| 4 | |
| 5 | This file keeps the ROCm-specific FP4 coverage: |
| 6 | - ROCm AITER FP4 env and enablement gates |
| 7 | - large-shape MXFP4 wrapper checks through vLLM's public helper op |
| 8 | - AITER Triton MXFP4 quant format and determinism |
| 9 | - FP4 GEMM, preshuffled-scale, hardware-quant, and skinny decode |
| 10 | paths |
| 11 | |
| 12 | Generic MXFP4 wrapper/reference checks already live in |
| 13 | ``tests/kernels/quantization/test_quark.py``. |
| 14 | """ |
| 15 | |
| 16 | import importlib |
| 17 | import warnings |
| 18 | |
| 19 | import pytest |
| 20 | import torch |
| 21 | |
| 22 | from tests.kernels.utils import _assert_accurate, _assert_deterministic |
| 23 | from tests.quantization.reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.platforms.rocm import on_gfx950 |
| 26 | |
| 27 | pytestmark = pytest.mark.skipif( |
| 28 | not current_platform.is_rocm(), reason="ROCm-specific tests" |
| 29 | ) |
| 30 | |
| 31 | LLAMA_SHAPES = [ |
| 32 | (64, 4096), |
| 33 | (64, 11008), |
| 34 | (32, 14336), |
| 35 | ] |
| 36 | GEMM_ATOL = 0.5 |
| 37 | GEMM_RTOL = 0.0 |
| 38 | GEMM_MAX_VIOLATION_FACTOR = 2.0 |
| 39 | A4W4_GEMM_PASS_RATES = { |
| 40 | (64, 128, 64): 0.9999, |
| 41 | (128, 256, 128): 0.9999, |
| 42 | (128, 4096, 4096): 0.9999, |
| 43 | (256, 4096, 11008): 0.9999, |
| 44 | (64, 8192, 28672): 0.99765, |
| 45 | } |
| 46 | SKINNY_GEMM_PASS_RATES = { |
| 47 | (4096, 4096): 0.9999, |
| 48 | (4096, 11008): 0.99235, |
| 49 | (8192, 8192): 0.99765, |
| 50 | } |
| 51 | PRESHUFFLED_SHAPES = [ |
| 52 | (64, 4096, 8192), |
| 53 | (32, 8192, 16384), |
| 54 | ] |
| 55 | |
| 56 | |
| 57 | def _reload_envs(): |
| 58 | import vllm.envs as envs |
| 59 | |
| 60 | return importlib.reload(envs) |
| 61 | |
| 62 | |
| 63 | @pytest.fixture(autouse=True) |
| 64 | def _restore_rocm_env_state(): |
| 65 | """Restore global env + AITER flag state after every test.""" |
| 66 | yield |
| 67 | _reload_envs() |
| 68 | from vllm._aiter_ops import rocm_aiter_ops |
| 69 | |
| 70 | rocm_aiter_ops.refresh_env_variables() |
| 71 | |
| 72 | |
| 73 | def _assert_aiter_supported() -> None: |
| 74 | from vllm._aiter_ops import is_aiter_found_and_supported |
| 75 | |
| 76 | assert is_aiter_found_and_supported(), ( |
| 77 | "aiter is required on supported ROCm hardware for this test" |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | def _format_observed_rate(count: int, total: int) -> str: |
| 82 | return f"{count / total:.4%} ({count}/{total})" |
| 83 | |
| 84 | |
| 85 | def _format_allowed_rate(rate: float, total: int) -> str: |
| 86 | allowed_count = int(rate * total) |
| 87 | return f"{rate:.4%} (<= {allowed_count}/{total})" |
| 88 | |
| 89 | |
| 90 | def _print_close_stats( |
| 91 | label: str, |
| 92 | actual: torch.Tensor, |
| 93 | expected: torch.Tensor, |
| 94 | *, |
| 95 | atol: float, |
| 96 | rtol: float, |
| 97 | pass_rate: float = 1.0, |
| 98 | max_atol: float | None = None, |
| 99 | ) -> None: |
| 100 | abs_diff = (actual - expected).abs().float().flatten() |
| 101 | expected_abs = expected.abs().float().flatten() |
| 102 | allowed = atol + rtol * expected_abs |
| 103 | within = abs_diff <= allowed |
| 104 | |
| 105 | total = abs_diff.numel() |
| 106 | passed = int(within.sum().item()) |
| 107 | failed = total - passed |
| 108 | allowed_fail_rate = 1.0 - pass_rate |
| 109 | |
| 110 | max_abs = abs_diff.max().item() |
| 111 | mean_abs = abs_diff.mean().item() |
| 112 | p99_abs = torch.quantile(abs_diff, 0.99).item() |
| 113 | p999_abs = torch.quantile(abs_diff, 0.999).item() |
| 114 | worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item() |
| 115 | |
| 116 | msg = ( |
| 117 | "[rocm_aiter_fp4] " |
| 118 | f"{label}: " |
| 119 | f"pass={passed / total:.4%} ({passed}/{total}) " |
| 120 | f"fail={_format_observed_rate(failed, total)} " |
| 121 | f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} " |
| 122 | f"atol={atol:g} " |
| 123 | f"rtol={rtol:g} " |
| 124 | ) |
| 125 | if max_atol is not None: |
| 126 | above_max_count = int((abs_diff > max_atol).sum().item()) |
| 127 | msg += ( |
| 128 | f"abs>{max_atol:g}={_format_observed_rate(above_max_count, total)} " |
| 129 | f"allowed_above_max={_format_allowed_rate(0.0, total)} " |
| 130 | ) |
| 131 | msg += ( |
| 132 | f"max_abs={max_abs:.6g} " |
| 133 | f"mean_abs={mean_abs:.6g} " |
| 134 | f"p99_abs={p99_abs:.6g} " |
| 135 | f"p999_abs={p999_abs:.6g} " |
| 136 | f"worst_ratio={worst_ratio:.6g}" |
| 137 | ) |
| 138 | print(msg) |
| 139 | if failed > 0: |
| 140 | warnings.warn(msg, stacklevel=2) |
| 141 | |
| 142 | |
| 143 | # Env and enablement tests ------------------------------------------------ |
| 144 | |
| 145 | |
| 146 | def test_fp4_env_defaults(): |
| 147 | """ROCm FP4 env defaults should stay stable for the AITER gates.""" |
| 148 | import vllm.envs as envs |
| 149 | |
| 150 | assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False |
| 151 | assert envs.VLLM_ROCM_USE_AITER_FP4BMM is True |
| 152 | |
| 153 | |
| 154 | @pytest.mark.parametrize( |
| 155 | ( |
| 156 | "use_aiter", |
| 157 | "use_fp4_asm_gemm", |
| 158 | "use_fp4bmm", |
| 159 | ), |
| 160 | [ |
| 161 | (True, True, True), |
| 162 | (True, True, False), |
| 163 | (True, False, True), |
| 164 | (False, True, True), |
| 165 | ], |
| 166 | ) |
| 167 | def test_rocm_aiter_fp4_enablement_follows_env_and_arch( |
| 168 | use_aiter, |
| 169 | use_fp4_asm_gemm, |
| 170 | use_fp4bmm, |
| 171 | monkeypatch, |
| 172 | ): |
| 173 | """The ROCm FP4 AITER gates should depend only on the env toggles and the |
| 174 | gfx950 hardware check. |
| 175 | |
| 176 | Expected values are derived from the *actual* platform rather than a patched |
| 177 | arch (gfx950 -> MI355, gfx942 -> MI300). A failure in this group therefore |
| 178 | points straight at the hardware it ran on instead of a mocked value. |
| 179 | """ |
| 180 | from vllm._aiter_ops import rocm_aiter_ops |
| 181 | |
| 182 | _assert_aiter_supported() |
| 183 | |
| 184 | on_gfx950_value = on_gfx950() |
| 185 | expected_asm_gemm = use_aiter and use_fp4_asm_gemm and on_gfx950_value |
| 186 | expected_fp4bmm = use_aiter and use_fp4bmm and on_gfx950_value |
| 187 | |
| 188 | with monkeypatch.context() as mp: |
| 189 | mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0") |
| 190 | mp.setenv( |
| 191 | "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", |
| 192 | "1" if use_fp4_asm_gemm else "0", |
| 193 | ) |
| 194 | mp.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1" if use_fp4bmm else "0") |
| 195 | _reload_envs() |
| 196 | rocm_aiter_ops.refresh_env_variables() |
| 197 | |
| 198 | assert ( |
| 199 | rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() is expected_asm_gemm |
| 200 | ) |
| 201 | assert rocm_aiter_ops.is_fp4bmm_enabled() is expected_fp4bmm |
| 202 | |
| 203 | |
| 204 | # Large-shape MXFP4 wrapper tests ----------------------------------------- |
| 205 | |
| 206 | |
| 207 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 208 | @pytest.mark.parametrize("shape", LLAMA_SHAPES) |
| 209 | def test_vllm_quant_dequant_mxfp4_matches_reference_on_large_shapes(shape, dtype): |
| 210 | """The public vLLM MXFP4 QDQ helper should stay exact against the torch |
| 211 | reference even on larger Llama-like shapes.""" |
| 212 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 213 | quant_dequant_mxfp4, |
| 214 | ) |
| 215 | |
| 216 | torch.set_default_device("cuda") |
| 217 | torch.manual_seed(0) |
| 218 | rows, cols = shape |
| 219 | x = torch.randn(rows, cols, dtype=dtype) |
| 220 | out = quant_dequant_mxfp4(x) |
| 221 | ref = qdq_mxfp4_torch(x, "even") |
| 222 | |
| 223 | assert out.shape == x.shape |
| 224 | assert out.dtype == dtype |
| 225 | assert torch.isfinite(out).all() |
| 226 | torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0) |
| 227 | |
| 228 | |
| 229 | def test_vllm_quant_dequant_mxfp4_is_deterministic(): |
| 230 | """The public vLLM MXFP4 QDQ helper should stay bitwise deterministic.""" |
| 231 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 232 | quant_dequant_mxfp4, |
| 233 | ) |
| 234 | |
| 235 | torch.set_default_device("cuda") |
| 236 | torch.manual_seed(3) |
| 237 | |
| 238 | x = torch.randn(64, 128, dtype=torch.bfloat16) |
| 239 | |
| 240 | _assert_deterministic(quant_dequant_mxfp4, x, n_runs=4) |
| 241 | |
| 242 | |
| 243 | # Triton MXFP4 quant tests ------------------------------------------------ |
| 244 | |
| 245 | |
| 246 | def test_aiter_dynamic_mxfp4_quant_output_format(): |
| 247 | """dynamic_mxfp4_quant returns packed uint8 FP4 values and E8M0 uint8 scales. |
| 248 | |
| 249 | OCP MXFP4: block_size=32, 2 FP4 E2M1 values packed per byte. |
| 250 | Scale shape: (M, K // 32); one E8M0 exponent byte per 32-element block. |
| 251 | """ |
| 252 | _assert_aiter_supported() |
| 253 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 254 | |
| 255 | torch.set_default_device("cuda") |
| 256 | M, K = 64, 256 |
| 257 | x = torch.randn(M, K, dtype=torch.bfloat16) |
| 258 | |
| 259 | x_fp4, x_scale = dynamic_mxfp4_quant(x) |
| 260 | |
| 261 | # FP4 values packed 2-per-byte; shape (M, K // 2) |
| 262 | assert x_fp4.shape == (M, K // 2), ( |
| 263 | f"Expected fp4 shape ({M}, {K // 2}), got {x_fp4.shape}" |
| 264 | ) |
| 265 | assert x_fp4.dtype == torch.uint8 |
| 266 | |
| 267 | # One E8M0 scale byte per 32-element block; shape (M, K // 32) |
| 268 | assert x_scale.shape == (M, K // 32), ( |
| 269 | f"Expected scale shape ({M}, {K // 32}), got {x_scale.shape}" |
| 270 | ) |
| 271 | assert x_scale.dtype == torch.uint8 |
| 272 | |
| 273 | |
| 274 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 275 | @pytest.mark.parametrize( |
| 276 | "shape", |
| 277 | [ |
| 278 | (128, 4096), # Llama-7B hidden |
| 279 | (256, 11008), # Llama-7B FFN |
| 280 | (64, 14336), # Llama-70B FFN |
| 281 | (512, 8192), # Llama-70B hidden |
| 282 | (32, 28672), # DeepSeek-style large FFN |
| 283 | ], |
| 284 | ) |
| 285 | def test_aiter_dynamic_mxfp4_quant_llama_shapes(shape, dtype): |
| 286 | """MXFP4 quantization output format is correct for Llama-class weight shapes. |
| 287 | |
| 288 | Targets B200 parity: NVIDIA nvfp4 tests parametrize over Llama shapes |
| 289 | (7168, 14336, 28672...). This covers gfx950 MXFP4 quant format validation. |
| 290 | """ |
| 291 | _assert_aiter_supported() |
| 292 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 293 | |
| 294 | torch.set_default_device("cuda") |
| 295 | M, K = shape |
| 296 | x = torch.randn(M, K, dtype=dtype) |
| 297 | |
| 298 | x_fp4, x_scale = dynamic_mxfp4_quant(x) |
| 299 | |
| 300 | assert x_fp4.shape == (M, K // 2) |
| 301 | assert x_fp4.dtype == torch.uint8 |
| 302 | assert x_scale.shape == (M, K // 32) |
| 303 | assert x_scale.dtype == torch.uint8 |
| 304 | assert not torch.any(torch.isnan(x_fp4.float())) |
| 305 | |
| 306 | |
| 307 | def test_aiter_dynamic_mxfp4_quant_determinism(): |
| 308 | """dynamic_mxfp4_quant is bitwise deterministic across repeated runs.""" |
| 309 | _assert_aiter_supported() |
| 310 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 311 | |
| 312 | torch.set_default_device("cuda") |
| 313 | torch.manual_seed(7) |
| 314 | x = torch.randn(128, 256, dtype=torch.bfloat16) |
| 315 | |
| 316 | _assert_deterministic(dynamic_mxfp4_quant, x, n_runs=4) |
| 317 | |
| 318 | |
| 319 | # gfx950 hardware FP4 GEMM tests ------------------------------------------ |
| 320 | |
| 321 | |
| 322 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 323 | @pytest.mark.parametrize( |
| 324 | "shape", |
| 325 | [ |
| 326 | (64, 128, 64), # small square-ish |
| 327 | (128, 256, 128), # medium |
| 328 | (128, 4096, 4096), # Llama-7B hidden square |
| 329 | (256, 4096, 11008), # Llama-7B FFN |
| 330 | (64, 8192, 28672), # Llama-70B FFN |
| 331 | ], |
| 332 | ) |
| 333 | def test_aiter_fp4_gemm_a4w4_accuracy(shape): |
| 334 | """AITER A4W4 FP4 GEMM output is close to matmul on dequantized weights. |
| 335 | |
| 336 | Tests gfx950 parity with B200 test_nvfp4_scaled_mm.py: E2M1 x E2M1 GEMM |
| 337 | with block-scaled quantization, multiple Llama-class shapes. |
| 338 | Requires gfx950 hardware; gracefully skips elsewhere. |
| 339 | """ |
| 340 | atol = GEMM_ATOL |
| 341 | rtol = GEMM_RTOL |
| 342 | pass_rate = A4W4_GEMM_PASS_RATES[shape] |
| 343 | _assert_aiter_supported() |
| 344 | from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 |
| 345 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 346 | |
| 347 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 348 | quant_dequant_mxfp4, |
| 349 | ) |
| 350 | |
| 351 | torch.set_default_device("cuda") |
| 352 | torch.manual_seed(0) |
| 353 | M, K, N = shape |
| 354 | |
| 355 | A = torch.randn(M, K, dtype=torch.bfloat16) |
| 356 | B = torch.randn(N, K, dtype=torch.bfloat16) # weight stored as (N, K) |
| 357 | |
| 358 | # Quantize |
| 359 | A_fp4, A_scale = dynamic_mxfp4_quant(A) |
| 360 | B_fp4, B_scale = dynamic_mxfp4_quant(B) |
| 361 | |
| 362 | # FP4 GEMM computes A_fp4 (M, K) @ B_fp4.T (K, N) into (M, N). |
| 363 | out = gemm_afp4wfp4(A_fp4, B_fp4, A_scale, B_scale) |
| 364 | |
| 365 | assert out.shape == (M, N), f"Expected ({M}, {N}), got {out.shape}" |
| 366 | assert out.dtype == torch.bfloat16 |
| 367 | assert not torch.any(torch.isnan(out)) |
| 368 | |
| 369 | # Reference: matmul on dequantized FP4 inputs (quark roundtrip) |
| 370 | A_dq = quant_dequant_mxfp4(A) |
| 371 | B_dq = quant_dequant_mxfp4(B) |
| 372 | ref = torch.matmul(A_dq.float(), B_dq.t().float()) |
| 373 | |
| 374 | # FP4 has only 3 mantissa bits; large K amplifies the accumulation error. |
| 375 | # These budgets were measured on gfx950 and keep p99/p999 drift visible |
| 376 | # without pretending the kernel is close to BF16 matmul everywhere. |
| 377 | _print_close_stats( |
| 378 | f"a4w4_gemm shape={shape}", |
| 379 | out.float(), |
| 380 | ref, |
| 381 | atol=atol, |
| 382 | rtol=rtol, |
| 383 | pass_rate=pass_rate, |
| 384 | max_atol=GEMM_MAX_VIOLATION_FACTOR * atol, |
| 385 | ) |
| 386 | _assert_accurate( |
| 387 | out.float(), |
| 388 | ref, |
| 389 | atol=atol, |
| 390 | rtol=rtol, |
| 391 | pass_rate=pass_rate, |
| 392 | max_violation_factor=GEMM_MAX_VIOLATION_FACTOR, |
| 393 | ) |
| 394 | |
| 395 | |
| 396 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 397 | @pytest.mark.parametrize( |
| 398 | "shape", |
| 399 | PRESHUFFLED_SHAPES, |
| 400 | ) |
| 401 | def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape): |
| 402 | """The preshuffled FP4 GEMM path should run on tuned shapes with the same |
| 403 | packed inputs vLLM uses in production. |
| 404 | |
| 405 | This kernel only advertises a narrow tuned `(N, K)` set on gfx950. The |
| 406 | important contract here is that those tuned shapes execute with the |
| 407 | production-style shuffled weights and stay deterministic. |
| 408 | """ |
| 409 | _assert_aiter_supported() |
| 410 | from aiter import per_1x32_f4_quant_hip |
| 411 | from aiter.ops.shuffle import shuffle_weight |
| 412 | from aiter.ops.triton.gemm_afp4wfp4 import ( |
| 413 | gemm_afp4wfp4_preshuffle, |
| 414 | ) |
| 415 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 416 | |
| 417 | from vllm._aiter_ops import rocm_aiter_ops |
| 418 | |
| 419 | torch.set_default_device("cuda") |
| 420 | torch.manual_seed(1) |
| 421 | M, K, N = shape |
| 422 | |
| 423 | assert M <= 64 |
| 424 | # The predicate takes K as packed bytes, matching `weight.shape[1]`. |
| 425 | assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K // 2) |
| 426 | |
| 427 | A = torch.randn(M, K, dtype=torch.bfloat16) |
| 428 | B = torch.randn(N, K, dtype=torch.bfloat16) |
| 429 | |
| 430 | B_fp4, B_scale = dynamic_mxfp4_quant(B) |
| 431 | A_q, A_s = per_1x32_f4_quant_hip(A, shuffle=M >= 32) |
| 432 | |
| 433 | scale_rows, scale_cols = B_scale.shape |
| 434 | B_scale = ( |
| 435 | B_scale.view(scale_rows // 32, 2, 16, scale_cols // 8, 2, 4, 1) |
| 436 | .permute(0, 3, 5, 2, 4, 1, 6) |
| 437 | .contiguous() |
| 438 | .view(scale_rows, scale_cols) |
| 439 | ) |
| 440 | B_fp4 = shuffle_weight(B_fp4, layout=(16, 16)) |
| 441 | |
| 442 | if M >= 32: |
| 443 | A_s = A_s.contiguous().view(torch.uint8).reshape(A_s.shape[0] // 32, -1) |
| 444 | else: |
| 445 | A_s = A_s[:M, ...].contiguous().view(torch.uint8) |
| 446 | B_scale = B_scale.contiguous().view(torch.uint8).reshape(B_scale.shape[0] // 32, -1) |
| 447 | |
| 448 | def run_preshuffled() -> torch.Tensor: |
| 449 | y = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) |
| 450 | return gemm_afp4wfp4_preshuffle( |
| 451 | A_q.contiguous().view(torch.uint8), |
| 452 | B_fp4.contiguous().view(torch.uint8).reshape(B_fp4.shape[0] // 16, -1), |
| 453 | A_s, |
| 454 | B_scale, |
| 455 | torch.bfloat16, |
| 456 | y, |
| 457 | ) |
| 458 | |
| 459 | out = run_preshuffled() |
| 460 | assert out.shape == (M, N) |
| 461 | assert out.dtype == torch.bfloat16 |
| 462 | assert torch.isfinite(out).all() |
| 463 | |
| 464 | _assert_deterministic(run_preshuffled, n_runs=3) |
| 465 | |
| 466 | |
| 467 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 468 | def test_aiter_fp4_gemm_a4w4_determinism(): |
| 469 | """AITER FP4 A4W4 GEMM is bitwise deterministic across repeated runs. |
| 470 | |
| 471 | Targets parity with B200 FP4 determinism requirements. |
| 472 | Requires gfx950. |
| 473 | """ |
| 474 | _assert_aiter_supported() |
| 475 | from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 |
| 476 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 477 | |
| 478 | torch.set_default_device("cuda") |
| 479 | torch.manual_seed(2) |
| 480 | |
| 481 | M, K, N = 128, 256, 128 |
| 482 | A = torch.randn(M, K, dtype=torch.bfloat16) |
| 483 | B = torch.randn(N, K, dtype=torch.bfloat16) |
| 484 | |
| 485 | A_fp4, A_scale = dynamic_mxfp4_quant(A) |
| 486 | B_fp4, B_scale = dynamic_mxfp4_quant(B) |
| 487 | |
| 488 | def run_gemm(): |
| 489 | return gemm_afp4wfp4(A_fp4, B_fp4, A_scale, B_scale) |
| 490 | |
| 491 | _assert_deterministic(run_gemm, n_runs=4) |
| 492 | |
| 493 | |
| 494 | # gfx950 hardware FP4 dynamic quantization -------------------------------- |
| 495 | |
| 496 | |
| 497 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 498 | @pytest.mark.parametrize( |
| 499 | "shape", |
| 500 | [ |
| 501 | (128, 256), |
| 502 | (256, 4096), |
| 503 | (64, 14336), |
| 504 | ], |
| 505 | ) |
| 506 | def test_aiter_hardware_fp4_dynamic_quant_format(shape): |
| 507 | """Aiter hardware FP4 dynamic quant produces correct output format. |
| 508 | |
| 509 | Tests gfx950 hardware-accelerated FP4 quantization (OCP MXFP4 E2M1). |
| 510 | Parity with B200 scaled_fp4_quant: block_size=32, packed uint8 output. |
| 511 | Requires gfx950. |
| 512 | """ |
| 513 | _assert_aiter_supported() |
| 514 | |
| 515 | from aiter import dynamic_per_group_scaled_quant_fp4 |
| 516 | |
| 517 | torch.set_default_device("cuda") |
| 518 | M, K = shape |
| 519 | group_size = 32 |
| 520 | |
| 521 | x = torch.randn(M, K, dtype=torch.bfloat16) |
| 522 | out_fp4 = torch.empty(M, K // 2, dtype=torch.uint8) |
| 523 | scales = torch.empty(M, K // group_size, dtype=torch.uint8) |
| 524 | |
| 525 | dynamic_per_group_scaled_quant_fp4(out_fp4, x, scales, group_size) |
| 526 | |
| 527 | assert out_fp4.shape == (M, K // 2), ( |
| 528 | f"Shape {shape}: expected fp4 ({M}, {K // 2}), got {out_fp4.shape}" |
| 529 | ) |
| 530 | assert scales.shape == (M, K // group_size), ( |
| 531 | f"Shape {shape}: expected scale ({M}, {K // group_size}), got {scales.shape}" |
| 532 | ) |
| 533 | |
| 534 | |
| 535 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 536 | def test_aiter_hardware_fp4_quant_vs_triton(): |
| 537 | """Hardware FP4 quant should dequantize back to the same MXFP4 values as |
| 538 | the Triton path. |
| 539 | |
| 540 | The packed bytes are not identical on gfx950, but with |
| 541 | ``shuffle_scale=False`` the dequantized results match tightly. |
| 542 | """ |
| 543 | _assert_aiter_supported() |
| 544 | |
| 545 | from aiter import dynamic_per_group_scaled_quant_fp4 |
| 546 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 547 | |
| 548 | torch.set_default_device("cuda") |
| 549 | torch.manual_seed(5) |
| 550 | |
| 551 | M, K = 128, 256 |
| 552 | group_size = 32 |
| 553 | x = torch.randn(M, K, dtype=torch.bfloat16) |
| 554 | |
| 555 | # Hardware path |
| 556 | out_hw = torch.empty(M, K // 2, dtype=torch.uint8) |
| 557 | scales_hw = torch.empty(M, K // group_size, dtype=torch.uint8) |
| 558 | dynamic_per_group_scaled_quant_fp4( |
| 559 | out_hw, x, scales_hw, group_size, shuffle_scale=False |
| 560 | ) |
| 561 | |
| 562 | # Triton path |
| 563 | out_triton, scales_triton = dynamic_mxfp4_quant(x) |
| 564 | |
| 565 | assert out_hw.shape == out_triton.shape |
| 566 | assert scales_hw.shape == scales_triton.shape |
| 567 | |
| 568 | dq_hw = dq_mxfp4_torch(out_hw, scales_hw, torch.bfloat16) |
| 569 | dq_triton = dq_mxfp4_torch(out_triton, scales_triton, torch.bfloat16) |
| 570 | |
| 571 | _print_close_stats( |
| 572 | "hardware_quant_vs_triton", |
| 573 | dq_hw.float(), |
| 574 | dq_triton.float(), |
| 575 | atol=0.25, |
| 576 | rtol=0.0, |
| 577 | ) |
| 578 | torch.testing.assert_close(dq_hw, dq_triton, atol=0.25, rtol=0.0) |
| 579 | |
| 580 | |
| 581 | # ROCm skinny GEMM FP4 tests ---------------------------------------------- |
| 582 | |
| 583 | |
| 584 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 585 | @pytest.mark.parametrize("M", [1, 2, 4, 8]) # decode / skinny batch sizes |
| 586 | @pytest.mark.parametrize("N, K", [(4096, 4096), (4096, 11008), (8192, 8192)]) |
| 587 | def test_aiter_fp4_gemm_skinny_shapes(M, N, K): |
| 588 | """FP4 GEMM accuracy for skinny (small-M) shapes (decode phase). |
| 589 | |
| 590 | Skinny GEMMs (M=1..8) are the bottleneck in decode phase. |
| 591 | Tests parity with B200 nvfp4_scaled_mm tests at decode batch sizes. |
| 592 | Requires gfx950. |
| 593 | """ |
| 594 | atol = GEMM_ATOL |
| 595 | rtol = GEMM_RTOL |
| 596 | pass_rate = SKINNY_GEMM_PASS_RATES[(N, K)] |
| 597 | _assert_aiter_supported() |
| 598 | from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 |
| 599 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 600 | |
| 601 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 602 | quant_dequant_mxfp4, |
| 603 | ) |
| 604 | |
| 605 | torch.set_default_device("cuda") |
| 606 | torch.manual_seed(0) |
| 607 | |
| 608 | A = torch.randn(M, K, dtype=torch.bfloat16) |
| 609 | B = torch.randn(N, K, dtype=torch.bfloat16) |
| 610 | |
| 611 | A_fp4, A_scale = dynamic_mxfp4_quant(A) |
| 612 | B_fp4, B_scale = dynamic_mxfp4_quant(B) |
| 613 | |
| 614 | out = gemm_afp4wfp4(A_fp4, B_fp4, A_scale, B_scale) |
| 615 | |
| 616 | assert out.shape == (M, N) |
| 617 | assert out.dtype == torch.bfloat16 |
| 618 | assert not torch.any(torch.isnan(out)) |
| 619 | |
| 620 | # Accumulate in FP32, then match the GEMM's BF16 output rounding. |
| 621 | A_dq = quant_dequant_mxfp4(A) |
| 622 | B_dq = quant_dequant_mxfp4(B) |
| 623 | ref = torch.matmul(A_dq.float(), B_dq.t().float()).to(torch.bfloat16).float() |
| 624 | |
| 625 | _print_close_stats( |
| 626 | f"skinny_gemm M={M} N={N} K={K}", |
| 627 | out.float(), |
| 628 | ref, |
| 629 | atol=atol, |
| 630 | rtol=rtol, |
| 631 | pass_rate=pass_rate, |
| 632 | max_atol=GEMM_MAX_VIOLATION_FACTOR * atol, |
| 633 | ) |
| 634 | _assert_accurate( |
| 635 | out.float(), |
| 636 | ref, |
| 637 | atol=atol, |
| 638 | rtol=rtol, |
| 639 | pass_rate=pass_rate, |
| 640 | max_violation_factor=GEMM_MAX_VIOLATION_FACTOR, |
| 641 | ) |
| 642 | |
| 643 | |
| 644 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 645 | @pytest.mark.parametrize( |
| 646 | ("shape", "supported"), |
| 647 | [ |
| 648 | ((2560, 80), True), # shared_expert gate/up_proj (in=2560 -> sn=80) |
| 649 | ((1280, 80), True), # gate/up_proj under TP=2 (out sharded) |
| 650 | ((2560, 8), True), # minimal aligned columns |
| 651 | ((32, 8), True), # minimal aligned rows and columns |
| 652 | ((2560, 20), False), # shared_expert down_proj (in=640 -> sn=20) |
| 653 | ((2560, 10), False), # down_proj under TP=2 (in 640->320 -> sn=10) |
| 654 | ((30, 8), False), # rows not a multiple of 32 |
| 655 | ], |
| 656 | ) |
| 657 | def test_asm_fp4_scale_swizzle_supported_shape_rules(shape, supported): |
| 658 | """The ASM swizzle needs rows % 32 == 0 and columns % 8 == 0.""" |
| 659 | from vllm.model_executor.kernels.linear.mxfp4.aiter import ( |
| 660 | _asm_fp4_scale_swizzle_supported, |
| 661 | ) |
| 662 | |
| 663 | weight_scale = torch.empty(shape, dtype=torch.uint8, device="cpu") |
| 664 | assert _asm_fp4_scale_swizzle_supported(weight_scale) is supported |
| 665 | |
| 666 | |
| 667 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 668 | def test_asm_fp4_scale_swizzle_rejects_non_2d(): |
| 669 | from vllm.model_executor.kernels.linear.mxfp4.aiter import ( |
| 670 | _asm_fp4_scale_swizzle_supported, |
| 671 | ) |
| 672 | |
| 673 | assert not _asm_fp4_scale_swizzle_supported( |
| 674 | torch.empty(2560, dtype=torch.uint8, device="cpu") |
| 675 | ) |
| 676 | assert not _asm_fp4_scale_swizzle_supported( |
| 677 | torch.empty(80, 2, 16, dtype=torch.uint8, device="cpu") |
| 678 | ) |
| 679 | |
| 680 | |
| 681 | @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") |
| 682 | def test_aiter_mxfp4_process_weights_falls_back_to_triton_for_misaligned_scale(): |
| 683 | """A misaligned weight_scale must disable ASM and take the Triton path.""" |
| 684 | from torch.nn.parameter import Parameter |
| 685 | |
| 686 | from vllm.model_executor.kernels.linear.mxfp4.aiter import AiterMxfp4LinearKernel |
| 687 | |
| 688 | kernel = object.__new__(AiterMxfp4LinearKernel) |
| 689 | kernel.use_asm_gemm = True |
| 690 | kernel.out_dtype = torch.bfloat16 |
| 691 | |
| 692 | layer = torch.nn.Module() |
| 693 | layer.weight_scale = Parameter( |
| 694 | torch.zeros(2560, 10, dtype=torch.uint8, device="cpu"), requires_grad=False |
| 695 | ) |
| 696 | |
| 697 | kernel.process_weights_after_loading(layer) |
| 698 | |
| 699 | assert kernel.use_asm_gemm is False |
| 700 | # Triton path stores the transposed, contiguous scale. |
| 701 | assert tuple(layer.weight_scale.shape) == (10, 2560) |
| 702 | assert layer.weight_scale.is_contiguous() |