Skip to content

vllm.model_executor.layers.quantization.marlin

logger module-attribute 

logger = init_logger(__name__)

MarlinConfig

Bases: QuantizationConfig

Config class for Marlin.

Reference: https://github.com/IST-DASLab/marlin/tree/master

Source code in vllm/model_executor/layers/quantization/marlin.py 
 24
 25
 26
 27
 28
 29
 30
 31
 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
class MarlinConfig(QuantizationConfig):
    """Config class for Marlin.

    Reference: https://github.com/IST-DASLab/marlin/tree/master
    """

    def __init__(
        self,
        group_size: int,
        lm_head_quantized: bool,
    ) -> None:
        super().__init__()

        # Group size for the quantization.
        self.group_size = group_size
        self.lm_head_quantized = lm_head_quantized
        if self.group_size != 128 and self.group_size != -1:
            raise ValueError(
                "Currently, only group size 128 and -1 (channelwise) "
                "is supported for Marlin, but got group_size of "
                f"{self.group_size}")

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

        # Tile size used by marlin kernels.
        self.tile_size = 16

        # Min out_features dim
        self.min_n_threads = 64

        # Min in_features dim
        self.min_k_threads = 128

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

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

    def __repr__(self) -> str:
        return (f"MarlinConfig(group_size={self.group_size}, "
                f"lm_head_quantized={self.lm_head_quantized})")

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

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

    @classmethod
    # Need to figure it out
    def get_min_capability(cls) -> int:
        return 80

    @classmethod
    def get_config_filenames(cls) -> list[str]:
        return ["quantize_config.json"]

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "MarlinConfig":
        group_size = cls.get_from_keys(config, ["group_size"])
        lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
                                                 default=False)
        return cls(group_size, lm_head_quantized)

    @classmethod
    def override_quantization_method(
            cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:
        # compat: autogptq >=0.8.0 use checkpoint_format: str
        # compat: autogptq <=0.7.1 is_marlin_format: bool
        is_marlin_format = (hf_quant_cfg.get("checkpoint_format") == "marlin"
                            or hf_quant_cfg.get("is_marlin_format", False))

        is_valid_user_quant = (user_quant is None or user_quant == "gptq"
                               or user_quant == "marlin")

        if is_marlin_format and is_valid_user_quant:
            msg = ("The model is serialized in {} format. Using {} kernel.".
                   format(cls.get_name(), cls.get_name()))
            logger.info(msg)
            return cls.get_name()

        return None

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

group_size instance-attribute 

group_size = group_size

lm_head_quantized instance-attribute 

lm_head_quantized = lm_head_quantized

max_parallel instance-attribute 

max_parallel = 16

min_k_threads instance-attribute 

min_k_threads = 128

min_n_threads instance-attribute 

min_n_threads = 64

pack_factor instance-attribute 

pack_factor = 32 // 4

perm_len instance-attribute 

perm_len = 1024

tile_size instance-attribute 

tile_size = 16

__init__ 

__init__(group_size: int, lm_head_quantized: bool) -> None
Source code in vllm/model_executor/layers/quantization/marlin.py 
30
31
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
def __init__(
    self,
    group_size: int,
    lm_head_quantized: bool,
) -> None:
    super().__init__()

    # Group size for the quantization.
    self.group_size = group_size
    self.lm_head_quantized = lm_head_quantized
    if self.group_size != 128 and self.group_size != -1:
        raise ValueError(
            "Currently, only group size 128 and -1 (channelwise) "
            "is supported for Marlin, but got group_size of "
            f"{self.group_size}")

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

    # Tile size used by marlin kernels.
    self.tile_size = 16

    # Min out_features dim
    self.min_n_threads = 64

    # Min in_features dim
    self.min_k_threads = 128

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

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

__repr__ 

__repr__() -> str
Source code in vllm/model_executor/layers/quantization/marlin.py 
65
66
67
def __repr__(self) -> str:
    return (f"MarlinConfig(group_size={self.group_size}, "
            f"lm_head_quantized={self.lm_head_quantized})")

from_config classmethod 

from_config(config: dict[str, Any]) -> MarlinConfig
Source code in vllm/model_executor/layers/quantization/marlin.py 
86
87
88
89
90
91
@classmethod
def from_config(cls, config: dict[str, Any]) -> "MarlinConfig":
    group_size = cls.get_from_keys(config, ["group_size"])
    lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
                                             default=False)
    return cls(group_size, lm_head_quantized)

get_config_filenames classmethod 

get_config_filenames() -> list[str]
Source code in vllm/model_executor/layers/quantization/marlin.py 
82
83
84
@classmethod
def get_config_filenames(cls) -> list[str]:
    return ["quantize_config.json"]

get_min_capability classmethod 

get_min_capability() -> int
Source code in vllm/model_executor/layers/quantization/marlin.py 
77
78
79
80
@classmethod
# Need to figure it out
def get_min_capability(cls) -> int:
    return 80

get_name classmethod 

get_name() -> QuantizationMethods
Source code in vllm/model_executor/layers/quantization/marlin.py 
69
70
71
@classmethod
def get_name(cls) -> QuantizationMethods:
    return "marlin"

get_quant_method 

get_quant_method(
    layer: Module, prefix: str
) -> Optional[MarlinLinearMethod]
Source code in vllm/model_executor/layers/quantization/marlin.py 
112
113
114
115
116
117
def get_quant_method(self, layer: torch.nn.Module,
                     prefix: str) -> Optional["MarlinLinearMethod"]:
    if (isinstance(layer, LinearBase) or
        (isinstance(layer, ParallelLMHead) and self.lm_head_quantized)):
        return MarlinLinearMethod(self)
    return None

get_supported_act_dtypes classmethod 

get_supported_act_dtypes() -> list[dtype]
Source code in vllm/model_executor/layers/quantization/marlin.py 
73
74
75
@classmethod
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
    return [torch.half]

override_quantization_method classmethod 

override_quantization_method(
    hf_quant_cfg, user_quant
) -> Optional[QuantizationMethods]
Source code in vllm/model_executor/layers/quantization/marlin.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
@classmethod
def override_quantization_method(
        cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:
    # compat: autogptq >=0.8.0 use checkpoint_format: str
    # compat: autogptq <=0.7.1 is_marlin_format: bool
    is_marlin_format = (hf_quant_cfg.get("checkpoint_format") == "marlin"
                        or hf_quant_cfg.get("is_marlin_format", False))

    is_valid_user_quant = (user_quant is None or user_quant == "gptq"
                           or user_quant == "marlin")

    if is_marlin_format and is_valid_user_quant:
        msg = ("The model is serialized in {} format. Using {} kernel.".
               format(cls.get_name(), cls.get_name()))
        logger.info(msg)
        return cls.get_name()

    return None

MarlinLinearMethod

Bases: LinearMethodBase

Linear method for Marlin.

Parameters:

Name Type Description Default
quant_config MarlinConfig

The Marlin quantization config.

 required
Source code in vllm/model_executor/layers/quantization/marlin.py 
120
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
class MarlinLinearMethod(LinearMethodBase):
    """Linear method for Marlin.

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

    def __init__(self, quant_config: MarlinConfig):
        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,
    ):
        del output_size  # Unused.
        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)

        # Determine if channelwise or not
        input_groups = (1 if self.quant_config.group_size == -1 else
                        input_size_per_partition //
                        self.quant_config.group_size)

        weight_scale_args = {
            "data":
            torch.empty(
                input_groups,
                output_size_per_partition,
                device="cuda",
                dtype=params_dtype,
            ),
            "weight_loader":
            weight_loader
        }
        if input_groups == 1:
            scales = ChannelQuantScaleParameter(output_dim=1,
                                                **weight_scale_args)
        else:
            scales = GroupQuantScaleParameter(output_dim=1,
                                              input_dim=0,
                                              **weight_scale_args)

        # 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", scales)
        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 = Parameter(layer.s.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
        scales = layer.s
        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 = scales.shape[1]

        output_2d = ops.marlin_gemm(x_2d, qweight, scales, 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: MarlinConfig)
Source code in vllm/model_executor/layers/quantization/marlin.py 
127
128
def __init__(self, quant_config: MarlinConfig):
    self.quant_config = quant_config

apply 

apply(
    layer: Module, x: Tensor, bias: Optional[Tensor] = None
) -> Tensor
Source code in vllm/model_executor/layers/quantization/marlin.py 
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
def apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    qweight = layer.B
    scales = layer.s
    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 = scales.shape[1]

    output_2d = ops.marlin_gemm(x_2d, qweight, scales, 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/marlin.py 
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
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,
):
    del output_size  # Unused.
    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)

    # Determine if channelwise or not
    input_groups = (1 if self.quant_config.group_size == -1 else
                    input_size_per_partition //
                    self.quant_config.group_size)

    weight_scale_args = {
        "data":
        torch.empty(
            input_groups,
            output_size_per_partition,
            device="cuda",
            dtype=params_dtype,
        ),
        "weight_loader":
        weight_loader
    }
    if input_groups == 1:
        scales = ChannelQuantScaleParameter(output_dim=1,
                                            **weight_scale_args)
    else:
        scales = GroupQuantScaleParameter(output_dim=1,
                                          input_dim=0,
                                          **weight_scale_args)

    # 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", scales)
    layer.register_parameter("workspace", workspace)

process_weights_after_loading 

process_weights_after_loading(layer: Module) -> None
Source code in vllm/model_executor/layers/quantization/marlin.py 
233
234
235
236
237
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 = Parameter(layer.s.data, requires_grad=False)
    layer.workspace = Parameter(layer.workspace.data, requires_grad=False)