Skip to content

vllm.model_executor.layers.quantization.qqq

MARLIN_QQQ_MAX_PARALLEL module-attribute 

MARLIN_QQQ_MAX_PARALLEL = 16

MARLIN_QQQ_MIN_THREAD_K module-attribute 

MARLIN_QQQ_MIN_THREAD_K = 128

MARLIN_QQQ_MIN_THREAD_N module-attribute 

MARLIN_QQQ_MIN_THREAD_N = 64

MARLIN_QQQ_SUPPORTED_GROUP_SIZES module-attribute 

MARLIN_QQQ_SUPPORTED_GROUP_SIZES = [-1, 128]

MARLIN_QQQ_SUPPORTED_NUM_BITS module-attribute 

MARLIN_QQQ_SUPPORTED_NUM_BITS = [4]

MARLIN_QQQ_SUPPORTED_SYM module-attribute 

MARLIN_QQQ_SUPPORTED_SYM = [True]

MARLIN_QQQ_TILE module-attribute 

MARLIN_QQQ_TILE = 16

logger module-attribute 

logger = init_logger(__name__)

QQQConfig

Bases: QuantizationConfig

Config class for QQQ

Reference: https://arxiv.org/pdf/2406.09904

Source code in vllm/model_executor/layers/quantization/qqq.py 
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
class QQQConfig(QuantizationConfig):
    """Config class for QQQ

    Reference: https://arxiv.org/pdf/2406.09904
    """

    def __init__(
        self,
        weight_bits: int,
        group_size: int,
        is_sym: bool = True,
    ) -> None:
        super().__init__()
        self.weight_bits = weight_bits
        self.group_size = group_size
        self.is_sym = is_sym

        # Verify
        if self.weight_bits not in MARLIN_QQQ_SUPPORTED_NUM_BITS:
            raise ValueError(
                f"QQQ does not support weight_bits = {self.weight_bits}. "
                f"Only weight_bits = {MARLIN_QQQ_SUPPORTED_NUM_BITS} "
                "are supported.")
        if self.group_size not in MARLIN_QQQ_SUPPORTED_GROUP_SIZES:
            raise ValueError(
                f"QQQ does not support group_size = {self.group_size}. "
                f"Only group_sizes = {MARLIN_QQQ_SUPPORTED_GROUP_SIZES} "
                "are supported.")
        if self.is_sym not in MARLIN_QQQ_SUPPORTED_SYM:
            raise ValueError(
                f"QQQ does not support is_sym = {self.is_sym}. "
                f"Only sym = {MARLIN_QQQ_SUPPORTED_SYM} are supported.")

        # 4 Bits packed into 32 bit datatype.
        self.pack_factor = 32 // self.weight_bits

        # Tile size used by QQQ kernels.
        self.tile_size = MARLIN_QQQ_TILE

        # Min out_features dim
        self.min_n_threads = MARLIN_QQQ_MIN_THREAD_N

        # Min in_features dim
        self.min_k_threads = MARLIN_QQQ_MIN_THREAD_K

        # Max parallel problems to solve at once (improves large
        # batch performance)
        self.max_parallel = MARLIN_QQQ_MAX_PARALLEL

        # Permutation length used by the QQQ kernels.
        self.perm_len = 1024

    def __repr__(self) -> str:
        return "QQQConfig(weight_bits={}, group_size={})".format(
            self.weight_bits, self.group_size)

    @classmethod
    def get_name(cls) -> QuantizationMethods:
        return "qqq"

    @classmethod
    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
        return [torch.half]

    @classmethod
    def get_min_capability(cls) -> int:
        return 80

    @classmethod
    def get_config_filenames(cls) -> list[str]:
        """List of filenames to search for in the model directory."""
        return [
            "quant_config.json",
            "quantize_config.json",
        ]

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "QQQConfig":
        weight_bits = cls.get_from_keys(config, ["wbits"])
        group_size = cls.get_from_keys(config, ["group_size"])
        return cls(weight_bits, group_size)

    def get_quant_method(self, layer: torch.nn.Module,
                         prefix: str) -> Optional["QQQLinearMethod"]:
        if isinstance(layer, LinearBase):
            return QQQLinearMethod(self)
        return None

