Skip to content

vllm.model_executor.layers.quantization.quark.schemes.quark_w8a8_int8

logger module-attribute 

logger = init_logger(__name__)

QuarkW8A8Int8

Bases: QuarkScheme

Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py 
 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
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
class QuarkW8A8Int8(QuarkScheme):
    _kernel_backends_being_used: set[str] = set()

    def __init__(
        self,
        qscheme: str,
        is_static_input_scheme: bool | None,
        input_symmetric: bool | None,
    ):
        self.qscheme = qscheme
        self.is_static_input_scheme = is_static_input_scheme
        self.input_symmetric = input_symmetric

    @classmethod
    def get_min_capability(cls) -> int:
        # turing and up
        return 75

    def create_weights(
        self,
        layer: torch.nn.Module,
        output_partition_sizes: list[int],
        input_size_per_partition: int,
        params_dtype: torch.dtype,
        weight_loader: Callable,
        **kwargs,
    ):
        layer.logical_widths = output_partition_sizes

        scaled_mm_linear_kernel_config = ScaledMMLinearLayerConfig(
            is_channelwise=(self.qscheme == "per_channel"),
            is_static_input_scheme=(self.is_static_input_scheme is True),
            input_symmetric=(self.input_symmetric is True),
        )

        kernel_type = choose_scaled_mm_linear_kernel(scaled_mm_linear_kernel_config)

        if kernel_type.__name__ not in self._kernel_backends_being_used:
            logger.info("Using %s for QuarkW8A8Int8", kernel_type.__name__)
            self._kernel_backends_being_used.add(kernel_type.__name__)

        # WEIGHT
        weight = ModelWeightParameter(
            data=torch.empty(
                sum(output_partition_sizes), input_size_per_partition, dtype=torch.int8
            ),
            input_dim=1,
            output_dim=0,
            weight_loader=weight_loader,
        )

        layer.register_parameter("weight", weight)

        # WEIGHT SCALE
        if self.qscheme == "per_channel":
            weight_scale = ChannelQuantScaleParameter(
                data=torch.empty((sum(output_partition_sizes)), dtype=torch.float32),
                output_dim=0,
                weight_loader=weight_loader,
            )
            ChannelQuantZPParameter = ChannelQuantScaleParameter
            weight_zero_point = ChannelQuantZPParameter(
                data=torch.empty((sum(output_partition_sizes)), dtype=torch.int8),
                output_dim=0,
                weight_loader=weight_loader,
            )
        else:
            assert self.qscheme == "per_tensor"
            weight_scale = PerTensorScaleParameter(
                data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
                weight_loader=weight_loader,
            )
            PerTensorZPParameter = PerTensorScaleParameter
            weight_zero_point = PerTensorZPParameter(
                data=torch.empty(len(output_partition_sizes), dtype=torch.int8),
                weight_loader=weight_loader,
            )
        layer.register_parameter("weight_scale", weight_scale)
        layer.register_parameter("weight_zero_point", weight_zero_point)

        # INPUT SCALE
        if self.is_static_input_scheme:
            input_scale = BasevLLMParameter(
                data=torch.empty(1, dtype=torch.float32), weight_loader=weight_loader
            )
            layer.register_parameter("input_scale", input_scale)

            input_zero_point = BasevLLMParameter(
                data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
            )
            layer.register_parameter("input_zero_point", input_zero_point)

        self.kernel = kernel_type(
            c=scaled_mm_linear_kernel_config,
            w_q_param_name="weight",
            w_s_param_name="weight_scale",
            i_s_param_name="input_scale",
            i_zp_param_name="input_zero_point",
            azp_adj_param_name="azp_adj",
        )

    # Checkpoints are serialized in quark format, which is
    # different from the format the kernel may want. Handle repacking here.
    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        layer.register_parameter("weight_zero_point", None)
        delattr(layer, "weight_zero_point")
        if self.input_symmetric:
            layer.register_parameter("input_zero_point", None)
            delattr(layer, "input_zero_point")

        self.kernel.process_weights_after_loading(layer)

    def apply_weights(
        self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
    ) -> torch.Tensor:
        return self.kernel.apply_weights(layer, x, bias)

_kernel_backends_being_used class-attribute instance-attribute 

_kernel_backends_being_used: set[str] = set()

input_symmetric instance-attribute 

input_symmetric = input_symmetric

is_static_input_scheme instance-attribute 

is_static_input_scheme = is_static_input_scheme

qscheme instance-attribute 

qscheme = qscheme

__init__ 

