Skip to content

refactor: decouple shared FFN components from model adapters - #562

Open
lirx-pd wants to merge 4 commits into
vllm-project:mainfrom
lirx-pd:refactor/model-ffn-boundaries
Open

lirx-pd wants to merge 4 commits into
vllm-project:mainfrom
lirx-pd:refactor/model-ffn-boundaries

Conversation

@lirx-pd

@lirx-pd lirx-pd commented Nov 9, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

Decouple shared FFN communication and compute from concrete model adapters. Qwen and Ascend V4 no longer import DeepSeek V2 for generic proxies, and Mega initializes directly from loaded expert tensors without constructing Layered or retaining the connector.

Issue

Scope

  • In scope: Shared remote FFN calls, expert-weight contracts, Layered/Mega construction, checkpoint filtering, forward context, staged metadata ownership, affected callers, tests, fakes, E2E observers and documentation.
  • Out of scope: Model or native-kernel numerical changes, new execution layers, new monkey patches, model factories or registries, support-matrix expansion, unrelated validation consolidation and graph transaction reorganization.

Implementation Notes

FFN component relationships

model_executor/models/ owns model structure, role-specific construction, checkpoint policy and forward semantics. model_executor/ffn/ owns FFN mechanisms shared across model families, including Attention-side remote calls and FFN-side compute. The execution chain remains Executor → Worker → ModelRunner; this package adds no execution layer.

File Responsibility Direct consumers and dependencies
ffn/remote.py Connector-independent RemoteFFNProxy, AFDAttentionFusedMoE and routed TP padding Used by DeepSeek V2, Qwen3, Qwen3.5 and CUDA/Ascend V4. Depends on connector contracts, public vLLM helpers, shared forward context and the existing DBO yield hook.
ffn/async_cam.py CAM model-side TP/SP layout conversion, padding, output restoration and native DBO receive handoff Used by the V2/V4 CAM adapters. Depends on connector contracts, TP collectives and native DBO hooks. Connector/Session retains transport windows, pending work and Push/Pull protocol state.
ffn/expert_weights.py QuantizedExpertWeights, homogeneous layout validation, geometry and graph signatures Produced by V4/DSpark weight exports; consumed by the Runner and both compute implementations. Holds tensor references without native operator lookup or compute workspace allocation.
ffn/compute.py Independent Layered and Mega implementations, native arguments and owned workspace The existing FFN ModelRunner constructs the selected implementation directly. Both consume expert_weights.py; Mega receives explicit configuration and later binds comm_args and group_name.

The four files are not a pipeline: Attention-side models call remote.py or the CAM helpers; FFN-side execution consumes exported weights through compute.py. remote.py, async_cam.py and compute.py do not import one another. The minimal ffn/__init__.py does not eagerly import them, and no FFN module imports a concrete model or Worker/ModelRunner class. Mega neither constructs Layered nor retains a connector reference.

CAM helpers live in ffn/ because they adapt model tensors to a remote FFN call across model families. One module keeps layout preparation, raw receive handoff and output restoration together while retaining their separate call sites. Gate, Hash routing, shared experts, dense-layer behavior, HC/scaling and DSpark target/draft identities stay in their concrete model adapters.

model_executor before and after

The before view reflects main at 22a1eccb142e4456ef577391524a1d977a7ac646. Arrows show dependencies; dashed arrows identify retained object references.

flowchart LR
    subgraph After["After"]
        AA["Model Attention adapter"] --> AP["ffn/remote.py<br/>Shared remote proxy"]
        AF["Model FFN adapter"] --> AW["ffn/expert_weights.py<br/>Shared weight contract"]
        AR["FFN ModelRunner"] --> AL["Layered compute"]
        AR --> AM["Mega compute"]
        AL --> AW
        AM --> AW
        AC["Connector"] -->|Bind communication parameters after initialization| AM
    end

    subgraph Before["Before"]
        BA["Qwen / NPU V4<br/>Attention adapters"] --> BD["DeepSeek V2 model module<br/>Generic remote proxies"]
        BF["V4 FFN adapter"] --> BW["LayeredFFNWeights<br/>Defined in the Layered module"]
        BR["FFN ModelRunner"] --> BL["Layered executor"]
        BR --> BM["Mega executor"]
        BL --> BW
        BM -->|Reuses weights and geometry| BL
        BM -->|Holds| BC["Connector instance"]
    end
