flashinfer.megamoe.MegaMoESymmBuffer

class flashinfer.megamoe.MegaMoESymmBuffer(group: ProcessGroup, num_experts: int, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, activation_clamp: float = 10.0, fast_math: bool = True)

Allocate symmetric input, expert-token and output-combine storage.

All ranks must use identical geometry and invoke kernels in matching order, with the same token count on every rank (the plan is selected from it). Allocation and rendezvous are collective and must precede CUDA graph capture.

Parameters:

  • group (torch.distributed.ProcessGroup) – Expert-parallel process group over NVLink-connected Hopper GPUs.

  • num_experts (int) – Global expert count, divisible by the number of ranks.

  • num_max_tokens_per_rank (int) – Maximum input token count on each rank.

  • num_topk (int) – Number of routed expert slots per token.

  • hidden (int) – Input and output width, a multiple of 256.

  • intermediate_hidden (int) – Expert intermediate width, a multiple of 128.

  • activation_clamp (float) – SwiGLU clamp value. Defaults to 10.0.

  • fast_math (bool) – Use the family’s fast-math policy. Defaults to True.

x, x_sf, topk_idx, topk_weights

This rank’s staging views (E4M3 [capacity, hidden], FP32 [capacity, hidden // 128], INT64 [capacity, num_topk] and FP32 [capacity, num_topk]). stage() copies into them; a producer that writes the first m rows directly calls commit_staged() instead.

Type:

torch.Tensor

__init__(group: ProcessGroup, num_experts: int, num_max_tokens_per_rank: int, num_topk: int, hidden: int, intermediate_hidden: int, activation_clamp: float = 10.0, fast_math: bool = True) → None

Methods

__init__(group, num_experts, ...[, ...])

commit_staged(num_tokens)

Mark rows written directly into the staging views as this call's input.

destroy()

Release this rank's views and symmetric-memory handle.

prewarm([kernel_family, family_threshold])

Build every NVFP4 module nvfp4_mega_moe() can select up to capacity.

stage(x_fp8, x_sf, topk_idx, topk_weights)

Copy this rank's tokens and routing into the symmetric buffer.