def precompute_cross_attention_kv(
        self,
        max_len: int,
        encoder_out: mx.array,  # (B, S, E)
        src_positions: Optional[mx.array],  # (B, S)
    ) -> List[KVCache]:
        """
        Computes the Key and Value tensors for cross-attention for each layer from the encoder output.
        """
        per_layer_kv_cache: List[KVCache] = []

        for layer in self.layers:
            cross_attn_module = layer.cross_attention
            k_proj = cross_attn_module.k_proj(encoder_out)
            v_proj = cross_attn_module.v_proj(encoder_out)

            k_proj = cross_attn_module.rotary_emb(k_proj, position=src_positions)
            k = mx.transpose(k_proj, (0, 2, 1, 3))  # equivalent to transpose(1, 2)
            v = mx.transpose(v_proj, (0, 2, 1, 3))  # equivalent to transpose(1, 2)

            # Create KVCache without device parameter
            per_layer_kv_cache.append(
                KVCache(
                    cross_attn_module.num_kv_heads,
                    max_len,
                    cross_attn_module.head_dim,
                    k=k,
                    v=v,
                )
            )

        return per_layer_kv_cache

        Computes the Key and Value tensors for cross-attention for each layer from the encoder output.
        
mlx_audio/tts/models/dia/layers.py