[IlyasMoutawwakil]

Depends on #49429.

DeepGEMM Mega MoE fuses the token dispatch, the experts and the combine in one kernel. It needs global expert ids, unmasked routing weights and the EP process group. Until now it only got them through the megamoe_router and megamoe_experts styles. On the default ep_dispatch_experts plan, which DeepSeek V4 ships, loading with experts_implementation="deepgemm_megamoe" failed at the first forward (DeepGEMM Mega MoE requires a process_group).

This PR adds a "deepgemm_megamoe" EP implementation. "auto" runs it for experts loaded with experts_implementation="deepgemm_megamoe", since they can only run through it; another explicit implementation raises. It requires ep_size == tp_size, like the old styles.

🚨 Breaking: the megamoe_router and megamoe_experts parallel styles are removed. A plan that names them should use ep_dispatch_experts on the experts and drop the router rule; loading with experts_implementation="deepgemm_megamoe" selects the implementation.

How it runs

The kernel treats each rank's rows as that rank's own tokens. Here every rank holds the same batch.

• Each rank hands the kernel an equal slice of the batch. The last slice is padded with skipped slots.
• The slices' outputs are all-gathered back, so every token goes through the experts once.
• The symmetric-memory buffers are kept per EP group and shared by every layer. Before, each experts module allocated its own.

Example with 4 tokens on 2 ranks:

GPU                    0                      1
batch           [t0, t1, t2, t3]       [t0, t1, t2, t3]
kernel input        [t0, t1]               [t2, t3]
kernel output       [y0, y1]               [y2, y3]
after all-gather         [y0, y1, y2, y3] on both ranks

Changes

File	Change
distributed/expert_parallel.py	MegaMoeEP, registered as "deepgemm_megamoe". It passes the EP group to the experts as process_group. AutoEP builds it for Mega MoE experts.
integrations/deepgemm.py	The Mega MoE forward keeps its symmetric buffers per EP group (process_group.group_name) and shape, shared by every layer. destroy_megamoe_buffers() frees them, like DeepEP's destroy_deepep_buffers() (#49638).
integrations/tensor_parallel.py, integrations/finegrained_fp8.py	Removed: the megamoe_router and megamoe_experts styles, and the FP8 plan rewrites to them.
Docs	A "deepgemm_megamoe" row in the EP implementation table. The two styles leave the parallel styles page.

Tests

Test	Result
tests/expert_parallel, 2 ranks on gloo. The new Mega MoE test checks that "auto" runs it and "all_to_all" is refused, the padded slices, global ids, the EP group and the gathered output; the DeepGEMM kernel tests check the buffers are cached per EP group.	8 passed
TP and EP model tests of Mixtral, Qwen3-MoE, GPT-OSS and DeepSeek-V3, on CPU	273 passed
The DeepGEMM Mega MoE tests pass the cache and check the buffer lands in it.	2 passed

🤖 Generated with Claude Code

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Pages changed in this PR:

• en/expert_parallelism
• en/perf_infer_gpu_multi

[3outeille]

why do we have to do that but not before ?

[IlyasMoutawwakil]

We don't strictly have to; it bounds the reallocations. Before, each experts module kept its own buffer, sized to the batch, so every new largest batch reallocated it, once per layer. Now every layer of an EP group shares one buffer per shape, and allocating it is a collective rendezvous over symmetric memory, so it's sized to the next power of two: a batch that grows (prompts of different lengths, more sequences) reallocates a few times instead of at every new size, for at most twice the memory.

[3outeille]

how should we decide with auto which one to use at inference between all_reduce and deepgemm_megamoe ?

[IlyasMoutawwakil]

deepgemm_megamoe ep impl is locked at load time if the model uses deepgemm_megamoe experts impl, it's not part of the auto dispatch

        # DeepGEMM Mega MoE exchanges the tokens inside its kernel, through the `deepgemm_megamoe` implementation's buffers
        distributed_config = getattr(model.config, "distributed_config", None)
        if (
            getattr(model.config, "_experts_implementation", None) == "deepgemm_megamoe"
            and distributed_config is not None
        ):
            if distributed_config.ep_implementation not in ("auto", "deepgemm_megamoe"):
                raise ValueError(
                    "DeepGEMM Mega MoE exchanges tokens inside its kernel and only runs with "
                    f'`ep_implementation="deepgemm_megamoe"`, got '
                    f"`ep_implementation={distributed_config.ep_implementation!r}`."
                )
            distributed_config.ep_implementation = "deepgemm_megamoe"

[vasqu]

yep shouldnt be a decision for auto only for specific usecases, something we can share as recipe for specific models

[IlyasMoutawwakil]

update: moved this out of the quantizer, AutoEP now builds MegaMoeEP itself when the experts use deepgemm_megamoe. it's forced by the experts impl (they can't run any other way), not a perf choice, so it fits auto. an explicit other ep_implementation raises in install_forward

[vasqu]

Actual question but what is the benefit of keeping deepgemm? I would expect we switch to deep ep anyways and deepgemm is hard to maintain with lots of restrictions.

Imo, it might be more worth it to drop hard restrictions like these? Or am i too eager to drop it lol

[IlyasMoutawwakil]

deepgemm might still be the best way to serve w4a8 mxfp4 experts (dsv4), once we have deepEP and fused moe kernels, we can compare and drop it if it loses 🫣

[vasqu]

ok fair enough, i just hate deepgemm atp (sorry not sorry 😂)

[vasqu]

Some small comments but overall fine

[vasqu]

yep shouldnt be a decision for auto only for specific usecases, something we can share as recipe for specific models

[vasqu]

breaking 🚨 but for good reasons

[IlyasMoutawwakil]

yep, added 🚨 to the title and a breaking note in the description

[vasqu]

oh interesting so we can make this less restrictive under DP? guess we need to check the dp mesh additionally then (additional kwarg)

[IlyasMoutawwakil]

yes, megamoe is built for dp attention + ep where each rank has its own batch, then neither the slice nor the gather is needed. inference has no dp yet, so it stays a todo until it does, and that's where the dp mesh check would go

[vasqu]

maybe a dedicated function, id expect deep ep to have something similar so anything that clears all buffers

[IlyasMoutawwakil]

done, destroy_megamoe_buffers() in integrations/deepgemm.py, same as deepep's destroy_deepep_buffers(). the buffers live there now, per ep group, so MegaMoeEP doesn't hold them anymore

[vasqu]

should we maybe just adjust the experts forward to allow

1. rows
2. arbitrary kwargs (like the buffers)

[IlyasMoutawwakil]

partly: it now only passes process_group= by name, the megamoe forward looks up its buffers by group itself. i'd keep the other experts forwards as they are though, the extra arg only exists for the megamoe kernel, and the rows are just a slice of the inputs so nothing to add there

[vasqu]

should we be stricter and disallow auto?

[IlyasMoutawwakil]

i'd keep auto: loading with experts_implementation="deepgemm_megamoe" should just work without also setting ep_implementation. AutoEP now builds MegaMoeEP itself for those experts (no more config rewrite in the quantizer), and any other explicit implementation raises when the plan is applied

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: finegrained_fp8