Remove the hack for SAC + FlexAttention #2118
Conversation
[ghstack-poisoned]
meta-cla
bot
added
the
CLA Signed
tianyu-l
left a comment
tianyu-l
left a comment
There was a problem hiding this comment.
LGTM. I wonder if you have verified it "works", i.e. it doesn't invalidate SAC / compile anymore?
| [activation_checkpoint] | ||
| mode = "selective" # ["none", "selective", "full"] | ||
| selective_ac_option = '2' # 'int' = ac every positive int layer or 'op', ac based on ops policy | ||
| selective_ac_option = 'op' # 'int' = ac every positive int layer or 'op', ac based on ops policy |
There was a problem hiding this comment.
Is this intended? I think it's OK to change this but want to confirm.
There was a problem hiding this comment.
oops, accidentally commit this.
|
Yes, I verified it works with llama3 and llama4. FlexAttention is compiled within SAC. |
| options={ | ||
| "wrap_inductor_compiled_regions": True, | ||
| "max_autotune": True, | ||
| "coordinate_descent_tuning": True, |
There was a problem hiding this comment.
Noob question: is this coordinate_descent_tuning also part of the "mode=max-autotune-no-cudagraphs" -> "options={...}" change?
There was a problem hiding this comment.
forgot to ask: what's the context of this change to "options={}"?
There was a problem hiding this comment.
We cannot do mode and options at the same. torch.compile forbid it. According to here, max-autotune-no-cudagraphs equals to these three options.
[ghstack-poisoned]
Stack from ghstack (oldest at bottom):
PyTorch can now support torch.compile inside the SAC region even if torch.compile is not used to wrap SAC. This PR removes the workaround to ensure torch.compile works with Flex