| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This file is used for /tests and /benchmarks.""" |
| 4 | |
| 5 | from collections.abc import Callable, Mapping |
| 6 | from dataclasses import dataclass |
| 7 | from types import MappingProxyType |
| 8 | from typing import TYPE_CHECKING, ClassVar, Literal, NamedTuple |
| 9 | |
| 10 | import numpy |
| 11 | import torch |
| 12 | from torch import fx |
| 13 | |
| 14 | from vllm.distributed.parallel_state import get_ep_group, get_tp_group |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.scalar_type import ScalarType, scalar_types |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from vllm.model_executor.layers.linear import LinearBase |
| 20 | |
| 21 | FP8_DTYPE = current_platform.fp8_dtype() |
| 22 | FP4_DTYPE = torch.uint8 |
| 23 | MXFP_SCALE_DTYPE = torch.uint8 |
| 24 | INT4_DTYPE = scalar_types.uint4b8 |
| 25 | INT8_DTYPE = scalar_types.uint8b128 |
| 26 | |
| 27 | |
| 28 | def _dtype_abbr(dtype: torch.dtype | ScalarType) -> str: |
| 29 | """Return a stable short name for torch and ScalarType dtypes.""" |
| 30 | if isinstance(dtype, ScalarType): |
| 31 | return str(dtype) |
| 32 | return fx.graph.dtype_abbrs[dtype] |
| 33 | |
| 34 | |
| 35 | def weight_amax( |
| 36 | weight: torch.Tensor, *, dim: int | None = None, keepdim: bool = False |
| 37 | ) -> torch.Tensor: |
| 38 | """``max(|weight|)``, without materializing a full-size ``abs()``.""" |
| 39 | lo, hi = weight.aminmax(dim=dim, keepdim=keepdim) |
| 40 | return torch.maximum(lo.abs(), hi.abs()) |
| 41 | |
| 42 | |
| 43 | def amax_for_tp_weight_quant(amax: torch.Tensor, is_sharded: bool) -> torch.Tensor: |
| 44 | """Reduce a weight ``amax`` over the TP group when the weight is sharded |
| 45 | along a dim the ``amax`` reduces over, so each shard derives the same scale |
| 46 | it would as part of the whole weight. |
| 47 | """ |
| 48 | if is_sharded: |
| 49 | torch.distributed.all_reduce( |
| 50 | amax, |
| 51 | op=torch.distributed.ReduceOp.MAX, |
| 52 | group=get_tp_group().device_group, |
| 53 | ) |
| 54 | return amax |
| 55 | |
| 56 | |
| 57 | def amax_for_moe_weight_quant(amax: torch.Tensor, moe_tp_size: int) -> torch.Tensor: |
| 58 | """Reduce a per-expert weight ``amax`` over the ranks that tensor-shard the |
| 59 | MoE weights. That sharding is flattened over DP x PCP x TP, exactly the EP |
| 60 | group's span. Under EP ``moe_tp_size`` is 1 and each rank owns whole |
| 61 | experts, so no reduction is needed. |
| 62 | """ |
| 63 | if moe_tp_size > 1: |
| 64 | torch.distributed.all_reduce( |
| 65 | amax, |
| 66 | op=torch.distributed.ReduceOp.MAX, |
| 67 | group=get_ep_group().device_group, |
| 68 | ) |
| 69 | return amax |
| 70 | |
| 71 | |
| 72 | def amax_for_moe_activation_quant( |
| 73 | a_scale: torch.Tensor, enable_eplb: bool |
| 74 | ) -> torch.Tensor: |
| 75 | """Reduce a per-expert activation scale to one value shared by all experts. |
| 76 | |
| 77 | Note: when EPLB is enabled and since this quantization scales get |
| 78 | folded into the per-expert dequantization alphas, we can only |
| 79 | ensure that the quant/dequant scales match by having a single |
| 80 | quantization scale shared across all ranks. |
| 81 | """ |
| 82 | a_max = a_scale.max().to(torch.float32) |
| 83 | if enable_eplb: |
| 84 | from vllm.distributed.parallel_state import get_ep_group |
| 85 | |
| 86 | torch.distributed.all_reduce( |
| 87 | a_max, |
| 88 | op=torch.distributed.ReduceOp.MAX, |
| 89 | group=get_ep_group().device_group, |
| 90 | ) |
| 91 | return a_max |
| 92 | |
| 93 | |
| 94 | def get_fp8_min_max() -> tuple[float, float]: |
| 95 | """Get the min and max values for FP8 quantization.""" |
| 96 | # Using the default value (240.0) from pytorch will cause accuracy |
| 97 | # issue on dynamic quantization models on ROCm. Here, use 224.0 for fnuz |
| 98 | # on ROCm platforms that use the torch.float8_e4m3fnuz dtype. |
| 99 | if current_platform.is_fp8_fnuz(): |
| 100 | return -224.0, 224.0 |
| 101 | finfo = torch.finfo(current_platform.fp8_dtype()) |
| 102 | return finfo.min, finfo.max |
| 103 | |
| 104 | |
| 105 | # Use proxy as NamedTuple direct subclasses cannot have static members |
| 106 | class _GroupShape(NamedTuple): |
| 107 | row: int |
| 108 | col: int |
| 109 | |
| 110 | |
| 111 | class GroupShape(_GroupShape): |
| 112 | """This class describes the quantization group shape. |
| 113 | It includes static members for common shapes (per-tensor, per-token). |
| 114 | """ |
| 115 | |
| 116 | # Aliases for common quantization group shapes |
| 117 | PER_TENSOR: ClassVar["GroupShape"] |
| 118 | PER_TOKEN: ClassVar["GroupShape"] |
| 119 | PER_CHANNEL: ClassVar["GroupShape"] |
| 120 | |
| 121 | def is_per_tensor(self) -> bool: |
| 122 | return self.row == -1 and self.col == -1 |
| 123 | |
| 124 | def is_per_token(self) -> bool: |
| 125 | return self.row == 1 and self.col == -1 |
| 126 | |
| 127 | def is_per_channel(self) -> bool: |
| 128 | return self.row == -1 and self.col == 1 |
| 129 | |
| 130 | def is_per_group(self) -> bool: |
| 131 | return self.row == 1 and self.col >= 1 |
| 132 | |
| 133 | |
| 134 | GroupShape.PER_TENSOR = GroupShape(-1, -1) |
| 135 | GroupShape.PER_TOKEN = GroupShape(1, -1) |
| 136 | GroupShape.PER_CHANNEL = GroupShape(-1, 1) |
| 137 | |
| 138 | |
| 139 | @dataclass(frozen=True) |
| 140 | class ScaleDesc: |
| 141 | """Class for describing a single quantization scaling factor. |
| 142 | dtype: data type of the scale |
| 143 | static: static scale if True, dynamic if False |
| 144 | group_shape: group shape of the scale |
| 145 | """ |
| 146 | |
| 147 | dtype: torch.dtype |
| 148 | static: bool |
| 149 | group_shape: GroupShape |
| 150 | |
| 151 | def __str__(self): |
| 152 | d = { |
| 153 | GroupShape.PER_TENSOR: "per_tensor", |
| 154 | GroupShape.PER_TOKEN: "per_token", |
| 155 | GroupShape.PER_CHANNEL: "per_channel", |
| 156 | } |
| 157 | group_shape = d.get(self.group_shape, str(self.group_shape)) |
| 158 | |
| 159 | return ( |
| 160 | f"{_dtype_abbr(self.dtype)}," |
| 161 | f"{'static' if self.static else 'dynamic'},{group_shape}" |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | @dataclass(frozen=True) |
| 166 | class QuantKey: |
| 167 | """Class for identifying the type of quantization. |
| 168 | dtype: quantized data type |
| 169 | scale: scale descriptor |
| 170 | scale2: second-level scale descriptor |
| 171 | symmetric: symmetric if True, asymmetric if False |
| 172 | """ |
| 173 | |
| 174 | # TODO: QuantKey.dtype is assumed to be `torch.dtype` in matcher_utils.py, |
| 175 | # but #37990 introduced e.g. `kInt4Static` that uses a `ScalarType` dtype, |
| 176 | # same for kMxfp6 that does not have a native torch representation. |
| 177 | # Logical dtype and storage (torch) dtype should be separated (see #48949). |
| 178 | dtype: torch.dtype | ScalarType |
| 179 | scale: ScaleDesc |
| 180 | scale2: ScaleDesc | None = None |
| 181 | symmetric: bool = True |
| 182 | |
| 183 | def __str__(self): |
| 184 | scale2_str = f"scale2({self.scale2})," if self.scale2 else "" |
| 185 | return ( |
| 186 | f"QuantKey({_dtype_abbr(self.dtype)}," |
| 187 | f"scale({self.scale}),{scale2_str}" |
| 188 | f"{'a' if not self.symmetric else ''}symmetric)" |
| 189 | ) |
| 190 | |
| 191 | |
| 192 | kStaticTensorScale = ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR) |
| 193 | kFp8StaticTensorSym = QuantKey(FP8_DTYPE, kStaticTensorScale, symmetric=True) |
| 194 | |
| 195 | kDynamicTensorScale = ScaleDesc(torch.float32, False, GroupShape.PER_TENSOR) |
| 196 | kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, kDynamicTensorScale, symmetric=True) |
| 197 | |
| 198 | kStaticTokenScale = ScaleDesc(torch.float32, True, GroupShape.PER_TOKEN) |
| 199 | kFp8StaticTokenSym = QuantKey(FP8_DTYPE, kStaticTokenScale, symmetric=True) |
| 200 | |
| 201 | kStaticChannelScale = ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL) |
| 202 | kFp8StaticChannelSym = QuantKey(FP8_DTYPE, kStaticChannelScale, symmetric=True) |
| 203 | |
| 204 | kDynamicTokenScale = ScaleDesc(torch.float32, False, GroupShape.PER_TOKEN) |
| 205 | kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, kDynamicTokenScale, symmetric=True) |
| 206 | |
| 207 | kNvfp4DynamicGroupScale = ScaleDesc(FP8_DTYPE, False, GroupShape(1, 16)) |
| 208 | kNvfp4Dynamic = QuantKey( |
| 209 | FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kStaticTensorScale |
| 210 | ) |
| 211 | kNvfp4DynamicToken = QuantKey( |
| 212 | FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kDynamicTokenScale |
| 213 | ) |
| 214 | |
| 215 | kNvfp4StaticGroupScale = ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)) |
| 216 | kNvfp4Static = QuantKey( |
| 217 | FP4_DTYPE, scale=kNvfp4StaticGroupScale, scale2=kStaticTensorScale |
| 218 | ) |
| 219 | |
| 220 | kDynamic128Scale = ScaleDesc(torch.float32, False, GroupShape(1, 128)) |
| 221 | kFp8Dynamic128Sym = QuantKey(FP8_DTYPE, kDynamic128Scale, symmetric=True) |
| 222 | |
| 223 | kStatic128BlockScale = ScaleDesc(torch.float32, True, GroupShape(128, 128)) |
| 224 | kFp8Static128BlockSym = QuantKey(FP8_DTYPE, kStatic128BlockScale, symmetric=True) |
| 225 | kFp8Static128BlockE8M0Sym = QuantKey( |
| 226 | FP8_DTYPE, |
| 227 | ScaleDesc(torch.float8_e8m0fnu, True, GroupShape(128, 128)), |
| 228 | symmetric=True, |
| 229 | ) |
| 230 | |
| 231 | kMxfp8StaticScale = ScaleDesc(torch.uint8, True, GroupShape(1, 32)) |
| 232 | kMxfp8Static = QuantKey(FP8_DTYPE, kMxfp8StaticScale, symmetric=True) |
| 233 | |
| 234 | kMxfp8DynamicScale = ScaleDesc(torch.uint8, False, GroupShape(1, 32)) |
| 235 | kMxfp8Dynamic = QuantKey(FP8_DTYPE, kMxfp8DynamicScale, symmetric=True) |
| 236 | |
| 237 | kDynamic64Scale = ScaleDesc(torch.float32, False, GroupShape(1, 64)) |
| 238 | kFp8Dynamic64Sym = QuantKey(FP8_DTYPE, kDynamic64Scale, symmetric=True) |
| 239 | |
| 240 | # TODO (zyongye): Convert all the torch.dtype to scale_dtype |
| 241 | # Changing that requires changing torch compile fused AR+Quant Quant key |
| 242 | # to avoid assertion error |
| 243 | kMxfp4DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) |
| 244 | kMxfp4Dynamic = QuantKey(FP4_DTYPE, scale=kMxfp4DynamicGroupScale, symmetric=True) |
| 245 | |
| 246 | kMxfp8DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) |
| 247 | kMxfp8Dynamic = QuantKey(FP8_DTYPE, scale=kMxfp8DynamicGroupScale, symmetric=True) |
| 248 | |
| 249 | kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) |
| 250 | kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True) |
| 251 | |
| 252 | kMxfp6E3M2StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) |
| 253 | kMxfp6E3M2Static = QuantKey( |
| 254 | scalar_types.float6_e3m2f, scale=kMxfp6E3M2StaticGroupScale, symmetric=True |
| 255 | ) |
| 256 | |
| 257 | kMxfp6E3M2DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) |
| 258 | kMxfp6E3M2Dynamic = QuantKey( |
| 259 | scalar_types.float6_e3m2f, scale=kMxfp6E3M2DynamicGroupScale, symmetric=True |
| 260 | ) |
| 261 | |
| 262 | kMxfp6E2M3StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) |
| 263 | kMxfp6E2M3Static = QuantKey( |
| 264 | scalar_types.float6_e2m3f, scale=kMxfp6E2M3StaticGroupScale, symmetric=True |
| 265 | ) |
| 266 | |
| 267 | kMxfp6E2M3DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) |
| 268 | kMxfp6E2M3Dynamic = QuantKey( |
| 269 | scalar_types.float6_e2m3f, scale=kMxfp6E2M3DynamicGroupScale, symmetric=True |
| 270 | ) |
| 271 | |
| 272 | # TODO: convert this to use SCALAR_TYPE. This is not right. |
| 273 | kInt4StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1)) |
| 274 | kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True) |
| 275 | kInt8StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1)) |
| 276 | kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True) |
| 277 | |
| 278 | kInt4Static32GroupScale = ScaleDesc(torch.float16, True, GroupShape(1, 32)) |
| 279 | kInt4Static32 = QuantKey(INT4_DTYPE, scale=kInt4Static32GroupScale, symmetric=True) |
| 280 | |
| 281 | kInt4StaticAsym = QuantKey( |
| 282 | scalar_types.uint4, scale=kInt4StaticGroupScale, symmetric=False |
| 283 | ) |
| 284 | kInt4Static32Asym = QuantKey( |
| 285 | scalar_types.uint4, scale=kInt4Static32GroupScale, symmetric=False |
| 286 | ) |
| 287 | |
| 288 | kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True) |
| 289 | kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True) |
| 290 | kInt8StaticTensorSym = QuantKey(torch.int8, kStaticTensorScale, symmetric=True) |
| 291 | kInt8DynamicTensorSym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=True) |
| 292 | kInt8DynamicTokenAsym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=False) |
| 293 | kInt8StaticTensorAsym = QuantKey(torch.int8, kStaticTensorScale, symmetric=False) |
| 294 | kInt8DynamicTensorAsym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=False) |
| 295 | |
| 296 | # INT4 W4A8 quantization keys |
| 297 | |
| 298 | # For group-wise quantization (e.g., group_size=128) |
| 299 | # Note: group_size will be specified at runtime, this is a generic group scale |
| 300 | kInt4W4A8StaticGroupScale128 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 128)) |
| 301 | kInt4W4A8StaticGroup128Sym = QuantKey( |
| 302 | torch.int8, kInt4W4A8StaticGroupScale128, symmetric=True |
| 303 | ) |
| 304 | |
| 305 | kInt4W4A8StaticGroupScale64 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 64)) |
| 306 | kInt4W4A8StaticGroup64Sym = QuantKey( |
| 307 | torch.int8, kInt4W4A8StaticGroupScale64, symmetric=True |
| 308 | ) |
| 309 | |
| 310 | kInt4W4A8StaticGroupScale32 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 32)) |
| 311 | kInt4W4A8StaticGroup32Sym = QuantKey( |
| 312 | torch.int8, kInt4W4A8StaticGroupScale32, symmetric=True |
| 313 | ) |
| 314 | |
| 315 | # Generic group-wise with flexible group size (per-token groups) |
| 316 | kInt4W4A8StaticGroupScale = ScaleDesc(torch.bfloat16, True, GroupShape(1, -1)) |
| 317 | kInt4W4A8StaticGroupSym = QuantKey( |
| 318 | torch.int8, kInt4W4A8StaticGroupScale, symmetric=True |
| 319 | ) |
| 320 | |
| 321 | kInt4W4A8StaticChannelSym = QuantKey( |
| 322 | torch.int8, |
| 323 | ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL), |
| 324 | symmetric=True, |
| 325 | ) |
| 326 | |
| 327 | |
| 328 | def create_fp8_quant_key( |
| 329 | static: bool, |
| 330 | group_shape: GroupShape, |
| 331 | symmetric: bool = True, |
| 332 | scale_dtype: torch.dtype = torch.float32, |
| 333 | ) -> QuantKey: |
| 334 | scale_desc = ScaleDesc(scale_dtype, static, group_shape) |
| 335 | return QuantKey(FP8_DTYPE, scale_desc, symmetric=symmetric) |
| 336 | |
| 337 | |
| 338 | # Normalize the group_shape to the full extent for any dims that are -1 |
| 339 | def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape): |
| 340 | # -1 means full extent |
| 341 | return ( |
| 342 | group_shape[0] if group_shape[0] > 0 else x.shape[-2], |
| 343 | group_shape[1] if group_shape[1] > 0 else x.shape[-1], |
| 344 | ) |
| 345 | |
| 346 | |
| 347 | # Useful when treating N-dimensional group scaling as extended numpy-style |
| 348 | # broadcasting in numpy simply stretches dimensions with an extent of 1 to match |
| 349 | # the target shape by repeating the data along that dimension (broadcasting) |
| 350 | # , we extend these semantics to say if the extent of a dimension in the |
| 351 | # source shape is not 1 and does not match the target shape we repeat each |
| 352 | # element along that dimension src_shape[dim] // target_shape[dim] times |
| 353 | # example if we have: |
| 354 | # a = [[1, 2], and target_shape = (2, 4) |
| 355 | # [3, 4]] |
| 356 | # then we would expand a to: |
| 357 | # a = [[1, 1, 2, 2], |
| 358 | # [3, 3, 4, 4]] |
| 359 | # NOTE this function does not explicitly broadcast dimensions |
| 360 | # with an extent of 1, since this can be done implicitly by pytorch |
| 361 | def group_broadcast(t, shape): |
| 362 | for i, s in enumerate(shape): |
| 363 | # If tensor has fewer dimensions than target shape, treat missing |
| 364 | # dimensions as size 1 (standard PyTorch broadcasting behavior) |
| 365 | t_dim_size = t.shape[i] if i < t.ndim else 1 |
| 366 | if t_dim_size != s and t_dim_size != 1: |
| 367 | assert s % t_dim_size == 0 |
| 368 | t = ( |
| 369 | t.unsqueeze(i + 1) |
| 370 | .expand(*t.shape[: i + 1], s // t_dim_size, *t.shape[i + 1 :]) |
| 371 | .flatten(i, i + 1) |
| 372 | ) |
| 373 | return t |
| 374 | |
| 375 | |
| 376 | def prep_scale_for_group_broadcast( |
| 377 | scale: torch.Tensor, |
| 378 | x: torch.Tensor, |
| 379 | group_shape: GroupShape | None, |
| 380 | ) -> torch.Tensor: |
| 381 | """Prepare the input quantization scale for group broadcasting. |
| 382 | |
| 383 | Args: |
| 384 | scale: The scale tensor (scalar or 1D). |
| 385 | x: Target tensor whose shape determines broadcast dimensions. |
| 386 | group_shape: GroupShape to broadcast over. |
| 387 | |
| 388 | Returns: |
| 389 | scale reshaped for correct broadcasting. |
| 390 | |
| 391 | """ |
| 392 | if scale.numel() == 1: |
| 393 | # For per-tensor quant, keep the scale as a scalar (not reshaped to (1, 1)). |
| 394 | # This avoids misclassifying it as channelwise quant in Fp8LinearOp.apply, |
| 395 | # where the "per_tensor_activations" check relies on "x_scale.dim() < 2": |
| 396 | # per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2 |
| 397 | # For all other cases, reshape scalar scales to (1, 1) for broadcasting. |
| 398 | return ( |
| 399 | scale |
| 400 | if group_shape is not None and group_shape.is_per_tensor() |
| 401 | else scale.reshape(1, 1) |
| 402 | ) |
| 403 | if scale.ndim == 1: |
| 404 | assert group_shape is not None, ( |
| 405 | "group_shape must be provided to correctly broadcast 1D scale" |
| 406 | ) |
| 407 | rows, cols = _normalize_quant_group_shape(x, group_shape) |
| 408 | # Determine broadcasting dimension: either rows or columns match group size |
| 409 | if rows == x.shape[-2]: |
| 410 | scale = scale.unsqueeze(-2) |
| 411 | elif cols == x.shape[-1]: |
| 412 | scale = scale.unsqueeze(-1) |
| 413 | else: |
| 414 | raise ValueError( |
| 415 | f"1D scale with shape {scale.shape} cannot be broadcast to x with shape" |
| 416 | f" {x.shape}, group_shape={(rows, cols)}" |
| 417 | ) |
| 418 | return scale |
| 419 | |
| 420 | |
| 421 | # Quantize assuming once scale per group of elements with shape group_shape, |
| 422 | # example group shapes: |
| 423 | # * (-1, -1) for per-tensor quantization |
| 424 | # * (1, -1) for per-row quantization |
| 425 | # * (-1, 1) for per-column quantization |
| 426 | # * (128, 128) for 128x128 deepseek style block quantization |
| 427 | # * (1, 128) for deepseek style activation quantization |
| 428 | # (i.e. per-token-per-group) |
| 429 | def scaled_quantize( |
| 430 | x: torch.Tensor, |
| 431 | group_shape: GroupShape, |
| 432 | quant_dtype: torch.dtype, |
| 433 | compute_dtype: torch.dtype | None = None, |
| 434 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 435 | """Args: |
| 436 | x: Input tensor to quantize |
| 437 | group_shape: Shape of quantization groups |
| 438 | quant_dtype: Target quantized dtype (e.g., torch.float8_e4m3fn) |
| 439 | compute_dtype: Optional dtype for intermediate computations. |
| 440 | If None, uses input dtype. Use torch.float32 for higher precision. |
| 441 | |
| 442 | """ |
| 443 | group_shape = _normalize_quant_group_shape(x, group_shape) |
| 444 | assert quant_dtype.is_floating_point, ( |
| 445 | "currently `scaled_quantize` only supports floating point dtypes " |
| 446 | "but could be extended to support other dtypes" |
| 447 | ) |
| 448 | |
| 449 | finfo = torch.finfo(quant_dtype) |
| 450 | |
| 451 | # Convert to compute dtype if specified |
| 452 | x_compute = x if compute_dtype is None else x.to(compute_dtype) |
| 453 | |
| 454 | # Reshape (M, N) into (BLK_M, BLOCK_SIZE_M, BLK_N, BLOCK_SIZE_N) |
| 455 | assert x.ndim == 2 |
| 456 | assert x.shape[0] % group_shape[0] == 0 and x.shape[1] % group_shape[1] == 0 |
| 457 | blk_m, blk_n = x.shape[0] // group_shape[0], x.shape[1] // group_shape[1] |
| 458 | x_blkd = x_compute.reshape(blk_m, group_shape[0], blk_n, group_shape[1]) |
| 459 | |
| 460 | # Permute to (BLK_M, BLK_N, BLOCK_SIZE_M, BLOCK_SIZE_N) |
| 461 | x_blkd_permd = x_blkd.permute(0, 2, 1, 3) |
| 462 | # Flatten to (BLK_M, BLK_N, BLOCK_SIZE_M * BLOCK_SIZE_N) |
| 463 | x_blkd_permd = x_blkd_permd.flatten(start_dim=2) |
| 464 | |
| 465 | # Compute scales |
| 466 | min_val, max_val = x_blkd_permd.aminmax(dim=-1) |
| 467 | amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) |
| 468 | _, fp8_max = get_fp8_min_max() |
| 469 | scale = fp8_max / amax |
| 470 | |
| 471 | # Apply scale and convert from: |
| 472 | # (BLK_M, BLK_N, BLOCK_SIZE_M * BLOCK_SIZE_N) to (M, N) |
| 473 | x_scl_sat = ( |
| 474 | (x_blkd_permd * scale.unsqueeze(-1)) |
| 475 | .clamp(min=finfo.min, max=finfo.max) |
| 476 | .reshape(blk_m, blk_n, group_shape[0], group_shape[1]) |
| 477 | .permute(0, 2, 1, 3) |
| 478 | .reshape(x.shape) |
| 479 | ) |
| 480 | |
| 481 | return x_scl_sat.to(quant_dtype).contiguous(), scale.float().reciprocal() |
| 482 | |
| 483 | |
| 484 | # inverses `scaled_quantize` |
| 485 | def scaled_dequantize( |
| 486 | x_q: torch.Tensor, |
| 487 | x_s: torch.Tensor, |
| 488 | group_shape: GroupShape | None = None, |
| 489 | out_dtype: torch.dtype = torch.float32, |
| 490 | ) -> torch.Tensor: |
| 491 | x_s = prep_scale_for_group_broadcast(x_s, x_q, group_shape) |
| 492 | if group_shape is not None: |
| 493 | assert x_s.shape[-1] == x_q.shape[-1] // group_shape[1] |
| 494 | assert x_s.shape[-2] == x_q.shape[-2] // group_shape[0] |
| 495 | x_s = group_broadcast(x_s.to(torch.float32), x_q.shape) |
| 496 | return (x_q.to(torch.float32) * x_s).to(out_dtype) |
| 497 | |
| 498 | |
| 499 | def get_attribute_fallback(obj, attributes: list[str]): |
| 500 | for attr in attributes: |
| 501 | if hasattr(obj, attr): |
| 502 | return getattr(obj, attr) |
| 503 | raise AttributeError(f"'{obj}' has no recognized attributes: {attributes}.") |
| 504 | |
| 505 | |
| 506 | def get_and_maybe_dequant_weights( |
| 507 | layer: "LinearBase", out_dtype: torch.dtype = torch.float32 |
| 508 | ): |
| 509 | """Return layer's unquantized weights in [out, in] layout.""" |
| 510 | from vllm.model_executor.layers.linear import UnquantizedLinearMethod |
| 511 | from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod |
| 512 | from vllm.model_executor.layers.quantization.online.fp8 import ( |
| 513 | Fp8PerTensorOnlineLinearMethod, |
| 514 | ) |
| 515 | |
| 516 | # LoRA linear wrappers store quantization metadata on `base_layer`. |
| 517 | # Unwrap here so callers can pass either a raw linear layer or its LoRA |
| 518 | # wrapper without special-casing. |
| 519 | while hasattr(layer, "base_layer") and hasattr(layer.base_layer, "quant_method"): |
| 520 | layer = layer.base_layer |
| 521 | |
| 522 | weight = get_attribute_fallback(layer, ["weight", "qweight", "weight_packed"]) |
| 523 | |
| 524 | # Unquantized layer: just return base weights |
| 525 | if layer.quant_method is None or isinstance( |
| 526 | layer.quant_method, UnquantizedLinearMethod |
| 527 | ): |
| 528 | return weight.to(out_dtype) |
| 529 | |
| 530 | # Simple Fp8 case: rescale with tensor or block weight scales |
| 531 | if ( |
| 532 | isinstance( |
| 533 | layer.quant_method, (Fp8LinearMethod, Fp8PerTensorOnlineLinearMethod) |
| 534 | ) |
| 535 | and not layer.quant_method.use_marlin |
| 536 | # DeepGEMM transforms the scales using `transform_sf_into_required_layout` into |
| 537 | # a layout that is not compatible with `scaled_dequantize`. |
| 538 | and not layer.quant_method.use_deep_gemm |
| 539 | ): |
| 540 | weight_scales = get_attribute_fallback( |
| 541 | layer, ["weight_scale", "weight_scale_inv"] |
| 542 | ) |
| 543 | dequant_weights = scaled_dequantize( |
| 544 | weight, |
| 545 | weight_scales, |
| 546 | group_shape=layer.weight_block_size, |
| 547 | out_dtype=out_dtype, |
| 548 | ) |
| 549 | # per-tensor scaling stores weights in [in, out] layout |
| 550 | if not layer.quant_method.block_quant: |
| 551 | dequant_weights = dequant_weights.T |
| 552 | return dequant_weights |
| 553 | |
| 554 | # NOTE: Most generic base case |
| 555 | # - Call the layer with identity matrix which returns unquantized weights. |
| 556 | # - Must be used with extra care when dealing with static activation quantization: |
| 557 | # quantizing 1.0 may lead to over/underflows |
| 558 | # - Should only be used offline, since it's O(N^3) |
| 559 | assert hasattr(layer, "input_size_per_partition") |
| 560 | eye = torch.eye( |
| 561 | layer.input_size_per_partition, |
| 562 | dtype=out_dtype, |
| 563 | device=weight.device, |
| 564 | ) |
| 565 | dequant_weights = layer.quant_method.apply(layer, eye, bias=None).to(out_dtype) |
| 566 | return dequant_weights.T |
| 567 | |
| 568 | |
| 569 | def pack_quantized_values_into_int32( |
| 570 | w_q: torch.Tensor, wtype: ScalarType, packed_dim: int = 0 |
| 571 | ): |
| 572 | # move dim to pack to the end |
| 573 | perm = (*[i for i in range(len(w_q.shape)) if i != packed_dim], packed_dim) |
| 574 | inv_perm = tuple(perm.index(i) for i in range(len(perm))) |
| 575 | w_q_perm = w_q.permute(perm) |
| 576 | |
| 577 | pack_factor = 32 // wtype.size_bits |
| 578 | mask = (1 << wtype.size_bits) - 1 |
| 579 | |
| 580 | new_shape_perm = list(w_q_perm.shape) |
| 581 | assert w_q_perm.shape[-1] % pack_factor == 0 |
| 582 | new_shape_perm[-1] //= pack_factor |
| 583 | |
| 584 | res = torch.zeros(new_shape_perm, dtype=torch.int32, device=w_q.device) |
| 585 | for i in range(pack_factor): |
| 586 | res |= (w_q_perm[..., i::pack_factor] & mask) << wtype.size_bits * i |
| 587 | |
| 588 | return res.permute(inv_perm) |
| 589 | |
| 590 | |
| 591 | def unpack_quantized_values_into_int32( |
| 592 | w_q: torch.Tensor, wtype: ScalarType, packed_dim: int = 0 |
| 593 | ): |
| 594 | # move dim to pack to the end |
| 595 | perm = (*[i for i in range(len(w_q.shape)) if i != packed_dim], packed_dim) |
| 596 | inv_perm = tuple(perm.index(i) for i in range(len(perm))) |
| 597 | w_q_perm = w_q.permute(perm) |
| 598 | |
| 599 | pack_factor = 32 // wtype.size_bits |
| 600 | mask = (1 << wtype.size_bits) - 1 |
| 601 | |
| 602 | new_shape_perm = list(w_q_perm.shape) |
| 603 | new_shape_perm[-1] *= pack_factor |
| 604 | |
| 605 | res = torch.zeros(new_shape_perm, dtype=torch.int32, device=w_q.device) |
| 606 | for i in range(pack_factor): |
| 607 | res[..., i::pack_factor] = (w_q_perm >> wtype.size_bits * i) & mask |
| 608 | |
| 609 | return res.permute(inv_perm) |
| 610 | |
| 611 | |
| 612 | def is_layer_skipped( |
| 613 | prefix: str, |
| 614 | ignored_layers: list[str], |
| 615 | fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), |
| 616 | *, |
| 617 | match_mode: Literal["exact", "substring", "suffix"] = "exact", |
| 618 | ) -> bool: |
| 619 | def prefix_full_match(prefix: str, ignored_layers: list[str]) -> bool: |
| 620 | return prefix in ignored_layers |
| 621 | |
| 622 | # For case like: ignored_layers = ["self_attn"] |
| 623 | def substr_match(prefix: str, ignored_layers: list[str]) -> bool: |
| 624 | return any(layer in prefix for layer in ignored_layers) |
| 625 | |
| 626 | # Match abbreviated module paths at a component boundary. For example, |
| 627 | # ``b_proj`` matches ``model.layers.0.self_attn.b_proj`` but not |
| 628 | # ``model.layers.0.self_attn.q_b_proj``. |
| 629 | def suffix_match(prefix: str, ignored_layers: list[str]) -> bool: |
| 630 | return any( |
| 631 | prefix == layer or prefix.endswith(f".{layer}") for layer in ignored_layers |
| 632 | ) |
| 633 | |
| 634 | if match_mode == "exact": |
| 635 | match_func = prefix_full_match |
| 636 | elif match_mode == "substring": |
| 637 | match_func = substr_match |
| 638 | elif match_mode == "suffix": |
| 639 | match_func = suffix_match |
| 640 | else: |
| 641 | raise ValueError(f"Unsupported layer skip match mode: {match_mode}") |
| 642 | |
| 643 | # prefix: model.layers.0.self_attn.q_proj |
| 644 | # proj_name: q_proj |
| 645 | proj_name = prefix.split(".")[-1] |
| 646 | |
| 647 | # Fused layers like gate_up_proj or qkv_proj will not be fused |
| 648 | # in the safetensors checkpoint. So, we convert the name |
| 649 | # from the fused version to unfused + check to make sure that |
| 650 | # each shard of the fused layer has the same scheme. |
| 651 | # |
| 652 | # Some checkpoints (e.g. block-FP8 Step-3.5-Flash) already list the |
| 653 | # fused name (e.g. ``self_attn.qkv_proj``) directly in |
| 654 | # ``modules_to_not_convert``. Honor that fused-name match first so |
| 655 | # those layers are still correctly skipped even when a |
| 656 | # ``packed_modules_mapping`` is registered on the model. |
| 657 | if proj_name in fused_mapping and match_func(prefix, ignored_layers): |
| 658 | is_skipped = True |
| 659 | elif proj_name in fused_mapping: |
| 660 | shard_prefixes = [ |
| 661 | prefix.replace(proj_name, shard_proj_name) |
| 662 | for shard_proj_name in fused_mapping[proj_name] |
| 663 | ] |
| 664 | |
| 665 | is_skipped = None |
| 666 | for shard_prefix in shard_prefixes: |
| 667 | is_shard_skipped = match_func(shard_prefix, ignored_layers) |
| 668 | |
| 669 | if is_skipped is None: |
| 670 | is_skipped = is_shard_skipped |
| 671 | elif is_shard_skipped != is_skipped: |
| 672 | raise ValueError( |
| 673 | f"Detected some but not all shards of {prefix} " |
| 674 | "are quantized. All shards of fused layers " |
| 675 | "to have the same precision." |
| 676 | ) |
| 677 | elif "experts" in prefix and match_mode == "exact": |
| 678 | expert_ignore_layers = filter( |
| 679 | lambda layer_name: "experts" in layer_name, ignored_layers |
| 680 | ) |
| 681 | return any(prefix in layer_name for layer_name in expert_ignore_layers) |
| 682 | else: |
| 683 | is_skipped = match_func(prefix, ignored_layers) |
| 684 | |
| 685 | assert is_skipped is not None |
| 686 | return is_skipped |
| 687 | |
| 688 | |
| 689 | def get_pack_factor(num_bits): |
| 690 | assert 32 % num_bits == 0, f"Unsupported num_bits = {num_bits}" |
| 691 | return 32 // num_bits |
| 692 | |
| 693 | |
| 694 | def quantize_weights( |
| 695 | w: torch.Tensor, |
| 696 | quant_type: ScalarType, |
| 697 | group_size: int | None, |
| 698 | zero_points: bool = False, |
| 699 | ref_zero_points_after_scales: bool = False, |
| 700 | ): |
| 701 | assert quant_type.is_integer(), ( |
| 702 | "Floating point quantization may work but has not been tested" |
| 703 | ) |
| 704 | assert not zero_points or group_size is not None, ( |
| 705 | "to have group zero points, group_size must be provided " |
| 706 | "(-1 group_size is channelwise)" |
| 707 | ) |
| 708 | |
| 709 | orig_device = w.device |
| 710 | orig_type = w.dtype |
| 711 | size_k, size_n = w.shape |
| 712 | |
| 713 | assert w.is_floating_point(), "w must be float" |
| 714 | |
| 715 | if group_size == -1: |
| 716 | group_size = size_k |
| 717 | |
| 718 | # Reshape to [groupsize, -1] |
| 719 | if group_size is not None and group_size < size_k: |
| 720 | w = w.reshape((-1, group_size, size_n)) |
| 721 | w = w.permute(1, 0, 2) |
| 722 | w = w.reshape((group_size, -1)) |
| 723 | |
| 724 | # Compute scale for each group |
| 725 | max_val = torch.max(w, 0, keepdim=True).values |
| 726 | min_val = torch.min(w, 0, keepdim=True).values |
| 727 | |
| 728 | max_q_val = quant_type.max() |
| 729 | min_q_val = quant_type.min() |
| 730 | |
| 731 | w_s = torch.Tensor([1.0]).to(w.device) # unscaled case |
| 732 | maybe_w_zp = None |
| 733 | if group_size is not None: |
| 734 | if zero_points: |
| 735 | assert not quant_type.is_signed() and quant_type.max() > 0 |
| 736 | w_s = (max_val - min_val).clamp(min=1e-5) / quant_type.max() |
| 737 | maybe_w_zp = ( |
| 738 | torch.round(torch.abs(min_val / w_s)).clamp(min_q_val, max_q_val).int() |
| 739 | ) |
| 740 | else: |
| 741 | # If the bias is such that there are no possible negative/positive |
| 742 | # values, set the max value to inf to avoid divide by 0 |
| 743 | w_s = torch.max( |
| 744 | abs(max_val / (max_q_val if max_q_val != 0 else torch.inf)), |
| 745 | abs(min_val / (min_q_val if min_q_val != 0 else torch.inf)), |
| 746 | ) |
| 747 | |
| 748 | # Quantize |
| 749 | w_q = torch.round(w / w_s).int() + (maybe_w_zp if zero_points else 0) |
| 750 | w_q = torch.clamp(w_q, min_q_val, max_q_val) |
| 751 | |
| 752 | # Compute ref (dequantized) |
| 753 | # For some kernels (namely Machete) the zero-points are applied after the |
| 754 | # scales are applied, for this case computing the reference in similar way |
| 755 | # allows us to use tighter error tolerances in our unit tests. |
| 756 | if ref_zero_points_after_scales and maybe_w_zp is not None: |
| 757 | w_ref = w_q.to(orig_type) * w_s - maybe_w_zp.to(orig_type) * w_s |
| 758 | else: |
| 759 | w_ref = (w_q - (maybe_w_zp if zero_points else 0)).to(orig_type) * w_s |
| 760 | |
| 761 | if quant_type.has_bias(): |
| 762 | w_q += quant_type.bias |
| 763 | |
| 764 | # Restore original shapes |
| 765 | if group_size is not None and group_size < size_k: |
| 766 | |
| 767 | def reshape_w(w): |
| 768 | w = w.reshape((group_size, -1, size_n)) |
| 769 | w = w.permute(1, 0, 2) |
| 770 | w = w.reshape((size_k, size_n)).contiguous() |
| 771 | return w |
| 772 | |
| 773 | w_q = reshape_w(w_q) |
| 774 | w_ref = reshape_w(w_ref) |
| 775 | w_s = w_s.reshape((-1, size_n)).contiguous() |
| 776 | |
| 777 | if maybe_w_zp is not None: |
| 778 | maybe_w_zp = maybe_w_zp.reshape((-1, size_n)).contiguous() |
| 779 | maybe_w_zp = maybe_w_zp.to(device=orig_device) |
| 780 | |
| 781 | return ( |
| 782 | w_ref.to(device=orig_device), |
| 783 | w_q.to(device=orig_device), |
| 784 | w_s if group_size is not None else None, |
| 785 | maybe_w_zp, |
| 786 | ) |
| 787 | |
| 788 | |
| 789 | SUPPORTED_GPTQ_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128] |
| 790 | SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] |
| 791 | |
| 792 | |
| 793 | def gptq_quantize_weights( |
| 794 | w: torch.Tensor, |
| 795 | quant_type: ScalarType, |
| 796 | group_size: int, |
| 797 | ): |
| 798 | size_k, _ = w.shape |
| 799 | |
| 800 | assert w.is_floating_point(), "w must be float" |
| 801 | assert quant_type in SUPPORTED_GPTQ_QUANT_TYPES, ( |
| 802 | f"Unsupported gptq type = {quant_type}" |
| 803 | ) |
| 804 | assert group_size in SUPPORTED_GROUP_SIZES + [size_k], ( |
| 805 | f"Unsupported groupsize = {group_size}" |
| 806 | ) |
| 807 | |
| 808 | w_ref, w_q, w_s, _ = quantize_weights(w, quant_type, group_size) |
| 809 | return w_ref, w_q, w_s |
| 810 | |
| 811 | |
| 812 | def pack_rows( |
| 813 | q_w: torch.Tensor, |
| 814 | num_bits: int, |
| 815 | size_k: int, |
| 816 | size_n: int, |
| 817 | ): |
| 818 | assert q_w.shape == (size_k, size_n) |
| 819 | |
| 820 | pack_factor = get_pack_factor(num_bits) |
| 821 | assert size_k % pack_factor == 0 |
| 822 | |
| 823 | orig_device = q_w.device |
| 824 | |
| 825 | q_w = q_w.cpu().numpy().astype(numpy.uint32) |
| 826 | |
| 827 | q_res = numpy.zeros((size_k // pack_factor, size_n), dtype=numpy.uint32) |
| 828 | |
| 829 | for i in range(pack_factor): |
| 830 | q_res |= q_w[i::pack_factor, :] << num_bits * i |
| 831 | |
| 832 | q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) |
| 833 | return q_res |
| 834 | |
| 835 | |
| 836 | def pack_cols( |
| 837 | q_w: torch.Tensor, |
| 838 | num_bits: int, |
| 839 | size_k: int, |
| 840 | size_n: int, |
| 841 | ): |
| 842 | assert q_w.shape == (size_k, size_n) |
| 843 | |
| 844 | pack_factor = get_pack_factor(num_bits) |
| 845 | assert size_n % pack_factor == 0 |
| 846 | |
| 847 | orig_device = q_w.device |
| 848 | |
| 849 | q_w = q_w.cpu().numpy().astype(numpy.uint32) |
| 850 | |
| 851 | q_res = numpy.zeros((size_k, size_n // pack_factor), dtype=numpy.uint32) |
| 852 | |
| 853 | for i in range(pack_factor): |
| 854 | q_res |= q_w[:, i::pack_factor] << num_bits * i |
| 855 | |
| 856 | q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) |
| 857 | q_res = q_res.contiguous() |
| 858 | |
| 859 | return q_res |
| 860 | |
| 861 | |
| 862 | def unpack_cols( |
| 863 | packed_q_w: torch.Tensor, |
| 864 | num_bits: int, |
| 865 | size_k: int, |
| 866 | size_n: int, |
| 867 | ): |
| 868 | pack_factor = get_pack_factor(num_bits) |
| 869 | assert size_n % pack_factor == 0 |
| 870 | assert packed_q_w.shape == (size_k, size_n // pack_factor), ( |
| 871 | "packed_q_w.shape = {} size_k = {}, size_n = {} pack_Factor = {}".format( |
| 872 | packed_q_w.shape, size_k, size_n, pack_factor |
| 873 | ) |
| 874 | ) |
| 875 | |
| 876 | orig_device = packed_q_w.device |
| 877 | |
| 878 | packed_q_w_cpu = packed_q_w.cpu().numpy().astype(numpy.uint32) |
| 879 | q_res = numpy.zeros((size_k, size_n), dtype=numpy.uint32) |
| 880 | |
| 881 | mask = (1 << num_bits) - 1 |
| 882 | for i in range(pack_factor): |
| 883 | vals = packed_q_w_cpu & mask |
| 884 | packed_q_w_cpu >>= num_bits |
| 885 | q_res[:, i::pack_factor] = vals |
| 886 | |
| 887 | q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device) |
| 888 | q_res = q_res.contiguous() |
| 889 | |
| 890 | return q_res |
| 891 | |
| 892 | |
| 893 | def gptq_pack( |
| 894 | q_w: torch.Tensor, |
| 895 | num_bits: int, |
| 896 | size_k: int, |
| 897 | size_n: int, |
| 898 | ): |
| 899 | return pack_rows(q_w, num_bits, size_k, size_n) |
| 900 | |
| 901 | |
| 902 | def awq_pack( |
| 903 | q_w: torch.Tensor, |
| 904 | num_bits: int, |
| 905 | size_k: int, |
| 906 | size_n: int, |
| 907 | ): |
| 908 | assert q_w.shape == (size_k, size_n) |
| 909 | |
| 910 | # Interleave column dim (for the dequantize code) and pack it to int32 |
| 911 | if num_bits == 4: |
| 912 | interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) |
| 913 | elif num_bits == 8: |
| 914 | interleave = numpy.array([0, 2, 1, 3]) |
| 915 | else: |
| 916 | raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) |
| 917 | |
| 918 | q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() |
| 919 | q_w = q_w.reshape((-1, size_n)).contiguous() |
| 920 | |
| 921 | return pack_cols(q_w, num_bits, size_k, size_n) |
| 922 | |
| 923 | |
| 924 | def convert_bf16_scales_to_fp8( |
| 925 | quant_fp8: Callable, scales: torch.Tensor |
| 926 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 927 | """Convert a BF16 scale tensor into the pair of (fp8_scales, channel_scales) |
| 928 | expected by W4A8 GEMM kernels. |
| 929 | """ |
| 930 | assert scales.is_contiguous(), ( |
| 931 | f"scale tensor must be contiguous, got {scales.stride()=}" |
| 932 | ) |
| 933 | assert scales.is_cuda, "scales must be on gpu" |
| 934 | |
| 935 | orig_shape = scales.shape |
| 936 | k_groups = orig_shape[-1] |
| 937 | flat_scales = scales.view(-1, k_groups) |
| 938 | |
| 939 | fp8_scales, chan_scales = quant_fp8(flat_scales) |
| 940 | fp8_scales = (fp8_scales.float() / 8.0).to(torch.float8_e4m3fn) |
| 941 | chan_scales *= 8.0 |
| 942 | |
| 943 | # restore original shape |
| 944 | fp8_scales = fp8_scales.view(orig_shape) |
| 945 | chan_scales = chan_scales.view(*orig_shape[:-1], -1) |
| 946 | |
| 947 | return fp8_scales, chan_scales |
| 948 | |
| 949 | |
| 950 | def convert_packed_uint4b8_to_signed_int4_inplace(t: torch.Tensor) -> torch.Tensor: |
| 951 | """Convert int4b8 (packed to int32) to signed int4.""" |
| 952 | assert t.is_cuda, "tensor must be on gpu" |
| 953 | assert t.dtype == torch.int32, f"expected int32 packed weights but got {t.dtype}" |
| 954 | |
| 955 | # loop through the 8 4-bit nibbles in each int32 entry |
| 956 | for i in range(8): |
| 957 | shift = 4 * i |
| 958 | # extract the i-th 4-bit nibble |
| 959 | nib = (t >> shift) & 0xF |
| 960 | # clear the original nibble by masking out |
| 961 | t &= ~(0xF << shift) |
| 962 | # convert int4b8 [0..15] to signed int4 [-8..7] by subtracting 8 |
| 963 | # and update in-place |
| 964 | t |= ((nib - 8) & 0xF) << shift |
| 965 | |
| 966 | return t |