[FlexAttention] Fix max-autotune bug with captured buffer grads #141531
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/141531
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 fb56ad9 with merge base f472b3a ( UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
drisspg
added
topic: not user facing
|
The problem here is that we are passing in the node an input That in the non-autotune case ends up being the last buffer argument while generating the args in max-autotune it ends up being the second to last because we have an explicit output node. I need to figure out how to reorder the output nodes from the kernel. Confirmed by doing a hacky swap in in autotune process: input_tensors = list(input_tensors)
tmp = input_tensors[-1]
input_tensors[-1] = output_tensor
output_tensor = tmp |
|
Yeah, we have this test for default compile, but can add for autotuning |
|
@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 |
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
|
@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 |
…rch#141531) # Summary Fix tensor argument ordering for autotuning flex attention, change how we enabled scatters codegen for triton. We used to go through the existing store_output triton codegen but now we just short circuit and generate the correct expression earlier on. This enables us to instead of relying on arg.python_defs to thread arguments through via input_buffers we can instead reuse the exact same mutated buffer infra as we did for multiple outputs before. Test cases added for both default and max-autotune-no-cudagraphs modes. Pull Request resolved: pytorch#141531 Approved by: https://github.com/Chillee
…rch#141531) # Summary Fix tensor argument ordering for autotuning flex attention, change how we enabled scatters codegen for triton. We used to go through the existing store_output triton codegen but now we just short circuit and generate the correct expression earlier on. This enables us to instead of relying on arg.python_defs to thread arguments through via input_buffers we can instead reuse the exact same mutated buffer infra as we did for multiple outputs before. Test cases added for both default and max-autotune-no-cudagraphs modes. Pull Request resolved: pytorch#141531 Approved by: https://github.com/Chillee
3 participants
Stack from ghstack (oldest at bottom):
Summary
Fix tensor argument ordering for autotuning flex attention, change how we enabled scatters codegen for triton. We used to go through the existing store_output triton codegen but now we just short circuit and generate the correct expression earlier on.
This enables us to instead of relying on arg.python_defs to thread arguments through via input_buffers we can instead reuse the exact same mutated buffer infra as we did for multiple outputs before.
Test cases added for both default and max-autotune-no-cudagraphs modes.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @Chillee @yanboliang @BoyuanFeng