[RFC]: Refactor AFD MoE forward around native MoERunner injection

[jiangkuaixue123]

Motivation

The current AFD model adaptation still carries model-owned MoE forward proxies and duplicated routing behavior:

• GPU and synchronous NPU Attention roles replace remote FFN/MoE execution with handwritten proxies such as RemoteFFNProxy and AFDAttentionFusedMoE;
• NPU Attention-side routing uses GateOnlyRemoteMoE plus a copied compute_gate_topk path;
• these implementations manually track upstream FusedMoE signatures, internal-router behavior, EPLB interfaces, routing options, scaling, and reduction semantics.

This is one of the highest-drift parts of model adaptation. Every vLLM or vLLM-Ascend upgrade requires the plugin to re-check behavior that upstream already owns.

vLLM 0.26 provides a native extension point: FusedMoE is a factory that accepts runner_cls and routed_experts_cls and returns a MoERunner. vLLM-Ascend already uses the same mechanism to inject AscendMoERunner.

This RFC proposes moving AFD's MoE communication boundary into role-specific runner implementations so model-level MoE forward remains native.

Goals

• Keep native model MoE forward as the source of truth.
• Remove handwritten AFD copies of routing and MoE forward contracts.
• Reuse the existing connector protocols without changing payloads or ordering.
• Preserve role-aware construction and weight loading without allocating remote expert weights.
• Share the synchronous remote-experts behavior across GPU P2P NCCL and NPU CAMP2p.
• Map NPU CAMAsync directly onto native routing plus remote dispatch/combine.
• Establish a model-independent FFN MoE computation entry point in a follow-up phase.

Non-goals

• No vLLM or vLLM-Ascend source-tree changes.
• No connector wire-protocol or control-state-machine changes.
• No behavior change for dense, non-MoE layers.
• No attempt to move FFN communication into a MoE runner.
• No new support claim for sequence-parallel MoE, unsupported graph modes, models that do not use the FusedMoE factory, or unvalidated backend/topology combinations.
• Phase 1 does not generalize the FFN-side model runner.

Proposed architecture

Inject an AFD runner through the existing FusedMoE(..., runner_cls=..., routed_experts_cls=...) factory boundary according to role, backend, connector, and gate placement.

Path	Attention-side runner behavior	FFN-side behavior
GPU P2pNcclAFDConnector, gate on FFN	AFDAttentionMoERunner(MoERunner): send hidden states, yield for DBO when required, receive the final MoE output	Existing driver receives at the layer boundary, runs native MoE, and sends the final output
NPU CAMP2pAFDConnector, gate on FFN	AFDNPUAttentionMoERunner(AscendMoERunner): same remote-experts behavior through a shared mixin	Existing driver and native MoE forward remain unchanged
NPU CAMAsyncAFDConnector, gate on Attention	AFDNPUAttentionGateMoERunner(AscendMoERunner): retain native gate and select_experts, replace dispatch/combine with CAM async send/receive, and skip local expert computation	Existing dispatched-expert kernel path remains unchanged in Phase 1

The GPU and CAMP2p implementations share an AFDRemoteExpertsForwardMixin. Their concrete classes differ only because one derives from MoERunner and the other from AscendMoERunner.

Attention roles also inject an AFDRemoteRoutedExperts implementation whose weight creation is a no-op, preventing allocation of expert weights that belong to the FFN role.

Why remote-experts runners override forward()

For gate-on-FFN paths, the connector returns the final FFN result after shared experts, routed scaling, and FFN-side reductions.

Overriding only _forward_impl() would send that final result through the remaining native MoERunner.forward() post-processing and could apply scaling or reductions twice. The remote-experts runners must therefore override the complete public forward(hidden_states, router_logits, input_ids=None) boundary and return the received result directly.

They must also report is_internal_router = True. The Attention role does not construct a gate for gate-on-FFN, but native model forward needs the internal-router branch so it does not attempt to call a missing local gate.

NPU CAMAsync mapping

CAMAsync already has dispatch/combine semantics:

