Bump to v2.1.2

.github/workflows/publish.yml

@@ -63,11 +63,17 @@ jobs:
             # Pytorch <= 2.0 only supports CUDA <= 11.8
             - torch-version: '1.12.1'
               cuda-version: '12.1.0'
+            - torch-version: '1.12.1'
+              cuda-version: '12.2.0'
             - torch-version: '1.13.1'
               cuda-version: '12.1.0'
+            - torch-version: '1.13.1'
+              cuda-version: '12.2.0'
             - torch-version: '2.0.1'
               cuda-version: '12.1.0'
-            # Pytorch >= 2.1 only supports CUDA 12.1
+            - torch-version: '2.0.1'
+              cuda-version: '12.2.0'
+            # Pytorch >= 2.1 only supports CUDA >= 12.1
             - torch-version: '2.1.0.dev20230731'
               cuda-version: '11.6.2'
             - torch-version: '2.1.0.dev20230731'

csrc/flash_attn/src/flash_fwd_kernel.h

@@ -947,7 +947,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params &params, cons
     for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) {
         float sum = scores_sum(mi);
         float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum;
-        lse(mi) = (sum == 0.f || sum != sum) ? INFINITY : scores_max(mi) * params.scale_softmax + __logf(sum);
+        lse(mi) = (sum == 0.f || sum != sum) ? -INFINITY : scores_max(mi) * params.scale_softmax + __logf(sum);
         float scale = inv_sum;
         #pragma unroll
         for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { acc_o_rowcol(mi, ni) *= scale; }

flash_attn/__init__.py

-__version__ = "2.1.1"
+__version__ = "2.1.2"
 
 from flash_attn.flash_attn_interface import (
     flash_attn_func,

training/Dockerfile

@@ -85,11 +85,11 @@ RUN pip install transformers==4.25.1 datasets==2.8.0 pytorch-lightning==1.8.6 tr
 RUN pip install git+https://github.com/mlcommons/logging.git@2.1.0
 
 # Install FlashAttention
-RUN pip install flash-attn==2.1.1
+RUN pip install flash-attn==2.1.2
 
 # Install CUDA extensions for cross-entropy, fused dense, layer norm
 RUN git clone https://github.com/HazyResearch/flash-attention \
-    && cd flash-attention && git checkout v2.1.1 \
+    && cd flash-attention && git checkout v2.1.2 \
     && cd csrc/fused_softmax && pip install . && cd ../../ \
     && cd csrc/rotary && pip install . && cd ../../ \
     && cd csrc/xentropy && pip install . && cd ../../ \