no need to .view(dtype) as kv and dkv are now the same dtype

ring_attention_pytorch/ring_flash_attention_cuda.py

@@ -1,7 +1,6 @@
 import math
 from functools import partial
 from typing import Optional, Tuple
-import packaging.version as pkg_version
 
 import torch
 from torch import nn, einsum, Tensor
@@ -251,7 +250,6 @@ def backward(ctx, do):
         assert k.dtype == v.dtype
         kv_dtype = k.dtype
 
-        k, v, dk, dv = map(lambda t: t.view(torch.float32), (k, v, dk, dv))
         kv_and_dkv = torch.stack((k, v, dk, dv))
 
         # receive buffers, to be alternated with sent buffer
@@ -267,10 +265,6 @@ def backward(ctx, do):
 
             k, v, dk, dv = kv_and_dkv
 
-            # view k, v, dk, dv as the correct type of either float16 or bfloat16
-
-            k, v, dk, dv = map(lambda t: t.view(kv_dtype), (k, v, dk, dv))
-
             # translate key padding mask to bias
 
             bias = None
@@ -336,8 +330,6 @@ def backward(ctx, do):
 
             dk, dv = dkv
 
-        dk, dv = map(lambda t: t.view(kv_dtype), (dk, dv))
-
         dq, dk, dv = map(lambda t: t.to(dtype), (dq, dk, dv))
 
         return dq, dk, dv, None, None, None, None, None, None, None

ring_attention_pytorch/triton_flash_attn.py

@@ -21,6 +21,8 @@ def is_contiguous(x: Tensor):
 
 # make sure triton 2.1+ is installed
 
+import packaging.version as pkg_version
+
 import importlib
 from importlib.metadata import version
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'ring-attention-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.3.19',
+  version = '0.3.20',
   license='MIT',
   description = 'Ring Attention - Pytorch',
   author = 'Phil Wang',