main
py 74 lines 2.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Humming GEMM as a mixed-precision WNA16Int linear kernel."""
4
5 import torch
6
7 from vllm.platforms import current_platform
8 from vllm.utils.import_utils import has_humming
9
10 from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
11
12
13 class HummingLinearKernel(MPLinearKernel):
14 @classmethod
15 def get_min_capability(cls) -> int:
16 return 75
17
18 @classmethod
19 def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
20 if not current_platform.is_cuda():
21 return False, "Humming is only supported on CUDA"
22 if not has_humming():
23 return False, "Humming is not installed"
24 return True, None
25
26 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
27 from vllm.model_executor.layers.quantization.utils.humming import (
28 convert_linear_layer_to_humming_standard,
29 get_humming_linear_compute_config,
30 prepare_humming_linear_layer_config,
31 )
32
33 name_map = {"weight": self.w_q_name, "weight_scale": self.w_s_name}
34 if self.w_zp_name is not None and hasattr(layer, self.w_zp_name):
35 name_map["zero_point"] = self.w_zp_name
36 group_size = self.config.group_size
37 quant_config = {
38 "quant_method": "humming",
39 "dtype": "int" + str(self.config.weight_type.size_bits),
40 "group_size": 0 if group_size == -1 else group_size,
41 "has_zero_point": self.config.zero_points,
42 }
43
44 if self.config.zero_points:
45 assert self.w_zp_name is not None
46 name_map["zero_point"] = self.w_zp_name
47 quant_config["has_zero_point"] = True
48
49 convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
50 input_quant_config = getattr(layer, "_humming_input_quant_config", None)
51 self.layer_config = prepare_humming_linear_layer_config(
52 layer, quant_config, input_quant_config=input_quant_config
53 )
54 self.compute_config = get_humming_linear_compute_config()
55 self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
56
57 def apply_weights(
58 self,
59 layer: torch.nn.Module,
60 x: torch.Tensor,
61 bias: torch.Tensor | None = None,
62 ) -> torch.Tensor:
63 from vllm.model_executor.layers.quantization.utils.humming import (
64 apply_humming_linear,
65 )
66
67 return apply_humming_linear(
68 layer,
69 x,
70 skip_bias_add=bias is None,
71 layer_config=self.layer_config,
72 compute_config=self.compute_config,
73 locks=self.locks,
74 )