add ring_flash_attn_cuda to __init__.py

ring_attention_pytorch/__init__.py

@@ -10,3 +10,8 @@
     ring_flash_attn,
     ring_flash_attn_
 )
+
+from ring_attention_pytorch.ring_flash_attention_cuda import (
+    ring_flash_attn_cuda,
+    ring_flash_attn_cuda_
+)

ring_attention_pytorch/ring_flash_attention_cuda.py

@@ -44,21 +44,6 @@ def pad_at_dim(t, pad: Tuple[int, int], *, dim = -1, value = 0.):
 def is_empty(t: Tensor):
     return t.numel() == 0
 
-# make sure triton is installed
-
-import importlib
-from importlib.metadata import version
-
-assert exists(importlib.util.find_spec('triton')), 'latest triton must be installed. `pip install triton -U` first'
-
-triton_version = version('triton')
-assert pkg_version.parse(triton_version) >= pkg_version.parse('2.1'), 'triton must be version 2.1 or above. `pip install triton -U` to upgrade'
-
-from ring_attention_pytorch.triton_flash_attn import (
-    flash_attn_backward,
-    flash_attn_forward
-)
-
 # ring + (flash) attention forwards and backwards
 
 # flash attention v1 - https://arxiv.org/abs/2205.14135
@@ -82,6 +67,8 @@ def forward(
         max_lookback_seq_len: Optional[int],
         ring_size: Optional[int]
     ):
+        from ring_attention_pytorch.triton_flash_attn import flash_attn_forward
+
         assert all([t.is_cuda for t in (q, k, v)]), 'inputs must be all on cuda'
 
         dtype = q.dtype
@@ -220,7 +207,8 @@ def forward(
     @staticmethod
     @torch.no_grad()
     def backward(ctx, do):
-        """ Algorithm 2 in the v2 paper """
+
+        from ring_attention_pytorch.triton_flash_attn import flash_attn_backward
 
         (
             causal,

ring_attention_pytorch/triton_flash_attn.py

@@ -7,8 +7,6 @@
 
 import torch
 from torch import Tensor
-import triton
-import triton.language as tl
 
 from einops import repeat
 
@@ -21,6 +19,21 @@ def default(val, d):
 def is_contiguous(x: Tensor):
     return x.stride(-1) == 1
 
+# make sure triton 2.1+ is installed
+
+import importlib
+from importlib.metadata import version
+
+assert exists(importlib.util.find_spec('triton')), 'latest triton must be installed. `pip install triton -U` first'
+
+triton_version = version('triton')
+assert pkg_version.parse(triton_version) >= pkg_version.parse('2.1'), 'triton must be version 2.1 or above. `pip install triton -U` to upgrade'
+
+import triton
+import triton.language as tl
+
+# kernels
+
 @triton.heuristics(
     {
         "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,