Mario/overlap moe kernel integration

benchmarks/mega_moe/bench_mega_moe.py

@@ -0,0 +1,176 @@
+"""Benchmark prime-mega-moe's fused Mega MoE kernels (bf16 forward + backward) against prime-rl's
+bf16 EP path (``TorchTokenDispatcher`` + ``BF16ExpertCompute``) for the same routed-expert
+workload, weights, and token counts, and check the gradients agree.
+
+Requires >=2 GPUs with symmetric-memory support (SM100/Blackwell), PyTorch >= 2.9, and the
+prime-mega-moe ``deep_gemm`` wheel from Prime Intellect installed over the public one.
+``--hidden`` must be a multiple of 256 and ``--intermediate`` a multiple of 128.
+
+Usage:
+    torchrun --nproc_per_node=8 benchmarks/mega_moe/bench_mega_moe.py \
+        --num-experts 128 --top-k 8 --hidden 2048 --intermediate 1024 --tokens-per-rank 2048
+"""
+
+import argparse
+import os
+
+import torch
+import torch.distributed as dist
+
+from prime_rl.trainer.distributed.token_dispatcher import FusedTokenDispatcher, TorchTokenDispatcher
+from prime_rl.trainer.models.layers.expert_compute import BF16ExpertCompute
+from prime_rl.trainer.models.layers.mega_moe import MegaMoEExpertCompute, mega_moe_available
+from prime_rl.trainer.models.layers.moe import GroupedExperts
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--num-experts", type=int, default=128)
+    parser.add_argument("--top-k", type=int, default=8)
+    parser.add_argument("--hidden", type=int, default=2048)
+    parser.add_argument("--intermediate", type=int, default=1024)
+    parser.add_argument("--tokens-per-rank", type=int, default=2048)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=50)
+    return parser.parse_args()
+
+
+def bench(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    dist.barrier()
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def rel_diff(a: torch.Tensor, b: torch.Tensor) -> float:
+    return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)).item()
+
+
+def main() -> None:
+    args = parse_args()
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend="nccl", device_id=torch.device("cuda", local_rank))
+    rank, world_size = dist.get_rank(), dist.get_world_size()
+    group = dist.group.WORLD
+
+    if not mega_moe_available():
+        raise RuntimeError(
+            "Mega MoE requires the prime-mega-moe deep_gemm wheel from Prime Intellect installed over the public one, on an SM100+ GPU."
+        )
+    if args.num_experts % world_size:
+        raise ValueError(f"num_experts ({args.num_experts}) must be divisible by world_size ({world_size}).")
+
+    torch.manual_seed(rank)
+    experts_per_rank = args.num_experts // world_size
+    device = torch.device("cuda", local_rank)
+
+    experts = GroupedExperts(
+        dim=args.hidden,
+        hidden_dim=args.intermediate,
+        num_experts=experts_per_rank,
+    ).to(device=device, dtype=torch.bfloat16)
+    experts.init_weights(init_std=0.02)
+    params = [p for p in experts.parameters() if p is not None]
+
+    x = torch.randn(args.tokens_per_rank, args.hidden, dtype=torch.bfloat16, device=device, requires_grad=True)
+    dy = torch.randn(args.tokens_per_rank, args.hidden, dtype=torch.bfloat16, device=device)
+    scores = torch.softmax(
+        torch.randn(args.tokens_per_rank, args.num_experts, device=device, dtype=torch.float32), dim=-1
+    )
+    # Keep the model's (num_tokens, top_k) routing shape: the Mega MoE backward returns dtop_scores in it.
+    top_scores, selected_experts_indices = torch.topk(scores, args.top_k, dim=-1)
+    top_scores = top_scores.contiguous().requires_grad_(True)
+    selected_experts_indices = selected_experts_indices.contiguous()
+
+    bf16_compute = BF16ExpertCompute()
+    baseline = (
+        TorchTokenDispatcher(
+            num_experts=args.num_experts,
+            top_k=args.top_k,
+            token_group_alignment=bf16_compute.token_group_alignment,
+            group=group,
+        ),
+        bf16_compute,
+    )
+    mega_compute = MegaMoEExpertCompute(
+        experts,
+        num_experts=args.num_experts,
+        top_k=args.top_k,
+        group=group,
+        max_tokens_per_rank=args.tokens_per_rank,
+    )
+    mega = (FusedTokenDispatcher(), mega_compute)
+
+    def run(path, backward: bool):
+        dispatcher, compute = path
+
+        def fn():
+            experts.compute = compute
+            y = dispatcher.run(x, top_scores, selected_experts_indices, experts, score_before_experts=False)
+            if backward:
+                y.backward(dy)
+            return y
+
+        return fn
+
+    def grads_after(fn):
+        for p in params:
+            p.grad = None
+        x.grad = None
+        top_scores.grad = None
+        y = fn()
+        torch.cuda.synchronize()
+        return y.detach().clone(), x.grad.clone(), top_scores.grad.clone(), [p.grad.clone() for p in params]
+
+    y_ref, dx_ref, ds_ref, dp_ref = grads_after(run(baseline, backward=True))
+    y_mega, dx_mega, ds_mega, dp_mega = grads_after(run(mega, backward=True))
+    diffs = {
+        "y": rel_diff(y_mega, y_ref),
+        "dx": rel_diff(dx_mega, dx_ref),
+        "dtop_scores": rel_diff(ds_mega, ds_ref),
+        "dexperts": max(rel_diff(a, b) for a, b in zip(dp_mega, dp_ref)),
+    }
+
+    def bench_fwd_bwd(path):
+        fwd = bench(run(path, backward=False), args.warmup, args.iters)
+        total = bench(run(path, backward=True), args.warmup, args.iters)
+        return fwd, total
+
+    base_fwd, base_total = bench_fwd_bwd(baseline)
+    mega_fwd, mega_total = bench_fwd_bwd(mega)
+
+    if rank == 0:
+        print(
+            f"world_size={world_size} num_experts={args.num_experts} (per rank={experts_per_rank}) "
+            f"top_k={args.top_k} hidden={args.hidden} intermediate={args.intermediate} "
+            f"tokens_per_rank={args.tokens_per_rank}"
+        )
+        print(f"{'':46} {'fwd':>10} {'bwd':>10} {'fwd+bwd':>10}")
+        print(
+            f"{'bf16 EP (TorchTokenDispatcher+BF16ExpertCompute)':46} "
+            f"{base_fwd:9.3f}ms {base_total - base_fwd:9.3f}ms {base_total:9.3f}ms"
+        )
+        print(
+            f"{'Mega MoE (fused bf16 fwd + fused bf16 bwd)':46} "
+            f"{mega_fwd:9.3f}ms {mega_total - mega_fwd:9.3f}ms {mega_total:9.3f}ms"
+        )
+        print(
+            f"{'speedup':46} {base_fwd / mega_fwd:9.2f}x "
+            f"{(base_total - base_fwd) / (mega_total - mega_fwd):9.2f}x {base_total / mega_total:9.2f}x"
+        )
+        print("relative diff vs bf16 EP path: " + ", ".join(f"{k}={v:.4f}" for k, v in diffs.items()))
+
+    mega_compute.buffer.destroy()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

