Diff Coverage

    if interleaved:
        cos, sin = (part[..., :tensor.shape[-1] // 2].repeat_interleave(2, dim=-1)
                    for part in (cos, sin))
    cos, sin = (part.expand(tensor.shape[0], -1, -1).transpose(0, 1) for part in (cos, sin))
    rotated = functional.apply_rotary_pos_emb_single(
        tensor.transpose(0, 1), cos, sin,
        rotary_interleaved=interleaved, use_fused_rotary_pos_emb=True,
    )
    if interleaved:

    """Expand and rotate the query through the managed TP-local projection module."""
    projected = module.q_b_proj(latent)
    projected = projected.view(*latent.shape[:2], heads, module.qk_head_dim)
    content, rotary = projected.split((module.qk_nope_head_dim, module.qk_rope_head_dim), dim=-1)
    rotary = rotary_fn(rotary, positions, interleaved=module.rotary_interleaved)
    return torch.cat((content, rotary), dim=-1).transpose(1, 2)


def _key_value(

            )
            lengths, key_lengths = partition.query_lengths, partition.key_lengths
        positions = _positions(position_embeddings, state, dimensions.rope_dim)
        _validate_state(module, self.plan, state)
        query, key, value = self._attention_inputs(module, state, positions, partition)
        attention_interface = (
            npu_fusion_attention_forward if self.attention_interface is None else self.attention_interface
        )
        output = attention_interface(
            module, query, key, value, attention_mask, scaling=module.scaling,
            is_causal=causal, actual_seq_qlen=lengths, actual_seq_kvlen=key_lengths,
            pre_tokens=2147483647, next_tokens=0 if causal else 2147483647,
            mla_head_offset=self.head_rank * self.plan.compute_heads,

        return output, None

    def _attention_inputs(self, module, state, positions, partition):
        """Prepare Q/K/V through the selected rotary callback and communication schedule."""
        degree = self.plan.cp_degree
        rotary_fn = apply_mla_rope if self.rotary_fn is None else self.rotary_fn
        rotated_key = rotary_fn(
            state.key_rope.unsqueeze(2), positions, interleaved=module.rotary_interleaved,
        ).squeeze(2)
        if self.plan.overlap and degree > 1:
            prepare = self._expanded_overlap if self.plan.strategy == "expanded_ulysses" else self._latent_overlap
            query, key, value = prepare(module, state, positions, rotated_key, partition, rotary_fn=rotary_fn)
        else:
            query = _query(module, state.query_latent, positions, self.plan.local_heads, rotary_fn)
            if self.plan.strategy == "expanded_ulysses" or degree == 1:
                query, key, value = self._expanded(module, state, query, rotated_key, partition)
            else:
                query, key, value = self._latent(module, state, query, rotated_key, partition)
        return query, key, value

    def _expanded(self, module, state, query, key_rope, partition=None):
        """Pack unequal-width QKV into one sequence-to-head exchange."""
        key, value = _key_value(module, state.kv_latent, key_rope, self.plan.local_heads, None)

        if use_ulysses:
            pending_kv = async_ulysses_seq_to_head_launch(payload, 1, 2, self.ulysses_mesh)
        else:
            pending_kv = self._expanded_kv_launch(payload, partition)
        query = _query(module, state.query_latent, positions, self.plan.local_heads, rotary_fn)
        pending_query = (async_ulysses_seq_to_head_launch(query.transpose(1, 2), 1, 2, self.ulysses_mesh)
                         if use_ulysses else None)
        payload = pending_kv.wait()
        if use_ulysses and self.plan.all_gather_degree > 1:

            pending_kv = _selective_kv_launch(payload, partition, self.cp_mesh, self.plan.ulysses_degree, True)
            partition = None
        else:
            pending_kv = async_cp_allgather_launch(payload, 1, self.cp_mesh)
        query = _query(module, state.query_latent, positions, self.plan.local_heads, rotary_fn)
        payload = pending_kv.wait()
        pending_query = (async_ulysses_seq_to_head_launch(query.transpose(1, 2), 1, 2, self.ulysses_mesh)
                         if self.plan.ulysses_degree > 1 else None)
        key, value = self._expand_latent(module, payload, partition)