Loading
Maintenance concern Before After
Reuse generic proxies Qwen and Ascend V4 import the DeepSeek V2 model module. Models import ffn/remote.py; generic callers do not load CAM helpers.
Change CAM call mechanics Layout and DBO helpers live under model adapters. One model-independent ffn/async_cam.py owns that boundary; model-specific routing and merge order remain explicit.
Change Mega or FFN graphs Mega obtains weights through a Layered instance and reads a full connector. Both compute implementations consume the same weight contract; the Runner owns selection, graph lifecycle and communication binding.
Add a model Shared checkpoint parsing/filtering is copied across adapters. Reuse parsing/filtering and remote components; implement the model's role policy and, for Async CAM native compute, its expert-weight export. Native constructor gaps still require model-specific adaptation.
Maintain runtime metadata Runners import context and stage-buffer helpers from models/. Shared context is package-level; Ascend metadata and stage RoPE lifetime are worker-owned.

Preserved behavior and integration

  • CUDA V4 keeps token-ID validation and its missing-context error in RemoteDeepseekV4FFN, while inheriting metadata construction, stage selection and send/yield/receive from RemoteFFNProxy. Its Attention-side mHC broadcast finalizer delegates to native vLLM; FFN skips finalization because it owns no mHC state.
  • Preserve native argument order and values, tensor references, layer IDs, quantization, clamp/scaling behavior and graph signatures, including native_mc2: and the negotiated W8A8 bias layout. Model registration and worker/runner class paths remain unchanged.
  • Bind Mega communication arguments after connector initialization and before warmup/capture. Binding failures reach peers through Startup and retain resources. Mega W8A8 avoids Layered's unused zero bias.
  • Replace five copied checkpoint parsers and five filtering loops with models/_weight_utils.py. Each adapter keeps its classifier; DSpark keeps its distinct draft-prefix policy. No copied upstream constructor is removed: existing constructor hooks are reused, and necessary role-aware copies remain where native hooks cannot omit unused modules.
  • afd_plugin/forward_context.py documents both runner-side metadata installation and model/proxy access, retaining the existing package-level getter as a lazy export. v1/worker/npu/attention_metadata.py uses responsibility-based names with explicit SFA/MLA/DSA handling.
  • Update callers, fakes, test ownership and device evidence readers together. The Mega E2E observer associates the selected compute instance with its connector in test code; production compute has no connector back-reference. No forwarding modules, factory, registry or new monkey patch is introduced.

Test Plan

  • Run the required hooks and git diff --check on the final change.
  • Run CPU unit suites with and without Torch, preserving dependency skips.
  • Exercise independent Mega startup, communication-binding failure publication/resource retention, shared checkpoint mechanics, weight geometry/signatures and existing compute contracts.
  • Compare W4A8/W8A8 geometry and graph signatures with the base implementation across quantization, clamp, layer-ID and scale-layout variations.
  • Run the affected model/proxy, Qwen and DSpark unit cases with the installed vLLM/Ascend dependencies.
  • Run matched native Layered/Mega × eager/FULL × W4A8/W8A8 finite cases, checking outputs, capture/replay, shutdown and idle drain.
  • Run matched DeepSeek-V4 Flash W4A8 real-model generation with Layered eager and Mega FULL, including the changed E2E observer.
  • Validate matched real-model requests with active Attention DBO.
  • Validate DSpark target/draft generation.
  • Validate GPU non-CAM callers.

Existing compute, layout, DBO and metadata cases are retained under their source owners. CAM cases are in tests/unit/model_executor/ffn/test_async_cam.py; staged metadata cases are in tests/unit/v1/worker/npu/test_attention_metadata.py. The FFN test package prevents its CAM module name from colliding with connector tests under the default pytest import mode. New host coverage checks the changed startup boundary and one-pass filtering; it does not simulate NPU correctness. Model construction and role-policy coverage is reused rather than replaced with source-text assertions.

