[WIP] feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary by ksiyuan · Pull Request #340 · vllm-project/afd-plugin

ksiyuan · 2026-09-14T01:46:55Z

Relationship to #324

This PR carries no commit from #324 and is based directly on main: 12 files, all plugin Python and tests. csrc/npu/a2e/** and csrc/npu/e2a/** are untouched.

#324 registers ascend950 for those operators. That matters for running this on Ascend 950, not for reviewing or merging it. main already registers ascend910_93 and already contains the a2e ids mode this PR drives; #324 adds the A5 architecture plumbing on top. So on A3 the operators build from main as-is, and nothing here is architecture-specific. Treat #324 as a runtime prerequisite for A5 only.

Commits

feat(npu): wire the CAMP2P token-id channel for token-keyed routers.
feat(npu): run DeepSeek V4 over the synchronous CAMP2P boundary.

Commit 1 introduces the ids channel, its connector wiring, and the receive side of the ids-mode handshake. Commit 2 turns DSV4 on over that boundary and carries the model, runner, and weight-filtering changes, including the send side's half of that handshake.

The operators themselves are untouched: the ids mode already exists in the a2e kernel on main, and this PR only drives it from the connector.

Purpose

Enable DeepSeek V4 through the synchronous CAMP2pAFDConnector. DSV4 was previously restricted to CAMAsyncAFDConnector, which needs CAM/UMDK operator packages: those ship for 910C today, and A5 has no such path, where the only AFD transport is the plugin's own a2e/e2a operator pair. Nothing in this change is architecture-specific — the connector and model code carry no device gate, and A3 builds the operators from main.

The blocker is narrow. A DSV4 Hash layer routes by token identity, only the Attention role holds input_ids, and the synchronous boundary has to carry them alongside the activations.

Issue

Related: [NPU] add A5/ascend950 support for A2E and E2A ops #324 registers ascend950 for the a2e/e2a operators, which this path needs to run on A5. It is not a merge prerequisite. feat(npu): A5 (Ascend950) support for CAMP2p via HCCL p2p #303 uses a different A5 transport and is not part of this change.
No issue is closed by this PR.

Scope

In scope: the CAMP2P token-id transport, DSV4 model wiring for the gate-on-FFN configuration, role-aware weight filtering for that configuration, and feature validation.
Out of scope: wider topologies, graph capture, accuracy, and performance. compute_gate_on_attention=true remains unsupported on CAMP2P.

Hardware scope. Everything above is verified on Ascend 950 (A5). Nothing in the plugin gates this path on device type, and the ids mode this PR drives sits outside the kernel's AFD_ARCH_A5 regions, so it compiles for A3 as well: on A3 the operators build from main and this PR needs nothing else. On A5 they need
#324's ascend950 registration and a build with SOC_VERSION=950, because the kernel selects different window addressing per architecture. A3 has not been run.

Implementation Notes

Connector (camp2p.py). The a2e operator already provides an ids channel: its expert_ids input slot is int32 and it returns an int32 ids tensor beside a same-shape float32 scales tensor. The plugin had left that path idle by sending None/None with a literal compute_gate=0. Attention now packs token ids with inert scales and raises compute_gate to 1; FFN exposes the received ids on CAMP2PTransferState.

Ids-mode handshake. compute_gate is a per-side argument, and the operator only writes its ids output in the ids mode, so both sides have to select it. The receive call previously passed a literal 0 while Attention sent with 1. The FFN rank then read its simulate_expert_ids output untouched — an at::empty allocation — and installed that as forward_context.input_ids, which the Hash router indexes with. The abort surfaced inside moe_gating_top_k_hash, two calls away from the actual defect. recv_attn_output now passes the mode it was asked for.

The two roles also have to agree on the mode before anything is sent, so Attention no longer drops ids when its forward context carries none: it raises instead. The FFN rank asks for the operator's ids channel on every layer, which makes a quiet activations-only fallback a deadlock rather than a degradation. There is no opt-out switch — a Hash layer routes by token identity, so a run without ids cannot produce correct output.

FFN side. vLLM-Ascend's fused-expert selector reads forward_context.input_ids for Hash routing, so ascend_forward_context() gained an input_ids parameter. The runner receives the transfer before building the forward context, because the ids arrive with that transfer, and forwards them to compute_ffn_output.