group_size instance-attribute 

group_size = group_size

is_sym instance-attribute 

is_sym = is_sym

max_parallel instance-attribute 

max_parallel = MARLIN_QQQ_MAX_PARALLEL

min_k_threads instance-attribute 

min_k_threads = MARLIN_QQQ_MIN_THREAD_K

min_n_threads instance-attribute 

min_n_threads = MARLIN_QQQ_MIN_THREAD_N

pack_factor instance-attribute 

pack_factor = 32 // weight_bits

perm_len instance-attribute 

perm_len = 1024

tile_size instance-attribute 

tile_size = MARLIN_QQQ_TILE

weight_bits instance-attribute 

weight_bits = weight_bits

__init__ 

__init__(
    weight_bits: int, group_size: int, is_sym: bool = True
) -> None
Source code in vllm/model_executor/layers/quantization/qqq.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def __init__(
    self,
    weight_bits: int,
    group_size: int,
    is_sym: bool = True,
) -> None:
    super().__init__()
    self.weight_bits = weight_bits
    self.group_size = group_size
    self.is_sym = is_sym

    # Verify
    if self.weight_bits not in MARLIN_QQQ_SUPPORTED_NUM_BITS:
        raise ValueError(
            f"QQQ does not support weight_bits = {self.weight_bits}. "
            f"Only weight_bits = {MARLIN_QQQ_SUPPORTED_NUM_BITS} "
            "are supported.")
    if self.group_size not in MARLIN_QQQ_SUPPORTED_GROUP_SIZES:
        raise ValueError(
            f"QQQ does not support group_size = {self.group_size}. "
            f"Only group_sizes = {MARLIN_QQQ_SUPPORTED_GROUP_SIZES} "
            "are supported.")
    if self.is_sym not in MARLIN_QQQ_SUPPORTED_SYM:
        raise ValueError(
            f"QQQ does not support is_sym = {self.is_sym}. "
            f"Only sym = {MARLIN_QQQ_SUPPORTED_SYM} are supported.")

    # 4 Bits packed into 32 bit datatype.
    self.pack_factor = 32 // self.weight_bits

    # Tile size used by QQQ kernels.
    self.tile_size = MARLIN_QQQ_TILE

    # Min out_features dim
    self.min_n_threads = MARLIN_QQQ_MIN_THREAD_N

    # Min in_features dim
    self.min_k_threads = MARLIN_QQQ_MIN_THREAD_K

    # Max parallel problems to solve at once (improves large
    # batch performance)
    self.max_parallel = MARLIN_QQQ_MAX_PARALLEL

    # Permutation length used by the QQQ kernels.
    self.perm_len = 1024

__repr__ 

__repr__() -> str
Source code in vllm/model_executor/layers/quantization/qqq.py 
84
85
86
def __repr__(self) -> str:
    return "QQQConfig(weight_bits={}, group_size={})".format(
        self.weight_bits, self.group_size)

from_config classmethod 

from_config(config: dict[str, Any]) -> QQQConfig
Source code in vllm/model_executor/layers/quantization/qqq.py 
108
109
110
111
112
@classmethod
def from_config(cls, config: dict[str, Any]) -> "QQQConfig":
    weight_bits = cls.get_from_keys(config, ["wbits"])
    group_size = cls.get_from_keys(config, ["group_size"])
    return cls(weight_bits, group_size)

get_config_filenames classmethod 

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

Source code in vllm/model_executor/layers/quantization/qqq.py 
100
101
102
103
104
105
106
@classmethod
def get_config_filenames(cls) -> list[str]:
    """List of filenames to search for in the model directory."""
    return [
        "quant_config.json",
        "quantize_config.json",
    ]

get_min_capability classmethod 

get_min_capability() -> int
Source code in vllm/model_executor/layers/quantization/qqq.py 
96
97
98
@classmethod
def get_min_capability(cls) -> int:
    return 80

get_name classmethod 