Test Result

  • Base: main / upstream/main at 22a1eccb142e4456ef577391524a1d977a7ac646. Branch: refactor/model-ffn-boundaries. Head commit: d91ff3bdc8ae438cf5a0c353a1472ded26a1922a. CPU/static validation and the A3 unit, finite-device, real-model, DBO and DSpark results cover this commit; its patch relative to 4ec0da3d17eefe4cffac571a8e71c2c25d88c155 has SHA256 ede56ed3f2618afba6afa815bc0f0a91b87e77d4e8c70139219e5d65f457e127 (source manifest SHA256 f64410688b50d9c9cebc402b272cdc2d8102094a5b6eae451cf0c36967d72042).
  • Required pre-commit hooks, including Ruff, markdownlint, SPDX and Python 3.11 mypy: passed. git diff --check: passed.
  • pytest -o addopts='' -q tests/unit -m 'not gpu and not vllm_runtime' -ra --tb=short without Torch: 661 passed, 91 skipped, 8 deselected, 1 warning, 10 subtests passed. Environment: pytest 9.0.3, setuptools 80.10.2; Torch/vLLM/Ascend dependencies are unavailable. The warning is an external Babel cgi deprecation.
  • The same suite with Torch 2.14.0+cpu, pytest 9.1.1 and setuptools 80.10.2: 786 passed, 71 skipped, 8 deselected, 10 subtests passed. vLLM, vLLM-Ascend, torch-npu, NumPy and CMake are unavailable. This CPU Torch version is not the target Ascend runtime.
  • CUDA V4 method-isolation comparison against commit 4ec0da3: 13 proxy cases and 16 mHC cases matched, including payload identity, stage fallback and retention across yield, original errors, broadcast values and stable buffer identity. These checks execute the actual method bodies with isolated dependencies; they do not validate full vLLM imports or device execution. The V4 proxy/construction pytest modules remain dependency-skipped locally because vLLM is unavailable.
  • Exact base/candidate signature and geometry comparison: 36 cases passed, covering W4A8 per-channel/per-group, W8A8, clamp 0/10, contiguous/gapped/reversed layer IDs, and contiguous/transposed/gapped scale layouts. This is CPU contract evidence only.
  • A3 target environment: vLLM 0.32.0+empty (the installed package does not record a source commit), vLLM-Ascend 0.19.1rc2.dev3100+gbd42eea01, torch 2.10.0+cpu, torch-npu 2.10.0.post7, CANN 9.2.0. The affected CPU unit selection passed 397 cases on the base and 406 on the candidate, with 6 deselected and 17 warnings each.
  • A3 finite synthetic-weight integration: 16/16 runs passed across W4A8/W8A8, Layered/Mega, eager/FULL and base/candidate. Two Attention ranks (TP2) and two FFN ranks used layers 2 and 5; all 80 candidate output tensors matched base eager bit for bit. All FULL FFN ranks captured and replayed the graph; all rank cleanup checks and Mega idle-drain checks passed. This does not validate checkpoint loading or model accuracy.
  • The native package was reused after verifying all 241 package files against two retained copies and all 47 recorded native hashes. Native sources are unchanged relative to the base; no fresh native build was run.
  • A3 real-checkpoint generation: DeepSeek-V4 Flash W4A8, Attention TP4 × DP1 and FFN EP8 × TP1, capacity 128. Base and candidate each ran Layered eager and Mega FULL against the same five fixed cases (six completions, eight generated tokens each). All four runs passed; generated token IDs matched the base Layered eager outputs exactly. The candidate Mega FULL run recorded Mega on all eight FFN ranks, each capturing one FULL graph with non-idle request-window replays. These are bounded functional checks, not an accuracy or performance benchmark.
  • The model harness terminated the service processes after its requests; the FFN process exited with code 1 after KeyboardInterrupt: terminated. Its cleanup check passed, but graceful model-service shutdown is not established. Finite device tests separately verified their own normal cleanup.
  • A3 real-checkpoint Attention DBO: DeepSeek-V4 Flash W4A8, Attention DP2 × TP4 and FFN EP8 × TP1, FFN eager. Base and candidate each answered 10 overlapping arithmetic requests correctly with identical prompts and answers; all eight Attention ranks recorded live two-stage execution. This is a bounded functional check, not an accuracy benchmark.
  • In both DBO runs, service termination produced EngineDeadError after shutdown began; the FFN process manager timed out and force-killed eight workers, and the resource tracker reported semaphore and shared-memory cleanup warnings. There were no error lines before shutdown in either run. These results establish request-time DBO behavior only; they do not establish graceful shutdown.
  • A3 real-checkpoint DSpark eager generation: DeepSeek-V4 Flash W4A8, Attention DP1 × TP1 and FFN EP4 × TP1, five draft tokens. Base and candidate both passed five fixed cases, including SWA boundaries, chunked prefill and a two-prompt batch; their generated token IDs matched exactly. Concurrent requests, stop and stream-abort/repeat checks also passed. Each FFN rank returned all 43 target and three draft layers; role evidence confirmed the expected Attention/FFN ownership. The candidate recorded 57 draft rounds and 266 accepted tokens. The observer synchronizes device data, so this is functional evidence only; DSpark graph numerics and performance were not measured.
  • DSpark test cleanup found no remaining owned processes and required no forced cleanup; FFN exited with code 1 after termination in both runs. Graceful service shutdown is not established.
  • GPU non-CAM validation at head d91ff3bdc8ae438cf5a0c353a1472ded26a1922a: 12 tests passed, none skipped on three NVIDIA L20X GPUs reserved through gpu run, with vLLM 0.32.0, PyTorch 2.13.0+cu130 and CUDA 13.0. The selection included two Attention routing cases and ten real-model E2E cases using the synchronous P2P connector.
  • DeepSeek-V2-Lite passed 1A1F eager/graph and 2A1F eager/graph/graph+DBO; Qwen3-30B-A3B passed 2A1F eager/graph/graph+DBO; Qwen3.6-35B-A3B passed 2A1F eager/graph. All graph cases recorded live replays on every participating rank; both graph+DBO cases recorded live two-stage graph replays on every rank.
  • GSM8K used eight-shot prompts and seven samples per ordinary case or 24 per DBO case; all selected cases met their thresholds. These are bounded functional checks, not full-dataset accuracy or performance qualification.

