Update modeling_neollm.py

modeling_neollm.py

@@ -8,8 +8,9 @@ generator, an optional Leviathan-JTok-M token-indexed modulation module,
 optional Spelling Bee Embeddings (Rabe et al., 2026), optional Context
 Re-Positioning (Li et al., 2026), optional REPO-GRAPE contextual group
 positioning (Li et al., 2026 + Zhang et al., 2026), optional GOAT-style
-factorised attention log-priors (Litman & Guo, 2026), and optional StackMemory
-(Zhang et al., NeurIPS 2025).
+factorised attention log-priors (Litman & Guo, 2026), optional StackMemory
+(Zhang et al., NeurIPS 2025), and optional HOLA-derived episodic memory for
+the FlashAttention-2 IHA seq-expand path (Cui, 2026).
 
 Attention stack (orthogonal, all active simultaneously when enabled):
     1. Gated Attention (use_gated_attention implicit via q_proj gate chunk):
@@ -65,6 +66,8 @@ References:
         arXiv:2601.15380.
     StackMemory / STACKTRANS: Zhang et al. (2025). "Recursive Transformer:
         Boosting Reasoning Ability with State Stack." NeurIPS 2025.
+    HOLA: Cui (2026). "A Hippocampus for Linear Attention: An Exact Memory
+        for What the Recurrent State Forgets." arXiv:2607.02303.
     TWEO: Liang et al. (2026). "Transformers Without Extreme Outliers Enables
         FP8 Training And Quantization For Dummies." arXiv:2511.23225.
 """
@@ -2571,6 +2574,7 @@ class NeoLLMAttention(nn.Module):
     Pipeline (all active simultaneously when enabled):
         FANLayer → q_proj(gate) → q_norm/k_norm → REPO/RoPE or REPO-GRAPE-M
         → Momentum → MEA(K,V) → LUCID(V) → v_ref → optional GOAT log-prior
+        → optional HOLA memory-augmented FA2 read (IHA seq-expand only)
         → Affine-Scaled SDPA
         → MEAHeadRMSNorm → XSA → Directional Routing → reshape
         → o_proj · sigmoid(gate) → dropout
@@ -2727,6 +2731,51 @@ class NeoLLMAttention(nn.Module):
         self.use_xsa = getattr(config, "use_xsa", False)
         self.xsa_eps = float(getattr(config, "xsa_eps", 1e-6))
 
+        # ── HOLA-derived episodic memory (Cui, 2026) ─────────────────────────
+        # This flag is deliberately scoped to the direct FlashAttention-2
+        # seq-expand branch used by paper-correct IHA (P>1).  It does not add a
+        # second attention output.  Instead, each chunk performs a single FA2 call
+        # over [selected episodic KV, current chunk KV].  Memory logits live in an
+        # augmented subspace so ordinary keys and memory keys can use different
+        # read geometries while still sharing one softmax.
+        #
+        # Exact HOLA writes use beta*||e|| from a recurrent delta state.  This
+        # Transformer branch has no such state and stock FA2 does not expose the
+        # pre-dropout accumulators needed for exact leave-one-out residuals.
+        # Therefore the FA2 path below uses a compile-friendly value-innovation
+        # proxy for episode eviction.  The read path itself is exact within the
+        # single augmented FA2 softmax.
+        self.use_hola_memory = bool(getattr(config, "use_hola_memory", False))
+        self.hola_memory_size = int(getattr(config, "hola_memory_size", 64))
+        self.hola_chunk_size = int(getattr(config, "hola_chunk_size", 256))
+        self.hola_score_eps = float(getattr(config, "hola_score_eps", 1e-6))
+
+        if self.use_hola_memory:
+            self.hola_q_gamma = nn.Parameter(
+                torch.full(
+                    (config.num_attention_heads, self.head_dim),
+                    float(getattr(config, "hola_rms_gamma_init", 1.0)),
+                )
+            )
+            self.hola_k_gamma = nn.Parameter(
+                torch.full(
+                    (config.num_key_value_heads, self.head_dim),
+                    float(getattr(config, "hola_rms_gamma_init", 1.0)),
+                )
+            )
+            # Tied at KV-head granularity because memory keys are stored once per
+            # KV component. Query heads that share a KV component share this prior.
+            self.hola_logit_bias = nn.Parameter(
+                torch.full(
+                    (config.num_key_value_heads,),
+                    float(getattr(config, "hola_logit_bias_init", -2.0)),
+                )
+            )
+        else:
+            self.hola_q_gamma = None
+            self.hola_k_gamma = None
+            self.hola_logit_bias = None
+
         if self.use_affine_scaled_attention:
             self.alpha_proj = nn.Linear(
                 config.hidden_size, config.num_attention_heads, bias=False
@@ -3248,6 +3297,207 @@ class NeoLLMAttention(nn.Module):
         q_rot, k_rot = apply_rotary_pos_emb(q, k, cos, sin)
         return q_rot, k_rot
 
+
+    def _hola_rms_gamma(
+        self,
+        x: torch.Tensor,
+        gamma: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        HOLA cache-read normalization used only for the memory subspace.
+
+        Args:
+            x:     [B, T, H, d] tensor in flash-attn layout.
+            gamma: [H, d] learned RMSNorm-γ scale.
+        """
+        rms = x.float().pow(2).mean(dim=-1, keepdim=True).add(self.hola_score_eps).rsqrt()
+        return (x.float() * rms * gamma.view(1, 1, gamma.shape[0], gamma.shape[1]).float()).to(x.dtype)
+
+    def _hola_value_innovation_scores(
+        self,
+        v: torch.Tensor,
+        gate_states: Optional[torch.Tensor],
+        S_orig: int,
+    ) -> torch.Tensor:
+        """
+        FA2-compatible HOLA-derived eviction proxy.
+
+        Full HOLA writes by beta*||e|| where e is the committed residual of a
+        recurrent delta state.  This Transformer/FA2 path intentionally avoids a
+        second attention pass and stock FA2 does not return the pre-dropout
+        leave-one-out accumulators needed for the exact score.  We therefore use
+        a causal value-innovation proxy:
+
+            score_t = || V_bundle(t) - V_bundle(t-1) ||_rms * gate_strength_t
+
+        The score is computed per *real* token and aggregates all pseudo-slots,
+        so TopK stores complete IHA episodes rather than isolated virtual slots.
+        """
+        B, H_kv, SP, d_head = v.shape
+        P = self.iha_P
+        v_tok = v.reshape(B, H_kv, S_orig, P, d_head).float()
+        prev = torch.cat([torch.zeros_like(v_tok[:, :, :1]), v_tok[:, :, :-1]], dim=2)
+        delta = v_tok - prev
+        scores = delta.pow(2).mean(dim=(1, 3, 4)).add(self.hola_score_eps).sqrt()  # [B, S]
+
+        if gate_states is not None:
+            gate = torch.sigmoid(gate_states).view(B, S_orig, self.config.num_attention_heads, d_head)
+            gate_strength = gate.float().pow(2).mean(dim=(2, 3)).add(self.hola_score_eps).sqrt()
+            scores = scores * gate_strength
+        return scores
+
+    def _hola_augmented_fa_call(
+        self,
+        q_cur: torch.Tensor,
+        k_cur: torch.Tensor,
+        v_cur: torch.Tensor,
+        k_mem: Optional[torch.Tensor],
+        v_mem: Optional[torch.Tensor],
+        dropout: float,
+        window_left: int,
+    ) -> torch.Tensor:
+        """
+        Single-FA2 HOLA read using augmented Q/K/V subspaces.
+
+        Ordinary keys use the first d channels:
+            <sqrt(D/d) q, k> / sqrt(D) = <q,k> / sqrt(d)
+
+        Memory keys use the second d channels plus one constant-bias channel:
+            <sqrt(D/d) q_H, k_H> / sqrt(D) + b
+
+        Values are padded with zeros and sliced back to d after FA2.  The result
+        is one softmax over [memory, current chunk], not a separate cache
+        attention output.
+        """
+        flash_attn_fn = self._iha_flash_attn_func
+        B, Tq, H_q, d_head = q_cur.shape
+        H_kv = k_cur.shape[2]
+        base_aug = 2 * d_head + 1
+        d_aug = ((base_aug + 7) // 8) * 8
+        pad_tail = d_aug - base_aug
+        scale = math.sqrt(float(d_aug) / float(d_head))
+        softmax_scale = 1.0 / math.sqrt(float(d_aug))
+
+        q_h = self._hola_rms_gamma(q_cur, self.hola_q_gamma)
+        one_q = torch.ones(B, Tq, H_q, 1, device=q_cur.device, dtype=q_cur.dtype)
+        q_parts = [q_cur * scale, q_h * scale, one_q]
+        if pad_tail > 0:
+            q_parts.append(torch.zeros(B, Tq, H_q, pad_tail, device=q_cur.device, dtype=q_cur.dtype))
+        q_aug = torch.cat(q_parts, dim=-1).contiguous()
+
+        zeros_cur = torch.zeros_like(k_cur)
+        zero_bias_cur = torch.zeros(B, k_cur.shape[1], H_kv, 1, device=k_cur.device, dtype=k_cur.dtype)
+        k_cur_parts = [k_cur, zeros_cur, zero_bias_cur]
+        if pad_tail > 0:
+            k_cur_parts.append(torch.zeros(B, k_cur.shape[1], H_kv, pad_tail, device=k_cur.device, dtype=k_cur.dtype))
+        k_cur_aug = torch.cat(k_cur_parts, dim=-1)
+
+        v_cur_pad = torch.zeros(B, v_cur.shape[1], H_kv, d_aug - d_head, device=v_cur.device, dtype=v_cur.dtype)
+        v_cur_aug = torch.cat([v_cur, v_cur_pad], dim=-1)
+
+        if k_mem is not None and k_mem.shape[1] > 0:
+            k_mem_h = self._hola_rms_gamma(k_mem, self.hola_k_gamma)
+            zeros_mem = torch.zeros_like(k_mem)
+            bias = (self.hola_logit_bias.view(1, 1, H_kv, 1).to(k_mem.dtype) * math.sqrt(float(d_aug)))
+            bias = bias.expand(B, k_mem.shape[1], H_kv, 1)
+            k_mem_parts = [zeros_mem, k_mem_h, bias]
+            if pad_tail > 0:
+                k_mem_parts.append(torch.zeros(B, k_mem.shape[1], H_kv, pad_tail, device=k_mem.device, dtype=k_mem.dtype))
+            k_aug = torch.cat([torch.cat(k_mem_parts, dim=-1), k_cur_aug], dim=1).contiguous()
+
+            v_mem_pad = torch.zeros(B, v_mem.shape[1], H_kv, d_aug - d_head, device=v_mem.device, dtype=v_mem.dtype)
+            v_mem_aug = torch.cat([v_mem, v_mem_pad], dim=-1)
+            v_aug = torch.cat([v_mem_aug, v_cur_aug], dim=1).contiguous()
+            mem_len = k_mem.shape[1]
+        else:
+            k_aug = k_cur_aug.contiguous()
+            v_aug = v_cur_aug.contiguous()
+            mem_len = 0
+
+        # When local IHA is active, memory tokens are placed immediately before
+        # the current chunk.  Enlarging the left window by mem_len makes all
+        # selected episodes visible while preserving causal masking inside the
+        # current chunk.  For global layers window_left is -1.
+        if window_left >= 0:
+            window = (mem_len + window_left, 0)
+        else:
+            window = (-1, -1)
+
+        out_aug = flash_attn_fn(
+            q_aug,
+            k_aug,
+            v_aug,
+            dropout_p=dropout if self.training else 0.0,
+            softmax_scale=softmax_scale,
+            causal=True,
+            window_size=window,
+        )
+        return out_aug[..., :d_head].contiguous()
+
+    def _iha_hola_seq_expand_flash_forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        dropout: float,
+        S_orig: int,
+        gate_states: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        HOLA-derived episodic memory for the dense IHA seq-expand FA2 path.
+
+        The sequence is processed in original-token chunks.  At the beginning of
+        a chunk, a fixed-size TopK memory is selected from strictly previous
+        tokens using the value-innovation proxy.  The chunk then performs a
+        single FlashAttention-2 call over [memory KV, chunk KV].  This preserves
+        causal semantics, keeps the memory frozen during the chunk, and avoids a
+        second attention output.
+        """
+        P = self.iha_P
+        B, H_q, SP, d_head = q.shape
+        H_kv = k.shape[1]
+        q_fa_all = q.transpose(1, 2).contiguous()  # [B, S*P, H_q, d]
+        k_fa_all = k.transpose(1, 2).contiguous()  # [B, S*P, H_kv, d]
+        v_fa_all = v.transpose(1, 2).contiguous()  # [B, S*P, H_kv, d]
+
+        scores = self._hola_value_innovation_scores(v, gate_states, S_orig)  # [B, S]
+        chunk = max(1, int(self.hola_chunk_size))
+        mem_k = max(0, int(self.hola_memory_size))
+        window_left = self._resolve_iha_window(S_orig) * P if self.iha_is_local else -1
+
+        outputs = []
+        for start in range(0, S_orig, chunk):
+            end = min(start + chunk, S_orig)
+            q_cur = q_fa_all[:, start * P : end * P]
+            k_cur = k_fa_all[:, start * P : end * P]
+            v_cur = v_fa_all[:, start * P : end * P]
+
+            if mem_k > 0 and start > 0:
+                k_select = min(mem_k, start)
+                _, token_idx = torch.topk(scores[:, :start], k=k_select, dim=1)
+                p_offsets = torch.arange(P, device=q.device, dtype=token_idx.dtype).view(1, 1, P)
+                virtual_idx = (token_idx.unsqueeze(-1) * P + p_offsets).reshape(B, k_select * P)
+                gather_idx = virtual_idx[:, :, None, None].expand(B, k_select * P, H_kv, d_head)
+                k_mem = torch.gather(k_fa_all, dim=1, index=gather_idx).contiguous()
+                v_mem = torch.gather(v_fa_all, dim=1, index=gather_idx).contiguous()
+            else:
+                k_mem = None
+                v_mem = None
+
+            outputs.append(
+                self._hola_augmented_fa_call(
+                    q_cur=q_cur,
+                    k_cur=k_cur,
+                    v_cur=v_cur,
+                    k_mem=k_mem,
+                    v_mem=v_mem,
+                    dropout=dropout,
+                    window_left=window_left,
+                )
+            )
+
+        return torch.cat(outputs, dim=1).contiguous()
+
     def _iha_seq_expand_flash_forward(
         self,
         q: torch.Tensor,
@@ -3260,6 +3510,7 @@ class NeoLLMAttention(nn.Module):
         S_orig: int,
         position_ids: Optional[torch.Tensor],
         attention_mask: Optional[torch.Tensor] = None,
+        gate_states: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, None]:
         """
         Flash attention for the IHA seq-expand path (P>1).
@@ -3308,6 +3559,8 @@ class NeoLLMAttention(nn.Module):
         dropout    : attention dropout probability
         S_orig     : original (pre-expand) sequence length S
         position_ids: [B, S] or None — for packing detection + cu_seqlens
+        gate_states : [B, S, H_q*d] or None — optional Gated-Attention gate
+            logits used only by the HOLA-derived eviction proxy.
 
         Returns
         -------
@@ -3408,15 +3661,30 @@ class NeoLLMAttention(nn.Module):
 
         # ── Dense (non-packed) path ───────────────────────────────────────────
         else:
-            flash_out = flash_attn_fn(
-                q_fa,
-                k_fa,
-                v_fa,
-                dropout_p=dropout if self.training else 0.0,
-                softmax_scale=self.scaling,
-                causal=True,
-                window_size=window,
-            )  # [B, S·P, H_q, d]
+            if self.use_hola_memory and self.iha_P > 1 and not use_affine:
+                # HOLA-derived memory is intentionally restricted to the dense
+                # direct-FA2 seq-expand branch.  Packed varlen and the affine
+                # inline decomposition keep the original baseline path until a
+                # dedicated fused kernel exposes the exact pre-dropout residual
+                # statistics.
+                flash_out = self._iha_hola_seq_expand_flash_forward(
+                    q,
+                    k,
+                    v,
+                    dropout=dropout,
+                    S_orig=S_orig,
+                    gate_states=gate_states,
+                )
+            else:
+                flash_out = flash_attn_fn(
+                    q_fa,
+                    k_fa,
+                    v_fa,
+                    dropout_p=dropout if self.training else 0.0,
+                    softmax_scale=self.scaling,
+                    causal=True,
+                    window_size=window,
+                )  # [B, S·P, H_q, d]
 
         # ── Affine-Scaled Attention inline (flash path) ───────────────────────
         # Decomposition used here:
@@ -3860,6 +4128,7 @@ class NeoLLMAttention(nn.Module):
                 S_orig=_S_orig,
                 position_ids=position_ids,
                 attention_mask=attention_mask,
+                gate_states=gate,
             )
         elif use_affine:
             if self.config._attn_implementation == "eager":