Update rope.py

fix: fallback to apply rotatory_emb

rope.py

@@ -16,36 +16,37 @@ def precompute_freqs_cis(
     freqs = torch.exp(1j * freqs)
     return torch.stack([freqs.real, freqs.imag], dim=-1)
 
-
-def apply_rotary_emb(x, freqs_cis, position_ids, num_heads, rot_dim=None, interleave=False):
-    B, H, T, D = x.shape
-    rot_half_from_freqs = freqs_cis.size(-2)
-    rd = rot_dim or (rot_half_from_freqs * 2)
-    rd = min(rd, D)
-
-    x_rot, x_pass = x[..., :rd], x[..., rd:]
-
-    # gather cos/sin for these positions
-    if torch.is_tensor(position_ids):
-        if position_ids.dim() == 2 and position_ids.size(0) == B:
-            freq = freqs_cis[position_ids]                                  # (B,T,rot_half,2)
-        elif position_ids.dim() == 1 and position_ids.size(0) == T:
-            freq = freqs_cis[position_ids].unsqueeze(0).expand(B, -1, -1, -1)
-        else:                                                                # scalar
-            pid = position_ids.view(()).long()
-            freq = freqs_cis[pid].unsqueeze(0).expand(B, T, -1, -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):
+    """
+    Applies RoPE to q, k.
+    Keeps shapes: (B, H, T, D) or (B, T, H, D). Cos/sin are cast to x.dtype.
+    """
+    # 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)
     else:
-        pid = torch.tensor(position_ids, device=x.device, dtype=torch.long)
-        freq = freqs_cis[pid].unsqueeze(0).expand(B, T, -1, -1)
-
-    rot_half = rd // 2
-    cos = freq[..., 0][..., :rot_half].unsqueeze(1)                          # (B,1,T,rot_half)
-    sin = freq[..., 1][..., :rot_half].unsqueeze(1)                          # (B,1,T,rot_half)
+        cos_, sin_ = cos, sin
 
-    x_rot = x_rot.view(B, H, T, rot_half, 2)
-    xr, xi = x_rot[..., 0], x_rot[..., 1]
-    yr = xr * cos - xi * sin
-    yi = xr * sin + xi * cos
-    y = torch.stack((yr, yi), dim=-1).flatten(-2)
+    q = (x_q * cos_) + (_rotate_half(x_q) * sin_)
+    k = (x_k * cos_) + (_rotate_half(x_k) * sin_)
+    return q, k
 
-    return torch.cat([y.to(x.dtype), x_pass], dim=-1)