@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 165b7864d0
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 1, 21:15 UTC

[ROCM] [FEAT] Integrate Aiter hipBLASLt GEMM online tuning (#40426)

Signed-off-by: hanlin12 <hanlin12@amd.com> Signed-off-by: Han Lin <hanlin12@amd.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

Han Lin committed Jun 5, 2026 at 14:45 UTC 165b7864d0e3208e09c33ed4ce459d2bca3c99ea
6 files changed +591
tests/rocm/aiter/test_aiter_hipb_mm_linear_kernel.py new
+406
@@ -0,0 +1,406 @@
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 + kFp8DynamicTokenSym,
22 + kFp8StaticChannelSym,
23 + kFp8StaticTensorSym,
24 +)
25 +from vllm.platforms import current_platform
26 +
27 +aiter_available = importlib.util.find_spec("aiter") is not None
28 +
29 +pytestmark = [
30 + pytest.mark.skipif(
31 + not (
32 + current_platform.is_rocm()
33 + and current_platform.supports_fp8()
34 + and aiter_available
35 + ),
36 + reason="Requires ROCm + FP8 support + aiter",
37 + ),
38 + pytest.mark.usefixtures("default_vllm_config"),
39 +]
40 +
41 +
42 +@pytest.fixture
43 +def enable_hipb_mm_kernel(monkeypatch: pytest.MonkeyPatch):
44 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
45 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")
46 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "1")
47 + rocm_aiter_ops.refresh_env_variables()
48 + yield
49 + rocm_aiter_ops.refresh_env_variables()
50 +
51 +
52 +def _make_config(
53 + *,
54 + weight_quant_key=kFp8StaticChannelSym,
55 + out_dtype: torch.dtype = torch.bfloat16,
56 + weight_shape: tuple[int, int] = (512, 4096),
57 +) -> FP8ScaledMMLinearLayerConfig:
58 + return FP8ScaledMMLinearLayerConfig(
59 + weight_quant_key=weight_quant_key,
60 + activation_quant_key=kFp8DynamicTokenSym,
61 + weight_shape=weight_shape,
62 + input_dtype=torch.bfloat16,
63 + out_dtype=out_dtype,
64 + )
65 +
66 +
67 +def _find_csv_row(path: str, m: int, n: int, k: int) -> dict | None:
68 + if not os.path.exists(path):
69 + return None
70 +
71 + with open(path, newline="") as f:
72 + reader = csv.DictReader(f, skipinitialspace=True)
73 + for row in reader:
74 + try:
75 + if (
76 + int(row.get("m", -1)) == m
77 + and int(row.get("n", -1)) == n
78 + and int(row.get("k", -1)) == k
79 + ):
80 + return dict(row)
81 + except (TypeError, ValueError):
82 + continue
83 + return None
84 +
85 +
86 +def _skip_if_no_hipb_mm_solution(exc: RuntimeError) -> None:
87 + if "hipblasLtMatmulAlgoGetHeuristic found 0 valid solutions" in str(exc):
88 + pytest.skip(
89 + "hipb_mm bpreshuffle path has no valid hipBLASLt solution on "
90 + "this ROCm stack."
91 + )
92 +
93 +
94 +def _check_bpreshuffle_runtime_support(weight_shape: tuple[int, int], num_tokens: int):
95 + import aiter
96 + from aiter.ops.shuffle import shuffle_weight
97 +
98 + x = torch.randn(num_tokens, weight_shape[1], dtype=torch.bfloat16, device="cuda")
99 + w = torch.randn(weight_shape, dtype=torch.bfloat16, device="cuda")
100 +
101 + aiter.hipb_create_extension()
102 + x_q, x_scale = aiter.pertoken_quant(x, quant_dtype=current_platform.fp8_dtype())
103 + w_q, w_scale = aiter.pertoken_quant(w, quant_dtype=current_platform.fp8_dtype())
104 +
105 + try:
106 + aiter.hipb_mm(
107 + x_q,
108 + shuffle_weight(w_q, layout=(16, 16)).t(),
109 + solution_index=-1,
110 + out_dtype=torch.bfloat16,
111 + scaleA=x_scale,
112 + scaleB=w_scale.t().contiguous(),
113 + scaleOut=None,
114 + bpreshuffle=True,
115 + )
116 + except RuntimeError as exc:
117 + _skip_if_no_hipb_mm_solution(exc)
118 + raise
119 +
120 +
121 +def test_hipb_mm_kernel_requires_hipbmm_flag(monkeypatch: pytest.MonkeyPatch):
122 + # The kernel rejects when `is_hip_fp8bmm_enabled()` is False. That helper
123 + # requires AITER + AITER_LINEAR + MI3xx, so dropping AITER_LINEAR exercises
124 + # the rejection branch.
125 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
126 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "0")
127 + monkeypatch.delenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", raising=False)
128 + rocm_aiter_ops.refresh_env_variables()
129 +
130 + is_supported, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.is_supported()
131 +
132 + assert not is_supported
133 + assert reason == (
134 + "requires setting `VLLM_ROCM_USE_AITER=1`, "
135 + "`VLLM_ROCM_USE_AITER_LINEAR=1`, "
136 + "and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`."
137 + )
138 +
139 +
140 +def test_hipb_mm_flag_enables_hip_online_tuning(
141 + monkeypatch: pytest.MonkeyPatch,
142 +):
143 + import vllm.envs as envs_mod
144 + import vllm.platforms.rocm as rocm_mod
145 +
146 + # The rocm.py gate requires all three AITER flags (and MI3xx) to auto-set
147 + # HIP_ONLINE_TUNING.
148 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
149 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")
150 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "1")
151 +
152 + try:
153 + importlib.reload(envs_mod)
154 + importlib.reload(rocm_mod)
155 + assert envs_mod.VLLM_ROCM_USE_AITER
156 + assert envs_mod.VLLM_ROCM_USE_AITER_LINEAR
157 + assert envs_mod.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
158 + assert os.environ.get("HIP_ONLINE_TUNING") == "1"
159 + finally:
160 + monkeypatch.undo()
161 + os.environ.pop("HIP_ONLINE_TUNING", None)
162 + importlib.reload(envs_mod)
163 + importlib.reload(rocm_mod)
164 + rocm_aiter_ops.refresh_env_variables()
165 +
166 +
167 +def test_hipb_mm_kernel_can_implement_success(enable_hipb_mm_kernel):
168 + can_implement, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.can_implement(
169 + _make_config()
170 + )
171 +
172 + assert can_implement
173 + assert reason is None
174 +
175 +
176 +@pytest.mark.parametrize(
177 + ("config", "expected_reason"),
178 + [
179 + (
180 + _make_config(weight_quant_key=kFp8StaticTensorSym),
181 + "requires per token activation scales and per channel weight scales.",
182 + ),
183 + (
184 + _make_config(out_dtype=torch.float16),
185 + "requires bfloat16 output dtype.",
186 + ),
187 + (
188 + _make_config(weight_shape=(8, 4090)),
189 + "requires N >= 16 and both N and K divisible by 16, "
190 + "received N=8 and K=4090.",
191 + ),
192 + ],
193 +)
194 +def test_hipb_mm_kernel_can_implement_rejects_unsupported_configs(
195 + enable_hipb_mm_kernel,
196 + config: FP8ScaledMMLinearLayerConfig,
197 + expected_reason: str,
198 +):
199 + can_implement, reason = AiterHipbMMPerTokenFp8ScaledMMLinearKernel.can_implement(
200 + config
201 + )
202 +
203 + assert not can_implement
204 + assert reason == expected_reason
205 +
206 +
207 +def test_hipb_mm_kernel_process_weights_after_loading_shuffles_weights(
208 + enable_hipb_mm_kernel,
209 +):
210 + weight_shape = (512, 4096)
211 + kernel = AiterHipbMMPerTokenFp8ScaledMMLinearKernel(
212 + _make_config(weight_shape=weight_shape),
213 + layer_param_names=("weight", "weight_scale", "input_scale", "input_scale_ub"),
214 + )
215 +
216 + layer = torch.nn.Module()
217 + layer.weight = torch.nn.Parameter(
218 + torch.rand(weight_shape, device="cuda").to(current_platform.fp8_dtype()).t(),
219 + requires_grad=False,
220 + )
221 + layer.weight_scale = torch.nn.Parameter(
222 + torch.rand((weight_shape[0], 1), dtype=torch.float32, device="cuda"),
223 + requires_grad=False,
224 + )
225 + layer.input_scale = None
226 + layer.input_scale_ub = None
227 +
228 + original_weight = layer.weight.detach().clone()
229 + original_weight_scale = layer.weight_scale.detach().clone()
230 +
231 + kernel.process_weights_after_loading(layer)
232 +
233 + # process_weights_after_loading now pre-applies the transposes that used
234 + # to live in _rocm_aiter_hipb_mm_fp8_impl, so the stored weight is the
235 + # shuffled tensor with a trailing `.t()` view, and the stored weight scale
236 + # is its transposed-contiguous form.
237 + expected_weight = rocm_aiter_ops.shuffle_weight(
238 + original_weight.t().contiguous()
239 + ).t()
240 + torch.testing.assert_close(layer.weight, expected_weight)
241 +
242 + expected_weight_scale = original_weight_scale.t().contiguous()
243 + torch.testing.assert_close(layer.weight_scale, expected_weight_scale)
244 +
245 +
246 +def test_hipb_mm_kernel_forward_matches_raw_aiter_hipb_mm(enable_hipb_mm_kernel):
247 + import aiter
248 +
249 + weight_shape = (512, 4096)
250 + _check_bpreshuffle_runtime_support(weight_shape, num_tokens=32)
251 +
252 + layer = TestFP8Layer(
253 + weight_shape=weight_shape,
254 + activation_quant_key=kFp8DynamicTokenSym,
255 + weight_quant_key=kFp8StaticChannelSym,
256 + input_dtype=torch.bfloat16,
257 + out_dtype=torch.bfloat16,
258 + device=torch.device("cuda"),
259 + force_kernel=AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
260 + )
261 +
262 + # hipb_mm uses a transposed-result GEMM internally, so the flattened token
263 + # count becomes the effective N dimension passed into hipBLASLt. Keep it
264 + # aligned to avoid heuristic failures for tiny N.
265 + x = torch.randn(2, 16, weight_shape[1], dtype=torch.bfloat16, device="cuda")
266 + bias = torch.randn(weight_shape[0], dtype=torch.bfloat16, device="cuda")
267 +
268 + try:
269 + out = layer(x, bias)
270 + except RuntimeError as exc:
271 + _skip_if_no_hipb_mm_solution(exc)
272 + raise
273 +
274 + x_2d = x.view(-1, x.shape[-1])
275 + x_q, x_scale = layer.kernel.quant_fp8(
276 + x_2d,
277 + layer.input_scale,
278 + layer.input_scale_ub,
279 + )
280 + try:
281 + # process_weights_after_loading already applies the trailing `.t()` on
282 + # the shuffled weight and the `.t().contiguous()` on the weight scale,
283 + # so the raw aiter call uses them directly.
284 + expected = aiter.hipb_mm(
285 + x_q,
286 + layer.weight,
287 + solution_index=-1,
288 + bias=bias,
289 + out_dtype=torch.bfloat16,
290 + scaleA=x_scale,
291 + scaleB=layer.weight_scale,
292 + scaleOut=None,
293 + bpreshuffle=True,
294 + ).view(*out.shape)
295 + except RuntimeError as exc:
296 + _skip_if_no_hipb_mm_solution(exc)
297 + raise
298 +
299 + assert isinstance(layer.kernel, AiterHipbMMPerTokenFp8ScaledMMLinearKernel)
300 + assert out.shape == (2, 16, weight_shape[0])
301 + torch.testing.assert_close(out, expected)
302 +
303 +
304 +def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel):
305 + """Kernel output should match a dequantized fp32 reference within
306 + fp8 per-token / per-channel quantization noise."""
307 + weight_shape = (512, 4096) # (N, K)
308 + num_tokens = 32
309 + _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens)
310 +
311 + fp8_dtype = current_platform.fp8_dtype()
312 + fp8_max = torch.finfo(fp8_dtype).max
313 + device = torch.device("cuda")
314 +
315 + # Build a bf16 weight and quantize per output channel (one scale per row).
316 + w_bf16 = torch.randn(weight_shape, dtype=torch.bfloat16, device=device)
317 + w_amax = w_bf16.abs().amax(dim=1, keepdim=True).to(torch.float32)
318 + w_scale = (w_amax / fp8_max).clamp(min=1e-12)
319 + w_fp8 = (w_bf16.to(torch.float32) / w_scale).clamp(-fp8_max, fp8_max).to(fp8_dtype)
320 + w_dequant = w_fp8.to(torch.float32) * w_scale
321 +
322 + bias = torch.randn(weight_shape[0], dtype=torch.bfloat16, device=device)
323 +
324 + layer = torch.nn.Module()
325 + # Pre-`process_weights_after_loading` convention: weight stored as the
326 + # `[K, N]` view of the fp8 tensor.
327 + layer.weight = torch.nn.Parameter(w_fp8.t(), requires_grad=False)
328 + layer.weight_scale = torch.nn.Parameter(w_scale, requires_grad=False)
329 + layer.input_scale = None
330 + layer.input_scale_ub = None
331 +
332 + kernel = AiterHipbMMPerTokenFp8ScaledMMLinearKernel(
333 + _make_config(weight_shape=weight_shape),
334 + layer_param_names=("weight", "weight_scale", "input_scale", "input_scale_ub"),
335 + )
336 + kernel.process_weights_after_loading(layer)
337 +
338 + x = torch.randn(num_tokens, weight_shape[1], dtype=torch.bfloat16, device=device)
339 +
340 + try:
341 + out = kernel.apply_weights(layer, x, bias)
342 + except RuntimeError as exc:
343 + _skip_if_no_hipb_mm_solution(exc)
344 + raise
345 +
346 + # Reference: quantize x per-token the same way the kernel does, then run
347 + # the matmul in fp32 against the dequantized weight. This isolates plumbing
348 + # / reduction bugs from inherent fp8 quantization noise.
349 + x_amax = x.abs().amax(dim=1, keepdim=True).to(torch.float32)
350 + x_scale_ref = (x_amax / fp8_max).clamp(min=1e-12)
351 + x_q = (x.to(torch.float32) / x_scale_ref).clamp(-fp8_max, fp8_max).to(fp8_dtype)
352 + x_dequant = x_q.to(torch.float32) * x_scale_ref
353 + expected = (x_dequant @ w_dequant.t() + bias.to(torch.float32)).to(torch.bfloat16)
354 +
355 + assert out.shape == (num_tokens, weight_shape[0])
356 + # K=4096 fp8 reduction leaves room for accumulation order drift and
357 + # catastrophic cancellation on near-zero outputs; tolerances are loose
358 + # enough to absorb that but tight enough to catch wrong layouts, missing
359 + # bias, swapped scales, etc.
360 + torch.testing.assert_close(out, expected, atol=5.0, rtol=0.1)
361 +
362 +
363 +def test_hipb_mm_kernel_online_tuning_writes_csv(
364 + enable_hipb_mm_kernel,
365 + monkeypatch: pytest.MonkeyPatch,
366 + tmp_path,
367 +):
368 + weight_shape = (256, 4096)
369 + cache_file = tmp_path / "hip_online_tuning_res.csv"
370 +
371 + _check_bpreshuffle_runtime_support(weight_shape, num_tokens=16)
372 +
373 + monkeypatch.setenv("HIP_ONLINE_TUNING", "1")
374 + monkeypatch.chdir(tmp_path)
375 +
376 + layer = TestFP8Layer(
377 + weight_shape=weight_shape,
378 + activation_quant_key=kFp8DynamicTokenSym,
379 + weight_quant_key=kFp8StaticChannelSym,
380 + input_dtype=torch.bfloat16,
381 + out_dtype=torch.bfloat16,
382 + device=torch.device("cuda"),
383 + force_kernel=AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
384 + )
385 +
386 + # The effective heuristic N dimension is the flattened token count.
387 + x = torch.randn(16, weight_shape[1], dtype=torch.bfloat16, device="cuda")
388 + try:
389 + out = layer(x)
390 + except RuntimeError as exc:
391 + _skip_if_no_hipb_mm_solution(exc)
392 + raise
393 + torch.accelerator.synchronize()
394 +
395 + assert out.shape == (16, weight_shape[0])
396 + assert cache_file.exists()
397 +
398 + # hipb_mm records the internal GEMM dimensions used by hipBLASLt after its
399 + # transposed-result transformation.
400 + row = _find_csv_row(
401 + str(cache_file),
402 + m=weight_shape[0],
403 + n=x.shape[0],
404 + k=weight_shape[1],
405 + )
406 + assert row is not None
vllm/_aiter_ops.py
+73
@@ -27,6 +27,16 @@ except ImportError:
27 # on ROCm the fp8_dtype always calls is_fp8_fnuz
28 # which is a host op, so we cache it once here.
29 FP8_DTYPE = current_platform.fp8_dtype()
30 +_HIPB_MM_INITIALIZED_DEVICES: set[int] = set()
31 +
32 +
33 +def _ensure_hipb_mm_extension_initialized() -> None:
34 + import aiter
35 +
36 + device = torch.accelerator.current_device_index()
37 + if device not in _HIPB_MM_INITIALIZED_DEVICES:
38 + aiter.hipb_create_extension()
39 + _HIPB_MM_INITIALIZED_DEVICES.add(device)
40
41
42 def is_aiter_found() -> bool:
@@ -625,6 +635,43 @@ def _rocm_aiter_preshuffled_per_token_w8a8_gemm_fake(
635 return torch.empty(m, n, dtype=output_dtype, device=A.device)
636
637
638 +def _rocm_aiter_hipb_mm_fp8_impl(
639 + A: torch.Tensor,
640 + B: torch.Tensor,
641 + As: torch.Tensor,
642 + Bs: torch.Tensor,
643 + bias: torch.Tensor | None = None,
644 + output_dtype: torch.dtype = torch.bfloat16,
645 +) -> torch.Tensor:
646 + from aiter import hipb_mm
647 +
648 + _ensure_hipb_mm_extension_initialized()
649 + return hipb_mm(
650 + A,
651 + B,
652 + solution_index=-1,
653 + bias=bias,
654 + out_dtype=output_dtype,
655 + scaleA=As,
656 + scaleB=Bs,
657 + scaleOut=None,
658 + bpreshuffle=True,
659 + )
660 +
661 +
662 +def _rocm_aiter_hipb_mm_fp8_fake(
663 + A: torch.Tensor,
664 + B: torch.Tensor,
665 + As: torch.Tensor,
666 + Bs: torch.Tensor,
667 + bias: torch.Tensor | None = None,
668 + output_dtype: torch.dtype = torch.bfloat16,
669 +) -> torch.Tensor:
670 + m = A.shape[0]
671 + n = B.shape[1]
672 + return torch.empty(m, n, dtype=output_dtype, device=A.device)
673 +
674 +
675 def _rocm_aiter_triton_gemm_a8w8_blockscale_impl(
676 A: torch.Tensor,
677 B: torch.Tensor,
@@ -1308,6 +1355,7 @@ class rocm_aiter_ops:
1355 # TODO: Consolidate under _LINEAR_ENABLED
1356 _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
1357 _FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
1358 + _LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
1359 # TODO: Consolidate under _LINEAR_ENABLED
1360 _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
1361 # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
@@ -1340,6 +1388,7 @@ class rocm_aiter_ops:
1388 cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
1389 cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
1390 cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
1391 + cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
1392 cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
1393 cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
1394 cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
@@ -1512,6 +1561,13 @@ class rocm_aiter_ops:
1561
1562 return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950()
1563
1564 + @classmethod
1565 + @if_aiter_supported
1566 + def is_linear_hipbmm_enabled(cls) -> bool:
1567 + from vllm.platforms.rocm import on_mi3xx
1568 +
1569 + return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED
1570 +
1571 @classmethod
1572 @if_aiter_supported
1573 def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:
@@ -1668,6 +1724,12 @@ class rocm_aiter_ops:
1724 fake_impl=_rocm_aiter_preshuffled_per_token_w8a8_gemm_fake,
1725 )
1726
1727 + direct_register_custom_op(
1728 + op_name="rocm_aiter_hipb_mm_fp8",
1729 + op_func=_rocm_aiter_hipb_mm_fp8_impl,
1730 + fake_impl=_rocm_aiter_hipb_mm_fp8_fake,
1731 + )
1732 +
1733 direct_register_custom_op(
1734 op_name="rocm_aiter_triton_gemm_a8w8_blockscale",
1735 op_func=_rocm_aiter_triton_gemm_a8w8_blockscale_impl,
@@ -1858,6 +1920,17 @@ class rocm_aiter_ops:
1920 A, B, As, Bs, bias, output_dtype
1921 )
1922
1923 + @staticmethod
1924 + def hipb_mm_fp8(
1925 + A: torch.Tensor,
1926 + B: torch.Tensor,
1927 + As: torch.Tensor,
1928 + Bs: torch.Tensor,
1929 + bias: torch.Tensor | None = None,
1930 + output_dtype: torch.dtype = torch.bfloat16,
1931 + ) -> torch.Tensor:
1932 + return torch.ops.vllm.rocm_aiter_hipb_mm_fp8(A, B, As, Bs, bias, output_dtype)
1933 +
1934 @staticmethod
1935 def triton_gemm_a8w8_blockscale(
1936 A: torch.Tensor,
vllm/envs.py
+4
@@ -115,6 +115,7 @@ if TYPE_CHECKING:
115 VLLM_ROCM_USE_AITER: bool = False
116 VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
117 VLLM_ROCM_USE_AITER_LINEAR: bool = True
118 + VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False
119 VLLM_ROCM_USE_AITER_MOE: bool = True
120 VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0
121 VLLM_ROCM_USE_AITER_RMSNORM: bool = True
@@ -1117,6 +1118,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1118 "VLLM_ROCM_USE_AITER_LINEAR": lambda: (
1119 os.getenv("VLLM_ROCM_USE_AITER_LINEAR", "True").lower() in ("true", "1")
1120 ),
1121 + "VLLM_ROCM_USE_AITER_LINEAR_HIPBMM": lambda: (
1122 + os.getenv("VLLM_ROCM_USE_AITER_LINEAR_HIPBMM", "False").lower() in ("true", "1")
1123 + ),
1124 # Whether to use aiter moe ops.
1125 # By default is enabled.
1126 "VLLM_ROCM_USE_AITER_MOE": lambda: (
vllm/model_executor/kernels/linear/__init__.py
+3
@@ -128,6 +128,7 @@ from vllm.model_executor.kernels.linear.scaled_mm import (
128 )
129 from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
130 AiterFp8BlockScaledMMKernel,
131 + AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
132 AiterInt8ScaledMMLinearKernel,
133 AiterPerTokenFp8ScaledMMLinearKernel,
134 AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
@@ -285,6 +286,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
286 ChannelWiseTorchFP8ScaledMMLinearKernel,
287 ],
288 PlatformEnum.ROCM: [
289 + AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
290 AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
291 AiterPerTokenFp8ScaledMMLinearKernel,
292 ROCmFP8ScaledMMLinearKernel,
@@ -1024,6 +1026,7 @@ __all__ = [
1026 "FP8ScaledMMLinearLayerConfig",
1027 "Int8ScaledMMLinearLayerConfig",
1028 "ScaledMMLinearLayerConfig",
1029 + "AiterHipbMMPerTokenFp8ScaledMMLinearKernel",
1030 "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel",
1031 "AiterPerTokenFp8ScaledMMLinearKernel",
1032 "NvFp4LinearKernel",
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+93
@@ -212,6 +212,99 @@ class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
212 )
213
214
215 +class AiterHipbMMPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
216 + @classmethod
217 + def is_supported(
218 + cls, compute_capability: int | None = None
219 + ) -> tuple[bool, str | None]:
220 + if not current_platform.is_rocm():
221 + return False, "requires ROCm."
222 +
223 + if not rocm_aiter_ops.is_linear_hipbmm_enabled():
224 + return (
225 + False,
226 + "requires setting `VLLM_ROCM_USE_AITER=1`, "
227 + "`VLLM_ROCM_USE_AITER_LINEAR=1`, "
228 + "and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`.",
229 + )
230 + try:
231 + import aiter # noqa: F401
232 + except Exception:
233 + return False, "requires aiter library to be installed."
234 +
235 + if not hasattr(aiter, "hipb_mm"):
236 + return False, "requires aiter hipb_mm support."
237 +
238 + return True, None
239 +
240 + @classmethod
241 + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
242 + is_ptpc = (
243 + c.activation_quant_key.scale.group_shape.is_per_token()
244 + and c.weight_quant_key.scale.group_shape.is_per_channel()
245 + )
246 + if c.weight_shape is None:
247 + return False, "weight_shape is required for Aiter kernels"
248 + N, K = c.weight_shape
249 +
250 + if c.out_dtype is not torch.bfloat16:
251 + return False, "requires bfloat16 output dtype."
252 +
253 + if not is_ptpc:
254 + return (
255 + False,
256 + "requires per token activation scales and per channel weight scales.",
257 + )
258 +
259 + if not (N >= 16 and N % 16 == 0 and K % 16 == 0):
260 + return (
261 + False,
262 + "requires N >= 16 and both N and K divisible by 16, "
263 + f"received N={N} and K={K}.",
264 + )
265 +
266 + return True, None
267 +
268 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
269 + w_name, w_s_name, *_ = self.layer_param_names
270 + w, w_s, *_ = self._get_layer_params(layer)
271 +
272 + # Pre-apply the transposes that used to live in
273 + # _rocm_aiter_hipb_mm_fp8_impl so the kernel can consume B/Bs directly.
274 + # The `.t()` on the shuffled weight is kept as a non-contiguous view —
275 + # materializing it with `.contiguous()` would re-arrange the bytes and
276 + # break the `bpreshuffle` layout.
277 + shuffled_w = rocm_aiter_ops.shuffle_weight(w.t().contiguous())
278 + replace_parameter(
279 + layer,
280 + w_name,
281 + torch.nn.Parameter(shuffled_w.t(), requires_grad=False),
282 + )
283 +
284 + if w_s.ndim > 1:
285 + replace_parameter(
286 + layer,
287 + w_s_name,
288 + torch.nn.Parameter(w_s.t().contiguous(), requires_grad=False),
289 + )
290 +
291 + def apply_scaled_mm(
292 + self,
293 + *,
294 + A: torch.Tensor,
295 + B: torch.Tensor,
296 + out_dtype: torch.dtype,
297 + As: torch.Tensor,
298 + Bs: torch.Tensor,
299 + bias: torch.Tensor | None,
300 + output_shape: list,
301 + ) -> torch.Tensor:
302 + output_shape[-1] = B.shape[1]
303 + return rocm_aiter_ops.hipb_mm_fp8(A, B, As, Bs, bias, out_dtype).view(
304 + *output_shape
305 + )
306 +
307 +
308 class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
309 @classmethod
310 def is_supported(
vllm/platforms/rocm.py
+12
@@ -134,6 +134,7 @@ def _sync_hip_cuda_env_vars():
134 # Sync at import time - catches misconfigurations from process start.
135 _sync_hip_cuda_env_vars()
136
137 +
138 # AMDSMI utils
139 # Note that NVML is not affected by `{CUDA/HIP}_VISIBLE_DEVICES`,
140 # all the related functions work on real physical device ids.
@@ -312,6 +313,17 @@ def on_gfx950() -> bool:
313 return _ON_GFX950
314
315
316 +# Enable HIP online tuning early, before hipBLASLt initializes.
317 +# Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM.
318 +if (
319 + envs.VLLM_ROCM_USE_AITER
320 + and envs.VLLM_ROCM_USE_AITER_LINEAR
321 + and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
322 + and on_mi3xx()
323 +):
324 + os.environ["HIP_ONLINE_TUNING"] = "1"
325 +
326 +
327 @cache
328 def use_rocm_custom_paged_attention(
329 qtype: torch.dtype,