Update modeling_neollm.py

modeling_neollm.py

@@ -2245,23 +2245,54 @@ class RepoGrapePositioning(nn.Module):
     r"""
     Minimal REPO-GRAPE-M positional operator.
 
-    This module intentionally keeps only the specific operator from the
-    proposal:
+    This module intentionally keeps only the specific multiplicative path:
 
         1. REPO predicts contextual coordinates z_i^(h)=f_phi^(h)(h_i).
         2. REPO-GRAPE uses u_i^(h)=z_i^(h), with no extra position mode.
         3. GRAPE-M applies canonical commuting SO(d) rotary planes to Q/K.
-        4. Each query head learns its own angular spectrum:
+        4. Each query head learns its own angular spectrum
 
-              theta_{h,r}=inv_freq_r*exp(s_{h,r}).
+              theta_{h,r}=inv_freq_r*exp(s_{h,r}),
 
-    Therefore
+           and therefore its own phase
+
+              phi_{b,h,i,r}=u_{b,h,i}*theta_{h,r}.
+
+    With one key per query head, the multiplicative action obeys the exact
+    relative law
 
         (G_h(u_i)q_i)^T(G_h(u_j)k_j)
-        = q_i^T G_h(u_j-u_i) k_j,
+        = q_i^T G_h(u_j-u_i) k_j.
+
+    GQA introduces a separate geometric restriction: several query heads
+    share one key, although their phases phi_{b,h,i,r} can differ.  A shared
+    key cannot realize every head-specific rotation simultaneously.  For a
+    KV group H_m and rotary plane r, this implementation therefore projects
+    the requested rotations onto the single optimal shared SO(2) rotation:
+
+        Z_{b,m,i,r} = (1/g) * sum_{h in H_m} exp(i*phi_{b,h,i,r})
+                    = C_{b,m,i,r} + i*S_{b,m,i,r},
+
+        psi*_{b,m,i,r} = arg Z_{b,m,i,r},
+
+    or, equivalently when |Z| > 0,
 
-    preserving the relative GRAPE law while staying compatible with the
-    existing attention backend.
+        cos(psi*) = C/|Z|,
+        sin(psi*) = S/|Z|.
+
+    This is the global minimizer of
+
+        sum_{h in H_m} ||R(phi_{b,h,i,r})k - R(psi)k||_2^2,
+
+    so it replaces the previous arithmetic approximation
+
+        mean(u_h) * mean(theta_h),
+
+    which neither averages the effective phase u_h*theta_h nor respects its
+    2*pi periodicity.  If the circular resultant is numerically zero, the
+    objective does not identify a preferred shared direction; the code uses
+    the identity rotation as a deterministic finite convention.  Under IHA,
+    the projection is computed independently inside every pseudo-slot.
     """
 
     def __init__(
@@ -2335,15 +2366,38 @@ class RepoGrapePositioning(nn.Module):
         inv_freq: torch.Tensor,
         attention_scaling: float,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
+        r"""
         Apply REPO-GRAPE-M rotation to Q/K using contextual REPO coordinates.
 
+        Query heads keep their individual phases
+
+            phi_{b,h,i,r} = z_{b,h,i} * theta_{h,r}.
+
+        For GQA, each KV head m serves a group H_m of query heads.  Its shared
+        rotation is the circular projection
+
+            C = mean_{h in H_m} cos(phi_h),
+            S = mean_{h in H_m} sin(phi_h),
+            (cos psi*, sin psi*) = (C, S) / sqrt(C^2 + S^2).
+
+        This is equivalent to psi*=arg(mean_h exp(i*phi_h)) and minimizes the
+        total squared discrepancy between the head-specific rotated keys and
+        one shared rotated key.  The reduction is performed independently for
+        every batch element, current sequence index, rotary plane and KV group.
+        Because seq-expand IHA stores each pseudo-slot at a different sequence
+        index, its pseudo-slots remain independent automatically.  No extra
+        mode or configuration flag is introduced.
+
         Args:
             q:                 [B, H_q_eff, S, head_dim]
             k:                 [B, H_k_eff, S, head_dim]
             z:                 [B, H_q_base, S]
             inv_freq:          [rotary_dim/2]
             attention_scaling: scaling from NeoLLMRotaryEmbedding
+
+        Returns:
+            Rotated ``(q, k)`` tensors with the same shapes and dtypes as the
+            inputs. Channels beyond ``rotary_dim`` pass through unchanged.
         """
         B, H_repo, S = z.shape
         H_q_eff = q.shape[1]
@@ -2372,11 +2426,6 @@ class RepoGrapePositioning(nn.Module):
 
         q_per_k_base = H_repo // H_k_base
         z_q = z_q_struct.reshape(B, H_q_eff, S)
-        z_k = (
-            z_q_struct.view(B, H_k_base, q_per_k_base, P, S)
-            .mean(dim=2)
-            .reshape(B, H_k_eff, S)
-        )
 
         rot_half = int(inv_freq.shape[0])
         rotary_dim = rot_half * 2
@@ -2384,26 +2433,66 @@ class RepoGrapePositioning(nn.Module):
         freq_repo = self._query_freq(inv_freq, H_repo)
         freq_q_struct = freq_repo.unsqueeze(1).expand(H_repo, P, rot_half)
         freq_q = freq_q_struct.reshape(H_q_eff, rot_half)
-        freq_k = (
-            freq_q_struct.view(H_k_base, q_per_k_base, P, rot_half)
-            .mean(dim=1)
-            .reshape(H_k_eff, rot_half)
-        )
 
-        freqs_q = z_q.float().unsqueeze(-1) * freq_q.to(device=q.device).view(
+        # Individual query phases are the fundamental quantity.  They are
+        # calculated in fp32 before trigonometric evaluation, as in the
+        # original REPO-GRAPE path.
+        phase_q = z_q.float().unsqueeze(-1) * freq_q.to(device=q.device).view(
             1, H_q_eff, 1, rot_half
         )
-        freqs_k = z_k.float().unsqueeze(-1) * freq_k.to(device=k.device).view(
-            1, H_k_eff, 1, rot_half
-        )
-
-        emb_q = torch.cat([freqs_q, freqs_q], dim=-1)
-        emb_k = torch.cat([freqs_k, freqs_k], dim=-1)
-
-        cos_q = (emb_q.cos() * attention_scaling).to(q.dtype)
-        sin_q = (emb_q.sin() * attention_scaling).to(q.dtype)
-        cos_k = (emb_k.cos() * attention_scaling).to(k.dtype)
-        sin_k = (emb_k.sin() * attention_scaling).to(k.dtype)
+        cos_q_half = phase_q.cos()
+        sin_q_half = phase_q.sin()
+
+        # GQA circular projection.  Restore the logical head factorization
+        # (KV group, query-within-group, retained head-expansion slot) and
+        # average only over query-within-group.  In the current seq-expand IHA
+        # path pseudo-slots live on the sequence axis, which is never reduced.
+        # Averaging cos/sin and normalizing the resultant is exactly equivalent
+        # to arg(mean(exp(i*phase))).
+        cos_q_grouped = cos_q_half.reshape(
+            B, H_k_base, q_per_k_base, P, S, rot_half
+        )
+        sin_q_grouped = sin_q_half.reshape(
+            B, H_k_base, q_per_k_base, P, S, rot_half
+        )
+        cos_k_half = cos_q_grouped.mean(dim=2)
+        sin_k_half = sin_q_grouped.mean(dim=2)
+
+        resultant_sq = cos_k_half.square() + sin_k_half.square()
+        well_defined = resultant_sq > torch.finfo(resultant_sq.dtype).eps
+        safe_resultant_sq = torch.where(
+            well_defined, resultant_sq, torch.ones_like(resultant_sq)
+        )
+        inv_resultant = torch.rsqrt(safe_resultant_sq)
+
+        # If the resultant vanishes, every shared angle has the same objective
+        # value.  Select identity (cos=1, sin=0) to keep forward/backward finite
+        # and deterministic; this is a numerical convention, not a new mode.
+        cos_k_half = torch.where(
+            well_defined,
+            cos_k_half * inv_resultant,
+            torch.ones_like(cos_k_half),
+        )
+        sin_k_half = torch.where(
+            well_defined,
+            sin_k_half * inv_resultant,
+            torch.zeros_like(sin_k_half),
+        )
+        cos_k_half = cos_k_half.reshape(B, H_k_eff, S, rot_half)
+        sin_k_half = sin_k_half.reshape(B, H_k_eff, S, rot_half)
+
+        cos_q = (
+            torch.cat([cos_q_half, cos_q_half], dim=-1) * attention_scaling
+        ).to(q.dtype)
+        sin_q = (
+            torch.cat([sin_q_half, sin_q_half], dim=-1) * attention_scaling
+        ).to(q.dtype)
+        cos_k = (
+            torch.cat([cos_k_half, cos_k_half], dim=-1) * attention_scaling
+        ).to(k.dtype)
+        sin_k = (
+            torch.cat([sin_k_half, sin_k_half], dim=-1) * attention_scaling
+        ).to(k.dtype)
 
         q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
         k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]