From 165b7864d0e3208e09c33ed4ce459d2bca3c99ea Mon Sep 17 00:00:00 2001 From: Han Lin Date: Fri, 5 Jun 2026 14:45:36 +0800 Subject: [PATCH] [ROCM] [FEAT] Integrate Aiter hipBLASLt GEMM online tuning (#40426) Signed-off-by: hanlin12 Signed-off-by: Han Lin Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: TJian --- .../aiter/test_aiter_hipb_mm_linear_kernel.py | 406 ++++++++++++++++++ vllm/_aiter_ops.py | 73 ++++ vllm/envs.py | 4 + .../model_executor/kernels/linear/__init__.py | 3 + .../kernels/linear/scaled_mm/aiter.py | 93 ++++ vllm/platforms/rocm.py | 12 + 6 files changed, 591 insertions(+) create mode 100644 tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py diff --git a/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py b/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py new file mode 100644 index 00000000000..92017e95cb7 --- /dev/null +++ b/tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import csv +import importlib +import importlib.util +import os + +import pytest +import torch + +from tests.utils import TestFP8Layer +from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, +) +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, + kFp8StaticChannelSym, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform + +aiter_available = importlib.util.find_spec("aiter") is not None + +pytestmark = [ + pytest.mark.skipif( + not ( + current_platform.is_rocm() + and current_platform.supports_fp8() + and aiter_available + ), + reason="Requires ROCm + FP8 support + aiter", + ), + pytest.mark.usefixtures("default_vllm_config"), +] + + +@pytest.fixture +def enable_hipb_mm_kernel(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "1") + rocm_aiter_ops.refresh_env_variables() + yield + rocm_aiter_ops.refresh_env_variables() + + +def _make_config( + *, + weight_quant_key=kFp8StaticChannelSym, + out_dtype: torch.dtype = torch.bfloat16, + weight_shape: tuple[int, int] = (512, 4096), +) -> FP8ScaledMMLinearLayerConfig: + return FP8ScaledMMLinearLayerConfig( + weight_quant_key=weight_quant_key, + activation_quant_key=kFp8DynamicTokenSym, + weight_shape=weight_shape, + input_dtype=torch.bfloat16, + out_dtype=out_dtype, + ) + + +def _find_csv_row(path: str, m: int, n: int, k: int) -> dict | None: + if not os.path.exists(path): + return None + + with open(path, newline="") as f: + reader = csv.DictReader(f, skipinitialspace=True) + for row in reader: + try: + if ( + int(row.get("m", -1)) == m + and int(row.get("n", -1)) == n + and int(row.get("k", -1)) == k + ): + return dict(row) + except (TypeError, ValueError): + continue + return None + + +def _skip_if_no_hipb_mm_solution(exc: RuntimeError) -> None: + if "hipblasLtMatmulAlgoGetHeuristic found 0 valid solutions" in str(exc): + pytest.skip( + "hipb_mm bpreshuffle path has no valid hipBLASLt solution on " + "this ROCm stack." + ) + + +def _check_bpreshuffle_runtime_support(weight_shape: tuple[int, int], num_tokens: int): + import aiter + from aiter.ops.shuffle import shuffle_weight + + x = torch.randn(num_tokens, weight_shape[1], dtype=torch.bfloat16, device="cuda") + w = torch.randn(weight_shape, dtype=torch.bfloat16, device="cuda") + + aiter.hipb_create_extension() + x_q, x_scale = aiter.pertoken_quant(x, quant_dtype=current_platform.fp8_dtype()) + w_q, w_scale = aiter.pertoken_quant(w, quant_dtype=current_platform.fp8_dtype()) + + try: + aiter.hipb_mm( + x_q, + shuffle_weight(w_q, layout=(16, 16)).t(), + solution_index=-1, + out_dtype=torch.bfloat16, + scaleA=x_scale, + scaleB=w_scale.t().contiguous(), + scaleOut=None, + bpreshuffle=True, + ) + except RuntimeError as exc: + _skip_if_no_hipb_mm_solution(exc) + raise + + +def test_hipb_mm_kernel_requires_hipbmm_flag(monkeypatch: pytest.MonkeyPatch): + # The kernel rejects when `is_hip_fp8bmm_enabled()` is False. That helper + # requires AITER + AITER_LINEAR + MI3xx, so dropping AITER_LINEAR exercises + # the rejection branch. + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "0") + monkeypatch.delenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", raising=False) + rocm_aiter_ops.refresh_env_variables() + + is_supported, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.is_supported() + + assert not is_supported + assert reason == ( + "requires setting `VLLM_ROCM_USE_AITER=1`, " + "`VLLM_ROCM_USE_AITER_LINEAR=1`, " + "and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`." + ) + + +def test_hipb_mm_flag_enables_hip_online_tuning( + monkeypatch: pytest.MonkeyPatch, +): + import vllm.envs as envs_mod + import vllm.platforms.rocm as rocm_mod + + # The rocm.py gate requires all three AITER flags (and MI3xx) to auto-set + # HIP_ONLINE_TUNING. + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "1") + + try: + importlib.reload(envs_mod) + importlib.reload(rocm_mod) + assert envs_mod.VLLM_ROCM_USE_AITER + assert envs_mod.VLLM_ROCM_USE_AITER_LINEAR + assert envs_mod.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM + assert os.environ.get("HIP_ONLINE_TUNING") == "1" + finally: + monkeypatch.undo() + os.environ.pop("HIP_ONLINE_TUNING", None) + importlib.reload(envs_mod) + importlib.reload(rocm_mod) + rocm_aiter_ops.refresh_env_variables() + + +def test_hipb_mm_kernel_can_implement_success(enable_hipb_mm_kernel): + can_implement, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.can_implement( + _make_config() + ) + + assert can_implement + assert reason is None + + +@pytest.mark.parametrize( + ("config", "expected_reason"), + [ + ( + _make_config(weight_quant_key=kFp8StaticTensorSym), + "requires per token activation scales and per channel weight scales.", + ), + ( + _make_config(out_dtype=torch.float16), + "requires bfloat16 output dtype.", + ), + ( + _make_config(weight_shape=(8, 4090)), + "requires N >= 16 and both N and K divisible by 16, " + "received N=8 and K=4090.", + ), + ], +) +def test_hipb_mm_kernel_can_implement_rejects_unsupported_configs( + enable_hipb_mm_kernel, + config: FP8ScaledMMLinearLayerConfig, + expected_reason: str, +): + can_implement, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.can_implement( + config + ) + + assert not can_implement + assert reason == expected_reason + + +def test_hipb_mm_kernel_process_weights_after_loading_shuffles_weights( + enable_hipb_mm_kernel, +): + weight_shape = (512, 4096) + kernel = AiterHipbMMPerTokenFp8ScaledMMLinearKernel( + _make_config(weight_shape=weight_shape), + layer_param_names=("weight", "weight_scale", "input_scale", "input_scale_ub"), + ) + + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.rand(weight_shape, device="cuda").to(current_platform.fp8_dtype()).t(), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + torch.rand((weight_shape[0], 1), dtype=torch.float32, device="cuda"), + requires_grad=False, + ) + layer.input_scale = None + layer.input_scale_ub = None + + original_weight = layer.weight.detach().clone() + original_weight_scale = layer.weight_scale.detach().clone() + + kernel.process_weights_after_loading(layer) + + # process_weights_after_loading now pre-applies the transposes that used + # to live in _rocm_aiter_hipb_mm_fp8_impl, so the stored weight is the + # shuffled tensor with a trailing `.t()` view, and the stored weight scale + # is its transposed-contiguous form. + expected_weight = rocm_aiter_ops.shuffle_weight( + original_weight.t().contiguous() + ).t() + torch.testing.assert_close(layer.weight, expected_weight) + + expected_weight_scale = original_weight_scale.t().contiguous() + torch.testing.assert_close(layer.weight_scale, expected_weight_scale) + + +def test_hipb_mm_kernel_forward_matches_raw_aiter_hipb_mm(enable_hipb_mm_kernel): + import aiter + + weight_shape = (512, 4096) + _check_bpreshuffle_runtime_support(weight_shape, num_tokens=32) + + layer = TestFP8Layer( + weight_shape=weight_shape, + activation_quant_key=kFp8DynamicTokenSym, + weight_quant_key=kFp8StaticChannelSym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + device=torch.device("cuda"), + force_kernel=AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + ) + + # hipb_mm uses a transposed-result GEMM internally, so the flattened token + # count becomes the effective N dimension passed into hipBLASLt. Keep it + # aligned to avoid heuristic failures for tiny N. + x = torch.randn(2, 16, weight_shape[1], dtype=torch.bfloat16, device="cuda") + bias = torch.randn(weight_shape[0], dtype=torch.bfloat16, device="cuda") + + try: + out = layer(x, bias) + except RuntimeError as exc: + _skip_if_no_hipb_mm_solution(exc) + raise + + x_2d = x.view(-1, x.shape[-1]) + x_q, x_scale = layer.kernel.quant_fp8( + x_2d, + layer.input_scale, + layer.input_scale_ub, + ) + try: + # process_weights_after_loading already applies the trailing `.t()` on + # the shuffled weight and the `.t().contiguous()` on the weight scale, + # so the raw aiter call uses them directly. + expected = aiter.hipb_mm( + x_q, + layer.weight, + solution_index=-1, + bias=bias, + out_dtype=torch.bfloat16, + scaleA=x_scale, + scaleB=layer.weight_scale, + scaleOut=None, + bpreshuffle=True, + ).view(*out.shape) + except RuntimeError as exc: + _skip_if_no_hipb_mm_solution(exc) + raise + + assert isinstance(layer.kernel, AiterHipbMMPerTokenFp8ScaledMMLinearKernel) + assert out.shape == (2, 16, weight_shape[0]) + torch.testing.assert_close(out, expected) + + +def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel): + """Kernel output should match a dequantized fp32 reference within + fp8 per-token / per-channel quantization noise.""" + weight_shape = (512, 4096) # (N, K) + num_tokens = 32 + _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens) + + fp8_dtype = current_platform.fp8_dtype() + fp8_max = torch.finfo(fp8_dtype).max + device = torch.device("cuda") + + # Build a bf16 weight and quantize per output channel (one scale per row). + w_bf16 = torch.randn(weight_shape, dtype=torch.bfloat16, device=device) + w_amax = w_bf16.abs().amax(dim=1, keepdim=True).to(torch.float32) + w_scale = (w_amax / fp8_max).clamp(min=1e-12) + w_fp8 = (w_bf16.to(torch.float32) / w_scale).clamp(-fp8_max, fp8_max).to(fp8_dtype) + w_dequant = w_fp8.to(torch.float32) * w_scale + + bias = torch.randn(weight_shape[0], dtype=torch.bfloat16, device=device) + + layer = torch.nn.Module() + # Pre-`process_weights_after_loading` convention: weight stored as the + # `[K, N]` view of the fp8 tensor. + layer.weight = torch.nn.Parameter(w_fp8.t(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(w_scale, requires_grad=False) + layer.input_scale = None + layer.input_scale_ub = None + + kernel = AiterHipbMMPerTokenFp8ScaledMMLinearKernel( + _make_config(weight_shape=weight_shape), + layer_param_names=("weight", "weight_scale", "input_scale", "input_scale_ub"), + ) + kernel.process_weights_after_loading(layer) + + x = torch.randn(num_tokens, weight_shape[1], dtype=torch.bfloat16, device=device) + + try: + out = kernel.apply_weights(layer, x, bias) + except RuntimeError as exc: + _skip_if_no_hipb_mm_solution(exc) + raise + + # Reference: quantize x per-token the same way the kernel does, then run + # the matmul in fp32 against the dequantized weight. This isolates plumbing + # / reduction bugs from inherent fp8 quantization noise. + x_amax = x.abs().amax(dim=1, keepdim=True).to(torch.float32) + x_scale_ref = (x_amax / fp8_max).clamp(min=1e-12) + x_q = (x.to(torch.float32) / x_scale_ref).clamp(-fp8_max, fp8_max).to(fp8_dtype) + x_dequant = x_q.to(torch.float32) * x_scale_ref + expected = (x_dequant @ w_dequant.t() + bias.to(torch.float32)).to(torch.bfloat16) + + assert out.shape == (num_tokens, weight_shape[0]) + # K=4096 fp8 reduction leaves room for accumulation order drift and + # catastrophic cancellation on near-zero outputs; tolerances are loose + # enough to absorb that but tight enough to catch wrong layouts, missing + # bias, swapped scales, etc. + torch.testing.assert_close(out, expected, atol=5.0, rtol=0.1) + + +def test_hipb_mm_kernel_online_tuning_writes_csv( + enable_hipb_mm_kernel, + monkeypatch: pytest.MonkeyPatch, + tmp_path, +): + weight_shape = (256, 4096) + cache_file = tmp_path / "hip_online_tuning_res.csv" + + _check_bpreshuffle_runtime_support(weight_shape, num_tokens=16) + + monkeypatch.setenv("HIP_ONLINE_TUNING", "1") + monkeypatch.chdir(tmp_path) + + layer = TestFP8Layer( + weight_shape=weight_shape, + activation_quant_key=kFp8DynamicTokenSym, + weight_quant_key=kFp8StaticChannelSym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + device=torch.device("cuda"), + force_kernel=AiterHipbMMPerTokenFp8ScaledMMLinearKernel, + ) + + # The effective heuristic N dimension is the flattened token count. + x = torch.randn(16, weight_shape[1], dtype=torch.bfloat16, device="cuda") + try: + out = layer(x) + except RuntimeError as exc: + _skip_if_no_hipb_mm_solution(exc) + raise + torch.accelerator.synchronize() + + assert out.shape == (16, weight_shape[0]) + assert cache_file.exists() + + # hipb_mm records the internal GEMM dimensions used by hipBLASLt after its + # transposed-result transformation. + row = _find_csv_row( + str(cache_file), + m=weight_shape[0], + n=x.shape[0], + k=weight_shape[1], + ) + assert row is not None diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 5a8b690433c..eb12bedd7bf 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -27,6 +27,16 @@ except ImportError: # on ROCm the fp8_dtype always calls is_fp8_fnuz # which is a host op, so we cache it once here. FP8_DTYPE = current_platform.fp8_dtype() +_HIPB_MM_INITIALIZED_DEVICES: set[int] = set() + + +def _ensure_hipb_mm_extension_initialized() -> None: + import aiter + + device = torch.accelerator.current_device_index() + if device not in _HIPB_MM_INITIALIZED_DEVICES: + aiter.hipb_create_extension() + _HIPB_MM_INITIALIZED_DEVICES.add(device) def is_aiter_found() -> bool: @@ -625,6 +635,43 @@ def _rocm_aiter_preshuffled_per_token_w8a8_gemm_fake( return torch.empty(m, n, dtype=output_dtype, device=A.device) +def _rocm_aiter_hipb_mm_fp8_impl( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + from aiter import hipb_mm + + _ensure_hipb_mm_extension_initialized() + return hipb_mm( + A, + B, + solution_index=-1, + bias=bias, + out_dtype=output_dtype, + scaleA=As, + scaleB=Bs, + scaleOut=None, + bpreshuffle=True, + ) + + +def _rocm_aiter_hipb_mm_fp8_fake( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + m = A.shape[0] + n = B.shape[1] + return torch.empty(m, n, dtype=output_dtype, device=A.device) + + def _rocm_aiter_triton_gemm_a8w8_blockscale_impl( A: torch.Tensor, B: torch.Tensor, @@ -1308,6 +1355,7 @@ class rocm_aiter_ops: # TODO: Consolidate under _LINEAR_ENABLED _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM _FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM + _LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM # TODO: Consolidate under _LINEAR_ENABLED _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE @@ -1340,6 +1388,7 @@ class rocm_aiter_ops: cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM + cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS @@ -1512,6 +1561,13 @@ class rocm_aiter_ops: return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950() + @classmethod + @if_aiter_supported + def is_linear_hipbmm_enabled(cls) -> bool: + from vllm.platforms.rocm import on_mi3xx + + return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED + @classmethod @if_aiter_supported def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool: @@ -1668,6 +1724,12 @@ class rocm_aiter_ops: fake_impl=_rocm_aiter_preshuffled_per_token_w8a8_gemm_fake, ) + direct_register_custom_op( + op_name="rocm_aiter_hipb_mm_fp8", + op_func=_rocm_aiter_hipb_mm_fp8_impl, + fake_impl=_rocm_aiter_hipb_mm_fp8_fake, + ) + direct_register_custom_op( op_name="rocm_aiter_triton_gemm_a8w8_blockscale", op_func=_rocm_aiter_triton_gemm_a8w8_blockscale_impl, @@ -1858,6 +1920,17 @@ class rocm_aiter_ops: A, B, As, Bs, bias, output_dtype ) + @staticmethod + def hipb_mm_fp8( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.bfloat16, + ) -> torch.Tensor: + return torch.ops.vllm.rocm_aiter_hipb_mm_fp8(A, B, As, Bs, bias, output_dtype) + @staticmethod def triton_gemm_a8w8_blockscale( A: torch.Tensor, diff --git a/vllm/envs.py b/vllm/envs.py index 5746ff89833..a32f055a028 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -115,6 +115,7 @@ if TYPE_CHECKING: VLLM_ROCM_USE_AITER: bool = False VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False VLLM_ROCM_USE_AITER_LINEAR: bool = True + VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 VLLM_ROCM_USE_AITER_RMSNORM: bool = True @@ -1117,6 +1118,9 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER_LINEAR": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_LINEAR", "True").lower() in ("true", "1") ), + "VLLM_ROCM_USE_AITER_LINEAR_HIPBMM": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "False").lower() in ("true", "1") + ), # Whether to use aiter moe ops. # By default is enabled. "VLLM_ROCM_USE_AITER_MOE": lambda: ( diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index cd2c9eb01be..8162acd5e8d 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -128,6 +128,7 @@ from vllm.model_executor.kernels.linear.scaled_mm import ( ) from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( AiterFp8BlockScaledMMKernel, + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, AiterInt8ScaledMMLinearKernel, AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, @@ -285,6 +286,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = ChannelWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ + AiterHipbMMPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, AiterPerTokenFp8ScaledMMLinearKernel, ROCmFP8ScaledMMLinearKernel, @@ -1024,6 +1026,7 @@ __all__ = [ "FP8ScaledMMLinearLayerConfig", "Int8ScaledMMLinearLayerConfig", "ScaledMMLinearLayerConfig", + "AiterHipbMMPerTokenFp8ScaledMMLinearKernel", "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel", "AiterPerTokenFp8ScaledMMLinearKernel", "NvFp4LinearKernel", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 5ded5ca798a..1b39491ab34 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -212,6 +212,99 @@ class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): ) +class AiterHipbMMPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_rocm(): + return False, "requires ROCm." + + if not rocm_aiter_ops.is_linear_hipbmm_enabled(): + return ( + False, + "requires setting `VLLM_ROCM_USE_AITER=1`, " + "`VLLM_ROCM_USE_AITER_LINEAR=1`, " + "and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`.", + ) + try: + import aiter # noqa: F401 + except Exception: + return False, "requires aiter library to be installed." + + if not hasattr(aiter, "hipb_mm"): + return False, "requires aiter hipb_mm support." + + return True, None + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + is_ptpc = ( + c.activation_quant_key.scale.group_shape.is_per_token() + and c.weight_quant_key.scale.group_shape.is_per_channel() + ) + if c.weight_shape is None: + return False, "weight_shape is required for Aiter kernels" + N, K = c.weight_shape + + if c.out_dtype is not torch.bfloat16: + return False, "requires bfloat16 output dtype." + + if not is_ptpc: + return ( + False, + "requires per token activation scales and per channel weight scales.", + ) + + if not (N >= 16 and N % 16 == 0 and K % 16 == 0): + return ( + False, + "requires N >= 16 and both N and K divisible by 16, " + f"received N={N} and K={K}.", + ) + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w_name, w_s_name, *_ = self.layer_param_names + w, w_s, *_ = self._get_layer_params(layer) + + # Pre-apply the transposes that used to live in + # _rocm_aiter_hipb_mm_fp8_impl so the kernel can consume B/Bs directly. + # The `.t()` on the shuffled weight is kept as a non-contiguous view — + # materializing it with `.contiguous()` would re-arrange the bytes and + # break the `bpreshuffle` layout. + shuffled_w = rocm_aiter_ops.shuffle_weight(w.t().contiguous()) + replace_parameter( + layer, + w_name, + torch.nn.Parameter(shuffled_w.t(), requires_grad=False), + ) + + if w_s.ndim > 1: + replace_parameter( + layer, + w_s_name, + torch.nn.Parameter(w_s.t().contiguous(), requires_grad=False), + ) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + output_shape[-1] = B.shape[1] + return rocm_aiter_ops.hipb_mm_fp8(A, B, As, Bs, bias, out_dtype).view( + *output_shape + ) + + class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): @classmethod def is_supported( diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 89471e844d8..7f6d8794c28 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -134,6 +134,7 @@ def _sync_hip_cuda_env_vars(): # Sync at import time - catches misconfigurations from process start. _sync_hip_cuda_env_vars() + # AMDSMI utils # Note that NVML is not affected by `{CUDA/HIP}_VISIBLE_DEVICES`, # all the related functions work on real physical device ids. @@ -312,6 +313,17 @@ def on_gfx950() -> bool: return _ON_GFX950 +# Enable HIP online tuning early, before hipBLASLt initializes. +# Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM. +if ( + envs.VLLM_ROCM_USE_AITER + and envs.VLLM_ROCM_USE_AITER_LINEAR + and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM + and on_mi3xx() +): + os.environ["HIP_ONLINE_TUNING"] = "1" + + @cache def use_rocm_custom_paged_attention( qtype: torch.dtype,