Update configuration_neollm.py

configuration_neollm.py

@@ -79,6 +79,29 @@ class NeoLLMConfig(PretrainedConfig):
             Whether to add a bias term to Q, K, V, and output projections.
         attention_dropout (:obj:`float`, *optional*, defaults to 0.1):
             Dropout probability applied to attention weights during training.
+        use_hola_memory (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable the FA2-compatible HOLA-derived episodic memory read path for
+            the IHA seq-expand branch. This is not a second attention branch: it
+            prepends a fixed-size selected KV memory to each FlashAttention-2
+            chunk and reads memory/current tokens inside one softmax. The write
+            policy uses a compile-friendly value-innovation proxy because the
+            stock FlashAttention-2 kernel does not expose the pre-dropout
+            accumulators required for the exact HOLA committed-residual score.
+        hola_memory_size (:obj:`int`, *optional*, defaults to 64):
+            Number of real-token episodes kept per layer/chunk when
+            ``use_hola_memory=True``. Each episode stores all IHA pseudo-slots.
+        hola_chunk_size (:obj:`int`, *optional*, defaults to 256):
+            Number of original tokens processed per HOLA memory chunk. The
+            memory is frozen while a chunk is processed and updated between
+            chunks, preserving causality with standard FlashAttention-2 masks.
+        hola_logit_bias_init (:obj:`float`, *optional*, defaults to -2.0):
+            Initial additive logit prior for memory entries in the augmented
+            FA2 subspace. Negative values make the model start close to the
+            baseline attention and learn to use memory gradually.
+        hola_rms_gamma_init (:obj:`float`, *optional*, defaults to 1.0):
+            Initial RMSNorm-γ scale for the memory-read query/key subspace.
+        hola_score_eps (:obj:`float`, *optional*, defaults to 1e-6):
+            Numerical epsilon for the HOLA-derived value-innovation score.
         head_dim (:obj:`int`, *optional*, defaults to 64):
             Dimensionality of each attention head.
         use_momentum_attention (:obj:`bool`, *optional*, defaults to ``True``):
@@ -666,6 +689,13 @@ class NeoLLMConfig(PretrainedConfig):
         partial_rotary_factor=0.25,
         attention_bias=False,
         attention_dropout=0.1,
+        # ── HOLA-derived episodic memory for FlashAttention-2 ─────────────
+        use_hola_memory=False,
+        hola_memory_size=64,
+        hola_chunk_size=256,
+        hola_logit_bias_init=-2.0,
+        hola_rms_gamma_init=1.0,
+        hola_score_eps=1e-6,
         head_dim=64,
         use_momentum_attention=True,
         momentum_gamma=0.10,
@@ -1075,6 +1105,13 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Attention ─────────────────────────────────────────────────────
         self.attention_bias = attention_bias
         self.attention_dropout = attention_dropout
+        # ── HOLA-derived episodic memory for FlashAttention-2 ─────────────
+        self.use_hola_memory = bool(use_hola_memory)
+        self.hola_memory_size = int(hola_memory_size)
+        self.hola_chunk_size = int(hola_chunk_size)
+        self.hola_logit_bias_init = float(hola_logit_bias_init)
+        self.hola_rms_gamma_init = float(hola_rms_gamma_init)
+        self.hola_score_eps = float(hola_score_eps)
         self.head_dim = head_dim
         self.use_momentum_attention = use_momentum_attention
         self.momentum_gamma = momentum_gamma