Mega MoE

Mega MoE is MATE’s distributed FP8 Mixture-of-Experts execution path for multi-rank expert-parallel workloads. Use it when your integration already has token activations, top-k expert indices, and top-k expert weights prepared and you need MATE to run the expert compute path across a process group.

For integrations that already target DeepGEMM-style distributed MoE execution, prefer the deep-gemm wrapper first. Use the direct MATE APIs below when wrapper coverage is not enough.

The examples here use the mate.deep_gemm aliases. The same entrypoints are also available from mate.mega_moe.

Use Mega MoE when

  • you are integrating a distributed expert-parallel MoE path across multiple ranks

  • your expert weights are stored as grouped FP8 tensors

  • you need the fused two-layer expert execution path, not just one grouped GEMM call

Use the GEMM APIs instead when you only need grouped or masked expert GEMM operators, including the 0.2.4 W4A8 mixed-dtype GEMM path.

Typical Flow

Mega MoE uses a shared runtime workspace that holds FP8 inputs, routing metadata, and intermediate activations for the fused execution path.

  1. Allocate the symmetric runtime buffer with get_symm_buffer_for_mega_moe.

  2. Transform logical FP8 expert weights with transform_weights_for_mega_moe.

  3. Launch the kernel with fp8_fp8_mega_moe.

Current Limits

  • Intranode process groups only, with up to 8 ranks.

  • num_experts must be divisible by the process-group size.

  • activation="swiglu" only.

  • use_fp8_dispatch=True only.

  • hidden must be divisible by 128.

  • intermediate_hidden must be divisible by 256.

  • fp8_fp8_mega_moe writes into a contiguous torch.bfloat16 output tensor.

API

mate.deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, use_fp8_dispatch=True, activation='swiglu')

Allocate the shared runtime workspace used by Mega MoE execution.

mate.deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights=None)

Transform grouped FP8 expert weights into the layout expected by the Mega MoE runtime.

mate.deep_gemm.fp8_fp8_mega_moe(y, l1_weights, l2_weights, sym_buffer, cumulative_local_expert_recv_stats=None, recipe=(1, 1, 32), activation='swiglu', activation_clamp=None, fast_math=True)

Run the FP8 Mega MoE expert compute path and write the result to y.