def precompute_cross_attn_cache(
        self,
        enc_out: torch.Tensor,  # (B, S, E)
        enc_positions: torch.Tensor,  # (B, S)
        k_padding_mask: torch.Tensor | None = None,
    ) -> 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(enc_out)
            v_proj = cross_attn_module.v_proj(enc_out)

            k_proj = cross_attn_module.rotary_emb(k_proj, position=enc_positions)
            k = k_proj.transpose(1, 2)
            v = v_proj.transpose(1, 2)
            if k_padding_mask is not None:
                k = k.masked_fill(~k_padding_mask.unsqueeze(1).unsqueeze(3), 0.0)

            per_layer_kv_cache.append(KVCache.from_kv(k, v))

        return per_layer_kv_cache

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