fix: restore RoPE functions that were lost during mHC refactor

single_shot_inference.py

@@ -159,6 +159,46 @@ class mHCBlock:
     def post_block(self, X_l, F_out, ctx):
         return self._impl.post_block(X_l, F_out, ctx)
 
+
+# =====================================================================
+# RoPE — partial, GPT-J interleaved, last rope_dim dims
+# =====================================================================
+
+def build_rope_cache(max_pos, rope_dim, device, theta=10000.0):
+    """Build cos/sin caches for partial RoPE.
+    Returns: (cos_cache, sin_cache) each (max_pos, rope_dim//2) BF16
+    """
+    half = rope_dim // 2
+    freqs = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
+    angles = torch.outer(torch.arange(max_pos, dtype=torch.float32), freqs)
+    return torch.cos(angles).bfloat16().to(device), torch.sin(angles).bfloat16().to(device)
+
+
+def apply_rope_partial(x, positions, cos_cache, sin_cache, head_dim, rope_dim):
+    """Apply partial GPT-J interleaved RoPE to the last rope_dim dims of each head."""
+    T, n_h, hd = x.shape
+    nope = hd - rope_dim
+    cos = cos_cache[positions].unsqueeze(1)  # (T, 1, half) BF16
+    sin = sin_cache[positions].unsqueeze(1)
+    out = x.clone()
+    x_rope = x[:, :, nope:]
+    out[:, :, nope:][..., 0::2] = x_rope[..., 0::2] * cos - x_rope[..., 1::2] * sin
+    out[:, :, nope:][..., 1::2] = x_rope[..., 0::2] * sin + x_rope[..., 1::2] * cos
+    return out
+
+
+def apply_inverse_rope(o, positions, cos_cache, sin_cache, head_dim, rope_dim):
+    """Apply inverse RoPE (conjugate rotation) to attention output."""
+    T, n_h, hd = o.shape
+    nope = hd - rope_dim
+    cos = cos_cache[positions].unsqueeze(1)
+    sin = sin_cache[positions].unsqueeze(1)
+    out = o.clone()
+    o_rope = o[:, :, nope:]
+    out[:, :, nope:][..., 0::2] =  o_rope[..., 0::2] * cos + o_rope[..., 1::2] * sin
+    out[:, :, nope:][..., 1::2] = -o_rope[..., 0::2] * sin + o_rope[..., 1::2] * cos
+    return out
+
 class SimpleKVCache:
     """Per-layer KV cache for decode. Stores BF16 K,V accumulated across steps.
     MQA: 1 KV head, so cache is (1, seq_len, hd) per layer."""