Add umas_flash execution mode: fused Triton backbone for UMA inference - #2185
Open
ArcelliEmanuele wants to merge 3 commits into
Open
ArcelliEmanuele wants to merge 3 commits into
ArcelliEmanuele wants to merge 3 commits into
Conversation
Adds a new execution mode, `umas_flash`, implementing the eSCN-MD backbone as a set of fused Triton kernels. Beyond the existing fast backends it avoids materializing the Wigner matrices, radial embeddings and rotated edge messages in HBM, so it cuts peak memory as well as runtime. The backend is inference-only and is validated against a fixed model shape (lmax=2, mmax=2, and the UMA-S channel widths); anything outside that is rejected up front with an explanatory error rather than being silently mis-executed. It requires merge_mole=True. Wigner rotations use a packed block-diagonal layout holding the 34 non-trivial entries of the l=1 and l=2 blocks, the l=0 scalar being identically 1. Integration points are additive: a `fused_backbone_features` dispatch hook in escn_md.py, the backend registration in execution_backends.py, and a doc comment for the new mode in the inference settings. Tests cover numerical agreement against the reference backend and the validation guards. Three tests require a second GPU and skip otherwise. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
wood-b
self-requested a review
September 2, 2026 19:36
mshuaibii
self-requested a review
September 8, 2026 15:04
mshuaibii
reviewed
Sep 11, 2026
| "merge_mole": True, | ||
| "activation_checkpointing": False, | ||
| "external_graph_gen": False, | ||
| "execution_mode": "umas_flash", |
Contributor
There was a problem hiding this comment.
Should we explicitly test compile:True here as well? if both are valid settings we should test both.
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
3 participants
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.
Adds a new execution mode,
umas_flash, implementing the eSCN-MD backbone as a set of fused Triton kernels. Beyond the existing fast backends it avoids materializing the Wigner matrices, radial embeddings and rotated edge messages in HBM, so it cuts peak memory as well as runtime.The backend is inference-only and is validated against a fixed model shape (lmax=2, mmax=2, and the UMA-S channel widths); anything outside that is rejected up front with an explanatory error rather than being silently mis-executed. It requires merge_mole=True.
Wigner rotations use a packed block-diagonal layout holding the 34 non-trivial entries of the l=1 and l=2 blocks, the l=0 scalar being identically 1.
Integration points are additive: a
fused_backbone_featuresdispatch hook in escn_md.py, the backend registration in execution_backends.py, and a doc comment for the new mode in the inference settings.Tests cover numerical agreement against the reference backend and the validation guards.