-
Notifications
You must be signed in to change notification settings - Fork 174
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
[Feat] Sage Attention Support for Triton kernel #929
base: develop
Are you sure you want to change the base?
Conversation
Thanks for your contribution! |
paddlemix/triton_ops/triton_ops.py
Outdated
PD_BUILD_OP(${op_name}) | ||
.Inputs({"x", "k_tensor", "v_tensor", "q_scale", "k_scale"}) | ||
.Outputs({"out_tensor", "lse_tensor"}) | ||
.Attrs({"output_dtype: std::string", "tensor_layout: std::string", "return_lse: int"}) |
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.
lse需要吗?如果确认推理不需要的话是否可以删掉呢?
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.
多卡并行推理可能需要打开,建议保留。
|
This PR was left clean, which is dedicated to sageattn triton kernel. Modifications on other files were removed, for the purpose of a clean PR and clean review. Annotations added, for params description and function usage. All code have passed test scripts. The test script will be shared on inflow, instead of in this PR. |
Sage attention Triton kernel Support
So far stagely support
is_casual=False
situation.is_casual=True
is implemented in another PRThe performance and comparison with FA2 was presented in inflow documents, and here we will save it.