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, 03:30 UTC
[Feature] A calibration-free RTN-based quantization for accurate and accelerated INT4/INT8 inference (#18768)
Signed-off-by: Alex Kogan <alex.kogan@oracle.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Alex Kogan committed
Jul 1, 2025 at 01:44 UTC
27949354faa06035645aa908cc73922500a80b17
3 files changed
+319
tests/quantization/test_rtn.py
new
+28
@@ -0,0 +1,28 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# Copyright © 2025, Oracle and/or its affiliates.
3
+"""Tests RTN quantization startup and generation,
4
+doesn't test correctness
5
+"""
6
+import pytest
7
+
8
+from tests.quantization.utils import is_quant_method_supported
9
+
10
+MODELS = ["microsoft/Phi-3-mini-4k-instruct"]
11
+
12
+
13
+@pytest.mark.skipif(not is_quant_method_supported("rtn"),
14
+ reason="RTN is not supported on this GPU type.")
15
+@pytest.mark.parametrize("model", MODELS)
16
+@pytest.mark.parametrize("dtype", ["bfloat16"])
17
+@pytest.mark.parametrize("max_tokens", [10])
18
+def test_model_rtn_startup(
19
+ hf_runner,
20
+ vllm_runner,
21
+ example_prompts,
22
+ model: str,
23
+ dtype: str,
24
+ max_tokens: int,
25
+) -> None:
26
+
27
+ with vllm_runner(model, dtype=dtype, quantization="rtn") as vllm_model:
28
+ vllm_model.generate_greedy(example_prompts, max_tokens)
vllm/model_executor/layers/quantization/__init__.py
+3
@@ -35,6 +35,7 @@ QuantizationMethods = Literal[
35
"moe_wna16",
36
"torchao",
37
"auto-round",
38
+ "rtn",
39
]
40
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
41
@@ -110,6 +111,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
111
from .neuron_quant import NeuronQuantConfig
112
from .ptpc_fp8 import PTPCFp8Config
113
from .qqq import QQQConfig
114
+ from .rtn import RTNConfig
115
from .torchao import TorchAOConfig
116
from .tpu_int8 import Int8TpuConfig
117
@@ -142,6 +144,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
144
"moe_wna16": MoeWNA16Config,
145
"torchao": TorchAOConfig,
146
"auto-round": AutoRoundConfig,
147
+ "rtn": RTNConfig
148
}
149
# Update the `method_to_config` with customized quantization methods.
150
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
vllm/model_executor/layers/quantization/rtn.py
new
+288
@@ -0,0 +1,288 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# Copyright © 2025, Oracle and/or its affiliates.
3
+
4
+import os
5
+from typing import Any, Optional
6
+
7
+import torch
8
+import torch.nn.functional as F
9
+from torch.nn.parameter import Parameter
10
+
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
13
+ set_weight_attrs)
14
+from vllm.model_executor.layers.quantization import QuantizationMethods
15
+from vllm.model_executor.layers.quantization.base_config import (
16
+ QuantizationConfig)
17
+
18
+logger = init_logger(__name__)
19
+"""By default, use 8 bit as target precision, but it can be
20
+overridden by setting the RTN_NUM_BITS envvar
21
+"""
22
+NUM_BITS = os.getenv('RTN_NUM_BITS', "8")
23
+"""By default, use group size of 128 parameters, but it can be
24
+overridden by setting the RTN_GROUP_SIZE envvar
25
+"""
26
+GROUP_SIZE = os.getenv('RTN_GROUP_SIZE', "128")
27
+
28
+
29
+class RTNConfig(QuantizationConfig):
30
+ """Config class for RTN.
31
+ """
32
+
33
+ def __init__(
34
+ self,
35
+ weight_bits: int = int(NUM_BITS),
36
+ group_size: int = int(GROUP_SIZE),
37
+ ) -> None:
38
+ self.weight_bits = weight_bits
39
+ self.group_size = group_size
40
+
41
+ if self.weight_bits != 4 and self.weight_bits != 8:
42
+ raise ValueError(
43
+ "Currently, only 4-bit or 8-bit weight quantization is "
44
+ f"supported for RTN, but got {self.weight_bits} bits.")
45
+
46
+ def __repr__(self) -> str:
47
+ return (f"RTNConfig(weight_bits={self.weight_bits}, "
48
+ f"group_size={self.group_size})")
49
+
50
+ @classmethod
51
+ def get_name(cls) -> QuantizationMethods:
52
+ return "rtn"
53
+
54
+ @classmethod
55
+ def get_supported_act_dtypes(cls) -> list[torch.dtype]:
56
+ return [torch.bfloat16, torch.half]
57
+
58
+ @classmethod
59
+ def get_min_capability(cls) -> int:
60
+ return 80
61
+
62
+ @classmethod
63
+ def get_config_filenames(cls) -> list[str]:
64
+ return []
65
+
66
+ @classmethod
67
+ def from_config(cls, config: dict[str, Any]) -> "RTNConfig":
68
+ weight_bits = cls.get_from_keys(config, ["bits"])
69
+ group_size = cls.get_from_keys(config, ["group_size"])
70
+ return cls(weight_bits, group_size)
71
+
72
+ def get_quant_method(self, layer: torch.nn.Module,
73
+ prefix: str) -> Optional["RTNLinearMethod"]:
74
+ if isinstance(layer, LinearBase):
75
+ return RTNLinearMethod(self)
76
+ return None
77
+
78
+
79
+class RTNTensor:
80
+ """A wrapper over Tensor that enables quantization on-the-fly by
81
+ overloading the copy_ method.
82
+ """
83
+
84
+ def __init__(self, data: torch.Tensor, scale: torch.Tensor,
85
+ quant_config: RTNConfig) -> None:
86
+ self.data = data
87
+ self.scale = scale
88
+ self.quant_config = quant_config
89
+
90
+ def narrow(self, dim, start, length):
91
+ factor = 1 if self.quant_config.weight_bits == 8 else 2
92
+ return RTNTensor(
93
+ self.data.narrow(dim, start // factor, length // factor),
94
+ self.scale.narrow(dim, start, length), self.quant_config)
95
+
96
+ @property
97
+ def shape(self):
98
+ shape = self.data.shape
99
+ factor = 1 if self.quant_config.weight_bits == 8 else 2
100
+ return torch.Size((shape[0] * factor, shape[1]))
101
+
102
+ def copy_(self, loaded_weight: torch.Tensor) -> None:
103
+ qweight, weight_scale = rtn_quantize(loaded_weight.cuda(),
104
+ self.quant_config.weight_bits,
105
+ self.quant_config.group_size)
106
+
107
+ self.data.copy_(qweight)
108
+ self.scale.data.copy_(weight_scale)
109
+
110
+
111
+class RTNParameter(Parameter):
112
+ """A wrapper over Parameter that returns RTNTensor (a wrapper over Tensor)
113
+ when its data is accessed. We need this wrapper for the data loading phase
114
+ only, so we can intercept a weight copying function (torch.Tensor.copy_)
115
+ and apply quantization on-the-fly.
116
+ """
117
+
118
+ def __new__(cls, data: torch.Tensor, **kwargs):
119
+ return super().__new__(cls, data=data, requires_grad=False)
120
+
121
+ def __init__(self, data: torch.Tensor, scale: torch.Tensor,
122
+ quant_config: RTNConfig) -> None:
123
+ self.scale = scale
124
+ self.quant_config = quant_config
125
+
126
+ @property
127
+ def data(self):
128
+ return RTNTensor(super().data, self.scale, self.quant_config)
129
+
130
+
131
+class RTNLinearMethod(LinearMethodBase):
132
+ """Linear method for RTN.
133
+
134
+ Args:
135
+ quant_config: The RTN quantization config.
136
+ """
137
+
138
+ def __init__(self, quant_config: RTNConfig):
139
+ self.quant_config = quant_config
140
+
141
+ def create_weights(
142
+ self,
143
+ layer: torch.nn.Module,
144
+ input_size_per_partition: int,
145
+ output_partition_sizes: list[int],
146
+ input_size: int,
147
+ output_size: int,
148
+ params_dtype: torch.dtype,
149
+ **extra_weight_attrs,
150
+ ):
151
+ output_size_per_partition = sum(output_partition_sizes)
152
+ num_groups_per_col = (input_size_per_partition //
153
+ self.quant_config.group_size
154
+ if self.quant_config.group_size != -1 else 1)
155
+
156
+ scale = Parameter(
157
+ torch.empty(output_size_per_partition,
158
+ num_groups_per_col,
159
+ dtype=params_dtype),
160
+ requires_grad=False,
161
+ )
162
+ factor = 1 if self.quant_config.weight_bits == 8 else 2
163
+
164
+ weight = RTNParameter(data=torch.empty(output_size_per_partition //
165
+ factor,
166
+ input_size_per_partition,
167
+ dtype=torch.int8),
168
+ scale=scale,
169
+ quant_config=self.quant_config)
170
+
171
+ layer.register_parameter("weight", weight)
172
+ set_weight_attrs(weight, {
173
+ **extra_weight_attrs,
174
+ "input_dim": 1,
175
+ "output_dim": 0,
176
+ })
177
+
178
+ layer.register_parameter("scale", scale)
179
+ layer.output_size_per_partition = output_size_per_partition
180
+
181
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
182
+ """torch.compile does not know how to deal with a Parameter subclass
183
+ (aka RTNParameter). As we don't really need RTNParameters for the
184
+ forward pass, we replace them with equivalent instances of Parameters.
185
+ """
186
+ old_weight = layer.weight
187
+ assert isinstance(old_weight, RTNParameter)
188
+ data = old_weight.data.data
189
+
190
+ delattr(layer, "weight")
191
+
192
+ new_weight = Parameter(data=data, requires_grad=False)
193
+ layer.register_parameter("weight", new_weight)
194
+
195
+ def apply(self,
196
+ layer: torch.nn.Module,
197
+ x: torch.Tensor,
198
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
199
+ qweight = layer.weight
200
+ scale = layer.scale
201
+
202
+ weight = rtn_dequantize(qweight, scale)
203
+ out = F.linear(x, weight)
204
+ del weight
205
+ if bias is not None:
206
+ out.add_(bias)
207
+
208
+ return out
209
+
210
+
211
+def rtn_quantize(tensor: torch.Tensor, num_bits: int,
212
+ group_size: int) -> tuple[torch.Tensor, torch.Tensor]:
213
+ """Quantize a tensor using per-group static scaling factor.
214
+
215
+ Args:
216
+ tensor: The input tensor.
217
+ num_bits: Target precision for the result (supported values are
218
+ 8 or 4).
219
+ group_size: Quantization granularity.
220
+ If equal to -1, each row in the input tensor is treated
221
+ as one group.
222
+ """
223
+
224
+ q_range = 2**num_bits
225
+ num_groups = (tensor.shape[0] * tensor.shape[1] //
226
+ group_size if group_size != -1 else tensor.shape[0])
227
+ """Calculate a scaling factor per input group.
228
+ """
229
+ input_flat = tensor.reshape(num_groups, -1)
230
+ input_min = torch.min(input_flat, dim=1, keepdim=True)[0]
231
+ input_max = torch.max(input_flat, dim=1, keepdim=True)[0]
232
+ input_max_abs = torch.max(input_min.abs(), input_max.abs())
233
+ scale = (input_max_abs * 2.0 / (q_range - 1))
234
+ """Scale each input group, truncate and round to the nearest integer.
235
+ """
236
+ scaled_input = input_flat / scale
237
+ scaled_input = scaled_input.clamp(-q_range // 2, q_range // 2 - 1)
238
+ scaled_input = scaled_input.round()
239
+
240
+ scale = scale.reshape(tensor.shape[0], -1).contiguous()
241
+ inputs_q = scaled_input.reshape(tensor.shape).to(torch.int8)
242
+ inputs_q = inputs_q.contiguous()
243
+
244
+ if num_bits == 4:
245
+ """Pack two 4-bit values into each byte.
246
+ """
247
+ inputs_q = (inputs_q[:, 1::2] << 4) | (inputs_q[:, ::2] & 0xf)
248
+ inputs_q = inputs_q.reshape(tensor.shape[0] // 2, tensor.shape[1])
249
+ inputs_q = inputs_q.contiguous()
250
+
251
+ return inputs_q, scale
252
+
253
+
254
+def rtn_dequantize(tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
255
+ """Dequantize a tensor using per-group static scaling factors.
256
+
257
+ Args:
258
+ tensor: The input tensor.
259
+ scale: The tensor with per-group scale factors.
260
+ """
261
+
262
+ num_groups = scale.size(0) * scale.size(1)
263
+ input_dim, output_dim = tensor.shape
264
+
265
+ num_bits = 8 if input_dim == scale.size(0) else 4
266
+ if num_bits == 4:
267
+ input_dim *= 2
268
+
269
+ data = torch.empty((input_dim, output_dim),
270
+ dtype=scale.dtype,
271
+ device=tensor.device)
272
+
273
+ if num_bits == 8:
274
+ data.copy_(tensor)
275
+ else:
276
+ """Unpack two 4-bit values from each byte.
277
+ """
278
+ tensor = tensor.reshape(input_dim, output_dim // 2)
279
+ for i in range(2):
280
+ data[:, i::2] = (tensor << 4 * (1 - i)) >> 4
281
+ """Scale each input group with its scaling factor.
282
+ """
283
+ scale = scale.reshape(num_groups, -1)
284
+ data = data.reshape(num_groups, -1)
285
+ data = torch.mul(data, scale)
286
+
287
+ input_deq = data.reshape((input_dim, output_dim)).contiguous()
288
+ return input_deq