Attention side. Token ids are selected through local_hash_input_ids, which is also used by the existing local Hash routing path, so the send-side slice and the local-routing slice cannot drift apart. A mismatch would silently route the wrong tokens rather than fail.

Weight filtering. Two role-aware fixes were needed because the upstream Ascend loader indexes its parameter dict by name without a membership check, so a path handed to a role that never registered it raises KeyError:

gate.tid2eid is registered only where the Hash MoE is built, which is FFN.
With the gate on FFN, Attention builds no router at all, so no gate.* parameter belongs to it.

Test Plan

CPU: uv run pytest, uv run ruff check .
Ascend: git pull on both roles and run a DSV4 completion round-trip. No operator rebuild is needed — the fix is confined to Python.

Test Result

CPU: ruff check and ruff format --check clean; affected unit suites pass; the full unit run's failures are unchanged from this branch's baseline.
Ascend (A5), verified: DSV4 completes a full request through the synchronous CAMP2pAFDConnector and produces sensible output. This was reached with a git pull on both roles and no operator rebuild.
Ascend (A5), observed while getting there: the rebuilt operators load and execute, DSV4 weight loading completes, and the A-to-F transfer reaches the FFN compute. Two weight-loading KeyErrors were found and fixed this way, and the routing abort was traced to the ids-mode mismatch described above.
Ascend, still open: accuracy. The output has not been compared against a native A5 run, and no quantitative evaluation (for example GSM8K) has been run for this path.

Docs Impact

None. The stale README/recipe text and the unused pytest markers are not part of this PR; it stays on the DSV4 requirement.
The A5 DSV4 boundary notes and the review of that plan are also not part of this PR.

Essential PR Checklist

Purpose is clear and linked to public context when possible.
Scope is bounded.
Compatibility with vLLM v0.26.0 is considered.
No changes are made to the vLLM source checkout.
Plugin-owned classes or explicit dotted class paths are preferred over monkey patches. The two compatibility patches touched here gained the version guard the other patches already had; no new patch was added.
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. The remaining gap (accuracy) is called out above.
Documentation impact is stated.

jiangkuaixue123 · 2026-09-14T06:42:39Z

cc @bjf-frz

bjf-frz · 2026-09-14T07:21:06Z

-core_module.EngineCoreProc.run_busy_loop = run_busy_loop
-core_module.DPEngineCoreProc.run_busy_loop = run_busy_loop
-core_module.logger.debug("AFD EngineCore patch applied")
+def _is_target_vllm_compatible() -> bool:


This function seems unnecessary and duplicates existing logic. Unless there’s a specific reason for it, I’d suggest removing it and cleaning up similar overly defensive code as well.