docs/scaling.md

@@ -107,6 +107,20 @@ transport = "bf16"
 
 For DeepEP, set `type = "deepep"` and tune `num_sms` plus optional `token_chunk_size` in the same dispatch table. Routed-expert precision is selected separately with `[trainer.model.moe.compute]` (`bf16`, `deepgemm_fp8`, or `mxfp8`).
 
+Mega MoE (`type = "mega_moe"`) fuses dispatch, the bf16 SwiGLU expert MLP, and combine into one persistent kernel per direction, overlapping the expert-parallel communication with the GEMMs; `[trainer.model.moe.compute]` does not apply to those layers. The kernels come from prime-mega-moe, Prime Intellect's closed-source DeepGEMM fork. It is distributed to partners as a prebuilt `deep_gemm` wheel that replaces the public one; install it after syncing with `uv pip install --no-deps <wheel>`. Without it, selecting `mega_moe` fails at startup with a message naming the missing kernels. Mega MoE needs EP > 1, SM100 GPUs, nvcc 13 on `CUDA_HOME` for the runtime JIT, gated experts without biases, `hidden_size` divisible by 256 and `moe_intermediate_size` by 128:
+
+```toml
+[trainer.model.moe.dispatch]
+type = "mega_moe"
+max_tokens_per_rank = 2048  # >= micro_batch_size * seq_len on every rank
+num_reserved_sms = 16       # SMs left free for NCCL; pair with NCCL_MAX_CTAS <= this
+
+[trainer.env_vars]
+NCCL_MAX_CTAS = "16"
+CUDA_HOME = "/usr/local/cuda-13.0"
+PATH = "/usr/local/cuda-13.0/bin:${PATH}"
+```
+
 ### Context Parallelism
 
 CP shards a single sequence across multiple GPUs along the token dimension — for long-context sequences. Prefer `ulysses`: it gets the most throughput and is the only style that works for hybrid linear-attention/Mamba models (Qwen3.5, NemotronH) and for VLMs. Each model class declares its supported styles in `cp_support`; an unsupported `cp_style`, or a model with no CP support at all, is rejected at setup.

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -253,8 +253,27 @@ class DeepEPMoEDispatchConfig(BaseConfig):
     """Optional chunk size used to pipeline dispatch with local expert compute."""
 
 
