perf(dsv4 decode swa): deferred attention RMSNorm and dep cleanup#791
perf(dsv4 decode swa): deferred attention RMSNorm and dep cleanup#791zhangqi-chen wants to merge 4 commits into
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:
📝 WalkthroughWalkthroughCSA and SWA decode paths now inline projection and normalization kernels with local Torch references. Sparse SWA task stages no longer use explicit dependency tokens for gather, RoPE, and merge operations. ChangesDecode kernel updates
Estimated code review effort: 4 (Complex) | ~45 minutes 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}
There was a problem hiding this comment.
Code Review
This pull request inlines the rms_norm, attn_norm_xgamma, and qkv_proj_rope functions directly into the CSA and SWA attention decode files, adds PyTorch golden reference implementations for testing, and cleans up unused task IDs and dependencies in SWA sparse attention. The review feedback suggests optimizing qr_rms_norm_quant in the CSA path by deferring inv_rms and avoiding expensive vector reciprocal instructions, as well as removing allow_early_resolve=True from q_rope_prepare for consistency and reduced overhead.
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.
- Copy RMSNorm and QKV projection/RoPE kernels into SWA and CSA decode modules - Specialize the copied signatures for the fixed decode token count - Keep Torch golden helpers local to remove shared module dependencies Standalone SWA and CSA tests pass with PTOAS 0.48. The combined decode layer still reports duplicate local kernel names and is intentionally held.
Port of the qkv-only deferred attention RMSNorm to swa-perf's self-contained decode_attention_swa. The pre-qkv inv_rms is a per-token positive scalar, and qkv_proj_rope re-RMSNorms BOTH qr and kv, so inv_rms cancels exactly in q/kv/qr /qr_scale. SWA's x_normed feeds only qkv, so we drop inv_rms entirely: new attn_norm_xgamma stores just xg = x*gamma, replacing rms_norm (which is removed -- it had no remaining caller). Being pure elementwise (no reduction), it needs no physical-8 tile-floor trick -- each core writes its own [1, XG_D_TILE] rows. Parallelized to one token per core (pl.spmd(t_dim)) with a wide XG_D_TILE=1024. attn_norm_xgamma Exec ~2.8us (vs rms_norm ~10.8us). decode_attention_swa a2a3 PASS (x_out, kv_cache), ptoas 0.48. hc_post finish ~272us.
Dependency-chain cleanup guided by deps.json (dep-gen): - decode_attention_swa: drop allow_early_resolve from q_rope_prepare; it is not on the qkv critical chain, so it need not be early-dispatched. - decode_sparse_attn_swa: remove explicit deps that duplicate tensormap (data) edges -- the scheduler already orders these through the tensor dependency: * qk_pv deps=[gather_tids[0]] (reads swa_kv_flat; affine writes tracked) * merge_norm deps=[qk_tid, rope_tid] (reads qk_pv sparse_blk_* + rope_cs rope_*) Dead tid captures (gather_tid/qk_tid/rope_tid) removed; merge_tid kept (proj_a_mm). deps.json confirms the explicit edges are gone (tensormap edges remain). decode_attention_swa a2a3 PASS (x_out, kv_cache) 5/5, ptoas 0.48.
Symmetric int8 quant of qr*gamma cancels the per-token inv_rms exactly, so the quant pass can use amax(qr*gamma) directly and inv_rms rides only qr_scale (dequant) -- same transform as gate ffn_norm 37d5dee. Drops the redundant row_expand_mul(inv_rms) from the quant pipeline; int8 output unchanged. Marginal (qr_rms_norm_quant ~5.1->~5.0us; within hc_post noise) but strictly less work. decode_attention_swa a2a3 PASS (x_out, kv_cache), ptoas 0.48.
1 participant
Summary
decode_attention_swaso the SWA decode path is self-contained.attn_norm_xgamma): the pre-qkvinv_rmsis a per-token positive scalar, andqkv_proj_ropere-RMSNorms BOTH the qr and kv projections downstream, soinv_rmscancels exactly in every qkv output (q / kv / qr int8 / qr_scale). SWA's normalized input feeds only qkv, soinv_rmsis dropped entirely and justxg = x*gammais stored — removing the sq_sum reduction, the rsqrt, and the whole normalize pass. Being pure elementwise (no reduction), it needs no physical-row tile-floor workaround, so it is parallelized to one token per core with a wide D-tile. The now-unused localrms_normis removed.inv_rms.inv_rmsinqr_rms_norm_quant(same transform as the gateffn_norm): symmetric int8 quant ofqr*gammacancelsinv_rmsexactly, so the quant usesamax(qr*gamma)directly andinv_rmsrides onlyqr_scale(dequant). int8 output unchanged.deps.json):allow_early_resolvefromq_rope_prepare(not on the qkv critical chain).qk_pv deps=[gather_tids[0]]andmerge_norm deps=[qk_tid, rope_tid].deps.jsonconfirms only the tensormap edges remain.Validated on real a2a3 (ptoas 0.48):
x_outandkv_cachePASS. The attention-norm task drops from ~10.8us (rms_norm) to ~2.6us (attn_norm_xgamma, 1 token/core). At start_pos=8k the hc_post finish is ~281us and stable (10 runs, 276.6-285.1us).Related Issues
None