Code

[Pallas] Support Flash Attention#6658

Code

Merged

alanwaketan merged 9 commits intomasterpytorch/xla:masterfrom alanwaketan/flash_attentionpytorch/xla:alanwaketan/flash_attentionCopy head branch name to clipboardMar 6, 2024

+41-3

Commits

Commits on Mar 6, 2024

• tmp

  alanwaketan

  committed
  406224a

• Revert some changes

  alanwaketan

  committed
  6e8f05d

• Revert more changes

  alanwaketan

  committed
  ebcfca9

• Turn off TPU layout by default

  alanwaketan

  committed
  d24a222

• Improve the test case

  alanwaketan

  committed
  cb0a1bb

• Turn tpu layout back on

  alanwaketan

  committed
  f424e2e

• Reenable torchvision

  alanwaketan

  committed
  1031d1d

• Fix linters

  alanwaketan

  committed
  c805109

• Fix more liners

  alanwaketan

  committed
  e60f07b