vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe ¶

def interleave_linear_and_gate(
    x: torch.Tensor,
    group_size: int = 64,
    dim: int = -1,
) -> torch.Tensor:
    """Interleave gate and linear weight rows for CuteDSL wrapper."""
    sizes = x.size()
    dim = dim % x.dim()
    assert sizes[dim] % (group_size * 2) == 0, (
        f"dim {dim} size {sizes[dim]} must be divisible by {group_size * 2}"
    )
    prev_sizes = sizes[:dim]
    post_sizes = sizes[dim + 1 :]
    x = x.view(*prev_sizes, 2, sizes[dim] // (group_size * 2), group_size, *post_sizes)
    x = x.transpose(dim, dim + 1).contiguous().view(*sizes)
    return x

def nvfp4_swizzled_scale_to_cutedsl_mma_view(scale: torch.Tensor) -> torch.Tensor:
    """View a swizzled (E, M_padded, K_sf_padded) block-scale tensor in the
    MMA layout expected by the CuteDSL MoE kernel.

    The returned tensor aliases `scale`'s storage, so in-place updates of the
    registered Parameter (weight reloads, EPLB rearrangement) are visible to
    the kernel with no extra bookkeeping.
    """
    from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout

    num_experts, m_padded, k_sf_padded = scale.shape
    mma_view = convert_sf_to_mma_layout(
        scale.reshape(num_experts * m_padded, k_sf_padded),
        m=m_padded,
        k=k_sf_padded * 16,
        num_groups=num_experts,
        sf_vec_size=16,
    )
    assert mma_view.data_ptr() == scale.data_ptr(), (
        "convert_sf_to_mma_layout no longer returns a view of its input; "
        "the quant config would go stale after weight updates."
    )
    return mma_view

def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
    layer: "RoutedExperts",
    w13: torch.Tensor,
    w13_scale: torch.Tensor,
    w13_scale_2: torch.Tensor,
    a13_scale: torch.Tensor,
    w2: torch.Tensor,
    w2_scale: torch.Tensor,
    w2_scale_2: torch.Tensor,
    a2_scale: torch.Tensor,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend.

    Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper,
    and interleaves w13 gate/linear rows for gated activations. Non-gated
    activations use a single w13 projection and keep its row order unchanged.
    """
    # Global scaling factors (same as other FlashInfer backends).
    num_experts = w13.shape[0]
    enable_eplb = layer.moe_config.moe_parallel_config.enable_eplb
    a13_scale = amax_for_moe_activation_quant(a13_scale, enable_eplb).repeat(
        num_experts
    )
    a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(num_experts)

    if layer.activation.is_gated:
        w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(
            layer.activation, w13, w13_scale
        )

        # Interleave up/gate rows for w13 weights and scales.
        w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
        w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)

    w13_scale = swizzle_blockscale(w13_scale)
    w2_scale = swizzle_blockscale(w2_scale)

    return (
        w13,
        w13_scale,
        w13_scale_2,
        a13_scale,
        w2,
        w2_scale,
        w2_scale_2,
        a2_scale,
    )

def reorder_w13_to_w31_for_flashinfer_cutedsl(
    activation: MoEActivation,
    w13: torch.Tensor,
    w13_scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Normalize gated w13 rows to the [up; gate] order used by FlashInfer."""
    if activation == MoEActivation.SWIGLUOAI:
        # gpt-oss checkpoints store w13 interleaved as [gate0, up0, gate1, ...].
        gate, up = w13[:, 0::2], w13[:, 1::2]
        gate_scale, up_scale = w13_scale[:, 0::2], w13_scale[:, 1::2]
        return (
            torch.cat([up, gate], dim=1).contiguous(),
            torch.cat([up_scale, gate_scale], dim=1).contiguous(),
        )

    half = w13.shape[1] // 2
    return (
        torch.cat([w13[:, half:], w13[:, :half]], dim=1).contiguous(),
        torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1).contiguous(),
    )

def reorder_w1w3_to_w3w1(
    weight: torch.Tensor, scale: torch.Tensor, dim: int = -2
) -> tuple[torch.Tensor, torch.Tensor]:
    """Re-order concatenated `[w1, w3]` tensors to `[w3, w1]` in-place.

    `weight` and `scale` must be contiguous; they remain contiguous on return.
    """
    assert weight.is_contiguous(), "weight must be contiguous"
    assert scale.is_contiguous(), "scale must be contiguous"
    size = weight.size(dim)
    assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
    half = size // 2
    d = dim % weight.dim()

    # 64 MB transient cap
    bytes_per_row = max(
        weight.numel() // size * weight.element_size(),
        scale.numel() // size * scale.element_size(),
    )
    chunk = max(1, min(half, (64 << 20) // max(bytes_per_row, 1)))

    fa, fb = [slice(None)] * weight.dim(), [slice(None)] * weight.dim()
    for off in range(0, half, chunk):
        end = min(off + chunk, half)
        fa[d], fb[d] = slice(off, end), slice(half + off, half + end)
        a, b = tuple(fa), tuple(fb)
        for t in (weight, scale):
            tmp = t[b].clone()
            t[b] = t[a]
            t[a] = tmp

    return weight, scale