Update modeling_neollm.py

modeling_neollm.py

@@ -5,8 +5,9 @@ full attention augmented with optional Momentum, MEA, and LUCID operators,
 Gated Attention (Qiu et al., 2025) combined with Affine-Scaled Attention
 (Bae et al., 2026), an optional Leviathan continuous token embedding
 generator, an optional Leviathan-JTok-M token-indexed modulation module,
-optional Spelling Bee Embeddings (Rabe et al., 2026), and optional Context
-Re-Positioning (Li et al., 2026).
+optional Spelling Bee Embeddings (Rabe et al., 2026), optional Context
+Re-Positioning (Li et al., 2026), and optional REPO-GRAPE contextual group
+positioning (Li et al., 2026 + Zhang et al., 2026).
 
 Attention stack (orthogonal, all active simultaneously when enabled):
     1. Gated Attention (use_gated_attention implicit via q_proj gate chunk):
@@ -55,6 +56,8 @@ References:
         Embeddings for Language Modeling." arXiv:2601.18030.
     Context Re-Positioning: Li, Zhao, Cai & Sproat (2026). "REPO: Language
         Models with Context Re-Positioning." arXiv:2512.14391.
+    GRAPE: Zhang et al. (2026). "Group Representational Position Encoding."
+        arXiv:2512.07805 / ICLR 2026.
 """
 
 import math
@@ -243,6 +246,28 @@ class REPOAnalysis:
     r_repr:    Optional[torch.Tensor] = None  # r [B, S, d_p] — shared repr
 
 
+
+@dataclass
+class RepoGrapeAnalysis:
+    """
+    Internals of the REPO-GRAPE positional operator.
+    Populated only when config.use_repo_grape=True and the current layer is
+    at or above repo_start_layer.
+
+    Minimal operator kept from the proposal:
+
+        u_i^(h) = z_i^(h)
+        theta_{h,r} = inv_freq_r * exp(s_{h,r})
+        logit_ij += q_i^T exp((u_j-u_i) Omega_h) k_j / sqrt(d)
+
+    The implementation uses canonical commuting rotary planes, learns the
+    per-head/per-plane angular spectrum, rotates Q/K once, and then calls the
+    existing attention backend.
+    """
+    positions: Optional[torch.Tensor] = None  # [B,H,S] u_i = z_i
+    freq:      Optional[torch.Tensor] = None  # [H,rotary_dim/2] theta_{h,r}
+
+
 @dataclass
 class AttentionAnalysis:
     """
@@ -325,6 +350,9 @@ class AttentionAnalysis:
     # ── REPO position prediction (conditional on use_repo) ────────────
     repo: Optional["REPOAnalysis"] = None  # None when layer_idx < repo_start_layer
 
+    # ── REPO-GRAPE contextual group positioning (conditional on use_repo_grape)
+    repo_grape: Optional["RepoGrapeAnalysis"] = None
+
     # ── IHA pseudo-head mixing (conditional on use_iha) ──────────────
     iha: Optional["IHAAnalysis"] = None   # None when use_iha=False
 
