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 15, 02:15 UTC
[V1] [ROCm] [AITER] Upgrade AITER to commit `916bf3c` and bugfix APIs (#20880)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Jul 13, 2025 at 08:19 UTC
80d38b8ac850fbb19b9a76e74cd53ff04573e58b
3 files changed
+48
-5
docker/Dockerfile.rocm_base
+1
-1
@@ -12,7 +12,7 @@ ARG PYTORCH_REPO="https://github.com/pytorch/pytorch.git"
12
ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
13
ARG FA_BRANCH="1a7f4dfa"
14
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
15
-ARG AITER_BRANCH="6487649"
15
+ARG AITER_BRANCH="916bf3c"
16
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
17
18
FROM ${BASE_IMAGE} AS base
vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py
+46
-3
@@ -8,11 +8,55 @@ import torch
8
import vllm.envs as envs
9
from vllm import _custom_ops as ops
10
from vllm.platforms import current_platform
11
+from vllm.utils import direct_register_custom_op
12
13
from .cutlass import CutlassScaledMMLinearKernel
14
from .ScaledMMLinearKernel import ScaledMMLinearLayerConfig
15
16
17
+def rocm_aiter_gemm_w8a8_impl(
18
+ A: torch.Tensor,
19
+ B: torch.Tensor,
20
+ As: torch.Tensor,
21
+ Bs: torch.Tensor,
22
+ bias: Optional[torch.Tensor] = None,
23
+ output_dtype: torch.dtype = torch.float16,
24
+) -> torch.Tensor:
25
+
26
+ from aiter import gemm_a8w8_CK
27
+
28
+ # gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
29
+ # a to be [M, K]
30
+ # b to be [N, K]
31
+ # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
32
+ return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)
33
+
34
+
35
+def rocm_aiter_gemm_w8a8_fake(
36
+ A: torch.Tensor,
37
+ B: torch.Tensor,
38
+ As: torch.Tensor,
39
+ Bs: torch.Tensor,
40
+ bias: Optional[torch.Tensor] = None,
41
+ output_dtype: torch.dtype = torch.float16,
42
+) -> torch.Tensor:
43
+
44
+ m = A.shape[0]
45
+ n = B.shape[0]
46
+ Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
47
+ return Y
48
+
49
+
50
+if current_platform.is_rocm():
51
+ direct_register_custom_op(
52
+ op_name="rocm_aiter_gemm_w8a8",
53
+ op_func=rocm_aiter_gemm_w8a8_impl,
54
+ mutates_args=[],
55
+ fake_impl=rocm_aiter_gemm_w8a8_fake,
56
+ dispatch_key=current_platform.dispatch_key,
57
+ )
58
+
59
+
60
class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
61
62
@classmethod
@@ -111,10 +155,9 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
155
" w8a8 scaled gemm. `AiterScaledMMLinearKernel` " +
156
"does not support AITER block scaled GEMM.")
157
114
- from aiter import gemm_a8w8_CK
115
-
158
# gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
159
# a to be [M, K]
160
# b to be [N, K]
161
# CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
120
- return gemm_a8w8_CK(x_q, w_q.t(), x_s, w_s, bias).to(out_dtype)
162
+ return torch.ops.vllm.rocm_aiter_gemm_w8a8(x_q, w_q.t(), x_s, w_s,
163
+ bias, out_dtype)
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1
-1
@@ -56,7 +56,7 @@ def rocm_aiter_gemm_w8a8_blockscale_impl(
56
) -> torch.Tensor:
57
import aiter as rocm_aiter
58
59
- return rocm_aiter.gemm_a8w8_blockscale_CK(A, B, As, Bs, dtype=output_dtype)
59
+ return rocm_aiter.gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
60
61
62
def rocm_aiter_gemm_w8a8_blockscale_fake(