Update modeling_neollm.py

modeling_neollm.py

@@ -26,16 +26,18 @@ Attention stack (orthogonal, all active simultaneously when enabled):
        increases attention entropy, and is complementary to Gated Attention
        (Bae et al. 2026, Table 2: Affine-Scaled + Gated > either alone).
 
-       Flash/SDPA path (exact, no approximation):
+       Flash/SDPA path:
            Expanding [α·softmax(QKᵀ)+β]V distributively yields two terms:
-               α · flash_attn(Q,K,V)    — Flash computes this directly
-             + β · Σ_{j≤i} V_j          — causal prefix-sum of V (V.cumsum)
-           The V_cumsum tensor [B,H,S,d_head] is the only extra memory vs a
-           standard flash call. Semantically identical to the eager path;
-           per-weight softmax tensors (attn_weights_pre/post_affine) are
-           unavailable in flash mode and remain None in AnalysisState.
+               α · attention(Q,K,V)      — backend computes this directly
+             + β · Σ_{j∈A(i)} V_j        — sum over the same valid keys A(i)
+           For global causal attention A(i) is a prefix and this is V.cumsum.
+           For sliding-window or packed sequences the sum is windowed or
+           segmented so β cannot leak outside the backend attention mask.
+           Per-weight tensors (attn_weights_pre/post_affine) are unavailable
+           in flash mode and remain None in AnalysisState.
 
-       Eager path: exact with full weight access, used for interpretability.
+       Eager path: applies the same valid-key mask to α·softmax + β and keeps
+           full weight access for interpretability.
 
 References:
     FANformer: "FANformer: Improving Large Language Models Through Effective
@@ -86,7 +88,7 @@ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_u
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, logging
-from .configuration_neollm import NeoLLMConfig
+from configuration_neollm import NeoLLMConfig
 
 from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
 torch._dynamo.config.capture_scalar_outputs = True
@@ -351,7 +353,7 @@ class AttentionAnalysis:
     beta_per_head:            Optional[torch.Tensor] = None  # β [B,H,S,1] moving-avg bias
     alpha_moving_avg:         Optional[torch.Tensor] = None  # alpha_ma EMA snapshot
     attn_weights_pre_affine:  Optional[torch.Tensor] = None  # softmax weights before α,β
-    attn_weights_post_affine: Optional[torch.Tensor] = None  # α·softmax + β
+    attn_weights_post_affine: Optional[torch.Tensor] = None  # valid_mask·(α·softmax + β)
 
     # ── Standard attention weights (eager non-affine path) ────────────
     attn_weights: Optional[torch.Tensor] = None  # softmax weights (None for flash/sdpa)
