@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 201 lines 7.12 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 import torch
6
7 from vllm import _custom_ops as ops
8 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
9 MARLIN_SUPPORTED_GROUP_SIZES,
10 apply_gptq_marlin_linear,
11 marlin_act_int8_process_scales,
12 marlin_make_empty,
13 marlin_pad_dim,
14 marlin_pad_qweight,
15 marlin_pad_scales,
16 marlin_padded_nk,
17 marlin_permute_bias,
18 marlin_permute_scales,
19 marlin_zero_points,
20 query_marlin_supported_quant_types,
21 unpack_cols,
22 )
23 from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
24 from vllm.platforms import current_platform
25 from vllm.scalar_type import scalar_types
26
27 from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
28
29
30 class MarlinLinearKernel(MPLinearKernel):
31 @classmethod
32 def get_min_capability(cls) -> int:
33 return 75
34
35 @classmethod
36 def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
37 # Marlin uses inline PTX, so it can only be compatible with Nvidia
38 if not current_platform.is_cuda():
39 return False, "Marlin only supported on CUDA"
40
41 quant_types = query_marlin_supported_quant_types(c.zero_points)
42 if c.weight_type not in quant_types:
43 return (
44 False,
45 f"Quant type ({c.weight_type}) not supported by"
46 f" Marlin, supported types are: {quant_types}",
47 )
48
49 if c.group_size not in MARLIN_SUPPORTED_GROUP_SIZES:
50 return (
51 False,
52 f"Group size ({c.group_size}) not supported by "
53 "Marlin, supported group sizes are: "
54 f"{MARLIN_SUPPORTED_GROUP_SIZES}",
55 )
56
57 # A group straddling TP ranks cannot be fixed by padding.
58 if (
59 c.group_size != -1
60 and c.group_size < c.full_weight_shape[0]
61 and c.partition_weight_shape[0] % c.group_size != 0
62 ):
63 return False, (
64 f"in_features per partition {c.partition_weight_shape[0]} is "
65 f"not divisible by group_size = {c.group_size}."
66 )
67
68 # Tile misalignment is fixed by zero-padding at weight prep.
69 return True, None
70
71 # note assumes that
72 # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
73 # `weight_scale` is: {input_dim = 0, output_dim = 1}
74 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
75 device = getattr(layer, self.w_q_name).device
76 c = self.config
77 is_a_8bit = c.act_type is not None and c.act_type.itemsize == 1
78
79 if is_a_8bit:
80 assert c.weight_type == scalar_types.uint4b8, (
81 "W8A8 is not supported by marlin kernel."
82 )
83
84 if c.act_type == torch.float8_e4m3fn:
85 ops.marlin_int4_fp8_preprocess(getattr(layer, self.w_q_name), inplace=True)
86 getattr(layer, self.w_s_name).data = (
87 getattr(layer, self.w_s_name).data * 512
88 )
89
90 size_k, size_n = c.partition_weight_shape
91 padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
92
93 # Default name since marlin requires empty parameter for zp,
94 # TODO: remove this requirement from marlin (allow optional tensors)
95 if self.w_zp_name is None:
96 self.w_zp_name = "w_zp"
97
98 def transform_w_q(x):
99 assert isinstance(x, BasevLLMParameter)
100 permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
101 x.data = ops.gptq_marlin_repack(
102 marlin_pad_qweight(
103 x.data.contiguous(), size_n, size_k, padded_n, padded_k
104 ),
105 size_k=padded_k,
106 size_n=padded_n,
107 num_bits=c.weight_type.size_bits,
108 is_a_8bit=is_a_8bit,
109 )
110 return x
111
112 def transform_w_s(x):
113 assert isinstance(x, BasevLLMParameter)
114 permute_param_layout_(x, input_dim=0, output_dim=1)
115 x.data = marlin_permute_scales(
116 marlin_pad_scales(
117 x.data.contiguous(),
118 size_n,
119 size_k,
120 padded_n,
121 padded_k,
122 c.group_size,
123 ),
124 size_k=padded_k,
125 size_n=padded_n,
126 group_size=c.group_size,
127 is_a_8bit=is_a_8bit,
128 )
129
130 if c.group_size == -1:
131 num_groups = 1
132 else:
133 num_groups = c.partition_weight_shape[0] // c.group_size
134
135 if c.act_type == torch.int8 and num_groups > 1:
136 x.data, input_global_scale = marlin_act_int8_process_scales(x.data)
137 layer.register_parameter(
138 "input_global_scale",
139 torch.nn.Parameter(input_global_scale, requires_grad=False),
140 )
141 else:
142 layer.input_global_scale = None
143 return x
144
145 if c.zero_points:
146 grouped_k = size_k // c.group_size if c.group_size != -1 else 1
147 padded_grouped_k = padded_k // c.group_size if c.group_size != -1 else 1
148 self._transform_param(
149 layer,
150 self.w_zp_name,
151 lambda x: marlin_zero_points(
152 marlin_pad_scales(
153 unpack_cols(
154 x.t(),
155 c.weight_type.size_bits,
156 grouped_k,
157 size_n,
158 ),
159 size_n,
160 size_k,
161 padded_n,
162 padded_k,
163 c.group_size,
164 ),
165 size_k=padded_grouped_k,
166 size_n=padded_n,
167 num_bits=c.weight_type.size_bits,
168 is_a_8bit=is_a_8bit,
169 ),
170 )
171 else:
172 setattr(layer, self.w_zp_name, marlin_make_empty(device))
173 self._transform_param(layer, self.w_q_name, transform_w_q)
174 self._transform_param(layer, self.w_s_name, transform_w_s)
175
176 if hasattr(layer, "bias") and layer.bias is not None:
177 layer.bias.data = marlin_permute_bias(
178 marlin_pad_dim(layer.bias, size_n, padded_n)
179 )
180
181 def apply_weights(
182 self,
183 layer: torch.nn.Module,
184 x: torch.Tensor,
185 bias: torch.Tensor | None = None,
186 ) -> torch.Tensor:
187 c = self.config
188 w_q, w_s, w_zp = self._get_weight_params(layer)
189 return apply_gptq_marlin_linear(
190 input=x,
191 weight=w_q,
192 weight_scale=w_s,
193 weight_zp=w_zp, # type: ignore
194 workspace=None,
195 wtype=c.weight_type,
196 input_size_per_partition=c.partition_weight_shape[0],
197 output_size_per_partition=c.partition_weight_shape[1],
198 input_global_scale=getattr(layer, "input_global_scale", None),
199 bias=bias,
200 input_dtype=c.act_type,
201 )