[IlyasMoutawwakil]

Depends on #49607.

MXFP4 GPT-OSS did not run under expert parallelism on main. Loading openai/gpt-oss-20b with DistributedConfig(tp_size=2, ep_size=2) fails twice:

• The load: rank 1 gets no expert scales (IndexError in Mxfp4Deserialize), and rank 0 gets every expert's scales. The expert params were declared with the 4-D shape of the blocks. So the loader took the 3-D scales for one expert's slice.
• The forward: the MXFP4 MLP calls the experts with Triton routing objects. ep_dispatch_experts expects (hidden_states, top_k_index, top_k_weights), so it fails with unexpected keyword argument 'scatter_idx'.

Changes

All in integrations/mxfp4.py:

• Mxfp4GptOssExperts.forward takes (hidden_states, top_k_index, top_k_weights), like every other experts class. routing_from_top_k builds the Triton routing metadata from them. It skips ids outside [0, num_experts), which are other ranks' routes. So every ep_implementation works.
• Under expert parallelism, the MLP calls the router, then the experts. It checks the experts module for the EP hook, as the EP plan installs on mlp.experts.
• On a single device, the MLP keeps the fused Triton routing, through Mxfp4GptOssExperts.forward_routed.
• The packed expert params are declared 3-D, (experts, rows, packed bytes). EP now shards the checkpoint's blocks and scales the same way, on the expert dim.
• routing_torch_dist is removed. It picked the local experts from LOCAL_RANK and the world size, not from the EP mesh.

🚨 Breaking changes

• Mxfp4GptOssExperts.forward takes (hidden_states, top_k_index, top_k_weights). The old Triton routing call is forward_routed.
• routing_torch_dist is removed.
• Mxfp4GptOssExperts.gate_up_proj and down_proj are declared (experts, rows, packed bytes) before loading, not (experts, rows, blocks, 16).

Tests

Test	Result
openai/gpt-oss-20b on 2 GPUs, ep_size=tp_size=2, against 1 GPU: 2 prompts, 48 greedy tokens	identical generations with "auto", "all_reduce" and "all_to_all"
The same with the old MLP hook check	illegal memory access, as the local experts get global ids
main and the parent PR	fail at load and at the first forward, as described above
tests/quantization/mxfp4 on 1 GPU	19 passed, 8 skipped

#49615 then removes the router-masked plan (ep_router).

🤖 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.

[github-actions]

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

run-slow: llama4

[vasqu]

Yes fine for now, imo we should deprecate them all in a follow up with finegrained but better to fix for now indeed

[vasqu]

Should be marked as 🚨 no?

[IlyasMoutawwakil]

yes done, the desc also lists the breaking changes now

[vasqu]

As mentioned internally, just let's make sure it doesnt interfere with #50245