tutorials icon indicating copy to clipboard operation
tutorials copied to clipboard

Add FlexAttention examples to SDPA tutorial

Open dorbittonn opened this issue 1 month ago • 1 comments

Summary

  • Add new section demonstrating flex_attention from PyTorch 2.5
  • Include score_mod examples: relative position bias and ALiBi
  • Show block_mask for sparse attention patterns (causal masking)
  • Demonstrate combining score_mod and block_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.compile works 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

dorbittonn avatar Nov 29 '25 06:11 dorbittonn