Skip to content

vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel

ScaledMMLinearKernel

Bases: ABC

Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
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
68
69
70
71
72
73
class ScaledMMLinearKernel(ABC):
    @classmethod
    @abstractmethod
    def get_min_capability(cls) -> int:
        raise NotImplementedError

    @classmethod
    @abstractmethod
    def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
        raise NotImplementedError

    def __init__(
        self,
        c: ScaledMMLinearLayerConfig,
        w_q_param_name: str,
        w_s_param_name: str,
        i_s_param_name: str,
        i_zp_param_name: str,
        azp_adj_param_name: str,
    ) -> None:
        assert self.can_implement(c)
        self.config = c
        self.w_q_name = w_q_param_name
        self.w_s_name = w_s_param_name
        self.i_s_name = i_s_param_name
        self.i_zp_name = i_zp_param_name
        self.azp_adj_name = azp_adj_param_name

    @abstractmethod
    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        raise NotImplementedError

    @abstractmethod
    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

    def _get_weight_params(
        self, layer: torch.nn.Module
    ) -> tuple[
        torch.Tensor,  # weight
        torch.Tensor,  # weight_scale
        torch.Tensor | None,  # input_scale,
        torch.Tensor | None,  # input_zp
        torch.Tensor | None,  # azp_adj
    ]:
        return (
            getattr(layer, self.w_q_name),
            getattr(layer, self.w_s_name),
            getattr(layer, self.i_s_name),
            getattr(layer, self.i_zp_name),
            getattr(layer, self.azp_adj_name),
        )

azp_adj_name instance-attribute 

azp_adj_name = azp_adj_param_name

config instance-attribute 

config = c

i_s_name instance-attribute 

i_s_name = i_s_param_name

i_zp_name instance-attribute 

i_zp_name = i_zp_param_name

w_q_name instance-attribute 

w_q_name = w_q_param_name

w_s_name instance-attribute 

w_s_name = w_s_param_name

__init__ 

__init__(
    c: ScaledMMLinearLayerConfig,
    w_q_param_name: str,
    w_s_param_name: str,
    i_s_param_name: str,
    i_zp_param_name: str,
    azp_adj_param_name: str,
) -> None
Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def __init__(
    self,
    c: ScaledMMLinearLayerConfig,
    w_q_param_name: str,
    w_s_param_name: str,
    i_s_param_name: str,
    i_zp_param_name: str,
    azp_adj_param_name: str,
) -> None:
    assert self.can_implement(c)
    self.config = c
    self.w_q_name = w_q_param_name
    self.w_s_name = w_s_param_name
    self.i_s_name = i_s_param_name
    self.i_zp_name = i_zp_param_name
    self.azp_adj_name = azp_adj_param_name

_get_weight_params 

_get_weight_params(
    layer: Module,
) -> tuple[
    Tensor,
    Tensor,
    Tensor | None,
    Tensor | None,
    Tensor | None,
]
Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def _get_weight_params(
    self, layer: torch.nn.Module
) -> tuple[
    torch.Tensor,  # weight
    torch.Tensor,  # weight_scale
    torch.Tensor | None,  # input_scale,
    torch.Tensor | None,  # input_zp
    torch.Tensor | None,  # azp_adj
]:
    return (
        getattr(layer, self.w_q_name),
        getattr(layer, self.w_s_name),
        getattr(layer, self.i_s_name),
        getattr(layer, self.i_zp_name),
        getattr(layer, self.azp_adj_name),
    )

apply_weights abstractmethod 

apply_weights(
    layer: Module, x: Tensor, bias: Tensor | None = None
) -> Tensor
Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
49
50
51
52
53
54
55
56
@abstractmethod
def apply_weights(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: torch.Tensor | None = None,
) -> torch.Tensor:
    raise NotImplementedError

can_implement abstractmethod classmethod 

can_implement(
    c: ScaledMMLinearLayerConfig,
) -> tuple[bool, str | None]
Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
23
24
25
26
@classmethod
@abstractmethod
def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
    raise NotImplementedError

get_min_capability abstractmethod classmethod 

get_min_capability() -> int
Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
18
19
20
21
@classmethod
@abstractmethod
def get_min_capability(cls) -> int:
    raise NotImplementedError

process_weights_after_loading abstractmethod 

process_weights_after_loading(layer: Module) -> None
Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
45
46
47
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    raise NotImplementedError

ScaledMMLinearLayerConfig dataclass

Source code in vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py 
10
11
12
13
14
@dataclass
class ScaledMMLinearLayerConfig:
    is_channelwise: bool
    is_static_input_scheme: bool
    input_symmetric: bool

input_symmetric instance-attribute 

input_symmetric: bool

is_channelwise instance-attribute 

is_channelwise: bool

is_static_input_scheme instance-attribute 

is_static_input_scheme: bool

__init__ 

__init__(
    is_channelwise: bool,
    is_static_input_scheme: bool,
    input_symmetric: bool,
) -> None