Skip to content

vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel

Classes:

MPLinearKernel

Bases: ABC

Methods:

Source code in vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py 
 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
class MPLinearKernel(ABC):
    @classmethod
    @abstractmethod
    def get_min_capability(cls) -> int:
        raise NotImplementedError

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

    def __init__(
        self,
        c: MPLinearLayerConfig,
        w_q_param_name: str,
        w_s_param_name: str,
        w_zp_param_name: str | None = None,
    ) -> 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
        if c.zero_points:
            assert w_zp_param_name is not None
        self.w_zp_name: str | None = w_zp_param_name

    def input_quant_key(self) -> QuantKey | None:
        """Mixed-precision kernels take unquantized inputs."""
        return None

    @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 _retire_unused_zp(self, layer: torch.nn.Module) -> None:
        name = self.w_zp_name
        if name is not None and getattr(layer, name, None) is not None:
            # Loaded but unused: keep the param and its data.
            retire_parameter(layer, name)
        self.w_zp_name = None

    def _transform_param(
        self, layer: torch.nn.Module, name: str | None, fn: Callable
    ) -> None:
        if name is not None and getattr(layer, name, None) is not None:
            # `fn` may retag its input; the kept loader param keeps its own dims.
            new_param = fn(_shallow_copy(getattr(layer, name)))
            replace_parameter(
                layer, name, torch.nn.Parameter(new_param.data, requires_grad=False)
            )

    def _get_weight_params(
        self, layer: torch.nn.Module
    ) -> tuple[
        torch.Tensor,  # w_q
        torch.Tensor,  # w_s
        torch.Tensor | None,  # w_zp
    ]:
        return (
            getattr(layer, self.w_q_name),
            getattr(layer, self.w_s_name),
            getattr(layer, self.w_zp_name or "", None),
        )

input_quant_key()

Mixed-precision kernels take unquantized inputs.

Source code in vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py 
61
62
63
def input_quant_key(self) -> QuantKey | None:
    """Mixed-precision kernels take unquantized inputs."""
    return None

_shallow_copy(param)

A new object of param's class sharing its data, with its own copy of the attributes (e.g. dim tags) for a transform to retag.

Source code in vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py 
16
17
18
19
20
21
def _shallow_copy(param: torch.Tensor) -> torch.Tensor:
    """A new object of `param`'s class sharing its data, with its own copy of
    the attributes (e.g. dim tags) for a transform to retag."""
    copy = torch.Tensor._make_subclass(type(param), param.data, param.requires_grad)
    copy.__dict__.update(param.__dict__)
    return copy