[310P] Avoid nonzero / per-element host syncs in GDN under TP=4 (v2, rebased on c460a15af) - #23
Open
adeepn wants to merge 1 commit into
Open
[310P] Avoid nonzero / per-element host syncs in GDN under TP=4 (v2, rebased on c460a15af)#23adeepn wants to merge 1 commit into
adeepn wants to merge 1 commit into
Conversation
On Atlas 300I Duo (310P3) under TP=4 with cross-card P2P, two GDN code paths fault or stall the host: 1. The boolean-mask index assignment initial_state[~has_initial_state, ...] = 0 lowers to aclnnNonzeroV2 plus a device->host sync; the op faults on some chips and stalls on others, the worker busy-spins in the NPU driver, and the host hangs (only an iBMC reset recovers it). Replace it with a nonzero-free torch.where helper (_zero_states_without_initial), which also writes an exact 0 for masked rows holding NaN/Inf (a state*mask multiply would leak NaN). 2. _iter_seq_ranges unpacked cu_seqlens with 2*(N-1) per-element .item() calls, each a device->host sync sharing the PHB path with cross-card P2P DMA. Collapse to a single bulk .tolist(). Add CPU unit tests for both helpers. Signed-off-by: Viacheslav Bocharov <v@baodeep.com>
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
1 participant
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
Re-derive the 310P TP=4 GDN host-stability fix onto current upstream (base
upstream-main@c460a15af).On Atlas 300I Duo (310P3) under TP=4 with cross-card P2P, two GDN paths fault or stall the host:
initial_state[~has_initial_state, ...] = 0(gdn_310.py) lowers toaclnnNonzeroV2plus adevice→host sync; it faults on some chips and stalls on others → the worker busy-spins inside the
NPU driver → RCU stall → host hang (only an iBMC reset recovers). Replaced with a nonzero-free
torch.wherehelper (_zero_states_without_initial) that also writes an exact 0 for masked rowsholding NaN/Inf (a
state * maskmultiply would leak NaN)._iter_seq_ranges(chunk_gated_delta_rule.py) unpackedcu_seqlenswith2*(N-1)per-element.item()syncs sharing the PHB path with cross-card P2P DMA. Collapsed to a single.tolist().Why a new PR
Supersedes #21. Upstream rewrote
gdn_310.py(+250lines); #21 was authored against an oldersnapshot and no longer applies. This re-derives the same fix onto
c460a15af(the commit baked intothe current 310P nightly image), preserving the original intent and tests.
Test
tests/ut/_310p/ops/test_gdn_310.py(torch.where vs theboolean-index assignment incl. NaN/Inf and a shape-mismatch guard),
tests/ut/_310p/ops/test_chunk_gated_delta_rule_310.py(_iter_seq_ranges). 23 ut pass.nightly-main-310p(c460a15af): loads,serves, tool-calling, prefix-caching (29.9% hit), decode 4.4 t/s, host stable across long runs
(no hang).