wehub-resource-sync
48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
from vllm.model_executor.layers.linear import LinearMethodBase
|
|
|
|
if TYPE_CHECKING:
|
|
from .schemes.inc_scheme import INCLinearScheme
|
|
|
|
|
|
class INCLinearMethod(LinearMethodBase):
|
|
def __init__(self, scheme: "INCLinearScheme") -> None:
|
|
self.scheme = scheme
|
|
|
|
def create_weights(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
input_size_per_partition: int,
|
|
output_partition_sizes: list[int],
|
|
input_size: int,
|
|
output_size: int,
|
|
params_dtype: torch.dtype,
|
|
**extra_weight_attrs,
|
|
):
|
|
return self.scheme.create_weights(
|
|
layer=layer,
|
|
input_size_per_partition=input_size_per_partition,
|
|
output_partition_sizes=output_partition_sizes,
|
|
input_size=input_size,
|
|
output_size=output_size,
|
|
params_dtype=params_dtype,
|
|
**extra_weight_attrs,
|
|
)
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
return self.scheme.process_weights_after_loading(layer)
|
|
|
|
def apply(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
return self.scheme.apply_weights(layer, x, bias)
|