axlearn icon indicating copy to clipboard operation
axlearn copied to clipboard
verified publisher icon apple

Sliding window local attention

Open changlan opened this issue 1 year ago • 0 comments

  • Added sliding window local attention feature to MultiheadAttention and FlashAttention.
  • Added or_masks and and_masks which would be useful to composite different mask functions.
  • Fixed a bug that causes TPU decoding always to fail.

changlan avatar Oct 09 '24 05:10 changlan