fix for latest flash-attn function signatures

facebookresearch · Aug 11, 2023 · e654647 · e654647
1 parent fc48943
commit e654647
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/xformers/ops/fmha/flash.py b/xformers/ops/fmha/flash.py
@@ -67,6 +67,7 @@ def _flash_fwd(
             out_padded,
             softmax_lse,
             p,
+            rng_state,
         ) = _C_flashattention.varlen_fwd(
             query,
             key,
@@ -83,7 +84,7 @@ def _flash_fwd(
             return_softmax,
             None,
         )
-        return out, softmax_lse, None
+        return out, softmax_lse, rng_state
 
     def _flash_bwd(
         grad,
@@ -123,6 +124,7 @@ def _flash_bwd(
             False,  # zero_tensors
             causal,
             None,
+            rng_state,
         )
         return dq