torch.compile compliant transformer decoder + MHA - speed goes vrrrrrrmmmm - #316
Merged
Conversation
…moment - crazy speed
There was a problem hiding this comment.
Pull request overview
This PR introduces a torch.compile-oriented path for transformer-decoder inference and refactors several recipe inference scripts to do a warm-up pass before the “real” run.
Changes:
- Add an
EOLE_TORCH_COMPILEenv-flag and route decoder forward through a@torch.compilewrapper when enabled. - Rework decoder/MHA cache handling to be more compilation-friendly (cache seqlens + scatter-based cache updates, flash-kvcache path).
- Refactor recipe inference scripts into
main()helpers and add a two-pass warm-up + run flow.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| recipes/pixtral/test_inference.py | Refactors into helper functions and adds warm-up + actual run passes. |
| recipes/mistral/test_inference.py | Same refactor + warm-up flow for mistral recipe script. |
| recipes/hunyuanocr/test_inference.py | Adds warm-up + actual run pass before printing results. |
| recipes/gemma3/test_inference.py | Refactors into helper functions and adds warm-up + actual run passes. |
| recipes/deepseekocr/test_inference.py | Refactors into helper functions and adds warm-up + actual run passes. |
| eole/predict/inference.py | Threads pos_ids_2d into decoder call and always initializes decoder cache at step 0. |
| eole/predict/generator.py | Disables decoder cache after decoding loop and clears CUDA cache. |
| eole/predict/init.py | Adds torch.compile-related runtime toggles and inductor/dynamo config when enabled. |
| eole/modules/rope.py | Refactors RoPE ops dispatch to a new eole.ops abstraction and updates RoPE update API. |
| eole/modules/rmsnorm.py | Refactors RMSNorm fused op dispatch to eole.ops and gates under torch-compile flag. |
| eole/modules/multi_headed_attn.py | Refactors flash-attn integration + cache update path intended to reduce graph breaks. |
| eole/encoders/vision.py | Updates RoPE update call signature for vision encoder. |
| eole/encoders/transformer.py | Changes encoder to reuse precomputed RoPE embeddings from rope.cos_sin. |
| eole/decoders/transformer.py | Adds compile/eager split, changes cache tracking to cache_seqlens, and updates flash eligibility logic. |
| eole/decoders/decoder.py | Renames cache initializer hook to _init_cache. |
| eole/constants.py | Refactors fused activation dispatch to eole.ops and gates under torch-compile flag. |
| eole/init.py | Introduces EOLE_TORCH_COMPILE env flag. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
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.
We now have a env-flag
EOLE_TORCH_COMPILEDefault is "0": no torch compile, uses the new cuda kernel, gives already x4 speed vs HF transformer "generate"
When set to "1", torch.compile will build the cudagraph for Transformer Decoder and subsequent modules: