Update modeling_neollm.py

modeling_neollm.py

@@ -2355,11 +2355,16 @@ def affine_scaled_flash_attention_forward(
         v_sum_t = value_sum.transpose(1, 2)                       # [B,S_q,H_q,d]
         output = alpha_t * flash_out + beta_t * v_sum_t
     else:
+        # Keep the Affine β branch as a causal prefix summary even when the
+        # attention backend itself uses a local sliding window.  This restores
+        # the original IHA/local behavior: local logits plus a cheap global
+        # causal value context.  Padding masks are still handled inside
+        # _affine_valid_value_sum; explicit 4D masks above remain exact.
         value_sum_kv = _affine_valid_value_sum(
             value,
             attention_mask=attention_mask,
             query_len=query.shape[-2],
-            sliding_window=kwargs.get("sliding_window", None),
+            sliding_window=None,
         )                                                         # [B,H_kv,S_q,d]
         v_sum_t = value_sum_kv.transpose(1, 2)                    # [B,S_q,H_kv,d]
         H_q = flash_out.shape[2]
@@ -3921,11 +3926,13 @@ class NeoLLMAttention(nn.Module):
             )                                                        # [B, S·P, H_q, d]
 
         # ── Affine-Scaled Attention inline (flash path) ───────────────────────
-        # Decomposition of [α·softmax(QKᵀ) + β]·V (Bae et al. 2026):
+        # Decomposition used here:
         #   Term 1: α · flash_out                      (standard flash output)
-        #   Term 2: β · Σ_{j∈A(i)} V_j                (same valid keys as flash)
-        # For local IHA, A(i) is the flash sliding window; for packed IHA it is
-        # additionally reset at each document boundary.
+        #   Term 2: β · Σ_{j≤i} V_j                    (causal prefix value summary)
+        # This intentionally restores the original local-IHA behavior: the
+        # attention logits may use a sliding window, but the β branch supplies a
+        # cheap global causal value context.  For packed IHA, the prefix is
+        # reset at each document boundary to avoid cross-document leakage.
         # α, β arrive pre-expanded to [B, H_q, S·P, 1]; permute to [B,S·P,H,1]
         # to match the flash output layout [B, S·P, H_q, d].
         if use_affine:
@@ -3934,18 +3941,20 @@ class NeoLLMAttention(nn.Module):
             # repeated query-head view:
             #   Sum(repeat_kv(V)) == repeat_kv(Sum(V)).
             # The broadcast to H_q happens only in the final affine combine.
-            window_left = window[0] if isinstance(window, tuple) else None
+            # β uses the full causal prefix, not the local flash window.
+            # Keep the sum in KV-head space first; GQA expansion happens only
+            # in the final affine combine below.
             if is_packed:
                 value_sum_kv = _segmented_causal_value_sum(
                     v,
                     segment_lengths=doc_lengths_p,
-                    window_left=window_left,
+                    window_left=None,
                 )                                                    # [1,H_kv,S·P,d]
             else:
                 value_sum_kv = _causal_value_sum(
                     v,
                     query_len=SP,
-                    window_left=window_left,
+                    window_left=None,
                 )                                                    # [B,H_kv,S·P,d]
 
             alpha_t = alpha.permute(0, 2, 1, 3)                      # [B,S·P,H_q,1]