-
Notifications
You must be signed in to change notification settings - Fork 21.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Inductor][Flex-attention] Support different sequence lengths for Query and Key/Value #126639
Conversation
test/inductor/test_flex_attention.py
Outdated
self._check_equal(golden_out, ref_out, compiled_out, fudge_factor, "Out") | ||
|
||
# Check gradients | ||
q_fudge_factor = 2.5 * fudge_factor |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we refactor this a bit lol. I think we have this code copied 3 different times in this file.
@@ -188,34 +188,36 @@ def build_subgraph_buffer( | |||
|
|||
Z = {{size("Q", 0)}} | |||
H = {{size("Q", 1)}} | |||
N_CTX = {{size("Q", 2)}} | |||
M = {{size("Q", 2)}} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's call this Q_LEN
and KV_LEN
?
closing this in favor of #127678 |
Fixes #ISSUE_NUMBER
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang