def __call__(
        self,
        Xq: mx.array,  # (B, T, D) T = 1 in AR generation
        Xkv: mx.array,  # (B, S, E) S = 1 in AR generation
        q_positions: mx.array,  # (B, T)
        kv_positions: Optional[mx.array] = None,  # (B, S)
        deterministic: bool = True,
        attn_mask: Optional[
            mx.array
        ] = None,  # None in Decoder Self Attention, Valid mask in Others
        cache: Optional[KVCache] = None,  # None in Encoder, KVCache in Decoder
        prefill: bool = False,  # True only when prefilling KV Cache
    ) -> Tuple[mx.array, Optional[Tuple[mx.array, mx.array]]]:
        """
        Performs attention calculation with optional KV caching.

        Args:
            Xq: Query tensor (B, T, D). T=1 during single-step decoding.
            Xkv: Key/Value source tensor (B, S, E). S=1 during single-step decoding for self-attn.
            q_positions: Positions for queries (B, T).
            kv_positions: Positions for keys/values (B, S). If None, uses q_positions.
            deterministic: If True, disable dropout.
            attn_mask: Attention mask.
            cache: KVCache.
            prefill: If True, use prefill mode.

        Returns:
            A tuple containing:
            - output: The attention output tensor (B, T, output_dim).
            - present_kv: The K/V state to be cached for the next step ((B, N, S_new, H), (B, N, S_new, H)).
              For self-attn, S_new = S_past + S. For cross-attn, S_new = S_kv.
        """
        if kv_positions is None:
            kv_positions = q_positions
        original_dtype = Xq.dtype

        Xq_BxTxNxH = self.q_proj(Xq)
        Xq_BxTxNxH = self.rotary_emb(Xq_BxTxNxH, position=q_positions)
        Xq_BxNxTxH = mx.transpose(Xq_BxTxNxH, (0, 2, 1, 3))

        # Input values into attention calculation
        attn_k = None
        attn_v = None

        # Decoder Cross Attention
        if self.is_cross_attn:
            # Directly use cache (no need to check index)
            attn_k, attn_v = cache.k, cache.v
            if (
                attn_k.shape[1] != self.num_query_heads
                or attn_v.shape[1] != self.num_query_heads
            ):
                raise ValueError(
                    f"Cross-attention cache head dimension ({attn_k.shape[1]}) "
                    f"does not match num_query_heads ({self.num_query_heads}). "
                    "Cache should be pre-repeated for GQA."
                )
        # Self Attention
        else:
            Xk_BxSxKxH = self.k_proj(Xkv)  # (B, S, K, H)
            Xv_BxSxKxH = self.v_proj(Xkv)  # (B, S, K, H)
            Xk_BxSxKxH = self.rotary_emb(
                Xk_BxSxKxH, position=kv_positions
            )  # (B, S, K, H)

            Xk_BxKxSxH = mx.transpose(Xk_BxSxKxH, (0, 2, 1, 3))  # (B, K, S, H)
            Xv_BxKxSxH = mx.transpose(Xv_BxSxKxH, (0, 2, 1, 3))  # (B, K, S, H)
            # S=1 for Decode Step

            if self.num_gqa_groups > 1:
                Xk_BxNxSxH = repeat(
                    Xk_BxKxSxH, "b k s h -> b (k g) s h", g=self.num_gqa_groups
                )
                Xv_BxNxSxH = repeat(
                    Xv_BxKxSxH, "b k s h -> b (k g) s h", g=self.num_gqa_groups
                )
            else:
                Xk_BxNxSxH = Xk_BxKxSxH
                Xv_BxNxSxH = Xv_BxKxSxH

            # Encoder Self Attention
            if cache is None:
                attn_k = Xk_BxNxSxH
                attn_v = Xv_BxNxSxH
            # Decoder Self Attention
            else:
                # In prefill mode, we fill in cache until prefill length
                if prefill:
                    attn_k, attn_v = Xk_BxNxSxH, Xv_BxNxSxH
                    cache.prefill_kv(attn_k, attn_v)
                # In decode step, we add current K/V to cache step by step
                else:
                    attn_k, attn_v = cache.update_and_fetch(Xk_BxNxSxH, Xv_BxNxSxH)

        # Attention Calculation
        attn_scores = mx.matmul(Xq_BxNxTxH, attn_k.swapaxes(2, 3))

        # Apply Scaling
        scale_factor = 1.0
        attn_scores = attn_scores * scale_factor

        # Apply Attention Mask
        if attn_mask is not None:
            # Add large negative value where mask is False/0
            attn_scores = mx.where(
                attn_mask, attn_scores, -1e9
            )  # Using -1e9 for numerical stability

        attn_weights = mx.softmax(attn_scores, axis=-1)
        attn_output = mx.matmul(attn_weights, attn_v)

        attn_output = mx.transpose(attn_output, (0, 2, 1, 3))  # (B, T, N, H)
        output = self.o_proj(attn_output)

        if output.dtype != original_dtype:
            output = output.astype(original_dtype)

        return output

        Performs attention calculation with optional KV caching.

        Args:
            Xq: Query tensor (B, T, D). T=1 during single-step decoding.
            Xkv: Key/Value source tensor (B, S, E). S=1 during single-step decoding for self-attn.
            q_positions: Positions for queries (B, T).
            kv_positions: Positions for keys/values (B, S). If None, uses q_positions.
            deterministic: If True, disable dropout.
            attn_mask: Attention mask.
            cache: KVCache.
            prefill: If True, use prefill mode.

        Returns:
            A tuple containing:
            - output: The attention output tensor (B, T, output_dim).
            - present_kv: The K/V state to be cached for the next step ((B, N, S_new, H), (B, N, S_new, H)).
              For self-attn, S_new = S_past + S. For cross-attn, S_new = S_kv.
        
mlx_audio/tts/models/dia/layers.py