| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE |
| 6 | from vllm.utils.torch_utils import direct_register_custom_op |
| 7 | |
| 8 | |
| 9 | def _quant_dequant_mxfp6( |
| 10 | x: torch.Tensor, |
| 11 | quant_dtype: str, |
| 12 | scale_calculation_mode: str = "even", |
| 13 | ) -> torch.Tensor: |
| 14 | try: |
| 15 | from quark.torch.kernel.hw_emulation.hw_emulation_interface import ( |
| 16 | fake_quantize_fp4_fp6_per_group_with_scale, |
| 17 | ) |
| 18 | from quark.torch.quantization.utils import even_round, reshape_to_blocks |
| 19 | except ImportError as err: |
| 20 | raise ImportError( |
| 21 | "The package `amd-quark` is required to use " |
| 22 | "MX-FP6 models. Please install it with `pip install " |
| 23 | "amd-quark`." |
| 24 | ) from err |
| 25 | |
| 26 | axis = -1 |
| 27 | block_x = reshape_to_blocks(x, OCP_MX_BLOCK_SIZE, axis) |
| 28 | amax, _ = torch.max(torch.abs(block_x), dim=-1, keepdim=True) |
| 29 | amax = amax.squeeze(-1) |
| 30 | |
| 31 | # TODO: there are other rounding strategies supported in quark and in the |
| 32 | # config.json that we do not check for here! |
| 33 | if scale_calculation_mode != "even": |
| 34 | raise NotImplementedError( |
| 35 | f"Scale calculation mode {scale_calculation_mode} is not yet " |
| 36 | "supported in MX-FP6 quantization" |
| 37 | ) |
| 38 | scale = even_round(amax, quant_dtype) |
| 39 | |
| 40 | # Apply dequantize(quantize(x)). |
| 41 | x = fake_quantize_fp4_fp6_per_group_with_scale( |
| 42 | x, |
| 43 | scale.to(x.device), |
| 44 | axis=axis, |
| 45 | group_size=OCP_MX_BLOCK_SIZE, |
| 46 | quant_dtype=quant_dtype, |
| 47 | ) |
| 48 | |
| 49 | return x |
| 50 | |
| 51 | |
| 52 | def _quant_dequant_mxfp6_fake( |
| 53 | x: torch.Tensor, |
| 54 | quant_dtype: str, |
| 55 | scale_calculation_mode: str = "even", |
| 56 | ) -> torch.Tensor: |
| 57 | return torch.empty_like(x) |
| 58 | |
| 59 | |
| 60 | def _dequant_mxfp6( |
| 61 | x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype, quant_dtype: str |
| 62 | ) -> torch.Tensor: |
| 63 | try: |
| 64 | from quark.torch.kernel.hw_emulation.hw_emulation_interface import ( |
| 65 | dequantize_fp4_fp6_per_group, |
| 66 | ) |
| 67 | from quark.torch.utils.pack import create_pack_method |
| 68 | except ImportError as e: |
| 69 | raise ImportError( |
| 70 | "The package `amd-quark` is required to use " |
| 71 | "MX-FP6 models. Please install it with `pip install " |
| 72 | "amd-quark`." |
| 73 | ) from e |
| 74 | |
| 75 | pack_method = create_pack_method(None, dtype=quant_dtype) |
| 76 | unpacked_x = pack_method.unpack(x, reorder=False) |
| 77 | |
| 78 | scale = 2 ** (scale.view(torch.uint8).to(torch.int16) - 127).to(float_dtype) |
| 79 | |
| 80 | # TODO: `dequantize_fp4_fp6_per_group` and `prepare_inputs_per_group` |
| 81 | # always return fp32. |
| 82 | return dequantize_fp4_fp6_per_group( |
| 83 | unpacked_x, |
| 84 | scale, |
| 85 | axis=-1, |
| 86 | group_size=OCP_MX_BLOCK_SIZE, |
| 87 | quant_dtype=quant_dtype, |
| 88 | ).to(float_dtype) |
| 89 | |
| 90 | |
| 91 | def _dequant_mxfp6_fake( |
| 92 | x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype, quant_dtype: str |
| 93 | ) -> torch.Tensor: |
| 94 | assert (x.shape[-1] * 4) % 3 == 0 |
| 95 | return torch.empty( |
| 96 | (*x.shape[:-1], (x.shape[-1] * 4) // 3), dtype=float_dtype, device=x.device |
| 97 | ) |
| 98 | |
| 99 | |
| 100 | # Protect these operations into a torch custom op to avoid errors as |
| 101 | # torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped |
| 102 | # Explanation: Dynamo does not know how to trace the builtin |
| 103 | # `kernel_ext.PyCapsule.dq_uint8_mxfp4_to_half.` This function is either a |
| 104 | # Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python |
| 105 | # extension (perhaps created with pybind). |
| 106 | # TODO: Make sure there is no way to avoid having these functions |
| 107 | # marked as skipped by dynamo. |
| 108 | try: |
| 109 | direct_register_custom_op( |
| 110 | op_name="quant_dequant_mxfp6", |
| 111 | op_func=_quant_dequant_mxfp6, |
| 112 | mutates_args=[], |
| 113 | fake_impl=_quant_dequant_mxfp6_fake, |
| 114 | ) |
| 115 | except AttributeError as error: |
| 116 | raise error |
| 117 | |
| 118 | |
| 119 | # Expose keyword arguments. |
| 120 | def quant_dequant_mxfp6( |
| 121 | x: torch.Tensor, |
| 122 | quant_dtype: str, |
| 123 | scale_calculation_mode: str = "even", |
| 124 | ) -> torch.Tensor: |
| 125 | return torch.ops.vllm.quant_dequant_mxfp6(x, quant_dtype, scale_calculation_mode) |
| 126 | |
| 127 | |
| 128 | try: |
| 129 | direct_register_custom_op( |
| 130 | op_name="dequant_mxfp6", |
| 131 | op_func=_dequant_mxfp6, |
| 132 | mutates_args=[], |
| 133 | fake_impl=_dequant_mxfp6_fake, |
| 134 | ) |
| 135 | except AttributeError as error: |
| 136 | raise error |
| 137 | |
| 138 | |
| 139 | def dequant_mxfp6( |
| 140 | x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype, quant_dtype: str |
| 141 | ) -> torch.Tensor: |
| 142 | return torch.ops.vllm.dequant_mxfp6(x, scale, float_dtype, quant_dtype) |