Only enable torch 2.1.0.dev and cuda 12.1

.github/workflows/publish.yml

@@ -43,9 +43,12 @@ jobs:
           # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
           # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
           os: [ubuntu-20.04]
-          python-version: ['3.7', '3.8', '3.9', '3.10']
-          torch-version: ['1.12.1', '1.13.1', '2.0.1', '2.1.0.dev20230731']
-          cuda-version: ['11.6.2', '11.7.1', '11.8.0', '12.1.0']
+          # python-version: ['3.7', '3.8', '3.9', '3.10']
+          # torch-version: ['1.12.1', '1.13.1', '2.0.1', '2.1.0.dev20230731']
+          # cuda-version: ['11.6.2', '11.7.1', '11.8.0', '12.1.0']
+          python-version: ['3.8', '3.9', '3.10']
+          torch-version: ['2.1.0.dev20230731']
+          cuda-version: ['12.1.0']
           # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
           # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
           # Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs)

flash_attn/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.0.6.post1"
+__version__ = "2.0.6.post2"
 
 from flash_attn.flash_attn_interface import flash_attn_func
 from flash_attn.flash_attn_interface import flash_attn_kvpacked_func