tutorials
tutorials copied to clipboard
Add FlexAttention examples to SDPA tutorial
Summary
- Add new section demonstrating
flex_attentionfrom PyTorch 2.5 - Include
score_modexamples: relative position bias and ALiBi - Show
block_maskfor sparse attention patterns (causal masking) - Demonstrate combining
score_modandblock_mask - Add performance comparison with standard SDPA
Test plan
- [x] All FlexAttention code tested with PyTorch 2.7.1
- [x] Verified output shapes are correct
- [x] Compilation with
torch.compileworks as expected
Checklist
- [x] The issue that is being fixed is referred in the description
- [x] Only one issue is addressed in this pull request
- [x] Labels from the issue that this PR is fixing are added to this pull request
- [x] No unnecessary issues are included into this pull request