1. native Ascend gate and select_experts produce topk_ids and topk_weights;
2. async_dispatch_send transfers routed tokens and routing metadata;
3. the FFN role computes experts from the dispatched layout;
4. async_combine_recv returns the weighted, combined result.

The AFD runner should reuse native Ascend routing and replace only dispatch, local expert execution, and combine. This removes GateOnlyRemoteMoE and the copied compute_gate_topk implementation.

The implementation must explicitly verify that routed scaling, top-k weighting, shared-expert output, and final reduction are each applied exactly once.

Construction and patch boundary

Plugin-owned:

• AFDRemoteExpertsForwardMixin;
• GPU and NPU Attention-side runner classes;
• AFDRemoteRoutedExperts;
• role/backend/gate-position selection;
• existing role-aware parameter filtering and connector operations.

Exact-version compatibility work:

• the scoped wrapper around vLLM-Ascend's _ascend_FusedMoE;
• any constructor patch needed to omit Attention-side gate/shared/expert allocation;
• load-order handling between the AFD and vLLM-Ascend factory wrappers.

Upstream-owned and unchanged:

• model-level DeepseekV2MoE.forward;
• native gate and expert selection;
• ordinary MoE execution on the FFN role;
• native parameter paths and supported runner interfaces.

Every compatibility patch must follow AGENTS.md: exact upstream signature and return type, an immediate patch-reason comment, marked PATCH START/END AFD deltas, focused drift tests, and a removal or upstream plan.

Dense-layer and communication boundary

This refactor applies only to MoE layers.

Dense layers retain their current behavior:

• GPU and CAMP2p Attention roles continue using the existing remote FFN proxy;
• CAMAsync dense layers remain local to Attention;
• the FFN driver retains its existing dense-layer handling.

Communication stays in the FFN driver loop. Phase 2 may unify the MoE computation entry point, but it must not move connector state or send/receive operations into FFN runner forward.

Delivery plan

Phase 1: Attention-side MoERunner reuse

1. Add AFDRemoteRoutedExperts and AFDRemoteExpertsForwardMixin.
2. Add the GPU remote-experts runner and replace AFDAttentionFusedMoE.
3. Add the CAMP2p remote-experts runner and construct an Attention-side MoE shell instead of replacing the complete MoE block.
4. Add the CAMAsync local-routing runner and remove GateOnlyRemoteMoE plus copied routing selection.
5. Validate all three connector paths against the existing FFN implementations.

Phase 1 is independently mergeable because it preserves connector payloads, ordering, and FFN behavior.

Phase 2: Model-independent FFN MoE computation

Introduce AFDFFNMoERunner(MoERunner) and AFDNPUFFNMoERunner(AscendMoERunner) so FFN drivers invoke registered MoE runners instead of model-specific methods.

For CAMAsync, the NPU runner may expose a dispatched-expert computation entry point using its own routed-expert weights, quantization state, activation, and shared experts. The driver continues to own communication state.

Phase 2 must prove model independence with at least one non-DeepSeek MoE model. Dense layers and unusual MoE implementations outside the native factory remain explicit exceptions.

Validation plan

CPU-safe tests:

• role-aware construction for GPU P2P, NPU CAMP2p, and NPU CAMAsync;
• the expected runner class is injected for each path;
• Attention roles allocate no expert/shared-expert parameters outside their required gate;
• remote-experts forward sends and receives once, preserves DBO/stage metadata, and returns the connector result without extra scaling or reduction;
• CAMAsync passes native topk_ids and topk_weights to dispatch;
• non-AFD construction and native registration remain unchanged;
• exact upstream signatures, wrapper ordering, and patch drift are checked.

Hardware validation:

• GPU P2P NCCL gate-on-FFN;
• NPU CAMP2p gate-on-FFN;
• NPU CAMAsync gate-on-Attention;
• eager plus currently supported CUDA/ACL/torchair graph paths;
• deterministic pre-/post-refactor logits comparison;
• repeated requests, DBO/uBatch where supported, failure cleanup, and process shutdown.

Risks and mitigations

