@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8d32dc603d
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 13:30 UTC

[Kernel] Support Microsoft Runtime Kernel Lib for our Low Precision Computation - BitBLAS (#6036)

Signed-off-by: xinyuxiao <xinyuxiao2024@gmail.com> Co-authored-by: xinyuxiao <xinyuxiao2024@gmail.com>

Lei Wang committed Apr 22, 2025 at 16:01 UTC 8d32dc603d03fa2d9b2fe538fbc53ef049433762
15 files changed +1864 -7
benchmarks/kernels/benchmark_bitblas.py new
+236
@@ -0,0 +1,236 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# Copyright (c) Microsoft Corporation.
3 +# Licensed under the MIT License.
4 +
5 +from vllm.model_executor.layers.quantization.utils.bitblas_utils import (
6 + MINIMUM_BITBLAS_VERSION)
7 +
8 +try:
9 + import bitblas
10 + if bitblas.__version__ < MINIMUM_BITBLAS_VERSION:
11 + raise ImportError("bitblas version is wrong. Please "
12 + f"install bitblas>={MINIMUM_BITBLAS_VERSION}")
13 +except ImportError as e:
14 + bitblas_import_exception = e
15 + raise ValueError("Trying to use the bitblas backend, but could not import"
16 + f"with the following error: {bitblas_import_exception}. "
17 + "Please install bitblas through the following command: "
18 + f"`pip install bitblas>={MINIMUM_BITBLAS_VERSION}`"
19 + ) from bitblas_import_exception
20 +
21 +from bitblas import Matmul, MatmulConfig, auto_detect_nvidia_target
22 +
23 +from vllm.utils import FlexibleArgumentParser
24 +
25 +parser = FlexibleArgumentParser(
26 + description="Benchmark BitBLAS int4 on a specific target.")
27 +
28 +# Add arguments to the parser
29 +parser.add_argument(
30 + "--target",
31 + type=str,
32 + default=auto_detect_nvidia_target(),
33 + help="Specify the target device for benchmarking.",
34 +)
35 +parser.add_argument("--group_size",
36 + type=int,
37 + default=None,
38 + help="Group size for grouped quantization.")
39 +parser.add_argument(
40 + "--A_dtype",
41 + type=str,
42 + default="float16",
43 + choices=["float16", "float32", "float64", "int32", "int8"],
44 + help="Data type of activation A.",
45 +)
46 +parser.add_argument(
47 + "--W_dtype",
48 + type=str,
49 + default="int4",
50 + choices=[
51 + "float16",
52 + "float32",
53 + "float64",
54 + "int32",
55 + "int8",
56 + "int4",
57 + "int2",
58 + "int1",
59 + "nf4",
60 + "fp4_e2m1",
61 + ],
62 + help="Data type of weight W.",
63 +)
64 +parser.add_argument(
65 + "--accum_dtype",
66 + type=str,
67 + default="float16",
68 + choices=["float16", "int32"],
69 + help="Data type for accumulation.",
70 +)
71 +parser.add_argument(
72 + "--out_dtype",
73 + type=str,
74 + default="float16",
75 + choices=["float16", "float32", "int32", "int8"],
76 + help="Data type for output.",
77 +)
78 +parser.add_argument(
79 + "--layout",
80 + type=str,
81 + default="nt",
82 + choices=["nt", "nn"],
83 + help="Matrix layout, 'nt' for non-transpose A and transpose W.",
84 +)
85 +parser.add_argument("--with_bias",
86 + action="store_true",
87 + help="Include bias in the benchmark.")
88 +parser.add_argument(
89 + "--with_scaling",
90 + action="store_true",
91 + help="Include scaling factor in the quantization.",
92 +)
93 +parser.add_argument("--with_zeros",
94 + action="store_true",
95 + help="Include zeros in the quantization.")
96 +parser.add_argument(
97 + "--zeros_mode",
98 + type=str,
99 + default=None,
100 + choices=["original", "rescale", "quantized"],
101 + help="Specify the mode for calculating zeros.",
102 +)
103 +
104 +# Parse the arguments
105 +args = parser.parse_args()
106 +
107 +# Assign arguments to variables
108 +target = args.target
109 +A_dtype = args.A_dtype
110 +W_dtype = args.W_dtype
111 +accum_dtype = args.accum_dtype
112 +out_dtype = args.out_dtype
113 +layout = args.layout
114 +with_bias = args.with_bias
115 +group_size = args.group_size
116 +with_scaling = args.with_scaling
117 +with_zeros = args.with_zeros
118 +zeros_mode = args.zeros_mode
119 +
120 +# Define a list of shared arguments that repeat in every config
121 +shared_args = [
122 + A_dtype,
123 + W_dtype,
124 + out_dtype,
125 + accum_dtype,
126 + layout,
127 + with_bias,
128 + group_size,
129 + with_scaling,
130 + with_zeros,
131 + zeros_mode,
132 +]
133 +
134 +# Define just the (M, K, N) shapes in a more compact list
135 +shapes = [
136 + # square test
137 + (1, 16384, 16384),
138 + # BLOOM-176B
139 + (1, 43008, 14336),
140 + (1, 14336, 14336),
141 + (1, 57344, 14336),
142 + (1, 14336, 57344),
143 + # OPT-65B
144 + (1, 9216, 9216),
145 + (1, 36864, 9216),
146 + (1, 9216, 36864),
147 + (1, 22016, 8192),
148 + # LLAMA-70B/65B
149 + (1, 8192, 22016),
150 + (1, 8192, 8192),
151 + (1, 28672, 8192),
152 + (1, 8192, 28672),
153 + # square test
154 + (16384, 16384, 16384),
155 + # BLOOM-176B
156 + (8192, 43008, 14336),
157 + (8192, 14336, 14336),
158 + (8192, 57344, 14336),
159 + (8192, 14336, 57344),
160 + # OPT-65B
161 + (8192, 9216, 9216),
162 + (8192, 36864, 9216),
163 + (8192, 9216, 36864),
164 + (8192, 22016, 8192),
165 + # LLAMA-70B/65B
166 + (8192, 8192, 22016),
167 + (8192, 8192, 8192),
168 + (8192, 28672, 8192),
169 + (8192, 8192, 28672),
170 +]
171 +
172 +# Build test shapes with all the shared arguments
173 +test_shapes = [(MatmulConfig, Matmul, (*shape, *shared_args))
174 + for shape in shapes]
175 +
176 +benchmark_sets = []
177 +benchmark_sets.extend(test_shapes)
178 +
179 +benchmark_results = {}
180 +for config_class, operator, input_args in benchmark_sets:
181 + config = config_class(*input_args)
182 + matmul = operator(config, target=target, enable_tuning=True)
183 + kernel_latency = matmul.profile_latency()
184 +
185 + print("Time cost is: {:.3f} ms".format(kernel_latency))
186 +
187 + profile_config = {
188 + f"{operator.__name__}-{'-'.join([str(i) for i in input_args])}": {
189 + "BitBLAS_top20_latency": kernel_latency,
190 + }
191 + }
192 +
193 + benchmark_results.update(profile_config)
194 +
195 +# Define headers for the table
196 +headers = [
197 + "PrimFunc",
198 + "Input Arguments",
199 + "BitBLAS Top20 Latency",
200 +]
201 +
202 +# Calculate column widths for pretty printing
203 +col_widths = [0, 0, 0]
204 +for config_key, values in benchmark_results.items():
205 + args_split = config_key.split("-")
206 + func_name = args_split[0]
207 + input_args_str = "-".join(args_split[1:])
208 + col_widths[0] = max(col_widths[0], len(func_name) + 2, len(headers[0]) + 2)
209 + col_widths[1] = max(col_widths[1],
210 + len(input_args_str) + 2,
211 + len(headers[1]) + 2)
212 + col_widths[2] = max(col_widths[2],
213 + len(f"{values['BitBLAS_top20_latency']:.3f} ms") + 2,
214 + len(headers[2]) + 2)
215 + # break only if you want to measure widths from a single example;
216 + # otherwise, let it loop over all items.
217 +
218 +# Print header
219 +for i, header in enumerate(headers):
220 + headers[i] = header.ljust(col_widths[i])
221 +print("".join(headers))
222 +print("-" * sum(col_widths))
223 +
224 +# Print rows
225 +for config_key, values in benchmark_results.items():
226 + args_split = config_key.split("-")
227 + func_name = args_split[0]
228 + input_args_str = "-".join(args_split[1:])
229 + row = [
230 + func_name,
231 + input_args_str,
232 + f"{values['BitBLAS_top20_latency']:.3f} ms",
233 + ]
234 + row_str = "".join(
235 + [str(cell).ljust(col_widths[idx]) for idx, cell in enumerate(row)])
236 + print(row_str)
docs/source/features/quantization/bitblas.md new
+40
@@ -0,0 +1,40 @@
1 +# BitBLAS
2 +
3 +vLLM now supports [BitBLAS](https://github.com/microsoft/BitBLAS) for more efficient and flexible model inference. Compared to other quantization frameworks, BitBLAS provides more precision combinations.
4 +
5 +Below are the steps to utilize BitBLAS with vLLM.
6 +
7 +```console
8 +pip install bitblas>=0.1.0
9 +```
10 +
11 +vLLM reads the model's config file and supports pre-quantized checkpoints.
12 +
13 +You can find pre-quantized models on:
14 +
15 +- [Hugging Face (BitBLAS)](https://huggingface.co/models?other=bitblas)
16 +- [Hugging Face (GPTQ)](https://huggingface.co/models?other=gptq)
17 +
18 +Usually, these repositories have a `quantize_config.json` file that includes a `quantization_config` section.
19 +
20 +## Read bitblas format checkpoint
21 +
22 +```python
23 +from vllm import LLM
24 +import torch
25 +
26 +# "hxbgsyxh/llama-13b-4bit-g-1-bitblas" is a pre-quantized checkpoint.
27 +model_id = "hxbgsyxh/llama-13b-4bit-g-1-bitblas"
28 +llm = LLM(model=model_id, dtype=torch.bfloat16, trust_remote_code=True, quantization="bitblas")
29 +```
30 +
31 +## Read gptq format checkpoint
32 +
33 +```python
34 +from vllm import LLM
35 +import torch
36 +
37 +# "hxbgsyxh/llama-13b-4bit-g-1" is a pre-quantized checkpoint.
38 +model_id = "hxbgsyxh/llama-13b-4bit-g-1"
39 +llm = LLM(model=model_id, dtype=torch.float16, trust_remote_code=True, quantization="bitblas", max_model_len=1024)
40 +```
docs/source/features/quantization/index.md
+1
@@ -11,6 +11,7 @@ Quantization trades off model precision for smaller memory footprint, allowing l
11 supported_hardware
12 auto_awq
13 bnb
14 +bitblas
15 gguf
16 gptqmodel
17 int4
docs/source/features/quantization/supported_hardware.md
+11
@@ -74,6 +74,17 @@ The table below shows the compatibility of various quantization implementations
74 * ❌
75 * ❌
76 * ❌
77 +- * BitBLAS (GPTQ)
78 + * ✅︎
79 + * ✅︎
80 + * ✅︎
81 + * ✅︎
82 + * ✅︎
83 + * ✅︎
84 + * ❌
85 + * ❌
86 + * ❌
87 + * ❌
88 - * AQLM
89 * ✅︎
90 * ✅︎
tests/models/test_bitblas.py new
+63
@@ -0,0 +1,63 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Compare the outputs of a GPTQ model to a bitblas model.
3 +
4 +Note: GPTQ and bitblas do not have bitwise correctness.
5 +As a result, in this test, we just confirm that the top selected tokens of the
6 +bitblas/GPTQ models are in the top 3 selections of each other.
7 +
8 +Note: bitblas internally uses locks to synchronize the threads. This can
9 +result in very slight nondeterminism for bitblas. As a result, we re-run the
10 +test up to 3 times to see if we pass.
11 +
12 +Run `pytest tests/models/test_bitblas.py`.
13 +"""
14 +from dataclasses import dataclass
15 +
16 +import pytest
17 +
18 +from .utils import check_logprobs_close
19 +
20 +
21 +@dataclass
22 +class ModelPair:
23 + model_bitblas: str
24 + model_gptq: str
25 +
26 +
27 +model_pairs = [
28 + ModelPair(model_bitblas="hxbgsyxh/opt-125m-4bit-128g-bitblas",
29 + model_gptq="hxbgsyxh/opt-125m-4bit-128g"),
30 +]
31 +
32 +
33 +@pytest.mark.flaky(reruns=2)
34 +@pytest.mark.skipif(True, reason="BitBLAS takes too much time for tuning.")
35 +@pytest.mark.parametrize("model_pair", model_pairs)
36 +@pytest.mark.parametrize("dtype", ["half"])
37 +@pytest.mark.parametrize("max_tokens", [32])
38 +@pytest.mark.parametrize("num_logprobs", [5])
39 +def test_models(
40 + vllm_runner,
41 + example_prompts,
42 + model_pair: ModelPair,
43 + dtype: str,
44 + max_tokens: int,
45 + num_logprobs: int,
46 +) -> None:
47 + with vllm_runner(model_pair.model_bitblas,
48 + dtype=dtype,
49 + quantization="bitblas") as bitblas_model:
50 + bitblas_outputs = bitblas_model.generate_greedy_logprobs(
51 + example_prompts, max_tokens, num_logprobs)
52 +
53 + with vllm_runner(model_pair.model_gptq, dtype=dtype,
54 + quantization="gptq") as gptq_model:
55 + gptq_outputs = gptq_model.generate_greedy_logprobs(
56 + example_prompts, max_tokens, num_logprobs)
57 +
58 + check_logprobs_close(
59 + outputs_0_lst=gptq_outputs,
60 + outputs_1_lst=bitblas_outputs,
61 + name_0="gptq",
62 + name_1="bitblas",
63 + )
tests/models/test_gptq_bitblas.py new
+61
@@ -0,0 +1,61 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Compare the outputs of a GPTQ model to a bitblas model.
3 +
4 +Note: GPTQ and bitblas do not have bitwise correctness.
5 +As a result, in this test, we just confirm that the top selected tokens of the
6 +bitblas/GPTQ models are in the top 3 selections of each other.
7 +
8 +Note: bitblas internally uses locks to synchronize the threads. This can
9 +result in very slight nondeterminism for bitblas. As a result, we re-run the
10 +test up to 3 times to see if we pass.
11 +
12 +Run `pytest tests/models/test_bitblas.py`.
13 +"""
14 +from dataclasses import dataclass
15 +
16 +import pytest
17 +
18 +from .utils import check_logprobs_close
19 +
20 +
21 +@dataclass
22 +class ModelPair:
23 + model_gptq: str
24 +
25 +
26 +model_pairs = [
27 + ModelPair(model_gptq="hxbgsyxh/opt-125m-4bit-128g"),
28 +]
29 +
30 +
31 +@pytest.mark.flaky(reruns=2)
32 +@pytest.mark.skipif(True, reason="BitBLAS takes too much time for tuning.")
33 +@pytest.mark.parametrize("model_pair", model_pairs)
34 +@pytest.mark.parametrize("dtype", ["half"])
35 +@pytest.mark.parametrize("max_tokens", [32])
36 +@pytest.mark.parametrize("num_logprobs", [5])
37 +def test_models(
38 + vllm_runner,
39 + example_prompts,
40 + model_pair: ModelPair,
41 + dtype: str,
42 + max_tokens: int,
43 + num_logprobs: int,
44 +) -> None:
45 + with vllm_runner(model_pair.model_gptq,
46 + dtype=dtype,
47 + quantization="bitblas") as bitblas_model:
48 + bitblas_outputs = bitblas_model.generate_greedy_logprobs(
49 + example_prompts, max_tokens, num_logprobs)
50 +
51 + with vllm_runner(model_pair.model_gptq, dtype=dtype,
52 + quantization="gptq") as gptq_model:
53 + gptq_outputs = gptq_model.generate_greedy_logprobs(
54 + example_prompts, max_tokens, num_logprobs)
55 +
56 + check_logprobs_close(
57 + outputs_0_lst=gptq_outputs,
58 + outputs_1_lst=bitblas_outputs,
59 + name_0="gptq",
60 + name_1="gptq_bitblas",
61 + )
vllm/config.py
+2 -1
@@ -750,7 +750,8 @@ class ModelConfig:
750 optimized_quantization_methods = [
751 "fp8", "marlin", "modelopt", "gptq_marlin_24", "gptq_marlin",
752 "awq_marlin", "fbgemm_fp8", "compressed_tensors",
753 - "compressed-tensors", "experts_int8", "quark", "nvfp4"
753 + "compressed-tensors", "experts_int8", "quark", "nvfp4", "bitblas",
754 + "gptq_bitblas"
755 ]
756 if self.quantization is not None:
757 self.quantization = self.quantization.lower()
vllm/model_executor/layers/linear.py
+16
@@ -31,6 +31,8 @@ logger = init_logger(__name__)
31
32 WEIGHT_LOADER_V2_SUPPORTED = [
33 "CompressedTensorsLinearMethod",
34 + "BitBLASLinearMethod",
35 + "GPTQBitBLASLinearMethod",
36 "AWQMarlinLinearMethod",
37 "AWQLinearMethod",
38 "GPTQMarlinLinearMethod",
@@ -50,6 +52,15 @@ WEIGHT_LOADER_V2_SUPPORTED = [
52 ]
53
54
55 +def adjust_bitblas_shard(param, shard_size, shard_offset):
56 + bitblas_tile_size = getattr(param, "bitblas_tile_size", None)
57 + if bitblas_tile_size is not None:
58 + return (shard_size // bitblas_tile_size,
59 + shard_offset // bitblas_tile_size)
60 +
61 + return shard_size, shard_offset
62 +
63 +
64 def adjust_marlin_shard(param, shard_size, shard_offset):
65 marlin_tile_size = getattr(param, "marlin_tile_size", None)
66 if marlin_tile_size is None:
@@ -615,6 +626,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
626 shard_size, shard_offset = adjust_marlin_shard(
627 param, shard_size, shard_offset)
628
629 + shard_size, shard_offset = adjust_bitblas_shard(
630 + param, shard_size, shard_offset)
631 +
632 if use_bitsandbytes_4bit:
633 index = list(itertools.accumulate([0] + self.output_sizes))
634 orig_offsets = {
@@ -646,6 +660,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
660 # Special case for Marlin.
661 shard_size, shard_offset = adjust_marlin_shard(
662 param, shard_size, shard_offset)
663 + shard_size, shard_offset = adjust_bitblas_shard(
664 + param, shard_size, shard_offset)
665
666 use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit",
667 False)
vllm/model_executor/layers/quantization/__init__.py
+7 -1
@@ -18,9 +18,11 @@ QUANTIZATION_METHODS: List[str] = [
18 # The order of gptq methods is important for config.py iteration over
19 # override_quantization_method(..)
20 "marlin",
21 + "bitblas",
22 "gguf",
23 "gptq_marlin_24",
24 "gptq_marlin",
25 + "gptq_bitblas",
26 "awq_marlin",
27 "gptq",
28 "compressed-tensors",
@@ -85,6 +87,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
87 from .aqlm import AQLMConfig
88 from .awq import AWQConfig
89 from .awq_marlin import AWQMarlinConfig
90 + from .bitblas import BitBLASConfig
91 from .bitsandbytes import BitsAndBytesConfig
92 from .compressed_tensors.compressed_tensors import ( # noqa: E501
93 CompressedTensorsConfig)
@@ -94,6 +97,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
97 from .fp8 import Fp8Config
98 from .gguf import GGUFConfig
99 from .gptq import GPTQConfig
100 + from .gptq_bitblas import GPTQBitBLASConfig
101 from .gptq_marlin import GPTQMarlinConfig
102 from .gptq_marlin_24 import GPTQMarlin24Config
103 from .hqq_marlin import HQQMarlinConfig
@@ -119,9 +123,11 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
123 # The order of gptq methods is important for config.py iteration over
124 # override_quantization_method(..)
125 "marlin": MarlinConfig,
126 + "bitblas": BitBLASConfig,
127 "gguf": GGUFConfig,
128 "gptq_marlin_24": GPTQMarlin24Config,
129 "gptq_marlin": GPTQMarlinConfig,
130 + "gptq_bitblas": GPTQBitBLASConfig,
131 "awq_marlin": AWQMarlinConfig,
132 "gptq": GPTQConfig,
133 "compressed-tensors": CompressedTensorsConfig,
@@ -146,4 +152,4 @@ __all__ = [
152 "QuantizationConfig",
153 "get_quantization_config",
154 "QUANTIZATION_METHODS",
149 -]
155 +]
\ No newline at end of file
vllm/model_executor/layers/quantization/bitblas.py new
+459
@@ -0,0 +1,459 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Any, Dict, List, Optional
3 +
4 +import torch
5 +
6 +from vllm.logger import init_logger
7 +from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
8 +from vllm.model_executor.layers.quantization.base_config import (
9 + QuantizationConfig)
10 +from vllm.model_executor.layers.quantization.utils.bitblas_utils import (
11 + BITBLAS_OPTIMIZE_FEATURES, BITBLAS_SUPPORTED_NUM_BITS,
12 + BITBLAS_SUPPORTED_SYM, MINIMUM_BITBLAS_VERSION)
13 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
14 +from vllm.model_executor.parameter import (BasevLLMParameter,
15 + ChannelQuantScaleParameter,
16 + GroupQuantScaleParameter,
17 + PackedvLLMParameter)
18 +from vllm.model_executor.utils import set_weight_attrs
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +class BitBLASConfig(QuantizationConfig):
24 + """Config class for BitBLAS.
25 +
26 + Reference: https://github.com/Microsoft/BitBLAS
27 + """
28 + TORCH_DTYPE = torch.float16
29 + STORAGE_DTYPE = "int8" # assume int8 storage
30 + TORCH_STORAGE_DTYPE = getattr(torch, STORAGE_DTYPE)
31 + # "original" or "rescale" or "quantized",
32 + # gptq_with_bitblas prefer "quantized implementation"
33 + ZEROS_MODE = "quantized"
34 +
35 + def __init__(
36 + self,
37 + weight_bits: int,
38 + group_size: Optional[int],
39 + desc_act: Optional[bool],
40 + is_sym: Optional[bool],
41 + quant_method: Optional[str],
42 + lm_head_quantized: bool,
43 + ) -> None:
44 + try:
45 + import bitblas
46 + if bitblas.__version__ < MINIMUM_BITBLAS_VERSION:
47 + raise ImportError(
48 + "bitblas version is wrong. Please "
49 + f"install bitblas>={MINIMUM_BITBLAS_VERSION}")
50 + except ImportError as e:
51 + bitblas_import_exception = e
52 + raise ValueError(
53 + "Trying to use the bitblas backend, but could not import"
54 + f"with the following error: {bitblas_import_exception}. "
55 + "Please install bitblas through the following command: "
56 + f"`pip install bitblas>={MINIMUM_BITBLAS_VERSION}`"
57 + ) from bitblas_import_exception
58 +
59 + if desc_act and group_size == -1:
60 + # In this case, act_order == True is the same as act_order == False
61 + # (since we have only one group per output channel)
62 + desc_act = False
63 +
64 + self.weight_bits = weight_bits
65 + self.group_size = group_size
66 + self.desc_act = desc_act
67 + self.is_sym = is_sym
68 + self.quant_method = quant_method
69 + self.lm_head_quantized = lm_head_quantized
70 +
71 + # Verify
72 + if self.weight_bits not in BITBLAS_SUPPORTED_NUM_BITS:
73 + raise ValueError(
74 + f"BitBLAS does not support weight_bits = {self.weight_bits}. "
75 + f"Only weight_bits = {BITBLAS_SUPPORTED_NUM_BITS} "
76 + "are supported.")
77 +
78 + if self.is_sym not in BITBLAS_SUPPORTED_SYM:
79 + raise ValueError(
80 + f"BitBLAS does not support is_sym = {self.is_sym}. "
81 + f"Only sym = {BITBLAS_SUPPORTED_SYM} are supported.")
82 +
83 + storage_dtype = self.STORAGE_DTYPE
84 + storage_nbit = int("".join(c for c in storage_dtype if c.isdigit()))
85 +
86 + self.storage_dtype = storage_dtype
87 + self.storage_torch_dtype = self.TORCH_STORAGE_DTYPE
88 + # 4 Bits packed into 32 bit datatype.
89 + self.pack_factor = storage_nbit // weight_bits
90 + self.nbits = weight_bits
91 +
92 + # Zeros type for the quantized weights.
93 + self.zeros_mode = self.ZEROS_MODE
94 +
95 + def __repr__(self) -> str:
96 + return (f"BitBLASConfig(weight_bits={self.weight_bits}, "
97 + f"group_size={self.group_size}, "
98 + f"desc_act={self.desc_act}, "
99 + f"is_sym={self.is_sym}, "
100 + f"quant_method={self.quant_method})")
101 +
102 + @classmethod
103 + def get_name(cls) -> str:
104 + return "bitblas"
105 +
106 + @classmethod
107 + def get_supported_act_dtypes(cls) -> List[torch.dtype]:
108 + return [torch.half, torch.bfloat16]
109 +
110 + @classmethod
111 + # Need to figure it out
112 + def get_min_capability(cls) -> int:
113 + return 70
114 +
115 + @classmethod
116 + def get_config_filenames(cls) -> List[str]:
117 + return ["quantize_config.json"]
118 +
119 + @staticmethod
120 + def get_from_keys(config: Dict[str, Any],
121 + keys: List[str],
122 + default: Any = None) -> Any:
123 + """Get a value from the model's quantization config."""
124 + for key in keys:
125 + if key in config:
126 + return config[key]
127 + return default
128 +
129 + @classmethod
130 + def from_config(cls, config: Dict[str, Any]) -> "BitBLASConfig":
131 + weight_bits = cls.get_from_keys(config, ["bits"])
132 + group_size = cls.get_from_keys(config, ["group_size"], -1)
133 + desc_act = cls.get_from_keys(config, ["desc_act"], False)
134 + is_sym = cls.get_from_keys(config, ["sym"], False)
135 + quant_method = cls.get_from_keys(config, ["quant_method"])
136 + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
137 + default=False)
138 + return cls(weight_bits, group_size, desc_act, is_sym, quant_method,
139 + lm_head_quantized)
140 +
141 + @classmethod
142 + def override_quantization_method(cls, hf_quant_cfg,
143 + user_quant) -> Optional[str]:
144 + # compat: autogptq >=0.8.0 use checkpoint_format: str
145 + # compat: autogptq <=0.7.1 is_bitblas_format: bool
146 + is_bitblas_format = (hf_quant_cfg.get("checkpoint_format") == "bitblas"
147 + or hf_quant_cfg.get("is_bitblas_format", False))
148 +
149 + is_valid_user_quant = (user_quant is None or user_quant == "gptq"
150 + or user_quant == "bitblas")
151 +
152 + if is_bitblas_format and is_valid_user_quant:
153 + msg = ("The model is serialized in {} format. Using {} kernel.".
154 + format(cls.get_name(), cls.get_name()))
155 + logger.info(msg)
156 + return cls.get_name()
157 +
158 + return None
159 +
160 + def get_quant_method(self, layer: torch.nn.Module,
161 + prefix: str) -> Optional["BitBLASLinearMethod"]:
162 + if isinstance(layer, LinearBase) or (isinstance(layer, ParallelLMHead)
163 + and self.lm_head_quantized):
164 + return BitBLASLinearMethod(self)
165 + return None
166 +
167 +
168 +class BitBLASLinearMethod(LinearMethodBase):
169 + """Linear method for BitBLAS.
170 +
171 + Args:
172 + quant_config: The BitBLAS quantization config.
173 + """
174 + # USE BITBLAS_OPTIMIZE_FEATURES_CONTIGUOUS
175 + # Instead of BITBLAS_OPTIMIZE_FEATURES
176 + # If you want to high contiguous batching
177 + # performance
178 + OPT_FEATURES = BITBLAS_OPTIMIZE_FEATURES
179 + ENABLE_TUNING = True
180 + BITBLAS_DTYPES = {
181 + torch.float32: "float32",
182 + torch.float16: "float16",
183 + torch.bfloat16: "bfloat16",
184 + torch.half: "float16",
185 + torch.int8: "int8",
186 + }
187 +
188 + def __init__(self, quant_config: BitBLASConfig):
189 + self.quant_config = quant_config
190 +
191 + def create_weights_gptq(
192 + self,
193 + layer: torch.nn.Module,
194 + input_size_per_partition: int,
195 + output_partition_sizes: List[int],
196 + input_size: int,
197 + output_size: int,
198 + params_dtype: torch.dtype,
199 + **extra_weight_attrs,
200 + ):
201 + """Creates quantized weights for use in linear operations.
202 +
203 + The function initializes and returns a dictionary containing quantized
204 + weights, scales, and zeros
205 + for performing quantized matrix multiplication operations.
206 +
207 + Args:
208 + input_size_per_partition: The size of the input partition.
209 + output_size_per_partition: The size of the output partition.
210 + input_size: The total size of the input (unused).
211 + output_size: The total size of the output (unused).
212 + params_dtype:
213 + The data type of the parameters (expected to be torch.float16).
214 +
215 + Returns:
216 + A dictionary containing the quantized weights ('qweight'),
217 + scales ('scales'), and zeros ('zeros').
218 +
219 + Raises:
220 + ValueError: If `params_dtype` is not `torch.float16` or if the
221 + input size per partition is not divisible by the group size in
222 + `quant_config`.
223 + """
224 + del input_size, output_size # Unused arguments.
225 + weight_loader = extra_weight_attrs["weight_loader"]
226 +
227 + if params_dtype not in self.quant_config.get_supported_act_dtypes():
228 + raise ValueError("Parameter data type must be torch.float16, "
229 + f"but got {params_dtype}")
230 + group_size = self.quant_config.group_size
231 + if group_size is None:
232 + group_size = -1
233 + # Validate output_size_per_partition
234 + output_size_per_partition = sum(output_partition_sizes)
235 + if (group_size != -1 and input_size_per_partition % group_size != 0):
236 + raise ValueError(
237 + f"Input size per partition ({input_size_per_partition}) must "
238 + f"be divisible by group size ({group_size}).")
239 +
240 + # Initialize or retrieve the BitBLAS matrix multiplication operator.
241 + self._configure_bitblas_matmul(
242 + input_size_per_partition,
243 + output_size_per_partition,
244 + params_dtype=params_dtype,
245 + enable_tuning=self.ENABLE_TUNING,
246 + bias=False,
247 + layout="nt",
248 + bits=self.quant_config.weight_bits,
249 + )
250 +
251 + # Initialize quantized weights with dimensions
252 + # Quantized 4Bit weights packed.
253 + qweight = PackedvLLMParameter(
254 + data=torch.empty(
255 + self.bitblas_matmul.retrieve_weight_shape(),
256 + device="cuda",
257 + dtype=self.quant_config.storage_torch_dtype,
258 + requires_grad=False,
259 + ),
260 + input_dim=1,
261 + output_dim=0,
262 + packed_dim=1,
263 + packed_factor=self.quant_config.pack_factor,
264 + bitblas_tile_size=(self.bitblas_matmul.retrieve_weight_shape()[-2]
265 + if self.bitblas_matmul.propagate_b else None),
266 + weight_loader=weight_loader,
267 + )
268 +
269 + # Compute the number of input groups for channel-wise quantization.
270 + input_groups = (1 if group_size == -1 else input_size_per_partition //
271 + group_size)
272 +
273 + # Initialize scales and zeros for the quantized weights.
274 + weight_scale_args = {
275 + "data":
276 + torch.empty(
277 + output_size_per_partition,
278 + input_groups,
279 + device="cuda",
280 + dtype=params_dtype,
281 + ),
282 + "weight_loader":
283 + weight_loader
284 + }
285 + if input_groups == 1:
286 + scales = ChannelQuantScaleParameter(output_dim=0,
287 + **weight_scale_args)
288 + else:
289 + scales = GroupQuantScaleParameter(output_dim=0,
290 + input_dim=1,
291 + **weight_scale_args)
292 +
293 + if self.quant_config.zeros_mode == "quantized":
294 + zeros = PackedvLLMParameter(
295 + data=torch.empty(
296 + input_groups,
297 + output_size_per_partition // self.quant_config.pack_factor,
298 + device="cuda",
299 + dtype=self.quant_config.storage_torch_dtype,
300 + requires_grad=False,
301 + ),
302 + input_dim=0,
303 + output_dim=1,
304 + packed_dim=1,
305 + packed_factor=self.quant_config.pack_factor,
306 + weight_loader=weight_loader,
307 + )
308 +
309 + else:
310 + zeros = BasevLLMParameter(
311 + torch.empty(output_size_per_partition,
312 + input_groups,
313 + device="cuda",
314 + dtype=params_dtype),
315 + weight_loader=weight_loader,
316 + )
317 + # Set attributes to indicate how scales and zeros are applied.
318 + set_weight_attrs(zeros, {
319 + "input_dim": None if input_groups == 1 else 1,
320 + "output_dim": 0,
321 + })
322 +
323 + layer.register_parameter("qweight", qweight)
324 + layer.register_parameter("scales", scales)
325 + layer.register_parameter("zeros", zeros)
326 +
327 + def create_weights(
328 + self,
329 + layer: torch.nn.Module,
330 + input_size_per_partition: int,
331 + output_partition_sizes: List[int],
332 + input_size: int,
333 + output_size: int,
334 + params_dtype: torch.dtype,
335 + **extra_weight_attrs,
336 + ):
337 + if self.quant_config.quant_method == "gptq":
338 + return self.create_weights_gptq(layer, input_size_per_partition,
339 + output_partition_sizes, input_size,
340 + output_size, params_dtype,
341 + **extra_weight_attrs)
342 + else:
343 + raise ValueError(
344 + f"Unsupported quant_method {self.quant_config.quant_method}")
345 +
346 + def _configure_bitblas_matmul(
347 + self,
348 + infeatures,
349 + outfeatures,
350 + params_dtype,
351 + enable_tuning,
352 + bias,
353 + layout,
354 + bits,
355 + out_dtype="float16",
356 + ):
357 + from bitblas import MatmulConfig
358 + bitblas_dtype = self.BITBLAS_DTYPES[params_dtype]
359 +
360 + with_scaling = False
361 + with_zeros = False
362 + group_size = self.quant_config.group_size
363 + zeros_mode = self.quant_config.zeros_mode
364 + if self.quant_config.quant_method == "gptq":
365 + with_scaling = True
366 + with_zeros = True
367 + W_dtype = f"uint{bits}"
368 + if self.quant_config.is_sym:
369 + with_zeros = False
370 + W_dtype = f"int{bits}"
371 + else:
372 + raise ValueError(
373 + f"Unsupported quant_method {self.quant_config.quant_method}")
374 +
375 + matmul_config = MatmulConfig(
376 + N=outfeatures,
377 + K=infeatures,
378 + A_dtype=bitblas_dtype,
379 + W_dtype=W_dtype,
380 + out_dtype=out_dtype,
381 + accum_dtype="int32" if bitblas_dtype == "int8" else bitblas_dtype,
382 + storage_dtype=self.quant_config.STORAGE_DTYPE,
383 + with_scaling=with_scaling,
384 + with_zeros=with_zeros,
385 + group_size=group_size,
386 + with_bias=bias,
387 + layout=layout,
388 + zeros_mode=zeros_mode,
389 + )
390 + self.bitblas_matmul = self._get_or_create_bitblas_operator(
391 + matmul_config, enable_tuning)
392 +
393 + def _get_or_create_bitblas_operator(self, config, enable_tuning):
394 + from bitblas import Matmul, auto_detect_nvidia_target
395 + from bitblas.cache import get_database_path, global_operator_cache
396 + BITBLAS_DATABASE_PATH = get_database_path()
397 + BITBLAS_TARGET = auto_detect_nvidia_target()
398 + if global_operator_cache.size() == 0:
399 + global_operator_cache.load_from_database(BITBLAS_DATABASE_PATH,
400 + BITBLAS_TARGET)
401 +
402 + bitblas_matmul = global_operator_cache.get(config)
403 + if bitblas_matmul is None:
404 + bitblas_matmul = Matmul(config,
405 + target=BITBLAS_TARGET,
406 + enable_tuning=False)
407 + if enable_tuning:
408 + TUNING_MESSAGE = (f"BitBLAS Operator {config} is tuning ...")
409 + logger.info(TUNING_MESSAGE)
410 + bitblas_matmul.hardware_aware_finetune(topk=20)
411 + global_operator_cache.add(config, bitblas_matmul)
412 + global_operator_cache.save_into_database(
413 + BITBLAS_DATABASE_PATH, BITBLAS_TARGET)
414 + TUNED_MESSAGE = (
415 + f"BitBLAS Operator {config} tuned and saved to database.")
416 + logger.info(TUNED_MESSAGE)
417 + else:
418 + _message = f"BitBLAS Operator {config} created."
419 + logger.info(_message)
420 + else:
421 + _message = (
422 + f"BitBLAS Operator {config} found in global_operator_cache.")
423 + logger.info(_message)
424 + return bitblas_matmul
425 +
426 + def apply_gptq(
427 + self,
428 + layer: torch.nn.Module,
429 + x: torch.Tensor,
430 + bias: Optional[torch.Tensor] = None,
431 + ) -> torch.Tensor:
432 + qweight = layer.qweight
433 + scales = layer.scales
434 + qzeros = layer.zeros
435 +
436 + x_2d = x.view(-1, x.shape[-1])
437 +
438 + if self.quant_config.is_sym:
439 + output_2d = self.bitblas_matmul(x_2d, qweight, scales)
440 + else:
441 + output_2d = self.bitblas_matmul(x_2d, qweight, scales, qzeros)
442 +
443 + output = output_2d.view(x.shape[:-1] + (output_2d.shape[1], ))
444 +
445 + if bias is not None:
446 + output.add_(bias) # In-place add
447 +
448 + return output
449 +
450 + def apply(
451 + self,
452 + *args: Any,
453 + **kwargs: Any,
454 + ) -> torch.Tensor:
455 + if self.quant_config.quant_method == "gptq":
456 + return self.apply_gptq(*args, **kwargs)
457 + else:
458 + raise ValueError(
459 + f"Unsupported quant_method {self.quant_config.quant_method}")
vllm/model_executor/layers/quantization/gptq_bitblas.py new
+438
@@ -0,0 +1,438 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Any, Dict, List, Optional, Set
3 +
4 +import torch
5 +from torch.nn.parameter import Parameter
6 +
7 +from vllm.logger import init_logger
8 +from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
9 + set_weight_attrs)
10 +from vllm.model_executor.layers.quantization.base_config import (
11 + QuantizationConfig)
12 +from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
13 + BitBLASLinearKernel, MPLinearLayerConfig)
14 +from vllm.model_executor.layers.quantization.utils.bitblas_utils import (
15 + BITBLAS_SUPPORTED_NUM_BITS as GPTQ_BITBLAS_SUPPORTED_NUM_BITS)
16 +from vllm.model_executor.layers.quantization.utils.bitblas_utils import (
17 + BITBLAS_SUPPORTED_SYM as GPTQ_BITBLAS_SUPPORTED_SYM)
18 +from vllm.model_executor.layers.quantization.utils.bitblas_utils import (
19 + MINIMUM_BITBLAS_VERSION, bitblas_repeat_scales_on_all_ranks,
20 + check_bitblas_supported, verify_bitblas_supported)
21 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
22 +from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
23 + GroupQuantScaleParameter,
24 + PackedColumnParameter,
25 + PackedvLLMParameter,
26 + RowvLLMParameter)
27 +from vllm.scalar_type import scalar_types
28 +
29 +logger = init_logger(__name__)
30 +
31 +
32 +class GPTQBitBLASConfig(QuantizationConfig):
33 + """Config class for GPTQ BitBLAS"""
34 +
35 + # (num_bits, is_sym) -> quant_type
36 + TYPE_MAP = {
37 + (4, True): scalar_types.uint4b8,
38 + (8, True): scalar_types.uint8b128,
39 + }
40 +
41 + TORCH_DTYPE = torch.float16
42 + GPTQ_CKPT_STORAGE_DTYPE = (
43 + "int32" # GPTQ Default Checkpoints use int32 as storage dtype
44 + )
45 + GPTQ_BITBLAS_STORAGE_DTYPE = "int8" # BitBLAS uses int8 as storage dtype
46 + TORCH_BITBLAS_STORAGE_DTYPE = getattr(torch, GPTQ_BITBLAS_STORAGE_DTYPE)
47 + # "original" or "rescale" or "quantized",
48 + # the gptq_bitblas prefer "quantized"
49 + ZEROS_MODE = "quantized"
50 +
51 + def __init__(
52 + self,
53 + weight_bits: int,
54 + group_size: int,
55 + desc_act: bool,
56 + is_sym: bool,
57 + quant_method: Optional[str],
58 + lm_head_quantized: bool,
59 + ) -> None:
60 +
61 + try:
62 + import bitblas
63 + if bitblas.__version__ < MINIMUM_BITBLAS_VERSION:
64 + raise ImportError(
65 + "bitblas version is wrong. Please "
66 + f"install bitblas>={MINIMUM_BITBLAS_VERSION}")
67 + except ImportError as e:
68 + bitblas_import_exception = e
69 + raise ValueError(
70 + "Trying to use the bitblas backend, but could not import"
71 + f"with the following error: {bitblas_import_exception}. "
72 + "Please install bitblas through the following command: "
73 + f"`pip install bitblas>={MINIMUM_BITBLAS_VERSION}`"
74 + ) from bitblas_import_exception
75 +
76 + if desc_act and group_size == -1:
77 + # In this case, act_order == True is the same as act_order == False
78 + # (since we have only one group per output channel)
79 + desc_act = False
80 +
81 + self.weight_bits = weight_bits
82 + self.group_size = group_size
83 + self.desc_act = desc_act
84 + self.is_sym = is_sym
85 + self.quant_method = quant_method
86 + self.lm_head_quantized = lm_head_quantized
87 +
88 + # Verify
89 + if self.weight_bits not in GPTQ_BITBLAS_SUPPORTED_NUM_BITS:
90 + raise ValueError(
91 + f"BitBLAS does not support weight_bits = {self.weight_bits}. "
92 + f"Only weight_bits = {GPTQ_BITBLAS_SUPPORTED_NUM_BITS} "
93 + "are supported.")
94 +
95 + if self.is_sym not in GPTQ_BITBLAS_SUPPORTED_SYM:
96 + raise ValueError(
97 + f"BitBLAS does not support is_sym = {self.is_sym}. "
98 + f"Only sym = {GPTQ_BITBLAS_SUPPORTED_SYM} are supported.")
99 +
100 + self.storage_dtype = self.GPTQ_BITBLAS_STORAGE_DTYPE
101 +
102 + storage_nbit = int("".join(c for c in self.GPTQ_CKPT_STORAGE_DTYPE
103 + if c.isdigit()))
104 +
105 + # 4 Bits packed into 32 bit datatype.
106 + self.pack_factor = storage_nbit // weight_bits
107 + self.nbits = weight_bits
108 +
109 + # Zeros type for the quantized weights.
110 + self.zeros_mode = self.ZEROS_MODE
111 +
112 + if (weight_bits, is_sym) not in self.TYPE_MAP:
113 + raise ValueError("Unsupported quantization config: "
114 + f"bits={weight_bits}, sym={is_sym}")
115 +
116 + self.quant_type = self.TYPE_MAP[(weight_bits, is_sym)]
117 +
118 + def __repr__(self) -> str:
119 + return (f"GPTQBitBLASConfig(weight_bits={self.weight_bits}, "
120 + f"group_size={self.group_size}, "
121 + f"desc_act={self.desc_act})"
122 + f"is_sym={self.is_sym}, "
123 + f"quant_method={self.quant_method})")
124 +
125 + @classmethod
126 + def get_name(cls) -> str:
127 + return "gptq_bitblas"
128 +
129 + @classmethod
130 + def get_supported_act_dtypes(cls) -> List[torch.dtype]:
131 + return [torch.half, torch.bfloat16]
132 +
133 + @classmethod
134 + def get_min_capability(cls) -> int:
135 + return 70
136 +
137 + @classmethod
138 + def get_config_filenames(cls) -> List[str]:
139 + return ["quantize_config.json"]
140 +
141 + @classmethod
142 + def from_config(cls, config: Dict[str, Any]) -> "GPTQBitBLASConfig":
143 + weight_bits = cls.get_from_keys(config, ["bits"])
144 + group_size = cls.get_from_keys(config, ["group_size"])
145 + desc_act = cls.get_from_keys(config, ["desc_act"])
146 + is_sym = cls.get_from_keys(config, ["sym"])
147 + quant_method = cls.get_from_keys(config, ["quant_method"])
148 + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
149 + default=False)
150 + return cls(weight_bits, group_size, desc_act, is_sym, quant_method,
151 + lm_head_quantized)
152 +
153 + @classmethod
154 + def override_quantization_method(cls, hf_quant_cfg,
155 + user_quant) -> Optional[str]:
156 + can_convert = cls.is_gptq_bitblas_compatible(hf_quant_cfg)
157 +
158 + is_valid_user_quant = (user_quant is None or user_quant == "bitblas"
159 + or user_quant == "gptq_bitblas")
160 +
161 + if can_convert and is_valid_user_quant:
162 + msg = ("The model is convertible to {} during runtime."
163 + " Using {} kernel.".format(cls.get_name(), cls.get_name()))
164 + logger.info(msg)
165 + return cls.get_name()
166 +
167 + if can_convert and user_quant == "gptq":
168 + logger.info("Detected that the model can run with gptq_bitblas"
169 + ", however you specified quantization=gptq explicitly,"
170 + " so forcing gptq. Use quantization=gptq_bitblas for"
171 + " faster inference")
172 + return None
173 +
174 + def get_quant_method(self, layer: torch.nn.Module,
175 + prefix: str) -> Optional["GPTQBitBLASLinearMethod"]:
176 + if isinstance(layer, LinearBase) or (isinstance(layer, ParallelLMHead)
177 + and self.lm_head_quantized):
178 + return GPTQBitBLASLinearMethod(self)
179 + return None
180 +
181 + @property
182 + def torch_storage_dtype(self) -> torch.dtype:
183 + return self.TORCH_BITBLAS_STORAGE_DTYPE
184 +
185 + @classmethod
186 + def is_gptq_bitblas_compatible(cls, quant_config: Dict[str, Any]):
187 + # Extract data from quant config.
188 + num_bits = quant_config.get("bits")
189 + group_size = quant_config.get("group_size")
190 + sym = quant_config.get("sym")
191 + desc_act = quant_config.get("desc_act")
192 +
193 + # If we cannot find the info needed in the config, cannot convert.
194 + if (num_bits is None or group_size is None or sym is None
195 + or desc_act is None):
196 + return False
197 +
198 + if (num_bits, sym) not in cls.TYPE_MAP:
199 + return False
200 +
201 + # If the capability of the device is too low, cannot convert.
202 + major, minor = torch.cuda.get_device_capability()
203 + device_capability = major * 10 + minor
204 + if device_capability < cls.get_min_capability():
205 + return False
206 +
207 + # Otherwise, can convert if model satisfies bitblas constraints.
208 + return check_bitblas_supported(quant_type=cls.TYPE_MAP[(num_bits,
209 + sym)],
210 + group_size=group_size)
211 +
212 +
213 +class GPTQBitBLASLinearMethod(LinearMethodBase):
214 + """Linear method for GPTQ BitBLAS.
215 +
216 + Args:
217 + quant_config: The GPTQ BitBLAS quantization config.
218 + """
219 +
220 + kernel_type = BitBLASLinearKernel
221 + _kernel_backends_being_used: Set[str] = set()
222 +
223 + def __init__(self, quant_config: GPTQBitBLASConfig) -> None:
224 + self.quant_config = quant_config
225 + # Verify supported on platform.
226 + verify_bitblas_supported(quant_type=self.quant_config.quant_type,
227 + group_size=self.quant_config.group_size)
228 +
229 + def create_weights(
230 + self,
231 + layer: torch.nn.Module,
232 + input_size_per_partition: int,
233 + output_partition_sizes: List[int],
234 + input_size: int,
235 + output_size: int,
236 + params_dtype: torch.dtype,
237 + **extra_weight_attrs,
238 + ) -> None:
239 + """Creates quantized weights for use in linear operations.
240 +
241 + The function initializes and returns a dictionary containing
242 + quantized weights, scales, and zeros
243 + for performing quantized matrix multiplication operations.
244 +
245 + Args:
246 + input_size_per_partition: The size of the input partition.
247 + output_partition_sizes: The size of the output partition.
248 + input_size: The total size of the input (unused).
249 + output_size: The total size of the output (unused).
250 + params_dtype:
251 + The data type of the parameters (expected to be torch.float16).
252 +
253 + Returns:
254 + A dictionary containing the quantized weights ('qweight'),
255 + scales ('scales'), and zeros ('zeros').
256 +
257 + Raises:
258 + ValueError: If `params_dtype` is not `torch.float16` or
259 + if the input size per partition is not divisible by the
260 + group size in `quant_config`.
261 + """
262 + if params_dtype != torch.float16:
263 + raise ValueError("Parameter data type must be torch.float16, "
264 + f"but got {params_dtype}")
265 +
266 + # Normalize group_size
267 + if self.quant_config.group_size != -1:
268 + group_size = self.quant_config.group_size
269 + else:
270 + group_size = input_size
271 +
272 + if input_size_per_partition % group_size != 0:
273 + raise ValueError(
274 + f"Input size per partition ({input_size_per_partition}) must "
275 + f"be divisible by group size ({self.quant_config.group_size})."
276 + )
277 +
278 + kernel_type = self.kernel_type
279 + # Validate output_size_per_partition
280 + output_size_per_partition = sum(output_partition_sizes)
281 +
282 + is_row_parallel = input_size != input_size_per_partition
283 + weight_loader = extra_weight_attrs.get("weight_loader")
284 +
285 + mp_linear_kernel_config = MPLinearLayerConfig(
286 + full_weight_shape=(input_size, output_size),
287 + partition_weight_shape=\
288 + (input_size_per_partition, output_size_per_partition),
289 + weight_type=self.quant_config.quant_type,
290 + act_type=params_dtype,
291 + group_size=self.quant_config.group_size,
292 + zero_points=False,
293 + has_g_idx=self.quant_config.desc_act
294 + )
295 +
296 + if kernel_type.__name__ not in self._kernel_backends_being_used:
297 + logger.info("Using %s for GPTQBitBLASLinearMethod",
298 + kernel_type.__name__)
299 + self._kernel_backends_being_used.add(kernel_type.__name__)
300 +
301 + # Normalize group_size
302 + if self.quant_config.group_size != -1:
303 + group_size = self.quant_config.group_size
304 + else:
305 + group_size = input_size
306 +
307 + # Determine sharding
308 + if bitblas_repeat_scales_on_all_ranks(self.quant_config.desc_act,
309 + self.quant_config.group_size,
310 + is_row_parallel):
311 + # By setting scale_dim == None, weight_loader will
312 + # repeat the scales on each GPU in TP>1 case.
313 + scales_and_zp_input_dim = None
314 + scales_and_zp_size = input_size // group_size
315 + else:
316 + # By setting scale_dim == 0, weight_loader will
317 + # shard the scales in TP>1 case.
318 + scales_and_zp_input_dim = 0
319 + scales_and_zp_size = input_size_per_partition // group_size
320 +
321 + # Init buffers
322 + # Quantized weights
323 + qweight = PackedvLLMParameter(
324 + data=torch.empty(
325 + input_size_per_partition // self.quant_config.pack_factor,
326 + output_size_per_partition,
327 + dtype=torch.int32,
328 + ),
329 + input_dim=0,
330 + output_dim=1,
331 + packed_dim=0,
332 + packed_factor=self.quant_config.pack_factor,
333 + weight_loader=weight_loader)
334 +
335 + # Activation order
336 + # Ignore warning from fused linear layers such as QKVParallelLinear.
337 + g_idx = RowvLLMParameter(data=torch.empty(
338 + input_size_per_partition,
339 + dtype=torch.int32,
340 + ),
341 + input_dim=0,
342 + weight_loader=weight_loader)
343 +
344 + # Scales
345 + scales = Parameter(
346 + torch.empty(
347 + scales_and_zp_size,
348 + output_size_per_partition,
349 + dtype=params_dtype,
350 + ),
351 + requires_grad=False,
352 + )
353 + set_weight_attrs(
354 + scales,
355 + {
356 + **extra_weight_attrs,
357 + "input_dim": scales_and_zp_input_dim,
358 + "output_dim": 1,
359 + },
360 + )
361 +
362 + # Quantized zero-points
363 + qzeros_args = {
364 + "data":
365 + torch.empty(
366 + scales_and_zp_size,
367 + output_size_per_partition // self.quant_config.pack_factor,
368 + dtype=torch.int32,
369 + ),
370 + "weight_loader":
371 + weight_loader
372 + }
373 + weight_scale_args = {
374 + "data":
375 + torch.empty(
376 + scales_and_zp_size,
377 + output_size_per_partition,
378 + dtype=params_dtype,
379 + ),
380 + "weight_loader":
381 + weight_loader
382 + }
383 +
384 + if scales_and_zp_input_dim is None:
385 + scales = ChannelQuantScaleParameter(output_dim=1,
386 + **weight_scale_args)
387 + qzeros = PackedColumnParameter(
388 + output_dim=1,
389 + packed_dim=1,
390 + packed_factor=self.quant_config.pack_factor,
391 + **qzeros_args)
392 +
393 + else:
394 + scales = GroupQuantScaleParameter(output_dim=1,
395 + input_dim=0,
396 + **weight_scale_args)
397 + qzeros = PackedvLLMParameter(
398 + input_dim=0,
399 + output_dim=1,
400 + packed_dim=1,
401 + packed_factor=self.quant_config.pack_factor,
402 + **qzeros_args)
403 +
404 + layer.register_parameter("qweight", qweight)
405 + layer.register_parameter("g_idx", g_idx)
406 + layer.register_parameter("scales", scales)
407 + layer.register_parameter("qzeros", qzeros)
408 +
409 + self.kernel = kernel_type(
410 + mp_linear_kernel_config,
411 + w_q_param_name="qweight",
412 + w_s_param_name="scales",
413 + w_zp_param_name="qzeros",
414 + w_gidx_param_name="g_idx",
415 + bitblas_quant_config=self.quant_config,
416 + )
417 +
418 + # Initialize or retrieve the BitBLAS matrix multiplication operator.
419 + self.kernel.configure_bitblas_matmul(
420 + input_size_per_partition,
421 + output_size_per_partition,
422 + params_dtype=params_dtype,
423 + bias=False,
424 + )
425 +
426 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
427 + self.kernel.process_weights_after_loading(layer)
428 +
429 + def apply(
430 + self,
431 + layer: torch.nn.Module,
432 + x: torch.Tensor,
433 + bias: Optional[torch.Tensor] = None,
434 + ) -> torch.Tensor:
435 + out = self.kernel.apply_gptq_bitblas_linear(layer, x)
436 + if bias is not None:
437 + out.add_(bias)
438 + return out
vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py
+4 -1
@@ -5,6 +5,8 @@ from typing import List, Optional, Type
5 import vllm.envs as envs
6 from vllm.model_executor.layers.quantization.kernels.mixed_precision.allspark import ( # noqa: E501
7 AllSparkLinearKernel)
8 +from vllm.model_executor.layers.quantization.kernels.mixed_precision.bitblas import ( # noqa: E501
9 + BitBLASLinearKernel)
10 from vllm.model_executor.layers.quantization.kernels.mixed_precision.exllama import ( # noqa: E501
11 ExllamaLinearKernel)
12 from vllm.model_executor.layers.quantization.kernels.mixed_precision.machete import ( # noqa: E501
@@ -20,6 +22,7 @@ _POSSIBLE_KERNELS: List[Type[MPLinearKernel]] = [
22 MacheteLinearKernel,
23 AllSparkLinearKernel,
24 MarlinLinearKernel,
25 + BitBLASLinearKernel,
26 ExllamaLinearKernel,
27 ]
28
@@ -76,4 +79,4 @@ def choose_mp_linear_kernel(
79 raise ValueError(
80 "Failed to find a kernel that can implement the "\
81 "WNA16 linear layer. Reasons: \n"
79 - + '\n'.join(failure_reasons))
82 + + '\n'.join(failure_reasons))
\ No newline at end of file
vllm/model_executor/layers/quantization/kernels/mixed_precision/bitblas.py new
+299
@@ -0,0 +1,299 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import Dict, List, Optional, Tuple
4 +
5 +import torch
6 +
7 +from vllm.logger import init_logger
8 +from vllm.model_executor.layers.quantization.base_config import (
9 + QuantizationConfig)
10 +from vllm.model_executor.layers.quantization.utils import replace_parameter
11 +from vllm.model_executor.layers.quantization.utils.bitblas_utils import (
12 + BITBLAS_OPTIMIZE_FEATURES, BITBLAS_SUPPORTED_GROUP_SIZES,
13 + MINIMUM_BITBLAS_VERSION, bitblas_make_empty_g_idx, bitblas_sort_g_idx,
14 + check_bitblas_supports_shape, query_bitblas_supported_quant_types,
15 + unpack_gptq_qweight, unpack_gptq_qzeros)
16 +
17 +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +class BitBLASLinearKernel(MPLinearKernel):
23 +
24 + OPT_FEATURES: List[int] = BITBLAS_OPTIMIZE_FEATURES
25 + ENABLE_TUNING: bool = True
26 + MATMUL_LAYOUT: str = "nt"
27 + BITBLAS_DTYPES: Dict[torch.dtype, str] = {
28 + torch.float32: "float32",
29 + torch.float16: "float16",
30 + torch.bfloat16: "bfloat16",
31 + torch.half: "float16",
32 + torch.int8: "int8",
33 + }
34 + bitblas_matmul: object = None
35 +
36 + def __init__(
37 + self,
38 + c: MPLinearLayerConfig,
39 + w_q_param_name: str,
40 + w_s_param_name: str,
41 + w_zp_param_name: Optional[str] = None,
42 + w_gidx_param_name: Optional[str] = None,
43 + bitblas_quant_config: Optional[QuantizationConfig] = None,
44 + ):
45 + self.quant_config = bitblas_quant_config
46 + super().__init__(c, w_q_param_name, w_s_param_name, w_zp_param_name,
47 + w_gidx_param_name)
48 +
49 + def repack_bitblas_from_gptq(
50 + self,
51 + b_q_weight: torch.Tensor,
52 + scales: torch.Tensor,
53 + qzeros: Optional[torch.Tensor] = None,
54 + ):
55 + from bitblas.quantization.utils import general_compress
56 + assert self.bitblas_matmul is not None, "bitblas_matmul is None"
57 +
58 + quant_config = self.quant_config
59 + # qweight in gptq old quant linear stored with
60 + # (outfeatures, infeatures), should be transposed.
61 + qweight = b_q_weight.T.contiguous().view(
62 + quant_config.torch_storage_dtype) # type: ignore[union-attr]
63 + intweight = unpack_gptq_qweight(
64 + qweight,
65 + quant_config.weight_bits).contiguous() # type: ignore[union-attr]
66 + if self.bitblas_matmul.weight_transform is not None: # type: ignore[attr-defined]
67 + qweight = self.bitblas_matmul.weight_transform( # type: ignore[attr-defined]
68 + intweight.cpu()).cuda()
69 + # scales in gptq old quant linear stored with
70 + # (infeatures // group_size, outfeatures), should be transposed.
71 + scales = scales.T.contiguous()
72 +
73 + if qzeros is None:
74 + return qweight, scales, None
75 +
76 + # qzeros should be de-quantized to int zeros.
77 + weight_bits = quant_config.weight_bits # type: ignore[union-attr]
78 + intzeros = unpack_gptq_qzeros(qzeros, weight_bits).T.contiguous()
79 + zeros: Optional[torch.Tensor] = None
80 + zeros_mode = self.bitblas_matmul.config.zeros_mode # type: ignore[attr-defined]
81 + if zeros_mode == "original":
82 + zeros = intzeros.to(torch.float16).contiguous()
83 + elif zeros_mode == "rescale":
84 + assert zeros is not None, "zeros should not be None"
85 + zeros[:, :] = intzeros.to(torch.float16)[:, :] * scales[:, :]
86 + elif zeros_mode == "quantized":
87 + zeros = (
88 + torch.Tensor(
89 + general_compress(
90 + intzeros.T.contiguous().cpu().numpy(),
91 + weight_bits,
92 + )).to(qweight.device).
93 + to(quant_config.torch_storage_dtype # type: ignore[union-attr]
94 + ).contiguous())
95 + else:
96 + raise ValueError("Unsupported zeros type: {}".format(zeros_mode))
97 +
98 + return qweight, scales, zeros
99 +
100 + @classmethod
101 + def get_min_capability(cls) -> int:
102 + return 70
103 +
104 + @classmethod
105 + def can_implement(cls,
106 + c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:
107 +
108 + is_bitblas_installed = True
109 +
110 + try:
111 + import bitblas
112 + if bitblas.__version__ < MINIMUM_BITBLAS_VERSION:
113 + raise ImportError(
114 + "bitblas version is wrong. Please "
115 + f"install bitblas>={MINIMUM_BITBLAS_VERSION}")
116 + except ImportError:
117 + is_bitblas_installed = False
118 +
119 + if not is_bitblas_installed:
120 + return False, "bitblas is not installed. Please install bitblas "\
121 + "by running `pip install bitblas>="\
122 + f"{MINIMUM_BITBLAS_VERSION}`"
123 +
124 + quant_types = query_bitblas_supported_quant_types(c.zero_points)
125 + if c.weight_type not in quant_types:
126 + return False, (f"Quant type ({c.weight_type}) not supported by"
127 + f" BitBLAS, supported types are: {quant_types}")
128 +
129 + if c.group_size not in BITBLAS_SUPPORTED_GROUP_SIZES:
130 + return False, (f"Group size ({c.group_size}) not supported by "
131 + "BitBLAS, supported group sizes are: "
132 + f"{BITBLAS_SUPPORTED_GROUP_SIZES}")
133 +
134 + return check_bitblas_supports_shape(
135 + c.partition_weight_shape[1], # out_features
136 + c.partition_weight_shape[0], # in_features
137 + c.full_weight_shape[0], # in_features
138 + c.group_size)
139 +
140 + # note assumes that
141 + # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
142 + # `weight_scale` is: {input_dim = 0, output_dim = 1}
143 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
144 + device = getattr(layer, self.w_q_name).device
145 + c = self.config
146 + quant_config = self.quant_config
147 +
148 + # Default names since bitblas requires empty parameters for these,
149 + # TODO: remove this requirement from bitblas (allow optional tensors)
150 + if self.w_gidx_name is None:
151 + self.w_gidx_name = "g_idx"
152 + if self.w_zp_name is None:
153 + self.w_zp_name = "qzeros"
154 +
155 + if c.has_g_idx:
156 + g_idx, g_idx_sort_indices = bitblas_sort_g_idx(
157 + getattr(layer, self.w_gidx_name))
158 + self._transform_param(layer, self.w_gidx_name, lambda _: g_idx)
159 + layer.g_idx_sort_indices = g_idx_sort_indices
160 + else:
161 + setattr(layer, self.w_gidx_name, bitblas_make_empty_g_idx(device))
162 + layer.g_idx_sort_indices = bitblas_make_empty_g_idx(device)
163 +
164 + if c.zero_points:
165 + raise NotImplementedError("Zero points not supported by BitBLAS")
166 + else:
167 + setattr(layer, self.w_zp_name, bitblas_make_empty_g_idx(device))
168 +
169 + # Repack weights
170 + bitblas_qweight, bitblas_scales, bitblas_qzeros = (
171 + self.repack_bitblas_from_gptq(
172 + layer.qweight,
173 + layer.scales,
174 + None if quant_config.is_sym else # type: ignore[union-attr]
175 + layer.qzeros, # type: ignore[union-attr]
176 + ))
177 + replace_parameter(layer, self.w_q_name, bitblas_qweight)
178 + replace_parameter(layer, self.w_s_name, bitblas_scales)
179 + if bitblas_qzeros is not None:
180 + replace_parameter(layer, self.w_zp_name, bitblas_qzeros)
181 +
182 + def configure_bitblas_matmul(
183 + self,
184 + infeatures: int,
185 + outfeatures: int,
186 + params_dtype: torch.dtype,
187 + bias: bool,
188 + ) -> None:
189 + enable_tuning = self.ENABLE_TUNING
190 + layout = self.MATMUL_LAYOUT
191 + bits = self.quant_config.weight_bits # type: ignore[union-attr]
192 + self._configure_bitblas_matmul(
193 + infeatures,
194 + outfeatures,
195 + params_dtype,
196 + enable_tuning,
197 + bias,
198 + layout,
199 + bits,
200 + )
201 +
202 + def _configure_bitblas_matmul(
203 + self,
204 + infeatures,
205 + outfeatures,
206 + params_dtype,
207 + enable_tuning,
208 + bias,
209 + layout,
210 + bits,
211 + ):
212 + from bitblas import MatmulConfig
213 + bitblas_dtype = self.BITBLAS_DTYPES[params_dtype]
214 + quant_config = self.quant_config
215 + with_scaling = False
216 + with_zeros = False
217 + group_size = quant_config.group_size # type: ignore[union-attr]
218 + zeros_mode = quant_config.zeros_mode # type: ignore[union-attr]
219 + if quant_config.quant_method == "gptq": # type: ignore[union-attr]
220 + with_scaling = True
221 + with_zeros = True
222 + W_dtype = f"uint{bits}"
223 + if quant_config.is_sym: # type: ignore[union-attr]
224 + with_zeros = False
225 + W_dtype = f"int{bits}"
226 + else:
227 + raise ValueError(
228 + f"Unsupported quant_method {quant_config.quant_method}" # type: ignore[union-attr]
229 + ) # type: ignore[union-attr]
230 +
231 + matmul_config = MatmulConfig(
232 + M=self.OPT_FEATURES,
233 + N=outfeatures,
234 + K=infeatures,
235 + A_dtype=bitblas_dtype,
236 + W_dtype=W_dtype,
237 + out_dtype=bitblas_dtype,
238 + accum_dtype="int32" if bitblas_dtype == "int8" else bitblas_dtype,
239 + storage_dtype=quant_config. # type: ignore[union-attr]
240 + storage_dtype, # type: ignore[union-attr]
241 + with_scaling=with_scaling,
242 + with_zeros=with_zeros,
243 + group_size=group_size,
244 + with_bias=bias,
245 + layout=layout,
246 + zeros_mode=zeros_mode,
247 + )
248 + self.bitblas_matmul = self._get_or_create_bitblas_operator(
249 + matmul_config, enable_tuning)
250 +
251 + def _get_or_create_bitblas_operator(self, config, enable_tuning):
252 + from bitblas import Matmul, auto_detect_nvidia_target
253 + from bitblas.cache import get_database_path, global_operator_cache
254 + BITBLAS_DATABASE_PATH = get_database_path()
255 + BITBLAS_TARGET = auto_detect_nvidia_target()
256 +
257 + if global_operator_cache.size() == 0:
258 + global_operator_cache.load_from_database(BITBLAS_DATABASE_PATH,
259 + BITBLAS_TARGET)
260 +
261 + bitblas_matmul = global_operator_cache.get(config)
262 + if bitblas_matmul is None:
263 + bitblas_matmul = Matmul(config,
264 + target=BITBLAS_TARGET,
265 + enable_tuning=False)
266 + if enable_tuning:
267 + bitblas_matmul.hardware_aware_finetune(topk=20)
268 + global_operator_cache.add(config, bitblas_matmul)
269 + global_operator_cache.save_into_database(
270 + BITBLAS_DATABASE_PATH, BITBLAS_TARGET)
271 + TUNING_MESSAGE = (
272 + f"BitBLAS Operator {config} tuned and saved to database.")
273 + logger.info(TUNING_MESSAGE)
274 + else:
275 + _message = f"BitBLAS Operator {config} created without tuning. "
276 + logger.info(_message)
277 + else:
278 + _message = f"BitBLAS Operator {config} retrieved from cache."
279 + logger.info(_message)
280 + return bitblas_matmul
281 +
282 + def apply_gptq_bitblas_linear(
283 + self,
284 + layer: torch.nn.Module,
285 + x: torch.Tensor,
286 + ) -> torch.Tensor:
287 + output_size_per_partition = self.config.partition_weight_shape[1]
288 + out_shape = x.shape[:-1] + (output_size_per_partition, )
289 + args = [x, layer.qweight, layer.scales]
290 + if self.bitblas_matmul.config.with_zeros: # type: ignore[attr-defined]
291 + args.append(layer.qzeros)
292 + output = self.bitblas_matmul(*args) # type: ignore[operator]
293 + return output.view(out_shape)
294 +
295 + def apply_weights(self, layer, x, bias=None):
296 + NOT_IMPLEMENT_MESSAGE = (
297 + f"{self.__class__.__name__}.apply_weights is not implemented. "
298 + "Please use BitBLASLinearKernel.apply_gptq_bitblas_linear instead")
299 + raise NotImplementedError(NOT_IMPLEMENT_MESSAGE)
vllm/model_executor/layers/quantization/utils/bitblas_utils.py new
+198
@@ -0,0 +1,198 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional, Tuple
3 +
4 +import torch
5 +
6 +from vllm.platforms import current_platform
7 +from vllm.scalar_type import ScalarType, scalar_types
8 +
9 +MINIMUM_BITBLAS_VERSION = "0.1.0"
10 +
11 +BITBLAS_MIN_WEIGHT_SIZE_N = 16
12 +BITBLAS_MIN_WEIGHT_SIZE_K = 16
13 +GPTQ_BITBLAS_MAX_PARALLEL = 16
14 +
15 +BITBLAS_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
16 +
17 +# For dynamic shape code generation
18 +BITBLAS_OPTIMIZE_FEATURES = [1, 16, 32, 64, 128, 256, 512, 1024]
19 +# If want to enable high performance for contiguous batching
20 +# Please use the following values
21 +BITBLAS_OPTIMIZE_FEATURES_CONTIGUOUS = [16, 32, 64, 128, 256, 512, 1024]
22 +
23 +BITBLAS_SUPPORTED_NUM_BITS = [1, 2, 4, 8]
24 +BITBLAS_SUPPORTED_SYM = [False, True]
25 +
26 +
27 +# Determines the supported quantization types for BitBLAS based on the
28 +# device's capability and whether zero-point (zp) is used.
29 +def query_bitblas_supported_quant_types(has_zp: bool,
30 + device_capability: Optional[int] = None
31 + ):
32 + if device_capability is None:
33 + capability_tuple = current_platform.get_device_capability()
34 + device_capability = (-1 if capability_tuple is None else
35 + capability_tuple.to_int())
36 +
37 + if device_capability < 70:
38 + return []
39 +
40 + if has_zp:
41 + # AWQ style, unsigned + runtime zero-point
42 + return [scalar_types.uint4, scalar_types.uint8]
43 + else:
44 + # GPTQ style, unsigned + symmetric bias
45 + # TODO: once fp8_bitblas is merged into "gptq_bitblas" we should be able
46 + # to add `scalar_types.float8_e4m3fn` here
47 + return [scalar_types.uint4b8, scalar_types.uint8b128]
48 +
49 +
50 +def _check_bitblas_supported(
51 + quant_type: ScalarType,
52 + group_size: Optional[int],
53 + has_zp: bool,
54 + device_capability: Optional[int] = None) -> Tuple[bool, Optional[str]]:
55 +
56 + if device_capability is None:
57 + capability_tuple = current_platform.get_device_capability()
58 + device_capability = (-1 if capability_tuple is None else
59 + capability_tuple.to_int())
60 +
61 + supported_types = query_bitblas_supported_quant_types(
62 + has_zp, device_capability)
63 +
64 + if quant_type not in supported_types:
65 + return (False, f"BitBLAS does not support weight_bits = {quant_type}. "
66 + f"Only types = {supported_types} "
67 + f"are supported (for group_size = {group_size}, "
68 + f"device_capability = {device_capability}, zp = {has_zp}).")
69 + if (group_size is None or group_size not in BITBLAS_SUPPORTED_GROUP_SIZES):
70 + return (False, f"BitBLAS does not support group_size = {group_size}. "
71 + f"Only group_sizes = {BITBLAS_SUPPORTED_GROUP_SIZES} "
72 + "are supported.")
73 +
74 + return True, None
75 +
76 +
77 +def check_bitblas_supported(quant_type: ScalarType,
78 + group_size: int,
79 + has_zp: bool = False,
80 + device_capability: Optional[int] = None) -> bool:
81 + cond, _ = _check_bitblas_supported(quant_type, group_size, has_zp,
82 + device_capability)
83 + return cond
84 +
85 +
86 +def verify_bitblas_supported(quant_type: ScalarType,
87 + group_size: int,
88 + has_zp: bool = False) -> None:
89 + cond, err_msg = _check_bitblas_supported(quant_type, group_size, has_zp)
90 + if not cond:
91 + assert err_msg is not None
92 + raise ValueError(err_msg)
93 +
94 +
95 +def verify_bitblas_supports_shape(output_size_per_partition: int,
96 + input_size_per_partition: int,
97 + input_size: int, group_size: int) -> None:
98 +
99 + # Validate output_size_per_partition
100 + if output_size_per_partition % BITBLAS_MIN_WEIGHT_SIZE_N != 0:
101 + raise ValueError(f"Weight output_size_per_partition = "
102 + f"{output_size_per_partition} is not divisible by "
103 + f" min_thread_n = {BITBLAS_MIN_WEIGHT_SIZE_N}. "
104 + "Consider reducing tensor_parallel_size or running "
105 + "with --quantization gptq.")
106 +
107 + # Validate input_size_per_partition
108 + if input_size_per_partition % BITBLAS_MIN_WEIGHT_SIZE_K != 0:
109 + raise ValueError(f"Weight input_size_per_partition = "
110 + f"{input_size_per_partition} is not divisible "
111 + f"by min_thread_k = {BITBLAS_MIN_WEIGHT_SIZE_K}. "
112 + "Consider reducing tensor_parallel_size or running "
113 + "with --quantization gptq.")
114 +
115 + if (group_size < input_size
116 + and input_size_per_partition % group_size != 0):
117 + raise ValueError(
118 + f"Weight input_size_per_partition = {input_size_per_partition}"
119 + f" is not divisible by group_size = {group_size}."
120 + "Consider reducing tensor_parallel_size or running "
121 + "with --quantization gptq.")
122 +
123 +
124 +def check_bitblas_supports_shape(output_size_per_partition: int,
125 + input_size_per_partition: int,
126 + input_size: int, group_size: int) \
127 + -> Tuple[bool, Optional[str]]:
128 + try:
129 + verify_bitblas_supports_shape(output_size_per_partition,
130 + input_size_per_partition, input_size,
131 + group_size)
132 + except ValueError as e:
133 + return False, e.__str__()
134 + return True, None
135 +
136 +
137 +def bitblas_is_k_full(act_order: bool, is_row_parallel: bool) -> bool:
138 + return (not act_order) or (act_order and not is_row_parallel)
139 +
140 +
141 +def bitblas_repeat_scales_on_all_ranks(act_order: bool, group_size: int,
142 + is_row_parallel: bool) -> bool:
143 + # Need to repeat scales on every rank if act_ordering or
144 + # channelwise and RowParallelLinear
145 + is_channelwise = group_size == -1
146 + return act_order or (is_channelwise and is_row_parallel)
147 +
148 +
149 +def bitblas_make_empty_g_idx(device: torch.device) -> torch.Tensor:
150 + return torch.nn.Parameter(torch.empty(0, dtype=torch.int, device=device),
151 + requires_grad=False)
152 +
153 +
154 +def bitblas_make_empty_zp(device: torch.device) -> torch.Tensor:
155 + return torch.nn.Parameter(torch.empty(0, dtype=torch.int, device=device),
156 + requires_grad=False)
157 +
158 +
159 +def bitblas_sort_g_idx(
160 + g_idx: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
161 + g_idx_sort_indices = torch.argsort(g_idx).to(torch.int)
162 + return g_idx[g_idx_sort_indices], g_idx_sort_indices
163 +
164 +
165 +def unpack_gptq_qzeros(qzeros, bits, is_gptq_v2=False) -> torch.Tensor:
166 + qzeros = qzeros.view(torch.int32)
167 + elems_per_int32 = 32 // bits
168 + unpacked_zeros = torch.zeros(
169 + (qzeros.shape[0], qzeros.shape[1] * elems_per_int32),
170 + dtype=torch.int8,
171 + device=qzeros.device,
172 + requires_grad=False,
173 + )
174 +
175 + for col in range(unpacked_zeros.shape[1]):
176 + i = col % elems_per_int32
177 + unpacked_zeros[:, col] = (qzeros[:, col // elems_per_int32] >>
178 + (bits * i)) & 0xF
179 + if not is_gptq_v2:
180 + return unpacked_zeros + 1
181 + return unpacked_zeros
182 +
183 +
184 +def unpack_gptq_qweight(qweight, bits):
185 + qweight = qweight.view(torch.int8)
186 + elems_per_int8 = 8 // bits
187 + unpacked_weight = torch.zeros(
188 + (qweight.shape[0], qweight.shape[1] * elems_per_int8),
189 + dtype=torch.int8,
190 + device=qweight.device,
191 + requires_grad=False,
192 + )
193 + for col in range(unpacked_weight.shape[1]):
194 + i = col % elems_per_int8
195 + unpacked_weight[:, col] = (qweight[:, col // elems_per_int8] >>
196 + (bits * i))
197 +
198 + return torch.bitwise_and(unpacked_weight, 2**bits - 1)
vllm/model_executor/parameter.py
+29 -4
@@ -282,10 +282,12 @@ class PackedColumnParameter(_ColumnvLLMParameter):
282 packed_factor: Union[int, Fraction],
283 packed_dim: int,
284 marlin_tile_size: Optional[int] = None,
285 + bitblas_tile_size: Optional[int] = None,
286 **kwargs):
287 self._packed_factor = packed_factor
288 self._packed_dim = packed_dim
289 self._marlin_tile_size = marlin_tile_size
290 + self._bitblas_tile_size = bitblas_tile_size
291 super().__init__(**kwargs)
292
293 @property
@@ -300,12 +302,17 @@ class PackedColumnParameter(_ColumnvLLMParameter):
302 def marlin_tile_size(self):
303 return self._marlin_tile_size
304
305 + @property
306 + def bitblas_tile_size(self):
307 + return self._bitblas_tile_size
308 +
309 def adjust_shard_indexes_for_packing(self, shard_size, shard_offset):
310 return _adjust_shard_indexes_for_packing(
311 shard_size=shard_size,
312 shard_offset=shard_offset,
313 packed_factor=self.packed_factor,
308 - marlin_tile_size=self.marlin_tile_size)
314 + marlin_tile_size=self.marlin_tile_size,
315 + bitblas_tile_size=self.bitblas_tile_size)
316
317
318 class PackedvLLMParameter(ModelWeightParameter):
@@ -323,10 +330,12 @@ class PackedvLLMParameter(ModelWeightParameter):
330 packed_factor: Union[int, Fraction],
331 packed_dim: int,
332 marlin_tile_size: Optional[int] = None,
333 + bitblas_tile_size: Optional[int] = None,
334 **kwargs):
335 self._packed_factor = packed_factor
336 self._packed_dim = packed_dim
337 self._marlin_tile_size = marlin_tile_size
338 + self._bitblas_tile_size = bitblas_tile_size
339 super().__init__(**kwargs)
340
341 @property
@@ -341,12 +350,17 @@ class PackedvLLMParameter(ModelWeightParameter):
350 def marlin_tile_size(self):
351 return self._marlin_tile_size
352
353 + @property
354 + def bitblas_tile_size(self):
355 + return self._bitblas_tile_size
356 +
357 def adjust_shard_indexes_for_packing(self, shard_size, shard_offset):
358 return _adjust_shard_indexes_for_packing(
359 shard_size=shard_size,
360 shard_offset=shard_offset,
361 packed_factor=self.packed_factor,
349 - marlin_tile_size=self.marlin_tile_size)
362 + marlin_tile_size=self.marlin_tile_size,
363 + bitblas_tile_size=self.bitblas_tile_size)
364
365
366 class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
@@ -421,8 +435,13 @@ def _adjust_shard_indexes_for_marlin(shard_size, shard_offset,
435 return shard_size * marlin_tile_size, shard_offset * marlin_tile_size
436
437
438 +def _adjust_shard_indexes_for_bitblas(shard_size, shard_offset,
439 + bitblas_tile_size):
440 + return shard_size // bitblas_tile_size, shard_offset // bitblas_tile_size
441 +
442 +
443 def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor,
425 - marlin_tile_size):
444 + marlin_tile_size, bitblas_tile_size):
445 shard_size = shard_size // packed_factor
446 shard_offset = shard_offset // packed_factor
447 if marlin_tile_size is not None:
@@ -430,4 +449,10 @@ def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor,
449 shard_size=shard_size,
450 shard_offset=shard_offset,
451 marlin_tile_size=marlin_tile_size)
433 - return shard_size, shard_offset
452 + elif bitblas_tile_size is not None:
453 + return _adjust_shard_indexes_for_bitblas(
454 + shard_size=shard_size,
455 + shard_offset=shard_offset,
456 + bitblas_tile_size=bitblas_tile_size)
457 +
458 + return shard_size, shard_offset
\ No newline at end of file