flashinfer.megamoe.mxfp4_mega_moe

flashinfer.megamoe.mxfp4_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) → Tensor

Run fused MXFP4 MegaMoE with chunk-major weights and FP8 activations.

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 MXFP4 weights from transform_mxfp4_weights_for_mega_moe().

  • l2_weights (torch.Tensor) – Packed MXFP4 weights from transform_mxfp4_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.

Returns:

The output tensor y.

Return type:

torch.Tensor

Notes

All tensors must be contiguous and on the buffer’s CUDA device. The rank count must be a power of two from 1 to 32. Every rank must invoke the kernel in matching order; MXFP4 always uses the fused family.