Repository navigation
Fix: run SPMD paged attention on A5 at the a2a3 shapes - #1898
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (6)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe change aligns A5 MIX paged-attention pipes with GM ring buffers, updates PV-step synchronization, adds explicit FIFO count casts, and revises A5 and A2A3 test selection for supported onboard platforms. ChangesSPMD paged-attention alignment
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This change corrects A5 paged-attention FIFO handling and expands coverage to the shared a2a3 shapes; the reported cases pass across both architectures, so no actionable merge-blocking risk remains beyond normal checks and review. Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
afb286c to
db02f6e
Compare
The A5 port of the SPMD paged-attention MIX kernel produced wrong results for any non-zero query, and hung outright once a workload spanned more than FIFO_DEPTH KV blocks. Two kernel defects, both in how the kernel drives the PTO-ISA TPUSH/TPOP pipes on A5: - The three FIFOs are GM ring buffers the orchestration allocates, but the pipes were declared with DIR_C2V / DIR_V2C, which on A5 select the on-chip VEC_FIFO path and ignore the GM base entirely. That path also strides local slots by a full SLOT_SIZE while the kernel budgets its UB at half that, so the second slot ran past SIJ_UB_BASE into OI_UB_BASE. Use DIR_C2V_GM / DIR_V2C_GM. - The Cube popped pij with TILE_NO_SPLIT. The split axis also selects the two-lane sync: setIntraBlockBySplit / waitIntraBlockBySplit only touch the second AIV lane's flag space when the axis is not NO_SPLIT. So free() released only lane 0's slot, stalling lane 1's third push, and wait() observed only lane 0's data-ready, letting the Cube read a pij tile lane 1 had not finished writing. The Mat-from-GM pop ignores the axis for addressing, so TILE_UP_DOWN changes synchronization only. With both fixed, A5 runs the a2a3 shapes unchanged at the stock FIFO_DEPTH, so the A5 cases become Case1 / Case2 / SmallCase1 with the same parameters and names as a2a3, and Case1 joins the Per-PR sweep as a2a3's does. The previous A5 case zeroed the query, which makes sij zero regardless of how K is loaded and therefore exercised only the PV path and the final normalization; it passed throughout while every non-zero query was wrong, so it is replaced rather than kept. The host_build_graph cases now mirror the tensormap_and_ringbuffer shapes on both architectures instead of carrying a single small zero-query case. Simulator coverage is dropped for this workload: the MIX kernel does not build for the CPU simulator, because PTO-ISA rejects the ColMajor TMULS in the softmax tail. The a2a3sim / a5sim platforms are removed from the case lists, and the __CPU_SIM compatibility shim, which only existed to get that build further, goes with them. The FIFO tensor shapes in both orchestrations get an explicit cast: the division by sizeof yields size_t, and narrowing it inside a braced initializer is an error for Clang. Verified onboard on a2a3 (Ascend910/9392) and a5 (Ascend950PR/9579): all 12 cases pass, and Case1 / Case2 / SmallCase1 each pass 5/5 on A5. Closes hw-native-sys#1816 Refs hw-native-sys#1832
2 participants
Summary
The A5 port of the SPMD paged-attention MIX kernel produced wrong results for any
non-zero query, and hung outright once a workload spanned more than
FIFO_DEPTHKV blocks. Both are defects in how the kernel drives the PTO-ISA TPUSH/TPOP pipes
on A5. With them fixed, A5 runs the a2a3 shapes at the unchanged
FIFO_DEPTH = 2,so the A5 cases become
Case1/Case2/SmallCase1with the same parametersand names as a2a3.
The two kernel defects
1. Pipe direction — the GM ring was never used. The three FIFOs are GM ring
buffers the orchestration allocates, but the pipes were declared with
DIR_C2V/DIR_V2C. On A5 those select the on-chip VEC_FIFO path, which ignores the GM baseentirely —
popTileFromVecFiFoonly ever readsC2V_CONSUMER_BUF. That path alsostrides local slots by a full
SLOT_SIZE, while the kernel budgets its UB at2 * SUB_QT * MAX_BLOCK_SIZE * sizeof(float). Forq_tile = 16both are 8192 B,so slot 1 starts exactly at
OI_UB_BASEand overlaps the entire oi buffer.A5 has
DIR_C2V_GM/DIR_V2C_GMfor GM-backed rings; they are behind#ifdef PTO_NPU_ARCH_A5and a2a3 has no such distinction, which is how the portinherited the wrong spelling.
examples/a5/tensormap_and_ringbuffer/bgemmis thecorrect use of the on-chip form — it passes
nullptras the GM base.2. Split axis on the Cube's pij pop. The Cube consumes the whole
Q_TILEtile, so
TILE_NO_SPLITlooked right. But the split axis also selects thetwo-lane sync:
setIntraBlockBySplit/waitIntraBlockBySplitonly touch thesecond AIV lane's flag space (
flagId + VEC_CORE_ID_OFFSET) when the axis is notNO_SPLIT. Sofree()released only lane 0's slot, andwait()observed onlylane 0's data-ready — letting the Cube read a pij tile lane 1 had not finished
writing.
popMatTileFromGMFiFoignores the axis for addressing, soTILE_UP_DOWNchanges synchronization only, not the data path.
That single omission produced both symptoms:
shouldWaitFreereturns false whiletileIndex < SlotNum, which is why theboundary sat exactly at
FIFO_DEPTH. The predicate is cumulative per hardwareblock, since the pipe's
tileIndexis not reset across the stride loop:(logical blocks on that hw block) * n_blocks > FIFO_DEPTH. At the stock depthof 2 the a2a3 shapes need 704 / 384 / 64 pushes.
Coverage changes
the query, which makes
sijzero regardless of how K is loaded, so it exercisedonly the PV path and the final normalization — it passed throughout while every
non-zero query was wrong. It is replaced rather than kept.
Case1isnon-manual, mirroring a2a3;
Case2andSmallCase1stay manual.host_build_graphmirrorstensormap_and_ringbufferon both architectures,instead of carrying a single small zero-query case. All HBG cases stay manual.
build for the CPU simulator: PTO-ISA rejects the ColMajor
TMULSin the softmaxtail.
a2a3sim/a5simare removed from the case lists, and the__CPU_SIMcompatibility shim, which only existed to get that build further, goes with it.
by
sizeofyieldssize_t, and narrowing it inside a braced initializer is anerror for Clang.
Verification
Onboard on a2a3 (
Ascend910 / 9392, CANN 9.0.0) and a5 (Ascend950PR / 9579,CANN 9.2.0), every run wrapped in
task-submit. All 12 declared cases pass:Both selections were exercised on each architecture:
--manual include(the fullsweep above) and the default Per-PR selection, which collects exactly one node per
architecture — the non-manual
Case1— and passes.Because the second defect was a race, the three shapes were re-run individually on
a5: 5/5 pass each, 15/15 total.
The diagnosis rests on these measurements, all on a5:
FIFO_DEPTH. At depth 2,n_blocks1 and 2 passed while3 and 4 returned 507014 and 8 returned 507018. At depth 4,
n_blocks3 and 4passed and 8 returned 507014. At depth 8,
batch = 64— 6 cumulative pushes —stopped hanging. It is not "merely slow": it still stalled under
PTO2_SCHEDULER_TIMEOUT_MS=60000, reported assched_error_code=100,sub_class=S1:running-stalled,completed=0/1.n_blocks = 2: batch 2 and 4passed 5/5; batch 8 passed 2/5; batch 10 passed 1/3; batch 9 failed 3/3; batches
12, 14, 16, 18, 20, 24 and 64 failed every run,
max_diff0.12–0.22. Thenon-monotonic, reproducibly flaky pattern is what identified it as a race rather
than a size limit.
Ruled out along the way, each by measurement rather than inspection: porting the
a2a3
setEntryOffsetlane-split fix-up (made it worse, 0.44 → 6.18); passingsub_block_idthrough theTPUSH/TPOPoverloads (bit-identical output, so A5already returns the correct lane id); flag-ID overlap between the three pipes (the
_GMdirections use onlyFlagIDandFlagID + 1); a missingTFREE(the_GMpop paths return
reqFree = trueandTPOP_IMPLalready frees);SPMD_BLOCK_NUMover-subscribing A5; and per-hardware-block GM FIFO regions overlapping.
Closes #1816
Refs #1832