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 prepare_static_weights_for_trtllm_fp4_moe(
    gemm1_weights: torch.Tensor,
    gemm2_weights: torch.Tensor,
    gemm1_scales_linear_fp4_bytes: torch.Tensor,
    gemm2_scales_linear_fp4_bytes: torch.Tensor,
    hidden_size: int,
    intermediate_size: int,
    num_experts: int,
    is_gated_activation: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Shuffle NVFP4 weights into the FlashInfer TRT-LLM layout."""
    from flashinfer import nvfp4_block_scale_interleave
    from flashinfer.fused_moe.core import (
        _maybe_get_cached_w3_w1_permute_indices,
        get_w2_permute_indices_with_cache,
    )

    permute_indices_cache: dict[tuple[object, ...], torch.Tensor] = {}
    epilogue_tile_m = 128  # FIXME: this depends on the kernel internals
    gemm1_intermediate_size = (
        2 * intermediate_size if is_gated_activation else intermediate_size
    )

    # Convert quantized weights to proper formats
    gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(
        num_experts, gemm1_intermediate_size, hidden_size // 2
    )  # packed fp4
    gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(
        torch.float8_e4m3fn
    ).reshape(
        num_experts, gemm1_intermediate_size, hidden_size // 16
    )  # fp8 scaling factors

    gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(
        num_experts, hidden_size, intermediate_size // 2
    )  # packed fp4
    gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(
        torch.float8_e4m3fn
    ).reshape(num_experts, hidden_size, intermediate_size // 16)  # fp8 scaling factors

    gemm1_weight_indices = _maybe_get_cached_w3_w1_permute_indices(
        permute_indices_cache,
        gemm1_weights_fp4[0].view(torch.uint8),
        epilogue_tile_m,
        is_gated_act_gemm=is_gated_activation,
    )
    gemm1_scale_indices = _maybe_get_cached_w3_w1_permute_indices(
        permute_indices_cache,
        gemm1_scales_linear_fp4[0].view(torch.uint8),
        epilogue_tile_m,
        num_elts_per_sf=16,
        is_gated_act_gemm=is_gated_activation,
    )
    gemm2_weight_indices = get_w2_permute_indices_with_cache(
        permute_indices_cache,
        gemm2_weights_fp4[0].view(torch.uint8),
        epilogue_tile_m,
    )
    gemm2_scale_indices = get_w2_permute_indices_with_cache(
        permute_indices_cache,
        gemm2_scales_linear_fp4[0].view(torch.uint8),
        epilogue_tile_m,
        num_elts_per_sf=16,
    )

    gemm1_weights_fp4_shuffled = torch.empty_like(gemm1_weights_fp4, dtype=torch.uint8)
    gemm1_scales_fp4_shuffled = torch.empty_like(gemm1_scales_linear_fp4)
    gemm2_weights_fp4_shuffled = torch.empty_like(gemm2_weights_fp4, dtype=torch.uint8)
    gemm2_scales_fp4_shuffled = torch.empty_like(gemm2_scales_linear_fp4)
    gemm1_scale_scratch = torch.empty_like(
        gemm1_scales_linear_fp4[0], dtype=torch.uint8
    )
    gemm2_scale_scratch = torch.empty_like(
        gemm2_scales_linear_fp4[0], dtype=torch.uint8
    )

    for expert_id in range(num_experts):
        torch.index_select(
            gemm1_weights_fp4[expert_id].view(torch.uint8),
            0,
            gemm1_weight_indices,
            out=gemm1_weights_fp4_shuffled[expert_id],
        )
        torch.index_select(
            gemm1_scales_linear_fp4[expert_id].view(torch.uint8),
            0,
            gemm1_scale_indices,
            out=gemm1_scale_scratch,
        )
        gemm1_scales_fp4_shuffled[expert_id].view(torch.uint8).reshape(-1).copy_(
            nvfp4_block_scale_interleave(gemm1_scale_scratch)
        )

        torch.index_select(
            gemm2_weights_fp4[expert_id].view(torch.uint8),
            0,
            gemm2_weight_indices,
            out=gemm2_weights_fp4_shuffled[expert_id],
        )
        torch.index_select(
            gemm2_scales_linear_fp4[expert_id].view(torch.uint8),
            0,
            gemm2_scale_indices,
            out=gemm2_scale_scratch,
        )
        gemm2_scales_fp4_shuffled[expert_id].view(torch.uint8).reshape(-1).copy_(
            nvfp4_block_scale_interleave(gemm2_scale_scratch)
        )

    return (
        gemm1_weights_fp4_shuffled,
        gemm1_scales_fp4_shuffled,
        gemm2_weights_fp4_shuffled,
        gemm2_scales_fp4_shuffled,
    )

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