From e1459ff3d0077ac90980483f0523d6dd85c6b7ce Mon Sep 17 00:00:00 2001
From: kabachuha <artemkhrapov2001@yandex.ru>
Date: Wed, 6 Mar 2024 12:18:03 +0300
Subject: [PATCH 1/2] option to use rebased linear attention

Former-commit-id: f1b39b34df5fbdef14ef3f6a0e6d0972a8570445
---
 opensora/models/diffusion/latte/latte.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py
index 5b6bf66b0..6ccb7de18 100644
--- a/opensora/models/diffusion/latte/latte.py
+++ b/opensora/models/diffusion/latte/latte.py
@@ -27,6 +27,12 @@ try:
 except:
     XFORMERS_IS_AVAILBLE = False
 
+try:
+    # needs to have https://github.com/corl-team/rebased/ installed
+    from fla.ops.triton.rebased_fast import parallel_rebased
+except:
+    REBASED_IS_AVAILABLE = False
+    
 # from timm.models.layers.helpers import to_2tuple
 # from timm.models.layers.trace_utils import _assert
 
@@ -38,7 +44,7 @@ def modulate(x, shift, scale):
 #################################################################################
 
 class Attention(nn.Module):
-    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_lora=False, attention_mode='math'):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_lora=False, attention_mode='math', eps=1e-12):
         super().__init__()
         assert dim % num_heads == 0, 'dim should be divisible by num_heads'
         self.num_heads = num_heads
@@ -49,6 +55,7 @@ class Attention(nn.Module):
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
+        self.eps = eps
 
     def forward(self, x):
         B, N, C = x.shape
@@ -70,6 +77,9 @@ class Attention(nn.Module):
             attn = self.attn_drop(attn)
             x = (attn @ v).transpose(1, 2).reshape(B, N, C)
 
+        elif self.attention_mode == 'rebased':
+            x = parallel_rebased(q, k, v, self.eps, True, True).reshape(B, N, C)
+
         else:
             raise NotImplemented
 
@@ -536,4 +546,4 @@ if __name__ == '__main__':
     # lora.mark_only_lora_as_trainable(network)
     # out = y_embeder(y, True)
     # out = network(img, t, y)
-    # print(out.shape)
\ No newline at end of file
+    # print(out.shape)

From abffb358f96686df6816837ae104d3b37bb3dc00 Mon Sep 17 00:00:00 2001
From: kabachuha <artemkhrapov2001@yandex.ru>
Date: Wed, 6 Mar 2024 12:24:10 +0300
Subject: [PATCH 2/2] support for ring attention

Former-commit-id: 0533c73b99217e88da121bd0f68e0949bc88016e
---
 opensora/models/diffusion/latte/latte.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py
index 6ccb7de18..5615897c8 100644
--- a/opensora/models/diffusion/latte/latte.py
+++ b/opensora/models/diffusion/latte/latte.py
@@ -32,6 +32,12 @@ try:
     from fla.ops.triton.rebased_fast import parallel_rebased
 except:
     REBASED_IS_AVAILABLE = False
+
+try:
+    # needs to have https://github.com/lucidrains/ring-attention-pytorch installed
+    from ring_attention_pytorch.ring_flash_attention_cuda import ring_flash_attn_cuda
+except:
+    RING_ATTENTION_IS_AVAILABLE = False
     
 # from timm.models.layers.helpers import to_2tuple
 # from timm.models.layers.trace_utils import _assert
@@ -44,7 +50,7 @@ def modulate(x, shift, scale):
 #################################################################################
 
 class Attention(nn.Module):
-    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_lora=False, attention_mode='math', eps=1e-12):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_lora=False, attention_mode='math', eps=1e-12, causal=True, ring_bucket_size=1024):
         super().__init__()
         assert dim % num_heads == 0, 'dim should be divisible by num_heads'
         self.num_heads = num_heads
@@ -56,6 +62,8 @@ class Attention(nn.Module):
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
         self.eps = eps
+        self.causal = causal
+        self.ring_bucket_size = ring_bucket_size
 
     def forward(self, x):
         B, N, C = x.shape
@@ -80,6 +88,9 @@ class Attention(nn.Module):
         elif self.attention_mode == 'rebased':
             x = parallel_rebased(q, k, v, self.eps, True, True).reshape(B, N, C)
 
+        elif self.attention_mode == 'ring':
+            x = ring_flash_attn_cuda(q, k, v, causal=self.causal, bucket_size=self.ring_bucket_size).reshape(B, N, C)
+
         else:
             raise NotImplemented