Update rope.py

fix: apply rotatory embed.

rope.py

@@ -16,37 +16,54 @@ def precompute_freqs_cis(
     freqs = torch.exp(1j * freqs)
     return torch.stack([freqs.real, freqs.imag], dim=-1)
 
-def apply_rotary_emb(x_q: torch.Tensor,
-                     x_k: torch.Tensor,
-                     cos: torch.Tensor,
-                     sin: torch.Tensor,
-                     position_ids: torch.Tensor | None = None):
+# rope.py
+import torch
+
+def apply_rotary_emb(
+    x: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    position_ids: torch.Tensor,
+    num_heads: int,
+    rot_dim: int = 32,
+    interleave: bool = False,
+) -> torch.Tensor:
     """
-    Applies RoPE to q, k.
-    Keeps shapes: (B, H, T, D) or (B, T, H, D). Cos/sin are cast to x.dtype.
+    RoPE as used in the original moondream2 text stack:
+      x: (B, H, T, D)
+      freqs_cis: (max_T, rot_dim//2, 2) where [...,0]=cos, [...,1]=sin
+      position_ids: (T,) or (B,T)
+      returns x with first rot_dim dims rotated.
     """
-    # Align dtypes to avoid fp16/bf16 → fp32 mismatches
-    cos = cos.to(x_q.dtype)
-    sin = sin.to(x_q.dtype)
-
-    def _rotate_half(x):
-        d = x.shape[-1] // 2
-        x1, x2 = x[..., :d], x[..., d:]
-        return torch.cat((-x2, x1), dim=-1)
-
-    if position_ids is not None:
-        # (B, T) → index cos/sin per token
-        # cos/sin expected as (L, D) or broadcastable to x
-        cos_ = cos.index_select(0, position_ids.view(-1)).view(*position_ids.shape, -1)
-        sin_ = sin.index_select(0, position_ids.view(-1)).view(*position_ids.shape, -1)
-        # reshape to broadcast over heads
-        while cos_.dim() < x_q.dim():
-            cos_ = cos_.unsqueeze(1)
-            sin_ = sin_.unsqueeze(1)
+    assert rot_dim == freqs_cis.shape[-2] * 2
+    assert num_heads == x.shape[1]
+
+    B, H, T, D = x.shape
+    rd = min(rot_dim, D)
+    x_rot, x_pass = x[..., :rd], x[..., rd:]
+
+    # split real/imag parts depending on layout
+    if interleave:
+        xr = x_rot.float().reshape(*x_rot.shape[:-1], -1, 2)[..., 0]
+        xi = x_rot.float().reshape(*x_rot.shape[:-1], -1, 2)[..., 1]
     else:
-        cos_, sin_ = cos, sin
+        d = x_rot.shape[-1] // 2
+        xr, xi = x_rot[..., :d], x_rot[..., d:]
+
+    # gather cos/sin for these positions
+    if position_ids.dim() == 2 and position_ids.size(0) == B:
+        freq = freqs_cis[position_ids]                         # (B, T, rd//2, 2)
+    else:  # (T,) or scalar
+        freq = freqs_cis[position_ids].unsqueeze(0).expand(B, -1, -1, -1)
+
+    rot_half = rd // 2
+    cos = freq[..., 0][..., :rot_half].unsqueeze(1).to(x.dtype)  # (B,1,T,rot_half)
+    sin = freq[..., 1][..., :rot_half].unsqueeze(1).to(x.dtype)
+
+    # complex multiply
+    yr = xr * cos - xi * sin
+    yi = xr * sin + xi * cos
+    y = torch.stack((yr, yi), dim=-1).flatten(-2)                # (B,H,T,rd)
+
+    return torch.cat([y, x_pass], dim=-1)
 
-    q = (x_q * cos_) + (_rotate_half(x_q) * sin_)
-    k = (x_k * cos_) + (_rotate_half(x_k) * sin_)
-    return q, k