| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | import vllm._custom_ops as ops |
| 8 | from vllm.logger import init_logger |
| 9 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 10 | USE_FP32_REDUCE_DEFAULT, |
| 11 | get_marlin_input_dtype, |
| 12 | marlin_moe_padded_intermediate, |
| 13 | marlin_pad_dim, |
| 14 | marlin_pad_qweight, |
| 15 | marlin_pad_scales, |
| 16 | marlin_padded_nk, |
| 17 | marlin_permute_bias, |
| 18 | marlin_permute_scales, |
| 19 | marlin_repacked_nk, |
| 20 | marlin_unpad_output, |
| 21 | should_use_atomic_add_reduce, |
| 22 | ) |
| 23 | from vllm.model_executor.utils import replace_parameter |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.scalar_type import scalar_types |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | |
| 30 | def is_fp8_marlin_supported(): |
| 31 | return current_platform.has_device_capability(75) |
| 32 | |
| 33 | |
| 34 | def fp8_fused_exponent_bias_into_scales(scales): |
| 35 | fp8_exponent = 4 |
| 36 | if scales.dtype == torch.half: |
| 37 | target_exponent = 5 |
| 38 | elif scales.dtype == torch.bfloat16: |
| 39 | target_exponent = 8 |
| 40 | # exponent_bias_fp16 = 2 ** 4 - 2 ** 3 = 8 |
| 41 | # exponent_bias_bf16 = 2 ** 7 - 2 ** 3 = 120 |
| 42 | exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp8_exponent - 1) |
| 43 | s = torch.ones_like(scales) * 2 |
| 44 | s = s**exponent_bias |
| 45 | return scales * s |
| 46 | |
| 47 | |
| 48 | def apply_fp8_marlin_linear( |
| 49 | input: torch.Tensor, |
| 50 | weight: torch.Tensor, |
| 51 | weight_scale: torch.Tensor, |
| 52 | workspace: torch.Tensor | None, |
| 53 | size_n: int, |
| 54 | size_k: int, |
| 55 | bias: torch.Tensor | None, |
| 56 | input_dtype: torch.dtype | None = None, |
| 57 | use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, |
| 58 | ) -> torch.Tensor: |
| 59 | # For GPUs that lack FP8 hardware support, we can leverage the |
| 60 | # Marlin kernel for fast weight-only FP8 quantization |
| 61 | |
| 62 | reshaped_x = input.reshape(-1, input.shape[-1]) |
| 63 | out_shape = input.shape[:-1] + (size_n,) |
| 64 | |
| 65 | padded_n, padded_k = marlin_repacked_nk(weight, num_bits=8) |
| 66 | reshaped_x = marlin_pad_dim(reshaped_x, size_k, padded_k) |
| 67 | |
| 68 | use_atomic_add = should_use_atomic_add_reduce( |
| 69 | m=reshaped_x.size(0), |
| 70 | n=padded_n, |
| 71 | k=padded_k, |
| 72 | device=input.device, |
| 73 | dtype=input.dtype, |
| 74 | ) |
| 75 | |
| 76 | inputs = reshaped_x |
| 77 | a_scales = None |
| 78 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 79 | # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) |
| 80 | raise RuntimeError("Marlin W8A8 is not supported.") |
| 81 | |
| 82 | output = torch.ops.vllm.marlin_gemm( |
| 83 | a=inputs, |
| 84 | b_q_weight=weight, |
| 85 | b_bias=bias, |
| 86 | b_scales=weight_scale, |
| 87 | a_scales=a_scales, |
| 88 | global_scale=None, |
| 89 | b_zeros=None, |
| 90 | workspace=workspace, |
| 91 | b_q_type_id=scalar_types.float8_e4m3fn.id, |
| 92 | size_m=reshaped_x.size(0), |
| 93 | size_n=padded_n, |
| 94 | size_k=padded_k, |
| 95 | use_atomic_add=use_atomic_add, |
| 96 | use_fp32_reduce=use_fp32_reduce, |
| 97 | ) |
| 98 | |
| 99 | output = marlin_unpad_output(output, size_n, padded_n) |
| 100 | return output.reshape(out_shape) |
| 101 | |
| 102 | |
| 103 | def prepare_fp8_layer_for_marlin( |
| 104 | layer: torch.nn.Module, |
| 105 | size_k_first: bool = True, |
| 106 | input_dtype: torch.dtype | None = None, |
| 107 | ) -> None: |
| 108 | logger.warning_once( |
| 109 | "Your GPU does not have native support for FP8 computation but " |
| 110 | "FP8 quantization is being used. Weight-only FP8 compression will " |
| 111 | "be used leveraging the Marlin kernel. This may degrade " |
| 112 | "performance for compute-heavy workloads." |
| 113 | ) |
| 114 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 115 | raise RuntimeError("Marlin W8A8 is not supported.") |
| 116 | |
| 117 | part_size_n = layer.output_size_per_partition |
| 118 | part_size_k = layer.input_size_per_partition |
| 119 | weight_block_size = getattr(layer, "weight_block_size", None) |
| 120 | group_size = -1 if weight_block_size is None else weight_block_size[1] |
| 121 | padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size) |
| 122 | |
| 123 | if size_k_first: |
| 124 | assert layer.weight.shape == (part_size_k, part_size_n) |
| 125 | else: |
| 126 | assert layer.weight.shape == (part_size_n, part_size_k) |
| 127 | |
| 128 | # WEIGHT |
| 129 | # Repack weights to marlin format |
| 130 | qweight = pack_fp8_to_int32(layer.weight, size_k_first) |
| 131 | if not size_k_first: |
| 132 | qweight = qweight.T.contiguous() |
| 133 | qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k) |
| 134 | |
| 135 | marlin_qweight = ops.gptq_marlin_repack( |
| 136 | b_q_weight=qweight, |
| 137 | size_k=padded_k, |
| 138 | size_n=padded_n, |
| 139 | num_bits=8, |
| 140 | ) |
| 141 | replace_parameter(layer, "weight", marlin_qweight) |
| 142 | |
| 143 | # WEIGHT SCALES |
| 144 | # Permute scales |
| 145 | if "weight_scale" in dir(layer): |
| 146 | scales = layer.weight_scale.to(layer.orig_dtype) |
| 147 | elif "weight_scale_inv" in dir(layer): |
| 148 | scales = layer.weight_scale_inv.to(layer.orig_dtype) |
| 149 | |
| 150 | # marlin kernel only support channel-wise and group-wise quantization |
| 151 | # we need to convert the scales |
| 152 | if weight_block_size is None: |
| 153 | logical_widths = getattr(layer, "logical_widths", []) |
| 154 | if scales.nelement() == 1: |
| 155 | # tensor-wise quantization -> channel-wise quantization |
| 156 | # (1, 1) =>(repeat)=> (1, size_n) |
| 157 | scales = scales.view(1, 1).repeat_interleave(part_size_n, 1) |
| 158 | elif scales.nelement() == len(logical_widths): |
| 159 | # tensor-wise quantization with logical_widths -> |
| 160 | # channel-wise quantization |
| 161 | assert sum(logical_widths) == part_size_n, ( |
| 162 | f"Sum of logical_widths ({sum(logical_widths)}) must be equal " |
| 163 | f"to part_size_n ({part_size_n})" |
| 164 | ) |
| 165 | lw_tensor = scales.new_tensor(logical_widths, dtype=torch.int64) |
| 166 | scales = scales.view(1, -1).repeat_interleave(lw_tensor, dim=1) |
| 167 | elif scales.nelement() > 1 and scales.nelement() != part_size_n: |
| 168 | assert part_size_n % scales.nelement() == 0 |
| 169 | s_size = scales.nelement() |
| 170 | # tensor-wise quantization (for gate-up proj) |
| 171 | # -> channel-wise quantization |
| 172 | # (1, s_size) =>(repeat)=> (1, size_n) |
| 173 | scales = scales.view(1, s_size) |
| 174 | scales = scales.repeat_interleave(part_size_n // s_size, 1) |
| 175 | else: |
| 176 | # channel-wise quantization |
| 177 | # (1, size_n) |
| 178 | scales = scales.view(1, part_size_n) |
| 179 | else: |
| 180 | # block-wise quantization -> group-wise quantization |
| 181 | # (size_k // block_size[1], ceil(size_n / block_size[0])) |
| 182 | # =>(repeat)=> (size_k // block_size[1], size_n) |
| 183 | if not size_k_first: |
| 184 | scales = scales.T.contiguous() |
| 185 | block_n = weight_block_size[0] |
| 186 | scales = scales.repeat_interleave(block_n, 1) |
| 187 | # size_n may not divisible by block_size[0] |
| 188 | scales = scales[:, :part_size_n] |
| 189 | |
| 190 | scales = marlin_pad_scales( |
| 191 | scales, part_size_n, part_size_k, padded_n, padded_k, group_size |
| 192 | ) |
| 193 | marlin_scales = marlin_permute_scales( |
| 194 | s=scales, size_k=padded_k, size_n=padded_n, group_size=group_size |
| 195 | ) |
| 196 | if input_dtype != torch.float8_e4m3fn: |
| 197 | marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales) |
| 198 | if hasattr(layer, "weight_scale"): |
| 199 | replace_parameter(layer, "weight_scale", marlin_scales) |
| 200 | elif hasattr(layer, "weight_scale_inv"): |
| 201 | replace_parameter(layer, "weight_scale_inv", marlin_scales) |
| 202 | |
| 203 | if hasattr(layer, "bias") and layer.bias is not None: |
| 204 | assert layer.bias.shape == (part_size_n,) |
| 205 | bias = marlin_permute_bias(marlin_pad_dim(layer.bias, part_size_n, padded_n)) |
| 206 | replace_parameter(layer, "bias", bias) |
| 207 | |
| 208 | |
| 209 | def _moe_pad_shard_rows(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor: |
| 210 | """Zero-pad each gate/up shard of a ``(E, 2 * n, ...)`` tensor to padded_n |
| 211 | rows. FP8 zero decodes to 0.0, so the padded rows contribute nothing.""" |
| 212 | if padded_n == n: |
| 213 | return x |
| 214 | e = x.size(0) |
| 215 | rest = x.shape[2:] |
| 216 | x = x.view(e, 2, n, *rest) |
| 217 | x = torch.nn.functional.pad(x, (0, 0) * len(rest) + (0, padded_n - n)) |
| 218 | return x.reshape(e, 2 * padded_n, *rest) |
| 219 | |
| 220 | |
| 221 | def _moe_pad_last(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor: |
| 222 | """Zero-pad the last dim of a ``(E, ..., n)`` tensor to padded_n.""" |
| 223 | if padded_n == n: |
| 224 | return x |
| 225 | return torch.nn.functional.pad(x, (0, padded_n - n)) |
| 226 | |
| 227 | |
| 228 | def prepare_fp8_moe_layer_for_marlin( |
| 229 | layer: torch.nn.Module, |
| 230 | w13_weight: torch.Tensor, |
| 231 | w2_weight: torch.Tensor, |
| 232 | w13_weight_scale: torch.Tensor, |
| 233 | w2_weight_scale: torch.Tensor, |
| 234 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 235 | """Shuffle weights and scales into marlin format.""" |
| 236 | logger.warning_once( |
| 237 | "Your GPU does not have native support for FP8 computation but " |
| 238 | "FP8 quantization is being used. Weight-only FP8 compression will " |
| 239 | "be used leveraging the Marlin kernel. This may degrade " |
| 240 | "performance for compute-heavy workloads." |
| 241 | ) |
| 242 | input_dtype = get_marlin_input_dtype() |
| 243 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 244 | raise NotImplementedError("Marlin W8A8 is not supported.") |
| 245 | |
| 246 | e = layer.num_experts |
| 247 | k = layer.hidden_size |
| 248 | n = layer.intermediate_size_per_partition |
| 249 | w13_n = w13_weight.size(1) |
| 250 | weight_block_size = getattr(layer, "weight_block_size", None) |
| 251 | group_size = -1 if weight_block_size is None else weight_block_size[1] |
| 252 | |
| 253 | # Pad a tile-misaligned intermediate size to a valid Marlin thread tile. |
| 254 | # FP8 zero decodes to 0.0, so padded weights drop out; the converted scales |
| 255 | # are padded to match below (the padded values are irrelevant). |
| 256 | padded_n = marlin_moe_padded_intermediate(n, group_size) |
| 257 | if padded_n != n: |
| 258 | w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n) |
| 259 | w2_weight = _moe_pad_last(w2_weight, n, padded_n) |
| 260 | |
| 261 | # WEIGHT |
| 262 | # Repack weights to marlin format |
| 263 | def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor: |
| 264 | tensor_list = [] |
| 265 | size_n, size_k = weight.size(1), weight.size(2) |
| 266 | for i in range(e): |
| 267 | qweight = pack_fp8_to_int32(weight[i], size_k_first=False) |
| 268 | qweight = qweight.T.contiguous() |
| 269 | |
| 270 | marlin_qweight = ops.gptq_marlin_repack( |
| 271 | b_q_weight=qweight, size_k=size_k, size_n=size_n, num_bits=8 |
| 272 | ) |
| 273 | tensor_list.append(marlin_qweight) |
| 274 | |
| 275 | return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 276 | |
| 277 | w13_weight = repack_weight("w13", w13_weight) |
| 278 | w2_weight = repack_weight("w2", w2_weight) |
| 279 | |
| 280 | # WEIGHT SCALES |
| 281 | # Permute scales (convert at the original size, then pad to the tile). |
| 282 | def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor: |
| 283 | scales = scales.to(layer.orig_dtype) |
| 284 | tensor_list = [] |
| 285 | if "w13" in name: |
| 286 | size_n, size_k = w13_n, k |
| 287 | else: |
| 288 | size_n, size_k = k, n |
| 289 | |
| 290 | # marlin kernel only support channel-wise and group-wise quantization |
| 291 | # we need to convert the scales |
| 292 | if weight_block_size is None: |
| 293 | if scales.nelement() == e: |
| 294 | # tensor-wise quantization -> channel-wise quantization |
| 295 | # (e, 1, 1) =>(repeat)=> (e, 1, size_n) |
| 296 | scales = scales.view(e, 1, 1).repeat_interleave(size_n, 2) |
| 297 | elif scales.nelement() > e and scales.nelement() != e * size_n: |
| 298 | assert (e * size_n) % scales.nelement() == 0 |
| 299 | s_size = scales.nelement() // e |
| 300 | # tensor-wise quantization (for gate-up proj) |
| 301 | # -> channel-wise quantization |
| 302 | # (e, 1, s_size) =>(repeat)=> (e, 1, size_n) |
| 303 | scales = scales.view(e, 1, s_size) |
| 304 | scales = scales.repeat_interleave(size_n // s_size, 2) |
| 305 | else: |
| 306 | # channel-wise quantization |
| 307 | # (e, 1, size_n) |
| 308 | scales = scales.view(e, 1, size_n) |
| 309 | else: |
| 310 | # block-wise quantization -> group-wise quantization |
| 311 | # (e, size_k // block_size[1], ceil(size_n / block_size[0])) |
| 312 | # =>(repeat)=> (e, size_k // block_size[1], size_n) |
| 313 | scales = scales.permute(0, 2, 1) |
| 314 | block_n = weight_block_size[0] |
| 315 | scales = scales.repeat_interleave(block_n, 2) |
| 316 | # size_n may not divisible by block_size[0] |
| 317 | scales = scales[..., :size_n].contiguous() |
| 318 | |
| 319 | # Pad the converted (E, G, size_n) scales to the padded thread tile. |
| 320 | if padded_n != n: |
| 321 | if "w13" in name: |
| 322 | g = scales.size(1) |
| 323 | scales = scales.view(e, g, 2, n) |
| 324 | scales = torch.nn.functional.pad(scales, (0, padded_n - n)) |
| 325 | scales = scales.reshape(e, g, 2 * padded_n) |
| 326 | size_n = 2 * padded_n |
| 327 | else: |
| 328 | if group_size > 0: |
| 329 | pad_groups = (padded_n - n) // group_size |
| 330 | scales = torch.nn.functional.pad(scales, (0, 0, 0, pad_groups)) |
| 331 | size_k = padded_n |
| 332 | |
| 333 | for i in range(e): |
| 334 | marlin_scales = marlin_permute_scales( |
| 335 | s=scales[i], size_k=size_k, size_n=size_n, group_size=group_size |
| 336 | ) |
| 337 | tensor_list.append(marlin_scales) |
| 338 | |
| 339 | scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 340 | if input_dtype != torch.float8_e4m3fn: |
| 341 | scales = fp8_fused_exponent_bias_into_scales(scales) |
| 342 | return scales |
| 343 | |
| 344 | w13_weight_scale = permute_scales(w13_weight_scale, "w13") |
| 345 | w2_weight_scale = permute_scales(w2_weight_scale, "w2") |
| 346 | |
| 347 | return w13_weight, w2_weight, w13_weight_scale, w2_weight_scale |
| 348 | |
| 349 | |
| 350 | def pack_fp8_to_int32( |
| 351 | fp8_tensor: torch.Tensor, size_k_first: bool = True |
| 352 | ) -> torch.Tensor: |
| 353 | """Repack FP8 weights to gptq format (packed int32 elements).""" |
| 354 | assert fp8_tensor.dtype == torch.float8_e4m3fn |
| 355 | assert fp8_tensor.ndim == 2 |
| 356 | |
| 357 | fp8_tensor = fp8_tensor.T if size_k_first else fp8_tensor |
| 358 | fp8_tensor = fp8_tensor.contiguous() |
| 359 | # fp8_tensor is contiguous and have shape (N, K) now |
| 360 | # with `.view(torch.int32)`, it become (N, K // 4) |
| 361 | int32_tensor = fp8_tensor.view(torch.int32) |
| 362 | return int32_tensor.T.contiguous() if size_k_first else int32_tensor |
| 363 | |
| 364 | |
| 365 | def mxfp8_marlin_process_scales(marlin_scales: torch.Tensor) -> torch.Tensor: |
| 366 | """Reorder scales for e8m0 kernel layout and convert to float8_e8m0fnu.""" |
| 367 | # fit the layout of fp8 dequantization |
| 368 | marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view( |
| 369 | marlin_scales.size(0), -1 |
| 370 | ) |
| 371 | marlin_scales = marlin_scales.to(torch.float8_e8m0fnu) |
| 372 | return marlin_scales |
| 373 | |
| 374 | |
| 375 | def apply_mxfp8_marlin_linear( |
| 376 | input: torch.Tensor, |
| 377 | weight: torch.Tensor, |
| 378 | weight_scale: torch.Tensor, |
| 379 | workspace: torch.Tensor | None, |
| 380 | size_n: int, |
| 381 | size_k: int, |
| 382 | bias: torch.Tensor | None = None, |
| 383 | use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, |
| 384 | ) -> torch.Tensor: |
| 385 | reshaped_x = input.reshape(-1, input.shape[-1]) |
| 386 | out_shape = input.shape[:-1] + (size_n,) |
| 387 | |
| 388 | padded_n, padded_k = marlin_repacked_nk(weight, num_bits=8) |
| 389 | reshaped_x = marlin_pad_dim(reshaped_x, size_k, padded_k) |
| 390 | |
| 391 | use_atomic_add = should_use_atomic_add_reduce( |
| 392 | m=reshaped_x.size(0), |
| 393 | n=padded_n, |
| 394 | k=padded_k, |
| 395 | device=input.device, |
| 396 | dtype=input.dtype, |
| 397 | ) |
| 398 | |
| 399 | output = torch.ops.vllm.marlin_gemm( |
| 400 | a=reshaped_x, |
| 401 | b_q_weight=weight, |
| 402 | b_bias=bias, |
| 403 | b_scales=weight_scale, |
| 404 | a_scales=None, |
| 405 | global_scale=None, |
| 406 | b_zeros=None, |
| 407 | workspace=workspace, |
| 408 | b_q_type_id=scalar_types.float8_e4m3fn.id, |
| 409 | size_m=reshaped_x.size(0), |
| 410 | size_n=padded_n, |
| 411 | size_k=padded_k, |
| 412 | use_atomic_add=use_atomic_add, |
| 413 | use_fp32_reduce=use_fp32_reduce, |
| 414 | ) |
| 415 | |
| 416 | output = marlin_unpad_output(output, size_n, padded_n) |
| 417 | return output.reshape(out_shape) |
| 418 | |
| 419 | |
| 420 | def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: |
| 421 | """Repack MXFP8 weights and scales into Marlin kernel format. |
| 422 | |
| 423 | Expects the layer to have: |
| 424 | - weight: [N, K] float8_e4m3fn |
| 425 | - weight_scale: [N, K//32] uint8 (e8m0 encoded) |
| 426 | - input_size_per_partition / output_size_per_partition |
| 427 | """ |
| 428 | part_size_n = layer.output_size_per_partition |
| 429 | part_size_k = layer.input_size_per_partition |
| 430 | group_size = 32 # MX standard block size |
| 431 | padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size) |
| 432 | |
| 433 | # WEIGHT - repack FP8 weights to Marlin format |
| 434 | qweight = pack_fp8_to_int32(layer.weight, size_k_first=False) |
| 435 | qweight = qweight.T.contiguous() |
| 436 | qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k) |
| 437 | |
| 438 | marlin_qweight = ops.gptq_marlin_repack( |
| 439 | b_q_weight=qweight, |
| 440 | size_k=padded_k, |
| 441 | size_n=padded_n, |
| 442 | num_bits=8, |
| 443 | ) |
| 444 | replace_parameter(layer, "weight", marlin_qweight) |
| 445 | |
| 446 | # WEIGHT SCALES |
| 447 | # Convert uint8 scales -> e8m0fnu -> param_dtype for permutation |
| 448 | # Scales are [N, K//32], need [K//32, N] for marlin_permute_scales |
| 449 | param_dtype = torch.get_default_dtype() |
| 450 | scales = layer.weight_scale.data[:part_size_n, : part_size_k // group_size] |
| 451 | scales = scales.contiguous() |
| 452 | scales = scales.view(torch.float8_e8m0fnu).to(param_dtype) |
| 453 | scales = scales.T.contiguous() |
| 454 | scales = marlin_pad_scales( |
| 455 | scales, part_size_n, part_size_k, padded_n, padded_k, group_size |
| 456 | ) |
| 457 | |
| 458 | # Permute scales to Marlin layout |
| 459 | marlin_scales = marlin_permute_scales( |
| 460 | s=scales, |
| 461 | size_k=padded_k, |
| 462 | size_n=padded_n, |
| 463 | group_size=group_size, |
| 464 | ) |
| 465 | |
| 466 | # Reorder for e8m0 kernel layout and convert back to e8m0fnu |
| 467 | marlin_scales = mxfp8_marlin_process_scales(marlin_scales) |
| 468 | replace_parameter(layer, "weight_scale", marlin_scales) |
| 469 | |
| 470 | # BIAS |
| 471 | if hasattr(layer, "bias") and layer.bias is not None: |
| 472 | assert layer.bias.shape == (part_size_n,) |
| 473 | bias = marlin_permute_bias(marlin_pad_dim(layer.bias, part_size_n, padded_n)) |
| 474 | replace_parameter(layer, "bias", bias) |
| 475 | |
| 476 | |
| 477 | def prepare_mxfp8_moe_layer_for_marlin( |
| 478 | layer: torch.nn.Module, |
| 479 | w13: torch.Tensor, |
| 480 | w2: torch.Tensor, |
| 481 | w13_scale: torch.Tensor, |
| 482 | w2_scale: torch.Tensor, |
| 483 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 484 | """Repack MXFP8 MoE weights and scales into Marlin kernel format. |
| 485 | |
| 486 | Args: |
| 487 | layer: MoE layer. |
| 488 | w13: [E, 2*N, K] float8_e4m3fn weights. |
| 489 | w2: [E, K, N] float8_e4m3fn weights. |
| 490 | w13_scale: [E, 2*N, K//32] uint8 e8m0 scales. |
| 491 | w2_scale: [E, K, N//32] uint8 e8m0 scales. |
| 492 | |
| 493 | Returns: |
| 494 | (w13, w2, w13_scale, w2_scale) in Marlin format. |
| 495 | |
| 496 | """ |
| 497 | group_size = 32 |
| 498 | e = w13.shape[0] |
| 499 | k = w13.shape[2] |
| 500 | n = w2.shape[2] |
| 501 | |
| 502 | # Pad a tile-misaligned intermediate size to a valid Marlin thread tile. |
| 503 | padded_n = marlin_moe_padded_intermediate(n, group_size) |
| 504 | if padded_n != n: |
| 505 | w13 = _moe_pad_shard_rows(w13, n, padded_n) |
| 506 | w13_scale = _moe_pad_shard_rows(w13_scale, n, padded_n) |
| 507 | w2 = _moe_pad_last(w2, n, padded_n) |
| 508 | w2_scale = _moe_pad_last(w2_scale, n // group_size, padded_n // group_size) |
| 509 | n = padded_n |
| 510 | w13_n = w13.shape[1] |
| 511 | |
| 512 | param_dtype = torch.get_default_dtype() |
| 513 | |
| 514 | def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: |
| 515 | if "w13" in name: |
| 516 | size_n, size_k = w13_n, k |
| 517 | else: |
| 518 | size_n, size_k = k, n |
| 519 | |
| 520 | assert weight.shape == (e, size_n, size_k) |
| 521 | |
| 522 | tensor_list = [] |
| 523 | for i in range(e): |
| 524 | qweight = pack_fp8_to_int32(weight[i], size_k_first=False) |
| 525 | qweight = qweight.T.contiguous() |
| 526 | marlin_qweight = ops.gptq_marlin_repack( |
| 527 | b_q_weight=qweight, |
| 528 | size_k=size_k, |
| 529 | size_n=size_n, |
| 530 | num_bits=8, |
| 531 | ) |
| 532 | tensor_list.append(marlin_qweight) |
| 533 | return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 534 | |
| 535 | w13 = repack_weight(w13, "w13") |
| 536 | w2 = repack_weight(w2, "w2") |
| 537 | |
| 538 | def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor: |
| 539 | if "w13" in name: |
| 540 | size_n, size_k = w13_n, k |
| 541 | else: |
| 542 | size_n, size_k = k, n |
| 543 | |
| 544 | tensor_list = [] |
| 545 | for i in range(e): |
| 546 | s = scales[i][:size_n, : size_k // group_size].contiguous() |
| 547 | s = s.view(torch.float8_e8m0fnu).to(param_dtype) |
| 548 | s = s.T.contiguous() |
| 549 | marlin_s = marlin_permute_scales( |
| 550 | s=s, |
| 551 | size_k=size_k, |
| 552 | size_n=size_n, |
| 553 | group_size=group_size, |
| 554 | ) |
| 555 | marlin_s = mxfp8_marlin_process_scales(marlin_s) |
| 556 | tensor_list.append(marlin_s) |
| 557 | return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 558 | |
| 559 | w13_scale = permute_scales(w13_scale, "w13") |
| 560 | w2_scale = permute_scales(w2_scale, "w2") |
| 561 | |
| 562 | return w13, w2, w13_scale, w2_scale |
| 563 | |
| 564 | |
| 565 | def marlin_quant_fp8_torch(weight, group_size, input_dtype=None): |
| 566 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 567 | if is_a_8bit: |
| 568 | assert input_dtype == torch.float8_e4m3fn |
| 569 | |
| 570 | size_n, size_k = weight.shape |
| 571 | if group_size != -1: |
| 572 | scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 448 |
| 573 | repeated_scales = scales.repeat_interleave(group_size, 1) |
| 574 | fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn) |
| 575 | weight_ref = fp8_weight.to(weight.dtype) * repeated_scales |
| 576 | else: |
| 577 | scales = weight.view(size_n, 1, group_size).abs().max(-1)[0] / 448 |
| 578 | repeated_scales = scales.repeat_interleave(size_k, 1) |
| 579 | fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn) |
| 580 | weight_ref = fp8_weight.to(weight.dtype) * repeated_scales |
| 581 | |
| 582 | packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous() |
| 583 | marlin_qweight = ops.gptq_marlin_repack( |
| 584 | b_q_weight=packed_weight, |
| 585 | size_k=size_k, |
| 586 | size_n=size_n, |
| 587 | num_bits=8, |
| 588 | is_a_8bit=is_a_8bit, |
| 589 | ) |
| 590 | |
| 591 | marlin_scales = marlin_permute_scales( |
| 592 | s=scales.T, |
| 593 | size_k=size_k, |
| 594 | size_n=size_n, |
| 595 | group_size=group_size, |
| 596 | is_a_8bit=is_a_8bit, |
| 597 | ) |
| 598 | |
| 599 | marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales) |
| 600 | |
| 601 | return weight_ref.T, marlin_qweight, marlin_scales |