| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from importlib.util import find_spec |
| 5 | from typing import Final |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ |
| 10 | from vllm.scalar_type import scalar_types |
| 11 | |
| 12 | from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig |
| 13 | |
| 14 | _CONCH_SUPPORTED_WEIGHT_TYPES: Final = [ |
| 15 | scalar_types.uint4, |
| 16 | scalar_types.uint8, |
| 17 | scalar_types.uint4b8, |
| 18 | scalar_types.uint8b128, |
| 19 | ] |
| 20 | _CONCH_SUPPORTED_GROUP_SIZES: Final = [-1, 128] |
| 21 | |
| 22 | |
| 23 | class ConchLinearKernel(MPLinearKernel): |
| 24 | @classmethod |
| 25 | def get_min_capability(cls) -> int: |
| 26 | return 80 |
| 27 | |
| 28 | @classmethod |
| 29 | def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: |
| 30 | if c.weight_type not in _CONCH_SUPPORTED_WEIGHT_TYPES: |
| 31 | error_msg = ( |
| 32 | f"Weight type ({c.weight_type}) not supported by " |
| 33 | "ConchLinearKernel, supported types are: " |
| 34 | f"{_CONCH_SUPPORTED_WEIGHT_TYPES}" |
| 35 | ) |
| 36 | return False, error_msg |
| 37 | |
| 38 | if c.group_size not in _CONCH_SUPPORTED_GROUP_SIZES: |
| 39 | error_msg = ( |
| 40 | f"Group size ({c.group_size}) not supported by " |
| 41 | "ConchLinearKernel, supported group sizes are: " |
| 42 | f"{_CONCH_SUPPORTED_GROUP_SIZES}" |
| 43 | ) |
| 44 | return False, error_msg |
| 45 | |
| 46 | if find_spec("conch") is None: |
| 47 | error_msg = ( |
| 48 | "conch-triton-kernels is not installed, please " |
| 49 | "install it via `pip install conch-triton-kernels` " |
| 50 | "and try again!" |
| 51 | ) |
| 52 | return False, error_msg |
| 53 | |
| 54 | return True, None |
| 55 | |
| 56 | # note assumes that |
| 57 | # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0} |
| 58 | # `weight_scale` is: {input_dim = 0, output_dim = 1} |
| 59 | # `weight_zero_point` is: {input_dim = 1, output_dim = 0, packed_dim = 0} |
| 60 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 61 | def transform_w_q(x): |
| 62 | assert isinstance(x, BasevLLMParameter) |
| 63 | permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) |
| 64 | x.data = x.data.contiguous() |
| 65 | return x |
| 66 | |
| 67 | def transform_w_s(x): |
| 68 | assert isinstance(x, BasevLLMParameter) |
| 69 | permute_param_layout_(x, input_dim=0, output_dim=1) |
| 70 | x.data = x.data.contiguous() |
| 71 | return x |
| 72 | |
| 73 | def transform_w_zp(x): |
| 74 | # Zero points are stored PACKED as [N//pack_factor, K//G] |
| 75 | # The Conch kernel expects UNPACKED zeros: [K//G, N] |
| 76 | # We need to unpack and reorder |
| 77 | assert isinstance(x, BasevLLMParameter) |
| 78 | packed = x.data # shape: [N//pack_factor, K//G], dtype: int32 |
| 79 | |
| 80 | # Determine packing based on weight bit width |
| 81 | size_bits = self.config.weight_type.size_bits |
| 82 | pack_factor = 32 // size_bits # 8 for 4-bit, 4 for 8-bit |
| 83 | mask = (1 << size_bits) - 1 # 0xF for 4-bit, 0xFF for 8-bit |
| 84 | |
| 85 | n_packed, k_groups = packed.shape |
| 86 | n_full = n_packed * pack_factor |
| 87 | |
| 88 | # Unpack using vectorized bitwise ops |
| 89 | # shifts = [0, size_bits, 2*size_bits, ...] for each packed position |
| 90 | shifts = torch.arange( |
| 91 | 0, 32, size_bits, dtype=torch.int32, device=packed.device |
| 92 | ) |
| 93 | # packed: [N//pack_factor, K//G] -> [N//pack_factor, K//G, 1] |
| 94 | # shifts: [pack_factor] -> [1, 1, pack_factor] |
| 95 | # Result: [N//pack_factor, K//G, pack_factor] |
| 96 | unpacked = (packed.unsqueeze(-1) >> shifts) & mask |
| 97 | |
| 98 | # Permute to [K//G, N//pack_factor, pack_factor] then reshape to [K//G, N] |
| 99 | unpacked = unpacked.permute(1, 0, 2).reshape(k_groups, n_full) |
| 100 | |
| 101 | x.data = unpacked.to(torch.uint8).contiguous() |
| 102 | |
| 103 | # Update metadata - zeros are no longer packed |
| 104 | if hasattr(x, "_input_dim"): |
| 105 | x._input_dim = 0 |
| 106 | if hasattr(x, "_output_dim"): |
| 107 | x._output_dim = 1 |
| 108 | if hasattr(x, "_packed_factor"): |
| 109 | x._packed_factor = 1 |
| 110 | return x |
| 111 | |
| 112 | self._transform_param(layer, self.w_q_name, transform_w_q) |
| 113 | self._transform_param(layer, self.w_s_name, transform_w_s) |
| 114 | if self.config.zero_points: |
| 115 | self._transform_param(layer, self.w_zp_name, transform_w_zp) |
| 116 | elif self.w_zp_name is not None: |
| 117 | layer.register_parameter(self.w_zp_name, None) |
| 118 | |
| 119 | def apply_weights( |
| 120 | self, |
| 121 | layer: torch.nn.Module, |
| 122 | x: torch.Tensor, |
| 123 | bias: torch.Tensor | None = None, |
| 124 | ) -> torch.Tensor: |
| 125 | from conch.ops.quantization.gemm import mixed_precision_gemm |
| 126 | |
| 127 | w_q, w_s, w_zp = self._get_weight_params(layer) |
| 128 | |
| 129 | # Map channelwise group_size=-1 to the actual input dimension K. |
| 130 | # The conch kernel computes stride_mul = block_k / group_size; |
| 131 | # passing -1 produces a negative stride that reads out-of-bounds |
| 132 | # scale values for all K-blocks after the first. |
| 133 | group_size = self.config.group_size |
| 134 | if group_size == -1: |
| 135 | group_size = x.shape[-1] |
| 136 | |
| 137 | x_2d = x.reshape(-1, x.shape[-1]) |
| 138 | out_shape = x.shape[:-1] + (self.config.partition_weight_shape[1],) |
| 139 | |
| 140 | output = mixed_precision_gemm( |
| 141 | x=x_2d, |
| 142 | w_q_packed=w_q.data, |
| 143 | w_s=w_s.data, |
| 144 | w_zp=w_zp.data if w_zp is not None else None, |
| 145 | weight_size_bits=self.config.weight_type.size_bits, |
| 146 | weight_bias=self.config.weight_type.bias, |
| 147 | group_size=group_size, |
| 148 | ) |
| 149 | |
| 150 | if bias is not None: |
| 151 | output.add_(bias) # In-place add |
| 152 | |
| 153 | return output.reshape(out_shape) |