fix(dsv4/sparse_attn): per-row inverse-RoPE to avoid a5 gather lowering bug#801
fix(dsv4/sparse_attn): per-row inverse-RoPE to avoid a5 gather lowering bug#801lwDavid wants to merge 3 commits into
Conversation
…ng bug On a5 (Ascend 950) the tensor.gather flat-index lowering miscompiles when the source tile spans more than one FP32 vector box (>8 rows). The inverse- RoPE in merge_norm gathered a [H_TILE=16, ROPE_DIM] tile of a strided slice of n_full, spanning two vector boxes, which corrupted the ROPE columns of every head except the first of each 8-head O_GROUP (only tile rows 0 and 8 received correct data). qkv_proj_rope avoids this because its rotation tile is exactly [8, ROPE_DIM] (one vector box). Fix: unroll the inverse-RoPE per head-row using 1-row [1, ROPE_DIM] tiles (mr_swap_idx / mr_cos_il / mr_sin_signed at [1, ROPE_DIM], loop n_hi over H_TILE, pl.mul on matched-shape 1-row tiles). Every gather now fits within a single FP32 vector box regardless of H_TILE. The rotation math is identical to the original (computed per row), so a2a3 is unaffected. Applied to all three sparse_attn variants (decode_sparse_attn, _hca, _swa). Verified on a5: pre-proj o_packed ROPE rel-L2 drops 0.417 -> 0.153 (matching NOPE) and the per-O_GROUP-head error pattern is eliminated. A separate, uniform ~15% base-attention precision divergence on a5 remains (qk_pv/merge_norm path) and is tracked independently; this commit addresses only the inverse-RoPE gather lowering bug.
|
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughInverse RoPE rotation and packing were changed from bulk ChangesInverse RoPE packing
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the inverse RoPE computation across several decode sparse attention files to process one head-row at a time (1-row tiles) instead of using a fused multi-row gather, working around an A5 compiler bug where large gathers miscompile. The feedback suggests optimizing decode_sparse_attn_swa.py by reusing the precomputed rope_swap_idx rather than recomputing it from scratch.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
models/deepseek/v4/decode_sparse_attn_swa.py (1)
321-338: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRemove now-dead
rope_swap_idxprecompute.Introducing the per-row
mr_swap_idxhere makes the earlierrope_swap_idxtensor (allocated and filled at lines 261-274, in the samerope_csblock) dead — it's no longer read anywhere in the function. That block still computesswap_ones/swap_col/... and writes intorope_swap_idxfor nothing, wasting vector-engine work in therope_cspl.atscope on every call.♻️ Proposed cleanup (lines 261-274)
- rope_swap_idx = pl.create_tensor([H_TILE, ROPE_DIM], dtype=pl.INT32) with pl.at(level=pl.Level.CORE_GROUP, name_hint="rope_cs") as rope_tid: - swap_ones = pl.full([H_TILE, ROPE_DIM], dtype=pl.FP32, value=1.0) - swap_range_i32 = pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32) - swap_range = pl.cast(swap_range_i32, target_type=pl.FP32) - swap_col = pl.col_expand_mul(swap_ones, swap_range) - swap_half = pl.mul(swap_col, 0.5) - swap_dup_i32 = pl.cast(swap_half, target_type=pl.INT32, mode="trunc") - swap_dup_f = pl.cast(swap_dup_i32, target_type=pl.FP32) - swap_lane = pl.sub(swap_col, pl.mul(swap_dup_f, 2.0)) - swap_next = pl.add(swap_col, 1.0) - swap_stride = pl.mul(swap_lane, 2.0) - swap_idx_f = pl.sub(swap_next, swap_stride) - rope_swap_idx[:, :] = pl.cast(swap_idx_f, target_type=pl.INT32) - cs_ones = pl.full([T, ROPE_INTERLEAVE_TILE], dtype=pl.FP32, value=1.0)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@models/deepseek/v4/decode_sparse_attn_swa.py` around lines 321 - 338, Remove the obsolete rope_swap_idx precomputation in the rope_cs block, including its swap_ones/swap_col and related intermediate calculations and the writes through pl.at. Keep the per-row mr_swap_idx logic used by the n_rope_bf16 loop unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@models/deepseek/v4/decode_sparse_attn_swa.py`:
- Around line 321-338: Remove the obsolete rope_swap_idx precomputation in the
rope_cs block, including its swap_ones/swap_col and related intermediate
calculations and the writes through pl.at. Keep the per-row mr_swap_idx logic
used by the n_rope_bf16 loop unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 966ba3a4-dce1-4166-87b6-cc6a84a5d214
📒 Files selected for processing (3)
models/deepseek/v4/decode_sparse_attn.pymodels/deepseek/v4/decode_sparse_attn_hca.pymodels/deepseek/v4/decode_sparse_attn_swa.py
…erge_norm swap idx Slice the head-invariant rope_swap_idx[0:1, :] instead of recomputing the [1, ROPE_DIM] swap-index pattern from scratch inside merge_norm. Every row of rope_swap_idx (filled in the rope_cs task) is identical, so the first row is mathematically equal to the rebuilt pattern. Slicing drops the redundant vector ops while keeping the per-row gather tile 1-row, so the a5 gather-lowering fix is preserved. A5-validated (decode_sparse_attn_swa, 4 runs): baseline 68.47/68.90% vs change 69.09/68.84% fail ratio -- within ~0.4% run-to-run noise; max abs diff unchanged (0.0011-0.0013). No correctness regression. Addresses gemini-code-assist review on PR hw-native-sys#801.
…ASS)
PV-cube fusion workaround (decode_sparse_attn{,_hca,_swa}.py, prefill_sparse_attn.py):
on a5 the composed qk_pv task mis-fuses the exp->cast(bf16)->PV-cube chain, losing
~15% precision (cube bf16xbf16->fp32 precise standalone at K=128/256/512 b_trans+
no-b_trans; softmax row_max+exp precise; KV gather correct; mi/li correct; only the
fused PV output drifts). Materializing qk_exp_bf16 to GM before the PV cube breaks
the fusion -> precision restored. +13 PASS (decode/prefill sparse_attn x4,
decode/prefill_attention_{csa,hca,swa} x6, decode_layer/decode_mtp/prefill_mtp x3).
Survey: 21 -> 34 PASS.
hc_pre: lift the a5 guard for the core-count-agnostic _hc_pre_separate impl (the
syncall impl still needs 910B). +1 PASS (decode+prefill).
Real fix pending in pypto/simpler: find the fusion that loses precision on a5; then
remove the 4-file GM workaround (perf cost: one GM round-trip per PV cube).
c8e391a to
00799b0
Compare
1 participant
What
Fixes an a5 (Ascend 950) correctness bug in the inverse-RoPE of DeepSeek-V4
sparse_attn(and the_hca/_swavariants).Root cause
On a5 the
tensor.gatherflat-index lowering (PR #2044 area) miscompiles when the source tile spans more than one FP32 vector box (>8 rows). The inverse-RoPE inmerge_normdid:H_TILE=16spans two vector boxes, so only tile rows 0 and 8 (the first head of each 8-headO_GROUP) received correct data — the ROPE columns of all other heads were corrupted.qkv_proj_ropeis unaffected because its rotation tile is exactly[8, ROPE_DIM](one box).Evidence (measured on a5)
Bypassing proj and emitting pre-proj
o_packeddirectly:Per-head ROPE rel-L2 before: ~0.15 for
n_hh==0heads vs ~0.44 forn_hh==1..7(the smoking gun). After fix: uniform ~0.15 across all heads.decode_sparse_attnfinal fail rate 59.5% → 47.7%.Fix
Unroll the inverse-RoPE per head-row using 1-row
[1, ROPE_DIM]tiles (mr_swap_idx/mr_cos_il/mr_sin_signedat[1, ROPE_DIM], loopn_hioverH_TILE,pl.mulon matched-shape 1-row tiles). Every gather now fits within a single FP32 vector box regardless ofH_TILE. The rotation math is identical to the original (computed per row), so a2a3 is unaffected.Applied to all three sparse_attn variants:
decode_sparse_attn.py,decode_sparse_attn_hca.py,decode_sparse_attn_swa.py.Scope note
A separate, uniform ~15% base-attention precision divergence on a5 remains in the
qk_pv/merge_normpath (individual ops pass but the composed attention is ~15% rel-L2 off vs <0.5% on a2a3). That is a different issue — likely a5 cube/vector precision at the compiler/simpler level, not a kernel lowering bug — and is tracked separately. This PR addresses only the inverse-RoPE gather lowering bug.