Add attention backend tests to more-tests.yml

[mikekgfb]

Tests from #1477 only, without the generate.py refactor
(Separate the code changes from the test)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

As of commit b52d1b4 with merge base 5684175 ():

NEW FAILURE - The following job has failed:

• Run parallel prefill / test-sdpa-backends / linux-job (gh)
  RuntimeError: Command docker exec -t 945a03c74265b068c53f069f9ee9a51d684d07e943337dbbe3bd5db391b225dc /exec failed with exit code 1

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

[mikekgfb]

Repro: python torchchat.py generate --checkpoint-path checkpoints/stories15M/stories15M.pt --attention-backend flash_attention --device cuda --dtype bfloat16 --temperature 0 fails because Flash Attention does not support non-null attn_mask.

This is with the current setup, i.e., no code change. The background is that the various accelerated SDPA implementations do not support all parameter combinations and dimensions, so restricting to a subset that does not include MATH can always lead to a fail when the particular input is not supported. (The only implementation guaranteed to handle all inputs is MATH)

2025-01-27T22:19:36.9882863Z ******* cuda bfloat16 flash_attention 
2025-01-27T22:19:36.9883664Z + python torchchat.py generate --checkpoint-path checkpoints/stories15M/stories15M.pt --attention-backend flash_attention --device cuda --dtype bfloat16 --temperature 0
2025-01-27T22:19:36.9884464Z NumExpr defaulting to 16 threads.
2025-01-27T22:19:36.9884780Z PyTorch version 2.7.0.dev20250124+cu124 available.
2025-01-27T22:19:36.9885286Z Warning: PTEModel (ExecuTorch) not available with exception: No module named 'executorch'
2025-01-27T22:19:36.9886208Z Unable to import torchao experimental quant_api with error:  [Errno 2] No such file or directory: '/pytorch/torchchat/torchao-build/src/ao/torchao/experimental/quant_api.py'
2025-01-27T22:19:36.9886946Z Using device=cuda NVIDIA A10G
2025-01-27T22:19:36.9887209Z Loading model...
2025-01-27T22:19:36.9887444Z Time to load model: 0.20 seconds
2025-01-27T22:19:36.9887746Z -----------------------------------------------------------
2025-01-27T22:19:36.9888968Z /pytorch/torchchat/torchchat/model.py:897: UserWarning: Memory efficient kernel not used because: (Triggered internally at /pytorch/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:784.)
2025-01-27T22:19:36.9889841Z   y = F.scaled_dot_product_attention(
2025-01-27T22:19:36.9890719Z /pytorch/torchchat/torchchat/model.py:897: UserWarning: Memory Efficient attention has been runtime disabled. (Triggered internally at /pytorch/aten/src/ATen/native/transformers/sdp_utils_cpp.h:548.)
2025-01-27T22:19:36.9891595Z   y = F.scaled_dot_product_attention(
2025-01-27T22:19:36.9892439Z /pytorch/torchchat/torchchat/model.py:897: UserWarning: Flash attention kernel not used because: (Triggered internally at /pytorch/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:786.)
2025-01-27T22:19:36.9893288Z   y = F.scaled_dot_product_attention(
2025-01-27T22:19:36.9894146Z /pytorch/torchchat/torchchat/model.py:897: UserWarning: Flash Attention does not support non-null attn_mask. (Triggered internally at /pytorch/aten/src/ATen/native/transformers/sdp_utils_cpp.h:261.)
2025-01-27T22:19:36.9895161Z   y = F.scaled_dot_product_attention(
2025-01-27T22:19:36.9896011Z /pytorch/torchchat/torchchat/model.py:897: UserWarning: CuDNN attention kernel not used because: (Triggered internally at /pytorch/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:788.)
2025-01-27T22:19:36.9896864Z   y = F.scaled_dot_product_attention(
2025-01-27T22:19:36.9897711Z /pytorch/torchchat/torchchat/model.py:897: UserWarning: CuDNN attention has been runtime disabled. (Triggered internally at /pytorch/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:535.)
2025-01-27T22:19:36.9898580Z   y = F.scaled_dot_product_attention(

[mikekgfb]

@yanbing-j @Jack-Khuu current way of setting a single attention implementation fails. Two directions: if we actually want the flash attention, we need to replace the causal mask with is_causal when we specify it (and other masks are not supported by flash at all). That's a prereq to use flash in any form shape or flavor.

The second one is how to think about a fallback because non of the sdpa kernels except math is guaranteed to work on all inputs. See #1477 for a solution that puts that backup MATH in there, but it just means that atm whenever we try to call flash we end up with MATH.