| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | BFLOAT16_EXP_BIAS = 127 |
| 6 | BFLOAT16_MANTISSA_BITS = 7 |
| 7 | BFLOAT16_EXP_BITS = 8 |
| 8 | |
| 9 | FLOAT16_EXP_BIAS = 15 |
| 10 | FLOAT16_MANTISSA_BITS = 10 |
| 11 | FLOAT16_EXP_BITS = 5 |
| 12 | |
| 13 | FLOAT8_E8M0_MAX_EXP = 127 |
| 14 | FLOAT4_EXP_BIAS = 1 |
| 15 | FLOAT4_MANTISSA_BITS = 1 |
| 16 | |
| 17 | FLOAT16_VAL_TO_ADD = 1 << (FLOAT16_MANTISSA_BITS - FLOAT4_MANTISSA_BITS - 1) |
| 18 | FLOAT16_SIGN_EXPONENT_MASK = ( |
| 19 | (1 << (FLOAT16_EXP_BITS + 1)) - 1 |
| 20 | ) << FLOAT16_MANTISSA_BITS |
| 21 | |
| 22 | BFLOAT16_VAL_TO_ADD = 1 << (BFLOAT16_MANTISSA_BITS - FLOAT4_MANTISSA_BITS - 1) |
| 23 | BFLOAT16_SIGN_EXPONENT_MASK = ( |
| 24 | (1 << (BFLOAT16_EXP_BITS + 1)) - 1 |
| 25 | ) << BFLOAT16_MANTISSA_BITS |
| 26 | |
| 27 | |
| 28 | def e8m0_to_half(scale, half_dtype: torch.dtype): |
| 29 | assert scale.dtype == torch.uint8 |
| 30 | |
| 31 | scale_exp = scale.to(torch.int16) - 127 |
| 32 | |
| 33 | # This can be implemented with bitwise operations in a proper kernel. |
| 34 | scale_half = 2.0 ** (scale_exp.to(torch.float)) |
| 35 | |
| 36 | return scale_half.to(half_dtype) |
| 37 | |
| 38 | |
| 39 | def upcast_fp4_to_fp16_or_bf16( |
| 40 | val, float_dtype: torch.dtype, half_exp_bias: int, half_mantissa_bits: int |
| 41 | ): |
| 42 | assert val.dtype == torch.uint8 |
| 43 | |
| 44 | unpacked = torch.zeros( |
| 45 | *val.shape[:-1], val.shape[-1] * 2, dtype=torch.uint8, device=val.device |
| 46 | ) |
| 47 | unpacked[..., 1::2] = (val >> 4) & 0x0F # Extract high 4 bits. |
| 48 | unpacked[..., ::2] = val & 0x0F # Extract low 4 bits. |
| 49 | |
| 50 | # Takes one float4 values represented as b0000xxxx, |
| 51 | # and converts it to the corresponding float16 value. |
| 52 | |
| 53 | sign = unpacked >> 3 |
| 54 | |
| 55 | exp = (unpacked >> 1) & 3 |
| 56 | new_mantissa = unpacked & 1 |
| 57 | |
| 58 | # if exp == 0 and new_mantissa == 0: |
| 59 | # new_exp = 0 |
| 60 | # else: |
| 61 | # new_exp = exp - FLOAT4_EXP_BIAS + FLOAT16_EXP_BIAS |
| 62 | |
| 63 | # int8_t works with float16, but may overflow with bfloat16. |
| 64 | new_exp = exp - FLOAT4_EXP_BIAS + half_exp_bias |
| 65 | |
| 66 | # Cast b0000 to 0. in fp16/bf16. |
| 67 | new_exp = new_exp * torch.logical_or(exp > 0, new_mantissa > 0) |
| 68 | |
| 69 | # Cast b0001 to 0.5 in fp16/bf16. |
| 70 | new_mantissa = torch.logical_and(new_mantissa, exp > 0) |
| 71 | |
| 72 | new_mantissa = new_mantissa.to(torch.int32) |
| 73 | new_exp = new_exp.to(torch.int32) |
| 74 | sign = sign.to(torch.int32) |
| 75 | |
| 76 | qdq_val = ( |
| 77 | (sign << 15) |
| 78 | + (new_exp << half_mantissa_bits) |
| 79 | + (new_mantissa << (half_mantissa_bits - 1)) |
| 80 | ) |
| 81 | |
| 82 | assert qdq_val.max() <= 65535 |
| 83 | assert qdq_val.min() >= 0 |
| 84 | qdq_val = qdq_val.to(torch.uint16) |
| 85 | |
| 86 | result = qdq_val.view(float_dtype) |
| 87 | |
| 88 | return result |
| 89 | |
| 90 | |
| 91 | def dq_mxfp4_torch( |
| 92 | x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype |
| 93 | ) -> torch.Tensor: |
| 94 | assert x.dtype == torch.uint8 |
| 95 | assert scale.dtype == torch.uint8 |
| 96 | |
| 97 | if float_dtype == torch.float16: |
| 98 | half_exp_bias = FLOAT16_EXP_BIAS |
| 99 | half_mantissa_bits = FLOAT16_MANTISSA_BITS |
| 100 | elif float_dtype == torch.bfloat16: |
| 101 | half_exp_bias = BFLOAT16_EXP_BIAS |
| 102 | half_mantissa_bits = BFLOAT16_MANTISSA_BITS |
| 103 | |
| 104 | scale_half = e8m0_to_half(scale, half_dtype=float_dtype) |
| 105 | |
| 106 | x_half = upcast_fp4_to_fp16_or_bf16( |
| 107 | x, |
| 108 | float_dtype=float_dtype, |
| 109 | half_exp_bias=half_exp_bias, |
| 110 | half_mantissa_bits=half_mantissa_bits, |
| 111 | ) |
| 112 | |
| 113 | x_half = x_half.reshape(*x_half.shape[:-1], -1, 32) |
| 114 | x_half = x_half * scale_half[..., None] |
| 115 | x_half = x_half.reshape(*x_half.shape[:-2], -1) |
| 116 | |
| 117 | return x_half |
| 118 | |
| 119 | |
| 120 | def fp16_to_fp4_simulate( |
| 121 | val, half_mantissa_bits: int, half_exp_bits: int, half_exp_bias: int |
| 122 | ): |
| 123 | # Casts an fp16/bf16 input to the restricted values of float4_e2m1, |
| 124 | # that is to say [0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, |
| 125 | # -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]. |
| 126 | |
| 127 | float_type = val.dtype |
| 128 | |
| 129 | # "rshift_cuda" not implemented for 'UInt16' |
| 130 | val_view = val.view(torch.int16) # .to(torch.int32) |
| 131 | |
| 132 | exp = val_view >> half_mantissa_bits |
| 133 | exp = exp & ((1 << half_exp_bits) - 1) |
| 134 | |
| 135 | exp = exp.view(torch.uint16).to(torch.int32) |
| 136 | |
| 137 | sign = (val_view >> (half_mantissa_bits + half_exp_bits)) & 1 |
| 138 | |
| 139 | mantissa_last = (val_view >> (half_mantissa_bits - 1)) & 1 |
| 140 | |
| 141 | exp_unbias = exp - half_exp_bias |
| 142 | new_exp = exp_unbias + FLOAT4_EXP_BIAS |
| 143 | |
| 144 | exp_shift = (new_exp <= 0) * (1 - new_exp) |
| 145 | |
| 146 | # Typically 9. |
| 147 | # Take the min to prevent overflow on `uint16_t half`. This is the case for |
| 148 | # very small values, correctly mapped to `round_close`. |
| 149 | tail_bits = half_mantissa_bits - FLOAT4_MANTISSA_BITS + exp_shift |
| 150 | tail_bits[tail_bits >= 16] = 16 |
| 151 | |
| 152 | mantissa_plus_one = val_view & ((1 << (half_mantissa_bits + 1)) - 1) |
| 153 | |
| 154 | half = 1 << (tail_bits - 1) |
| 155 | |
| 156 | tail = mantissa_plus_one & ((1 << tail_bits) - 1) |
| 157 | |
| 158 | round_close = tail < half # round towards 0 |
| 159 | round_away = tail > half # round away from 0 |
| 160 | tie = tail == half |
| 161 | |
| 162 | new_mantissa_close = torch.zeros(val.shape, device=val.device, dtype=torch.bool) |
| 163 | new_exp_close = torch.zeros(val.shape, device=val.device, dtype=torch.uint16) |
| 164 | |
| 165 | new_mantissa_away = torch.zeros(val.shape, device=val.device, dtype=torch.bool) |
| 166 | new_exp_away = torch.zeros(val.shape, device=val.device, dtype=torch.uint16) |
| 167 | |
| 168 | new_exp_tie = torch.zeros(val.shape, device=val.device, dtype=torch.uint16) |
| 169 | |
| 170 | # 1. round down |
| 171 | # if new_exp == 0: # case [0.5, 0.749999] |
| 172 | # new_mantissa = 0 |
| 173 | # elif new_exp < 0: # case [0, 0.24999] |
| 174 | # new_mantissa = 0 |
| 175 | # else: |
| 176 | # new_mantissa = mantissa_last |
| 177 | |
| 178 | new_mantissa_close = (new_exp > 0) * mantissa_last |
| 179 | new_exp_close = exp |
| 180 | |
| 181 | # # 2. round up |
| 182 | # if new_exp <= 0: # case [0.250001, 0.499999] and [0.75001, 0.99999] |
| 183 | # new_mantissa = 0 |
| 184 | # new_exp += 1 |
| 185 | # elif mantissa_last == 0: |
| 186 | # new_mantissa = 1 |
| 187 | # else: |
| 188 | # new_mantissa = 0 |
| 189 | # new_exp += 1 |
| 190 | |
| 191 | new_mantissa_away = torch.logical_and(new_exp > 0, mantissa_last == 0) |
| 192 | new_exp_away = exp + torch.logical_or(new_exp <= 0, mantissa_last == 1) |
| 193 | |
| 194 | # # 3. tie |
| 195 | # 0.25 -> 0. (handled by `exp > (half_exp_bias - 2)`) |
| 196 | # 0.75 -> 1. |
| 197 | # 1.25 -> 1. |
| 198 | # 1.75 -> 2. |
| 199 | # 2.5 -> 2. |
| 200 | # 3.5 -> 4. |
| 201 | # 5. -> 4. |
| 202 | new_exp_tie = (exp > (half_exp_bias - 2)) * (exp + (mantissa_last == 1)) |
| 203 | |
| 204 | # Gather round up, round down and tie. |
| 205 | new_exp = ( |
| 206 | round_away * new_exp_away + round_close * new_exp_close + tie * new_exp_tie |
| 207 | ) |
| 208 | |
| 209 | new_mantissa = round_away * new_mantissa_away + round_close * new_mantissa_close |
| 210 | |
| 211 | # if new_exp > 3: |
| 212 | # new_mantissa = 1 |
| 213 | new_mantissa = new_mantissa + (new_exp > (2 + half_exp_bias)) * (new_mantissa == 0) |
| 214 | |
| 215 | # Clamp the exponent to acceptable values. |
| 216 | new_exp = (new_exp >= (half_exp_bias - 2)) * torch.clamp( |
| 217 | new_exp, half_exp_bias - 2, half_exp_bias + 2 |
| 218 | ) |
| 219 | |
| 220 | sign = sign.to(torch.int32) |
| 221 | new_mantissa = new_mantissa.to(torch.int32) |
| 222 | |
| 223 | qdq_val = ( |
| 224 | (sign << 15) |
| 225 | + (new_exp << half_mantissa_bits) |
| 226 | + (new_mantissa << (half_mantissa_bits - 1)) |
| 227 | ) |
| 228 | |
| 229 | assert qdq_val.max() <= 65535 |
| 230 | assert qdq_val.min() >= 0 |
| 231 | assert qdq_val.dtype == torch.int32 |
| 232 | qdq_val = qdq_val.to(torch.uint16) |
| 233 | |
| 234 | result = qdq_val.view(float_type) |
| 235 | return result |
| 236 | |
| 237 | |
| 238 | def qdq_mxfp4_torch( |
| 239 | x: torch.Tensor, scale_calculation_mode: str = "even" |
| 240 | ) -> torch.Tensor: |
| 241 | half_dtype = x.dtype |
| 242 | |
| 243 | if half_dtype == torch.float16: |
| 244 | half_mantissa_bits = FLOAT16_MANTISSA_BITS |
| 245 | half_exp_bits = FLOAT16_EXP_BITS |
| 246 | half_exp_bias = FLOAT16_EXP_BIAS |
| 247 | val_to_add = FLOAT16_VAL_TO_ADD |
| 248 | sign_exponent_mask = FLOAT16_SIGN_EXPONENT_MASK |
| 249 | elif half_dtype == torch.bfloat16: |
| 250 | half_mantissa_bits = BFLOAT16_MANTISSA_BITS |
| 251 | half_exp_bits = BFLOAT16_EXP_BITS |
| 252 | half_exp_bias = BFLOAT16_EXP_BIAS |
| 253 | val_to_add = BFLOAT16_VAL_TO_ADD |
| 254 | sign_exponent_mask = BFLOAT16_SIGN_EXPONENT_MASK |
| 255 | else: |
| 256 | raise ValueError("not implemented") |
| 257 | |
| 258 | x = x.reshape(*x.shape[:-1], -1, 32) |
| 259 | |
| 260 | block_max = torch.max(torch.abs(x), dim=-1).values |
| 261 | |
| 262 | block_max = block_max.view(torch.uint16).to(torch.int32) |
| 263 | |
| 264 | block_max_uint = torch.bitwise_and(block_max + val_to_add, sign_exponent_mask) |
| 265 | |
| 266 | assert block_max_uint.max() <= 65535 |
| 267 | assert block_max_uint.min() >= 0 |
| 268 | assert block_max_uint.dtype == torch.int32 |
| 269 | block_max_uint = block_max_uint.to(torch.uint16) |
| 270 | |
| 271 | block_max = block_max_uint.view(half_dtype) |
| 272 | |
| 273 | scale_exp = ( |
| 274 | FLOAT8_E8M0_MAX_EXP + torch.floor(torch.log2(block_max)).to(torch.int32) - 2 |
| 275 | ) |
| 276 | |
| 277 | scale_exp = torch.clamp(scale_exp, 0, 2 * FLOAT8_E8M0_MAX_EXP) |
| 278 | |
| 279 | scale = 2.0 ** (scale_exp - FLOAT8_E8M0_MAX_EXP) |
| 280 | scale = scale.to(half_dtype) |
| 281 | |
| 282 | x = x / scale[..., None] |
| 283 | |
| 284 | x_fp4 = fp16_to_fp4_simulate( |
| 285 | x, |
| 286 | half_exp_bits=half_exp_bits, |
| 287 | half_mantissa_bits=half_mantissa_bits, |
| 288 | half_exp_bias=half_exp_bias, |
| 289 | ) |
| 290 | |
| 291 | x_fp4 = x_fp4 * scale[..., None] |
| 292 | return x_fp4.reshape(*x_fp4.shape[:-2], -1) |