| 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 import _custom_ops as ops |
| 8 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 9 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 10 | GroupShape, |
| 11 | convert_bf16_scales_to_fp8, |
| 12 | convert_packed_uint4b8_to_signed_int4_inplace, |
| 13 | ) |
| 14 | from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.scalar_type import scalar_types |
| 17 | |
| 18 | from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig |
| 19 | |
| 20 | |
| 21 | class CutlassW4A8LinearKernel(MPLinearKernel): |
| 22 | def __init__(self, *args, **kwargs): |
| 23 | super().__init__(*args, **kwargs) |
| 24 | # dynamic per-tok fp8 activation quantization |
| 25 | self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN) |
| 26 | |
| 27 | @classmethod |
| 28 | def get_min_capability(cls) -> int: |
| 29 | return 90 |
| 30 | |
| 31 | @classmethod |
| 32 | def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: |
| 33 | if not current_platform.is_cuda(): |
| 34 | return False, "CUTLASS only supported on CUDA" |
| 35 | |
| 36 | if not current_platform.is_device_capability(90): |
| 37 | return False, "CUTLASS W4A8 requires compute capability of 90 (Hopper)" |
| 38 | |
| 39 | if c.act_type != torch.float8_e4m3fn: |
| 40 | return False, "CUTLASS W4A8 only supports FP8 (e4m3) activations" |
| 41 | |
| 42 | if c.zero_points: |
| 43 | return False, "Zero points not supported by CUTLASS W4A8" |
| 44 | |
| 45 | if c.weight_type != scalar_types.int4: |
| 46 | return ( |
| 47 | False, |
| 48 | f"Quant type ({c.weight_type}) not supported by " |
| 49 | "CUTLASS W4A8, only supported int4", |
| 50 | ) |
| 51 | |
| 52 | if c.group_size != 128: |
| 53 | return False, "Only group_size 128 is supported" |
| 54 | |
| 55 | in_features, out_features = c.partition_weight_shape |
| 56 | if in_features % 128 or out_features % 128: |
| 57 | return ( |
| 58 | False, |
| 59 | f"K and N must be divisible by 128, got {c.partition_weight_shape}", |
| 60 | ) |
| 61 | |
| 62 | if c.out_type != torch.bfloat16: |
| 63 | return ( |
| 64 | False, |
| 65 | f"Only bfloat16 output type currently supportedgot {c.out_type=}", |
| 66 | ) |
| 67 | |
| 68 | return True, None |
| 69 | |
| 70 | # note assumes that |
| 71 | # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0} |
| 72 | # `weight_scale` is: {input_dim = 0, output_dim = 1} |
| 73 | def process_weights_after_loading(self, layer: torch.nn.Module): |
| 74 | def transform_w_q(x): |
| 75 | assert isinstance(x, BasevLLMParameter) |
| 76 | convert_packed_uint4b8_to_signed_int4_inplace(x.data) |
| 77 | torch.accelerator.synchronize() |
| 78 | permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) |
| 79 | x.data = ops.cutlass_encode_and_reorder_int4b(x.data.t().contiguous().t()) |
| 80 | return x |
| 81 | |
| 82 | def transform_w_s(x): |
| 83 | assert isinstance(x, BasevLLMParameter) |
| 84 | permute_param_layout_(x, input_dim=0, output_dim=1) |
| 85 | x.data = x.data.contiguous().to(torch.float8_e4m3fn) |
| 86 | x.data = ops.cutlass_pack_scale_fp8(x.data) |
| 87 | return x |
| 88 | |
| 89 | w_s = getattr(layer, self.w_s_name) |
| 90 | fp8_scales, chan_scales = convert_bf16_scales_to_fp8(self.quant_fp8, w_s.data) |
| 91 | w_s.data = fp8_scales |
| 92 | |
| 93 | # register per-channel scales |
| 94 | layer.register_parameter( |
| 95 | "weight_chan_scale", torch.nn.Parameter(chan_scales, requires_grad=False) |
| 96 | ) |
| 97 | |
| 98 | # Encode/reorder weights and pack scales |
| 99 | self._transform_param(layer, self.w_q_name, transform_w_q) |
| 100 | self._transform_param(layer, self.w_s_name, transform_w_s) |
| 101 | |
| 102 | def apply_weights( |
| 103 | self, |
| 104 | layer: torch.nn.Module, |
| 105 | x: torch.Tensor, |
| 106 | bias: torch.Tensor | None = None, |
| 107 | ) -> torch.Tensor: |
| 108 | c = self.config |
| 109 | w_q, w_s, _ = self._get_weight_params(layer) |
| 110 | w_ch_s = layer.weight_chan_scale |
| 111 | |
| 112 | x_2d = x.reshape(-1, x.shape[-1]) |
| 113 | out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) |
| 114 | |
| 115 | x_2d, act_scales = self.quant_fp8(x_2d) |
| 116 | output = ops.cutlass_w4a8_mm( |
| 117 | a=x_2d, |
| 118 | b_q=w_q, |
| 119 | b_group_scales=w_s, |
| 120 | b_group_size=c.group_size, |
| 121 | a_token_scales=act_scales, |
| 122 | b_channel_scales=w_ch_s, |
| 123 | ) |
| 124 | |
| 125 | if bias is not None: |
| 126 | output.add_(bias) # In-place add |
| 127 | |
| 128 | return output.reshape(out_shape) |