perf(dsv4): qproj wq_b transpose + serving contract; sparse_attn/indexer tuning#676
perf(dsv4): qproj wq_b transpose + serving contract; sparse_attn/indexer tuning#676Hzfengsy wants to merge 4 commits into
Conversation
…b_trans support Add .claude/rules/optimization-rules.md: - Rule 1: express a cube K-reduction as ONE stage=2 pipeline with an if-first branch (k0==0 -> matmul, else matmul_acc). Folding chunk 0 into the pipeline lets its prologue prefetch chunk 1, so an MTE2-bound cube stays busy from iter 0. - Rule 2: keep L2-traversing tile.load/store innermost dim >= the backend granularity (PH001 TileInnermostDimGranularity; 512B on a2a3, 128B on a5). cube-tile-tuning skill: add --b-trans (transposed [N,K], K-contiguous weight) support to tile_budget.py and hint_l1_tile.py with a layout-aware cache-line floor (N-contig for [K,N], K-contig for [N,K]); document the transpose lever in SKILL.md.
Fold each cube K-reduction into ONE stage=2 pipeline that starts at chunk 0 with an if-first branch (k==0 -> matmul, else matmul_acc), instead of peeling the first matmul outside the loop. Iteration 0 stays a real matmul (matmul_acc from a zero carry trips the TLOAD DN->NZ assertion, pypto#1540) but now sits inside the pipeline so its stage=2 prologue prefetches iter 1 -- the INT8 cube stays busy from the first chunk. Measured -15.7% proj_b_mm (722 -> 609us aggregate exec). Reference impl for optimization-rules Rule 1.
…t; qr_proj save-cores tile Store wq_b transposed [H*HEAD_DIM, Q_LORA] and consume it via b_trans=True so the qproj MTE2 weight stream is K-contiguous (DN2ZN load): fewer/longer bursts, measurably more MTE2-efficient. qproj_matmul 55.8 -> 31.9us, decode span ~-15us. Per-output-channel scale is transpose-invariant. Propagated through decode/prefill attention, layer, and fwd; the serving-framework weight-layout contract is consolidated under a grep-able [WEIGHT-PERMUTE] block in decode_layer.py and prefill_layer.py (serving must emit wq_b as [N,K], scale unchanged). Also qr_proj_matmul: QR_N_TILE 128 -> 256 (ON=4, OK=2 -> 8 spmd tasks) to free ~8 of the 24 cube cores for other work. This is a DELIBERATE core-headroom tradeoff: measured +5.6us on the isolated qkv leaf (seeded decode a2a3, 161.5 -> 167.1us median); the 2-scope core-budget sweep confirmed it, and the measurement is recorded inline.
Hoist the per-token query hadamard-transform + INT8 quant into the head of each score task (score is its sole consumer and IDX_N_HEADS == QH_QUANT_TILE gives a 1:1 per-token map), removing the standalone qr_hadamard_quant scope. Golden PASS (score + topk_idxs). MEASURED: this is a net regression vs the 2-scope baseline -- ~128us median (110-134, unstable) vs ~114us (112-123) on a2a3 decode. A cube B-operand must route through GM, so the fused intra-task query write->read serializes worse than the 2-scope form, which pipelines that vector->cube handoff across tasks. Kept as the default path per request; WIP pending the on-chip fix (feed the query as the matmul A-operand + transposed score epilogue). The 2-scope baseline is preserved locally at decode_indexer.py.bak.
📝 WalkthroughWalkthroughThis PR adds performance-rule documentation and extends cube-tile-tuning scripts with a ChangesDocumentation and Tuning Tools
Estimated code review effort: 2 (Simple) | ~12 minutes wq_b Transposed Layout Migration
Estimated code review effort: 4 (Complex) | ~60 minutes Pipeline Refactor and Indexer Fusion
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ScoreTask as score task
participant Hadamard as Hadamard transform
participant Quant as INT8 quantizer
participant QBuffer as qr_hadamard_i8
participant Matmul as score matmul
ScoreTask->>Hadamard: compute qh_nope/qh_rope from qr_proj_flat
Hadamard->>Quant: per-head amax reduction
Quant->>Quant: compute dequant scale
Quant->>QBuffer: write INT8 rows at qh_row0
ScoreTask->>Matmul: slice QBuffer[qh_row0] vs cached KV tile
Matmul-->>ScoreTask: accumulated score
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 introduces performance optimizations and structural updates for PyPTO-Lib kernels, specifically targeting DeepSeek-V4. It establishes a new weight-layout contract where the Q-LoRA up-projection weight wq_b is stored transposed as [H * HEAD_DIM, Q_LORA] and consumed via b_trans=True to ensure K-contiguous DMA transfers, significantly improving MTE2 efficiency. Correspondingly, optimization rules, tuning scripts (hint_l1_tile.py, tile_budget.py), and attention layer signatures across decode and prefill paths have been updated to support this transposed layout. Additionally, the matmul K-reductions in qkv_proj_rope.py and decode_sparse_attn.py were refactored into a single-pipeline if k==0 form to maximize load/compute overlap, and the query-side hadamard-quantization scope in decode_indexer.py was fused into the score task to reduce dispatch overhead. There are no review comments to assess, so I have no feedback to provide.
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.
![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: 9b887ccbd8
ℹ️ 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".
| # task: the INT8 query lands in the qr_hadamard_i8 GM buffer (a cube B-operand | ||
| # must come via GM) and its per-head dequant scale qh_scale_s [1, IDX_N_HEADS] | ||
| # stays in UB for reuse. Merging the scope drops the ~7us dispatch gap into score. | ||
| qh_row0 = qb + s * IDX_N_HEADS |
There was a problem hiding this comment.
For CSA decode layers, this newly inlined hadamard/quant path now runs unconditionally inside every score task, but the measurements added just above report the fused form as a net regression (~128us median vs ~114us for the previous two-scope path) because the same-task GM read-after-write serializes the query handoff. That means production decode takes the slower path by default; keep the two-scope implementation as the default or gate this WIP until the on-chip query rework is available.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
models/deepseek/v4/decode_indexer.py (1)
155-171: 🚀 Performance & Scalability | 🔵 TrivialFused path is a documented net regression shipping as the default decode path.
The block clearly documents that this fused form regresses (~128us vs ~114us baseline) and is intentional/WIP per request, with the 2-scope baseline preserved in
decode_indexer.py.bak. No correctness concern (golden passes), but shipping a known decode-latency regression as the default is worth tracking so the follow-up rework (keep the query on-chip as the matmul A-operand + transposed score epilogue) isn't lost.Want me to open a tracking issue capturing the on-chip A-operand rework and the
.bakbaseline reference?🤖 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_indexer.py` around lines 155 - 171, The default decode path in score is a documented latency regression, so make the fused qr_hadamard_quant-in-score flow non-default or revert the default to the preserved two-scope baseline. Update the score path selection in decode_indexer.py so the known-regressing fused implementation is behind a flag/WIP guard, and keep the baseline path as the default until the on-chip A-operand rework and transposed epilogue are implemented. Preserve the existing qr_hadamard_quant reference and the .bak baseline as the fallback/tracking anchor.
🤖 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 @.claude/skills/cube-tile-tuning/tile_budget.py:
- Around line 144-158: The `k_fit` computation in `tile_budget.py` uses
`a.cache_line` as if it were an element count in the `a.b_trans` branch, but
`k_max` is in elements; update the rounding to use the already computed
`cl_elems` (elements per cache line) in the `if a.b_trans` path of the
tile-budget logic. Keep the existing `n_fit` behavior, and make sure the
follow-up log message in the `print(...)` block still reports the corrected
`k_fit` value for `--b-trans` mode.
---
Nitpick comments:
In `@models/deepseek/v4/decode_indexer.py`:
- Around line 155-171: The default decode path in score is a documented latency
regression, so make the fused qr_hadamard_quant-in-score flow non-default or
revert the default to the preserved two-scope baseline. Update the score path
selection in decode_indexer.py so the known-regressing fused implementation is
behind a flag/WIP guard, and keep the baseline path as the default until the
on-chip A-operand rework and transposed epilogue are implemented. Preserve the
existing qr_hadamard_quant reference and the .bak baseline as the
fallback/tracking anchor.
🪄 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: 26104726-cae2-4024-a397-03d28e3d23dc
📒 Files selected for processing (17)
.claude/rules/optimization-rules.md.claude/skills/cube-tile-tuning/SKILL.md.claude/skills/cube-tile-tuning/hint_l1_tile.py.claude/skills/cube-tile-tuning/tile_budget.pymodels/deepseek/v4/decode_attention_csa.pymodels/deepseek/v4/decode_attention_hca.pymodels/deepseek/v4/decode_attention_swa.pymodels/deepseek/v4/decode_fwd.pymodels/deepseek/v4/decode_indexer.pymodels/deepseek/v4/decode_layer.pymodels/deepseek/v4/decode_sparse_attn.pymodels/deepseek/v4/prefill_attention_csa.pymodels/deepseek/v4/prefill_attention_hca.pymodels/deepseek/v4/prefill_attention_swa.pymodels/deepseek/v4/prefill_fwd.pymodels/deepseek/v4/prefill_layer.pymodels/deepseek/v4/qkv_proj_rope.py
| if a.b_trans: | ||
| # K is contiguous -> keep K on a cache-line multiple; N is free (16-multiple). | ||
| n_fit = max(0, int(n_max // 16 * 16)) | ||
| k_fit = max(0, int(k_max // a.cache_line * a.cache_line)) | ||
| else: | ||
| # N is contiguous -> keep N >= cache-line floor; K is free (16-multiple). | ||
| n_fit = max(0, int(n_max // 16 * 16)) | ||
| k_fit = max(0, int(k_max // 16 * 16)) | ||
| print(f" Mat is the wall: at K={a.K}, N<={n_fit} fits; at N={a.N}, K<={k_fit} fits.") | ||
| print(f" -> trade K down to {k_fit} (>= cache-line {a.cache_line}) to keep N={a.N}, " | ||
| f"or shrink N to {n_fit}.") | ||
| if a.b_trans: | ||
| print(f" -> weight is K-contiguous: keep K>={a.cache_line}B (K>={cl_elems} elems), " | ||
| f"trade N down to {n_fit} to keep K={a.K}, or shrink K to {k_fit}.") | ||
| else: | ||
| print(f" -> weight is N-contiguous: keep N>={a.cache_line}B (N>={cl_elems} elems), " | ||
| f"trade K down to {k_fit} to keep N={a.N}, or shrink N to {n_fit} (>= {cl_elems}).") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
k_fit rounds by byte count instead of element count in the --b-trans branch.
k_max is an element count, but Line 147 rounds it to a multiple of a.cache_line (bytes), not cl_elems (elements-per-cache-line, already computed at Line 133). This only happens to be correct when bytes_in == 1 (INT8, the CLI default); for BF16 (bytes_in=2) or FP32 weights, cl_elems is 256/128 while the code rounds to 512, silently producing wrong (often 0) tile-shrink suggestions in --b-trans mode — exactly the mode this PR adds support for.
🐛 Proposed fix
if a.b_trans:
# K is contiguous -> keep K on a cache-line multiple; N is free (16-multiple).
n_fit = max(0, int(n_max // 16 * 16))
- k_fit = max(0, int(k_max // a.cache_line * a.cache_line))
+ k_fit = max(0, int(k_max // cl_elems * cl_elems))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if a.b_trans: | |
| # K is contiguous -> keep K on a cache-line multiple; N is free (16-multiple). | |
| n_fit = max(0, int(n_max // 16 * 16)) | |
| k_fit = max(0, int(k_max // a.cache_line * a.cache_line)) | |
| else: | |
| # N is contiguous -> keep N >= cache-line floor; K is free (16-multiple). | |
| n_fit = max(0, int(n_max // 16 * 16)) | |
| k_fit = max(0, int(k_max // 16 * 16)) | |
| print(f" Mat is the wall: at K={a.K}, N<={n_fit} fits; at N={a.N}, K<={k_fit} fits.") | |
| print(f" -> trade K down to {k_fit} (>= cache-line {a.cache_line}) to keep N={a.N}, " | |
| f"or shrink N to {n_fit}.") | |
| if a.b_trans: | |
| print(f" -> weight is K-contiguous: keep K>={a.cache_line}B (K>={cl_elems} elems), " | |
| f"trade N down to {n_fit} to keep K={a.K}, or shrink K to {k_fit}.") | |
| else: | |
| print(f" -> weight is N-contiguous: keep N>={a.cache_line}B (N>={cl_elems} elems), " | |
| f"trade K down to {k_fit} to keep N={a.N}, or shrink N to {n_fit} (>= {cl_elems}).") | |
| if a.b_trans: | |
| # K is contiguous -> keep K on a cache-line multiple; N is free (16-multiple). | |
| n_fit = max(0, int(n_max // 16 * 16)) | |
| k_fit = max(0, int(k_max // cl_elems * cl_elems)) | |
| else: | |
| # N is contiguous -> keep N >= cache-line floor; K is free (16-multiple). | |
| n_fit = max(0, int(n_max // 16 * 16)) | |
| k_fit = max(0, int(k_max // 16 * 16)) | |
| print(f" Mat is the wall: at K={a.K}, N<={n_fit} fits; at N={a.N}, K<={k_fit} fits.") | |
| if a.b_trans: | |
| print(f" -> weight is K-contiguous: keep K>={a.cache_line}B (K>={cl_elems} elems), " | |
| f"trade N down to {n_fit} to keep K={a.K}, or shrink K to {k_fit}.") | |
| else: | |
| print(f" -> weight is N-contiguous: keep N>={a.cache_line}B (N>={cl_elems} elems), " | |
| f"trade K down to {k_fit} to keep N={a.N}, or shrink N to {n_fit} (>= {cl_elems}).") |
🤖 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 @.claude/skills/cube-tile-tuning/tile_budget.py around lines 144 - 158, The
`k_fit` computation in `tile_budget.py` uses `a.cache_line` as if it were an
element count in the `a.b_trans` branch, but `k_max` is in elements; update the
rounding to use the already computed `cl_elems` (elements per cache line) in the
`if a.b_trans` path of the tile-budget logic. Keep the existing `n_fit`
behavior, and make sure the follow-up log message in the `print(...)` block
still reports the corrected `k_fit` value for `--b-trans` mode.
1 participant
Summary
DeepSeek-V4 decode perf work plus the
cube-tile-tuningtooling behind it. Four independent commits.1.
docs(cube-tile-tuning)— rules + b_trans tooling.claude/rules/optimization-rules.md: Rule 1 (single-pipeline if-first cube K-reduction) + Rule 2 (PH001 L2 innermost-dim granularity, 512 B a2a3 / 128 B a5).--b-trans(transposed[N,K]K-contiguous weight) intile_budget.py/hint_l1_tile.py; transpose-lever section inSKILL.md.2.
perf(dsv4 sparse_attn)— if-first pipeline for proj_a/proj_bFold each cube K-reduction into one
stage=2pipeline (k==0 → matmul, elsematmul_acc); iter 0 stays inside the pipeline so its prologue prefetches iter 1. −15.7 % proj_b_mm (722 → 609 µs aggregate). Reference impl for Rule 1.3.
perf(dsv4 qproj)— wq_b transpose + serving contract + save-cores ✅ main winwq_btransposed[H*HEAD_DIM, Q_LORA], consume viab_trans=True→ K-contiguous MTE2 (DN2ZN load). qproj_matmul 55.8 → 31.9 µs; decode span ~−15 µs. Per-output-channel scale is transpose-invariant. Propagated through decode/prefill attention + layer + fwd.wq_bas[N,K](scale unchanged). Contract consolidated under a grep-able[WEIGHT-PERMUTE]block indecode_layer.py/prefill_layer.py.qr_proj_matmulQR_N_TILE128 → 256 (8 tasks) to free ~8 cube cores — a deliberate +5.6 µs qkv-leaf tradeoff for core headroom (seeded a2a3, 161.5 → 167.1 µs median; recorded inline).4.⚠️ WIP / regression
perf(dsv4 indexer)— fuse qr_hadamard_quant into scoreHoist the per-token query hadamard-transform + INT8 quant into the
scoretask head (sole consumer;IDX_N_HEADS == QH_QUANT_TILE→ 1:1 per-token map). Golden PASS (score + topk_idxs), but a measured net regression — ~128 µs median (110–134, unstable) vs ~114 µs (112–123). A cube B-operand must route through GM, so the fused intra-task write→read serializes worse than the 2-scope pipeline. Kept per request; WIP pending the on-chip fix (query as matmul A-operand + transposed score epilogue). Baseline preserved locally atdecode_indexer.py.bak(==HEAD~1).Why draft
Commit 4 is an intentional WIP regression. Commit 3's wq_b transpose is the ready-to-merge win but shares this branch; can split on request.
Validation
**PASS** (exit 0).