__init__(
    qscheme: str,
    is_static_input_scheme: bool | None,
    input_symmetric: bool | None,
)
Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py 
27
28
29
30
31
32
33
34
35
def __init__(
    self,
    qscheme: str,
    is_static_input_scheme: bool | None,
    input_symmetric: bool | None,
):
    self.qscheme = qscheme
    self.is_static_input_scheme = is_static_input_scheme
    self.input_symmetric = input_symmetric

apply_weights 

apply_weights(
    layer: Module, x: Tensor, bias: Tensor | None
) -> Tensor
Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py 
136
137
138
139
def apply_weights(
    self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
) -> torch.Tensor:
    return self.kernel.apply_weights(layer, x, bias)

create_weights 

create_weights(
    layer: Module,
    output_partition_sizes: list[int],
    input_size_per_partition: int,
    params_dtype: dtype,
    weight_loader: Callable,
    **kwargs,
)
Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py 
 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
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def create_weights(
    self,
    layer: torch.nn.Module,
    output_partition_sizes: list[int],
    input_size_per_partition: int,
    params_dtype: torch.dtype,
    weight_loader: Callable,
    **kwargs,
):
    layer.logical_widths = output_partition_sizes

    scaled_mm_linear_kernel_config = ScaledMMLinearLayerConfig(
        is_channelwise=(self.qscheme == "per_channel"),
        is_static_input_scheme=(self.is_static_input_scheme is True),
        input_symmetric=(self.input_symmetric is True),
    )

    kernel_type = choose_scaled_mm_linear_kernel(scaled_mm_linear_kernel_config)

    if kernel_type.__name__ not in self._kernel_backends_being_used:
        logger.info("Using %s for QuarkW8A8Int8", kernel_type.__name__)
        self._kernel_backends_being_used.add(kernel_type.__name__)

    # WEIGHT
    weight = ModelWeightParameter(
        data=torch.empty(
            sum(output_partition_sizes), input_size_per_partition, dtype=torch.int8
        ),
        input_dim=1,
        output_dim=0,
        weight_loader=weight_loader,
    )

    layer.register_parameter("weight", weight)

    # WEIGHT SCALE
    if self.qscheme == "per_channel":
        weight_scale = ChannelQuantScaleParameter(
            data=torch.empty((sum(output_partition_sizes)), dtype=torch.float32),
            output_dim=0,
            weight_loader=weight_loader,
        )
        ChannelQuantZPParameter = ChannelQuantScaleParameter
        weight_zero_point = ChannelQuantZPParameter(
            data=torch.empty((sum(output_partition_sizes)), dtype=torch.int8),
            output_dim=0,
            weight_loader=weight_loader,
        )
    else:
        assert self.qscheme == "per_tensor"
        weight_scale = PerTensorScaleParameter(
            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
            weight_loader=weight_loader,
        )
        PerTensorZPParameter = PerTensorScaleParameter
        weight_zero_point = PerTensorZPParameter(
            data=torch.empty(len(output_partition_sizes), dtype=torch.int8),
            weight_loader=weight_loader,
        )
    layer.register_parameter("weight_scale", weight_scale)
    layer.register_parameter("weight_zero_point", weight_zero_point)

    # INPUT SCALE
    if self.is_static_input_scheme:
        input_scale = BasevLLMParameter(
            data=torch.empty(1, dtype=torch.float32), weight_loader=weight_loader
        )
        layer.register_parameter("input_scale", input_scale)

        input_zero_point = BasevLLMParameter(
            data=torch.empty(1, dtype=torch.int8), weight_loader=weight_loader
        )
        layer.register_parameter("input_zero_point", input_zero_point)

    self.kernel = kernel_type(
        c=scaled_mm_linear_kernel_config,
        w_q_param_name="weight",
        w_s_param_name="weight_scale",
        i_s_param_name="input_scale",
        i_zp_param_name="input_zero_point",
        azp_adj_param_name="azp_adj",
    )

get_min_capability classmethod 

get_min_capability() -> int
Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py 
37
38
39
40
@classmethod
def get_min_capability(cls) -> int:
    # turing and up
    return 75

process_weights_after_loading 

process_weights_after_loading(layer: Module) -> None
Source code in vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py 
127
128
129
130
131
132
133
134
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    layer.register_parameter("weight_zero_point", None)
    delattr(layer, "weight_zero_point")
    if self.input_symmetric:
        layer.register_parameter("input_zero_point", None)
        delattr(layer, "input_zero_point")

    self.kernel.process_weights_after_loading(layer)