| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.logger import init_logger |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.triton_utils import tl, triton |
| 10 | from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels |
| 11 | from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer |
| 12 | |
| 13 | logger = init_logger(__name__) |
| 14 | |
| 15 | |
| 16 | def should_use_cdna4_mx_scale_swizzle() -> bool: |
| 17 | """Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950. |
| 18 | |
| 19 | CDNA4 swizzle requires BLOCK_K%256==0; at TP>=4 the A8W4 dispatch |
| 20 | picks BK<256 tiles for the smaller per-rank shapes, so swizzle must |
| 21 | be off. Used by both the weight-load swizzle in `_swizzle_mxfp4` and |
| 22 | the kernel-argument gate in `aiter_mxfp4_w4a8_moe`; they must agree. |
| 23 | """ |
| 24 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 25 | from vllm.platforms.rocm import on_gfx950 |
| 26 | |
| 27 | return on_gfx950() and get_tensor_model_parallel_world_size() <= 2 |
| 28 | |
| 29 | |
| 30 | def weight_mx_scale(precision_config): |
| 31 | """mxfp4 weight scale from a triton_kernels PrecisionConfig: 3.8 exposes it |
| 32 | as b_mx_scale, 3.5.1/3.6 as weight_scale. |
| 33 | """ |
| 34 | if get_triton_kernels_version() == "3.8": |
| 35 | return precision_config.b_mx_scale |
| 36 | return precision_config.weight_scale |
| 37 | |
| 38 | |
| 39 | def mx_scale_kwargs(scale): |
| 40 | """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size, |
| 41 | 3.5.1/3.6 use weight_scale. |
| 42 | """ |
| 43 | if get_triton_kernels_version() == "3.8": |
| 44 | return {"b_mx_scale": scale, "b_microblock_size": 32} |
| 45 | return {"weight_scale": scale} |
| 46 | |
| 47 | |
| 48 | def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): |
| 49 | """Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel.""" |
| 50 | assert has_triton_kernels() |
| 51 | if get_triton_kernels_version() == "3.8": |
| 52 | # 3.8: matmul_ogs_details -> matmul_details |
| 53 | import triton_kernels.matmul_details.opt_flags as opt_flags |
| 54 | else: |
| 55 | import triton_kernels.matmul_ogs_details.opt_flags as opt_flags |
| 56 | from triton_kernels.numerics import InFlexData |
| 57 | from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor |
| 58 | from triton_kernels.tensor_details import layout |
| 59 | from triton_kernels.tensor_details.layout import StridedLayout |
| 60 | |
| 61 | value_layout_opts: dict[str, Any] = {} |
| 62 | scale_layout_opts: dict[str, Any] = {} |
| 63 | |
| 64 | if ( |
| 65 | current_platform.is_cuda() |
| 66 | and current_platform.is_device_capability(90) |
| 67 | and not is_torch_equal_or_newer("2.8.1") |
| 68 | ): |
| 69 | logger.warning_once( |
| 70 | "Mxfp4 on hopper is running on torch < 2.8.1, " |
| 71 | "this cause swizling to be disabled, which may " |
| 72 | "cause performance degradation. Please upgrade to torch nightly" |
| 73 | ) |
| 74 | value_layout = StridedLayout |
| 75 | scale_layout = StridedLayout |
| 76 | elif current_platform.is_rocm() and get_triton_kernels_version() == "3.8": |
| 77 | # StridedLayout(-2) is triton_kernels' default CDNA mxfp4 weight layout |
| 78 | # (make_default_matmul_mxfp4_w_layout); 3.8 wants layout instances. |
| 79 | from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout |
| 80 | |
| 81 | value_layout = StridedLayout(-2) |
| 82 | scale_layout = ( |
| 83 | CDNA4MXScaleLayout() |
| 84 | if should_use_cdna4_mx_scale_swizzle() |
| 85 | else StridedLayout() |
| 86 | ) |
| 87 | elif current_platform.is_rocm(): |
| 88 | value_layout = StridedLayout |
| 89 | if should_use_cdna4_mx_scale_swizzle(): |
| 90 | try: |
| 91 | # triton < 3.6 |
| 92 | from triton_kernels.tensor_details.layout import GFX950MXScaleLayout |
| 93 | |
| 94 | scale_layout = GFX950MXScaleLayout |
| 95 | except ImportError: |
| 96 | # triton >= 3.6 |
| 97 | from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout |
| 98 | |
| 99 | scale_layout = CDNA4MXScaleLayout |
| 100 | else: |
| 101 | scale_layout = StridedLayout |
| 102 | else: |
| 103 | value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout( |
| 104 | mx_axis=1 |
| 105 | ) |
| 106 | scale_layout, scale_layout_opts = ( |
| 107 | layout.make_default_matmul_mxfp4_w_scale_layout( |
| 108 | mx_axis=1, num_warps=num_warps |
| 109 | ) |
| 110 | ) |
| 111 | if current_platform.is_cuda(): |
| 112 | if current_platform.is_device_capability(90): |
| 113 | constraints = { |
| 114 | "split_k": 1, |
| 115 | } |
| 116 | opt_flags.update_opt_flags_constraints(constraints) |
| 117 | # Patches #47303: pad K (num scale groups) to 0 mod 4 |
| 118 | # TODO: Remove once we upgrade to Triton 3.8.0+ kernels |
| 119 | if scale.numel() > 0: |
| 120 | K = scale.shape[-1] |
| 121 | pad_k = -K % 4 |
| 122 | scale = torch.nn.functional.pad(scale, (0, pad_k)) |
| 123 | elif current_platform.is_device_capability_family(100): |
| 124 | constraints = { |
| 125 | "is_persistent": True, |
| 126 | "epilogue_subtile": 1, |
| 127 | } |
| 128 | opt_flags.update_opt_flags_constraints(constraints) |
| 129 | # transpose the tensor so that the quantization axis is on dim1 |
| 130 | quant_tensor = quant_tensor.transpose(-2, -1) |
| 131 | scale = scale.transpose(-2, -1) |
| 132 | quant_tensor = convert_layout( |
| 133 | wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts |
| 134 | ) |
| 135 | scale = convert_layout(wrap_torch_tensor(scale), scale_layout, **scale_layout_opts) |
| 136 | return quant_tensor, InFlexData(), scale |
| 137 | |
| 138 | |
| 139 | def _dequant_mxfp4( |
| 140 | x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype |
| 141 | ) -> torch.Tensor: |
| 142 | try: |
| 143 | from quark.torch.kernel import mx |
| 144 | except ImportError as err: |
| 145 | raise ImportError( |
| 146 | "The package `amd-quark` is required to use " |
| 147 | "MX-FP4 models. Please install it with `pip install " |
| 148 | "amd-quark`." |
| 149 | ) from err |
| 150 | |
| 151 | return mx.dq_mxfp4(x, scale, float_dtype) |
| 152 | |
| 153 | |
| 154 | def _dequant_mxfp4_fake( |
| 155 | x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype |
| 156 | ) -> torch.Tensor: |
| 157 | return torch.empty( |
| 158 | (*x.shape[:-1], x.shape[-1] * 2), dtype=float_dtype, device=x.device |
| 159 | ) |
| 160 | |
| 161 | |
| 162 | def _quant_dequant_mxfp4( |
| 163 | x: torch.Tensor, scale_calculation_mode: str = "even" |
| 164 | ) -> torch.Tensor: |
| 165 | try: |
| 166 | from quark.torch.kernel import mx |
| 167 | except ImportError as err: |
| 168 | raise ImportError( |
| 169 | "The package `amd-quark` is required to use " |
| 170 | "MX-FP4 models. Please install it with `pip install " |
| 171 | "amd-quark`." |
| 172 | ) from err |
| 173 | |
| 174 | return mx.qdq_mxfp4(x, scale_calculation_mode) |
| 175 | |
| 176 | |
| 177 | def _quant_dequant_mxfp4_fake( |
| 178 | x: torch.Tensor, scale_calculation_mode: str = "even" |
| 179 | ) -> torch.Tensor: |
| 180 | return torch.empty_like(x) |
| 181 | |
| 182 | |
| 183 | # Protect these operations into a torch custom op to avoid errors as |
| 184 | # torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped |
| 185 | # Explanation: Dynamo does not know how to trace the builtin |
| 186 | # `kernel_ext.PyCapsule.dq_uint8_mxfp4_to_half.` This function is either a |
| 187 | # Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python |
| 188 | # extension (perhaps created with pybind). |
| 189 | # TODO: Make sure there is no way to avoid having these functions |
| 190 | # marked as skipped by dynamo. |
| 191 | try: |
| 192 | direct_register_custom_op( |
| 193 | op_name="dequant_mxfp4", |
| 194 | op_func=_dequant_mxfp4, |
| 195 | fake_impl=_dequant_mxfp4_fake, |
| 196 | ) |
| 197 | dequant_mxfp4 = torch.ops.vllm.dequant_mxfp4 |
| 198 | except AttributeError as error: |
| 199 | raise error |
| 200 | |
| 201 | try: |
| 202 | direct_register_custom_op( |
| 203 | op_name="quant_dequant_mxfp4", |
| 204 | op_func=_quant_dequant_mxfp4, |
| 205 | fake_impl=_quant_dequant_mxfp4_fake, |
| 206 | ) |
| 207 | quant_dequant_mxfp4 = torch.ops.vllm.quant_dequant_mxfp4 |
| 208 | except AttributeError as error: |
| 209 | raise error |
| 210 | |
| 211 | |
| 212 | def xpu_mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 213 | return torch.ops.vllm.xpu_mxfp4_quantize(x) |
| 214 | |
| 215 | |
| 216 | def permute_shape(shape: tuple[int, ...], permute_order: list[int]) -> tuple[int, ...]: |
| 217 | return tuple(shape[i] for i in permute_order) |
| 218 | |
| 219 | |
| 220 | @triton.jit |
| 221 | def _compute_quant_and_scale(src_tensor, valid_src_mask): |
| 222 | BLOCK_SIZE_OUT_DIM: tl.constexpr = src_tensor.shape[0] |
| 223 | BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1] |
| 224 | BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // 32 |
| 225 | |
| 226 | # Explicit cast to fp32 since most ops are not supported on bfloat16. We |
| 227 | # avoid needless conversions to and from bf16 |
| 228 | f32_tensor = src_tensor.to(tl.float32) |
| 229 | abs_tensor = tl.abs(f32_tensor) |
| 230 | abs_tensor = tl.where( |
| 231 | valid_src_mask, abs_tensor, -1.0 |
| 232 | ) # Don't consider padding tensors in scale computation |
| 233 | abs_tensor = tl.reshape( |
| 234 | abs_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32] |
| 235 | ) |
| 236 | max_val = tl.max(abs_tensor, axis=2, keep_dims=True) |
| 237 | # Round-to-nearest-even block scale (OCP MX spec), matching hardware/aiter |
| 238 | # MXFP4 quantizers: compute 2 ** (floor(log2(rounding(max_abs(v)))) - max_exp). |
| 239 | max_val = max_val.to(tl.int32, bitcast=True) |
| 240 | max_val = (max_val + 0x200000).to(tl.uint32, bitcast=True) & 0x7F800000 |
| 241 | max_val = max_val.to(tl.float32, bitcast=True) |
| 242 | # Add epsilon to prevent log2(0) = -inf when max_val is zero (all-zero input block) |
| 243 | eps = tl.where(max_val == 0.0, 2 ** (-126), 0.0) |
| 244 | scale_e8m0_unbiased = tl.log2(max_val + eps).floor() - 2 |
| 245 | scale_e8m0_unbiased = tl.clamp(scale_e8m0_unbiased, min=-127, max=127) |
| 246 | dequant_scale_rounded = tl.exp2(scale_e8m0_unbiased) |
| 247 | dequant_scale_exponent = dequant_scale_rounded.to(tl.uint32, bitcast=True) |
| 248 | |
| 249 | dequant_scale_rounded = dequant_scale_exponent.to(tl.float32, bitcast=True) |
| 250 | quant_scale = tl.where(dequant_scale_rounded == 0, 0, 1.0 / dequant_scale_rounded) |
| 251 | |
| 252 | f32_tensor = tl.reshape( |
| 253 | f32_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32] |
| 254 | ) |
| 255 | quant_tensor = f32_tensor * quant_scale |
| 256 | |
| 257 | # Reshape the tensors after scaling |
| 258 | quant_tensor = quant_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM]) |
| 259 | # Set the invalid portions of the tensor to 0. This will ensure that any |
| 260 | # padding tensors are 0 in the mx format. |
| 261 | quant_tensor = tl.where(valid_src_mask, quant_tensor, 0) |
| 262 | dequant_scale_exponent = dequant_scale_exponent.reshape( |
| 263 | [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE] |
| 264 | ) |
| 265 | |
| 266 | # First, we simply extract the exponent part of the scales and store the result |
| 267 | dequant_scale_exponent = (dequant_scale_exponent >> 23).to(tl.uint8) |
| 268 | # Now we must convert the tensor to packed e2m1 (mxfp4). |
| 269 | quant_tensor = quant_tensor.to(tl.uint32, bitcast=True) |
| 270 | signs = quant_tensor & 0x80000000 |
| 271 | exponents = (quant_tensor >> 23) & 0xFF |
| 272 | mantissas = quant_tensor & 0x7FFFFF |
| 273 | |
| 274 | # 0.25 <= x < 0.75 maps to 0.5, a denormal number |
| 275 | E8_BIAS = 127 |
| 276 | E2_BIAS = 1 |
| 277 | # Move implicit bit 1 at the beginning to mantissa for denormals. Bits |
| 278 | # shifted out are OR-reduced into the LSB as a sticky bit, so the |
| 279 | # round-to-even logic below can tell an exact tie from a value that is |
| 280 | # merely close to one. |
| 281 | adjusted_exponents = tl.core.sub(E8_BIAS, exponents + 1, sanitize_overflow=False) |
| 282 | folded_mantissas = 0x400000 | (mantissas >> 1) |
| 283 | denormal_mantissas = folded_mantissas >> adjusted_exponents |
| 284 | denormal_mantissas |= ( |
| 285 | (denormal_mantissas << adjusted_exponents) != folded_mantissas |
| 286 | ).to(tl.uint32) |
| 287 | mantissas = tl.where(exponents < E8_BIAS, denormal_mantissas, mantissas) |
| 288 | |
| 289 | # For normal numbers, we change the bias from 127 to 1, and for subnormals, |
| 290 | # we keep exponent as 0. |
| 291 | exponents = tl.maximum(exponents, E8_BIAS - E2_BIAS) - (E8_BIAS - E2_BIAS) |
| 292 | |
| 293 | # Combine sign, exponent, and mantissa, rounding to nearest with ties |
| 294 | # broken to even (OCP MX spec, matching hardware/aiter MXFP4 |
| 295 | # quantizers) via the standard "add a rounding bias, then truncate" |
| 296 | # trick: the bias's LSB is the kept mantissa bit, so an exact tie |
| 297 | # rounds to whichever candidate is even, while any lower (sticky) bit |
| 298 | # being set always rounds up. |
| 299 | rounding_bias = ((mantissas >> 22) & 1) + 0x1FFFFF |
| 300 | rounded_mantissa = (mantissas + rounding_bias) >> 22 # in {0, 1, 2} |
| 301 | e2m1_tmp = tl.minimum((exponents << 1) + rounded_mantissa, 0x7) |
| 302 | e2m1_value = ((signs >> 28) | e2m1_tmp).to(tl.uint8) |
| 303 | |
| 304 | e2m1_value = tl.reshape( |
| 305 | e2m1_value, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2] |
| 306 | ) |
| 307 | evens, odds = tl.split(e2m1_value) |
| 308 | out_tensor = evens | (odds << 4) |
| 309 | |
| 310 | return out_tensor, dequant_scale_exponent |
| 311 | |
| 312 | |
| 313 | @triton.jit |
| 314 | def _downcast_to_mxfp( |
| 315 | mx_tensor_ptr, |
| 316 | stride_mxt_outer, |
| 317 | stride_mxt_quant: tl.constexpr, |
| 318 | mx_scale_ptr, |
| 319 | stride_mx_scale_outer, |
| 320 | stride_mx_scale_quant, |
| 321 | src_ptr, |
| 322 | stride_src_outer, |
| 323 | stride_src_quant, |
| 324 | outer_dim, |
| 325 | quant_dim, |
| 326 | BLOCK_SIZE_OUT_DIM: tl.constexpr, |
| 327 | BLOCK_SIZE_QUANT_DIM: tl.constexpr, |
| 328 | ): |
| 329 | tl.static_assert( |
| 330 | stride_mxt_quant == 1, f"Output stride, {stride_mxt_quant=} must be 1." |
| 331 | ) |
| 332 | tl.static_assert( |
| 333 | BLOCK_SIZE_QUANT_DIM % 32 == 0, |
| 334 | f"{BLOCK_SIZE_QUANT_DIM=} must be a multiple of 32", |
| 335 | ) |
| 336 | |
| 337 | # uint8 signifies two fp4 e2m1 values packed into a single byte |
| 338 | tl.static_assert( |
| 339 | mx_tensor_ptr.dtype.element_ty == tl.uint8, |
| 340 | f"Invalid {mx_tensor_ptr.dtype.element_ty=}. Must be uint8.", |
| 341 | ) |
| 342 | |
| 343 | src_dtype: tl.constexpr = src_ptr.dtype.element_ty |
| 344 | tl.static_assert( |
| 345 | mx_scale_ptr.dtype.element_ty == tl.uint8, |
| 346 | f"{mx_scale_ptr.dtype.element_ty=} must be uint8", |
| 347 | ) |
| 348 | tl.static_assert( |
| 349 | (src_dtype == tl.bfloat16) or (src_dtype == tl.float16), |
| 350 | f"{src_dtype=} must be bfloat16 or float16", |
| 351 | ) |
| 352 | |
| 353 | outer_block = tl.program_id(0).to(tl.int64) |
| 354 | quant_block = tl.program_id(1).to(tl.int64) |
| 355 | |
| 356 | BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 32 |
| 357 | BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 2 |
| 358 | |
| 359 | start_src_quant = quant_block * BLOCK_SIZE_QUANT_DIM |
| 360 | start_mx_scale_quant = quant_block * BLOCK_SIZE_QUANT_MX_SCALE |
| 361 | start_mx_quant = quant_block * BLOCK_SIZE_QUANT_MX_TENSOR |
| 362 | start_out = outer_block * BLOCK_SIZE_OUT_DIM |
| 363 | |
| 364 | src_ptr += start_src_quant * stride_src_quant + start_out * stride_src_outer |
| 365 | mx_scale_ptr += ( |
| 366 | start_mx_scale_quant * stride_mx_scale_quant + start_out * stride_mx_scale_outer |
| 367 | ) |
| 368 | mx_tensor_ptr += start_mx_quant * stride_mxt_quant + start_out * stride_mxt_outer |
| 369 | |
| 370 | offs_src_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :].to(tl.int64) |
| 371 | offs_mxt_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :].to(tl.int64) |
| 372 | offs_scale_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :].to(tl.int64) |
| 373 | offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None].to(tl.int64) |
| 374 | |
| 375 | mask_src_quant = start_src_quant + offs_src_quant < quant_dim |
| 376 | mask_n = start_out + offs_outer < outer_dim |
| 377 | full_mask_src = mask_src_quant and mask_n |
| 378 | |
| 379 | mask_mxt_quant = start_mx_quant + offs_mxt_quant < tl.cdiv(quant_dim, 2) |
| 380 | full_mask_mxt = mask_mxt_quant and mask_n |
| 381 | |
| 382 | scale_mask_k = start_mx_scale_quant + offs_scale_quant < tl.cdiv(quant_dim, 32) |
| 383 | full_scale_mask = scale_mask_k and mask_n |
| 384 | |
| 385 | src_tensor_offsets = ( |
| 386 | offs_src_quant * stride_src_quant + offs_outer * stride_src_outer |
| 387 | ) |
| 388 | mx_scale_offsets = ( |
| 389 | offs_scale_quant * stride_mx_scale_quant + offs_outer * stride_mx_scale_outer |
| 390 | ) |
| 391 | mx_tensor_offsets = ( |
| 392 | offs_mxt_quant * stride_mxt_quant + offs_outer * stride_mxt_outer |
| 393 | ) |
| 394 | src_tensor = tl.load(src_ptr + src_tensor_offsets, mask=full_mask_src) |
| 395 | |
| 396 | out_tensor, scale_tensor = _compute_quant_and_scale(src_tensor, full_mask_src) |
| 397 | |
| 398 | tl.store(mx_scale_ptr + mx_scale_offsets, scale_tensor, mask=full_scale_mask) |
| 399 | tl.store(mx_tensor_ptr + mx_tensor_offsets, out_tensor, mask=full_mask_mxt) |
| 400 | |
| 401 | |
| 402 | def downcast_to_mxfp( |
| 403 | src_tensor: torch.Tensor, |
| 404 | axis: int, |
| 405 | out_quant_tensor: torch.Tensor | None = None, |
| 406 | out_scale: torch.Tensor | None = None, |
| 407 | BLOCK_OUT_DIM: int = 128, |
| 408 | BLOCK_QUANT_DIM: int = 32, |
| 409 | ): |
| 410 | """Convert the src weights to MXFP4. The src weight is quantized along the |
| 411 | axis dimension into packed e2m1 values (torch.uint8, two values per |
| 412 | byte), so the size of that dimension in the output is half of the |
| 413 | logical (unpacked) size. |
| 414 | """ |
| 415 | out_quant_type = torch.uint8 |
| 416 | ndim = src_tensor.ndim |
| 417 | assert -ndim <= axis < ndim, f"Invalid axis {axis=}" |
| 418 | axis = axis if axis >= 0 else axis + ndim |
| 419 | |
| 420 | L = src_tensor.shape[axis] |
| 421 | # We make this assertion since we can't track if the "real" shape was odd, |
| 422 | # and we padded it to be even. |
| 423 | # We want to maintain the property dequant(quant(x)).shape == x.shape |
| 424 | assert L % 2 == 0, f"axis dim must be divisible by 2 for e2m1. Got {L}" |
| 425 | |
| 426 | device = src_tensor.device |
| 427 | |
| 428 | packed_quant_dim = triton.cdiv(L, 2) |
| 429 | out_scale_dim = triton.cdiv(L, 32) |
| 430 | |
| 431 | # Move the quantization axis to the end for the kernel, then permute back. |
| 432 | permute_order = list(range(ndim)) |
| 433 | permute_order[axis], permute_order[-1] = permute_order[-1], permute_order[axis] |
| 434 | |
| 435 | prmted_quant_tensor_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + ( |
| 436 | packed_quant_dim, |
| 437 | ) |
| 438 | prmted_scale_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + ( |
| 439 | out_scale_dim, |
| 440 | ) |
| 441 | prmted_src_tensor = src_tensor.permute(permute_order) |
| 442 | |
| 443 | if out_quant_tensor is None: |
| 444 | out_quant_tensor = torch.empty( |
| 445 | prmted_quant_tensor_shape, dtype=out_quant_type, device=device |
| 446 | ) |
| 447 | else: |
| 448 | expected_shape = ( |
| 449 | src_tensor.shape[:axis] + (packed_quant_dim,) + src_tensor.shape[axis + 1 :] |
| 450 | ) |
| 451 | assert out_quant_tensor.shape == expected_shape, ( |
| 452 | f"{out_quant_tensor.shape=} != {expected_shape=}" |
| 453 | ) |
| 454 | assert out_quant_tensor.dtype == out_quant_type, ( |
| 455 | f"{out_quant_tensor.dtype=} != {out_quant_type=}" |
| 456 | ) |
| 457 | assert out_quant_tensor.stride(axis) == 1, ( |
| 458 | f"{out_quant_tensor.stride(axis)=} != 1" |
| 459 | ) |
| 460 | # We expect the axis dimension to be last, so permute the tensor |
| 461 | out_quant_tensor = out_quant_tensor.permute(permute_order) |
| 462 | |
| 463 | if out_scale is None: |
| 464 | out_scale = torch.empty(prmted_scale_shape, dtype=torch.uint8, device=device) |
| 465 | else: |
| 466 | expected_scale_shape = permute_shape(prmted_scale_shape, permute_order) |
| 467 | assert out_scale.shape == expected_scale_shape, ( |
| 468 | f"{out_scale.shape=} {expected_scale_shape=}" |
| 469 | ) |
| 470 | assert out_scale.dtype == torch.uint8, f"{out_scale.dtype=} != torch.uint8" |
| 471 | out_scale = out_scale.permute(permute_order) |
| 472 | |
| 473 | # Flatten input tensor for kernel. This will typically make a copy |
| 474 | reshaped_src_tensor = prmted_src_tensor.reshape(-1, L) |
| 475 | blocks_quant_dim = triton.cdiv(reshaped_src_tensor.shape[-1], BLOCK_QUANT_DIM) |
| 476 | blocks_out_dim = triton.cdiv(reshaped_src_tensor.shape[0], BLOCK_OUT_DIM) |
| 477 | |
| 478 | # Flatten the output tensors for the kernel, this should be a view always |
| 479 | kernel_quant_tensor = out_quant_tensor.reshape(-1, packed_quant_dim) |
| 480 | kernel_scale = out_scale.reshape(-1, out_scale_dim) |
| 481 | assert kernel_quant_tensor.data_ptr() == out_quant_tensor.data_ptr() |
| 482 | assert kernel_scale.data_ptr() == out_scale.data_ptr() |
| 483 | |
| 484 | _downcast_to_mxfp[(blocks_out_dim, blocks_quant_dim)]( |
| 485 | kernel_quant_tensor, |
| 486 | *kernel_quant_tensor.stride(), |
| 487 | kernel_scale, |
| 488 | *kernel_scale.stride(), |
| 489 | reshaped_src_tensor, |
| 490 | *reshaped_src_tensor.stride(), |
| 491 | *reshaped_src_tensor.shape, |
| 492 | BLOCK_OUT_DIM, |
| 493 | BLOCK_QUANT_DIM, |
| 494 | num_warps=8, |
| 495 | ) |
| 496 | |
| 497 | out_quant_tensor = out_quant_tensor.permute(permute_order) |
| 498 | out_scale = out_scale.permute(permute_order).contiguous() |
| 499 | return out_quant_tensor, out_scale, permute_shape(prmted_scale_shape, permute_order) |
| 500 | |
| 501 | |
| 502 | def mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 503 | """Quantize a bf16/fp16 tensor to MXFP4 along its last dimension. |
| 504 | |
| 505 | Dispatches to the fastest backend available on the current platform: |
| 506 | |
| 507 | - the native XPU custom op on XPU, |
| 508 | - aiter on ROCm when `aiter` is installed, |
| 509 | - and the portable Triton kernel (``downcast_to_mxfp``) otherwise. |
| 510 | |
| 511 | Returns packed FP4 values (uint8, two values per byte) and per-block |
| 512 | (group-32) e8m0 scales (uint8). |
| 513 | """ |
| 514 | if current_platform.is_xpu(): |
| 515 | return xpu_mxfp4_quantize(x) |
| 516 | |
| 517 | from vllm._aiter_ops import is_aiter_found_and_supported |
| 518 | |
| 519 | if is_aiter_found_and_supported() and x.dtype == torch.bfloat16: |
| 520 | from vllm.model_executor.layers.quantization.quark.utils import ( |
| 521 | quark_quantize_weight_to_mxfp4, |
| 522 | ) |
| 523 | |
| 524 | return quark_quantize_weight_to_mxfp4(x) |
| 525 | |
| 526 | quant_tensor, scale, _ = downcast_to_mxfp(x, axis=-1) |
| 527 | |
| 528 | return quant_tensor, scale |