Support INT8 SDPA template for CPU #2148
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2148
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 71e17d0 with merge base 137b079 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
Valentine233
added
the
topic: not user facing
Could you further elaborate the benefit in the summary for why the template based solution needed in this PR? |
Thanks, done. |
| @@ -56,7 +56,7 @@ | |||
| tags=[torch._C.Tag.needs_fixed_stride_order], | |||
| ) | |||
| lib.define( | |||
| "scaled_dot_product_int8(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, float scale=0.0, float q_scale=1.0, int q_zp=0, float k_scale=1.0, int k_zp=0, float v_scale=1.0, int v_zp=0, float a_scale=1.0, int a_zp=0, float o_scale=1.0, int o_zp=0) -> Tensor" | |||
| "scaled_dot_product_int8(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, float? scale=None, float q_scale=1.0, int q_zp=0, float k_scale=1.0, int k_zp=0, float v_scale=1.0, int v_zp=0, float a_scale=1.0, int a_zp=0, float o_scale=1.0, int o_zp=0) -> Tensor" | |||
There was a problem hiding this comment.
why need to change the OP schema? It seems not related to this PR.
There was a problem hiding this comment.
Yes, it's an additional modification.
This is to keep consistent with the SDPA schema in pytorch func: scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> Tensor.
| from torchao.utils import TORCH_VERSION_AT_LEAST_2_7 | ||
|
|
||
| if TORCH_VERSION_AT_LEAST_2_7: | ||
| # TORCH_VERSION_AT_LEAST_2_7 is needed for int8 sdpa lowering |
There was a problem hiding this comment.
Add more notes for why at least 2.7 is needed?
There was a problem hiding this comment.
Some functions used in INT8 SDPA lowering are only available for higher version PyTorch, like get_fill_order.
| o_zp=o_zp, | ||
| ) | ||
|
|
||
| if len(choices) == 0: |
There was a problem hiding this comment.
Just to confirm, does that mean prioritizing using a template-based solution when available, rather than relying on tuning?
There was a problem hiding this comment.
Thanks, you are absolutely right. The two algorithms are almost the same now. In the future, the template one is expected to have a better perf.
A follow-up of #1372, with RFC pytorch/pytorch#144941.
Here, we support the template-based method for INT8 SDPA, for potential better performance. With the template, we can generate the corresponding template during compiling and implement the kernel with high flexibility. By taking advantage of this, it would be more flexible to tune the optimized kernel with different parallel strategies or block sizes through benchmarking in the future.
This PR creates a CPP Int8 SDPA template
CppInt8SdpaTemplateby inheritingCppFlexAttentionTemplate, and implement a template-based kernel. For max-autotune, we have two choices for Int8 SDPA: the template one by default, and the explicit OPtorch.ops.torchao.scaled_dot_product_int8as fallback. As the two algorithms are almost the same, there is no need to do the actual tuning and choose one. We assume the template one has better perf, and we will keep optimizing it in the future.