@@ -1920,6 +1922,122 @@ def infer_key_validity(
     return valid
 
 
+
+
+def _expand_pair_valid_mask(
+    attention_mask: Optional[torch.Tensor],
+    batch_size: int,
+    num_heads: int,
+    query_len: int,
+    key_len: int,
+) -> Optional[torch.Tensor]:
+    """
+    Return a boolean [B,H,Sq,Sk] mask for positions that attention may read.
+
+    Supported contracts used in this file:
+      • 2D padding mask [B,Sk], with truthy values for valid keys.
+      • 4D additive mask [B,1|H,Sq,Sk], with 0 on valid pairs and -inf/large
+        negative values on invalid causal, local or padding pairs.
+    """
+    if attention_mask is None:
+        return None
+
+    if attention_mask.ndim == 2:
+        if attention_mask.shape[-1] < key_len:
+            return None
+        valid = attention_mask[:, :key_len].to(dtype=torch.bool)  # [B,Sk]
+        valid = valid[:, None, None, :].expand(batch_size, num_heads, query_len, key_len)
+        return valid
+
+    if attention_mask.ndim != 4:
+        return None
+    if attention_mask.shape[-2] < query_len or attention_mask.shape[-1] < key_len:
+        return None
+
+    valid = attention_mask[:, :, :query_len, :key_len] == 0
+    if valid.shape[1] == 1 and num_heads != 1:
+        valid = valid.expand(-1, num_heads, -1, -1)
+    elif valid.shape[1] != num_heads:
+        valid = valid[:, :1, :, :].expand(-1, num_heads, -1, -1)
+    return valid
+
+
+def _masked_value_sum_from_pairs(
+    value_expanded: torch.Tensor,
+    pair_valid: torch.Tensor,
+) -> torch.Tensor:
+    """Compute Σ_{j∈A(i)} V_j from an explicit pair-valid mask."""
+    return torch.matmul(pair_valid.to(value_expanded.dtype), value_expanded)
+
+
+def _causal_value_sum(
+    value_expanded: torch.Tensor,
+    query_len: int,
+    window_left: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Compute Σ_{j∈A(i)} V_j for causal or sliding-window causal attention.
+
+    value_expanded: [B,H,Sk,d].  The returned tensor is [B,H,Sq,d].
+    For cached decoding, Sq may be smaller than Sk; query rows are aligned to
+    the last Sq key positions, matching causal self-attention with cache.
+    """
+    key_len = value_expanded.shape[2]
+    prefix = value_expanded.cumsum(dim=2)  # [B,H,Sk,d]
+
+    if window_left is None or window_left < 0 or window_left >= key_len:
+        sums_all = prefix
+    else:
+        # Valid keys for absolute position i are max(0, i-window_left) ... i.
+        prefix_padded = F.pad(prefix, (0, 0, 1, 0))  # [B,H,Sk+1,d]
+        end_idx = torch.arange(key_len, device=value_expanded.device) + 1
+        start_idx = torch.clamp(end_idx - (int(window_left) + 1), min=0)
+        end_vals = prefix_padded.index_select(2, end_idx)
+        start_vals = prefix_padded.index_select(2, start_idx)
+        sums_all = end_vals - start_vals
+
+    if query_len == key_len:
+        return sums_all
+    start = max(key_len - query_len, 0)
+    return sums_all[:, :, start:start + query_len, :]
+
+
+def _affine_valid_value_sum(
+    value_expanded: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    query_len: int,
+    sliding_window: Optional[Union[int, Tuple[int, int]]] = None,
+) -> torch.Tensor:
+    """
+    Compute the β term support S_i = Σ_{j∈A(i)} V_j for Affine attention.
+
+    It first respects an explicit 4D additive mask exactly.  Otherwise it uses
+    causal prefix sums, optionally restricted by a left sliding-window size.
+    """
+    batch_size, num_heads, key_len, _ = value_expanded.shape
+
+    pair_valid = None
+    if attention_mask is not None and attention_mask.ndim == 4:
+        pair_valid = _expand_pair_valid_mask(
+            attention_mask, batch_size, num_heads, query_len, key_len
+        )
+    if pair_valid is not None:
+        return _masked_value_sum_from_pairs(value_expanded, pair_valid)
+
+    # 2D padding masks zero invalid keys before the prefix/window sum.
+    if attention_mask is not None and attention_mask.ndim == 2:
+        if attention_mask.shape[-1] >= key_len:
+            valid = attention_mask[:, :key_len].to(value_expanded.dtype)
+            value_expanded = value_expanded * valid[:, None, :, None]
+
+    window_left = None
+    if isinstance(sliding_window, tuple):
+        window_left = int(sliding_window[0]) if len(sliding_window) > 0 else None
+    elif sliding_window is not None:
+        window_left = int(sliding_window)
+
+    return _causal_value_sum(value_expanded, query_len, window_left)
+
 def head_linear_compose(
     hidden_states: torch.Tensor, mixing_matrix: torch.Tensor
 ) -> torch.Tensor:
@@ -2061,21 +2179,38 @@ def affine_scaled_eager_attention_forward(
     value_states = repeat_kv(value, module.num_key_value_groups)
 
     attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    pair_valid = _expand_pair_valid_mask(
+        attention_mask,
+        batch_size=query.shape[0],
+        num_heads=key_states.shape[1],
+        query_len=query.shape[-2],
+        key_len=key_states.shape[-2],
+    )
     if attention_mask is not None:
-        attn_weights = attn_weights + attention_mask[:, :, :, : key_states.shape[-2]]
+        if attention_mask.ndim == 4:
+            attn_weights = attn_weights + attention_mask[:, :, : query.shape[-2], : key_states.shape[-2]]
+        elif attention_mask.ndim == 2:
+            key_valid = attention_mask[:, : key_states.shape[-2]].to(dtype=torch.bool)
+            additive = torch.zeros_like(attn_weights)
+            additive = additive.masked_fill(~key_valid[:, None, None, :], torch.finfo(attn_weights.dtype).min)
+            attn_weights = attn_weights + additive
 
     attn_weights_softmax = nn.functional.softmax(
         attn_weights, dim=-1, dtype=torch.float32
     ).to(query.dtype)
 
+    if pair_valid is not None:
+        attn_weights_softmax = attn_weights_softmax.masked_fill(~pair_valid, 0.0)
+
     if attn_analysis is not None:
         attn_analysis.attn_weights_pre_affine = attn_weights_softmax.detach()
 
-    # Affine reweighting: relaxes the unit-sum constraint of softmax.
-    # α ∈ [0, 1] scales down the softmax distribution (input-adaptively per head).
-    # β offsets to prevent collapse when α deviates from its running mean.
-    # Shapes: α, β are [B, H, S_q, 1], attn_weights is [B, H, S_q, S_k] → broadcast.
+    # Affine reweighting over valid keys only.  β must obey the same causal,
+    # local and padding mask as softmax; otherwise invalid V_j would leak in.
+    # Shapes: α, β are [B, H, S_q, 1], weights are [B, H, S_q, S_k].
     attn_weights_affine = alpha * attn_weights_softmax + beta
+    if pair_valid is not None:
+        attn_weights_affine = attn_weights_affine.masked_fill(~pair_valid, 0.0)
 
     if attn_analysis is not None:
         attn_analysis.attn_weights_post_affine = attn_weights_affine.detach()
@@ -2101,7 +2236,7 @@ def affine_scaled_flash_attention_forward(
     """
     Affine-Scaled Attention — flash/sdpa path.
 
-    Exact mathematical decomposition of [α·softmax(QKᵀ)+β]V using only the
+    Mathematical decomposition of [α·softmax(QKᵀ)+β]V using only the
     public flash/sdpa interface — no kernel modification required.
 
     Derivation
@@ -2109,12 +2244,12 @@ def affine_scaled_flash_attention_forward(
     The paper formula expands distributively:
 
         [α · softmax(QKᵀ/√dk) + β] · V
-        = α · [softmax(QKᵀ/√dk) · V]   ← term 1: standard flash output
-        + β · [Σ_{j≤i} V_j]             ← term 2: causal prefix-sum of V
+        = α · [softmax(QKᵀ/√dk) · V]   ← term 1: standard backend output
+        + β · [Σ_{j∈A(i)} V_j]          ← term 2: same valid keys A(i)
 
-    Term 2 follows because β is a scalar per query (broadcast over all S_k),
-    and softmax weights sum to 1 over the causal window:
-        β_i · Σ_j w_{i,j} · V_j = β_i · Σ_{j≤i} V_j  (since Σ w_{i,j} = 1)
+    For global causal attention, A(i) is the prefix j≤i and term 2 is a
+    cumsum. For sliding-window or packed/additive-mask paths, term 2 must use
+    the corresponding windowed or explicit valid-key sum.
 
     Dropout
     -------
@@ -2126,20 +2261,16 @@ def affine_scaled_flash_attention_forward(
     regularisation that achieves the same intent without the intermediate
     weight matrix. During inference dropout=0 so the paths are identical.
 
-    Padding mask
-    ------------
-    The V_cumsum must not accumulate values from padding positions, since the
-    flash kernel zeros those out internally but a plain cumsum does not.
-    Depending on the backend contract, attention_mask can arrive as either a
-    2D padding mask [B, S] (flash_attention_2/3) or a 4D additive mask
-    [B, 1, S, S] (eager/sdpa-style paths). We normalize both to a per-token
-    validity mask before the cumsum.
+    Mask handling
+    -------------
+    The β value-sum must respect the same valid keys as the backend attention.
+    A 2D padding mask zeros invalid keys before prefix/window sums.  A 4D
+    additive mask is used directly to compute Σ_{j∈A(i)} V_j exactly.
 
     Memory overhead vs standard flash call
     ---------------------------------------
-    One extra tensor of shape [B, H_q, S, d_head] for V_cumsum.
-    At (B=1, H=8, S=2048, d_head=64, bf16): ≈ 4 MB per call, ≈ 48 MB total
-    across 12 layers. Allocated and freed within each forward call.
+    Global causal attention adds one [B, H_q, S, d_head] value-sum tensor.
+    Exact 4D-mask fallback may materialise a [B, H_q, S_q, S_k] boolean mask.
 
     Attention-weight analysis fields (attn_weights_pre_affine,
     attn_weights_post_affine) remain None — those tensors are never
@@ -2161,28 +2292,18 @@ def affine_scaled_flash_attention_forward(
     )
     # flash_out: [B, S, H_q, d_head]  — HF wrappers all return this layout
 
-    # ── Term 2: β · causal prefix-sum of V ───────────────────────────────
-    # Expand V from KV heads to query heads for GQA.
-    value_expanded = repeat_kv(value, module.num_key_value_groups)  # [B, H_q, S, d_head]
-
-    # Zero out padded positions so they don't accumulate into the prefix sum.
-    # flash_attention_2/3 passes a 2D padding mask, while eager/sdpa-style
-    # paths may still pass a 4D additive mask.
-    if attention_mask is not None:
-        if attention_mask.ndim == 2:
-            valid = attention_mask.to(value_expanded.dtype)      # [B, S]
-            valid = valid.unsqueeze(1).unsqueeze(-1)             # [B, 1, S, 1]
-            value_expanded = value_expanded * valid              # [B, H_q, S, d_head]
-        elif attention_mask.ndim == 4:
-            diag = attention_mask.diagonal(dim1=-2, dim2=-1)     # [B, 1, S]
-            valid = (diag == 0).to(value_expanded.dtype)         # [B, 1, S]
-            valid = valid.unsqueeze(-1)                          # [B, 1, S, 1]
-            value_expanded = value_expanded * valid              # [B, H_q, S, d_head]
-
-    # Causal prefix-sum: position i accumulates all valid j ≤ i.
-    v_cumsum   = value_expanded.cumsum(dim=2)                    # [B, H_q, S, d_head]
+    # ── Term 2: β · Σ_{j∈A(i)} V_j ──────────────────────────────────────
+    # Expand V from KV heads to query heads for GQA, then sum only over the
+    # same valid keys used by the backend attention path.
+    value_expanded = repeat_kv(value, module.num_key_value_groups)  # [B, H_q, S_k, d_head]
+    value_sum = _affine_valid_value_sum(
+        value_expanded,
+        attention_mask=attention_mask,
+        query_len=query.shape[-2],
+        sliding_window=kwargs.get("sliding_window", None),
+    )                                                             # [B, H_q, S_q, d_head]
     # Transpose to match flash_out layout.
-    v_cumsum_t = v_cumsum.transpose(1, 2).contiguous()          # [B, S, H_q, d_head]
+    v_cumsum_t = value_sum.transpose(1, 2).contiguous()           # [B, S_q, H_q, d_head]
 
     # α, β: [B, H_q, S, 1] → [B, S, H_q, 1] to broadcast over d_head.
     alpha_t = alpha.permute(0, 2, 1, 3)   # [B, S, H_q, 1]
@@ -3596,8 +3717,8 @@ class NeoLLMAttention(nn.Module):
 
         Calls flash_attn_func (dense path) or flash_attn_varlen_func (packed
         path) directly.  Affine-Scaled Attention is also computed inline here
-        so that both Term 1 (flash output) and Term 2 (V prefix-sum) operate
-        on the S·P expanded sequence consistently.
+        so that Term 2 sums V over the same valid keys as flash attention on
+        the S·P expanded sequence.
 
         Packing detection
         -----------------
@@ -3614,12 +3735,11 @@ class NeoLLMAttention(nn.Module):
         Affine-Scaled Attention (inline)
         ---------------------------------
         When use_affine=True:
-          output = α · flash_out + β · V_cumsum_t
+          output = α · flash_out + β · Σ_{j∈A(i)} V_j
 
         α, β are [B, H_q, S·P, 1] (already expanded over the seq dim before
-        this call).  V_cumsum is the causal prefix-sum of V over the S·P
-        expanded sequence — semantically correct because virtual token (s, p)
-        accumulates all virtual tokens before it in interleaved order.
+        this call).  The V sum is causal prefix-sum for global IHA, windowed
+        for local IHA, and segmented when packed position_ids reset.
 
         Args
         ----
@@ -3735,16 +3855,34 @@ class NeoLLMAttention(nn.Module):
             )                                                        # [B, S·P, H_q, d]
 
         # ── Affine-Scaled Attention inline (flash path) ───────────────────────
-        # Exact decomposition of [α·softmax(QKᵀ) + β]·V (Bae et al. 2026):
+        # Decomposition of [α·softmax(QKᵀ) + β]·V (Bae et al. 2026):
         #   Term 1: α · flash_out                      (standard flash output)
-        #   Term 2: β · Σ_{j≤i} V_j                   (causal prefix-sum of V)
-        # Both terms operate on the S·P virtual sequence — semantically correct
-        # because token (s,p) causally precedes (s,p+1) and all of (s+1,·).
+        #   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.
         # α, β 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:
             value_expanded = repeat_kv(v, self.num_key_value_groups)  # [B,H_q,S·P,d]
-            v_cumsum_t = value_expanded.cumsum(dim=2).transpose(1, 2) # [B,S·P,H_q,d]
+            window_left = window[0] if isinstance(window, tuple) else None
+            if is_packed:
+                value_sum = torch.empty_like(value_expanded)
+                start = 0
+                for _seg_len in doc_lengths_p.to(dtype=torch.long).tolist():
+                    end = start + int(_seg_len)
+                    value_sum[:, :, start:end, :] = _causal_value_sum(
+                        value_expanded[:, :, start:end, :],
+                        query_len=int(_seg_len),
+                        window_left=window_left,
+                    )
+                    start = end
+            else:
+                value_sum = _causal_value_sum(
+                    value_expanded,
+                    query_len=SP,
+                    window_left=window_left,
+                )
+            v_cumsum_t = value_sum.transpose(1, 2).contiguous()        # [B,S·P,H_q,d]
             alpha_t    = alpha.permute(0, 2, 1, 3)                    # [B,S·P,H_q,1]
             beta_t     = beta.permute(0, 2, 1, 3)                     # [B,S·P,H_q,1]
             attn_out   = alpha_t * flash_out + beta_t * v_cumsum_t    # [B,S·P,H_q,d]
@@ -4095,7 +4233,7 @@ class NeoLLMAttention(nn.Module):
         # Active whenever use_affine_scaled_attention=True, regardless of
         # attention backend.  Two code paths — same math, different execution:
         #   eager : full weight access, attn_weights_pre/post_affine captured.
-        #   flash/sdpa: α·flash_out + β·V_cumsum, no weight tensors materialised.
+        #   flash/sdpa: α·backend_out + β·Σ valid V, no weight tensors materialised.
         alpha = None
         beta  = None
         use_affine = self.use_affine_scaled_attention
@@ -4131,8 +4269,7 @@ class NeoLLMAttention(nn.Module):
             if _iha_seq_expand:
                 alpha = alpha.repeat_interleave(self.iha_P, dim=2)  # [B,H_q,S·P,1]
                 beta  = beta.repeat_interleave(self.iha_P, dim=2)   # [B,H_q,S·P,1]
-                # Recompute beta with the correct N = S·P (denominator scales
-                # with the number of virtual keys the query can attend to).
+                # Recompute beta with N = S·P for the expanded attention axis.
                 # alpha_ma shape [1,H,1,1] broadcasts correctly against expanded alpha.
                 N_expanded = k.shape[-2]   # S·P after IHA expansion
                 beta = (self.alpha_ma.to(alpha.dtype) - alpha) / max(N_expanded, 1)
@@ -4147,7 +4284,7 @@ class NeoLLMAttention(nn.Module):
             #   3. attention_mask shaped [B,S] is incompatible with S·P tokens
             # _iha_seq_expand_flash_forward calls flash_attn_func /
             # flash_attn_varlen_func directly, handling packing, sliding window,
-            # and the Affine-Scaled inline decomposition (α·flash + β·V_cumsum).
+            # and the Affine-Scaled inline decomposition (α·flash + β·Σ valid V).
             # alpha/beta are [B, H_q, S·P, 1] at this point (already expanded).
             attn_out, attn_weights = self._iha_seq_expand_flash_forward(
                 q, k, v,
@@ -4169,7 +4306,7 @@ class NeoLLMAttention(nn.Module):
                     **kwargs,
                 )
             else:
-                # Flash / SDPA: exact formula via V.cumsum, no weight tensors.
+                # Flash / SDPA: valid-key value sum, no weight tensors.
                 # attn_weights_pre/post_affine remain None in AnalysisState.
                 backend_kwargs = kwargs
                 if flash_local_sliding_window is not None: