perf(dsv4 hca decode): early-resolve gather_kv/merge_norm + L2 swimlane level flag - #785
Conversation
…ne level flag - decode_sparse_attn_hca.py: allow_early_resolve on hca_gather_kv and merge_norm (complete the qk_pv / proj_a_mm producer sets so those flagged consumers can pre-stage). - decode_attention_hca.py: --enable-l2-swimlane takes a granularity level (1=AICore timing, 2=+dispatch/fanout, 3=+sched phases, 4=+orch phases; bare flag = 4/full).
📝 WalkthroughWalkthroughThe PR adds optional L2 swimlane performance levels and enables earlier dependency resolution for two sparse attention SPMD stages. ChangesDecode HCA updates
Estimated code review effort: 2 (Simple) | ~10 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 |
There was a problem hiding this comment.
Code Review
This pull request updates the --enable-l2-swimlane argument in decode_attention_hca.py to support new granularity levels (0-4) and adds allow_early_resolve=True to SPMD contexts in decode_sparse_attn_hca.py to enable early resolution. The reviewer recommends restoring the choices constraint for --enable-l2-swimlane with the updated range (0, 1, 2, 3, 4) to ensure robust command-line validation.
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.
| parser.add_argument( | ||
| "--enable-l2-swimlane", | ||
| nargs="?", | ||
| const=4, | ||
| default=0, | ||
| type=int, | ||
| metavar="PERF_LEVEL", | ||
| help="Enable L2 swimlane perf capture at the given granularity level. Bare flag " | ||
| "= level 4 (full). Levels: 1=AICore timing, 2=+dispatch/fanout, 3=+sched " | ||
| "phases, 4=+orch phases; 0 (default) disables.", | ||
| ) |
There was a problem hiding this comment.
The choices constraint was removed when updating --enable-l2-swimlane to support the new granularity levels. To prevent invalid inputs (e.g., values outside the 0-4 range) and ensure robust command-line validation, please restore the choices constraint with the updated valid range (0, 1, 2, 3, 4).
| parser.add_argument( | |
| "--enable-l2-swimlane", | |
| nargs="?", | |
| const=4, | |
| default=0, | |
| type=int, | |
| metavar="PERF_LEVEL", | |
| help="Enable L2 swimlane perf capture at the given granularity level. Bare flag " | |
| "= level 4 (full). Levels: 1=AICore timing, 2=+dispatch/fanout, 3=+sched " | |
| "phases, 4=+orch phases; 0 (default) disables.", | |
| ) | |
| parser.add_argument( | |
| "--enable-l2-swimlane", | |
| nargs="?", | |
| const=4, | |
| default=0, | |
| type=int, | |
| choices=(0, 1, 2, 3, 4), | |
| metavar="PERF_LEVEL", | |
| help="Enable L2 swimlane perf capture at the given granularity level. Bare flag " | |
| "= level 4 (full). Levels: 1=AICore timing, 2=+dispatch/fanout, 3=+sched " | |
| "phases, 4=+orch phases; 0 (default) disables.", | |
| ) |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/deepseek/v4/decode_attention_hca.py`:
- Around line 660-669: Restrict the --enable-l2-swimlane PERF_LEVEL argument in
the parser configuration to the documented values 0 through 4. Restore
choices=(0, 1, 2, 3, 4) or add equivalent validation before runtime_cfg is
constructed, while preserving the bare-flag default of level 4.
🪄 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: 32ac0e85-9ecf-4ec5-a934-946dd7d3d34f
📒 Files selected for processing (2)
models/deepseek/v4/decode_attention_hca.pymodels/deepseek/v4/decode_sparse_attn_hca.py
| parser.add_argument( | ||
| "--enable-l2-swimlane", | ||
| nargs="?", | ||
| const=4, | ||
| default=0, | ||
| type=int, | ||
| metavar="PERF_LEVEL", | ||
| help="Enable L2 swimlane perf capture at the given granularity level. Bare flag " | ||
| "= level 4 (full). Levels: 1=AICore timing, 2=+dispatch/fanout, 3=+sched " | ||
| "phases, 4=+orch phases; 0 (default) disables.", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Restrict PERF_LEVEL to the documented range.
Removing choices means values outside the supported levels, such as 5 or -1, are accepted and forwarded to run_jit. Restore choices=(0, 1, 2, 3, 4) or validate the range before constructing runtime_cfg.
Proposed fix
const=4,
default=0,
type=int,
+ choices=(0, 1, 2, 3, 4),
metavar="PERF_LEVEL",📝 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.
| parser.add_argument( | |
| "--enable-l2-swimlane", | |
| nargs="?", | |
| const=4, | |
| default=0, | |
| type=int, | |
| metavar="PERF_LEVEL", | |
| help="Enable L2 swimlane perf capture at the given granularity level. Bare flag " | |
| "= level 4 (full). Levels: 1=AICore timing, 2=+dispatch/fanout, 3=+sched " | |
| "phases, 4=+orch phases; 0 (default) disables.", | |
| parser.add_argument( | |
| "--enable-l2-swimlane", | |
| nargs="?", | |
| const=4, | |
| default=0, | |
| type=int, | |
| choices=(0, 1, 2, 3, 4), | |
| metavar="PERF_LEVEL", | |
| help="Enable L2 swimlane perf capture at the given granularity level. Bare flag " | |
| "= level 4 (full). Levels: 1=AICore timing, 2=+dispatch/fanout, 3=+sched " | |
| "phases, 4=+orch phases; 0 (default) disables.", |
🤖 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_attention_hca.py` around lines 660 - 669, Restrict
the --enable-l2-swimlane PERF_LEVEL argument in the parser configuration to the
documented values 0 through 4. Restore choices=(0, 1, 2, 3, 4) or add equivalent
validation before runtime_cfg is constructed, while preserving the bare-flag
default of level 4.
Summary
decode_sparse_attn_hca.py: addallow_early_resolvetohca_gather_kvandmerge_norm, completing the producer sets for the already-flaggedqk_pv/proj_a_mmconsumers so they can pre-stage (early dispatch).decode_attention_hca.py:--enable-l2-swimlanenow takes a granularity level (1=AICore timing, 2=+dispatch/fanout, 3=+sched phases, 4=+orch phases; bare flag = 4/full).Performance
hca_gather_kvis the main lever — it completesqk_pv's producer set soqk_pvcan pre-stage (its dispatch bubble → ~0) and stabilizes the HCA decode AICore time.merge_normis neutral (proj_a_mmis an O_GROUPS parallel reduction — core-contention-bound rather than dispatch-fanin-bound); kept for producer-set completeness.6-run AICore end-to-end on this config (a2a3,
--start-pos 8192): 310.9 / 312.9 / 314.4 / 319.0 / 320.7 / 326.7 us (median ~317, σ≈6).Validation
--start-pos 8192.