perf(qwen3-14b decode): fold RoPE into the CANN FAI attention extern

[ChaoWao]

Summary

• Fold QK-norm and RoPE into the CANN FAI attention extern, removing the
  standalone RoPE dispatch and its producer-to-consumer dependency edge.
• Run the generated RoPE body on its original 32 logical AIV workers; the
  remaining AIV workers skip RoPE but still participate in the mixed-core
  barrier and attention phase.
• Remove the two redundant dsb(DSB_DDR) calls around the RoPE-to-attention
  SyncAll<false> boundary while retaining the cache-backed metadata barrier.
• Add a contract regression test that ties the C++ worker guard to the Python
  ROPE_CORES setting and generated RoPE stride.
• Add a reusable guide for on-device InCore timestamp capture, exact phase
  partitioning, and reconciliation with L2 task timing.
• Record the C220/CANN 9.0.0 no-DSB evidence, validation matrix, near-tie
  diagnosis, and re-evaluation boundary in a dated investigation note.
• Validate the no-DSB candidate in ten synthetic 40-layer-plus-LM-head device
  runs: 9 clean, 1 reference near-tie, 0 hard failures, sampled token 160/160,
  all 24,330,240 logits within 5e-2, and no NaN or Inf. The synthetic harness
  repeats one generated layer and is not a real-checkpoint accuracy test.

Related Issues

None.

[coderabbitai]

📝 Walkthrough

Walkthrough

Qwen3 decode attention now uses a fused CANN paged-attention entry that performs QK-norm and RoPE phase processing, publishes rotated Q/K/V data, and then executes attention with updated scheduling dependencies and inputs.

Changes

Qwen3 fused RoPE attention

Layer / File(s)	Summary
Attention contract and decode wiring models/qwen3/14b/paged_attention_cce.py, models/qwen3/14b/decode_fwd.py	Adds and exports paged_attention_rope_cce, then routes decode attention through it with raw projections, normalization weights, RoPE tables, inverse-RMS state, and paging inputs.
Kernel entry and attention orchestration models/qwen3/14b/kernels/paged_attention_cce/attention_rope/entry.cpp, models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp	Adds the CANN entry point, metadata handling, optional RoPE execution, typed tensor access, and mixed all-core synchronization before attention.
Phase-0 normalization and RoPE execution models/qwen3/14b/kernels/paged_attention_cce/kernel/rope_phase0.hpp	Implements inverse-RMS/QK normalization, RoPE rotation, Q output generation, and paged K/V cache writes for assigned KV-head and batch items.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant DecodeLayer
  participant AttentionBridge
  participant KernelEntry
  participant RopePhase0
  participant PagedAttention
  DecodeLayer->>AttentionBridge: Submit projections and paging inputs
  AttentionBridge->>KernelEntry: Launch fused attention kernel
  KernelEntry->>RopePhase0: Normalize and rotate query and key
  RopePhase0->>KernelEntry: Return completed phase zero writes
  KernelEntry->>PagedAttention: Synchronize and execute attention

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#762: Updates the same Qwen3 decode attention scheduling and RoPE/QK-norm path.
• hw-native-sys/pypto-lib#789: Reworks the paged-attention pipeline to consume projections and RoPE/norm inputs directly.
• hw-native-sys/pypto-lib#792: Changes the interaction between inverse-RMS normalization and RoPE rotation.

Poem

I’m a bunny in the kernel, hopping Q and K,
Norm then spin the RoPE wreath before attention’s way.
Cache the heads and signal all,
Let fused computations call—
Clean decode paths now bloom today!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: folding RoPE into the CANN attention extern for Qwen3-14B decode.
Description check	✅ Passed	The description is related to the changeset and accurately describes the RoPE/attention fusion at a high level.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request refactors the Qwen3-14B decode forward pass by folding the QK-norm and RoPE operations directly into the external paged attention kernel as an in-kernel phase 0, replacing the standalone rope_qkv block with paged_attention_rope_cce. Feedback on the implementation points out a critical issue in rope_phase0.hpp where AscendC::ReduceSum is called with overlapping srcLocal and workLocal buffers (sq), which violates Ascend C constraints and can cause undefined behavior. It is recommended to use the idle normed tensor as the scratch workspace instead.

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]

Actionable comments posted: 2

🤖 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.

