FA 2.4.2 is falling unitest on A6000 and A5880

[BoxiangW]

I am using NVIDIA's 24.10-py3 container. This is how to reproduce:

git clone https://github.com/Dao-AILab/flash-attention.git
Verify the installed version of python -c 'import flash_attn; print(flash_attn.__version__)'
cd flash-attention && git checkout v2.4.2 && cd .. # for example, make sure you are not in the flash-attention dir when you run the test
py.test flash-attention/tests/test_flash_attn.py::test_flash_attn_qkvpacked[0.17-2048-160-True-True-False-False-dtype0]

I tested that this only happens on A5880 and A6000. Could be affecting A100 as well. H100 is not affected.