Skip to content

vllm.model_executor.layers.quantization.quark.schemes.quark_scheme

Classes:

  • QuarkScheme –

    Abstract class used to describe the weight creation and forward pass

QuarkScheme

Bases: ABC

Abstract class used to describe the weight creation and forward pass of different quantization schemes supported by Quark.

Methods:

Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class QuarkScheme(ABC):
    """Abstract class used to describe the weight creation and forward pass
    of different quantization schemes supported by Quark.
    """

    supported_activation_quant_keys: list[QuantKey | None] = []
    supported_weight_quant_keys: list[QuantKey] = []

    def __init__(
        self,
        weight_quant_key: QuantKey,
        activation_quant_key: QuantKey | None,
    ):
        if activation_quant_key not in self.supported_activation_quant_keys:
            raise ValueError(
                f"Unsupported activation quant key: {activation_quant_key}"
            )
        if weight_quant_key not in self.supported_weight_quant_keys:
            raise ValueError(f"Unsupported weight quant key: {weight_quant_key}")
        self.activation_quant_key = activation_quant_key
        self.weight_quant_key = weight_quant_key

    @classmethod
    @abstractmethod
    def get_min_capability(cls) -> int:
        """Get minimum device capability."""
        raise NotImplementedError

    @abstractmethod
    def create_weights(self, *args, **kwargs):
        """Weight creation for the particular scheme. Inputs to this function."""
        raise NotImplementedError

    @abstractmethod
    def apply_weights(
        self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
    ):
        """Run the forward pass for the particular scheme. This is where
        scheme-specific dequant/quant steps/kernels should be applied.

        Args:
            layer: torch.nn.Module with the registered weights and
                other parameters relevant to the particular scheme.
            x: input to the layer
            bias: bias parameter

        """
        raise NotImplementedError

    @abstractmethod
    def process_weights_after_loading(self, layer: torch.nn.Module):
        """Called after weight loading is complete for any cleanup that
        needs to occur.
        """
        raise NotImplementedError

apply_weights(layer, x, bias) abstractmethod

Run the forward pass for the particular scheme. This is where scheme-specific dequant/quant steps/kernels should be applied.

Parameters:

  • layer

    (Module) –

    torch.nn.Module with the registered weights and other parameters relevant to the particular scheme.

  • x

    (Tensor) –

    input to the layer

  • bias

    (Tensor | None) –

    bias parameter

Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
@abstractmethod
def apply_weights(
    self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
):
    """Run the forward pass for the particular scheme. This is where
    scheme-specific dequant/quant steps/kernels should be applied.

    Args:
        layer: torch.nn.Module with the registered weights and
            other parameters relevant to the particular scheme.
        x: input to the layer
        bias: bias parameter

    """
    raise NotImplementedError

create_weights(*args, **kwargs) abstractmethod

Weight creation for the particular scheme. Inputs to this function.

Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py 
41
42
43
44
@abstractmethod
def create_weights(self, *args, **kwargs):
    """Weight creation for the particular scheme. Inputs to this function."""
    raise NotImplementedError

get_min_capability() abstractmethod classmethod

Get minimum device capability.

Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py 
35
36
37
38
39
@classmethod
@abstractmethod
def get_min_capability(cls) -> int:
    """Get minimum device capability."""
    raise NotImplementedError

process_weights_after_loading(layer) abstractmethod

Called after weight loading is complete for any cleanup that needs to occur.

Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py 
62
63
64
65
66
67
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module):
    """Called after weight loading is complete for any cleanup that
    needs to occur.
    """
    raise NotImplementedError