Removed rather than consolidated — you were right that it should not be here at all. The guard, the captured _afd_original_* originals and the _AFD_*_PATCH_APPLIED flags are gone, and afd_plugin/compat/patches/** is byte-identical to main again.

bjf-frz · 2026-09-14T07:38:06Z

+    if not getattr(model, "afd_requires_input_ids", False):
+        return False
+
+    from afd_plugin.model_executor.models.npu.deepseek_v4 import (
+        transport_input_ids_enabled,
+    )
+
+    return transport_input_ids_enabled()


This generic runner shouldn't import DSV4-specific logic—let the model adapter declare whether remote token IDs are needed and cache that capability instead of re-checking it inside the layer loop.

Fixed. _model_requires_input_ids is gone; _ffn_forward reads getattr(self.model, "afd_requires_input_ids", False) once above the layer loop, as the generic runner already did. The DSV4 import is out of the runner.

Models whose router is keyed by token identity route on input_ids, which only the Attention role holds. The a2e operator already provides the whole transport for this: its expert_ids input slot is int32, and it returns a (base_batch_size, topk) int32 ids tensor plus a same-shape float32 scales tensor. The plugin left that path idle by sending None/None with a literal compute_gate=0, and the FFN side consumed only expand_x. Wire it up: * send_attn_output accepts an optional token-aligned input_ids. When present it packs ids and inert zero scales, and raises compute_gate to 1. * the FFN side exposes the received ids on CAMP2PTransferState, requested explicitly through recv_attn_output(recv_input_ids=True), and passes that mode on to the operator instead of leaving it at the default. compute_gate selects the operator's layout on each side independently, so both roles have to agree on it before anything is sent. The receiving rank declares it through recv_input_ids. The ids output is only written in the ids mode: reading it in the other mode yields an untouched at::empty allocation, and a token-keyed router turns that into an out-of-range table read. That failure surfaces as a device abort inside the routing operator rather than at the actual defect, which is why the mode is passed through explicitly rather than inferred from the returned tensor. Add prepare_token_id_transfer and received_token_ids with CPU tests. The received_token_ids token-count check is the alignment invariant for this transport: if the ids that arrive do not describe exactly the tokens the FFN rank computes on, a token-keyed router would silently select experts for the wrong tokens, so it fails loudly instead. Tests for the ids mode itself live in tests/unit/connectors/test_camp2p_connector.py, which requires torch_npu and therefore does not run on a CPU-only host; the helper tests above cover the transport contract everywhere. Signed-off-by: ksiyuan <ksiyuan@umich.edu>

DSV4 was restricted to CAMAsyncAFDConnector, which needs CAM/UMDK operator packages that ship for 910C today. That made DSV4 unrunnable through AFD on Ascend 950 (A5), whose only AFD transport is the plugin's own a2e/e2a operator pair. The blocker is narrow: a DSV4 Hash layer routes by token identity, so whichever role owns the gate needs the tokens' ids, and only Attention holds input_ids. The FFN side consumes them through the forward context. vLLM-Ascend's fused-expert selector reads forward_context.input_ids for Hash routing, so installing the transported ids there lets the native MoE pick them up without changes to its internals. Attention: * AFDDeepseekV4RemoteMoE replaces the plain RemoteFFNProxy for gate-on-FFN layers and sends the rank-local ids alongside the activations. Connectors that do not transport ids ignore the extra argument. * local_hash_input_ids lifts the id selection out of the Hash routing path so the send side reuses one implementation. The send-side slice and the local-routing slice have to agree, and a second copy would eventually drift into routing the wrong tokens. * hash_input_ids_from_context reuses that helper for the send side and raises when the forward context carries no ids. Whether ids cross the boundary is a run-level decision the two roles share through the model's afd_requires_input_ids declaration, not a per-layer one: the FFN role waits for the operator's ids channel, so a quiet activations-only fallback would leave it reading a slot the operator never wrote. FFN: * ascend_forward_context gains input_ids and installs it. * The runner receives the transfer before building the forward context rather than inside it, because the ids arrive with that transfer; the previous order could not have installed them. It resolves recv_input_ids once per forward from the model's afd_requires_input_ids, outside the layer loop, and compute_ffn_output forwards the ids to the MoE. Role-aware weight filtering. The upstream Ascend loader indexes its parameter dict by name without a membership check, so a checkpoint path handed to a role that never registered it raises KeyError instead of being skipped. Two paths were misassigned for the gate-on-FFN configuration and are fixed here: * gate.tid2eid is registered only where the Hash MoE is built, which is FFN, so Attention must not receive it. * With the gate on FFN, Attention builds no router at all, so no gate.* parameter belongs to it. Gate ownership now follows the configured placement rather than assuming a shared router. Validation admits CAMP2P for DSV4. Per-connector checks still apply, so CAMP2P keeps requiring compute_gate_on_attention=false and quant_mode=0, while CAM async keeps requiring gate-on-Attention. A CAMP2P_CONNECTOR constant replaces the remaining duplicate connector-name literal. Verified on Ascend 950 (A5): DSV4 completes a full request through the synchronous CAMP2P boundary and produces sensible output. Accuracy has not been compared against a native A5 run, so that remains open. Signed-off-by: ksiyuan <ksiyuan@umich.edu>

bjf-frz · 2026-09-14T07:43:20Z

+AFD_DSV4_SKIP_INPUT_IDS_ENV = "AFD_DSV4_SKIP_INPUT_IDS"
+AFD_DSV4_PARAM_DUMP_ENV = "AFD_DSV4_PARAM_DUMP"


looks like test temp var, plz clean up related code

Removed, with related code.

bjf-frz · 2026-09-14T07:53:38Z

Please go through the code and remove anything unnecessary: overly defensive checks, unnecessary patches/helpers, and similar workarounds. In particular, reduce model-specific code in generic infrastructures. Keep the change minimal and focused on the actual requirement. If this is only for local testing or not ready for review, please add [WIP] to the PR title.

ksiyuan · 2026-09-14T08:13:02Z

Please go through the code and remove anything unnecessary: overly defensive checks, unnecessary patches/helpers, and similar workarounds. In particular, reduce model-specific code in generic infrastructures. Keep the change minimal and focused on the actual requirement. If this is only for local testing or not ready for review, please add [WIP] to the PR title.

Thanks for reviewing! I'm still working on this and I will add [WIP] to it until it's ready.

ksiyuan requested review from hsliuustc0106 and jiangkuaixue123 as code owners September 14, 2026 01:46

ksiyuan changed the title ~~Npu a5 a2e e2a~~ feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary (A5) Sep 14, 2026

ksiyuan force-pushed the npu-a5-a2e-e2a branch 2 times, most recently from e936946 to 0d5a198 Compare September 14, 2026 02:34

hsliuustc0106 added the high priority Needs maintainer attention — author re-pushed after review label Sep 14, 2026

ksiyuan force-pushed the npu-a5-a2e-e2a branch from 0d5a198 to 2276e26 Compare September 14, 2026 06:42

bjf-frz reviewed Sep 14, 2026

View reviewed changes

ksiyuan added 2 commits September 14, 2026 15:40

ksiyuan force-pushed the npu-a5-a2e-e2a branch from 2276e26 to 3b1fe68 Compare September 14, 2026 07:40

bjf-frz reviewed Sep 14, 2026

View reviewed changes

ksiyuan force-pushed the npu-a5-a2e-e2a branch 3 times, most recently from bda0d89 to b639f37 Compare September 14, 2026 08:11

ksiyuan changed the title ~~feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary (A5)~~ [WIP]feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary (A5) Sep 14, 2026

ksiyuan changed the title ~~[WIP]feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary (A5)~~ [WIP] feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary (A5) Sep 14, 2026

ksiyuan force-pushed the npu-a5-a2e-e2a branch from b639f37 to 9375099 Compare September 14, 2026 08:27

ksiyuan changed the title ~~[WIP] feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary (A5)~~ [WIP] feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary Sep 14, 2026

Navigation Menu

Uh oh!

[WIP] feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary - #340

[WIP] feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary#340
ksiyuan wants to merge 2 commits into
vllm-project:mainfrom
ksiyuan:npu-a5-a2e-e2a

ksiyuan commented Sep 14, 2026 •

edited

Loading

Uh oh!

jiangkuaixue123 commented Sep 14, 2026

Uh oh!

bjf-frz Sep 14, 2026

Uh oh!

ksiyuan Sep 14, 2026

Uh oh!

bjf-frz Sep 14, 2026

Uh oh!

ksiyuan Sep 14, 2026

Uh oh!

bjf-frz Sep 14, 2026

Uh oh!

ksiyuan Sep 14, 2026 •

edited

Loading

Uh oh!

bjf-frz commented Sep 14, 2026

Uh oh!

ksiyuan commented Sep 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		AFD_DSV4_SKIP_INPUT_IDS_ENV = "AFD_DSV4_SKIP_INPUT_IDS"
		AFD_DSV4_PARAM_DUMP_ENV = "AFD_DSV4_PARAM_DUMP"

Uh oh!

Conversation

ksiyuan commented Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Relationship to #324

Commits

Purpose

Issue

Scope

Implementation Notes

Test Plan

Test Result

Docs Impact

Uh oh!

jiangkuaixue123 commented Sep 14, 2026

Uh oh!

bjf-frz Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

ksiyuan Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

bjf-frz Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

ksiyuan Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

bjf-frz Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

ksiyuan Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bjf-frz commented Sep 14, 2026

Uh oh!

ksiyuan commented Sep 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ksiyuan commented Sep 14, 2026 •

edited

Loading

ksiyuan Sep 14, 2026 •

edited

Loading