vllm.models.muse_spark.nvidia.ops.prepare_megamoe_nvfp4

Triton input staging for MuseSpark1x on the NVFP4 MegaMoE kernel.

MegaMoE reads linear scales, not scaled_fp4_quant's swizzled 128x4 layout.

Functions:

• prepare_megamoe_nvfp4_inputs –

  Stage hidden_states and routing into MegaMoE's symmetric buffer.

• prepare_megamoe_nvfp4_with_norm –

  Normalize BF16 pooled states and write NVFP4 into MegaMoE's buffer.

prepare_megamoe_nvfp4_inputs(hidden_states, topk_weights, topk_ids, x_packed, x_sf, topk_idx_out, topk_weights_out, is_padding=None)

Stage hidden_states and routing into MegaMoE's symmetric buffer.

Source code in vllm/models/muse_spark/nvidia/ops/prepare_megamoe_nvfp4.py

def prepare_megamoe_nvfp4_inputs(
    hidden_states: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    x_packed: torch.Tensor,
    x_sf: torch.Tensor,
    topk_idx_out: torch.Tensor,
    topk_weights_out: torch.Tensor,
    is_padding: torch.Tensor | None = None,
) -> None:
    """Stage ``hidden_states`` and routing into MegaMoE's symmetric buffer."""
    num_tokens, hidden_size = hidden_states.shape
    if num_tokens == 0:
        return
    if hidden_size % 64 != 0:
        raise ValueError(
            "MuseSpark1x NVFP4 MegaMoE staging requires hidden_size to be a "
            f"multiple of 64, got {hidden_size}."
        )
    if topk_weights.shape != topk_ids.shape:
        raise ValueError(
            "MuseSpark1x NVFP4 MegaMoE staging requires topk_weights and "
            "topk_ids to have the same shape."
        )

    top_k = topk_ids.shape[1]
    num_elements = num_tokens * hidden_size
    if num_elements <= 512 * 1024:
        block_k, num_warps = 256, 4
    elif num_elements < 2 * 1024 * 1024:
        block_k, num_warps = 512, 2
    else:
        block_k, num_warps = 1024, 1
    block_k = min(block_k, hidden_size)
    while hidden_size % block_k != 0:
        block_k //= 2

    enable_pdl = current_platform.is_cuda() and current_platform.is_arch_support_pdl()
    _prepare_megamoe_nvfp4_kernel[(num_tokens, hidden_size // block_k)](
        hidden_states,
        x_packed,
        x_sf,
        topk_ids,
        topk_weights,
        is_padding,
        topk_idx_out,
        topk_weights_out,
        hidden_states.stride(0),
        x_packed.stride(0),
        x_sf.stride(0),
        topk_ids.stride(0),
        topk_weights.stride(0),
        is_padding.stride(0) if is_padding is not None else 0,
        topk_idx_out.stride(0),
        topk_weights_out.stride(0),
        hidden_size=hidden_size,
        top_k=top_k,
        BLOCK_K=block_k,
        BLOCK_TOPK=triton.next_power_of_2(top_k),
        USE_PDL=enable_pdl,
        num_warps=num_warps,
        **({"launch_pdl": True} if enable_pdl else {}),
    )

prepare_megamoe_nvfp4_with_norm(hidden_states, norm_weight, norm_eps, gain_center, x_packed, x_sf)

Normalize BF16 pooled states and write NVFP4 into MegaMoE's buffer.

The caller must join the routing aux stream before launching MegaMoE.

Source code in vllm/models/muse_spark/nvidia/ops/prepare_megamoe_nvfp4.py

def prepare_megamoe_nvfp4_with_norm(
    hidden_states: torch.Tensor,
    norm_weight: torch.Tensor,
    norm_eps: float,
    gain_center: float,
    x_packed: torch.Tensor,
    x_sf: torch.Tensor,
) -> None:
    """Normalize BF16 pooled states and write NVFP4 into MegaMoE's buffer.

    The caller must join the routing aux stream before launching MegaMoE.
    """
    num_tokens, hidden_size = hidden_states.shape
    assert hidden_states.dtype == torch.bfloat16 and hidden_states.is_contiguous()
    assert hidden_size == 4096
    assert norm_weight.shape == (hidden_size,) and norm_weight.is_contiguous()
    if num_tokens == 0:
        return
    enable_pdl = current_platform.is_arch_support_pdl()
    _norm_quant_megamoe_nvfp4_kernel[(num_tokens, hidden_size // 1024)](
        hidden_states,
        norm_weight,
        x_packed,
        x_sf,
        hidden_states.stride(0),
        x_packed.stride(0),
        x_sf.stride(0),
        norm_eps,
        gain_center,
        hidden_size,
        1024,
        USE_PDL=enable_pdl,
        num_warps=8,
        enable_fp_fusion=False,
        **({"launch_pdl": True} if enable_pdl else {}),
    )