Repository navigation
perf(pt): fuse the mixing-stack per-layer backward chain (opt-in) - #6060
long-yi-2019 wants to merge 4 commits into
Conversation
…D_FUSED_INFER) Each gated layer's backward runs as recompute -> pointwise -> GEMM, materializing the sigmoid surface and the pre-activation gradient to HBM and reading both straight back. With the flag on (and the supported regime: fp32, narrow channels, no higher-order graph), one kernel walks the chain with both surfaces in registers, matching the relay bit-for-bit; the m = 0 and |m| = 1 column segments launch separately to keep the lighter segment off the heavy one's register profile. The 2.15 GB standing scratch never allocates; on the N=4096 eval workload it does not coincide with the frame's peak.
for more information, see https://pre-commit.ci
|
Important Review skippedWe couldn't safely recover the incremental review. No full review was started, and the last reviewed checkpoint was preserved. Retry later, or explicitly request a full review by commenting You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 20 functions across 3 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely identifies the main change: an opt-in performance optimization that fuses the PyTorch mixing-stack per-layer backward chain. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches: 💡 1
-
🧪 Generate unit tests (beta)
-
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
-
- Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Agent: dot
Reviewed the three-file delta from dbca0b1 to 4455e67, including the existing recompute/pointwise/GEMM implementations, launch configuration lookup, autograd routing, and new tests. One nonblocking regression-test issue is noted inline; source inspection did not establish a production kernel correctness defect.
Validation: eight isolated checks of the exact-source environment parser passed. Repository tests and CUDA/Triton numerical or compilation checks were not run: this executor has neither PyTorch nor Triton. The author-reported CUDA results were not independently reproduced. At review time all seven PR-triggered Actions runs reported action_required (the CUDA run has no jobs); pre-commit passed and Read the Docs reported failure. No CI runs were activated or rerun.
|
|
||
| def test_non_fp32_fall_back(self): | ||
| """The fused kernel assumes fp32 rounding; other dtypes keep the relay.""" | ||
| calls, out = self._record_launches(3, 32, dtype=torch.bfloat16) |
There was a problem hiding this comment.
Agent: dot
[P3] Enable fusion before asserting that unsupported inputs fall back
This test and test_wide_channels_fall_back never set DP_STACK_BWD_FUSED_INFER, and _record_launches does not set it either. In the normal default-off test run, fused short-circuits at use_stack_bwd_fused(), so these assertions pass without exercising the unsupported-dtype/channel dispatch contract. In particular, removing the fp32 guard would still leave this bf16 regression test green. Please scope both calls with mock.patch.dict(os.environ, {"DP_STACK_BWD_FUSED_INFER": "1"}) (or equivalent) so the tests actually request the feature they expect to reject.
|
@OutisLi @wanghan-iapcm @njzjz requesting a review, please. Following the 2026-10-13 discussion we validated the combined stack of our six performance PRs (#6056, #6057, #6060, #6070, #6071, #6072) with a multi-scale gate: diamond + bcc-Cu at ~500 / ~2000 / ~4400 atoms, three columns (pure torch.compile with all custom kernels off / current master / our changes on). All accuracy cells sit within 2.2x the pure baseline's own run-to-run band (criterion: 3x), and no tested atom count regresses in time or memory — full numbers in the PR descriptions. Suggested review order for the six: #6056 -> #6057 -> #6060 (small and self-contained) -> #6070 -> #6071 -> #6072. |
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
I reviewed the exact head. The existing test-contract finding at source/tests/pt/model/test_descriptor_sezm_triton.py:2272 remains valid: the fallback tests never enable DP_STACK_BWD_FUSED_INFER, so they can pass without exercising the unsupported-input guards. Please scope the tests with the flag enabled and rerun the failing documentation check.
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
I reviewed the exact head. The existing test-contract finding at source/tests/pt/model/test_descriptor_sezm_triton.py:2272 remains valid: the fallback tests never enable DP_STACK_BWD_FUSED_INFER, so they can pass without exercising the unsupported-input guards. Please scope the tests with the flag enabled and rerun the failing documentation check.
|
Thank you for the review. Addressed in ad9af59: both fallback tests now run with the flag enabled (renamed On the documentation check: the RTD failure is repo-wide (the same red check shows on other open PRs, e.g. #6058); this diff touches no docs inputs. |
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Reviewed ad9af590248897b720b7401e03269f09bbb68181.
This PR removes intermediate sigmoid/pre-activation-gradient surfaces from the fp32 mixing-stack backward by recomputing them inside the matrix-multiply kernels. The optimization has a real use case, but the public control, tuning integration, duplicated mathematics, and test cost need changes before merge.
Changes required
-
Use the existing inference levels; remove the new environment variable.
DP_STACK_BWD_FUSED_INFERis a fine-grained development switch and should not become a public interface. Do not replace it with another switch or introduce another level. This optimization belongs within the existingDP_TRITON_INFER >= 2value-path selection: level 2 already covers table-tuned kernels, while level 3 adds validated fp16x3 arithmetic. Keep supported-shape and profitability decisions internal. See the new selector and the existing level contract. -
Integrate the optimization into freeze-time autotuning, rather than shipping one fixed launch configuration.
Extend the existingsweep_tile_configs.pymachinery and its registration/lookup path so thatfreeze_pt2._tune_triton_configs()can discover and tune the new choice before tracing. A standalone benchmark that freeze never invokes is not sufficient. Compare numerically valid fused candidates against the current optimized relay, select the fastest valid candidate among those evaluated, and retain the existing path when fusion does not win. Record that outcome so repeated tuning does not redo completed work. Existing coverage for the old GEMM family must not incorrectly imply coverage for this new choice. See the fixed configuration, the sweep registry, and the freeze hook.Please provide results from running the added autotuner on your own available hardware: GPU/software versions, tested shapes, selected configurations or fallback, numerical agreement, and timings against the existing path. You are not required to populate the existing built-in configuration tables for hardware you do not have. The automatic tuning integration and evidence from hardware you can actually test are required.
The need for selection is measurable. On an RTX PRO 6000 Blackwell Server Edition with PyTorch 2.13.0+cu130 / Triton 3.7.1, the same synthetic fp32 backward workload (
E=300003, Cf=32, F=2, three gated layers plus the final linear layer, alpha enabled) gave:lmax Existing path Fused path Time change 3 11.58 ms 10.49 ms -9.4% 6 26.82 ms 45.39 ms +69.3% These are backward-operator microbenchmarks, not whole-model timings: identical inputs, CUDA-graph timing, off/on/on/off order, and the average of two median measurements per state. Outputs were bitwise equal. The degree-6 fused binaries also reached 255 registers with
n_spills=68–92; the fixed configuration does not universally avoid spills. -
Consolidate the duplicated derivative implementation.
The new kernel repeats the existing SiLU/gate derivative and scalar-gradient accumulation while requiring bitwise agreement. Shared mathematics and control flow should have one implementation wherever they can be unified; keep only genuinely different execution schedules separate. Small inlined Triton helpers are one option, subject to numerical and performance validation. Avoidable duplication is a maintainability issue to fix, not optional polish, and replacing it with a wrapper-heavy framework would not address that issue. See the duplicated chain. -
Reduce redundant test work and keep ordinary GPU CI inexpensive.
Remove the new environment-parser tests along with the switch, and extend existing level/configuration-selection tests instead. The dispatch-only cases need not all allocate the shared 20,000-edge fixture. Use small representative inputs for numerical checks where they still exercise the intended kernel and tile boundaries; reuse existing higher-order/compile coverage rather than adding another heavy suite. Keep production-scale benchmarks in the tuning/benchmark workflow, not ordinary GPU CI.Do not mechanically shrink the higher-order case: its current
2 * 20000 * 320 * 4-byte surface is just above the existing 48 MiB traversal threshold, so a smaller fixture changes the branch exercised. Preserve the intended coverage explicitly when reducing cost. See the fixture and the threshold.
Verification
Seven targeted tests passed, including 15 subtests and the SO2Convolution integration comparison. Additional off/on probes across 11 shape cases were bitwise equal; the first-order autograd and symbolic FX-to-Inductor checks also agreed bitwise, with fused dispatch confirmed. Ruff and diff whitespace checks passed. No numerical correctness defect was established within this coverage; the author's whole-model 1.7% result was not independently reproduced.
The earlier comment about fallback tests failing to enable the switch has been addressed in this head and is not being raised again. The requested changes above concern the architecture and acceptance criteria of the optimization, not that already-fixed test issue.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
|
||
| def _record_launches(self, lmax, focus_dim, apply_alpha=False, dtype=torch.float32): | ||
| """Run the traversal with a recording stand-in for the fused launch.""" | ||
| import deepmd.pt_expt.kernels.triton.sezm.so2_value_path as vp |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6060 +/- ##
==========================================
- Coverage 77.79% 77.49% -0.31%
==========================================
Files 1171 1171
Lines 141562 141663 +101
Branches 5066 5065 -1
==========================================
- Hits 110135 109785 -350
- Misses 29550 30004 +454
+ Partials 1877 1874 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Motivation
Each gated layer of the mixing-stack backward runs as a three-kernel relay: recompute the gate sigmoids, pointwise backward writing the pre-activation gradient
gz, then the transposed-weight GEMM readinggzback. At the SeZM N=4096 eval workload (E ~ 646k, Cf = 32, lmax = 3) the relay materializes two surfaces per stack and walks them through HBM every gated layer:sig(F, E, L*Cf)fp32, ~0.5 GB, written by the recompute kernel and read straight back by the pointwise kernel;gz(F, E, ROW), ~1.65 GB, written by the pointwise kernel and read straight back by the GEMM.Change
Opt-in via
DP_STACK_BWD_FUSED_INFER(default OFF, read when the backward traversal runs viause_stack_bwd_fused(), matching the repository's convention of construction/call-time reads for these switches). When enabled, and in the supported regime — fp32,Cf < GATE_BMM_MIN_FOCUS_DIM(no cuBLAS gate path),Cf == BLOCK_Kso one K tile is exactly one row group, no weight gradients, no second-order surfaces, no upstream cotangents, and no higher-order gradient graph being built — a single kernel per layer replaces the relay: each program recomputes its row group's gate sigmoid, forms thegztile in registers with exactly the pointwise kernel's operator sequence, and accumulates the GEMM with exactly the backward GEMM kernel'stl.dotorder, so results match the relay bit-for-bit.The
m = 0block (which carries the full scalar chain) and the|m| = 1stripes launch as two PART passes over disjoint column segments: a shared kernel would push the light segment onto the heavy one's register profile (BLOCK_M=64/BN=64 spills 226 slots; the shipped 32/64 config stays in budget).Outside the regime the dispatch falls back to the existing relay unchanged. The fused kernel registers no next-derivative formula;
create_graph=True(e.g. force-loss training) keeps the operator that carries the hand-derived second order — the official backward already routes those differentiations structurally, andnot torch.is_grad_enabled()at the dispatch is a second belt.Measurements
SeZM model, N=4096 atoms (~646k edges), RTX 5090, torch 2.12.1+cu130,
DP_COMPILE_INFER=1,DP_TRITON_INFER=3, throughput mode (16 evals, P50):Peak VRAM is unchanged at this workload: the eliminated 2.15 GB standing scratch does not coincide with the frame's peak allocation on this shape (the forward does); what it removes is the allocation itself and the per-layer HBM round-trip of both surfaces. Bit-identity of the fused kernel against the relay is pinned by
TestSeZMStackBwdFused(fp32,apply_alphaon and off); two-state compiled-eval accuracy at N=4096: energy rel diff 3.7e-10, max force diff 4.9e-6 with max |f| ~ 3.9.Tests
TestSeZMStackBwdFused(CUDA): fused-vs-relay bitwise cross-check; dispatch follows the env flag (recording stand-in); higher-order backward preserved with the flag on; wide-channel (Cf >= GATE_BMM_MIN_FOCUS_DIM) and non-fp32 fallbacks.TestSeZMStackBwdFlagParser(CPU): env parsing.source/tests/pt/model/test_descriptor_sezm_triton.py(43 passed) andtest_descriptor_sezm.py(61 passed) on the branch with the flag off (default).Summary by CodeRabbit
DP_STACK_BWD_FUSED_INFERsetting; it is disabled by default. Unsupported configurations and higher-order gradient calculations continue to use the existing backward path.