Docs Impact

  • Files updated: Model integration, FFN runtime, Attention runtime, execution platforms, connector contracts, compatibility/patches and NPU testing documentation.
  • Document each FFN module and its dependencies, before/after model-executor architecture diagrams, maintenance benefits, direct Mega construction and communication binding, retained constructor copies, and updated source/test ownership paths.
  • If none, reason: Not applicable.

Essential PR Checklist
  • Purpose is clear and linked to public context when possible.
  • Scope is bounded.
  • Compatibility with vLLM v0.32.0 and, for NPU changes, the pinned vLLM-Ascend source is considered, as applicable.
  • No changes are made to the vLLM source checkout.
  • Plugin-owned classes or explicit dotted class paths are preferred over monkey patches.
  • Any compat shim or monkey patch is isolated, idempotent, version-guarded, documented, and tested.
  • Imports remain CPU-safe; CUDA-heavy work is delayed or GPU-gated.
  • Validation evidence is included, including skipped GPU tests when applicable.
  • Documentation impact is stated.

@github-actions github-actions Bot added documentation Improvements or additions to documentation Ascend Ascend NPU platform and related changes NVIDIA NVIDIA GPU, CUDA, and related changes labels Nov 9, 2026
@lirx-pd
lirx-pd force-pushed the refactor/model-ffn-boundaries branch from d91ff3b to 96f66dd Compare November 9, 2026 07:12
@ronenkat

ronenkat commented Nov 9, 2026

Copy link
Copy Markdown
Collaborator

hi @lirx-pd

  1. can you connect the PR to an issue to help with a proper scope.
  2. Is this the NPU counterpart of refactor(moe): share remote MoE runners for GPU P2P #406 ? (including common sections)
  3. There are significant NPU specific code, could we differentiate with a ffn/npu/ prefix, e.g., ffn/npu/compute.py, and have the common code (architecture independent) in the ffn/ folder?

@lirx-pd
lirx-pd force-pushed the refactor/model-ffn-boundaries branch 2 times, most recently from d768908 to df8e91a Compare November 10, 2026 01:26
@lirx-pd

lirx-pd commented Nov 10, 2026

Copy link
Copy Markdown
Contributor Author
  1. Add related issue [RFC]: Clarify AFD connector, model, and validation boundaries #565
  2. This Pr is not related with refactor(moe): share remote MoE runners for GPU P2P #406
  3. The file structure has been changed

@lirx-pd
lirx-pd force-pushed the refactor/model-ffn-boundaries branch from df8e91a to 2f44642 Compare November 10, 2026 01:33
Mohammad-nassar10 added a commit to ms-llmd/afd-plugin that referenced this pull request Nov 10, 2026
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>
from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights


def _iter_role_weights(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this is the last role-filter copy. Could it use filter_role_weights from _weight_utils.py, with this MTP rule as its classify function?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

@lirx-pd
lirx-pd force-pushed the refactor/model-ffn-boundaries branch from 2f44642 to 11b8891 Compare November 12, 2026 02:22
Mohammad-nassar10 added a commit to ms-llmd/afd-plugin that referenced this pull request Nov 12, 2026
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>
@ronenkat

Copy link
Copy Markdown
Collaborator

please fix merge conflicts.
Looks good.

Signed-off-by: lirx-pd <616517220@qq.com>
Signed-off-by: lirx-pd <616517220@qq.com>
Signed-off-by: lirx-pd <616517220@qq.com>
Signed-off-by: lirx-pd <616517220@qq.com>
@lirx-pd
lirx-pd force-pushed the refactor/model-ffn-boundaries branch from 11b8891 to afff5cb Compare November 14, 2026 05:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Ascend Ascend NPU platform and related changes documentation Improvements or additions to documentation NVIDIA NVIDIA GPU, CUDA, and related changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants