-
Notifications
You must be signed in to change notification settings - Fork 32
Add TLX attention (WS pipelined pingpong hopper) #320
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
base: main
Are you sure you want to change the base?
Conversation
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.
See my comment but otherwise this looks good. Thanks!
@@ -299,6 +311,16 @@ def triton_tutorial_flash_v2_tma( | |||
q, k, v, self.causal, self.sm_scale, "tma" | |||
) | |||
|
|||
@register_benchmark(enabled=HAS_TLX) |
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.
Is the intention to run on CI? If you just want to trigger this test directly you can do with enabled=False
and without any risk of it running on hardware you don't expect (e.g. MI300), which it will by default right now.
|
||
import triton | ||
import triton.language as tl | ||
import triton.language.extra.tlx as tlx |
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.
I'm assuming we have tested this doesn't crash without tlx because of the guards in the operator.py file.
No description provided.