vllm.model_executor.layers.quantization.utils.mxfp4_utils ¶

def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
    """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
    assert has_triton_kernels()
    import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
    from triton_kernels.numerics import InFlexData
    from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
    from triton_kernels.tensor_details import layout
    from triton_kernels.tensor_details.layout import StridedLayout

    value_layout_opts: dict[str, Any] = {}
    scale_layout_opts: dict[str, Any] = {}

    if (
        current_platform.is_cuda()
        and current_platform.is_device_capability(90)
        and not is_torch_equal_or_newer("2.8.1")
    ):
        logger.warning_once(
            "Mxfp4 on hopper is running on torch < 2.8.1, "
            "this cause swizling to be disabled, which may "
            "cause performance degradation. Please upgrade to torch nightly"
        )
        value_layout = StridedLayout
        scale_layout = StridedLayout
    elif current_platform.is_rocm():
        value_layout = StridedLayout
        if should_use_cdna4_mx_scale_swizzle():
            try:
                # triton < 3.6
                from triton_kernels.tensor_details.layout import GFX950MXScaleLayout

                scale_layout = GFX950MXScaleLayout
            except ImportError:
                # triton >= 3.6
                from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout

                scale_layout = CDNA4MXScaleLayout
        else:
            scale_layout = StridedLayout
    else:
        value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(
            mx_axis=1
        )
        scale_layout, scale_layout_opts = (
            layout.make_default_matmul_mxfp4_w_scale_layout(
                mx_axis=1, num_warps=num_warps
            )
        )
    if current_platform.is_cuda():
        if current_platform.is_device_capability(90):
            constraints = {
                "split_k": 1,
            }
            opt_flags.update_opt_flags_constraints(constraints)
            # Patches #47303: pad K (num scale groups) to 0 mod 4
            # TODO: Remove once we upgrade to Triton 3.8.0+ kernels
            if scale.numel() > 0:
                K = scale.shape[-1]
                pad_k = -K % 4
                scale = torch.nn.functional.pad(scale, (0, pad_k))
        elif current_platform.is_device_capability_family(100):
            constraints = {
                "is_persistent": True,
                "epilogue_subtile": 1,
            }
            opt_flags.update_opt_flags_constraints(constraints)
    # transpose the tensor so that the quantization axis is on dim1
    quant_tensor = quant_tensor.transpose(-2, -1)
    scale = scale.transpose(-2, -1)
    quant_tensor = convert_layout(
        wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts
    )
    scale = convert_layout(wrap_torch_tensor(scale), scale_layout, **scale_layout_opts)
    return quant_tensor, InFlexData(), scale

def downcast_to_mxfp(
    src_tensor: torch.Tensor,
    axis: int,
    out_quant_tensor: torch.Tensor | None = None,
    out_scale: torch.Tensor | None = None,
    BLOCK_OUT_DIM: int = 128,
    BLOCK_QUANT_DIM: int = 32,
):
    """
    Convert the src weights to MXFP4. The src weight is quantized along the
    axis dimension into packed e2m1 values (torch.uint8, two values per
    byte), so the size of that dimension in the output is half of the
    logical (unpacked) size.
    """
    out_quant_type = torch.uint8
    ndim = src_tensor.ndim
    assert -ndim <= axis < ndim, f"Invalid axis {axis=}"
    axis = axis if axis >= 0 else axis + ndim

    L = src_tensor.shape[axis]
    # We make this assertion since we can't track if the "real" shape was odd,
    # and we padded it to be even.
    # We want to maintain the property dequant(quant(x)).shape == x.shape
    assert L % 2 == 0, f"axis dim must be divisible by 2 for e2m1. Got {L}"

    device = src_tensor.device

    packed_quant_dim = triton.cdiv(L, 2)
    out_scale_dim = triton.cdiv(L, 32)

    # Move the quantization axis to the end for the kernel, then permute back.
    permute_order = list(range(ndim))
    permute_order[axis], permute_order[-1] = permute_order[-1], permute_order[axis]

    prmted_quant_tensor_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + (
        packed_quant_dim,
    )
    prmted_scale_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + (
        out_scale_dim,
    )
    prmted_src_tensor = src_tensor.permute(permute_order)

    if out_quant_tensor is None:
        out_quant_tensor = torch.empty(
            prmted_quant_tensor_shape, dtype=out_quant_type, device=device
        )
    else:
        expected_shape = (
            src_tensor.shape[:axis] + (packed_quant_dim,) + src_tensor.shape[axis + 1 :]
        )
        assert out_quant_tensor.shape == expected_shape, (
            f"{out_quant_tensor.shape=} != {expected_shape=}"
        )
        assert out_quant_tensor.dtype == out_quant_type, (
            f"{out_quant_tensor.dtype=} != {out_quant_type=}"
        )
        assert out_quant_tensor.stride(axis) == 1, (
            f"{out_quant_tensor.stride(axis)=} != 1"
        )
        # We expect the axis dimension to be last, so permute the tensor
        out_quant_tensor = out_quant_tensor.permute(permute_order)

    if out_scale is None:
        out_scale = torch.empty(prmted_scale_shape, dtype=torch.uint8, device=device)
    else:
        expected_scale_shape = permute_shape(prmted_scale_shape, permute_order)
        assert out_scale.shape == expected_scale_shape, (
            f"{out_scale.shape=} {expected_scale_shape=}"
        )
        assert out_scale.dtype == torch.uint8, f"{out_scale.dtype=} != torch.uint8"
        out_scale = out_scale.permute(permute_order)

    # Flatten input tensor for kernel. This will typically make a copy
    reshaped_src_tensor = prmted_src_tensor.reshape(-1, L)
    blocks_quant_dim = triton.cdiv(reshaped_src_tensor.shape[-1], BLOCK_QUANT_DIM)
    blocks_out_dim = triton.cdiv(reshaped_src_tensor.shape[0], BLOCK_OUT_DIM)

    # Flatten the output tensors for the kernel, this should be a view always
    kernel_quant_tensor = out_quant_tensor.reshape(-1, packed_quant_dim)
    kernel_scale = out_scale.reshape(-1, out_scale_dim)
    assert kernel_quant_tensor.data_ptr() == out_quant_tensor.data_ptr()
    assert kernel_scale.data_ptr() == out_scale.data_ptr()

    _downcast_to_mxfp[(blocks_out_dim, blocks_quant_dim)](
        kernel_quant_tensor,
        *kernel_quant_tensor.stride(),
        kernel_scale,
        *kernel_scale.stride(),
        reshaped_src_tensor,
        *reshaped_src_tensor.stride(),
        *reshaped_src_tensor.shape,
        BLOCK_OUT_DIM,
        BLOCK_QUANT_DIM,
        num_warps=8,
    )

    out_quant_tensor = out_quant_tensor.permute(permute_order)
    out_scale = out_scale.permute(permute_order).contiguous()
    return out_quant_tensor, out_scale, permute_shape(prmted_scale_shape, permute_order)

def mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize a bf16/fp16 tensor to MXFP4 along its last dimension.

    Dispatches to the fastest backend available on the current platform:

    - the native XPU custom op on XPU,
    - aiter on ROCm when `aiter` is installed,
    - and the portable Triton kernel (``downcast_to_mxfp``) otherwise.

    Returns packed FP4 values (uint8, two values per byte) and per-block
    (group-32) e8m0 scales (uint8).
    """
    if current_platform.is_xpu():
        return xpu_mxfp4_quantize(x)

    from vllm._aiter_ops import is_aiter_found_and_supported

    if is_aiter_found_and_supported() and x.dtype == torch.bfloat16:
        from vllm.model_executor.layers.quantization.quark.utils import (
            quark_quantize_weight_to_mxfp4,
        )

        return quark_quantize_weight_to_mxfp4(x)

    quant_tensor, scale, _ = downcast_to_mxfp(x, axis=-1)

    return quant_tensor, scale

def should_use_cdna4_mx_scale_swizzle() -> bool:
    """Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950.

    CDNA4 swizzle requires BLOCK_K%256==0; at TP>=4 the A8W4 dispatch
    picks BK<256 tiles for the smaller per-rank shapes, so swizzle must
    be off. Used by both the weight-load swizzle in `_swizzle_mxfp4` and
    the kernel-argument gate in `aiter_mxfp4_w4a8_moe`; they must agree.
    """
    from vllm.distributed import get_tensor_model_parallel_world_size
    from vllm.platforms.rocm import on_gfx950

    return on_gfx950() and get_tensor_model_parallel_world_size() <= 2