Repository navigation
perf(pt): assemble Wigner-D behind a custom op; inverse rotation as a view - #6070
Open
long-yi-2019 wants to merge 2 commits into
Open
long-yi-2019 wants to merge 2 commits into
long-yi-2019 wants to merge 2 commits into
Conversation
added 2 commits
October 10, 2026 15:44
Functionalization expands the in-place block scatter into a chain of full-size zeros/index_put adds plus a realized transpose (~4 GB of HBM traffic per frame at the SeZM N=4096 workload). A wigner_d_assemble custom op keeps the eager scatter and hands the graph two contiguous tensors; the traced assembly takes it whenever a fake-tensor mode is active, so the eager path is untouched.
The flash and rotation-back kernels already address the rotation through explicit strides, so the contiguous transpose the assembly produced was pure storage: one 660 MB copy per frame that existed only to be read back. The assemble op now returns D_full alone and the caller hands out a transpose view; consumers read it in place and no transpose copy ever exists.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to No actionable merge-blocking defect was established. The untested PyTorch 2.11 compiled path remains a validation consideration, not a demonstrated failure. Pre-merge checks |
  |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 1 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
✅ Passed checks (4 passed)
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the two main changes: custom-op Wigner-D assembly and inverse-rotation handling through a view. 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.
- Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
-
🧪 Generate unit tests (beta)
-
- Create a new PR
🛠️ Fix failing CI checks 💡
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
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 @coderabbitai help to get the list of available commands.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Each eval frame assembles the block-diagonal Wigner-D rotation matrices by
scattering the per-degree blocks into a dense
(E, D, D)tensor. Tracedthrough Inductor, those in-place slice assignments functionalize into a
chain of full-size
zeros + index_put + addtensors plus a realizedtranspose -- five
(E, D, D)surfaces combined by one kernel, ~4 GB of HBMtraffic per frame at the SeZM N=4096 workload (E ~ 646k, D = 16). On top
of that, the inverse rotation
Dt_fullwas produced as a full contiguouscopy even though every consumer can address it through strides.
Change
Two commits, one battle:
Assembly behind a custom op.
sezm_nn::wigner_d_assembletakes theper-degree blocks (plus the polynomial tail when
lmax >= 11) and doesthe eager scatter: one zeros buffer, in-place block writes. The traced
forward takes the op whenever a fake-tensor mode is active (
is_compiling()is not enough -- this module is traced through make_fx, where the flag
stays False; a fake tensor carries its mode, so the forward probes for
it). The eager path is untouched. An exact autograd formula returns
each block its D-slice.
The inverse rotation is a view. The flash and rotation-back kernels
already address the rotation through explicit stride arguments
(
_rotation_strides/w_se, w_sr, w_sk), verified bit-for-bit againstthe contiguous copy with non-symmetric matrices. The op now returns
D_fullalone and the caller hands outD_full.transpose(-1, -2); thetranspose copy never exists. The one consumer that requires contiguous
storage (
so2_rotate_mix) readsD_fullitself, which stays contiguous.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):The step-1 saving is time (the ~4 GB/frame add-chain disappears; the
compiler-fragment family drops 38.3 -> 25.8 ms of kernel time). The step-2
saving is VRAM: deleting the 660 MB transpose copy costs nothing overall
because the stride-read overhead and the removed copy cancel out. Two-state
compiled-eval accuracy at N=4096: energy rel diff 5.8e-10, max force diff
4.8e-6 with max |f| ~ 3.9.
Tests
source/tests/pt/model/test_descriptor_sezm.py: 61 passed (default config).source/tests/pt/model/test_descriptor_sezm_triton.py: 37 passed, 1 skipped (pre-existing).source/tests/pt/model/test_sezm_model.py::TestSeZMModelCompile: 9 passed on this branch head (RTX 5090, torch 2.12.1), including the eval cache/shape-change, first-frame and force-loss double-backward tests.(view vs materialized inverse) matches exactly.
Summary by CodeRabbit