vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe

Utility helpers for NVFP4 + FlashInfer fused-MoE path.

Functions:

• reorder_w1w3_to_w3w1 –

  Re-order concatenated [w1, w3] tensors to [w3, w1] in-place.

_shard_nvfp4_moe_weights_per_die(layer, res, w13, w13_scale, w2, w2_scale)

Shard FC1 intermediate rows and FC2 hidden rows across locality domains.

Slice weights and scales along N in [E, N, K] before swizzling each shard.

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

def _shard_nvfp4_moe_weights_per_die(layer, res, w13, w13_scale, w2, w2_scale):
    """Shard FC1 intermediate rows and FC2 hidden rows across locality domains.

    Slice weights and scales along N in [E, N, K] before swizzling each shard.
    """
    n = res.num_split
    assert w13.shape[1] % n == 0 and w2.shape[1] % n == 0
    n13, n2 = w13.shape[1] // n, w2.shape[1] // n
    # Keep each interleaved 64-row up/gate pair in the same shard.
    assert n13 % 128 == 0, (
        f"per-die w13 rows {n13} must be a multiple of 128 (a whole up/gate "
        "pair after interleave_linear_and_gate(group_size=64))"
    )
    # Weight-derived slice offsets require matching scale rows.
    assert w13_scale.shape[1] == w13.shape[1], (
        f"w13_scale dim1 {w13_scale.shape[1]} != w13 dim1 {w13.shape[1]}; "
        "the per-die slice offsets would be wrong"
    )
    assert w2_scale.shape[1] == w2.shape[1], (
        f"w2_scale dim1 {w2_scale.shape[1]} != w2 dim1 {w2.shape[1]}; "
        "the per-die slice offsets would be wrong"
    )

    shards = []
    for i in range(n):
        pool = res.pools[i]
        with torch.cuda.use_mem_pool(pool):
            a, b = i * n13, i * n2
            shard = {
                # A single-expert slice is already contiguous; always allocate
                # independent storage in this domain's pool.
                "w13_weight": w13[:, a : a + n13, :].clone(
                    memory_format=torch.contiguous_format
                ),
                "w2_weight": w2[:, b : b + n2, :].clone(
                    memory_format=torch.contiguous_format
                ),
            }
            # Register contiguous, expert-major scales for EPLB, not the MMA view.
            # Keep each source alive until its asynchronous swizzle completes.
            scale_shard = w13_scale[:, a : a + n13, :].contiguous()
            shard["w13_weight_scale_param"] = swizzle_blockscale(scale_shard)
            torch.cuda.current_stream(w13_scale.device).synchronize()
            scale_shard = w2_scale[:, b : b + n2, :].contiguous()
            shard["w2_weight_scale_param"] = swizzle_blockscale(scale_shard)
            torch.cuda.current_stream(w2_scale.device).synchronize()

        # Register shards for EPLB; Parameter wrapping preserves their pool storage.
        for name in _LOCALIZED_SHARD_PARAMS:
            param = torch.nn.Parameter(shard[name], requires_grad=False)
            layer.register_parameter(f"localized{i}_{name}", param)
            shard[name] = param

        # Views preserve visibility of EPLB's in-place updates.
        shard["w13_weight_scale"] = nvfp4_swizzled_scale_to_cutedsl_mma_view(
            shard["w13_weight_scale_param"]
        )
        shard["w2_weight_scale"] = nvfp4_swizzled_scale_to_cutedsl_mma_view(
            shard["w2_weight_scale_param"]
        )
        shards.append(shard)
    return shards

interleave_linear_and_gate(x, group_size=64, dim=-1)

Interleave gate and linear weight rows for CuteDSL wrapper.

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

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

nvfp4_swizzled_scale_to_cutedsl_mma_view(scale)

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.

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

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

prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(layer, w13, w13_scale, w13_scale_2, a13_scale, w2, w2_scale, w2_scale_2, a2_scale)

Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend.

Pads the runtime expert tensors to the kernel's GEMM alignment, converts weight scale factors to the MMA layout expected by CuteDslMoEWrapper, and interleaves w13 gate/linear rows for gated activations.

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

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.

    Pads the runtime expert tensors to the kernel's GEMM alignment, converts
    weight scale factors to the MMA layout expected by CuteDslMoEWrapper, and
    interleaves w13 gate/linear rows for gated activations.
    """
    # 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)

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

    # GEMM1's output dimension must be a multiple of 128: 2I for gated
    # activations (also required by interleaving), but only I for non-gated.
    # Keep the checkpoint tensors unchanged and pad only the kernel's runtime
    # representation. Zero rows also make the padded GEMM2 contraction a no-op.
    w13, w13_scale, w2, w2_scale, padded_intermediate = align_fp4_moe_weights_for_fi(
        w13,
        w13_scale,
        w2,
        w2_scale,
        is_act_and_mul=gated,
        min_alignment=64 if gated else 128,
    )
    layer.moe_config.intermediate_size_per_partition = padded_intermediate

    # GEMM1 gathers full 256-element K tiles of activations and their scales
    # without a K-tail predicate. Pad runtime weights to match; the MoE runner
    # pads activations and trims the output using the original hidden size.
    w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi(
        w13, w13_scale, w2, w2_scale
    )
    if layer.moe_config.hidden_dim_unpadded is None:
        layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim
    layer.moe_config.hidden_dim = padded_hidden

    if gated:
        # Interleave up/gate rows for the fused gated activation.
        w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
        w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)

    # Shard after interleaving and before scale swizzling.
    from vllm.distributed.localization import get_localized_resources

    _localized_resources = get_localized_resources()
    # FlashInfer's localized half-GEMM requires a gated activation.
    if _localized_resources is not None and layer.activation.is_gated:
        layer.localized_shards = _shard_nvfp4_moe_weights_per_die(
            layer, _localized_resources, w13, w13_scale, w2, w2_scale
        )

        # Free originals before memory profiling so reclaimed storage is available
        # for KV cache. Retain shape metadata; EPLB uses the registered shards.
        if envs.VLLM_FLASHINFER_CUTEDSL_MOE_LOCALIZE_FREE_FULL_WEIGHTS:
            # Surface asynchronous sharding errors before releasing source storage.
            torch.accelerator.synchronize()
            freed = 0
            for _t in (w13, w2):
                _st = _t.untyped_storage()
                _n = _st.nbytes()
                if _n:
                    _st.resize_(0)
                    freed += _n
            if freed:
                # Exclude original weights from EPLB in favor of shards.
                layer.localized_freed_full_weight_names = (
                    "w13_weight",
                    "w2_weight",
                )
                logger.info_once(
                    "[locality domains] freeing full-width MoE weights per layer "
                    "(%.2f GiB here); shards keep the data",
                    freed / (1024**3),
                )

    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,
    )

prepare_static_weights_for_trtllm_fp4_moe(gemm1_weights, gemm2_weights, gemm1_scales_linear_fp4_bytes, gemm2_scales_linear_fp4_bytes, hidden_size, intermediate_size, num_experts, is_gated_activation)

Shuffle NVFP4 weights into the FlashInfer TRT-LLM layout.

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

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,
    )

reorder_w13_to_w31_for_flashinfer_cutedsl(activation, w13, w13_scale)

Normalize gated w13 rows to the [up; gate] order used by FlashInfer.

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

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(),
    )

reorder_w1w3_to_w3w1(weight, scale, dim=-2)

Re-order concatenated [w1, w3] tensors to [w3, w1] in-place.

weight and scale must be contiguous; they remain contiguous on return.

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

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