Fix a number of flexattention issues (cse, cudagraph, etc.) #145059
Conversation
[ghstack-poisoned]
Chillee
requested review from
albanD,
jbschlosser and
mikaylagawarecki
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/145059
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 1b78c80 with merge base 8d91bfd ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
| // do nothing | ||
| } else if (indices_batched && (self_bdim.has_value() || values_bdim.has_value())) { | ||
| auto arange_index = at::arange(0, batch_size); | ||
| auto arange_index = at::arange(batch_size, options.dtype(kLong)); |
There was a problem hiding this comment.
This is the cause for why bm creation wasn't cudagraphable (#143872), and also probably why fixing compilation of torch.compile made such a huge difference haha.
Basically, aten.index actually works fine if the indices are on different devices - it just does an implicit DtoH copy. But obviously it breaks cudagraphs
torch/_inductor/codegen/triton.py
Outdated
| not config.triton.codegen_upcast_to_fp32 | ||
| and isinstance(var, CSEVariable) | ||
| and var.dtype in (torch.float16, torch.bfloat16) | ||
| return isinstance(var, CSEVariable) and var.dtype in ( |
There was a problem hiding this comment.
did you remove the config?
There was a problem hiding this comment.
I don't think it makes sense to gate this upcast on this config. This needs_upcast is needed for correctness (in the case that we don't upcast everything immediately to fp32), but the original PR added this presumably to limit the blast radius of the potential change.
There was a problem hiding this comment.
This is fine. Previously it was gated because there was no situation we expected to see bf16 vars but there's no reason we need the gate. In the general case we should be upcasting before libdevice ops which dont support bf16.
There was a problem hiding this comment.
I actually ended up reverting this change because it is actually also used for downcasting.
There was a problem hiding this comment.
I don't follow , would you say more ? For ops which dont support low precision, it should upcast, run op, then downcast after.
| for arg in args: | ||
| if isinstance(arg, str): | ||
| # TODO: fix the flex attention instances, enable internally | ||
| if not config.is_fbcode(): |
There was a problem hiding this comment.
are there other fbcode issues this will pop now?
There was a problem hiding this comment.
Not sure, we'll see!
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov