| 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.platforms import current_platform |
| 9 | |
| 10 | |
| 11 | def cutlass_fp8_supported() -> bool: |
| 12 | if not current_platform.is_cuda(): |
| 13 | return False |
| 14 | |
| 15 | capability_tuple = current_platform.get_device_capability() |
| 16 | capability = -1 if capability_tuple is None else capability_tuple.to_int() |
| 17 | |
| 18 | return ops.cutlass_scaled_mm_supports_fp8(capability) |
| 19 | |
| 20 | |
| 21 | def cutlass_block_fp8_supported() -> bool: |
| 22 | if not current_platform.is_cuda(): |
| 23 | return False |
| 24 | |
| 25 | capability_tuple = current_platform.get_device_capability() |
| 26 | capability = -1 if capability_tuple is None else capability_tuple.to_int() |
| 27 | |
| 28 | return ops.cutlass_scaled_mm_supports_block_fp8(capability) |
| 29 | |
| 30 | |
| 31 | def cutlass_group_gemm_supported() -> bool: |
| 32 | if not current_platform.is_cuda(): |
| 33 | return False |
| 34 | |
| 35 | capability_tuple = current_platform.get_device_capability() |
| 36 | capability = -1 if capability_tuple is None else capability_tuple.to_int() |
| 37 | |
| 38 | return ops.cutlass_group_gemm_supported(capability) |
| 39 | |
| 40 | |
| 41 | CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported() |
| 42 | |
| 43 | |
| 44 | def per_tensor_dequantize( |
| 45 | tensor: torch.Tensor, inv_scale: float | torch.Tensor |
| 46 | ) -> torch.Tensor: |
| 47 | fake_qweight = tensor.to(torch.float16) |
| 48 | dq_weight = fake_qweight * inv_scale |
| 49 | return dq_weight |
| 50 | |
| 51 | |
| 52 | def all_close_1d(x: torch.Tensor) -> bool: |
| 53 | assert len(x.shape) == 1 |
| 54 | return all(torch.allclose(x[0], x[i]) for i in range(x.shape[0])) |
| 55 | |
| 56 | |
| 57 | def convert_to_channelwise( |
| 58 | weight_scale: torch.Tensor, logical_widths: list[int] |
| 59 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 60 | # Create channelwise buffer |
| 61 | weight_scale_channel = torch.empty( |
| 62 | (sum(logical_widths), 1), dtype=torch.float32, device=weight_scale.device |
| 63 | ) |
| 64 | |
| 65 | # Expand each scale to match the size of each logical matrix. |
| 66 | start = 0 |
| 67 | for idx, logical_width in enumerate(logical_widths): |
| 68 | end = start + logical_width |
| 69 | weight_scale_channel[start:end, :] = weight_scale[idx] |
| 70 | start = end |
| 71 | |
| 72 | return weight_scale_channel |
| 73 | |
| 74 | |
| 75 | def requantize_with_max_scale( |
| 76 | weight: torch.Tensor, weight_scale: torch.Tensor, logical_widths: list[int] |
| 77 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 78 | # Max scale to be used for requanitzation. |
| 79 | max_w_scale = weight_scale.max() |
| 80 | |
| 81 | # QKV / MLP is fused in the on disk checkpoint if any of the |
| 82 | # weight scales are still set to the default since we initialize |
| 83 | # N weight scales for N shards but we only load 1 weight scale |
| 84 | # from disk in this case. Skip requantization in this case (since) |
| 85 | # we already are quantized with the single scale. |
| 86 | # * Sample Model: nm-testing/Phi-3-mini-128k-instruct-FP8 |
| 87 | # |
| 88 | # Extra note: upon weight reloading weight_scale.ndim == 0 |
| 89 | unfused_module_in_checkpoint = ( |
| 90 | weight_scale.ndim != 0 |
| 91 | and weight_scale[-1] > torch.finfo(torch.float8_e4m3fn).min |
| 92 | ) |
| 93 | |
| 94 | # If unfused checkpoint, need requanize with the single scale. |
| 95 | if unfused_module_in_checkpoint: |
| 96 | start = 0 |
| 97 | for idx, logical_width in enumerate(logical_widths): |
| 98 | # Skip any component with zero width. |
| 99 | if logical_width == 0: |
| 100 | continue |
| 101 | end = start + logical_width |
| 102 | weight_dq = per_tensor_dequantize(weight[start:end, :], weight_scale[idx]) |
| 103 | weight[start:end, :], _ = ops.scaled_fp8_quant(weight_dq, max_w_scale) |
| 104 | start = end |
| 105 | |
| 106 | return max_w_scale, weight |
| 107 | |
| 108 | |
| 109 | def normalize_e4m3fn_to_e4m3fnuz( |
| 110 | weight: torch.Tensor, |
| 111 | weight_scale: torch.Tensor, |
| 112 | input_scale: torch.Tensor | None = None, |
| 113 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| 114 | assert weight.dtype == torch.float8_e4m3fn |
| 115 | # The bits pattern 10000000(-128) represents zero in e4m3fn |
| 116 | # but NaN in e4m3fnuz. So here we set it to 0. |
| 117 | # https://onnx.ai/onnx/technical/float8.html |
| 118 | weight_as_int8 = weight.view(torch.int8) |
| 119 | ROCM_FP8_NAN_AS_INT = -128 |
| 120 | weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0 |
| 121 | weight = weight_as_int8.view(torch.float8_e4m3fnuz) |
| 122 | |
| 123 | # For the same bits representation, e4m3fnuz value is half of |
| 124 | # the e4m3fn value, so we should double the scaling factor to |
| 125 | # get the same dequantized value. |
| 126 | # https://onnx.ai/onnx/technical/float8.html |
| 127 | weight_scale = weight_scale * 2.0 |
| 128 | if input_scale is not None: |
| 129 | input_scale = input_scale * 2.0 |
| 130 | return weight, weight_scale, input_scale |