flashinfer.megamoe.nvfp4_mega_moe

flashinfer.megamoe.nvfp4_mega_moe(y: Tensor, l1_weights: Tensor, l2_weights: Tensor, symm_buffer: MegaMoESymmBuffer, num_tokens: int | None = None, cumulative_local_expert_recv_stats: Tensor | None = None, l1_global_scales: Tensor | None = None, l2_global_scales: Tensor | None = None, force_push_dispatch: bool | None = None, force_no_clean_barrier: bool | None = None, kernel_family: str = 'auto', family_threshold: int = 256) → Tensor

Run NVFP4 MegaMoE using the fused or split family.

Parameters:

  • y (torch.Tensor) – Contiguous BF16 output with shape [at least num_tokens, hidden]. The kernel writes the first num_tokens rows in place.

  • l1_weights (torch.Tensor) – Packed NVFP4 weights from transform_nvfp4_weights_for_mega_moe().

  • l2_weights (torch.Tensor) – Packed NVFP4 weights from transform_nvfp4_weights_for_mega_moe().

  • symm_buffer (MegaMoESymmBuffer) – Buffer populated by MegaMoESymmBuffer.stage() on every rank.

  • num_tokens (int, optional) – Tokens on this rank. Defaults to y.shape[0]. Must be equal on every rank: the kernel family, tile and dispatch mode are selected from it, and all ranks must launch the same kernel. Pad short ranks with rows routed to -1.

  • cumulative_local_expert_recv_stats (torch.Tensor, optional) – INT32 vector with one entry per local expert. Receive counts are accumulated in place; initialize the vector before first use.

  • l1_global_scales (torch.Tensor, optional) – FP32 vectors with one multiplier per local expert. Omitted scales use 1.

  • l2_global_scales (torch.Tensor, optional) – FP32 vectors with one multiplier per local expert. Omitted scales use 1.

  • force_push_dispatch (bool, optional) – Override the fused counter-dispatch policy. None uses the default policy. Counter dispatch is limited to the validated small-M plans (num_tokens <= 64 and BLOCK_M <= 24); outside them the overrides are ignored and the barrier path is used. These overrides do not apply to the split family.

  • force_no_clean_barrier (bool, optional) – Override the fused counter-dispatch policy. None uses the default policy. Counter dispatch is limited to the validated small-M plans (num_tokens <= 64 and BLOCK_M <= 24); outside them the overrides are ignored and the barrier path is used. These overrides do not apply to the split family.

  • kernel_family (str) – "auto", "fused" or "split". Auto selects fused through family_threshold tokens per rank, then split.

  • family_threshold (int) – Token threshold used by auto selection. Defaults to 256.

Returns:

The output tensor y.

Return type:

torch.Tensor

Notes

All tensors must be contiguous and on the buffer’s CUDA device. NVFP4 supports four or eight ranks, which must call matching families in matching order.