@LuFinch , why do we need to add comput_log_sumexp? It breaks the ABI backward compatibility.

Ideally, we can check input tensors' attr like compute_logsumexp = query.requires_grad() || key.requires_grad() || value.requires_grad() to decide whether compute logsumexp. And this checking works in eager mode.

However, in torch.compile mode, the input tensors query/key/value require_grad()==True at the beginning would become require_grad()==False in the op after aot_autograd in some models. Hence, it needs a bool flag to indicate this op should compute logsumexp. I am not an expert of aot_autograd and not sure why it acts like this. But cudnn and efficient attention also has this parameter. I guess they meet the same issue, otherwise they should be able to move this check into op.

@LuFinch , why do we need to add comput_log_sumexp? It breaks the ABI backward compatibility.

By default comput_log_sumexp=False, this should not break API-level BC, right?

Eikan recommends me to move this argument as the last argument, then it will not break BC.

Cool!

Already moved this argument as the last argument.

According to the implementation details, it returns True when

• Grad mode is not enabled
• All input tensors do not require a gradient
• Not Group Query Attention and the attention mask do not require a gradient

@LuFinch , is my understanding correct? If so, I would suggest refining the name of check_grad a little bit. Something could be like is_onednn_attention_backward_supported to illustrate your idea clearly.

This function should be used to check the grad requirements of inputs to determine whether they are suitable for supporting overrideable SDPA on XPU in the future.

As guangye saying, it is use to determine whether use overrideable SDPA. If return True, then it can use overrideable SDPA.

In OneDNN v3.9, SDPA training forward and backward don't support GQA and won't output grad for attn_mask.

Hence this function means:

• If Grad mode is not enabled, we can use overrideable SDPA to run OneDNN SDPA inference forward graph.
• Grad mode is enabled but none of q/k/v needs grad, we can use overrideable SDPA to run OneDNN SDPA inference forward graph.
• If we need to compute grad, it is not GQA and attn_mask don't require gard, then we can use overrideable SDPA to run OneDNN SDPA training forward graph.
• Otherwise, it should fallback to MATH backend.

Since it has been gqa, why does this function return false?

In OneDNN v3.9, SDPA training forward and backward don't support GQA and won't compute grad for attn_mask.

ditto

In OneDNN v3.9, SDPA training forward and backward don't support GQA and won't compute grad for attn_mask.

gpu_float_sdpa_backward has been defined? Does it mean the backward function only supports float?

I have the same question.

It supports FP32/FP16/BF16. I directly copy this function name from SDPA inference. We could rename it.

renamed as sdpa_backward

What's the meaning of (B)?

Copy from forward code. It just means batchsize. Already removed the bracket.

This variable shadows a similar declaration at line 972. It's fine but not good. I recommend renaming this to avoid name shadowing.

Thanks for your patient review. Renamed it.

at::numeric_limits<>::lower_bound is better.

This file already uses -std::numeric_limits<float>::infinity, I replaced -INFINITY with -std::numeric_limits<float>::infinity, too.

Use a macro to reduce the duplicated code. such as

#define ADD_INPUT(variable) \
inputs.emplace_back(l_inputs[i++], eng, variable.data_ptr())
ADD_INPUT(grad_out);
ADD_INPUT(query);
...
#undef ADD_INPUT

Cool. Done.

Done.

This is fine. But I recommend using a name like kBitFloat32 instead of hardcoded number 3. Another one is that kBitFloat32 could be shared with find_or_create_graph_partition

Good idea. Done.

With this code change, sdpa will not support broadcast anymore. Is this a BC breaking, any impact for the old script.

After offline discussion, I add back broadcast for QKV. However, the output tensor attention and logsumexp are allocated by us and they should not be broadcasted. Hence I don't add broadcast for these two tensors.