| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 9 | |
| 10 | __all__ = ["QuarkScheme"] |
| 11 | |
| 12 | |
| 13 | class QuarkScheme(ABC): |
| 14 | """Abstract class used to describe the weight creation and forward pass |
| 15 | of different quantization schemes supported by Quark. |
| 16 | """ |
| 17 | |
| 18 | supported_activation_quant_keys: list[QuantKey | None] = [] |
| 19 | supported_weight_quant_keys: list[QuantKey] = [] |
| 20 | |
| 21 | def __init__( |
| 22 | self, |
| 23 | weight_quant_key: QuantKey, |
| 24 | activation_quant_key: QuantKey | None, |
| 25 | ): |
| 26 | if activation_quant_key not in self.supported_activation_quant_keys: |
| 27 | raise ValueError( |
| 28 | f"Unsupported activation quant key: {activation_quant_key}" |
| 29 | ) |
| 30 | if weight_quant_key not in self.supported_weight_quant_keys: |
| 31 | raise ValueError(f"Unsupported weight quant key: {weight_quant_key}") |
| 32 | self.activation_quant_key = activation_quant_key |
| 33 | self.weight_quant_key = weight_quant_key |
| 34 | |
| 35 | @classmethod |
| 36 | @abstractmethod |
| 37 | def get_min_capability(cls) -> int: |
| 38 | """Get minimum device capability.""" |
| 39 | raise NotImplementedError |
| 40 | |
| 41 | @abstractmethod |
| 42 | def create_weights(self, *args, **kwargs): |
| 43 | """Weight creation for the particular scheme. Inputs to this function.""" |
| 44 | raise NotImplementedError |
| 45 | |
| 46 | @abstractmethod |
| 47 | def apply_weights( |
| 48 | self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None |
| 49 | ): |
| 50 | """Run the forward pass for the particular scheme. This is where |
| 51 | scheme-specific dequant/quant steps/kernels should be applied. |
| 52 | |
| 53 | Args: |
| 54 | layer: torch.nn.Module with the registered weights and |
| 55 | other parameters relevant to the particular scheme. |
| 56 | x: input to the layer |
| 57 | bias: bias parameter |
| 58 | |
| 59 | """ |
| 60 | raise NotImplementedError |
| 61 | |
| 62 | @abstractmethod |
| 63 | def process_weights_after_loading(self, layer: torch.nn.Module): |
| 64 | """Called after weight loading is complete for any cleanup that |
| 65 | needs to occur. |
| 66 | """ |
| 67 | raise NotImplementedError |