Repository navigation
Conversation
lirx-pd
requested review from
hsliuustc0106,
jiangkuaixue123 and
ronenkat
as code owners
November 9, 2026 06:30
lirx-pd
force-pushed
the
refactor/model-ffn-boundaries
branch
from
November 9, 2026 07:12
d91ff3b to
96f66dd
Compare
Collaborator
|
hi @lirx-pd
|
lirx-pd
force-pushed
the
refactor/model-ffn-boundaries
branch
2 times, most recently
from
November 10, 2026 01:26
d768908 to
df8e91a
Compare
Contributor
Author
|
lirx-pd
force-pushed
the
refactor/model-ffn-boundaries
branch
from
November 10, 2026 01:33
df8e91a to
2f44642
Compare
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( |
There was a problem hiding this comment.
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?
lirx-pd
force-pushed
the
refactor/model-ffn-boundaries
branch
from
November 12, 2026 02:22
2f44642 to
11b8891
Compare
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>
2 tasks done
Mohammad-nassar10
approved these changes
Nov 12, 2026
Collaborator
|
please fix merge conflicts. |
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
force-pushed
the
refactor/model-ffn-boundaries
branch
from
November 14, 2026 05:33
11b8891 to
afff5cb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
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.ffn/remote.pyRemoteFFNProxy,AFDAttentionFusedMoEand routed TP paddingffn/async_cam.pyffn/expert_weights.pyQuantizedExpertWeights, homogeneous layout validation, geometry and graph signaturesffn/compute.pyexpert_weights.py; Mega receives explicit configuration and later bindscomm_argsandgroup_name.The four files are not a pipeline: Attention-side models call
remote.pyor the CAM helpers; FFN-side execution consumes exported weights throughcompute.py.remote.py,async_cam.pyandcompute.pydo not import one another. The minimalffn/__init__.pydoes 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
mainat22a1eccb142e4456ef577391524a1d977a7ac646. 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"] endffn/remote.py; generic callers do not load CAM helpers.ffn/async_cam.pyowns that boundary; model-specific routing and merge order remain explicit.models/.Preserved behavior and integration
RemoteDeepseekV4FFN, while inheriting metadata construction, stage selection and send/yield/receive fromRemoteFFNProxy. Its Attention-side mHC broadcast finalizer delegates to native vLLM; FFN skips finalization because it owns no mHC state.native_mc2:and the negotiated W8A8 bias layout. Model registration and worker/runner class paths remain unchanged.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.pydocuments both runner-side metadata installation and model/proxy access, retaining the existing package-level getter as a lazy export.v1/worker/npu/attention_metadata.pyuses responsibility-based names with explicit SFA/MLA/DSA handling.Test Plan
git diff --checkon the final change.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 intests/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
main/upstream/mainat22a1eccb142e4456ef577391524a1d977a7ac646. 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 to4ec0da3d17eefe4cffac571a8e71c2c25d88c155has SHA256ede56ed3f2618afba6afa815bc0f0a91b87e77d4e8c70139219e5d65f457e127(source manifest SHA256f64410688b50d9c9cebc402b272cdc2d8102094a5b6eae451cf0c36967d72042).git diff --check: passed.pytest -o addopts='' -q tests/unit -m 'not gpu and not vllm_runtime' -ra --tb=shortwithout 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 Babelcgideprecation.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.0.32.0+empty(the installed package does not record a source commit), vLLM-Ascend0.19.1rc2.dev3100+gbd42eea01, torch2.10.0+cpu, torch-npu2.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.KeyboardInterrupt: terminated. Its cleanup check passed, but graceful model-service shutdown is not established. Finite device tests separately verified their own normal cleanup.EngineDeadErrorafter 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.d91ff3bdc8ae438cf5a0c353a1472ded26a1922a: 12 tests passed, none skipped on three NVIDIA L20X GPUs reserved throughgpu 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.Docs Impact
Essential PR Checklist