Skip to content

[draft] refactor(moe): share remote MoE runners for GPU P2P. - #55

Draft
Mohammad-nassar10 wants to merge 1 commit into
stack/pr-562from
refactor/gpu-remote-moe-runner
Draft

[draft] refactor(moe): share remote MoE runners for GPU P2P.#55
Mohammad-nassar10 wants to merge 1 commit into
stack/pr-562from
refactor/gpu-remote-moe-runner

Conversation

@Mohammad-nassar10

@Mohammad-nassar10 Mohammad-nassar10 commented Nov 9, 2026 •
edited
Loading

Copy link
Copy Markdown

It's stacked on PR 562 (the base branch is a copy of its head), so the diff shows only our commit. Once PR 562 lands, we'll rebase on main and open the real PR upstream, replacing PR 406.

Stacked on vllm-project#562.

Summary

On the Attention role, MoE layers are now built by vLLM's own FusedMoEFactory instead of our hand-written AFDAttentionFusedMoE, which is removed. We plug two small pieces into the factory: experts that allocate no weights, and a runner that sends the hidden states to the FFN role (plus the router logits when the gate runs on Attention). vLLM's native MoE forward now works as is, so DeepSeek's copied forward is gone too. Outputs don't change.

This is @lirx-pd PR vllm-project#406 , ported to vLLM 0.32 and to the ffn/ layout from PR 562. The main difference: in vLLM 0.32 the runner owns the gate, so with the gate on Attention our runner applies it before sending.

  • The new code is in ffn/remote.py, used by DeepSeek V2 and Qwen3.5/3.6.
  • CUDA with P2pNcclAFDConnector only. Anything else (other devices or connectors, Attention-side EPLB, routed-expert capture) is rejected when the layer is built.
  • DeepEP M2N keeps its own path; it only gets a small hook so it keeps its RemoteFFNProxy.
  • NPU, the FFN side and the connectors are untouched.
  • Qwen3 MoE is unchanged: it never used the old class, because on Attention RemoteFFNProxy sends its whole MoE block to the FFN role.

Validation

  • tests/unit passes (1462 passed, 47 skipped), including the new ffn/test_remote.py; pre-commit passes on the changed files.
  • A new GPU test checks that the split MoE matches vLLM's native MoE bit for bit (rtol=0, atol=0), for both gate placements. It passes locally on an RTX A4500.
  • We also ran main's code and this branch on the same weights and inputs: identical outputs.
  • Not run yet: the GPU CI jobs (they only run upstream), and E2E runs with the gate on Attention and with DeepEP M2N.

Documentation

Updated the model-integration design page, one sentence in the NCCL P2P guide, and the review-pr skill notes.

@github-actions github-actions Bot added documentation Improvements or additions to documentation NVIDIA labels Nov 9, 2026
@Mohammad-nassar10
Mohammad-nassar10 marked this pull request as draft November 9, 2026 21:42
@Mohammad-nassar10
Mohammad-nassar10 force-pushed the refactor/gpu-remote-moe-runner branch from 8642321 to 08d297a Compare November 10, 2026 07:52
@Mohammad-nassar10 @lirx-pd @claude
refactor(moe): share remote MoE runners for GPU P2P
b3c24c2 
On CUDA with P2pNcclAFDConnector, the Attention role now builds the MoE
layers of the DeepSeek V2 family and Qwen3.5/3.6 through vLLM's own
FusedMoEFactory, with an AFD runner (runner_cls) and parameter-free
routed experts (routed_experts_cls), instead of the hand-made
AFDAttentionFusedMoE stand-in. This is phase 1 of RFC vllm-project#225 for GPU P2P.
It ports vllm-project#406 to vLLM v0.32.0 (5df63a72) and to the ffn/ layout of vllm-project#562.

ffn/remote.py:
- AFDRemoteMoEMethod allocates no expert weights and runs no kernel;
  AFDRemoteRoutedExperts selects it.
- AFDRemoteMoERunner sends the hidden states through the existing
  RemoteFFNProxy transport (send, DBO yield, receive).
  AFDExternalRoutingMoERunner also applies the gate it holds, as native
  MoERunner does, and sends the router logits. Both reject input_ids
  and shared_experts_input.
- build_attention_moe_runner() accepts only CUDA with
  P2pNcclAFDConnector and takes the model's gate exactly when
  compute_gate_on_attention is set. It rejects Attention-local EPLB,
  redundant experts, routed-expert capture (aux_output_config in vLLM
  0.32) and a non-finite routed_scaling_factor, and fixes the arguments
  that describe local experts.
- AFDAttentionFusedMoE is deleted.

Model adapters:
- AFDDeepseekV2RemoteExpertsMoE takes vllm_config, keeps its gate
  (assigned before the experts, so the mlp.gate.* checkpoint names hold)
  and passes it to the factory with the native 0.32 arguments. Its
  forward override is gone; the shell runs vLLM's native forward.
- The decoder-layer EPLB check moves into the factory. The DeepEP M2N
  shell keeps a parameter-free RemoteFFNProxy transport and rejects
  EPLB itself.
- AFDQwen3_5RemoteExpertsMoE matches the native constructor signature
  and builds the gate-on-FFN runner.

Outputs are unchanged: a new CUDA test checks that the Attention MoE
sends the payload its gate placement requires and returns the output of
the unsplit native MoE bit for bit (rtol=0, atol=0), for both gate
placements and for 1 and 7 tokens.

Co-authored-by: lirx-pd <616517220@qq.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Mohammad <mohammad.nassar@ibm.com>
@Mohammad-nassar10
Mohammad-nassar10 force-pushed the refactor/gpu-remote-moe-runner branch from 08d297a to b3c24c2 Compare November 12, 2026 08:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ronenkat ronenkat Awaiting requested review from ronenkat ronenkat is a code owner

Assignees

No one assigned

Labels

documentation Improvements or additions to documentation NVIDIA

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Mohammad-nassar10