From b8c8da91fe1a7c77849db86b156291a399878d71 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 04:10:51 +0000 Subject: [PATCH] fix: restore RoPE functions that were lost during mHC refactor --- single_shot_inference.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/single_shot_inference.py b/single_shot_inference.py index 8c50bce8..8ca6c526 100644 --- a/single_shot_inference.py +++ b/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."""