+class MegaMoeMoEDispatchConfig(BaseConfig):
+    type: Literal["mega_moe"] = "mega_moe"
+    precision: Literal["bf16", "mxfp8"] = "bf16"
+    """Precision of the fused expert GEMMs. ``mxfp8`` runs forward and backward on the MXFP8 kernels
+    (E4M3 data with 1x32 UE8M0 block scales, like the ``mxfp8`` expert compute); expert weights and
+    routed activations are quantized per call, gradients stay BF16/FP32. SM100 only."""
+
+    max_tokens_per_rank: int = Field(8192, ge=1)
+    """Upper bound on routed tokens per rank per forward call, used to size Mega MoE's symmetric
+    buffer once at startup. Must be >= the largest `bs * slen` any rank will pass through a MoE
+    layer; raise it if you hit a "buffer is sized for N tokens/rank" error."""
+
+    num_reserved_sms: int = Field(16, ge=0)
+    """SMs left free for concurrent NCCL kernels (FSDP all-gathers etc.). The Mega MoE kernels are
+    persistent grids that synchronize across ranks; if they occupied every SM while an NCCL kernel
+    on another stream was waiting for a peer, the two would deadlock. Pair with ``NCCL_MAX_CTAS``
+    <= this value in ``env_vars`` so every NCCL kernel fits in the reserved SMs."""
+
+
 MoEDispatchConfig: TypeAlias = Annotated[
-    TorchMoEDispatchConfig | DeepEPMoEDispatchConfig,
+    TorchMoEDispatchConfig | DeepEPMoEDispatchConfig | MegaMoeMoEDispatchConfig,
     Field(discriminator="type"),
 ]
 
@@ -384,6 +403,8 @@ def validate_moe_runtime(self):
         if isinstance(dispatch, DeepEPMoEDispatchConfig):
             if isinstance(compute, MXFP8MoEComputeConfig):
                 raise ValueError("MXFP8 expert compute does not support DeepEP dispatch.")
+        elif isinstance(dispatch, MegaMoeMoEDispatchConfig):
+            pass  # Compute is fused into dispatch; `model.moe.compute` is ignored for these layers.
         elif dispatch.transport == "mxfp8":
             if not isinstance(compute, MXFP8MoEComputeConfig):
                 raise ValueError("MXFP8 transport requires model.moe.compute.type='mxfp8'.")

src/prime_rl/trainer/activation_checkpointing.py

@@ -69,6 +69,7 @@
         "aten::mm",
         "prime_rl::dsv4_sparse_attn",
         "prime_rl::fp8_blockwise_mm",
+        "prime_rl::mega_moe_forward",
         "prime_rl::grouped_fp8_gemm",
         "prime_rl::sparse_mla",
         "quack::gemm_act_out",

src/prime_rl/trainer/distributed/token_dispatcher.py

@@ -14,7 +14,7 @@
 
 
 class ExpertFunction(Protocol):
-    def __call__(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: ...
+    def __call__(self, x: torch.Tensor, *routing: torch.Tensor) -> torch.Tensor: ...
 
 
 class TokenDispatcher(Protocol):
@@ -31,6 +31,28 @@ def run(
     def synchronize(self) -> None: ...
 
 
+class FusedTokenDispatcher:
+    """Passes undispatched tokens and their routing to an expert compute that dispatches and combines them itself."""
+
+    def run(
+        self,
+        x: torch.Tensor,
+        top_scores: torch.Tensor,
+        selected_experts_indices: torch.Tensor,
+        experts: ExpertFunction,
+        *,
+        score_before_experts: bool,
+    ) -> torch.Tensor:
+        if score_before_experts:
+            raise ValueError(
+                "Fused dispatch applies router weights at combine time and requires score_before_experts=False."
+            )
+        return experts(x, top_scores, selected_experts_indices)
+
+    def synchronize(self) -> None:
+        return None
+
+
 DispatchState = TypeVar("DispatchState")
 
 

src/prime_rl/trainer/models/layers/expert_compute.py

@@ -26,6 +26,25 @@ def __call__(
     ) -> torch.Tensor: ...
 
 
+class FusedDispatchExpertCompute(Protocol):
+    """Expert compute that also dispatches tokens to their experts and combines the results.
+
+    It takes this rank's undispatched tokens and their routing, paired with a ``FusedTokenDispatcher``.
+    """
+
+    token_group_alignment: int
+
+    def validate(self, experts: "GroupedExperts") -> None: ...
+
+    def __call__(
+        self,
+        experts: "GroupedExperts",
+        x: torch.Tensor,
+        top_scores: torch.Tensor,
+        selected_experts_indices: torch.Tensor,
+    ) -> torch.Tensor: ...
+
+
 def broadcast_expert_bias(
     bias: torch.Tensor,
     num_tokens_per_expert: torch.Tensor,