get_name() -> QuantizationMethods
Source code in vllm/model_executor/layers/quantization/qqq.py 
88
89
90
@classmethod
def get_name(cls) -> QuantizationMethods:
    return "qqq"

get_quant_method 

get_quant_method(
    layer: Module, prefix: str
) -> Optional[QQQLinearMethod]
Source code in vllm/model_executor/layers/quantization/qqq.py 
114
115
116
117
118
def get_quant_method(self, layer: torch.nn.Module,
                     prefix: str) -> Optional["QQQLinearMethod"]:
    if isinstance(layer, LinearBase):
        return QQQLinearMethod(self)
    return None

get_supported_act_dtypes classmethod 

get_supported_act_dtypes() -> list[dtype]
Source code in vllm/model_executor/layers/quantization/qqq.py 
92
93
94
@classmethod
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
    return [torch.half]

QQQLinearMethod

Bases: LinearMethodBase

Linear method for QQQ.

Parameters:

Name Type Description Default
quant_config QQQConfig

The QQQ quantization config.

 required
Source code in vllm/model_executor/layers/quantization/qqq.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
class QQQLinearMethod(LinearMethodBase):
    """Linear method for QQQ.

    Args:
        quant_config: The QQQ quantization config.
    """

    def __init__(self, quant_config: QQQConfig):
        self.quant_config = quant_config

    def create_weights(
        self,
        layer: torch.nn.Module,
        input_size_per_partition: int,
        output_partition_sizes: list[int],
        input_size: int,
        output_size: int,
        params_dtype: torch.dtype,
        **extra_weight_attrs,
    ):
        weight_loader = extra_weight_attrs["weight_loader"]
        if params_dtype != torch.float16:
            raise ValueError(
                f"The params dtype must be float16, but got {params_dtype}")

        # Validate output_size_per_partition
        output_size_per_partition = sum(output_partition_sizes)
        if output_size_per_partition % self.quant_config.min_n_threads != 0:
            raise ValueError(
                f"Weight output_size_per_partition = "
                f"{output_size_per_partition} is not divisible by "
                f"min_n_threads = {self.quant_config.min_n_threads}.")
        if output_size_per_partition % self.quant_config.pack_factor != 0:
            raise ValueError(
                f"Weight output_size_per_partition = "
                f"{output_size_per_partition} is not divisible by "
                f"pack_factor = {self.quant_config.pack_factor}.")

        # Validate input_size_per_partition
        if input_size_per_partition % self.quant_config.min_k_threads != 0:
            raise ValueError(
                f"Weight input_size_per_partition = "
                f"{input_size_per_partition} is not divisible by "
                f"min_k_threads = {self.quant_config.min_k_threads}.")
        if (self.quant_config.group_size != -1 and
                input_size_per_partition % self.quant_config.group_size != 0):
            raise ValueError(f"Weight input_size_per_partition = "
                             f"{input_size_per_partition} is not divisible by "
                             f"group_size = {self.quant_config.group_size}.")

        # Check that we have at least 4 tiles horizontally in the shard
        num_tiles_per_perm = self.quant_config.perm_len // (
            self.quant_config.tile_size**2)
        if output_size_per_partition % num_tiles_per_perm != 0:
            raise ValueError(
                "Each permutation group must reside on the same gpu")

        # Quantized 4Bit weights packed into Int32.
        qweight = PackedvLLMParameter(
            data=torch.empty(
                input_size_per_partition // self.quant_config.tile_size,
                output_size_per_partition * self.quant_config.tile_size //
                self.quant_config.pack_factor,
                device="cuda",
                dtype=torch.int32,
            ),
            input_dim=0,
            output_dim=1,
            packed_dim=1,
            packed_factor=self.quant_config.pack_factor,
            marlin_tile_size=self.quant_config.tile_size,
            weight_loader=weight_loader)

        s_channel = ChannelQuantScaleParameter(data=torch.empty(
            1,
            output_size_per_partition,
            device="cuda",
            dtype=torch.float,
        ),
                                               weight_loader=weight_loader,
                                               output_dim=1)

        if self.quant_config.group_size == -1:
            s_group_data = torch.tensor(
                [],
                device="cuda",
                dtype=torch.half,
            )
        else:
            s_group_data = torch.empty(
                input_size_per_partition // self.quant_config.group_size,
                output_size_per_partition,
                device="cuda",
                dtype=torch.half,
            )

        s_group_attr = {"data": s_group_data, "weight_loader": weight_loader}

        if self.quant_config.group_size == -1:
            s_group = BasevLLMParameter(**s_group_attr)
        else:
            s_group = GroupQuantScaleParameter(output_dim=1,
                                               input_dim=0,
                                               **s_group_attr)

        # Allocate workspace (Used for internal locking mechanism)
        max_workspace_size = (
            output_size_per_partition //
            self.quant_config.min_n_threads) * self.quant_config.max_parallel

        workspace = BasevLLMParameter(data=torch.zeros(max_workspace_size,
                                                       device="cuda",
                                                       dtype=torch.int),
                                      weight_loader=weight_loader)

        layer.register_parameter("B", qweight)
        layer.register_parameter("s_channel", s_channel)
        layer.register_parameter("s_group", s_group)
        layer.register_parameter("workspace", workspace)

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        # required by torch.compile
        layer.B = Parameter(layer.B.data, requires_grad=False)
        layer.s_channel = Parameter(layer.s_channel.data, requires_grad=False)
        layer.s_group = Parameter(layer.s_group.data, requires_grad=False)
        layer.workspace = Parameter(layer.workspace.data, requires_grad=False)

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        qweight = layer.B
        s_ch = layer.s_channel
        s_group = layer.s_group
        workspace = layer.workspace

        x_2d = x.view(-1, x.shape[-1])

        size_m = x_2d.shape[0]
        size_k = x_2d.shape[1]
        size_n = s_ch.shape[1]

        x_int8, s_tok, _ = ops.scaled_int8_quant(x_2d)

        output_2d = ops.marlin_qqq_gemm(x_int8, qweight, s_tok, s_ch, s_group,
                                        workspace, size_m, size_n, size_k)

        output = output_2d.view(x.shape[:-1] + (output_2d.shape[1], ))

        if bias is not None:
            output.add_(bias)  # In-place add

        return output

