Bump flash-attn to v2.0.4

[tmm1]

What does this PR do?

Fixes #805

see https://github.com/Dao-AILab/flash-attention/commits/main for recent fixes

cc #712 Dao-AILab/flash-attention#359 Dao-AILab/flash-attention#334
cc #795 @danthe3rd

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

[codecov-commenter]

Codecov Report

Patch coverage: 96.15% and project coverage change: +0.12% 🎉

Comparison is base (f525106) 81.73% compared to head (e115d8e) 81.85%.
Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #816      +/-   ##
==========================================
+ Coverage   81.73%   81.85%   +0.12%     
==========================================
  Files          96       96              
  Lines        6401     6427      +26     
==========================================
+ Hits         5232     5261      +29     
+ Misses       1169     1166       -3     

Flag	Coverage Δ
Python	81.85% <96.15%> (+0.12%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files Changed	Coverage Δ
xformers/ops/fmha/flash.py	48.46% <80.00%> (-1.17%)	⬇️
xformers/ops/fmha/triton.py	70.37% <100.00%> (+14.67%)	⬆️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ghunkins]

Would love to see this added!

[danthe3rd]

Hi,
Thanks for opening this PR!
Happy to merge once you have reverted the changes to nvcc_flags in setup.py

[danthe3rd]

Can we revert these changes? I don't think we need to use O2 for Flash-attention (and also might make performance worse)

[tmm1]

sure, i pulled out that commit. thanks for the feedback.

[danthe3rd]

Awesome! Thank you for your contribution

[fmassa]

Looks like this PR gives wrong results for flashattn backend. We need to investigate this, but in the meantime we might revert it

[tmm1]

Is there a failing test somewhere, or what exactly is the issue?

[fmassa]

There are some failing tests (which run on internal infra), but not as many as I originally thought.

Here they are

[danthe3rd]

Here is a simple repro (I get the failures on A100 at least):

$ python -m pytest tests/test_mem_eff_attention.py -k "test_backward[flshattBv2-cuda-torch.float16-BlockDiagonalCausalMask"

===================================================================================== short test summary info ======================================================================================
FAILED tests/test_mem_eff_attention.py::test_backward[flshattBv2-cuda-torch.float16-BlockDiagonalCausalMask-1-256-2-1-32-32-False-BMHK] - AssertionError: cutlassF+flshattBv2:query: out=nan and ref=0.0 (diff=nan > 0) at (0, 174, 0, 0) of shape (1, 256, 1, 32) / atol=0.09, rtol=0.02/ total failing elements: 0, percentage=0.0
FAILED tests/test_mem_eff_attention.py::test_backward[flshattBv2-cuda-torch.float16-BlockDiagonalCausalMask-1-256-2-1-32-32-True-BMHK] - AssertionError: cutlassF+flshattBv2:query: out=nan and ref=0.0 (diff=nan > 0) at (0, 174, 0, 0) of shape (1, 256, 1, 32) / atol=0.09, rtol=0.02/ total failing elements: 0, percentage=0.0
FAILED tests/test_mem_eff_attention.py::test_backward[flshattBv2-cuda-torch.float16-BlockDiagonalCausalMask-1-256-15-1-32-32-False-BMHK] - AssertionError: flshattFv2+flshattBv2:query: out=nan and ref=0.0 (diff=nan > 0) at (0, 245, 0, 0) of shape (1, 256, 1, 32) / atol=0.09, rtol=0.02/ total failing elements: 0, percentage=0.0
FAILED tests/test_mem_eff_attention.py::test_backward[flshattBv2-cuda-torch.float16-BlockDiagonalCausalMask-1-256-15-1-32-32-True-BMHK] - AssertionError: flshattFv2+flshattBv2:query: out=nan and ref=0.0 (diff=nan > 0) at (0, 245, 0, 0) of shape (1, 256, 1, 32) / atol=0.09, rtol=0.02/ total failing elements: 0, percentage=0.0
==================================================================== 4 failed, 120 passed, 8 skipped, 14123 deselected in 8.11s ====================================================================

I'll investigate this and send an MVP to Tri if it's a bug in Flash-Attention

[danthe3rd]

This is indeed a bug in Flash-Attention - I opened an issue in Dao-AILab/flash-attention#443

[tmm1]

@danthe3rd in your bug repro it still fails with bfloat16 for me, but the tests for bfloat16 pass here?

pytest tests/test_mem_eff_attention.py -vk "test_backward and flshattBv2-cuda and torch.bfloat16 and BlockDiagonalCausalMask"

[tmm1]

Tests pass again if we revert 698532d as well.

[danthe3rd]

Yes indeed. I was hopping that Tri might have some insight on what is causing this bug if there is a narrower condition maybe. Might also be specific to variable sequence length (BlockDiagonalCausalMask)

[jinqiua]

I try to change xops.fmha.cutlass.FwOp() to xops.fmha.flash.FwOp(), but it doesn`t get any speedup result.
(I use Xformer 0.0.21+ba5b449.d20230817), @tmm1