• Double scaling or reduction: remote-experts paths override complete forward(); CAMAsync tests each weighting and reduction boundary.
• Unexpected Attention-side weight allocation: inject no-op routed experts and assert role parameter inventories.
• Incorrect internal-router branch: both gate-on-FFN runners explicitly return is_internal_router = True.
• vLLM-Ascend factory-wrapper ordering: install a scoped, exact-version wrapper after platform initialization and test non-AFD passthrough.
• Graph/compiler regressions: preserve registered connector custom operations and validate every currently claimed graph mode.
• CAMP2p boundary shift: verify that the experts-call input is the same post-norm hidden state previously sent at the MLP boundary.
• Premature model independence claim: require a non-DeepSeek Phase 2 validation and list dense/non-factory exceptions.

Acceptance criteria

Phase 1

• Native model MoE forward is unchanged for all three connector paths.
• AFDAttentionFusedMoE, GateOnlyRemoteMoE, and copied compute_gate_topk behavior are removed or no longer used.
• GPU P2P and NPU CAMP2p share one remote-experts forward implementation.
• NPU CAMAsync reuses native Ascend routing and replaces only remote dispatch/combine and expert execution.
• Attention roles do not allocate FFN-owned expert/shared-expert weights.
• Connector payloads, ordering, FFN driver behavior, and dense-layer behavior remain unchanged.
• CPU contracts and all three backend/connector E2E paths pass with pre-/post-refactor parity.

Phase 2

• FFN MoE computation is invoked through reusable runner objects rather than DeepSeek-specific helper methods.
• Communication state remains in the FFN driver.
• One non-DeepSeek native FusedMoE model demonstrates reuse.
• Dense layers, sequence parallelism, and non-factory MoE exceptions are documented and fail closed where necessary.

Related work

• [RFC]: Refactor DeepSeek model adapter around native vLLM lifecycle #143 — native vLLM lifecycle boundary for the DeepSeek adapter.
• [RFC]: Add Qwen3 MoE CUDA AFD support on vLLM 0.26.0 #215 — Qwen3 MoE CUDA adaptation; Phase 2 should reduce future FFN-side model-specific work.
• [RFC]: Add GPU ModelRunnerV2 Support for AFD #223 — GPU ModelRunnerV2 migration; this RFC concerns the model's MoE runner boundary and is independent of scheduler/model-runner generation.

Feedback requested

1. Is the public MoERunner.forward() boundary the correct seam for remote-experts paths?
2. Is wrapping vLLM-Ascend's factory acceptable if the wrapper is role-scoped, exact-version guarded, and transparent for non-AFD construction?
3. Should Phase 1 be split into GPU/CAMP2p first and CAMAsync second?
4. Is the native MoE registration order sufficient for Phase 2, or should the plugin maintain an explicit per-layer runner registry?
5. Which exact GPU/NPU model, topology, and graph cells must gate each phase?

Before submitting

• I searched existing issues and RFCs.
• I identified plugin-owned, compat-patch, and upstream-owned behavior.

[jiangkuaixue123]

The design doc is here.

afd_moe_runner_forward_reuse.md

@ShwStone @sinpeyw @shanyulu @jiaran-king @specture724

[jiaran-king]

We reviewed the original design against the current implementations of afd-plugin, vLLM, and vLLM-Ascend, and adjusted several implementation details where the pinned source behavior differs from the initial assumptions.

Most of the original design is retained, including:

• the MoERunner-injection direction;
• the two-phase delivery structure;
• the existing connector payload and ordering contracts;
• FFN-driver and dense-layer ownership;
• the existing CAM async-MoE scheduler and Graph/DBO support boundaries.

The main corrections are:

• implementing parameter-free experts through the quant-method create_weights() seam;
• using one ordinary remote MoERunner for GPU P2P and CAMP2p;
• making NPU Attention construction and MoE discovery role-aware without changing FFN-side Ascend EPLB;
• using the native DeepSeek gate and pinned Ascend selector for CAM;
• overriding the complete public CAM forward() while preserving the existing async scheduler.

No new connector state machine, recovery protocol, support claim, or unrelated feature is introduced.

The finalized adjusted design is attached for review and implementation.

afd_moe_runner_revision5_rationale_then_design_comment.md