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, 17:15 UTC
[FEAT] [ROCm] Add AITER int8 scaled gemm kernel (#15433)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Mar 29, 2025 at 18:33 UTC
4965ec42d28830f0c30756dea19e14b45cdbe5b1
4 files changed
+202
-5
tests/quantization/test_compressed_tensors.py
+72
-4
@@ -20,6 +20,23 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
sparse_cutlass_supported)
21
from vllm.platforms import current_platform
22
23
+# AITER only supports per-channel-per-channel INT8 gemm
24
+# and per-tensor-per-tensor INT8 GEMM.
25
+# It does not support mix precision MM and mix quantization scheme.
26
+ROCM_AITER_SUPPORTED_INT8_MODEL = [
27
+ "neuralmagic/Llama-3.2-1B-quantized.w8a8",
28
+ "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2"
29
+]
30
+
31
+# TritonScaledMMLinearKernel only supports symmetric quantization.
32
+ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL = [
33
+ "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",
34
+ "nm-testing/tinyllama-oneshot-w8-channel-a8-tensor",
35
+ "neuralmagic/Llama-3.2-1B-quantized.w8a8",
36
+ "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2",
37
+ "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2",
38
+]
39
+
40
41
@pytest.fixture(scope="function", autouse=True)
42
def use_v0_only(monkeypatch):
@@ -57,6 +74,11 @@ def use_v0_only(monkeypatch):
74
)
75
def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
76
model_path, strategy, quant_type, shape_0, is_symmetric = model_args
77
+
78
+ if current_platform.is_rocm(
79
+ ) and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL:
80
+ pytest.skip(f"Skip model {model_path} as it is not support on ROCm.")
81
+
82
with vllm_runner(model_path, enforce_eager=True) as llm:
83
84
def check_model(model):
@@ -123,6 +145,8 @@ def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
145
)
146
@pytest.mark.parametrize("max_tokens", [32])
147
@pytest.mark.parametrize("num_logprobs", [10])
148
+@pytest.mark.parametrize(
149
+ "use_aiter", [True, False] if current_platform.is_rocm() else [False])
150
def test_compressed_tensors_w8a8_logprobs(
151
hf_runner,
152
vllm_runner,
@@ -130,7 +154,21 @@ def test_compressed_tensors_w8a8_logprobs(
154
model_path,
155
max_tokens,
156
num_logprobs,
157
+ use_aiter,
158
+ monkeypatch,
159
):
160
+
161
+ if current_platform.is_rocm(
162
+ ) and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL:
163
+ pytest.skip(f"Skip model {model_path} as it is not support on ROCm.")
164
+
165
+ if use_aiter:
166
+ if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL:
167
+ pytest.skip(
168
+ f"Skip model {model_path} as it is not support by aiter.")
169
+ # this will enable VLLM_ROCM_USE_AITER_LINEAR
170
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
171
+
172
dtype = "bfloat16"
173
174
# skip language translation prompt for the static per tensor asym model
@@ -154,6 +192,9 @@ def test_compressed_tensors_w8a8_logprobs(
192
name_1="vllm",
193
)
194
195
+ if current_platform.is_rocm():
196
+ torch.cuda.synchronize()
197
+
198
199
def test_compressed_tensors_no_enforce_eager(vllm_runner):
200
model_path = "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change"
@@ -177,8 +218,27 @@ def test_compressed_tensors_no_enforce_eager(vllm_runner):
218
),
219
],
220
)
180
-def test_compressed_tensors_w8a8_dynamic_per_token(vllm_runner, model_args):
221
+@pytest.mark.parametrize(
222
+ "use_aiter", [True, False] if current_platform.is_rocm() else [False])
223
+def test_compressed_tensors_w8a8_dynamic_per_token(
224
+ vllm_runner,
225
+ model_args,
226
+ use_aiter,
227
+ monkeypatch,
228
+):
229
model_path, strategy = model_args
230
+
231
+ if current_platform.is_rocm(
232
+ ) and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL:
233
+ pytest.skip(f"Skip model {model_path} as it is not support on ROCm.")
234
+
235
+ if use_aiter:
236
+ if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL:
237
+ pytest.skip(
238
+ f"Skip model {model_path} as it is not support by aiter.")
239
+ # this will enable VLLM_ROCM_USE_AITER_LINEAR
240
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
241
+
242
with vllm_runner(model_path, dtype=torch.float16) as llm:
243
244
def check_model(model):
@@ -207,6 +267,8 @@ def test_compressed_tensors_w8a8_dynamic_per_token(vllm_runner, model_args):
267
("nm-testing/tinyllama-oneshot-w8a16-per-channel", "channel", None, 4),
268
],
269
)
270
+@pytest.mark.skipif(not current_platform.is_cuda(),
271
+ reason="The tests are skipped on non-CUDA platform.")
272
def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
273
model, strategy, group, pack_factor = wNa16_args
274
with vllm_runner(model) as llm:
@@ -231,6 +293,8 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
293
assert output
294
295
296
+@pytest.mark.skipif(not current_platform.is_cuda(),
297
+ reason="This test is skipped on non-CUDA platform.")
298
def test_compressed_tensors_w4a16_marlin24(vllm_runner):
299
model_path = "nm-testing/llama7b-one-shot-2_4-w4a16-marlin24-t"
300
with vllm_runner(model_path) as llm:
@@ -271,7 +335,7 @@ def test_compressed_tensors_fp8(vllm_runner):
335
336
if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8):
337
assert len(qkv_proj.input_scale.shape) == 0
274
- assert qkv_proj.weight.dtype is torch.float8_e4m3fn
338
+ assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
339
assert qkv_proj.weight_scale.dtype is torch.float32
340
assert len(qkv_proj.weight_scale.shape) == 0
341
@@ -281,6 +345,8 @@ def test_compressed_tensors_fp8(vllm_runner):
345
assert output
346
347
348
+@pytest.mark.skipif(not current_platform.is_cuda(),
349
+ reason="This test is skipped on non-CUDA platform.")
350
def test_compressed_tensors_kv_cache(vllm_runner):
351
model_path = "nm-testing/TinyLlama-1.1B-compressed-tensors-kv-cache-scheme"
352
with vllm_runner(model_path, kv_cache_dtype="fp8") as llm:
@@ -309,7 +375,8 @@ def _test_2of4_quant_models(qkv_proj,
375
376
377
@pytest.mark.skipif(
312
- not current_platform.has_device_capability(90),
378
+ not current_platform.is_cuda()
379
+ or not current_platform.has_device_capability(90),
380
reason="Sparse FP8 is not yet supported on this GPU type.",
381
)
382
@pytest.mark.parametrize(
@@ -356,7 +423,8 @@ def test_compressed_tensors_2of4_quant_fp8(vllm_runner, args_2of4):
423
424
425
@pytest.mark.skipif(
359
- not current_platform.has_device_capability(90),
426
+ not current_platform.is_cuda()
427
+ or not current_platform.has_device_capability(90),
428
reason="Sparse FP8 is not yet supported on this GPU type.",
429
)
430
@pytest.mark.parametrize(
vllm/envs.py
+8
@@ -75,6 +75,7 @@ if TYPE_CHECKING:
75
VLLM_DISABLED_KERNELS: list[str] = []
76
VLLM_USE_V1: bool = True
77
VLLM_ROCM_USE_AITER: bool = False
78
+ VLLM_ROCM_USE_AITER_LINEAR: bool = True
79
VLLM_ROCM_USE_AITER_MOE: bool = True
80
VLLM_ROCM_USE_AITER_FP8_BLOCK_SCALED_MOE: bool = False
81
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
@@ -524,6 +525,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
525
lambda: (os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in
526
("true", "1")),
527
528
+ # use aiter linear op if aiter ops are enabled
529
+ # The following list of related ops
530
+ # - scaled_mm (per-tensor / rowwise)
531
+ "VLLM_ROCM_USE_AITER_LINEAR":
532
+ lambda: (os.getenv("VLLM_ROCM_USE_AITER_LINEAR", "True").lower() in
533
+ ("true", "1")),
534
+
535
# Whether to use aiter moe ops.
536
# By default is enabled.
537
"VLLM_ROCM_USE_AITER_MOE":
vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py
+3
-1
@@ -3,6 +3,8 @@
3
import os
4
from typing import Dict, List, Optional, Type
5
6
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (
7
+ AiterScaledMMLinearKernel)
8
from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
9
CutlassScaledMMLinearKernel)
10
from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
@@ -17,7 +19,7 @@ from vllm.platforms import PlatformEnum, current_platform
19
_POSSIBLE_KERNELS: Dict[PlatformEnum, List[Type[ScaledMMLinearKernel]]] = {
20
PlatformEnum.CPU: [CutlassScaledMMLinearKernel],
21
PlatformEnum.CUDA: [CutlassScaledMMLinearKernel],
20
- PlatformEnum.ROCM: [TritonScaledMMLinearKernel],
22
+ PlatformEnum.ROCM: [AiterScaledMMLinearKernel, TritonScaledMMLinearKernel],
23
PlatformEnum.TPU: [XLAScaledMMLinearKernel],
24
}
25
vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py
new
+119
@@ -0,0 +1,119 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from typing import Optional, Tuple
4
+
5
+import torch
6
+
7
+import vllm.envs as envs
8
+from vllm import _custom_ops as ops
9
+from vllm.platforms import current_platform
10
+
11
+from .cutlass import CutlassScaledMMLinearKernel
12
+from .ScaledMMLinearKernel import ScaledMMLinearLayerConfig
13
+
14
+
15
+class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
16
+
17
+ @classmethod
18
+ def get_min_capability(cls) -> int:
19
+ return 90
20
+
21
+ @classmethod
22
+ def can_implement(
23
+ cls, c: ScaledMMLinearLayerConfig) -> Tuple[bool, Optional[str]]:
24
+ if not current_platform.is_rocm():
25
+ return (
26
+ False,
27
+ "AiterScaledMMLinearKernel requires `aiter` which is not " +
28
+ "currently supported on non-ROCm platform.")
29
+
30
+ try:
31
+ import aiter # noqa: F401 # deliberately attempt to import aiter
32
+ except Exception:
33
+ return (
34
+ False,
35
+ "AiterScaledMMLinearKernel requires `aiter` which is not " +
36
+ "installed on ROCm.")
37
+ # Check if rocm_aiter_gemm_w8a8_scaled_mm is enabled
38
+ if not (
39
+ envs.VLLM_ROCM_USE_AITER_LINEAR \
40
+ and envs.VLLM_ROCM_USE_AITER
41
+ ):
42
+ return (False, "AiterScaledMMLinearKernel is disabled. " +
43
+ "Enable by setting `VLLM_ROCM_USE_AITER=1` " +
44
+ "and `VLLM_ROCM_USE_AITER_LINEAR=1`. " +
45
+ "`VLLM_ROCM_USE_AITER_LINEAR` default is True.")
46
+
47
+ if not c.input_symmetric:
48
+ return (False,
49
+ "AiterScaledMMLinearKernel only supports symmetric " +
50
+ "quantization.")
51
+ return True, None
52
+
53
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
54
+ super().process_weights_after_loading(layer)
55
+
56
+ def apply_weights(self,
57
+ layer: torch.nn.Module,
58
+ x: torch.Tensor,
59
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
60
+ """
61
+ `AiterScaledMMLinearKernel` implements a fused version of
62
+ `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
63
+ where scale_a * a and scale_b * b are implemented using numpy-style
64
+ broadcasting.
65
+ Currently only support per-tensor-per-tensor GEMM
66
+ and per-token-per-channel GEMM through AITER
67
+ w8a8 scaled gemm. `AiterScaledMMLinearKernel` also does not support
68
+ ATIER block scaled GEMM and mix-precision GEMM.
69
+ """
70
+ w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)
71
+
72
+ # ops.scaled_int8_quant supports both dynamic and static quant:
73
+ # * dynamic, i_s is None and x_s computed from x.
74
+ # * static, i_s is scalar and x_s is i_s.
75
+ symmetric = azp_adj is None
76
+ assert symmetric, ("AiterScaledMMLinearKernel only supports"
77
+ " symmetric quantization.")
78
+ x_q, x_s, x_zp = ops.scaled_int8_quant(x,
79
+ i_s,
80
+ i_zp,
81
+ symmetric=symmetric)
82
+
83
+ assert x_zp is None, ("AiterScaledMMLinearKernel only supports"
84
+ " symmetric quantization.")
85
+ out_dtype = x.dtype
86
+
87
+ assert (w_q.shape[0] % 16 == 0 and w_q.shape[1] % 16 == 0)
88
+ assert (out_dtype is torch.bfloat16 or out_dtype is torch.float16)
89
+ assert bias is None or bias.shape[0] == w_q.shape[
90
+ 1] and bias.dtype == out_dtype
91
+
92
+ m = x_q.shape[0] # a
93
+ n = w_q.shape[1] # b
94
+
95
+ per_tensor_scale_a = (x_s.numel() == 1)
96
+ per_tensor_scale_b = (w_s.numel() == 1)
97
+ per_token_scale_a = (x_s.numel() == m)
98
+ per_channel_scale_b = (w_s.numel() == n)
99
+
100
+ # @TODO:
101
+ # Maybe broadcast the per-tensor-scale into per-channel-scale
102
+ # if one of the scale is a per-channel-scale.
103
+ # For now, it only supports:
104
+ # - per-tensor-per-tensor a8w8 scaled GEMM, and
105
+ # - per-token-per-channel a8w8 scaled GEMM
106
+ assert ((per_tensor_scale_a and per_tensor_scale_b)
107
+ or (per_token_scale_a and per_channel_scale_b)), (
108
+ "Currently only support per-tensor-per-tensor GEMM " +
109
+ " and per-token-per-channel GEMM through AITER"
110
+ " w8a8 scaled gemm. `AiterScaledMMLinearKernel` " +
111
+ "does not support AITER block scaled GEMM.")
112
+
113
+ from aiter import gemm_a8w8_CK
114
+
115
+ # gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
116
+ # a to be [M, K]
117
+ # b to be [N, K]
118
+ # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
119
+ return gemm_a8w8_CK(x_q, w_q.t(), x_s, w_s, bias).to(out_dtype)