Grouped Query Attention

[jainapurva]

Approach: Using the current function declaration

Constraint: Q_Heads % KV_Heads == 0

Major change:

• Added a new argument enable_gqa: bool to sdpa function call
• It adds a meaning to the last third dimension.

Sample use cases this would enable:
LLama3

# LLama3 8b call to SDPA
query = torch.rand(batch, 32, seq_len_q, D)
key = torch.rand(batch, 8, seq_len_kv, D)
value = torch.rand(batch, 8, seq_len_kv, D)

output = scaled_dot_product_attention(query, key, value, is_causal=True, enable_gqa=True)

# Output Shape
(batch, 32, seq_len_q, D)

Design Choice:

• Check if Query.size(-3) == Key.size(-3) == Value.size(-3) or, Query.size(-3) % Key.size(-3) == 0
• The function adjusts the key and value tensors to match the query tensor's head dimension by using repeat_interleave if their number of heads are not equal, facilitating correct and efficient computation in attention mechanisms.
• By default the enable_gqa flag is set to False, which ensures that regular sdpa functionality remains unchanged.

Benchmarks:

• sdpa.py: Gqa benchmark #130634
  For different batch sizes enable_gqa=True shows a substansial improvement in the run_time of sdpa

batch_size	q_num_heads	kv_num_heads	q_seq_len	kv_seq_len	embed_dim	forward_time when enable_gqa=True	forward_time when enable_gqa=False
1	32	8	2048	2048	2048	100.71	119.70
8	32	8	2048	2048	2048	539.78	628.83
16	32	8	2048	2048	2048	1056.81	1225.48
32	32	8	2048	2048	2048	2099.54	2440.45

• TorchTitan: Updates to support new sdpa function torchtitan#458

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128898

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 95e724e with merge base e6cddc9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jainapurva]

@pytorch-bot rebase

[jainapurva]

@pytorchbot rebase

[drisspg]

Overall looks really good, left a few comments

I think it would be helpful to run this script: and get some micro benchmark data: https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py

[jainapurva]

@pytorchbot merge

[pytorchmergebot]

This PR updates submodules third_party/fmt

If those updates are intentional, please add "submodule" keyword to PR title/description.

[jainapurva]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here