Diff Coverage

    if compress_ratio <= 0:
        raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
    if query_chunk_size <= 0:
        raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
    batch_size, sequence_length, _, head_dim = query.shape
    if key.ndim != 3 or key.shape[0] != batch_size or key.shape[2] != head_dim:
        raise ValueError(
            "compressed index key must have shape [batch, compressed_sequence, head_dim]"
        )

                batch_indices = torch.arange(batch_size, device=index_query.device).view(-1, 1, 1)

                query_chunk = index_query[:, start:end].float()
                weight_chunk = merge_weight[:, start:end].float()
                selected_index_key = _gather_keys(
                    index_key, batch_indices, safe_indices).float()
                index_dots = torch.einsum(
                    "bcid,bckd->bcik", query_chunk, selected_index_key
                )

                if reduce_sum is not None:
                    index_scores = reduce_sum(index_scores)
                index_scores.masked_fill_(~valid, -1.0e9)

                selected_attention_key = _gather_keys(
                    compressed_key, batch_indices, safe_indices).float()
                attention_scores = torch.einsum(
                    "bhcd,bckd->bhck",
                    attention_query[:, :, start:end].float(),