Inline comments:
In `@models/qwen3/14b/kernels/paged_attention_cce/kernel/rope_phase0.hpp`:
- Around line 23-28: Correct the phase-0 documentation to distinguish Q/K RoPE
rotation from V scaling: update
models/qwen3/14b/kernels/paged_attention_cce/kernel/rope_phase0.hpp lines 23-28
to describe rotated Q/K and inverse-RMS-scaled V, update
models/qwen3/14b/decode_fwd.py lines 471-475 to say “RoPE'd K and scaled V,” and
make the same distinction in
models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp lines 194-196
before the synchronization description.
- Around line 35-37: Update the documentation comment near the paged cache row
calculation to identify cache_row_base as args[17], matching the fused ABI and
the implementation’s scalar index. Keep the cache-row formula and surrounding
explanation unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: f408691c-b8c0-4562-a20b-fdcd40f6b1b2

📥 Commits

Reviewing files that changed from the base of the PR and between f894b63 and 00cd688.

📒 Files selected for processing (5)

• models/qwen3/14b/decode_fwd.py
• models/qwen3/14b/kernels/paged_attention_cce/attention_rope/entry.cpp
• models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp
• models/qwen3/14b/kernels/paged_attention_cce/kernel/rope_phase0.hpp
• models/qwen3/14b/paged_attention_cce.py

[ChaoWao]

Investigation snapshot — attention still garbage (BLOCKED)

Pushed docs/investigations/2026-07-qwen3-rope-fai-fusion-accuracy.md with the full localization. TL;DR:

Decisively proven (a2a3, partial dump via pl.dump_tag + enable_dump_args=1):

• The RoPE writes bit-exact-correct K/V to the paged cache, and metadata.kv_lengths/cumulative_q are correct — yet the FA output is garbage. So it's not the rope, coherency, metadata, or write location.
• RoPE no-op isolation (if (false) rope_qkv(...), keep barrier+FA): the FA becomes perfect (attn_out == seed V bit-exact). → the fusion structure is fine; the RoPE's execution corrupts the FA compute, globally.

Ruled out on device (no change): read-coherency (full dcci = no-op), FD-combine (force NonFd), UB/TPipe reset, vector-mask reset, noinline, restricting rope to lanes 0-31, and removing SyncAll<false>. The compiled FA path uses resource.ptoTopology (not the broken native GetSubBlockIdx).

Masking traps (why it looked fine): fixture wo is tiny (out≈residual) + LM-head argmax/logits are robust → single-layer golden false-passes even with attention broken (guide §6). The sharp signal is a strict ratio_allclose on attn_out (99.7% wrong).

Sub-result: hand-written rope replaced with the verbatim pypto-generated rope_qkv (rope now bit-exact); the FA blocker is independent of the rope source.

Next: in-rope_qkv return; bisection (guide §7.3) + partial-dump of attn_out to bracket the corrupting statement, then reset the shared pto↔AscendC hw/SPR/framework state it touches.

[ChaoWao]

✅ Root-caused, fixed, and verified on-device (a2a3)

Confirming the fix in 8b123fe from an independent debugging pass (partial
dump_tag on the fused op output vs. what out_proj reads):

• The FA output buffer was always correct; out_proj consumed a different
  buffer that was bit-exact RoPE'd Q (q_tnd). A pl.jit.extern returns its
  first mutable arg; folding the rope made query: pl.InOut lead the
  signature, so attn_out = paged_attention_rope_cce(...) bound to query, not
  out. Placing out first (+ WithRope-conditional ABI remap) fixes it.
• Ruled out en route (each had no effect, so not the cause): a writer-side
  PipeBarrier<PIPE_ALL>+dsb drain, dropping allow_early_resolve, and out: pl.InOut — the wrong value was deterministic, i.e. a wrong-buffer bind, not a
  write-back race or early dispatch.

Strict per-element out golden (ratio_allclose atol=rtol=3e-3, ≤2% outliers, not the argmax/logits proxy):

case	result
seq_len=1, seed 1234	PASS
seq_len=1, seed 7	PASS
varied multi-block KV (seq_lens up to 3756), seed 1234	PASS

test_fused_attention_declares_real_output_first passes.

⚠️ Unrelated pre-existing CI note

test_loaded_kernel_signatures_match_contract_arg_counts and
test_compile_arg_builders_follow_loaded_stage_specs fail on the prefill
stage (prefill_fwd has 25 params, its contract spec lists 24). Both test
functions pre-date this PR (they exist on main); this PR touches neither
prefill_fwd nor the contract registry, so the drift is pre-existing (last
prefill change on main is #799) and should be triaged separately — it does
not gate this fix.