perf(closed_loop): torch.compile the evaluation forward pass - #7
Merged
HansRobo merged 1 commit intoAug 12, 2026
Merged
Conversation
Compile the encoder and the DiT with the cudagraphs backend for the duration of a closed-loop evaluation. Full-scale evaluation runs about 2x faster. A compiled model cannot take nn.MultiheadAttention's fused fastpath, so it rounds differently in the last float32 bit. A closed loop feeds that back into its own next input, so closed-loop metrics shift once.
yhisaki
approved these changes
Aug 12, 2026
HansRobo
added a commit
that referenced
this pull request
Aug 27, 2026
Compile the encoder and the DiT with the cudagraphs backend for the duration of a closed-loop evaluation. Full-scale evaluation runs about 2x faster. A compiled model cannot take nn.MultiheadAttention's fused fastpath, so it rounds differently in the last float32 bit. A closed loop feeds that back into its own next input, so closed-loop metrics shift once.
HansRobo
added a commit
that referenced
this pull request
Aug 27, 2026
Compile the encoder and the DiT with the cudagraphs backend for the duration of a closed-loop evaluation. Full-scale evaluation runs about 2x faster. A compiled model cannot take nn.MultiheadAttention's fused fastpath, so it rounds differently in the last float32 bit. A closed loop feeds that back into its own next input, so closed-loop metrics shift once.
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
2 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.
Original PR: tier4#342
Compiles the encoder and the DiT for the duration of a closed-loop evaluation. Full-scale evaluation goes from 2,869 s to 1,443 s — 1.99x.
Measured over four sites and both object modes, 64,227 steps, identical in both arms, with eight evaluations running concurrently on one node with 28 cores and one GPU each — the configuration closed-loop validation runs in. Spread over four repeats was 0.6% before and 0.1% after. The baseline ran with the MHA fastpath already disabled; that switch on its own moved wall time within ±0.3%.
Closed-loop metrics shift by the last float32 bit — a compiled model cannot take
nn.MultiheadAttention's fused fastpath, and the loop amplifies the difference into its own next input.What changed
New
scenario_generation/inference_compile.py.compiled_for_inference()compiles and restores the model on exit — the training loop hands over its live model, and a leftover wrapper would put_orig_mod.prefixes in the next checkpoint.mark_inference_step()opens one cudagraph step per inference.For reviewers
The fastpath switch is entered inside
compiled_for_inference()rather than left to the caller: it is not in dynamo's guard set, so setting it after compilation leaves the traced graph in place while the flag reads back as changed.inductorwas tried and rejected — a second source of numeric difference on top, and no faster.