vllm.model_executor.layers.quantization.utils.mxfp8_utils

Functions:

• dequant_mxfp8_to_bf16 –

  Dequantize MXFP8 tensor to BF16.

• mxfp8_e4m3_quantize_fake –

  Fake implementation for torch.compile tracing.

• requant_block_fp8_to_ue8m0 –

  Requantize block-FP8 weights to E4M3 with one UE8M0 scale per 1x32.

• swizzle_mxfp8_scale –

  Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout.

_flashinfer_mxfp8_quantize() cached

FlashInfer's MXFP8 quantize when this device uses it, else None.

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

@functools.cache
def _flashinfer_mxfp8_quantize():
    """FlashInfer's MXFP8 quantize when this device uses it, else None."""
    from vllm.platforms import current_platform
    from vllm.utils.flashinfer import has_flashinfer

    if current_platform.has_device_capability(100) and has_flashinfer():
        from flashinfer import mxfp8_quantize

        return mxfp8_quantize
    return None

_mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False)

Naive MXFP8 quantization. For each block of 32 elements along the last dimension, compute a shared e8m0 scale that fits the block-wise amax into the finite float8_e4m3fn range, and quantize each element to float8_e4m3fn.

Returns (quantized_values [same shape, fp8], scales uint8). Scale shape depends on is_sf_swizzled_layout: False -> [..., K//32] (row-major 2D) True -> [flat swizzled 1D]

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

def _mxfp8_e4m3_quantize_torch(
    x: torch.Tensor,
    is_sf_swizzled_layout: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Naive MXFP8 quantization.
    For each block of 32 elements along the last dimension, compute a
    shared e8m0 scale that fits the block-wise amax into the finite
    float8_e4m3fn range, and quantize each element to float8_e4m3fn.

    Returns (quantized_values [same shape, fp8], scales uint8).
    Scale shape depends on is_sf_swizzled_layout:
      False -> [..., K//32]  (row-major 2D)
      True  -> [flat swizzled 1D]
    """
    assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0
    orig_shape = x.shape
    num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE

    x_fp32 = x.to(torch.float32)
    x_blocked = x_fp32.view(*orig_shape[:-1], num_blocks, MXFP8_BLOCK_SIZE)

    amax = x_blocked.abs().amax(dim=-1)
    amax = amax.clamp(min=torch.finfo(torch.float32).tiny)
    fp8_max = torch.finfo(MXFP8_VALUE_DTYPE).max
    scale_biased = torch.ceil(torch.log2(amax / fp8_max)) + 127.0
    scale_biased = scale_biased.clamp(0, 254)
    scales_uint8 = scale_biased.to(torch.uint8)

    descale = torch.exp2(scale_biased - 127.0)
    x_scaled = x_blocked / descale.unsqueeze(-1)

    x_fp8 = x_scaled.view(orig_shape).to(MXFP8_VALUE_DTYPE)

    if x.ndim == 2:
        M, K = x.shape
        scales_uint8 = scales_uint8.view(M, -1)
        if is_sf_swizzled_layout:
            scales_uint8 = swizzle_mxfp8_scale(scales_uint8, M=M, K=K)
    elif x.ndim == 3:
        B, M, K = x.shape
        scales_uint8 = scales_uint8.view(B, M, -1)
        if is_sf_swizzled_layout:
            swizzled = []
            for i in range(B):
                swizzled.append(swizzle_mxfp8_scale(scales_uint8[i], M=M, K=K))
            scales_uint8 = torch.cat(swizzled)

    return x_fp8, scales_uint8

_mxfp8_e4m3_quantize_triton(x)

Fused 2D MXFP8 quant (non-swizzled, row-major [M, K//32] scales).

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

def _mxfp8_e4m3_quantize_triton(
    x: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fused 2D MXFP8 quant (non-swizzled, row-major [M, K//32] scales)."""
    from vllm.triton_utils import triton

    global _MXFP8_QUANT_KERNEL
    if _MXFP8_QUANT_KERNEL is None:
        _MXFP8_QUANT_KERNEL = _mxfp8_quant_triton_kernel()

    M, K = x.shape
    x = x.contiguous()
    xq = torch.empty((M, K), dtype=MXFP8_VALUE_DTYPE, device=x.device)
    scales = torch.empty(
        (M, K // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device
    )
    BLOCK_M = _MXFP8_M1_BLOCK_M.get(K, 64) if M == 1 else 64
    grid = (triton.cdiv(M, BLOCK_M), K // MXFP8_BLOCK_SIZE)
    _MXFP8_QUANT_KERNEL[grid](
        x,
        xq,
        scales,
        M,
        K,
        x.stride(0),
        x.stride(1),
        xq.stride(0),
        xq.stride(1),
        scales.stride(0),
        scales.stride(1),
        BLOCK_M=BLOCK_M,
        FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
        TINY=float(torch.finfo(torch.float32).tiny),
        num_warps=4,
    )
    return xq, scales

_mxfp8_quant_triton_kernel()

Lazily-built Triton kernel: per-32-block E8M0 scale + FP8-E4M3 quant.

Fuses what _mxfp8_e4m3_quantize_torch does in several elementwise passes into one launch. Each program handles [BLOCK_M, 32] (one MX block).

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

def _mxfp8_quant_triton_kernel():
    """Lazily-built Triton kernel: per-32-block E8M0 scale + FP8-E4M3 quant.

    Fuses what ``_mxfp8_e4m3_quantize_torch`` does in several elementwise passes
    into one launch. Each program handles ``[BLOCK_M, 32]`` (one MX block).
    """
    from vllm.triton_utils import tl, triton

    @triton.jit
    def _kernel(
        x_ptr,
        xq_ptr,
        s_ptr,
        M,
        K,
        sxm,
        sxk,
        sqm,
        sqk,
        ssm,
        ssk,
        BLOCK_M: tl.constexpr,
        FP8_MAX: tl.constexpr,
        TINY: tl.constexpr,
    ):
        pid_m = tl.program_id(0)
        pid_b = tl.program_id(1)  # which 32-element block along K
        offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_k = pid_b * 32 + tl.arange(0, 32)
        m_mask = offs_m < M
        x = tl.load(
            x_ptr + offs_m[:, None] * sxm + offs_k[None, :] * sxk,
            mask=m_mask[:, None],
            other=0.0,
        ).to(tl.float32)
        # Mirror _mxfp8_e4m3_quantize_torch: the scale has to put the block amax
        # at the top of the e4m3 range rather than at 1.0, or small elements of
        # the block end up in the subnormals.
        amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY)  # [BLOCK_M]
        sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
        sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
        # Scale by the reciprocal rather than dividing by ``exp2(sb - 127)``:
        # sb == 0 (an all-zero / denormal block) makes that divisor 2**-127,
        # which is subnormal in fp32 and flushes to zero on CDNA, turning the
        # block's zeros into 0/0 == NaN. The reciprocal is normal for every
        # reachable sb, and both forms are exact powers of two, so nothing
        # else changes.
        rescale = tl.exp2(127.0 - sb)
        xq = (x * rescale[:, None]).to(xq_ptr.dtype.element_ty)
        tl.store(
            xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk,
            xq,
            mask=m_mask[:, None],
        )
        tl.store(s_ptr + offs_m * ssm + pid_b * ssk, sb.to(tl.uint8), mask=m_mask)

    return _kernel

dequant_mxfp8_to_bf16(x, scales)

Dequantize MXFP8 tensor to BF16.

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
    """Dequantize MXFP8 tensor to BF16."""
    x_float = x.to(torch.float32)

    num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE
    x_blocked = x_float.view(*x.shape[:-1], num_blocks, MXFP8_BLOCK_SIZE)

    descale = torch.exp2(scales.to(torch.float32) - 127.0)

    dequantized = x_blocked * descale.unsqueeze(-1)

    dequantized = dequantized.view(*x.shape)

    return dequantized.to(torch.bfloat16)

mxfp8_e4m3_quantize_fake(x, is_sf_swizzled_layout=False, alignment=0)

Fake implementation for torch.compile tracing.

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

def mxfp8_e4m3_quantize_fake(
    x: torch.Tensor,
    is_sf_swizzled_layout: bool = False,
    alignment: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fake implementation for torch.compile tracing."""
    fp_data = torch.empty_like(x, dtype=MXFP8_VALUE_DTYPE)

    block_size = MXFP8_BLOCK_SIZE

    if x.ndim == 2:
        M, N = x.shape
        K = (N + block_size - 1) // block_size
        if is_sf_swizzled_layout:
            M_padded = ((M + 127) // 128) * 128
            K_padded = ((K + 3) // 4) * 4
            scales = torch.empty(
                M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device
            )
        else:
            scales = torch.empty((M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device)
    elif x.ndim == 3:
        B, M, N = x.shape
        K = (N + block_size - 1) // block_size
        if is_sf_swizzled_layout:
            M_padded = ((M + 127) // 128) * 128
            K_padded = ((K + 3) // 4) * 4
            scales = torch.empty(
                B * M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device
            )
        else:
            scales = torch.empty((B, M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device)
    else:
        scale_shape = list(x.shape)
        scale_shape[-1] = (x.shape[-1] + block_size - 1) // block_size
        scales = torch.empty(scale_shape, dtype=MXFP8_SCALE_DTYPE, device=x.device)

    return fp_data, scales

requant_block_fp8_to_ue8m0(weight, scale)

Requantize block-FP8 weights to E4M3 with one UE8M0 scale per 1x32.

The SM100 MegaMoE feeds its scales to the block-scaled MMA, so they have to be powers of two, one per row and 32 columns. A checkpoint with float32 scales per block (e.g. 128x128) is rounded once more to get there.

Parameters:

• weight

  (Tensor) –

  (n, mn, k) E4M3 values.

• scale

  (Tensor) –

  (n, mn / block_m, k / block_k) float32 dequantization scales (real = weight * scale).

Returns:

• Tensor –

  (weight, sf): the requantized E4M3 weights and their UE8M0 scale

• Tensor –

  bytes, (n, mn, k / 32) uint8.

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

def requant_block_fp8_to_ue8m0(
    weight: torch.Tensor, scale: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Requantize block-FP8 weights to E4M3 with one UE8M0 scale per 1x32.

    The SM100 MegaMoE feeds its scales to the block-scaled MMA, so they have to
    be powers of two, one per row and 32 columns. A checkpoint with float32
    scales per block (e.g. 128x128) is rounded once more to get there.

    Args:
        weight: ``(n, mn, k)`` E4M3 values.
        scale: ``(n, mn / block_m, k / block_k)`` float32 dequantization
            scales (``real = weight * scale``).

    Returns:
        ``(weight, sf)``: the requantized E4M3 weights and their UE8M0 scale
        bytes, ``(n, mn, k / 32)`` uint8.

    """
    n, mn, k = weight.shape
    _, blocks_m, blocks_k = scale.shape
    real = weight.to(torch.float32).view(n, blocks_m, -1, blocks_k, k // blocks_k)
    real.mul_(scale[:, :, None, :, None])
    return _mxfp8_e4m3_quantize_torch(real.view(n, mn, k))

swizzle_mxfp8_scale(sf, M, K)

Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout.

Source code in vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

def swizzle_mxfp8_scale(sf: torch.Tensor, M: int, K: int) -> torch.Tensor:
    """Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout."""
    scaling_vector_size = MXFP8_BLOCK_SIZE  # 32 for MXFP8
    factor = scaling_vector_size * 4  # 128

    num_m_tiles = (M + 127) // 128
    num_k_tiles = (K + factor - 1) // factor

    m_padded = num_m_tiles * 128
    k_scale_padded = num_k_tiles * 4

    scale_cols = K // scaling_vector_size
    sf_padded = torch.zeros(
        (m_padded, k_scale_padded), dtype=sf.dtype, device=sf.device
    )
    sf_padded[:M, :scale_cols] = sf

    sf_reshaped = sf_padded.view(num_m_tiles, 4, 32, num_k_tiles, 4)

    sf_swizzled = sf_reshaped.transpose(1, 3)

    return sf_swizzled.contiguous().view(-1)