| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.linear import LinearMethodBase |
| 9 | |
| 10 | if TYPE_CHECKING: |
| 11 | from .schemes.inc_scheme import INCLinearScheme |
| 12 | |
| 13 | |
| 14 | class INCLinearMethod(LinearMethodBase): |
| 15 | def __init__(self, scheme: "INCLinearScheme") -> None: |
| 16 | self.scheme = scheme |
| 17 | |
| 18 | def create_weights( |
| 19 | self, |
| 20 | layer: torch.nn.Module, |
| 21 | input_size_per_partition: int, |
| 22 | output_partition_sizes: list[int], |
| 23 | input_size: int, |
| 24 | output_size: int, |
| 25 | params_dtype: torch.dtype, |
| 26 | **extra_weight_attrs, |
| 27 | ): |
| 28 | return self.scheme.create_weights( |
| 29 | layer=layer, |
| 30 | input_size_per_partition=input_size_per_partition, |
| 31 | output_partition_sizes=output_partition_sizes, |
| 32 | input_size=input_size, |
| 33 | output_size=output_size, |
| 34 | params_dtype=params_dtype, |
| 35 | **extra_weight_attrs, |
| 36 | ) |
| 37 | |
| 38 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 39 | return self.scheme.process_weights_after_loading(layer) |
| 40 | |
| 41 | def apply( |
| 42 | self, |
| 43 | layer: torch.nn.Module, |
| 44 | x: torch.Tensor, |
| 45 | bias: torch.Tensor | None = None, |
| 46 | ) -> torch.Tensor: |
| 47 | return self.scheme.apply_weights(layer, x, bias) |