vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 ¶

def _nvfp4_compute_scale_factor(
    marlin_scales: torch.Tensor,
    a_dtype: torch.dtype | None = None,
) -> float:
    """Compute the power-of-2 scale_factor needed so that all non-zero
    values in marlin_scales * 2^7 are >= 2 after rescaling.
    Returns a Python float (power of 2, >= 1.0)."""

    # Since half has a smaller dynamic range compared to bfloat16,
    # no rescaling is applied here if active dtype is half.
    if a_dtype is not None and a_dtype == torch.half:
        return 1.0

    ws_float = marlin_scales.float() * (2**7)
    nonzero_mask = ws_float > 0
    if nonzero_mask.any():
        max_val = ws_float[nonzero_mask].max()
        if max_val < 448 * (2**7):
            sf = (448 * (2**7) / max_val).log2().floor().exp2()
            return sf.item()
    return 1.0

def nvfp4_marlin_process_scales(
    marlin_scales: torch.Tensor,
    scale_factor: float | None = None,
    a_dtype: torch.dtype | None = None,
) -> tuple[torch.Tensor, float]:
    """Process NVFP4 weight scales into the special S0E5M3 format for Marlin.

    Args:
        marlin_scales: Weight scales tensor in half precision, already
            permuted for the Marlin kernel layout.
        scale_factor: Optional power-of-2 rescaling factor. If None, the
            factor is computed automatically so that every non-zero scale
            satisfies ``scale * 2^7 >= 2`` (i.e., the MSB of the S0E5M3
            representation is always 1). When provided (e.g., for MoE
            layers where all experts must share the same factor), the
            given value is used directly. The caller is responsible for
            dividing ``global_scale`` by the returned ``scale_factor`` to
            preserve numerical correctness.

    Returns:
        A tuple of (processed_scales, scale_factor).
    """
    if not (marlin_scales >= 0).all():
        logger.warning_once(
            "NVFP4 Marlin assumes the scales to be >=0, but has encountered "
            "negative scales. Accuracy will likely be degraded. This is "
            "because it changes the scales from FP8-S1E4M3 to a special "
            "FP8-S0E5M3 format to speedup the dequantization."
        )

    # convert to half first, we would convert to fp8 later
    marlin_scales = marlin_scales.to(torch.half)

    # fit the layout of fp8 dequantization
    marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
        marlin_scales.size(0), -1
    )

    # We assume that weight_scale (FP8-S1E4M3) is always greater
    # than or equal to 0. So we can convert
    # (weight_scale * (2 ** 7) to a special FP8-S0E5M3 format.
    # After multiplying by 2 ** 7, the top bit of FP8-S0E5M3 would always be 1
    # when weight_scale > 0. This allows us to have an exponent bias
    # closer to zero after dequantization.

    # Rescale weight_scale so that all non-zero values have MSB=1
    # after multiplying by 2^7 (i.e., weight_scale * 2^7 >= 2).
    # This is needed for models whose E4M3 scales were not normalized
    # to fully utilize the E4M3 dynamic range (e.g., global_scale=1).
    # The caller must compensate by dividing global_scale by scale_factor.
    if scale_factor is None:
        scale_factor = _nvfp4_compute_scale_factor(marlin_scales, a_dtype)
    if scale_factor > 1.0:
        marlin_scales = (marlin_scales.float() * scale_factor).to(torch.half)

    marlin_scales = marlin_scales * (2**7)
    marlin_scales[marlin_scales < 2] = 0
    marlin_scales = marlin_scales.view(torch.int16) << 1
    marlin_scales = marlin_scales.view(torch.float8_e4m3fn)
    marlin_scales = marlin_scales[:, 1::2].contiguous()

    return marlin_scales, scale_factor

def prepare_moe_mxfp4_layer_for_marlin(
    layer: torch.nn.Module,
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    w13_bias: torch.Tensor | None,
    w2_bias: torch.Tensor | None,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor | None,
    torch.Tensor | None,
]:
    """Pure-function version of prepare_moe_fp4_layer_for_marlin for MXFP4.

    Takes weight tensors as inputs and returns transformed tensors.
    Does NOT modify the layer in-place.
    """
    input_dtype = get_marlin_input_dtype()
    if (
        input_dtype is not None
        and input_dtype.itemsize == 1
        and input_dtype != torch.float8_e4m3fn
    ):
        raise RuntimeError("MXFP4 weight + INT8 activation is not supported.")

    group_size = 32  # MXFP4 block size

    # Derive dimensions from actual weight shapes to handle rounded/padded
    # sizes correctly (e.g., Mxfp4MoEMethod rounds up hidden_dim).
    # w13 shape: (E, 2*N, K//2)
    e = w13.shape[0]
    n = w13.shape[1] // 2  # intermediate_size_per_partition
    k = w13.shape[2] * 2  # hidden_size

    device = w13.device
    param_dtype = layer.params_dtype
    is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
    perm = torch.empty(0, dtype=torch.int, device=device)

    # WEIGHT: Repack weights to marlin format
    def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
        tensor_list = []
        if "w13" in name:
            size_n, size_k = n * 2, k
        else:
            size_n, size_k = k, n

        assert weight.shape == (e, size_n, size_k // 2)

        for i in range(e):
            qweight = weight[i].view(torch.int32).T.contiguous()
            marlin_qweight = ops.gptq_marlin_repack(
                b_q_weight=qweight,
                perm=perm,
                size_k=size_k,
                size_n=size_n,
                num_bits=4,
                is_a_8bit=is_a_8bit,
            )
            tensor_list.append(marlin_qweight)
        return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)

    w13 = repack_weight(w13, "w13")
    w2 = repack_weight(w2, "w2")

    # WEIGHT SCALES: Permute scales
    def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor:
        scales = scales.view(torch.float8_e8m0fnu)
        scales = scales.to(param_dtype)

        tensor_list = []
        if "w13" in name:
            size_n, size_k = n * 2, k
        else:
            size_n, size_k = k, n

        for i in range(e):
            scale = scales[i].T
            marlin_scales = marlin_permute_scales(
                s=scale,
                size_k=size_k,
                size_n=size_n,
                group_size=group_size,
                is_a_8bit=is_a_8bit,
            )
            marlin_scales = mxfp4_marlin_process_scales(
                marlin_scales, input_dtype=input_dtype
            )
            tensor_list.append(marlin_scales)
        return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)

    w13_scale = permute_scales(w13_scale, "w13")
    w2_scale = permute_scales(w2_scale, "w2")

    # BIAS: Permute bias
    def permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None:
        if bias is None:
            return None
        bias = bias.to(param_dtype)
        tensor_list = []
        for i in range(e):
            tensor_list.append(marlin_permute_bias(bias[i]))
        return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)

    w13_bias = permute_bias(w13_bias)
    w2_bias = permute_bias(w2_bias)

    return w13, w2, w13_scale, w2_scale, w13_bias, w2_bias