fix(pt): pass PACKED and Wigner strides in the flash-bwd tuning sweep - #6015
fix(pt): pass PACKED and Wigner strides in the flash-bwd tuning sweep#6015iProzd wants to merge 1 commit into
Conversation
`dp --pt freeze` of a DPA4/SeZM checkpoint crashed in the Triton tuning
sweep at the default `DP_TRITON_INFER=2` on GPUs whose built-in table has
no `flash_bwd_edge` family (for example NVIDIA H20):
TypeError: dynamic_func() missing 1 required positional argument: 'PACKED'
deepmodeling#6001 added the `PACKED` rotation-layout constexpr to `_flash_bwd_kernel`
and six Wigner-stride arguments to `_flash_bwd_block_kernel`, and updated
the production launcher `_launch_backward`, but the two raw launches in
`sweep_flash_bwd` kept the old argument lists: `launch_edge` never passed
`PACKED`, and `launch` still passed 20 of the 26 positional arguments.
Derive `packed` and the Wigner strides with `_rotation_strides`, exactly
as `_launch_backward` does, and pass them to both launches. Add a GPU
regression test that runs the real `flash_bwd` sweep on a tiny shape
through `tune_missing_configs`; the existing tuning test replaces every
sweep with a fake and could not catch the signature drift.
There was a problem hiding this comment.
🟢 Approval recommended
The fixes are covered by regression tests and no unresolved blocking issues were identified.
Pull request overview
Fixes Triton flash-attention backward tuning sweeps by passing required PACKED metadata and Wigner strides, with GPU regression coverage.
Changes:
- Corrects backward kernel launch arguments.
- Adds real sweep regression testing.
- Validates configuration registration.
File summaries
| File | Summary |
|---|---|
source/tests/pt/model/test_descriptor_sezm_triton.py |
Adds regression coverage for the flash-backward sweep. |
deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py |
Corrects backward tuning launch arguments. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughThe flash-attention backward sweep now derives Wigner strides and packed-layout metadata for both launch paths. A GPU-gated test runs the real sweep and validates its registered tile configurations. ChangesFlash backward sweep
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~12 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to The backward sweep metadata fix and regression coverage are merge-ready with no remaining actionable risk. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6015 +/- ##
==========================================
- Coverage 77.25% 76.99% -0.26%
==========================================
Files 1153 1153
Lines 138930 138933 +3
Branches 5056 5056
==========================================
- Hits 107328 106976 -352
- Misses 29717 30075 +358
+ Partials 1885 1882 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The sweep launch arguments match the current backward-kernel signatures and the production launcher. The added regression test exercises the real sweep and configuration registration.
The with-comm freeze failure mentioned in the follow-up notes is addressed in #6016.
Summary
_flash_bwd_kernel: newPACKEDconstexpr;_flash_bwd_block_kernel: 20 → 26 positional args) and updated the launcher_launch_backward, but not the two raw launches in the tuning sweepsweep_flash_bwd, which omitPACKEDand the six Wigner strides.dp --pt freezeatDP_TRITON_INFER >= 2(the perf(pt): DPA4-family performance optimizations #6001 default) when the GPU's built-in table lacks the family. perf(pt): DPA4-family performance optimizations #6001 filled the new families only for the RTX PRO 6000, so on other GPUs (e.g. H20) every default DPA4 freeze fails:TypeError: dynamic_func() missing 1 required positional argument: 'PACKED'.packedand the strides with_rotation_strides, as_launch_backwarddoes; no kernel, table or policy change. Test: a GPU regression test runs the realflash_bwdsweep (fails on master, passes with the fix).Validation (H20)
test_descriptor_sezm_triton.py: 37 passed.dp --pt freeze(examples/water/dpa4): all six tuning groups complete.test_sezm_export.py+test_dpa4_export.py: 45 passed; ruff and pre-commit clean.Follow-ups
nlocconstraint violation at level 0): separate PR.Summary by CodeRabbit
Bug Fixes
Tests