quant_config instance-attribute 

quant_config = quant_config

__init__ 

__init__(quant_config: QQQConfig)
Source code in vllm/model_executor/layers/quantization/qqq.py 
128
129
def __init__(self, quant_config: QQQConfig):
    self.quant_config = quant_config

apply 

apply(
    layer: Module, x: Tensor, bias: Optional[Tensor] = None
) -> Tensor
Source code in vllm/model_executor/layers/quantization/qqq.py 
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
def apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    qweight = layer.B
    s_ch = layer.s_channel
    s_group = layer.s_group
    workspace = layer.workspace

    x_2d = x.view(-1, x.shape[-1])

    size_m = x_2d.shape[0]
    size_k = x_2d.shape[1]
    size_n = s_ch.shape[1]

    x_int8, s_tok, _ = ops.scaled_int8_quant(x_2d)

    output_2d = ops.marlin_qqq_gemm(x_int8, qweight, s_tok, s_ch, s_group,
                                    workspace, size_m, size_n, size_k)

    output = output_2d.view(x.shape[:-1] + (output_2d.shape[1], ))

    if bias is not None:
        output.add_(bias)  # In-place add

    return output

create_weights 

create_weights(
    layer: Module,
    input_size_per_partition: int,
    output_partition_sizes: list[int],
    input_size: int,
    output_size: int,
    params_dtype: dtype,
    **extra_weight_attrs,
)
Source code in vllm/model_executor/layers/quantization/qqq.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
def create_weights(
    self,
    layer: torch.nn.Module,
    input_size_per_partition: int,
    output_partition_sizes: list[int],
    input_size: int,
    output_size: int,
    params_dtype: torch.dtype,
    **extra_weight_attrs,
):
    weight_loader = extra_weight_attrs["weight_loader"]
    if params_dtype != torch.float16:
        raise ValueError(
            f"The params dtype must be float16, but got {params_dtype}")

    # Validate output_size_per_partition
    output_size_per_partition = sum(output_partition_sizes)
    if output_size_per_partition % self.quant_config.min_n_threads != 0:
        raise ValueError(
            f"Weight output_size_per_partition = "
            f"{output_size_per_partition} is not divisible by "
            f"min_n_threads = {self.quant_config.min_n_threads}.")
    if output_size_per_partition % self.quant_config.pack_factor != 0:
        raise ValueError(
            f"Weight output_size_per_partition = "
            f"{output_size_per_partition} is not divisible by "
            f"pack_factor = {self.quant_config.pack_factor}.")

    # Validate input_size_per_partition
    if input_size_per_partition % self.quant_config.min_k_threads != 0:
        raise ValueError(
            f"Weight input_size_per_partition = "
            f"{input_size_per_partition} is not divisible by "
            f"min_k_threads = {self.quant_config.min_k_threads}.")
    if (self.quant_config.group_size != -1 and
            input_size_per_partition % self.quant_config.group_size != 0):
        raise ValueError(f"Weight input_size_per_partition = "
                         f"{input_size_per_partition} is not divisible by "
                         f"group_size = {self.quant_config.group_size}.")

    # Check that we have at least 4 tiles horizontally in the shard
    num_tiles_per_perm = self.quant_config.perm_len // (
        self.quant_config.tile_size**2)
    if output_size_per_partition % num_tiles_per_perm != 0:
        raise ValueError(
            "Each permutation group must reside on the same gpu")

    # Quantized 4Bit weights packed into Int32.
    qweight = PackedvLLMParameter(
        data=torch.empty(
            input_size_per_partition // self.quant_config.tile_size,
            output_size_per_partition * self.quant_config.tile_size //
            self.quant_config.pack_factor,
            device="cuda",
            dtype=torch.int32,
        ),
        input_dim=0,
        output_dim=1,
        packed_dim=1,
        packed_factor=self.quant_config.pack_factor,
        marlin_tile_size=self.quant_config.tile_size,
        weight_loader=weight_loader)

    s_channel = ChannelQuantScaleParameter(data=torch.empty(
        1,
        output_size_per_partition,
        device="cuda",
        dtype=torch.float,
    ),
                                           weight_loader=weight_loader,
                                           output_dim=1)

    if self.quant_config.group_size == -1:
        s_group_data = torch.tensor(
            [],
            device="cuda",
            dtype=torch.half,
        )
    else:
        s_group_data = torch.empty(
            input_size_per_partition // self.quant_config.group_size,
            output_size_per_partition,
            device="cuda",
            dtype=torch.half,
        )

    s_group_attr = {"data": s_group_data, "weight_loader": weight_loader}

    if self.quant_config.group_size == -1:
        s_group = BasevLLMParameter(**s_group_attr)
    else:
        s_group = GroupQuantScaleParameter(output_dim=1,
                                           input_dim=0,
                                           **s_group_attr)

    # Allocate workspace (Used for internal locking mechanism)
    max_workspace_size = (
        output_size_per_partition //
        self.quant_config.min_n_threads) * self.quant_config.max_parallel

    workspace = BasevLLMParameter(data=torch.zeros(max_workspace_size,
                                                   device="cuda",
                                                   dtype=torch.int),
                                  weight_loader=weight_loader)

    layer.register_parameter("B", qweight)
    layer.register_parameter("s_channel", s_channel)
    layer.register_parameter("s_group", s_group)
    layer.register_parameter("workspace", workspace)

process_weights_after_loading 

process_weights_after_loading(layer: Module) -> None
Source code in vllm/model_executor/layers/quantization/qqq.py 
241
242
243
244
245
246
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    # required by torch.compile
    layer.B = Parameter(layer.B.data, requires_grad=False)
    layer.s_channel = Parameter(layer.s_channel.data, requires_grad=False)
    layer.s_group = Parameter(layer.s_group.data, requires_grad=False)
    layer.workspace = Parameter(layer.workspace.data, requires_grad=False)