Batching rule for `aten::_scaled_dot_product_efficient_attention`

[lengstrom]

🚀 The feature, motivation and pitch

Hi, I am trying to take batched gradients of a vector output given by _scaled_dot_product_efficient_attention but saw the error

/site-packages/optimum/bettertransformer/models/attention.py:56: UserWarning: There is a performance drop because we have not yet implemented the batching rul
e for aten::_scaled_dot_product_efficient_attention. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at ../aten/src/ATen/functorch/BatchedFallback.cpp:82.)
  sdpa_result = torch.nn.functional.scaled_dot_product_attention

when running my code. Implementing this would really increase throughput in our application! I would also be happy to take a stab at implementing it, if there is a document describing what I need to do at a high level.

Alternatives

No response

Additional context

No response

cc @ezyang @gchanan @zou3519 @kadeng @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @mikaylagawarecki @Chillee @samdow @kshitij12345 @janeyx99

[cpuhrsch]

cc @drisspg

[drisspg]

@zou3519 do you know if we have good documentation for adding a batching rule?

[lezcano]

In fact, we do!
https://github.com/pytorch/pytorch/blob/main/functorch/writing_batching_rules.md

[drisspg]

Thanks @lezcano!
@lengstrom if you wanted to take a crack at it I would be happy to review the PR. The cliff note is I think we will want to make whatever input be mapped to (batch, num_head) seq_len, head_dim. Adding size 1 for num_head where appropriate

[lengstrom]

great thank you so much! I will take a crack at it this weekend..

[lengstrom]

Hi, I started by trying to batch one of the downstream functions that is necessary: _chunk_grad_outputs_efficient_attention. As a point of start, I added the following function definition in aten/src/ATen/functorch/BatchRulesFactory.cpp:

bool _chunk_grad_outputs_efficient_attention_batch_rule(
  const Tensor& query, optional<int64_t> query_bdim,
  const Tensor& key,  optional<int64_t> key_bdim,
  const Tensor& value, optional<int64_t> value_bdim,
  bool is_causal
) {
  assert(false);
  return true;
}

and modified the TORCH_LIBRARY_IMPL block to include the statement

  VMAP_SUPPORT(_chunk_grad_outputs_efficient_attention, _chunk_grad_outputs_efficient_attention_batch_rule)

at the beginning; I was both trying to match the code I saw (in particular for the other boolean output valued function _has_same_storage_numel) and the guide posted helpfully above.

When I compile, I get errors about "generated_plumbing", e.g.:

/mnt/xfs/home/engstrom/src/pytorch/aten/src/ATen/functorch/BatchRulesFactory.cpp:122:16: error: ‘_chunk_grad_outputs_efficient_attention_generated_plumbing’ was not declared in this scope; did you mean ‘_chunk_grad_outputs_efficient_attention_batch_rule’?

Do you know what I am doing wrong here?

Thank you in advance!

[lengstrom]

(full output here: https://gist.github.com/lengstrom/c99cd6decb002f511a9efc51c881bc12)

[drisspg]

This function actually is not used anywhere, There is a comment in native_functions.yaml

# THIS FUNCTION iS DEPRECATED AND SHOULD BE REMOVED
- func: _chunk_grad_outputs_efficient_attention(Tensor query, Tensor key, Tensor value, bool is_causal=False) -> bool
  dispatch:
    CUDA: _chunk_grad_outputs_efficient_attention
    

I honestly wouldnt' even try debug this

[drisspg]

Another good, easy PR would be to just remove all traces of it if you are trying to explore codebase a little

[lengstrom]

Hi, I added a PR here: #103159 (comment)

obviously it is not ready for the prime time -- but is this the right idea? I tested locally a bit and it seems to mimic the "auto-batched" behavior.

[proger]

Thanks for working on this! Would love to use fast attention with vmap.