"""RoPE、GQA 的 repeat_kv 与增量 KV Cache 形状演示。"""
import torch


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x[..., ::2], x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)


def apply_rope(x: torch.Tensor, positions: torch.Tensor, base: float = 10_000.0):
    # x: (B, H, T, Dh)，Dh 必须为偶数。
    half = x.shape[-1] // 2
    inv = base ** (-torch.arange(half, device=x.device) / half)
    angle = positions[:, None] * inv[None, :]
    cos = torch.repeat_interleave(angle.cos(), 2, dim=-1)[None, None]
    sin = torch.repeat_interleave(angle.sin(), 2, dim=-1)[None, None]
    return x * cos + rotate_half(x) * sin


def repeat_kv(x: torch.Tensor, q_heads: int) -> torch.Tensor:
    # x: (B, Hkv, T, Dh)
    h_kv = x.shape[1]
    if q_heads % h_kv:
        raise ValueError("q_heads must be divisible by kv_heads")
    return x.repeat_interleave(q_heads // h_kv, dim=1)


def append_cache(cache: torch.Tensor | None, new: torch.Tensor) -> torch.Tensor:
    return new if cache is None else torch.cat([cache, new], dim=2)


if __name__ == "__main__":
    q = torch.randn(1, 8, 3, 16)
    k = torch.randn(1, 2, 3, 16)
    pos = torch.arange(3)
    q_rot = apply_rope(q, pos)
    k_rot = apply_rope(k, pos)
    assert q_rot.shape == q.shape
    assert repeat_kv(k_rot, 8).shape == q.shape
    assert append_cache(k_rot[:, :, :2], k_rot[:, :, 2:]).shape[2] == 3