[Feature Request][torch.compile] Improve UX when autotuning repeated layer
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🚀 The feature, motivation and pitch
Currently, there is a huge difference when we compile at the Decoder level or the DecoderLayer level, and it makes sense.
In my scenario, before any inductor cache effect, it is something like 60 seconds vs 7 seconds for 24 layers.
The reason is, when compiling at layer level it compiles once and reuse.
we can see for triton tuning:
```
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1123] [0/0] Overview info of inductor aten mms:
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1124] [0/0] Name | B | M | N | K | Count
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1129] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 4096 | 1024 | 1
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 2048 | 1
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 7168 | 1024 | 1
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:39:39.802000 455772 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 3584 | 1
I0128 14:39:39.803000 455772 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
```
When we compile at Decoder level, it will perform the same autotuning 24 times.
```
I0128 14:53:23.017000 458948 site-packages/torch/_inductor/compile_fx.py:1123] [0/0] Overview info of inductor aten mms:
I0128 14:53:23.017000 458948 site-packages/torch/_inductor/compile_fx.py:1124] [0/0] Name | B | M | N | K | Count
I0128 14:53:23.017000 458948 site-packages/torch/_inductor/compile_fx.py:1129] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 4096 | 1024 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 2048 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 7168 | 1024 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1132] [0/0] aten.mm | - | 4 | 1024 | 3584 | 24
I0128 14:53:23.018000 458948 site-packages/torch/_inductor/compile_fx.py:1133] [0/0] ----------------------------------------------------------------------------------------------------------------------------------
```
I think we still want to keep the best performance of compiling at the Decoder level, but I am sure the pipeline can detect that we re-use the same matmul and perform some optimization only once.
### Alternatives
_No response_
### Additional context
_No response_
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @Lucaskabela
Contributor guide
Assessment
This issue has not been assessed yet.