@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 159 lines 5.97 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.model_executor.layers.quantization.utils import replace_parameter
8 from vllm.platforms import CpuArchEnum, current_platform
9 from vllm.scalar_type import scalar_types
10
11 from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
12
13
14 # This implementation is for the KleidiAI-accelerated w4a8int quantization
15 # scheme on Arm CPUs:
16 # torch.ops.aten._dyn_quant_matmul_4bit performs dynamic quantized matmul
17 # it takes:
18 # - int4 weights packed along with bias/scales by
19 # torch.ops.aten._dyn_quant_pack_4bit_weight
20 # - float32/bfloat16 activations
21 # then it leverages KleidiAI ukernels that:
22 # - dynamically quantize the activations to int8
23 # - unpack the int4 weights to int8
24 # - perform int8 x int8 -> int32 matmul
25 # - dequantize the int32 output to float32/bfloat16 outputs
26 class Dynamic4bitLinearKernel(MPLinearKernel):
27 SUPPORTED_QUANT_TYPES = [scalar_types.int4]
28
29 @classmethod
30 def get_min_capability(cls) -> int:
31 return 1
32
33 @classmethod
34 def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
35 if not current_platform.is_cpu():
36 return False, "Only CPU is supported"
37 if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
38 return False, f"Unsupported quant type {c.weight_type}"
39 if (
40 current_platform.get_cpu_architecture() == CpuArchEnum.ARM
41 and c.act_type
42 not in [
43 torch.float32,
44 torch.bfloat16,
45 torch.float16,
46 ]
47 ):
48 return (
49 False,
50 "Dynamic4bitLinearKernel on Arm requires Float32 or"
51 " BFloat16 or Float16 activations",
52 )
53 if c.full_weight_shape[0] % c.group_size != 0:
54 return (
55 False,
56 f"Group size ({c.group_size}) does not evenly divide"
57 " the number of input features "
58 f"({c.full_weight_shape[0]})",
59 )
60 if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
61 try:
62 # Attempt to retrieve the operation
63 _ = torch.ops.aten._dyn_quant_matmul_4bit
64 except AttributeError:
65 return (
66 False,
67 f"PyTorch {torch.__version__} does not support"
68 " _dyn_quant_matmul_4bit. Install a newer version",
69 )
70 return True, None
71
72 def process_weights_after_loading(self, layer: torch.nn.Module):
73 c = self.config
74 packed_weight = getattr(layer, self.w_q_name)
75 packed_weight = packed_weight.add(8)
76 uint8_packed = (packed_weight[::, 1::2] << 4 | packed_weight[::, ::2]).to(
77 torch.uint8
78 )
79
80 scales = getattr(layer, self.w_s_name)
81 block_size = c.group_size
82
83 # Handle scaling factors for partitioned weights
84 if block_size == c.partition_weight_shape[0]:
85 scales = scales.to(
86 torch.float32
87 ) # Float32 & Bfloat16 variants requires float32 scales
88 scales = scales.view(-1, 1) # Channel-wise scales
89 if layer.bias is not None:
90 # Float32 & Bfloat16 variants requires float32 bias
91 replace_parameter(
92 layer,
93 "bias",
94 torch.nn.Parameter(
95 layer.bias.to(torch.float32), requires_grad=False
96 ),
97 )
98 else:
99 # KleidiAI kernel requires bfloat16 scales with groupwise scheme
100 scales = scales.to(torch.bfloat16)
101
102 # Repack weights as per kernel requirement
103 w = torch.ops.aten._dyn_quant_pack_4bit_weight(
104 uint8_packed,
105 scales,
106 layer.bias,
107 block_size,
108 c.partition_weight_shape[0],
109 c.partition_weight_shape[1],
110 )
111 replace_parameter(
112 layer, self.w_q_name, torch.nn.Parameter(w, requires_grad=False)
113 )
114 setattr(layer, self.w_s_name, None)
115
116 def apply_weights(
117 self,
118 layer: torch.nn.Module,
119 x: torch.Tensor,
120 bias: torch.Tensor | None = None,
121 ) -> torch.Tensor:
122 # PyTorch / KleidiAI kernels natively support the following configs:
123 # - channelwise with bfloat16 / float32 activations
124 # - groupwise with float32 activations
125 # To support:
126 # - groupwise with bfloat16/float16 activations: we need to upcast
127 # activations to float32 before matmul and downcast back to bfloat16/float16
128 # - channelwise with float16 activations, we need to upcast activations to
129 # float32 before matmul and downcast back to float16
130 # Note: these activations will be dynamically quantized to int8 by the kernel.
131
132 c = self.config
133 is_groupwise = c.group_size != c.partition_weight_shape[0]
134 # dtype of activations before they get dynamically quantized to int8
135 original_pre_quant_act_dtype = x.dtype
136 pre_quant_act_dtype = original_pre_quant_act_dtype
137 if (
138 is_groupwise and pre_quant_act_dtype == torch.bfloat16
139 ) or pre_quant_act_dtype == torch.float16:
140 pre_quant_act_dtype = torch.float32
141
142 x_2d = x.reshape(-1, x.shape[-1])
143 if pre_quant_act_dtype != original_pre_quant_act_dtype:
144 x_2d = x_2d.to(pre_quant_act_dtype)
145
146 out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
147
148 w_q = getattr(layer, self.w_q_name)
149 output = torch.ops.aten._dyn_quant_matmul_4bit(
150 x_2d,
151 w_q,
152 c.group_size,
153 c.partition_weight_shape[0],
154 c.partition_weight_shape[1],
155 ).reshape(out_shape)
156
157 if pre_quant_act_dtype != original_pre_quant_act_dtype:
158 output = output.to(original_pre_quant_act_dtype)
159 return output