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 5, 12:15 UTC
[ROCm][Quantization] Add asymmetric INT8 quantization support to TritonInt8ScaledMMLinearKernel (#38501)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Apr 5, 2026 at 20:42 UTC
780ba37458362bdc0596c6511e17749d44b145fc
7 files changed
+133
-14
.buildkite/lm-eval-harness/configs/Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml
+3
@@ -1,6 +1,9 @@
1
# For hf script, without -t option (tensor parallel size).
2
# bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5
3
model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
4
+required_gpu_arch:
5
+ - gfx942
6
+ - gfx950
7
tasks:
8
- name: "mmlu_pro"
9
metrics:
.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
+3
@@ -1,6 +1,9 @@
1
# For vllm script, with -t option (tensor parallel size)
2
# bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic -l 1319 -t 1
3
model_name: "RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic"
4
+required_gpu_arch:
5
+ - gfx942
6
+ - gfx950
7
tasks:
8
- name: "gsm8k"
9
metrics:
.buildkite/lm-eval-harness/configs/Qwen3-235B-A22B-Instruct-2507-FP8.yaml
+3
@@ -1,4 +1,7 @@
1
model_name: "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8"
2
+required_gpu_arch:
3
+ - gfx942
4
+ - gfx950
5
tasks:
6
- name: "mmlu_pro"
7
metrics:
.buildkite/lm-eval-harness/configs/models-small-rocm.txt
+1
@@ -1,5 +1,6 @@
1
Qwen2.5-1.5B-Instruct.yaml
2
Meta-Llama-3.2-1B-Instruct-INT8-compressed-tensors.yaml
3
+Meta-Llama-3-8B-Instruct-INT8-compressed-tensors-asym.yaml
4
Meta-Llama-3-8B-Instruct-nonuniform-compressed-tensors.yaml
5
Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
6
Qwen1.5-MoE-W4A16-compressed-tensors.yaml
.buildkite/lm-eval-harness/test_lm_eval_correctness.py
+32
@@ -13,6 +13,7 @@ import os
13
from contextlib import contextmanager
14
15
import lm_eval
16
+import pytest
17
import yaml
18
19
from vllm.platforms import current_platform
@@ -89,9 +90,40 @@ def launch_lm_eval(eval_config, tp_size):
90
return results
91
92
93
+def _check_rocm_gpu_arch_requirement(eval_config):
94
+ """Skip the test if the model requires a ROCm GPU arch not present.
95
+
96
+ Model YAML configs can specify::
97
+
98
+ required_gpu_arch:
99
+ - gfx942
100
+ - gfx950
101
+
102
+ The check only applies on ROCm. On other platforms (e.g. CUDA) the
103
+ field is ignored so that shared config files work for both NVIDIA and
104
+ AMD CI pipelines.
105
+ """
106
+ required_archs = eval_config.get("required_gpu_arch")
107
+ if not required_archs:
108
+ return
109
+
110
+ if not current_platform.is_rocm():
111
+ return
112
+
113
+ from vllm.platforms.rocm import _GCN_ARCH # noqa: E402
114
+
115
+ if not any(arch in _GCN_ARCH for arch in required_archs):
116
+ pytest.skip(
117
+ f"Model requires GPU arch {required_archs}, "
118
+ f"but detected arch is '{_GCN_ARCH}'"
119
+ )
120
+
121
+
122
def test_lm_eval_correctness_param(config_filename, tp_size):
123
eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
124
125
+ _check_rocm_gpu_arch_requirement(eval_config)
126
+
127
results = launch_lm_eval(eval_config, tp_size)
128
129
rtol = eval_config.get("rtol", DEFAULT_RTOL)
.buildkite/test-amd.yaml
+18
@@ -2690,6 +2690,24 @@ steps:
2690
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
2691
2692
2693
+- label: LM Eval Small Models (MI325) # TBD
2694
+ timeout_in_minutes: 180
2695
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
2696
+ agent_pool: mi325_1
2697
+ working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
2698
+ source_file_dependencies:
2699
+ - csrc/
2700
+ - vllm/model_executor/layers/quantization
2701
+ - vllm/model_executor/models/
2702
+ - vllm/model_executor/model_loader/
2703
+ - vllm/v1/attention/backends/
2704
+ - vllm/v1/attention/selector.py
2705
+ - vllm/_aiter_ops.py
2706
+ - vllm/platforms/rocm.py
2707
+ commands:
2708
+ - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt
2709
+
2710
+
2711
- label: LM Eval Small Models (B200-MI325) # TBD
2712
timeout_in_minutes: 180
2713
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
vllm/model_executor/kernels/linear/scaled_mm/triton.py
+73
-14
@@ -31,8 +31,6 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
31
32
@classmethod
33
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
34
- if not c.input_symmetric:
35
- return False, "supports symmetric input only."
34
return True, None
35
36
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -62,17 +60,59 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
60
# INPUT SCALE
61
if self.config.is_static_input_scheme:
62
assert i_s is not None
65
- replace_parameter(
66
- layer,
67
- i_s_name,
68
- torch.nn.Parameter(i_s.max(), requires_grad=False),
69
- )
70
- setattr(layer, i_zp_name, None)
63
+
64
+ if self.config.input_symmetric:
65
+ replace_parameter(
66
+ layer,
67
+ i_s_name,
68
+ torch.nn.Parameter(i_s.max(), requires_grad=False),
69
+ )
70
+ setattr(layer, i_zp_name, None)
71
+ else:
72
+ input_zero_point = getattr(layer, i_zp_name)
73
+
74
+ # Reconstruct the ranges to find a single scale and azp
75
+ int8_traits = torch.iinfo(torch.int8)
76
+ azps = input_zero_point.to(dtype=torch.int32)
77
+ range_max = (i_s * (int8_traits.max - azps)).max()
78
+ range_min = (i_s * (int8_traits.min - azps)).min()
79
+
80
+ scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
81
+ replace_parameter(
82
+ layer,
83
+ i_s_name,
84
+ torch.nn.Parameter(scale, requires_grad=False),
85
+ )
86
+
87
+ # AZP loaded as int8 but used as int32
88
+ azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32)
89
+ replace_parameter(
90
+ layer,
91
+ i_zp_name,
92
+ torch.nn.Parameter(azp, requires_grad=False),
93
+ )
94
else:
95
setattr(layer, i_s_name, None)
96
setattr(layer, i_zp_name, None)
97
75
- setattr(layer, azp_adj_name, None)
98
+ # azp_adj is the AZP adjustment term, used to account for weights.
99
+ # It does not depend on scales or azp, so it is the same for
100
+ # static and dynamic quantization.
101
+ # See csrc/quantization/w8a8/cutlass/Epilogues.md for the math.
102
+ if not self.config.input_symmetric:
103
+ weight = getattr(layer, w_q_name)
104
+ # weight is already transposed to [K, N], sum over K (dim=0)
105
+ azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)
106
+ if self.config.is_static_input_scheme:
107
+ # Fold azp into azp_adj for the per-tensor case
108
+ azp_adj = getattr(layer, i_zp_name) * azp_adj
109
+ setattr(
110
+ layer,
111
+ azp_adj_name,
112
+ torch.nn.Parameter(azp_adj, requires_grad=False),
113
+ )
114
+ else:
115
+ setattr(layer, azp_adj_name, None)
116
117
def apply_weights(
118
self,
@@ -80,14 +120,33 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
120
x: torch.Tensor,
121
bias: torch.Tensor | None = None,
122
) -> torch.Tensor:
83
- w_q, w_s, i_s, i_zp, _ = self._get_layer_params(layer)
123
+ w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
124
125
+ symmetric = azp_adj is None
126
x_q, x_s, x_zp = ops.scaled_int8_quant(
86
- x.contiguous(), i_s, i_zp, symmetric=True
127
+ x.contiguous(), i_s, i_zp, symmetric=symmetric
128
)
129
89
- assert x_zp is None, "Triton kernel only supports symmetric quantization"
90
-
91
- return triton_scaled_mm(
130
+ out = triton_scaled_mm(
131
x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
132
)
133
+
134
+ if azp_adj is not None:
135
+ # Asymmetric quantization: subtract the zero-point correction.
136
+ # D = scale_a * scale_b * (A_q @ B_q - azp * azp_adj) + bias
137
+ # triton_scaled_mm already computed scale_a * scale_b * (A_q @ B_q) + bias
138
+ # so we subtract scale_a * scale_b * azp * azp_adj
139
+ #
140
+ # x_s: [M, 1] or scalar, w_s: [N, 1] or scalar, azp_adj: [1, N]
141
+ # Reshape w_s from [N, 1] to [1, N] for proper broadcasting.
142
+ w_s_row = w_s.view(1, -1) if w_s.dim() > 0 else w_s
143
+ static = i_zp is not None
144
+ if not static and x_zp is not None:
145
+ # Dynamic per-token: azp is per-token, azp_adj is per-channel
146
+ # x_zp: [M, 1], azp_adj: [1, N]
147
+ out -= x_s * w_s_row * (x_zp * azp_adj).to(x.dtype)
148
+ else:
149
+ # Static per-tensor: azp already folded into azp_adj
150
+ out -= (x_s * w_s_row * azp_adj).to(x.dtype)
151
+
152
+ return out