| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.scalar_type import ScalarType, scalar_types |
| 8 | |
| 9 | MACHETE_PREPACKED_BLOCK_SHAPE = [64, 128] |
| 10 | |
| 11 | |
| 12 | def query_machete_supported_quant_types(zero_points: bool) -> list[ScalarType]: |
| 13 | if zero_points: |
| 14 | return [scalar_types.uint4, scalar_types.uint8] |
| 15 | else: |
| 16 | return [scalar_types.uint4b8, scalar_types.uint8b128] |
| 17 | |
| 18 | |
| 19 | def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]: |
| 20 | """Queries the supported group sizes for Machete based on the activation type. |
| 21 | |
| 22 | Args: |
| 23 | act_type: The activation data type (torch.float16, torch.bfloat16). |
| 24 | |
| 25 | Returns: |
| 26 | A list of supported group sizes. The group size must |
| 27 | be divisible by `TileShapeK = 128 * 8 // num_bits(act_type)`. |
| 28 | -1 indicates per-channel quantization. |
| 29 | |
| 30 | """ |
| 31 | if act_type in [torch.float16, torch.bfloat16]: |
| 32 | return [-1, 64, 128] |
| 33 | else: |
| 34 | return [-1, 128] |
| 35 | |
| 36 | |
| 37 | def check_machete_supports_shape( |
| 38 | in_features: int, out_features: int |
| 39 | ) -> tuple[bool, str | None]: |
| 40 | if in_features % MACHETE_PREPACKED_BLOCK_SHAPE[0] != 0: |
| 41 | return ( |
| 42 | False, |
| 43 | "Input features size must be divisible by " |
| 44 | f"{MACHETE_PREPACKED_BLOCK_SHAPE[0]}", |
| 45 | ) |
| 46 | if out_features % MACHETE_PREPACKED_BLOCK_SHAPE[1] != 0: |
| 47 | return ( |
| 48 | False, |
| 49 | "Output features size must be divisible by " |
| 50 | f"{MACHETE_PREPACKED_BLOCK_SHAPE[1]}", |
| 51 | ) |
| 52 | return True, None |