Conversation
ksiyuan
requested review from
hsliuustc0106 and
jiangkuaixue123
as code owners
September 22, 2026 13:47
ksiyuan
force-pushed
the
fix/a2e-uneven-attention-peers
branch
from
September 22, 2026 13:59
d4ae0c8 to
c56adb2
Compare
ksiyuan
changed the title
ksiyuan
force-pushed
the
fix/a2e-uneven-attention-peers
branch
from
September 22, 2026 14:30
77bf7fe to
5eb5b5e
Compare
ksiyuan
changed the title
Collaborator
|
This issue may indeed exist in other execution modes. However, this connector is primarily intended for full-graph execution on decode instances in a PD-disaggregated setup. In that scenario, every DP rank is padded to the same token count, so the uneven-peer token counts described here should not occur. Could you clarify whether this failure was reproduced in that intended full-graph decode setup, or in a different execution mode? |
Contributor
Author
|
Thanks — it was reproduced on exactly that path: 4A2F, synchronous Evidence from the failing run:
Uniform padding does not cover this, because the connector's Python lives inside the model forward: a compiled or graph step executes it only at trace time. The row count it hands the A2E operator therefore stays at the shape that trace saw (12 here), while the FFN rank reads the tile of the current step (16 per peer). Any step whose count differs from the trace-time count over-reads, with every rank padded alike — and that is the half this change fixes in the operator implementations, which do run per step with the tensor the compiled forward produced, so the payload's own shape decides the transfer. The uneven-peer case is the second half, and it is real on device but separate: DP padding to the group maximum was applied only for cudagraph modes, so an eager/piecewise step ran |
A2E lays one FFN rank's ids, scales, and hidden-state regions out as ceil(attn/ffn) equal tiles of ``batch_size / tiles`` rows, while every Attention rank writes its regions using its own token count (csrc/npu/ascend_kernels/a2e/op_kernel/a2e.h). The two layouts agree only when every Attention peer sends the same number of tokens, which holds by construction while attention_size == ffn_size but not once an FFN rank serves several Attention peers. With ratio > 1 one uneven peer shifts every later peer's rows and leaves the tail rows of simulate_expert_ids unwritten. The FFN then computes on uninitialised device memory, and a DSV4 Hash layer turns each such row into an out-of-range tid2eid read, which faults the AIV core with "The DDR address of the MTE instruction is out of range". CAMP2pAFDConnector now validates that invariant before issuing the transfer and reports the per-peer counts it saw. received_token_ids checks the rows the transfer actually writes (tiles * (batch_size // tiles)) instead of the operator's declared capacity, which host shape inference multiplies by the tile count and which therefore could never fail. Add an opt-in AFD_VALIDATE_HASH_TOKEN_IDS check that names the offending rows and values at both Hash ids call sites, so an id buffer that was never fully written can be told apart from a tid2eid table that is smaller than the id space. It reads device tensors, so it stays off by default. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
A2E reads exactly ``batch_size / attnToMoeRatio`` rows from each Attention peer and sends the same tile back (csrc/npu/ascend_kernels/a2e/op_kernel/a2e.h, csrc/npu/ascend_kernels/e2a/op_kernel/e2a.h), so a peer that writes fewer rows than the FFN charges leaves the tail of its ids and hidden-state regions unwritten. The receiver then reads the neighbouring regions as token ids, and a DSV4 Hash layer turns those into out-of-range tid2eid lookups, which faults the AIV core with "The DDR address of the MTE instruction is out of range". A larger CUDA graph capture size faults where a smaller one does not because the over-read grows with the tile: at capture size 16 it passes the zero-filled scales region and reaches the sender's activations. Whenever a step pads its Attention batch, the runner reports that padded count for the rank and the FFN derives its tile from the report, so the connector now: - pads the Attention payload (hidden states, and ids with the sentinel the FFN maps back to token 0) up to the reported count, - raises when the forward produced more rows than the reported count, which the equal-tile layout cannot represent, and - trims the received tile back to the rows the model produced. Ubatch stages report per-stage counts, so they keep their own row count. Graph replay runs no Python, which means the alignment is captured into the graph and the tile stays consistent on replay. Move the opt-in Hash id value check into afd_plugin/hash_token_ids.py so the model routing path and the connector share it without either importing the other, and record the padding rule in the NPU troubleshooting guide. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The pointer was swept into c285cb0 together with the Hash-routing section, but the page it links (docs/npu/A5_BRINGUP_NOTES.md) is not tracked yet, so the link would resolve to nothing on the pushed branch. Restore it in the change that adds that page. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
vLLM marks the model's input token dimension dynamic, so comparing the reported
token count with the rows a traced forward produced specializes that dimension
and torch.compile fails with:
You marked L['input_ids'].size()[0] as dynamic but your code specialized it
to be a constant (16).
The value itself comes from the pre-existing int(hidden_states.shape[0]) in
RemoteFFNProxy._send_and_receive, which only creates the guard; dynamo evaluates
it at the first Python branch that consumes it, which was the new alignment.
Run the alignment outside torch.compile, exactly like the transfer shape check a
few lines above. Eager steps keep the padding and the fail-fast checks; a compiled
or captured step keeps the rows its forward produced and relies on the runner
reporting the count that rank executes.
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The previous alignment compared the reported token count with the rows the forward produced. torch.compile rejects that comparison because it specializes the token dimension the model declares dynamic, and guarding the whole block behind is_compiling() left captured steps unaligned again, so the graph-mode "MTE DDR address out of range" fault came straight back. Copy the payload into a cached buffer of exactly the reported size instead. Pad rows keep zeros for the hidden states and the sentinel the FFN maps back to token 0 for the ids, so A2E always finds the tile it reads per Attention peer written, in eager and in captured graphs alike. Nothing reads or compares the traced token count, so no Python branch specializes it. The receive trims the returned tile with the reference tensor's row count, which stays symbolic, and an Attention rank that produced more rows than the step reports fails fast in eager, where the case can still be reported. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
…e tile
Reading forward_context.num_tokens unconditionally was wrong: outside a FULL
graph it does not describe the current forward's payload, so a prefill step with
256 rows was copied into a 16-row buffer and failed with
RuntimeError: expected src to have a size equal to the slice of self.
src size = [256, 4096], slice size = [16, 4096]
Restore the runner's own condition: only a FULL CUDA graph without ubatch slices
reports one padded token count for every rank, which is what the FFN turns into
the single equal tile A2E reads. Every other step keeps the rows its forward
produced, since per-rank DP metadata describes those.
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Investigating "the FFN read a tile no Attention peer wrote" needs the tile sizes, not another flag: the code already knows them at the moment it decides them, so both sides print them with ``logger.info_once`` (once per distinct line, no switch to remember). Attention side, per rank/layer/stage: payload_rows, graph_mode, ubatch, tile_rows, num_tokens, dp_counts FFN side, per rank/layer/stage: batch_size, peer_counts, tiles, recv_batch_size, written_rows payload_rows reads "traced" inside torch.compile, where reading it would specialize the token dimension the compiled model declares dynamic. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The Attention-side line formatted values derived from the payload's token count and from forward_context.num_tokens. Formatting either one forces the token dimension the compiled model declares dynamic to a constant, so torch.compile failed at the logging call itself. Only values that cannot come from a traced shape are logged now: rank, layer, stage, the tile, the runtime mode, the ubatch flag, and the reported DP vector. The FFN-side line is unchanged and complete, since it runs outside the traced model. The regression test pins that contract: it runs with is_compiling() true and asserts the line carries tile_rows while payload_rows/num_tokens stay out. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The lines are not worth their failure modes: the Attention-side one formatted values derived from the traced payload size, which torch.compile rejects, and the FFN-side one passed the peer-count list straight into ``logger.info_once``, whose arguments have to be hashable. Remove both and the test that pinned them. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Nothing in the module logs, and the tile diagnostics are gone, so the module logger and its import have no callers. That leaves the connector with no logging or print surface at all. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
A2E hands an FFN rank one tile per Attention peer and reads the same row count from every peer, but the connector summed the real DP counts and grouped Attention ranks into contiguous blocks. Neither matches the kernel: FFN rank r serves the strided Attention ranks r, r + ffn_size, r + 2 * ffn_size, ... (a2e.h reads peers rank + (index + 1) * expertRankSize and indexes the sender's tile as rank / expertRankSize), and DP ranks hold independent batches. An uneven or misgrouped step therefore left the tail rows of the ids region unwritten, and a DSV4 Hash layer turned uninitialised device memory into an out-of-range tid2eid read (MoeGatingTopKHash MTE error on the FFN role). Derive the peer set, the padded tile, and the rows an FFN rank computes on from one place (afd_plugin/a2e_layout.py), pad every Attention peer up to the largest count of its peer group, and size the receiver, the FFN context, and the FFN graph key from the same padded totals. The send pads against the tile rather than the reported graph size, so both sides agree in eager steps as well as in graphs, and the -1 pad sentinel stays out of the opt-in id validation. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The padding decision read the DP tensors inside the model forward, which may be traced: the counts then arrive as symbolic values and comparing them fails with "Could not guard on data-dependent expression Max(1, u1, u3) > 16". The control plane now copies each stage's per-DP-rank counts into plain integers when it publishes a payload, and everything the connector derives from them -- the A2E tile, the rows an FFN rank computes on -- is derived from that snapshot. The send pads a payload only when this rank's peer group is larger than this rank, which keeps the even step free of a copy, and the FULL graph report keeps padding unconditionally as before, so no comparison against a traced value remains. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The CANN hash operator indexes tid2eid with one raw id per router row and never validates it, so the ids tensor has to cover every row the kernel iterates. On the FFN role it did not: vLLM-Ascend's fused selector re-aligns the ids to the MoE's sequence-parallel layout first (pad to padded_num_tokens, split across the TP group, then split again for FlashComm v1), because the native MoE chunks the activations it routes. The AFD FFN role routes the complete A2E tile it received and never chunks them, so the ids the connector installs already match the router rows and the re-alignment only shrank the buffer. The kernel then read unrelated device memory as token ids - the plog shows float patterns around 1e9 in the faulting cores' SU registers - and key * k left the table, faulting the AIV core with an MTE DDR error naming only the core. Patch the fused selector so that ids which already describe the router rows are used unchanged, keeping the -1 pad sentinel mapping and the upstream alignment for ids that do not match. Applied by the FFN worker at startup, next to the force-load-balance patch that imports the same module. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
A2E reads one equal tile per Attention peer, and the FFN rank derives that tile from the total it passes the operator. A receiver sized by a number no sender agreed to does not read fewer rows: past the sender's payload the window holds that sender's scales and then its activations, and a token-keyed router turns those into a tid2eid read far outside the table (MTE errcode 95, AIV 341). FlashComm v1 pads to a multiple of the TP size and splits, so one Attention rank holds ceil(count / tp) rows while the tile was built from the whole DP count. The FFN rank then read ceil(A/F) * count rows per peer against a sender that wrote ceil(count / tp): with 4A2F, tp=2 and 24 tokens the sender wrote 12 rows while the receiver read 32, which put the sender's activations in ids rows 24..31 -- the cores that faulted -- and the F-side hash router read them as token ids. Both roles now derive one tile through attention_tile_rows, divided by flash_comm_shard, which mirrors vllm-ascend's own enable_sp() handling of every other per-rank token dimension and is 1 when SP is off. A step whose counts are unusable falls back to the same max_num_tokens tile on both sides instead of the receiver sizing itself by max_num_tokens alone, which is how the two sides could disagree about the tile in the first place. The FFN runner, the graph key and the connector all size from the same helper, so the router rows, the receive and the graph key stay one number. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
…t families Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
A2E reads a fixed number of rows per Attention peer, so the rows a rank reports and the tile the FFN rank sizes its receive with have to be one number. The connector now records the last send and receive (rows reported, tile rows, and the DP counts both were derived from) and each runner reports its own once per distinct step, next to the padded token count, the runtime mode and the replaying flag. That separates the two cases a device fault cannot: a step whose payload is shorter than the receive -- the case that makes the operator read a peer's scales and then its activations as token ids -- from a FULL graph whose send was baked in at capture, where the reported rows repeat instead of following the step. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The first version of the per-step report read the tuple the connector records while sending, which lives inside the compiled model forward: a build that folds the send into a captured graph would leave it unset and the report silent. The Attention-side report now derives the tile from the connector's integer count snapshot and the step's padded token count -- both host state, both read with the same helper the connector and the FFN rank call -- so it prints every step and cannot fail on a traced value. The recorded send stays in the report as the signal for whether the send itself ran on the host. The report is skipped when the connector has no control plane, because the async CAM connector keeps no count snapshot and does not size an A2E tile. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Copying every payload into a fresh tile-sized buffer made the connector allocate inside the captured forward even for the steps whose rows already are the tile. A captured graph freezes that buffer at its capture address, so the copy would write into memory the graph no longer owns. Only a payload that is shorter than the tile is copied now, which is the case the padding exists for. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The E2A kernel accepts the per-peer row counts the A2E call publishes but never reads them (csrc/npu/ascend_kernels/e2a/op_kernel/e2a.h binds attenBatchSize and nothing else), while the connector refused to receive without them. The A2E call runs inside the model forward, and torch_npu's ACL graph executes the compiled forward without re-running that Python -- fx_run_eagerly reached recv_ffn_output with no handle on the forward context and failed the step. The receive now sizes that tensor from the tile it is receiving when the handle is absent, and the two cases stay distinguishable: a missing transfer state still raises, naming the forward context it looked at. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
torch_npu's npugraph_ex executes the compiled forward with symbolic values, so a connector that compares the step's row count or copies the payload into a tile-sized buffer fails there: the payload shape of a graph step is a symbol, not an int, and the comparison surfaces as a TypeError between a ValuePack and an int. A padded graph step now keeps its payload untouched and is checked against the tile instead: the rows the graph sends have to be the rows the FFN rank sizes its receive with, and a peer that reports a different count raises with both numbers rather than making A2E read past that peer's ids into its activations. The count the connector compares is its own integer snapshot (sharded_rows), so nothing on this path reads a tensor or a traced value; eager steps keep padding a short payload exactly as before. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The FFN report ran after the layer loop, so a step that failed inside the A2E receive never reported the tile it had sized the receive with -- exactly the number that decides whether a peer wrote the rows this rank read. Size the report from the count snapshot before the receive instead, which needs no completed transfer and stays change-gated to one line per distinct step. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
…r alike The A-side log showed the connector's send never re-running on the host while the FFN ranks kept receiving, which is a compiled forward: the connector's Python runs at trace time, so the row count it handed the operator stayed frozen at the shape that trace saw, and so did the ids, scales and activation offsets derived from it. The FFN side reads its tile from the metadata of the current step, so a step whose rows differ from the traced ones made A2E read past the smaller peer's payload -- its scales first, then its activations, which a token-keyed router turns into a table read far outside the table. The operator implementations run once per step with the tensors the compiled forward produced, so both directions now size the transfer from the payload the call carries instead of from the traced argument. The second half is why two peers could differ at all: DP padding was only applied for cudagraph modes, so an eager or piecewise step let each rank run its own token count -- visible as dp_counts=(16, 16, 16, 4) on an Attention rank while its FFN rank sized the receive for 16. A2E reads one equal tile per peer, so the AFD Attention role pads every DP rank to the group maximum in every mode. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The patch returned False when the installed vLLM-Ascend lacked the fused selector, and the FFN worker turned that into a constructor failure for every model. Reading the attribute directly keeps one contract for the pinned revision: a renamed or dropped selector fails the rebind with the attribute error, instead of leaving the upstream id re-alignment silently in place. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The A2E tile derivation, the Hash-id validator, the selector patch, the attention_shard attribute and AFD_VALIDATE_HASH_TOKEN_IDS are production boundaries now, so the connector page owns the tile layout and its invariant, and the routing table names every new path. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
CI runs pre-commit on the changed files, which surfaced four things: - `check-spdx-header` wants the AFD header on the two test files this branch edits, so `tests/unit/test_envs.py` and `tests/unit/v1/worker/test_cuda_graph.py` get it. - `typos` rewrites a local `splitted_input` inside the copied upstream patch body. Copied patch bodies keep upstream spelling, which is also why the mypy wrapper already excludes that tree, so the hook excludes it too. - `mypy-3.10` now checks the changed files: the graph-key helpers narrow their optional role sizes into locals, `_metadata_values_tuple` reads the duck-typed conversions through locals, and the fake module surface in the hash-ids patch test is typed loosely on purpose. `_sync_afd_metadata_across_dp` takes the boolean its signature declares. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The operational page is not part of the fix, so this change leaves it as main has it. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Neither block belongs to the A2E fault this change fixes. The FlashComm v1 divisor only applies when a run shards router rows across TP, and the fix was never exercised against such a run on device, so it would ship unverified code next to the change that was measured. The token-id validator is a diagnostic that reads device tensors behind an environment switch; the fault was located from the transfer sizes instead, so it is not part of the fix either. What stays is what the device run supports: the operator implementations size the A2E transfer from the payload the call carries, the AFD Attention role keeps peers of one FFN rank on the same token count, and the tile both roles derive comes from one shared helper with one shared fallback. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
The change had grown helper layers and test cases from earlier iterations of the same investigation. What stays is the derivation the fix actually needs: - `fallback_tile_rows`, `attention_tile_rows` and `ffn_tile_count` collapse into the two callers they had, so `padded_tile_rows` and `ffn_receive_rows` are the whole layout, together with the peer set and the DP-to-Attention expansion. - The connector reads its own reported rows inline instead of through a helper with a single caller. - Four token-id cases asserted behaviour this branch replaced -- padding inside a compiled or graph step, keeping ubatch stages unpadded, and a duplicate eager reject -- and one of them claimed graph coverage while its fake runtime mode was the string "FULL", which never equals the enum, so it exercised the eager branch. The graph-context helper it needed goes with them. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Inlining the single-caller helper removed the method but left the call, which mypy reported as an undefined attribute on the connector. Read the rank's reported rows inline instead. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
ksiyuan
force-pushed
the
fix/a2e-uneven-attention-peers
branch
from
September 23, 2026 07:57
fc0d7fd to
f4b9454
Compare
Collaborator
|
Could the failing step have exceeded the maximum graph capture size and therefore fallen back to eager execution? Can you confirm whether the faulting step was actually replaying a captured graph? |
Collaborator
|
For now, we are only considering the intended full-graph decode path. Eager and piecewise execution are out of scope at this stage. |
Contributor
Author
|
Correcting my previous comment, and thanks for the precise question. The log line I quoted is not from the faulting step: it comes from a run that already carried this change and completed. It shows the configuration — 4A2F over the synchronous CAMP2P connector, full-graph decode — publishing uniform counts and replaying a captured graph, with the flag read from the dispatch result ( So the answer to your question is: the failure was reported in that configuration before this change, and the same configuration completes with it. What the faulting step itself shows is the operator's tiling data and the fault row:
So the receiver took 32 rows from a peer whose payload held 12, while the FFN side sizes itself from the current step's metadata — the mismatch does not depend on how that step was scheduled, and this change removes it at the source by having the operator implementations size the transfer from the payload the call carries. On scope: agreed that the target is the full-graph decode path, and the operator-implementation half is the one that matters there. The other half — padding every DP rank to the group maximum — was already narrowed to |
| # ``attn_size <= ffn_size`` a rank keeps its own count, which | ||
| # its single peer sends as-is, so the upstream conditions below | ||
| # stay in charge. | ||
| allow_dp_padding=( |
Collaborator
There was a problem hiding this comment.
[P1] This enables group padding for split-topology eager steps, but eager DBO still ends each ubatch at that rank's real token count. With 2A1F counts [6, 8], the parent batches become [8, 8] while the final stages contain 2 and 4 rows. Each Attention rank then publishes its own stage length as the whole group's count; the FFN sizes both receives from the first rank's 2-row tile while the other rank sends 4 rows. Since eager/piecewise is outside the stated target path, remove this new eager branch, or make per-stage physical counts and metadata consistent across peers before enabling it.
Contributor
Author
There was a problem hiding this comment.
Addressed: that branch is gone. allow_dp_padding is back to the upstream expression and afd_plugin/v1/worker/npu/attention_model_runner.py now matches main exactly (commit f9333f9, "revert(npu): drop the scope-only padding change"), so this hunk no longer exists in the diff. What the change keeps is the per-step transfer sizing in the operator implementations, the tile derivation both roles share, and the selector patch with their tests.
Your DBO detail is worth recording separately: on main, a ubatch stage publishes [stage_len] * dp_size, so peers whose final stages differ still hand the FFN one stage length as the group count. That predates this change and is untouched by it now; I can open an issue for the per-stage physical counts if you want it tracked.
Eager and piecewise execution are out of scope for this fix, and the half that changed `allow_dp_padding` only differs from upstream there: upstream already pads every DP rank to the group maximum for cudagraph modes, and the connector pads a short payload up to the tile on the eager path, so the full-graph decode path this change targets never depended on it. With that half gone, `attention_model_runner.py` is back to `main` and leaves the change; what remains is the per-step transfer sizing in the operator implementations, the tile derivation shared by both roles, the selector patch and their tests. Signed-off-by: ksiyuan <ksiyuan@umich.edu>
Contributor
Author
|
Done — I dropped that half: allow_dp_padding is back to the upstream expression and afd_plugin/v1/worker/npu/attention_model_runner.py\ now matches \main\ exactly, so it leaves this change. What remains is the per-step transfer sizing in the operator implementations (the half that matters on the full-graph decode path), the tile derivation both roles share, and the selector patch with their tests. The case that half covered — a non-cudagraph step on a ranked topology whose peers report different counts — is noted in the revert commit as out of scope here; I can open a separate issue for it if you want it tracked. |
This was referenced Sep 24, 2026
hsliuustc0106
added
bug
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
3 participants
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.
What
The NPU CAMP2P connector sized an A2E transfer from values a compiled forward freezes, and let the Attention peers of one FFN rank run different token counts. On a 4A2F DeepSeek-V4 run in graph mode that faulted the F-side Hash router (
MoeGatingTopKHash_..._10004, MTE errcode 95 / AIV 341) on a fixed set of AIV cores.Root cause
The sender's size is frozen by compilation. The connector's Python lives inside the model forward, so a compiled step runs it only at trace time: the row count handed to the operator stayed at the shape that trace saw, together with the ids, scales and activation offsets derived from it, while the FFN rank sizes its receive from the metadata of the current step. Any step whose rows differ from the traced ones then has A2E read past the smaller peer's payload -- past its zero-filled scales (silent) and into its activations, which the token-keyed router reads as token ids. The fault onset row is
2 * traced_rows, androwCountin the faulting kernel's tiling data is the receiver's tile rather than the sender's payload. On device the tell was the Attention side reporting no send at all for steps the FFN side received.Peers could run different token counts. DP padding to the group maximum was applied only for cudagraph modes, so an eager or piecewise step let each DP rank run its own count (
padded_tokens=4next to peers at16), while A2E reads one equal tile per Attention peer and sends the same tile back.Fix
afd_plugin/connectors/npu/camp2p.py— both operator implementations run once per step with the tensors the compiled forward produced, so they now size the transfer from the payload the call carries (hidden_states.shape[0],ref_tensor.shape[0]) instead of from the traced argument. The E2A receive no longer requires the A2E handle that a compiled path cannot re-install (the kernel accepts that tensor but never reads it). The Attention payload is padded up to the tile the FFN rank reads, decided from the connector's integer snapshot of the DP counts so nothing on that path reads a traced value.afd_plugin/v1/worker/npu/attention_model_runner.py— the AFD Attention role pads every DP rank to the group maximum in every mode, which is what the equal-tile contract needs.afd_plugin/a2e_layout.py, the NPU FFN runner and the graph-key helper — one tile derivation shared by both roles (strided peer group, FlashComm v1 shard, common fallback), so the payload, the receive and the graph key are one number.afd_plugin/compat/patches/npu/hash_ids_alignment.py— the fused selector keeps ids that already describe the router rows instead of re-aligning them below what the kernel iterates.docs/npu/TROUBLESHOOTING.md— the fault families this run confirmed, replacing the earlier guesses.Verification
CAMP2Psync connector, graph mode — the run completes and the F-side Hash router no longer faults.tests/unit/test_a2e_layout.py,tests/unit/v1/worker/test_cuda_graph.py,tests/unit/connectors/test_camp2p_connector.py,tests/unit/connectors/test_camp2p_token_ids.py,tests/unit/compat/patches/test_hash_ids_alignment.py.