Fuse QKV prologue into Qwen3-14B paged attention#789
Open
Little-oil wants to merge 3 commits into
Open
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe PR introduces an attention-only CCE kernel entrypoint, adds ABI-specific tensor indexing and fused RoPE/QKV preparation, routes decode callers through the new bridge, and updates decode orchestration to pass SPMD outputs directly into fused attention. ChangesQwen3 attention execution
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Decode
participant SPMD
participant CCEBridge
participant AttentionKernel
Decode->>SPMD: produce q_proj, k_proj, v_proj, and normalization state
Decode->>CCEBridge: submit projections, RoPE tables, and paging metadata
CCEBridge->>AttentionKernel: invoke attention-only entry
AttentionKernel-->>Decode: return attention output
Possibly related PRs
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}
Contributor
There was a problem hiding this comment.
Code Review
This pull request fuses the RoPE and QK-norm prologue directly into the CCE paged attention kernel (paged_attention_cce), removing the standalone rope_qkv SPMD dispatch in decode_fwd.py. It also introduces paged_attention_only_cce for the attention-only path. Feedback on the changes highlights a critical Write-After-Read (WAR) hazard in the newly added run_qwen_rope_qkv kernel function on Ascend C. Because the MTE2 and Vector coprocessor queues operate asynchronously, the local buffers x and weight can be prematurely overwritten by subsequent iterations. To resolve this, it is recommended to implement and call a Vector-to-MTE2 synchronization helper (qwen_wait_v_to_mte2()) at the end of each processing step.
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.
Little-oil
pushed a commit
to Little-oil/pypto-lib
that referenced
this pull request
Jul 16, 2026
- synchronize Vector reads before reusing RoPE UB via MTE2 - clamp serving warmup slot sentinel to the scratch page
Little-oil
pushed a commit
to Little-oil/pypto-lib
that referenced
this pull request
Jul 16, 2026
- preserve Vector RAW ordering throughout fused RMSNorm and RoPE - reduce all 128 head elements with an explicit count mask - clamp serving profile KV slot sentinels in prefill
Contributor
Author
|
cast没向0取整 |
Little-oil
force-pushed
the
feat/qwen3-14b-fused-attention-prologue
branch
3 times, most recently
from
July 17, 2026 01:16
4cf093d to
fa4832a
Compare
added 3 commits
July 17, 2026 01:18
…t on a2a3) First half of fusing rope_qkv into paged_attention_cce: a standalone CCEC extern `paged_attention_rope_only_cce` that runs ONLY the rope_qkv producer (per-head QK-norm + NeoX RoPE + paged KV-cache writes), with a torch golden test (test_rope_qkv_cce.py) validating query + the written k_cache/v_cache rows byte-exact against decode_fwd.py's rope_qkv. decode_fwd.py and the attention extern are untouched (serving path unchanged); this only adds the isolated kernel + test. Stage 2 (fuse into attention + cross-core sync) is pending. The C++ rope compute (`run_qwen_rope_qkv` in fai_body.hpp, vec-only, AIC no-op) mirrors decode_fwd.py: x=proj*inv_rms; inv_head=rsqrt(sum(x^2)/128+eps); x=(x*weight)*inv_head; NeoX half-split rotate; cast BF16; write paged cache. Four C220 (a2a3) hazards were pinned by dumping intermediates to GM via the on-device golden test (iterated locally through task-submit) and fixed: 1. Dynamic-index GetValue on a UB LocalTensor misreads (~128x). The S-side sum loop `for i square.GetValue(i)` returned ss~7.3 vs the true tree-sum 0.057 (inv_head 4.19 vs 84.9). Constant-index UB GetValue (weight/cos GetValue(0)) and dynamic-index GM GetValue (inv_rms.GetValue(batch)) are both fine. Fixed by summing the 8 tree partials with literal constant-index GetValues. 2. C220 V-pipe does NOT serialize dependent ops on the same buffer (intermittent RAW; surfaces as a clean ratio when it fails). Added PipeBarrier<PIPE_V>() between every dependent V op: Muls(x,inv_rms)->Mul(square,x,x) (was reading unscaled x, ss inflated 1/inv_rms^2); the reduce Muls/Adds/Rsqrt chain; the rope lo/hi Mul->Sub/Add; and the V-path Muls(x,inv_rms)->Cast. 3. count<8 (sub-block) AscendC::Add HANGS on C220 (507018). The tree uses 64/32/16/8; the final 8->1 sum is done in S (constant-index GetValue), not via more V Adds. 4. MTE3->MTE2 DMA tear: MTE2 (GM->UB load) and MTE3 (UB->GM store) share the external-memory DMA path, so a load right after a store with only a qwen_mte3_to_v (MTE3->V) fence tears the in-flight store (intermittent whole-head GM garbage, proportional to write count; key_cache was spared because its write->read gap ran through V's Muls+Cast and drained naturally). Added qwen_mte3_to_mte2() (HardEvent::MTE3_MTE2) after every UB->GM store. Result: test_rope_qkv_cce.py PASSES query+key_cache+value_cache on a2a3.
Fuses the QKV prologue (QK-norm + RoPE + paged KV-cache writes) into the mixed-core attention extern so the decode path no longer runs a separate rope_qkv SPMD task. Stage 1's isolated rope compute (proven byte-exact on a2a3, commit fa4832a) is now the AIV-side prologue of the attention launch. fai_body.hpp - run_qwen_fai arg indices switched by QWEN_FAI_ATTENTION_ONLY_ABI: standalone attention keeps 0..7; the fused extern prepends the rope inputs (q_proj..rope_sin 0..9, block_table 10) so the attention tensors shift to out=11, query/k/v=12..14, workspace/metadata=15..16, cache_row_offset=17. - rope_qkv producer guarded by !QWEN_FAI_ATTENTION_ONLY_ABI (compiled for both the fused entry and the rope_only isolation entry); its output arg indices switch by QWEN_FAI_ROPE_ONLY_ABI. The Stage 1 compute fixes (UB constant GetValue, V-pipe PipeBarrier<PIPE_V> between every dependent op incl MTE3->MTE2) carry over. attention/entry.cpp (fused) - acquire_metadata -> run_qwen_rope_qkv -> sync_qwen_rope_producers -> run_qwen_fai. - sync_qwen_rope_producers: AIV lanes barrier over the rope-ready metadata slots (pto::SYNCALL_SOFT_AIV_BARRIER), then one sub-lane per AIC pair signals its AIC via CrossCoreSetFlag<0x2, PIPE_MTE3>; the AIC CrossCoreWaitFlag before attention. PIPE_MTE3 so the signal fires only after rope's UB->GM stores flush (the cross-pipe GM visibility a plain SYNCALL<Mix> lacks). MODE=0x2 + pto primitives verified valid on a2a3. attention_only/entry.cpp (NEW) - defines QWEN_FAI_ATTENTION_ONLY_ABI -> run_qwen_fai only (no rope/sync); preserves the standalone attention ABI for test_paged_attention_cce.py. paged_attention_cce.py - paged_attention_cce is now the fused 18-arg extern (attention/entry.cpp). - paged_attention_only_cce (NEW) = the old 8-arg attention ABI (attention_only entry); qwen_decode_attention_cce + cache_offset_test driver route to it. - METADATA_BYTES bumped for the rope-ready soft-barrier region (27840 -> 29376, matches metadata_layout.h kMetadataBytes). metadata_layout.h + tiling/entry.cpp - rope-ready slot region constants (48 lanes x 32B) after the aligned FAI barrier; the tiling kernel zeroes it (clear_rope_ready) each launch. decode_fwd.py - deleted the standalone rope_qkv SPMD task; fa_fused now depends on [q_proj,k_proj,v_proj,rms,pa_tiling,attn_out_seed,mlp_out_seed] and calls the fused paged_attention_cce passing the raw projections + norm/rope/slot inputs, with attn_out_tnd as the Out return and q_tnd/k_cache/v_cache as InOut (written by the AIV prologue). Local validation: test_rope_qkv_cce (rope_only path) still PASSES on a2a3; both extern paths codegen-clean. The fused serving path (cross-core sync) is validated via CI (test_qwen3_accuracy).
The Stage2 sync_qwen_rope_producers (pto::SYNCALL<Soft,Mix>) deadlocked on
a2a3: pto-isa's no-arg topology builtins (get_subblockdim/get_subblockid)
read 0 in this pypto mixed+dual_aiv launch -- the dispatch sets no subblockdim
register -- so SYNCALL_GET_MIX_PARTICIPANT_COUNT under-counted (24 not 72)
and collided the AIV participant indices (507018 running-stalled, stuck
fa_fused). The vendor FAI kernel avoids this by using only the args-based
getters (get_block_idx(args)/get_sub_block_id(args)/get_block_num(args)).
Replace the MIX syncall with an args-based handoff that matches the vendor's
topology view:
- AIV lanes soft-barrier via pto::SYNCALL_SOFT_AIV_BARRIER with idx
get_block_idx(args)*2 + get_sub_block_id(args) (48 lanes);
- each AIC polls the 48 rope-ready GM slots (DCCI+dsb) until all non-zero.
Reverts the rope-ready region to 48 AIV slots (METADATA_BYTES 29376).
Fused decode_fwd (rope_qkv inside the mixed-core attention extern) now
PASSES on a2a3: exit=0, swimlane complete, Total Test Time ~3631us.
Little-oil
force-pushed
the
feat/qwen3-14b-fused-attention-prologue
branch
3 times, most recently
from
July 18, 2026 11:15
2fa426a to
f2c48ae
Compare
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.
Summary
Related Issues