modify sdpa implementation

modeling_moss_audio_tokenizer.py

@@ -972,8 +972,55 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         batch_size, max_seqlen, _ = x.shape
         q, k, v = self._project_qkv(x)
         q, k = self._apply_dense_rope(q, k)
-        attn_bias = self._build_non_streaming_sdpa_bias(input_lengths, max_seqlen, x.device)
-        out = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+
+        # The native SDPA path accepts the full local-causal mask, but on long
+        # codec sequences that materializes a large T x T bias and can also hit
+        # backend-specific silent errors. Query chunking keeps the exact same
+        # non-streaming attention pattern while bounding each SDPA call.
+        query_chunk_size = 1500
+        if max_seqlen <= query_chunk_size:
+            attn_bias = self._build_non_streaming_sdpa_bias(input_lengths, max_seqlen, x.device)
+            out = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+        else:
+            out = torch.empty_like(q)
+            all_positions = torch.arange(max_seqlen, device=x.device, dtype=torch.long)
+            for q_start in range(0, max_seqlen, query_chunk_size):
+                q_end = min(q_start + query_chunk_size, max_seqlen)
+                if self.causal:
+                    k_end = q_end
+                else:
+                    k_end = max_seqlen
+
+                if self.context is not None:
+                    k_start = max(0, q_start - self.context + 1)
+                else:
+                    k_start = 0
+
+                q_positions = all_positions[q_start:q_end]
+                k_positions = all_positions[k_start:k_end]
+                valid_k = k_positions.view(1, 1, -1) < input_lengths.view(-1, 1, 1)
+                if not self.causal and self.context is None:
+                    attn_bias = valid_k[:, None, :, :].expand(-1, 1, q_end - q_start, -1)
+                else:
+                    delta = q_positions.view(1, -1, 1) - k_positions.view(1, 1, -1)
+                    attn_bias = torch.ones(
+                        (1, q_end - q_start, k_end - k_start),
+                        device=x.device,
+                        dtype=torch.bool,
+                    )
+                    if self.causal:
+                        attn_bias = attn_bias & (delta >= 0)
+                    if self.context is not None:
+                        attn_bias = attn_bias & (delta < self.context)
+                    attn_bias = (attn_bias & valid_k)[:, None, :, :]
+
+                out[:, :, q_start:q_end, :] = F.scaled_dot_product_attention(
+                    q[:, :, q_start:q_end, :],
+                    k[:, :, k_start:k_end, :],
+                    v[:, :, k_start:k_end, :],
+                    attn_bias,
+                    dropout_p=0.0,
+                )
         valid_q = (torch.arange(max_seqlen, device=x.device).view(1, max_seqlen) < input_lengths.view(-1, 1)).view(
             batch_size, 1, max_seqlen, 1
         )