@@ -2272,13 +2300,201 @@ def _apply_repo_rope(
     return q_embed, k_embed
 
 
+
+class RepoGrapePositioning(nn.Module):
+    r"""
+    Minimal REPO-GRAPE-M positional operator.
+
+    This module intentionally keeps only the specific operator from the
+    proposal:
+
+        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:
+
+              theta_{h,r}=inv_freq_r*exp(s_{h,r}).
+
+    Therefore
+
+        (G_h(u_i)q_i)^T(G_h(u_j)k_j)
+        = q_i^T G_h(u_j-u_i) k_j,
+
+    preserving the relative GRAPE law while staying compatible with the
+    existing attention backend.
+    """
+
+    def __init__(
+        self,
+        config: NeoLLMConfig,
+        layer_idx: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+    ):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_attention_heads = int(num_attention_heads)
+        self.num_key_value_heads = int(num_key_value_heads)
+        self.head_dim = int(head_dim)
+        self.max_rot_half = self.head_dim // 2
+
+        # The single REPO-GRAPE parameter: a learned log-scale over the
+        # base RoPE/GRAPE frequencies, per query head and rotary plane.
+        self.freq_log_scale = nn.Parameter(
+            torch.zeros(self.num_attention_heads, self.max_rot_half, dtype=torch.float32),
+            requires_grad=True,
+        )
+
+    def _apply(self, fn):
+        # Keep the small spectral parameters in fp32 under model.to(dtype).
+        super()._apply(fn)
+        p = self.freq_log_scale
+        p.data = p.data.float()
+        if p.grad is not None:
+            p.grad.data = p.grad.data.float()
+        return self
+
+    def transform_positions(
+        self,
+        z: torch.Tensor,
+        analysis: Optional[RepoGrapeAnalysis] = None,
+    ) -> torch.Tensor:
+        """Use the REPO coordinate directly: u_i^(h)=z_i^(h)."""
+        if analysis is not None:
+            analysis.positions = z.detach()
+        return z
+
+    def _query_freq(
+        self,
+        inv_freq: torch.Tensor,
+        H_repo: int,
+        analysis: Optional[RepoGrapeAnalysis],
+    ) -> torch.Tensor:
+        rot_half = int(inv_freq.shape[0])
+        if rot_half > self.max_rot_half:
+            raise ValueError(
+                f"REPO-GRAPE rotary_dim/2={rot_half} exceeds head_dim/2={self.max_rot_half}."
+            )
+        if H_repo != self.num_attention_heads:
+            raise ValueError(
+                f"REPO-GRAPE expected H_repo={self.num_attention_heads}, got {H_repo}."
+            )
+
+        base = inv_freq.float().view(1, rot_half)
+        log_scale = self.freq_log_scale[:, :rot_half].float()
+        freq = base * torch.exp(log_scale)
+        freq = freq.to(device=inv_freq.device)
+
+        if analysis is not None:
+            analysis.freq = freq.detach()
+        return freq
+
+    def apply_multiplicative(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        z: torch.Tensor,
+        inv_freq: torch.Tensor,
+        attention_scaling: float,
+        analysis: Optional[RepoGrapeAnalysis] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply REPO-GRAPE-M rotation to Q/K using contextual REPO coordinates.
+
+        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
+        """
+        B, H_repo, S = z.shape
+        H_q_eff = q.shape[1]
+        H_k_eff = k.shape[1]
+
+        if H_q_eff % H_repo != 0:
+            raise ValueError(
+                f"REPO-GRAPE expected q heads divisible by z heads, "
+                f"got H_q={H_q_eff} vs H_repo={H_repo}."
+            )
+        P = H_q_eff // H_repo
+
+        if H_k_eff % max(P, 1) != 0:
+            raise ValueError(
+                f"REPO-GRAPE expected k heads divisible by pseudo factor P, "
+                f"got H_k={H_k_eff} vs P={P}."
+            )
+
+        z_q_struct = z.unsqueeze(2).expand(B, H_repo, P, S)  # [B,H_q_base,P,S]
+        H_k_base = H_k_eff // max(P, 1)
+        if H_repo % H_k_base != 0:
+            raise ValueError(
+                f"REPO-GRAPE expected original query heads divisible by key groups, "
+                f"got H_repo={H_repo} vs H_k_base={H_k_base}."
+            )
+
+        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
+
+        freq_repo = self._query_freq(inv_freq, H_repo, analysis=analysis)
+        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(
+            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)
+
+        q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+        k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+        q_embed = torch.cat(
+            [(q_rot * cos_q) + (rotate_half(q_rot) * sin_q), q_pass],
+            dim=-1,
+        )
+        k_embed = torch.cat(
+            [(k_rot * cos_k) + (rotate_half(k_rot) * sin_k), k_pass],
+            dim=-1,
+        )
+        return q_embed, k_embed
+
+
+
 class NeoLLMAttention(nn.Module):
     """
     Full attention with FANformer, SeeDNorm, ResFormer, Learnable Multipliers,
     optional Momentum, MEA head-level composition, optional LUCID preconditioning,
     optional Affine-Scaled Attention, optional Exclusive Self Attention,
-    optional Directional Routing (Taylor, 2026), and optional Context
-    Re-Positioning (Li et al., 2026).
+    optional Directional Routing (Taylor, 2026), optional Context
+    Re-Positioning (Li et al., 2026), and optional REPO-GRAPE contextual group
+    positioning (Li et al., 2026 + Zhang et al., 2026).
 
     Directional Routing inserts at position C — post-XSA, pre-reshape — where
     the output is already normalized (MEAHeadSeeDNorm) and has auto-position
@@ -2286,8 +2502,8 @@ class NeoLLMAttention(nn.Module):
     orthogonal to the self-position already cleaned by XSA.
 
     Pipeline (all active simultaneously when enabled):
-        FANLayer → q_proj(gate) → q_norm/k_norm → REPO/RoPE → Momentum
-        → MEA(K,V) → LUCID(V) → v_ref → Affine-Scaled SDPA
+        FANLayer → q_proj(gate) → q_norm/k_norm → REPO/RoPE or REPO-GRAPE-M
+        → Momentum → MEA(K,V) → LUCID(V) → v_ref → Affine-Scaled SDPA
         → MEAHeadSeeDNorm → XSA → Directional Routing → reshape
         → o_proj · sigmoid(gate) → dropout
 
@@ -2495,15 +2711,19 @@ class NeoLLMAttention(nn.Module):
             self.direction_vecs   = None
             self.direction_router = None
 
-        # ── Context Re-Positioning (Li et al., 2026) ─────────────────────
-        # Active for layers at or above repo_start_layer only.
-        # Layers below repo_start_layer use standard integer RoPE positions.
-        # inv_freq is accessed from the model's rotary_emb at forward time;
-        # stored here as a non-persistent buffer reference set by NeoLLMModel.
+        # ── Context Re-Positioning / REPO-GRAPE ─────────────────────────────
+        # REPO-GRAPE reuses the REPO coordinate module f_phi and replaces only
+        # the positional action applied to Q/K. Enabling use_repo_grape=True
+        # therefore activates the REPO coordinate path even when use_repo=False.
+        _repo_requested = bool(getattr(config, "use_repo", False))
+        _repo_grape_requested = bool(getattr(config, "use_repo_grape", False))
+        _repo_start_layer = getattr(config, "repo_start_layer", config.num_hidden_layers // 3)
         self.use_repo = (
-            getattr(config, "use_repo", False)
-            and layer_idx >= getattr(config, "repo_start_layer", config.num_hidden_layers // 3)
+            (_repo_requested or _repo_grape_requested)
+            and layer_idx >= _repo_start_layer
         )
+        self.use_repo_grape = _repo_grape_requested and self.use_repo
+
         if self.use_repo:
             _d_p = getattr(config, "repo_d_p", config.hidden_size // 8)
             self.repo_module = REPOModule(
@@ -2514,6 +2734,17 @@ class NeoLLMAttention(nn.Module):
         else:
             self.repo_module = None
 
+        if self.use_repo_grape:
+            self.repo_grape = RepoGrapePositioning(
+                config=config,
+                layer_idx=layer_idx,
+                num_attention_heads=config.num_attention_heads,
+                num_key_value_heads=config.num_key_value_heads,
+                head_dim=self.head_dim,
+            )
+        else:
+            self.repo_grape = None
+
         # ── Interleaved Head Attention (Duvvuri et al., 2026) ─────────────────
         # Implementa cross-head mixing antes de RoPE: cada pseudo-head h,p
         # es combinación lineal aprendida de los H input-heads originales.
@@ -3294,24 +3525,34 @@ class NeoLLMAttention(nn.Module):
             # no meta-tensor issue on lm_eval / to(device) paths.
             # (Li et al., 2026, §3.2 — Eq. 6–7)
             repo_a = attn_analysis.repo if attn_analysis is not None else None
+            repo_grape_a = attn_analysis.repo_grape if attn_analysis is not None else None
             z = self.repo_module(hidden_states, repo_analysis=repo_a)  # [B, H, S]
+            if self.use_repo_grape:
+                z = self.repo_grape.transform_positions(z, analysis=repo_grape_a)
             inv_freq, attn_scaling = repo_rope_args
             if _iha_seq_expand:
-                # IHA seq-expand + REPO: each real position s has P virtual slots.
-                # Assign continuous REPO position  z[b,h,s]·P + p  to slot (s,p),
-                # scaling z by P to create room for P sub-positions while preserving
-                # REPO's relative ordering across tokens.
-                # _apply_repo_rope handles the flat S·P sequence correctly when
-                # H_q_eff == H_repo (P_derived=1 inside the function).
+                # IHA seq-expand + REPO/REPO-GRAPE: each real position s has P virtual slots.
+                # Assign continuous position z[b,h,s]·P + p to slot (s,p), scaling z by P
+                # to create room for P sub-positions while preserving relative ordering.
                 P = self.iha_P
                 p_offsets = torch.arange(P, device=z.device, dtype=z.dtype)
                 # z[B,H,S] → [B,H,S,1]*P + [1,1,1,P] → [B,H,S,P] → [B,H,S·P]
                 z_expanded = (z.unsqueeze(-1) * P + p_offsets.view(1, 1, 1, P)
                               ).reshape(z.shape[0], z.shape[1], _S_orig * P)
-                q, k = _apply_repo_rope(q, k, z_expanded, inv_freq, attn_scaling)
+                if self.use_repo_grape:
+                    q, k = self.repo_grape.apply_multiplicative(
+                        q, k, z_expanded, inv_freq, attn_scaling, analysis=repo_grape_a
+                    )
+                else:
+                    q, k = _apply_repo_rope(q, k, z_expanded, inv_freq, attn_scaling)
             else:
                 # Standard REPO path (P=1 or no IHA): z covers original S positions.
-                q, k = _apply_repo_rope(q, k, z, inv_freq, attn_scaling)
+                if self.use_repo_grape:
+                    q, k = self.repo_grape.apply_multiplicative(
+                        q, k, z, inv_freq, attn_scaling, analysis=repo_grape_a
+                    )
+                else:
+                    q, k = _apply_repo_rope(q, k, z, inv_freq, attn_scaling)
         elif _iha_seq_expand:
             # IHA seq-expand + standard integer RoPE:
             # Build interleaved positions  pos_iha[s·P+p] = position_ids[s]·P + p
@@ -5115,18 +5356,22 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         Called once per layer per forward when analysis is active.
         """
         cfg = self.config
+        _repo_requested = bool(getattr(cfg, "use_repo", False))
+        _repo_grape_requested = bool(getattr(cfg, "use_repo_grape", False))
         _repo_active = (
-            getattr(cfg, "use_repo", False)
+            (_repo_requested or _repo_grape_requested)
             and layer_idx >= getattr(cfg, "repo_start_layer", cfg.num_hidden_layers // 3)
         )
+        _repo_grape_active = _repo_grape_requested and _repo_active
         _versatile = getattr(cfg, "use_versatile_ffn", False)
         return LayerAnalysis(
             seednorm_pre_attn  = SeeDNormAnalysis(),
             seednorm_post_attn = SeeDNormAnalysis(),
             attention          = AttentionAnalysis(
-                fan      = FANAnalysis(),
-                hadamard = HadamardAnalysis() if getattr(cfg, "use_hadamard_o_proj", False) else None,
-                repo     = REPOAnalysis() if _repo_active else None,
+                fan        = FANAnalysis(),
+                hadamard   = HadamardAnalysis() if getattr(cfg, "use_hadamard_o_proj", False) else None,
+                repo       = REPOAnalysis() if _repo_active else None,
+                repo_grape = RepoGrapeAnalysis() if _repo_grape_active else None,
             ),
             mlp                = MLPAnalysis(
                 fan      = FANAnalysis() if not _versatile else None,
@@ -5250,9 +5495,10 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             getattr(self.config, "use_iha", False)
             and int(getattr(self.config, "iha_num_pseudo_heads", 1)) > 1
         )
+        _repo_grape_needs_inv_freq = bool(getattr(self.config, "use_repo_grape", False))
         repo_rope_args = (
             (self.rotary_emb.inv_freq, self.rotary_emb.attention_scaling)
-            if (getattr(self.config, "use_repo", False) or _iha_needs_inv_freq)
+            if (getattr(self.config, "use_repo", False) or _repo_grape_needs_inv_freq or _iha_needs_inv_freq)
             else None
         )