Diff Coverage

        raise RuntimeError(
            "use_fused_rope=True requires cann_ops_transformer.inplace_partial_rotary_mul. "
            "See hyper_parallel/components/functional/README.md for installation and supported inputs."
        ) from err
    return inplace_partial_rotary_mul


class _PartialRotaryMul(torch.autograd.Function):
    """Keep caller storage intact and preserve the reference's branch-wise casts."""

def _validate_rotary_coefficients(tensor: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> None:
    """Check positional constants before preparing the native broadcast layout."""
    if cos.ndim != 3 or sin.shape != cos.shape:
        raise ValueError("Partial RoPE requires matching cos/sin shapes [B, S, R/2]")
    if cos.requires_grad or sin.requires_grad:
        raise ValueError("Partial RoPE requires constant cos/sin without gradients")
    if cos.dtype != torch.float32 or sin.dtype != torch.float32:
        raise ValueError("Partial RoPE requires FP32 cos/sin to preserve V4.1 rounding")
    if cos.device != tensor.device or sin.device != tensor.device:
        raise ValueError("Partial RoPE input and cos/sin must share a device")


def partial_rotary_mul(
        tensor: torch.Tensor,

        ValueError: Shapes, positional gradients or coefficient dtypes are unsupported.
        RuntimeError: Inputs are not on NPU or the required extension is missing.
    """
    if tensor.ndim != 4 or unsqueeze_dim not in (1, 2):
        raise ValueError("Partial RoPE requires a 4D tensor and unsqueeze_dim 1 or 2")
    _validate_rotary_coefficients(tensor, cos, sin)
    if tensor.device.type != "npu":
        raise RuntimeError(f"Fused partial RoPE requires NPU inputs; got device={tensor.device}")
    value = tensor if unsqueeze_dim == 2 else tensor.transpose(1, 2)
    if cos.shape[0] not in (1, value.shape[0]) or cos.shape[1] not in (1, value.shape[1]):
        raise ValueError("Partial RoPE cos/sin batch and sequence dimensions must broadcast to the input")
    rotary_dim = cos.shape[-1] * 2
    if rotary_dim > tensor.shape[-1]:
        raise ValueError("Partial RoPE frequencies cannot exceed the input head dimension")
    if not tensor.numel() or not rotary_dim:
        return tensor.clone()
    cosine = cos.expand(value.shape[0], -1, -1).repeat_interleave(2, dim=-1).unsqueeze(2).contiguous()
    sine = sin.expand(value.shape[0], -1, -1).repeat_interleave(2, dim=-1).unsqueeze(2).contiguous()
    output = _PartialRotaryMul.apply(value, cosine, sine)
    return output if unsqueeze_dim == 2 else output.transpose(1, 2)

        use_fused_rope: bool = False,
) -> torch.Tensor:
    """Apply V4 interleaved RoPE to the trailing rotary channels."""
    if use_fused_rope:
        return partial_rotary_mul(tensor, cos, sin, unsqueeze_dim)
    cos = cos.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
    sin = sin.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
    rope_dim = cos.shape[-1]
    pass_through, rotary = tensor.split((tensor.shape[-1] - rope_dim, rope_dim), dim=-1)

                f"Indexer query width {query.shape[-1]} is not divisible by head_dim {self.head_dim}"
            )
        query = query.view(*query_residual.shape[:2], -1, self.head_dim)
        if sequence_major:
            return _apply_v41_rope(query, cos, sin, unsqueeze_dim=2, use_fused_rope=self.use_fused_rope)
        return _apply_v41_rope(
            query.transpose(1, 2), cos, sin, use_fused_rope=self.use_fused_rope,
        ).transpose(1, 2)

                query = projected.new_empty((batch, length, projected.shape[-1]))
            query[:, start:end].copy_(projected[:, :end - start])
            del projected, chunk
        query = query.view(batch, length, query.shape[-1] // self.head_dim, self.head_dim)
        return _apply_v41_rope(query, cos, sin, unsqueeze_dim=2, use_fused_rope=self.use_fused_rope)

    def _project_index_key(
            self,
            latent: torch.Tensor | None,

                key_cos = cos[:, :compressed_length * self.compress_ratio:self.compress_ratio]
                key_sin = sin[:, :compressed_length * self.compress_ratio:self.compress_ratio]
            else:
                key_cos, key_sin = key_position_embeddings
            key = _apply_v41_rope(
                key.unsqueeze(1), key_cos, key_sin, use_fused_rope=self.use_fused_rope,
            ).squeeze(1)
            key_handle = None
            if cp_context is not None: