| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.utils.torch_utils import direct_register_custom_op |
| 7 | |
| 8 | # MXFP8 constants |
| 9 | MXFP8_VALUE_DTYPE = torch.float8_e4m3fn |
| 10 | MXFP8_SCALE_DTYPE = torch.uint8 |
| 11 | MXFP8_BLOCK_SIZE = 32 |
| 12 | |
| 13 | |
| 14 | def swizzle_mxfp8_scale(sf: torch.Tensor, M: int, K: int) -> torch.Tensor: |
| 15 | """Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout.""" |
| 16 | scaling_vector_size = MXFP8_BLOCK_SIZE # 32 for MXFP8 |
| 17 | factor = scaling_vector_size * 4 # 128 |
| 18 | |
| 19 | num_m_tiles = (M + 127) // 128 |
| 20 | num_k_tiles = (K + factor - 1) // factor |
| 21 | |
| 22 | m_padded = num_m_tiles * 128 |
| 23 | k_scale_padded = num_k_tiles * 4 |
| 24 | |
| 25 | scale_cols = K // scaling_vector_size |
| 26 | sf_padded = torch.zeros( |
| 27 | (m_padded, k_scale_padded), dtype=sf.dtype, device=sf.device |
| 28 | ) |
| 29 | sf_padded[:M, :scale_cols] = sf |
| 30 | |
| 31 | sf_reshaped = sf_padded.view(num_m_tiles, 4, 32, num_k_tiles, 4) |
| 32 | |
| 33 | sf_swizzled = sf_reshaped.transpose(1, 3) |
| 34 | |
| 35 | return sf_swizzled.contiguous().view(-1) |
| 36 | |
| 37 | |
| 38 | def _mxfp8_e4m3_quantize_torch( |
| 39 | x: torch.Tensor, |
| 40 | is_sf_swizzled_layout: bool = False, |
| 41 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 42 | """Naive MXFP8 quantization. |
| 43 | For each block of 32 elements along the last dimension, compute a |
| 44 | shared e8m0 scale that fits the block-wise amax into the finite |
| 45 | float8_e4m3fn range, and quantize each element to float8_e4m3fn. |
| 46 | |
| 47 | Returns (quantized_values [same shape, fp8], scales uint8). |
| 48 | Scale shape depends on is_sf_swizzled_layout: |
| 49 | False -> [..., K//32] (row-major 2D) |
| 50 | True -> [flat swizzled 1D] |
| 51 | """ |
| 52 | assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0 |
| 53 | orig_shape = x.shape |
| 54 | num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE |
| 55 | |
| 56 | x_fp32 = x.to(torch.float32) |
| 57 | x_blocked = x_fp32.view(*orig_shape[:-1], num_blocks, MXFP8_BLOCK_SIZE) |
| 58 | |
| 59 | amax = x_blocked.abs().amax(dim=-1) |
| 60 | amax = amax.clamp(min=torch.finfo(torch.float32).tiny) |
| 61 | fp8_max = torch.finfo(MXFP8_VALUE_DTYPE).max |
| 62 | scale_biased = torch.ceil(torch.log2(amax / fp8_max)) + 127.0 |
| 63 | scale_biased = scale_biased.clamp(0, 254) |
| 64 | scales_uint8 = scale_biased.to(torch.uint8) |
| 65 | |
| 66 | descale = torch.exp2(scale_biased - 127.0) |
| 67 | x_scaled = x_blocked / descale.unsqueeze(-1) |
| 68 | |
| 69 | x_fp8 = x_scaled.view(orig_shape).to(MXFP8_VALUE_DTYPE) |
| 70 | |
| 71 | if x.ndim == 2: |
| 72 | M, K = x.shape |
| 73 | scales_uint8 = scales_uint8.view(M, -1) |
| 74 | if is_sf_swizzled_layout: |
| 75 | scales_uint8 = swizzle_mxfp8_scale(scales_uint8, M=M, K=K) |
| 76 | elif x.ndim == 3: |
| 77 | B, M, K = x.shape |
| 78 | scales_uint8 = scales_uint8.view(B, M, -1) |
| 79 | if is_sf_swizzled_layout: |
| 80 | swizzled = [] |
| 81 | for i in range(B): |
| 82 | swizzled.append(swizzle_mxfp8_scale(scales_uint8[i], M=M, K=K)) |
| 83 | scales_uint8 = torch.cat(swizzled) |
| 84 | |
| 85 | return x_fp8, scales_uint8 |
| 86 | |
| 87 | |
| 88 | def _mxfp8_quant_triton_kernel(): |
| 89 | """Lazily-built Triton kernel: per-32-block E8M0 scale + FP8-E4M3 quant. |
| 90 | |
| 91 | Fuses what ``_mxfp8_e4m3_quantize_torch`` does in several elementwise passes |
| 92 | into one launch. Each program handles ``[BLOCK_M, 32]`` (one MX block). |
| 93 | """ |
| 94 | from vllm.triton_utils import tl, triton |
| 95 | |
| 96 | @triton.jit |
| 97 | def _kernel( |
| 98 | x_ptr, |
| 99 | xq_ptr, |
| 100 | s_ptr, |
| 101 | M, |
| 102 | K, |
| 103 | sxm, |
| 104 | sxk, |
| 105 | sqm, |
| 106 | sqk, |
| 107 | ssm, |
| 108 | ssk, |
| 109 | BLOCK_M: tl.constexpr, |
| 110 | FP8_MAX: tl.constexpr, |
| 111 | TINY: tl.constexpr, |
| 112 | ): |
| 113 | pid_m = tl.program_id(0) |
| 114 | pid_b = tl.program_id(1) # which 32-element block along K |
| 115 | offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) |
| 116 | offs_k = pid_b * 32 + tl.arange(0, 32) |
| 117 | m_mask = offs_m < M |
| 118 | x = tl.load( |
| 119 | x_ptr + offs_m[:, None] * sxm + offs_k[None, :] * sxk, |
| 120 | mask=m_mask[:, None], |
| 121 | other=0.0, |
| 122 | ).to(tl.float32) |
| 123 | # Mirror _mxfp8_e4m3_quantize_torch: the scale has to put the block amax |
| 124 | # at the top of the e4m3 range rather than at 1.0, or small elements of |
| 125 | # the block end up in the subnormals. |
| 126 | amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY) # [BLOCK_M] |
| 127 | sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0 |
| 128 | sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) |
| 129 | # Scale by the reciprocal rather than dividing by ``exp2(sb - 127)``: |
| 130 | # sb == 0 (an all-zero / denormal block) makes that divisor 2**-127, |
| 131 | # which is subnormal in fp32 and flushes to zero on CDNA, turning the |
| 132 | # block's zeros into 0/0 == NaN. The reciprocal is normal for every |
| 133 | # reachable sb, and both forms are exact powers of two, so nothing |
| 134 | # else changes. |
| 135 | rescale = tl.exp2(127.0 - sb) |
| 136 | xq = (x * rescale[:, None]).to(xq_ptr.dtype.element_ty) |
| 137 | tl.store( |
| 138 | xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk, |
| 139 | xq, |
| 140 | mask=m_mask[:, None], |
| 141 | ) |
| 142 | tl.store(s_ptr + offs_m * ssm + pid_b * ssk, sb.to(tl.uint8), mask=m_mask) |
| 143 | |
| 144 | return _kernel |
| 145 | |
| 146 | |
| 147 | _MXFP8_QUANT_KERNEL = None |
| 148 | |
| 149 | |
| 150 | def _mxfp8_e4m3_quantize_triton( |
| 151 | x: torch.Tensor, |
| 152 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 153 | """Fused 2D MXFP8 quant (non-swizzled, row-major [M, K//32] scales).""" |
| 154 | from vllm.triton_utils import triton |
| 155 | |
| 156 | global _MXFP8_QUANT_KERNEL |
| 157 | if _MXFP8_QUANT_KERNEL is None: |
| 158 | _MXFP8_QUANT_KERNEL = _mxfp8_quant_triton_kernel() |
| 159 | |
| 160 | M, K = x.shape |
| 161 | x = x.contiguous() |
| 162 | xq = torch.empty((M, K), dtype=MXFP8_VALUE_DTYPE, device=x.device) |
| 163 | scales = torch.empty( |
| 164 | (M, K // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device |
| 165 | ) |
| 166 | BLOCK_M = 64 |
| 167 | grid = (triton.cdiv(M, BLOCK_M), K // MXFP8_BLOCK_SIZE) |
| 168 | _MXFP8_QUANT_KERNEL[grid]( |
| 169 | x, |
| 170 | xq, |
| 171 | scales, |
| 172 | M, |
| 173 | K, |
| 174 | x.stride(0), |
| 175 | x.stride(1), |
| 176 | xq.stride(0), |
| 177 | xq.stride(1), |
| 178 | scales.stride(0), |
| 179 | scales.stride(1), |
| 180 | BLOCK_M=BLOCK_M, |
| 181 | FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max), |
| 182 | TINY=float(torch.finfo(torch.float32).tiny), |
| 183 | ) |
| 184 | return xq, scales |
| 185 | |
| 186 | |
| 187 | def _mxfp8_e4m3_quantize_impl( |
| 188 | x: torch.Tensor, |
| 189 | is_sf_swizzled_layout: bool = False, |
| 190 | alignment: int = 0, |
| 191 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 192 | from vllm.platforms import current_platform |
| 193 | from vllm.utils.flashinfer import has_flashinfer |
| 194 | |
| 195 | if current_platform.has_device_capability(100) and has_flashinfer(): |
| 196 | from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize |
| 197 | |
| 198 | x_q, x_scales = flashinfer_mxfp8_quantize( |
| 199 | x, |
| 200 | is_sf_swizzled_layout=is_sf_swizzled_layout, |
| 201 | alignment=alignment if alignment > 0 else 32, |
| 202 | backend="cute-dsl", |
| 203 | ) |
| 204 | if x_scales.ndim == 1 and x.ndim == 2 and not is_sf_swizzled_layout: |
| 205 | x_scales = x_scales.view(x.size(0), -1) |
| 206 | return x_q, x_scales |
| 207 | |
| 208 | # ROCm: a single fused Triton kernel beats the multi-pass torch path for the |
| 209 | # common 2D, non-swizzled activation-quant case (used by the native MX |
| 210 | # linear/MoE). Falls back to torch otherwise (3D weights, swizzled layout). |
| 211 | if ( |
| 212 | current_platform.is_rocm() |
| 213 | and not is_sf_swizzled_layout |
| 214 | and x.ndim == 2 |
| 215 | and x.shape[-1] % MXFP8_BLOCK_SIZE == 0 |
| 216 | ): |
| 217 | return _mxfp8_e4m3_quantize_triton(x) |
| 218 | |
| 219 | return _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout) |
| 220 | |
| 221 | |
| 222 | def mxfp8_e4m3_quantize( |
| 223 | x: torch.Tensor, |
| 224 | is_sf_swizzled_layout: bool = False, |
| 225 | alignment: int = 0, |
| 226 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 227 | return torch.ops.vllm.mxfp8_quantize(x, is_sf_swizzled_layout, alignment) |
| 228 | |
| 229 | |
| 230 | def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: |
| 231 | """Dequantize MXFP8 tensor to BF16.""" |
| 232 | x_float = x.to(torch.float32) |
| 233 | |
| 234 | num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE |
| 235 | x_blocked = x_float.view(*x.shape[:-1], num_blocks, MXFP8_BLOCK_SIZE) |
| 236 | |
| 237 | descale = torch.exp2(scales.to(torch.float32) - 127.0) |
| 238 | |
| 239 | dequantized = x_blocked * descale.unsqueeze(-1) |
| 240 | |
| 241 | dequantized = dequantized.view(*x.shape) |
| 242 | |
| 243 | return dequantized.to(torch.bfloat16) |
| 244 | |
| 245 | |
| 246 | def mxfp8_e4m3_quantize_fake( |
| 247 | x: torch.Tensor, |
| 248 | is_sf_swizzled_layout: bool = False, |
| 249 | alignment: int = 0, |
| 250 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 251 | """Fake implementation for torch.compile tracing.""" |
| 252 | fp_data = torch.empty_like(x, dtype=MXFP8_VALUE_DTYPE) |
| 253 | |
| 254 | block_size = MXFP8_BLOCK_SIZE |
| 255 | |
| 256 | if x.ndim == 2: |
| 257 | M, N = x.shape |
| 258 | K = (N + block_size - 1) // block_size |
| 259 | if is_sf_swizzled_layout: |
| 260 | M_padded = ((M + 127) // 128) * 128 |
| 261 | K_padded = ((K + 3) // 4) * 4 |
| 262 | scales = torch.empty( |
| 263 | M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device |
| 264 | ) |
| 265 | else: |
| 266 | scales = torch.empty((M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device) |
| 267 | elif x.ndim == 3: |
| 268 | B, M, N = x.shape |
| 269 | K = (N + block_size - 1) // block_size |
| 270 | if is_sf_swizzled_layout: |
| 271 | M_padded = ((M + 127) // 128) * 128 |
| 272 | K_padded = ((K + 3) // 4) * 4 |
| 273 | scales = torch.empty( |
| 274 | B * M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device |
| 275 | ) |
| 276 | else: |
| 277 | scales = torch.empty((B, M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device) |
| 278 | else: |
| 279 | scale_shape = list(x.shape) |
| 280 | scale_shape[-1] = (x.shape[-1] + block_size - 1) // block_size |
| 281 | scales = torch.empty(scale_shape, dtype=MXFP8_SCALE_DTYPE, device=x.device) |
| 282 | |
| 283 | return fp_data, scales |
| 284 | |
| 285 | |
| 286 | direct_register_custom_op( |
| 287 | op_name="mxfp8_quantize", |
| 288 | op_func=_mxfp8_e4m3_quantize_impl, |
| 289 | fake_impl=mxfp8_e4m3_quantize_fake, |
| 290 | ) |
| 291 | |
| 292 | |
| 293 | def xpu_mxfp8_quantize( |
| 294 | x: torch.Tensor, dtype: torch.dtype | None = None |
| 295 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 296 | return torch.ops.vllm.xpu_mxfp8_quantize(x, dtype) |