| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import csv |
| 5 | import importlib |
| 6 | import importlib.util |
| 7 | import os |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from tests.utils import TestFP8Layer |
| 13 | from vllm._aiter_ops import rocm_aiter_ops |
| 14 | from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( |
| 15 | AiterHipbMMPerTokenFp8ScaledMMLinearKernel, |
| 16 | ) |
| 17 | from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( |
| 18 | FP8ScaledMMLinearLayerConfig, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 21 | get_fp8_min_max, |
| 22 | kFp8DynamicTokenSym, |
| 23 | kFp8StaticChannelSym, |
| 24 | kFp8StaticTensorSym, |
| 25 | ) |
| 26 | from vllm.platforms import current_platform |
| 27 | |
| 28 | aiter_available = importlib.util.find_spec("aiter") is not None |
| 29 | |
| 30 | pytestmark = [ |
| 31 | pytest.mark.skipif( |
| 32 | not ( |
| 33 | current_platform.is_rocm() |
| 34 | and current_platform.supports_fp8() |
| 35 | and aiter_available |
| 36 | ), |
| 37 | reason="Requires ROCm + FP8 support + aiter", |
| 38 | ), |
| 39 | pytest.mark.usefixtures("default_vllm_config"), |
| 40 | ] |
| 41 | |
| 42 | |
| 43 | @pytest.fixture |
| 44 | def enable_hipb_mm_kernel(monkeypatch: pytest.MonkeyPatch): |
| 45 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 46 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1") |
| 47 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "1") |
| 48 | rocm_aiter_ops.refresh_env_variables() |
| 49 | yield |
| 50 | rocm_aiter_ops.refresh_env_variables() |
| 51 | |
| 52 | |
| 53 | def _make_config( |
| 54 | *, |
| 55 | weight_quant_key=kFp8StaticChannelSym, |
| 56 | out_dtype: torch.dtype = torch.bfloat16, |
| 57 | weight_shape: tuple[int, int] = (512, 4096), |
| 58 | ) -> FP8ScaledMMLinearLayerConfig: |
| 59 | return FP8ScaledMMLinearLayerConfig( |
| 60 | weight_quant_key=weight_quant_key, |
| 61 | activation_quant_key=kFp8DynamicTokenSym, |
| 62 | weight_shape=weight_shape, |
| 63 | input_dtype=torch.bfloat16, |
| 64 | out_dtype=out_dtype, |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | def _find_csv_row(path: str, m: int, n: int, k: int) -> dict | None: |
| 69 | if not os.path.exists(path): |
| 70 | return None |
| 71 | |
| 72 | with open(path, newline="") as f: |
| 73 | reader = csv.DictReader(f, skipinitialspace=True) |
| 74 | for row in reader: |
| 75 | try: |
| 76 | if ( |
| 77 | int(row.get("m", -1)) == m |
| 78 | and int(row.get("n", -1)) == n |
| 79 | and int(row.get("k", -1)) == k |
| 80 | ): |
| 81 | return dict(row) |
| 82 | except (TypeError, ValueError): |
| 83 | continue |
| 84 | return None |
| 85 | |
| 86 | |
| 87 | def _skip_if_no_hipb_mm_solution(exc: RuntimeError) -> None: |
| 88 | if "hipblasLtMatmulAlgoGetHeuristic found 0 valid solutions" in str(exc): |
| 89 | pytest.skip( |
| 90 | "hipb_mm bpreshuffle path has no valid hipBLASLt solution on " |
| 91 | "this ROCm stack." |
| 92 | ) |
| 93 | |
| 94 | |
| 95 | def _check_bpreshuffle_runtime_support(weight_shape: tuple[int, int], num_tokens: int): |
| 96 | import aiter |
| 97 | from aiter.ops.shuffle import shuffle_weight |
| 98 | |
| 99 | x = torch.randn(num_tokens, weight_shape[1], dtype=torch.bfloat16, device="cuda") |
| 100 | w = torch.randn(weight_shape, dtype=torch.bfloat16, device="cuda") |
| 101 | |
| 102 | aiter.hipb_create_extension() |
| 103 | x_q, x_scale = aiter.pertoken_quant(x, quant_dtype=current_platform.fp8_dtype()) |
| 104 | w_q, w_scale = aiter.pertoken_quant(w, quant_dtype=current_platform.fp8_dtype()) |
| 105 | |
| 106 | try: |
| 107 | aiter.hipb_mm( |
| 108 | x_q, |
| 109 | shuffle_weight(w_q, layout=(16, 16)).t(), |
| 110 | solution_index=-1, |
| 111 | out_dtype=torch.bfloat16, |
| 112 | scaleA=x_scale, |
| 113 | scaleB=w_scale.t().contiguous(), |
| 114 | scaleOut=None, |
| 115 | bpreshuffle=True, |
| 116 | ) |
| 117 | except RuntimeError as exc: |
| 118 | _skip_if_no_hipb_mm_solution(exc) |
| 119 | raise |
| 120 | |
| 121 | |
| 122 | def test_hipb_mm_kernel_requires_hipbmm_flag(monkeypatch: pytest.MonkeyPatch): |
| 123 | # The kernel rejects when `is_hip_fp8bmm_enabled()` is False. That helper |
| 124 | # requires AITER + AITER_LINEAR + MI3xx, so dropping AITER_LINEAR exercises |
| 125 | # the rejection branch. |
| 126 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 127 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "0") |
| 128 | monkeypatch.delenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", raising=False) |
| 129 | rocm_aiter_ops.refresh_env_variables() |
| 130 | |
| 131 | is_supported, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.is_supported() |
| 132 | |
| 133 | assert not is_supported |
| 134 | assert reason == ( |
| 135 | "requires setting `VLLM_ROCM_USE_AITER=1`, " |
| 136 | "`VLLM_ROCM_USE_AITER_LINEAR=1`, " |
| 137 | "and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`." |
| 138 | ) |
| 139 | |
| 140 | |
| 141 | def test_hipb_mm_flag_enables_hip_online_tuning( |
| 142 | monkeypatch: pytest.MonkeyPatch, |
| 143 | ): |
| 144 | import vllm.envs as envs_mod |
| 145 | import vllm.platforms.rocm as rocm_mod |
| 146 | |
| 147 | # The rocm.py gate requires all three AITER flags (and MI3xx) to auto-set |
| 148 | # HIP_ONLINE_TUNING. |
| 149 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 150 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1") |
| 151 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "1") |
| 152 | |
| 153 | try: |
| 154 | importlib.reload(envs_mod) |
| 155 | importlib.reload(rocm_mod) |
| 156 | assert envs_mod.VLLM_ROCM_USE_AITER |
| 157 | assert envs_mod.VLLM_ROCM_USE_AITER_LINEAR |
| 158 | assert envs_mod.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM |
| 159 | assert os.environ.get("HIP_ONLINE_TUNING") == "1" |
| 160 | finally: |
| 161 | monkeypatch.undo() |
| 162 | os.environ.pop("HIP_ONLINE_TUNING", None) |
| 163 | importlib.reload(envs_mod) |
| 164 | importlib.reload(rocm_mod) |
| 165 | rocm_aiter_ops.refresh_env_variables() |
| 166 | |
| 167 | |
| 168 | def test_hipb_mm_kernel_can_implement_success(enable_hipb_mm_kernel): |
| 169 | can_implement, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.can_implement( |
| 170 | _make_config() |
| 171 | ) |
| 172 | |
| 173 | assert can_implement |
| 174 | assert reason is None |
| 175 | |
| 176 | |
| 177 | @pytest.mark.parametrize( |
| 178 | ("config", "expected_reason"), |
| 179 | [ |
| 180 | ( |
| 181 | _make_config(weight_quant_key=kFp8StaticTensorSym), |
| 182 | "requires per token activation scales and per channel weight scales.", |
| 183 | ), |
| 184 | ( |
| 185 | _make_config(out_dtype=torch.float16), |
| 186 | "requires bfloat16 output dtype.", |
| 187 | ), |
| 188 | ( |
| 189 | _make_config(weight_shape=(8, 4090)), |
| 190 | "requires N >= 16 and both N and K divisible by 16, " |
| 191 | "received N=8 and K=4090.", |
| 192 | ), |
| 193 | ], |
| 194 | ) |
| 195 | def test_hipb_mm_kernel_can_implement_rejects_unsupported_configs( |
| 196 | enable_hipb_mm_kernel, |
| 197 | config: FP8ScaledMMLinearLayerConfig, |
| 198 | expected_reason: str, |
| 199 | ): |
| 200 | can_implement, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.can_implement( |
| 201 | config |
| 202 | ) |
| 203 | |
| 204 | assert not can_implement |
| 205 | assert reason == expected_reason |
| 206 | |
| 207 | |
| 208 | def test_hipb_mm_kernel_process_weights_after_loading_shuffles_weights( |
| 209 | enable_hipb_mm_kernel, |
| 210 | ): |
| 211 | weight_shape = (512, 4096) |
| 212 | kernel = AiterHipbMMPerTokenFp8ScaledMMLinearKernel( |
| 213 | _make_config(weight_shape=weight_shape), |
| 214 | layer_param_names=("weight", "weight_scale", "input_scale", "input_scale_ub"), |
| 215 | ) |
| 216 | |
| 217 | layer = torch.nn.Module() |
| 218 | layer.weight = torch.nn.Parameter( |
| 219 | torch.rand(weight_shape, device="cuda").to(current_platform.fp8_dtype()).t(), |
| 220 | requires_grad=False, |
| 221 | ) |
| 222 | layer.weight_scale = torch.nn.Parameter( |
| 223 | torch.rand((weight_shape[0], 1), dtype=torch.float32, device="cuda"), |
| 224 | requires_grad=False, |
| 225 | ) |
| 226 | layer.input_scale = None |
| 227 | layer.input_scale_ub = None |
| 228 | |
| 229 | original_weight = layer.weight.detach().clone() |
| 230 | original_weight_scale = layer.weight_scale.detach().clone() |
| 231 | |
| 232 | kernel.process_weights_after_loading(layer) |
| 233 | |
| 234 | # process_weights_after_loading now pre-applies the transposes that used |
| 235 | # to live in _rocm_aiter_hipb_mm_fp8_impl, so the stored weight is the |
| 236 | # shuffled tensor with a trailing `.t()` view, and the stored weight scale |
| 237 | # is its transposed-contiguous form. |
| 238 | expected_weight = rocm_aiter_ops.shuffle_weight( |
| 239 | original_weight.t().contiguous() |
| 240 | ).t() |
| 241 | torch.testing.assert_close(layer.weight, expected_weight) |
| 242 | |
| 243 | expected_weight_scale = original_weight_scale.t().contiguous() |
| 244 | torch.testing.assert_close(layer.weight_scale, expected_weight_scale) |
| 245 | |
| 246 | |
| 247 | def test_hipb_mm_kernel_forward_matches_raw_aiter_hipb_mm(enable_hipb_mm_kernel): |
| 248 | import aiter |
| 249 | |
| 250 | weight_shape = (512, 4096) |
| 251 | _check_bpreshuffle_runtime_support(weight_shape, num_tokens=32) |
| 252 | |
| 253 | layer = TestFP8Layer( |
| 254 | weight_shape=weight_shape, |
| 255 | activation_quant_key=kFp8DynamicTokenSym, |
| 256 | weight_quant_key=kFp8StaticChannelSym, |
| 257 | input_dtype=torch.bfloat16, |
| 258 | out_dtype=torch.bfloat16, |
| 259 | device=torch.device("cuda"), |
| 260 | force_kernel=AiterHipbMMPerTokenFp8ScaledMMLinearKernel, |
| 261 | ) |
| 262 | |
| 263 | # hipb_mm uses a transposed-result GEMM internally, so the flattened token |
| 264 | # count becomes the effective N dimension passed into hipBLASLt. Keep it |
| 265 | # aligned to avoid heuristic failures for tiny N. |
| 266 | x = torch.randn(2, 16, weight_shape[1], dtype=torch.bfloat16, device="cuda") |
| 267 | bias = torch.randn(weight_shape[0], dtype=torch.bfloat16, device="cuda") |
| 268 | |
| 269 | try: |
| 270 | out = layer(x, bias) |
| 271 | except RuntimeError as exc: |
| 272 | _skip_if_no_hipb_mm_solution(exc) |
| 273 | raise |
| 274 | |
| 275 | x_2d = x.view(-1, x.shape[-1]) |
| 276 | x_q, x_scale = layer.kernel.quant_fp8( |
| 277 | x_2d, |
| 278 | layer.input_scale, |
| 279 | layer.input_scale_ub, |
| 280 | ) |
| 281 | try: |
| 282 | # process_weights_after_loading already applies the trailing `.t()` on |
| 283 | # the shuffled weight and the `.t().contiguous()` on the weight scale, |
| 284 | # so the raw aiter call uses them directly. |
| 285 | expected = aiter.hipb_mm( |
| 286 | x_q, |
| 287 | layer.weight, |
| 288 | solution_index=-1, |
| 289 | bias=bias, |
| 290 | out_dtype=torch.bfloat16, |
| 291 | scaleA=x_scale, |
| 292 | scaleB=layer.weight_scale, |
| 293 | scaleOut=None, |
| 294 | bpreshuffle=True, |
| 295 | ).view(*out.shape) |
| 296 | except RuntimeError as exc: |
| 297 | _skip_if_no_hipb_mm_solution(exc) |
| 298 | raise |
| 299 | |
| 300 | assert isinstance(layer.kernel, AiterHipbMMPerTokenFp8ScaledMMLinearKernel) |
| 301 | assert out.shape == (2, 16, weight_shape[0]) |
| 302 | torch.testing.assert_close(out, expected) |
| 303 | |
| 304 | |
| 305 | def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel): |
| 306 | """Kernel output should match a dequantized fp32 reference within |
| 307 | fp8 per-token / per-channel quantization noise.""" |
| 308 | weight_shape = (512, 4096) # (N, K) |
| 309 | num_tokens = 32 |
| 310 | _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens) |
| 311 | |
| 312 | fp8_dtype = current_platform.fp8_dtype() |
| 313 | fp8_max = get_fp8_min_max()[1] |
| 314 | device = torch.device("cuda") |
| 315 | |
| 316 | # Build a bf16 weight and quantize per output channel (one scale per row). |
| 317 | w_bf16 = torch.randn(weight_shape, dtype=torch.bfloat16, device=device) |
| 318 | w_amax = w_bf16.abs().amax(dim=1, keepdim=True).to(torch.float32) |
| 319 | w_scale = (w_amax / fp8_max).clamp(min=1e-12) |
| 320 | w_fp8 = (w_bf16.to(torch.float32) / w_scale).clamp(-fp8_max, fp8_max).to(fp8_dtype) |
| 321 | w_dequant = w_fp8.to(torch.float32) * w_scale |
| 322 | |
| 323 | bias = torch.randn(weight_shape[0], dtype=torch.bfloat16, device=device) |
| 324 | |
| 325 | layer = torch.nn.Module() |
| 326 | # Pre-`process_weights_after_loading` convention: weight stored as the |
| 327 | # `[K, N]` view of the fp8 tensor. |
| 328 | layer.weight = torch.nn.Parameter(w_fp8.t(), requires_grad=False) |
| 329 | layer.weight_scale = torch.nn.Parameter(w_scale, requires_grad=False) |
| 330 | layer.input_scale = None |
| 331 | layer.input_scale_ub = None |
| 332 | |
| 333 | kernel = AiterHipbMMPerTokenFp8ScaledMMLinearKernel( |
| 334 | _make_config(weight_shape=weight_shape), |
| 335 | layer_param_names=("weight", "weight_scale", "input_scale", "input_scale_ub"), |
| 336 | ) |
| 337 | kernel.process_weights_after_loading(layer) |
| 338 | |
| 339 | x = torch.randn(num_tokens, weight_shape[1], dtype=torch.bfloat16, device=device) |
| 340 | |
| 341 | try: |
| 342 | out = kernel.apply_weights(layer, x, bias) |
| 343 | except RuntimeError as exc: |
| 344 | _skip_if_no_hipb_mm_solution(exc) |
| 345 | raise |
| 346 | |
| 347 | # Reference: quantize x per-token the same way the kernel does, then run |
| 348 | # the matmul in fp32 against the dequantized weight. This isolates plumbing |
| 349 | # / reduction bugs from inherent fp8 quantization noise. |
| 350 | x_amax = x.abs().amax(dim=1, keepdim=True).to(torch.float32) |
| 351 | x_scale_ref = (x_amax / fp8_max).clamp(min=1e-12) |
| 352 | x_q = (x.to(torch.float32) / x_scale_ref).clamp(-fp8_max, fp8_max).to(fp8_dtype) |
| 353 | x_dequant = x_q.to(torch.float32) * x_scale_ref |
| 354 | expected = (x_dequant @ w_dequant.t() + bias.to(torch.float32)).to(torch.bfloat16) |
| 355 | |
| 356 | assert out.shape == (num_tokens, weight_shape[0]) |
| 357 | # K=4096 fp8 reduction leaves room for accumulation order drift and |
| 358 | # catastrophic cancellation on near-zero outputs; tolerances are loose |
| 359 | # enough to absorb that but tight enough to catch wrong layouts, missing |
| 360 | # bias, swapped scales, etc. |
| 361 | torch.testing.assert_close(out, expected, atol=5.0, rtol=0.1) |
| 362 | |
| 363 | |
| 364 | def test_hipb_mm_kernel_online_tuning_writes_csv( |
| 365 | enable_hipb_mm_kernel, |
| 366 | monkeypatch: pytest.MonkeyPatch, |
| 367 | tmp_path, |
| 368 | ): |
| 369 | weight_shape = (256, 4096) |
| 370 | cache_file = tmp_path / "hip_online_tuning_res.csv" |
| 371 | |
| 372 | _check_bpreshuffle_runtime_support(weight_shape, num_tokens=16) |
| 373 | |
| 374 | monkeypatch.setenv("HIP_ONLINE_TUNING", "1") |
| 375 | monkeypatch.chdir(tmp_path) |
| 376 | |
| 377 | layer = TestFP8Layer( |
| 378 | weight_shape=weight_shape, |
| 379 | activation_quant_key=kFp8DynamicTokenSym, |
| 380 | weight_quant_key=kFp8StaticChannelSym, |
| 381 | input_dtype=torch.bfloat16, |
| 382 | out_dtype=torch.bfloat16, |
| 383 | device=torch.device("cuda"), |
| 384 | force_kernel=AiterHipbMMPerTokenFp8ScaledMMLinearKernel, |
| 385 | ) |
| 386 | |
| 387 | # The effective heuristic N dimension is the flattened token count. |
| 388 | x = torch.randn(16, weight_shape[1], dtype=torch.bfloat16, device="cuda") |
| 389 | try: |
| 390 | out = layer(x) |
| 391 | except RuntimeError as exc: |
| 392 | _skip_if_no_hipb_mm_solution(exc) |
| 393 | raise |
| 394 | torch.accelerator.synchronize() |
| 395 | |
| 396 | assert out.shape == (16, weight_shape[0]) |
| 397 | assert cache_file.exists() |
| 398 | |
| 399 | # hipb_mm records the internal GEMM dimensions used by hipBLASLt after its |
| 400 | # transposed-result transformation. |
| 401 | row = _find_csv_row( |
| 402 | str(cache_file), |
| 403 | m=weight_shape[0], |
| 404 | n=x.shape[0], |
| 405 | k=weight_shape[1], |
| 406 | ) |
| 407 | assert row is not None |