From e4f886a91305b215202dd23d8eae147714982e21 Mon Sep 17 00:00:00 2001
From: Jaeyong Sung <jaeyong0201@gmail.com>
Date: Wed, 28 May 2025 09:36:00 -0400
Subject: [PATCH] fix cross attention mask

---
 dia/layers.py |  3 ++-
 dia/state.py  | 13 ++++++++-----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/dia/layers.py b/dia/layers.py
index b57afc88..f9aed506 100644
--- a/dia/layers.py
+++ b/dia/layers.py
@@ -333,7 +333,7 @@ def __init__(
         self.kv_output_dim = num_kv_heads * kv_head_dim
         self.linear = nn.Linear(in_features, out_features, bias=bias)
 
-    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+    def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         x = self.linear(inputs)
 
         q, k, v = x.split([self.q_output_dim, self.kv_output_dim, self.kv_output_dim], dim=-1)
@@ -723,6 +723,7 @@ def forward(
             Xq=x_norm,
             q_positions=state.dec_positions,
             kv_positions=state.enc_positions,
+            attn_mask=state.cross_attn_mask,
             cache=cross_attn_cache,
         )
         x = residual + ca_out
diff --git a/dia/state.py b/dia/state.py
index be16e4ad..172ec52c 100644
--- a/dia/state.py
+++ b/dia/state.py
@@ -68,6 +68,9 @@ def new(cls, config: DiaConfig, cond_src: torch.Tensor) -> "EncoderInferenceStat
 
 
 class KVCache(torch.nn.Module):
+    k: torch.Tensor
+    v: torch.Tensor
+
     def __init__(
         self,
         batch_size: int,
@@ -83,7 +86,6 @@ def __init__(
         v = torch.zeros((2 * batch_size, num_heads, max_len, head_dim), dtype=dtype, device=device) if v is None else v
         super().__init__()
 
-        self.current_idx = torch.tensor(0)
         self.register_buffer("k", k)
         self.register_buffer("v", v)
 
@@ -104,15 +106,12 @@ def update(self, k: torch.Tensor, v: torch.Tensor, current_idx: torch.Tensor) ->
         k_out, v_out = self.k, self.v
         k_out[:, :, current_idx, :] = k
         v_out[:, :, current_idx, :] = v
-        # self.current_idx += 1
-        # return self.k[:, :, : self.current_idx, :], self.v[:, :, : self.current_idx, :]
         return self.k, self.v
 
-    def prefill(self, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def prefill(self, k: torch.Tensor, v: torch.Tensor):
         prefill_len = k.shape[2]
         self.k[:, :, :prefill_len, :] = k
         self.v[:, :, :prefill_len, :] = v
-        self.current_idx = prefill_len - 1
 
 
 @dataclass
@@ -127,6 +126,7 @@ class DecoderInferenceState:
     self_attn_cache: list[KVCache]
     cross_attn_cache: list[KVCache]
     casual_attn_mask: torch.Tensor
+    cross_attn_mask: torch.Tensor
 
     @classmethod
     def new(
@@ -145,6 +145,8 @@ def new(
 
         dec_positions = torch.full((2 * batch_size, 1), fill_value=0, dtype=torch.int32, device=device)
         causal_mask = torch.tril(torch.ones(max_audio_len, max_audio_len, dtype=torch.bool, device=device))
+        dec_mask = torch.ones((2 * batch_size, 1), dtype=torch.bool, device=device)
+        cross_attn_mask = create_attn_mask(dec_mask, enc_state.padding_mask, device, is_causal=False)
 
         self_attn_cache = [
             KVCache(
@@ -167,6 +169,7 @@ def new(
             self_attn_cache=self_attn_cache,
             cross_attn_cache=dec_cross_attn_cache,
             casual_attn_mask=causal_mask,
+            cross_attn_mask=cross_attn_mask,
         )
 
     def prepare_step(self, step_from: int, step_to: int | None = None) -> None: