Update modeling_neollm.py

modeling_neollm.py

@@ -1971,30 +1971,36 @@ def _masked_value_sum_from_pairs(
 
 
 def _causal_value_sum(
-    value_expanded: torch.Tensor,
+    value_states: 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].
+    value_states: [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.
+
+    The sliding-window path uses prefix differences with slices, not padded
+    gathers.  This keeps the operation exact while avoiding extra index tensors
+    and repeated index_select calls in the hot path.
     """
-    key_len = value_expanded.shape[2]
-    prefix = value_expanded.cumsum(dim=2)  # [B,H,Sk,d]
+    key_len = value_states.shape[2]
+    prefix = value_states.cumsum(dim=2)  # [B,H,Sk,d]
 
-    if window_left is None or window_left < 0 or window_left >= key_len:
+    if window_left is None or window_left < 0 or int(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
+        left = int(window_left)
+        sums_all = prefix.clone()
+        if left + 1 < key_len:
+            # For i > left:
+            #   sum_{j=i-left}^{i} V_j = prefix[i] - prefix[i-left-1].
+            # The first left+1 rows are already the causal prefixes.
+            sums_all[:, :, left + 1 :, :] = (
+                prefix[:, :, left + 1 :, :] - prefix[:, :, : -(left + 1), :]
+            )
 
     if query_len == key_len:
         return sums_all
@@ -2002,6 +2008,41 @@ def _causal_value_sum(
     return sums_all[:, :, start:start + query_len, :]
 
 
+def _segmented_causal_value_sum(
+    value_states: torch.Tensor,
+    segment_lengths: torch.Tensor,
+    window_left: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Compute segmented causal/window value sums without a Python loop over
+    documents.  Used only for packed IHA where B=1 and position_ids reset.
+
+    value_states:    [1,H,Sk,d]
+    segment_lengths: [num_segments] lengths in the same expanded sequence axis
+    return:          [1,H,Sk,d]
+    """
+    key_len = value_states.shape[2]
+    if key_len == 0:
+        return value_states
+
+    lengths = segment_lengths.to(device=value_states.device, dtype=torch.long)
+    seg_ends = lengths.cumsum(0)
+    seg_starts = seg_ends - lengths
+    seg_start_per_pos = torch.repeat_interleave(seg_starts, lengths)
+
+    pos = torch.arange(key_len, device=value_states.device, dtype=torch.long)
+    if window_left is None or window_left < 0:
+        start_idx = seg_start_per_pos
+    else:
+        start_idx = torch.maximum(seg_start_per_pos, pos - int(window_left))
+    end_idx = pos + 1
+
+    prefix = F.pad(value_states.cumsum(dim=2), (0, 0, 1, 0))  # [1,H,Sk+1,d]
+    end_vals = prefix.index_select(2, end_idx)
+    start_vals = prefix.index_select(2, start_idx)
+    return end_vals - start_vals
+
+
 def _affine_valid_value_sum(
     value_expanded: torch.Tensor,
     attention_mask: Optional[torch.Tensor],
@@ -2293,24 +2334,49 @@ def affine_scaled_flash_attention_forward(
     # flash_out: [B, S, H_q, d_head]  — HF wrappers all return this layout
 
     # ── 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 = value_sum.transpose(1, 2).contiguous()           # [B, S_q, H_q, d_head]
+    # Fast path: compute the valid-key value sum in KV-head space, then
+    # broadcast it to query heads during the final affine combine.  Since the
+    # sum is linear, WindowSum(repeat_kv(V)) == repeat_kv(WindowSum(V)); this
+    # avoids doing cumsum/window arithmetic on the expanded H_q heads.
+    alpha_t = alpha.permute(0, 2, 1, 3)   # [B, S_q, H_q, 1]
+    beta_t  = beta.permute(0, 2, 1, 3)    # [B, S_q, H_q, 1]
 
-    # α, β: [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]
-    beta_t  = beta.permute(0, 2, 1, 3)    # [B, S, H_q, 1]
+    if attention_mask is not None and attention_mask.ndim == 4:
+        # Arbitrary 4D additive masks may be expressed in query-head space, so
+        # keep the exact fallback on repeated V.  This path is not the regular
+        # causal/window hot path.
+        value_expanded = repeat_kv(value, module.num_key_value_groups)
+        value_sum = _affine_valid_value_sum(
+            value_expanded,
+            attention_mask=attention_mask,
+            query_len=query.shape[-2],
+            sliding_window=kwargs.get("sliding_window", None),
+        )
+        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:
+        value_sum_kv = _affine_valid_value_sum(
+            value,
+            attention_mask=attention_mask,
+            query_len=query.shape[-2],
+            sliding_window=kwargs.get("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]
+        H_kv = v_sum_t.shape[2]
+        groups = module.num_key_value_groups
+        if groups == 1 or H_q == H_kv:
+            output = alpha_t * flash_out + beta_t * v_sum_t
+        else:
+            Bsz, Sq, _, Dh = flash_out.shape
+            flash_g = flash_out.reshape(Bsz, Sq, H_kv, groups, Dh)
+            alpha_g = alpha_t.reshape(Bsz, Sq, H_kv, groups, 1)
+            beta_g  = beta_t.reshape(Bsz, Sq, H_kv, groups, 1)
+            output = (alpha_g * flash_g + beta_g * v_sum_t.unsqueeze(3)).reshape(
+                Bsz, Sq, H_q, Dh
+            )
 
     # ── Combine and apply dropout to the full affine output ───────────────
-    output = alpha_t * flash_out + beta_t * v_cumsum_t           # [B, S, H_q, d_head]
     output = nn.functional.dropout(output, p=dropout, training=module.training)
 
     # attn_weights is None — flash never exposes the softmax weight matrix.
@@ -3863,30 +3929,40 @@ class NeoLLMAttention(nn.Module):
         # α, β 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]
+            # Compute the β value-sum in KV-head space first.  For GQA this is
+            # exact because the valid-key sum is linear and independent of the
+            # 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
             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
+                value_sum_kv = _segmented_causal_value_sum(
+                    v,
+                    segment_lengths=doc_lengths_p,
+                    window_left=window_left,
+                )                                                    # [1,H_kv,S·P,d]
             else:
-                value_sum = _causal_value_sum(
-                    value_expanded,
+                value_sum_kv = _causal_value_sum(
+                    v,
                     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]
-            attn_out   = torch.nn.functional.dropout(
+                )                                                    # [B,H_kv,S·P,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]
+            value_sum_t = value_sum_kv.transpose(1, 2)               # [B,S·P,H_kv,d]
+
+            if self.num_key_value_groups == 1 or H_q == H_kv:
+                attn_out = alpha_t * flash_out + beta_t * value_sum_t
+            else:
+                flash_g = flash_out.reshape(B, SP, H_kv, self.num_key_value_groups, d_head)
+                alpha_g = alpha_t.reshape(B, SP, H_kv, self.num_key_value_groups, 1)
+                beta_g  = beta_t.reshape(B, SP, H_kv, self.num_key_value_groups, 1)
+                attn_out = (
+                    alpha_g * flash_g + beta_g * value_sum_t.unsqueeze(3)
+                ).reshape(B, SP, H_q, d_head)
+
+            attn_out = torch.nn.functional.dropout(
                 attn_out, p=dropout, training=self.training
             )
         else: