Upload BD3LM

modeling_bd3lm.py

@@ -547,9 +547,12 @@ class DITBackbone(nn.Module):
     if self.adaln:
       c = F.silu(self.sigma_map(sigma))
     if self.cross_attn:
+      mask = self.mask.to(x.device)
       n = self.mask.shape[-1] // 2
       # use block-causal mask only during sampling
-      if sample_mode:
+      if not sample_mode:
+        rotary_cos_sin = self.rotary_emb(x[:, :self.n])
+      else:
         if self.blocks[0].kv_cache is not None:
           mask = None
           accum_length = self.blocks[0].cache_idx + self.block_size
@@ -558,12 +561,9 @@ class DITBackbone(nn.Module):
             x.shape[0], accum_length, x.shape[2]), device=x.device)
           rotary_cos_sin = self.rotary_emb(x_full)
         else:
-          mask = self.mask.to(x.device)
           rotary_cos_sin = self.rotary_emb(x[:, :n])
           mask = mask[
             n:n+x.shape[1], n:n+x.shape[1]]
-      else:
-        rotary_cos_sin = self.rotary_emb(x[:, :self.n])
     else:
       mask = None
       rotary_cos_sin = self.rotary_emb(x)