Perf: hoist DSv4-flash RoPE interleave tables out of per-block scopes - #845
Perf: hoist DSv4-flash RoPE interleave tables out of per-block scopes#845Hzfengsy wants to merge 1 commit 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 Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughChangesThe DeepSeek-V4 Flash RoPE path now preprocesses half-width cosine and sine tables into full-width interleaved and sign-folded buffers. Decode compressors and the indexer consume this ABI, while KV RoPE reuses previously prepared tables and swap indices. RoPE preprocessing and decode integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant q_rope_prepare
participant kv_rms_norm_rope
participant rope_interleave
participant indexer
participant compressor_ratio4
q_rope_prepare->>rope_interleave: prepare interleaved cos and sign-folded sin
q_rope_prepare->>kv_rms_norm_rope: provide prepared tables and swap indices
kv_rms_norm_rope->>kv_rms_norm_rope: rotate KV lanes
rope_interleave->>indexer: provide full-width RoPE tables
rope_interleave->>compressor_ratio4: provide full-width RoPE tables
indexer->>indexer: apply lane swap and rotation
compressor_ratio4->>compressor_ratio4: apply lane swap and rotation
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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 488dd7628f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| sin_signed[0:B, 0:ROPE_HEAD_DIM] = pl.mul( | ||
| pl.gather(sin_half[0:B, 0:HALF_ROPE], dim=-1, index=il_dup_idx), il_sign) |
There was a problem hiding this comment.
Every new decode path invokes rope_interleave as a @pl.jit.inline function, but the function reaches its end after this assignment without returning a value. PyPTO's inline parser requires each inline call expression to produce a result, so compiling any of these callers fails before the kernel can run; return one of the in-place output tensors, as other inline helpers do.
AGENTS.md reference: AGENTS.md:L8-L14
Useful? React with 👍 / 👎.
| cos_il = pl.create_tensor([B, ROPE_HEAD_DIM], dtype=pl.FP32) | ||
| sin_signed = pl.create_tensor([B, ROPE_HEAD_DIM], dtype=pl.FP32) | ||
| rope_interleave(cos, sin, cos_il, sin_signed) |
There was a problem hiding this comment.
compressor_test declares and binds B_DYN, but these new buffers and rope_interleave are fixed to DECODE_BATCH (4), and the helper unconditionally reads cos_half[0:4] and sin_half[0:4]. When the compiled dynamic entry is invoked with a smaller batch, the helper reads beyond the input rows; a larger batch would leave output rows unmaterialized. Derive the buffer/helper extent from the runtime batch dimension instead of the static fixture size.
AGENTS.md reference: AGENTS.md:L8-L14
Useful? React with 👍 / 👎.
![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-flash/decode_indexer.py (1)
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSame j^1 swap-index arithmetic chain duplicated across four kernel files. Each site independently rebuilds the identical
arange -> *0.5 -> trunc -> lane -> j^1sequence (necessary since pypto has no tensor-level bitwise XOR), purely as a function ofROPE_HEAD_DIM/tile row count. This PR already extracted the analogous cos/sin interleave logic intorope_interleave.py; the swap-index chain is a good candidate for the same treatment.
models/deepseek/v4-flash/decode_indexer.py#L169-189: extract thesw_col/sw_dup_f/sw_lane/swap-index computation into a shared helper (e.g.rope_swap_index(rows)), keep the existing one-shot hoist-and-reuse structure.models/deepseek/v4-flash/decode_compressor_ratio4.py#L234-245: replace the inlinerope_col/rope_dup_f/rope_lane/rope_swap_idxblock with a call to the shared helper.models/deepseek/v4-flash/decode_compressor_ratio128.py#L268-279: replace the inlinerope_col/rope_dup_f/rope_lane/rope_swap_idxblock with a call to the shared helper.models/deepseek/v4-flash/decode_indexer_compressor.py#L247-258: replace the inlinerope_col/rope_dup_f/rope_lane/rope_swap_idxblock with a call to the shared helper.🤖 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-flash/decode_indexer.py` at line 1, Extract the duplicated swap-index arithmetic into a shared rope_swap_index(rows) helper alongside the existing rope interleave utilities, preserving the arange, scaling, truncation, lane, and j^1 computation. Update decode_indexer.py to hoist and reuse one helper result, and replace the equivalent inline blocks in decode_compressor_ratio4.py, decode_compressor_ratio128.py, and decode_indexer_compressor.py with calls to the helper.
🤖 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-flash/decode_indexer.py`:
- Line 1: Extract the duplicated swap-index arithmetic into a shared
rope_swap_index(rows) helper alongside the existing rope interleave utilities,
preserving the arange, scaling, truncation, lane, and j^1 computation. Update
decode_indexer.py to hoist and reuse one helper result, and replace the
equivalent inline blocks in decode_compressor_ratio4.py,
decode_compressor_ratio128.py, and decode_indexer_compressor.py with calls to
the helper.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 954020e8-efcd-4f33-8e84-395c8bccf4cb
📒 Files selected for processing (8)
models/deepseek/v4-flash/decode_attention_csa.pymodels/deepseek/v4-flash/decode_attention_hca.pymodels/deepseek/v4-flash/decode_compressor_ratio128.pymodels/deepseek/v4-flash/decode_compressor_ratio4.pymodels/deepseek/v4-flash/decode_indexer.pymodels/deepseek/v4-flash/decode_indexer_compressor.pymodels/deepseek/v4-flash/qkv_proj_rope.pymodels/deepseek/v4-flash/rope_interleave.py
The A3 interleaved rotation needs cos_il[j] = cos_half[j>>1] and a sign-folded sin per rope column. Every consumer rebuilt that j>>1 dup-gather inside its own scope; pl.gather lowers to a per-row TGATHER loop, so the cost scaled with (blocks x rows) instead of table size -- the indexer's qr_rope alone spent 16 blocks x 32 rows x 2 tables per layer rebuilding one small position-invariant table. - Add rope_interleave: builds the interleaved cos and sign-folded sin once per layer over B rows - Widen the cos/sin ABI of compressor_ratio4 / ratio128, indexer and indexer_compressor to full ROPE_HEAD_DIM; attention_csa and attention_hca build the tables once and pass them in - indexer qr_rope: hoist the block-invariant j^1 swap index out of the spmd, fold the cos/sin row broadcast into col_expand_mul, and pack the nope/rope halves into one store - qkv_proj_rope kv_rms_norm_rope: reuse q_rope_prepare's cos_il / sign-folded sin / swap_idx instead of re-deriving them - indexer_compressor: pipeline the two rmsnorm_rope reduce loops Standalone test entries call rope_interleave themselves, so fixtures and goldens keep the half-width cos/sin ABI. Isolated decode_indexer qr_rope 235 -> 97 core-us summed over its 16 spmd blocks (decode, B=4 S=2, a2a3, 3-run median). The single-block rmsnorm_rope scopes, where the hoisted table replaces only a 16-row gather, move the other way at 7 -> 15 core-us, so layer wall time is unchanged; the net effect is freed core occupancy.
488dd76 to
afbea89
Compare
1 participant
Summary
The A3 interleaved rotation needs
cos_il[j] = cos_half[j>>1]and a sign-folded sin per rope column. Every consumer rebuilt thatj>>1dup-gather inside its own scope, andpl.gatherlowers to a per-rowTGATHERloop — so the cost scaled with (blocks x rows) rather than with the table size. The indexer'sqr_ropealone spent 16 blocks x 32 rows x 2 tables per layer rebuilding one small position-invariant table.rope_interleave: builds the interleaved cos and sign-folded sin once per layer overBrows. Folding the sign into sin is exact — multiplying by +/-1 only flips the sign bit, so(x*sign)*sinandx*(sin*sign)are bit-identical.compressor_ratio4/compressor_ratio128,indexerandindexer_compressorfromROPE_HEAD_DIM//2to fullROPE_HEAD_DIM;attention_csaandattention_hcabuild the tables once and pass them in.indexerqr_rope: hoist the block-invariantj^1swap index out of the spmd into its own scope, fold the cos/sin row broadcast intocol_expand_mulso nocos_il/sin_iltile is materialized, and pack the nope/rope halves into a single store.qkv_proj_ropekv_rms_norm_rope: reuseq_rope_prepare'scos_il/ sign-folded sin /swap_idxinstead of re-deriving the same arange chain and re-gathering the same two tables.indexer_compressor: pipeline the twormsnorm_ropereduce loops, and allow early resolve onscatter_softmax_pool.Standalone test entries call
rope_interleavethemselves, so fixtures and goldens keep the half-width cos/sin ABI.Isolated
decode_indexerqr_rope235 -> 97 core-us summed over its 16 spmd blocks (decode, B=4 S=2, a2a3, 3-run median). The single-blockrmsnorm_ropescopes, where the hoisted table replaces only a 16-row gather, move the other way at 7 -> 15 core-us, so layer wall time is unchanged; the net effect is freed core occupancy.