Upload BD3LM

modeling_bd3lm.py

@@ -299,7 +299,6 @@ class DDiTBlock(nn.Module):
     else:
       return bias_dropout_add_scale_fused_inference
   
-
   def get_qkv(self, x, rotary_cos_sin, store_kv=False):
     # compute qkv (potentially use cache)
     if self.kv_cache is not None:
@@ -307,11 +306,10 @@ class DDiTBlock(nn.Module):
       qkv = torch.cat((self.kv_cache, new_qkv), dim=1)
     else:
       qkv = self.attn_qkv(x)
-    
     # store kv cache in a sliding window (can't exceed context len)
     if store_kv:
-      self.kv_cache = qkv
-      self.kv_cache = self.kv_cache[:, -self.n:]
+      self.kv_cache = qkv[:, -(self.n-self.block_size):]
+
     qkv = einops.rearrange(
       qkv,
       'b s (three h d) -> b s three h d',