Update rope.py

feat: allow support for batched

rope.py

@@ -17,32 +17,44 @@ def precompute_freqs_cis(
     return torch.stack([freqs.real, freqs.imag], dim=-1)
 
 
-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:
-    assert rot_dim == freqs_cis.shape[-2] * 2
-    assert num_heads == x.shape[1]
-
-    x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
-
-    if interleave:
-        xq_r = x_rot.float().reshape(*x_rot.shape[:-1], -1, 2)[..., 0]
-        xq_i = x_rot.float().reshape(*x_rot.shape[:-1], -1, 2)[..., 1]
+def apply_rotary_emb(x, freqs_cis, position_ids, num_heads, rot_dim=None, interleave=False):
+    """
+    x:            (B, num_heads, q_len, head_dim)
+    freqs_cis:    (max_seq, rot_half, 2)   # [..., cos/sin]
+    position_ids: (q_len,) or (B, q_len) or scalar
+    num_heads:    int (unused here; kept for API compatibility)
+    rot_dim:      optional; if None we use min(D, 2*rot_half)
+    """
+    B, H, T, D = x.shape
+    rot_half_from_freqs = freqs_cis.size(-2)                 # available rotary half
+    rd = rot_dim or (rot_half_from_freqs * 2)
+    rd = min(rd, D)                                          # don't exceed head_dim
+
+    x_rot, x_pass = x[..., :rd], x[..., rd:]                 # (B,H,T,rd), (B,H,T,D-rd)
+
+    # Gather cos/sin for each position; result (B,T,rot_half_from_freqs)
+    if torch.is_tensor(position_ids):
+        if position_ids.dim() == 2 and position_ids.size(0) == B:         # (B,T)
+            freq = freqs_cis[position_ids]                                # (B,T,rot_half,2)
+        elif position_ids.dim() == 1 and position_ids.size(0) == T:       # (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)
     else:
-        d_q = x_rot.shape[-1] // 2
-        xq_r, xq_i = x_rot[..., :d_q], x_rot[..., d_q:]
-
-    freqs_cos = freqs_cis[..., 0][position_ids, :].unsqueeze(0).unsqueeze(0)
-    freqs_sin = freqs_cis[..., 1][position_ids, :].unsqueeze(0).unsqueeze(0)
-
-    # Complex multiplication: (a + bi) * (c + di) = (ac - bd) + (ad + bc)i
-    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
-    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
-    xq_out = torch.stack((xq_out_r, xq_out_i), dim=-1).flatten(-2)
-
-    return torch.cat([xq_out.to(x.dtype), x_pass], dim=-1)
+        pid = torch.tensor(position_ids, device=x.device, dtype=torch.long)
+        freq = freqs_cis[pid].unsqueeze(0).expand(B, T, -1, -1)
+
+    # Trim freqs to rd//2 if needed
+    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)
+
+    # Split real/imag and apply rotation
+    x_rot = x_rot.view(B, H, T, rot_half, 2)
+    xr, xi = x_rot[..., 0], x_rot[..., 1]                    # (B,H,T,rot_half)
+    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.to(x.dtype), x_pass], dim=-1)