Add attention_backend to let user choose

[yanbing-j]

This PR is submitted for #1452, to add an argument attention_backend to let user choose the backend used in SDPA.

For CPU, only math and flash_attention are valid. For CUDA, all 4 backends are valid.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 2 Unrelated Failures

As of commit a9dfc0e with merge base 2fc98f7 ():

NEW FAILURE - The following job has failed:

• pull / test-build-runner-et-android / linux-job (gh)
  RuntimeError: Command docker exec -t c280bdd05e52a41ccfca921f81ff91498a56c0f0acec2b0ffc66b73308d8adb3 /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-tinystories-executorch (16-core-ubuntu) (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-tinystories-executorch (macos-14-xlarge) (gh) (trunk failure)
  RuntimeError: Failed to load ET compiled /Users/runner/work/torchchat/torchchat/stories15M.pte

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

[yanbing-j]

cc @mingfeima @lucylq @angelayi @Jack-Khuu

[lucylq]

Thanks for making these changes! This looks good to me.

cc @Jack-Khuu

[Jack-Khuu]

Looks like the ET pin needs bumping so I'll spin that up

[Jack-Khuu]

Can we bump this into the constructor of BuilderArgs?

[Jack-Khuu]

Thanks for adding the change!! Mind bumping the casting to the actual kernel into the body here?

That way we can do any sanity checks early and aren't passing around a string

[yanbing-j]

Have updated. Please review again!

[Jack-Khuu]

@angelayi I recall you mentioning that there might be export issue with the other SDP backend?

Do you know if this issue was addressed? Or should be gate the backends to non-export?

[Jack-Khuu]

Seems like yall had a chance to explore the export issue in the tagged issue

Let's add a check in export.py that flags when someone tries to use not Math.SDPA, to avoid users accidentally stumbling into it

[Jack-Khuu]

Can we push this into the builder.py as well?

[yanbing-j]

Done.

[mikekgfb]

@angelayi I recall you mentioning that there might be export issue with the other SDP backend?

Do you know if this issue was addressed? Or should be gate the backends to non-export?

FLASH seemed to work for AOTI. Maybe we can at least allow a subset, and then methodically put them all thru a new test patter -- which we should anyway, to test the new options for generate as well. (suggest as separate PR, not piggyback here...)

for backend in math flash cudnn efficient
do 
  # cuda
  py3 tc generate stories15M --attention-backend $backend --prompt "hello there"
  py3 tc generate stories15M --attention-backend $backend --prompt "hello there" --compile --compile-prefill
  py3 tc export stories15M --attention-backend $backend --output-dso s.dso
  py3 tc generate stories15M --attention-backend $backend --prompt "hello there" -dso s.dso
  # cpu
  py3 tc generate stories15M --attention-backend $backend --prompt "hello there" -device cpu
  py3 tc generate stories15M --attention-backend $backend --prompt "hello there" --compile --compile-prefill -device cpu
  py3 tc export stories15M --attention-backend $backend --output-dso s.dso -device cpu
  py3 tc generate stories15M --attention-backend $backend --prompt "hello there" -dso s.dso -device cpu
done

[Jack-Khuu]

Thanks again @yanbing-j, merging this in

We can work on creating a support matrix separately