Diff Coverage

            @staticmethod
            def apply(*args: Any, **kwargs: Any) -> Any:
                """Apply the original KL loss with saved global sequence tensors."""
                arguments = _DSA_KL_SIGNATURE.bind(*args, **kwargs).arguments
                saved = _dsa_tensor_context.get()
                if saved is not None:
                    _dsa_tensor_context.set(None)
                    length = saved.query.size(1)
                    arguments.update(
                        query=saved.query.flatten(0, 1),
                        key=saved.key.flatten(0, 1),
                        query_rope=saved.q_pe.flatten(0, 1),
                        key_rope=saved.k_pe.flatten(0, 1),

def _head_tail_peer_rank(cp_mesh: Any) -> int:
    """Return the mirror rank in the CP mesh."""
    rank_list = list(cp_mesh.rank_list)
    return rank_list[cp_mesh.size() - 1 - rank_list.index(dist.get_rank())]


def _run_head_tail_half(
        attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],

        attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],
        query: Tensor, key: Tensor, value: Tensor,
        call_kwargs: dict[str, Any]) -> Tensor:
    """Run attention for one query half and check its output contract."""
    output = attention_fn(query, key, value, call_kwargs)
    if not isinstance(output, Tensor):
        raise TypeError(
            "Head-Tail load balance requires the attention callable to "
            f"return a Tensor, got {type(output).__name__}"
        )
    return output


def head_tail_load_balance_attention(
        attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],

            f"length, got {local_q_len}; pad the global sequence to a "
            f"multiple of 2 * cp_size ({2 * cp_mesh.size()})"
        )

    peer_rank = _head_tail_peer_rank(cp_mesh)
    query_peer = _collectives.p2p_exchange(
        query.narrow(2, local_q_len // 2, local_q_len // 2), peer_rank)
    global_key, global_value = flex_cp_allgather(key, value, 2, cp_mesh)
    keep_output = _run_head_tail_half(
        attention_fn, query.narrow(2, 0, local_q_len // 2),
        global_key, global_value, attention_kwargs)
    peer_output = _run_head_tail_half(
        attention_fn, query_peer, global_key, global_value,
        attention_kwargs if peer_attention_kwargs is None else peer_attention_kwargs,
    )
    return torch.cat(
        [keep_output, _collectives.p2p_exchange(peer_output, peer_rank)], dim=2)


def _cp_offset_causal_mask(q_len: int, kv_len: int, lo: int,