Add sublinear attention v2 improvements

README.md

@@ -83,6 +83,6 @@ and a tighter structured-sublinear attention profile (`window=128`, `stride=128`
 an ETA around 326 days, under the 1y+90d target, while keeping ctx=1280, B=2,
 DiffusionBlocks, gradient-checkpointed blocks, tied heads, and structured masks.
 
-Sublinear coverage update 2026-05-29: the saved AGILLM-4 trainer snapshot now matches the live run's improved sparse global memory. When the strided anchor list is larger than `--sublinear_max_anchors`, it keeps anchors evenly spread across the whole sequence instead of only the most recent anchors, and it always includes a small set of first-token attention sinks. At the live 128/128/128 profile this preserves full-span coverage for 32k-style contexts instead of going blind to the deep past after the 16k anchor horizon, with essentially the same key budget. See `sublinear_improved_snippet.py` for the minimal drop-in block and `sublinear_improved.py` for the coverage demo/standalone selector.
+Sublinear coverage update 2026-05-29: the saved AGILLM-4 trainer snapshot now matches the live v2 sparse global memory path. It fixes gathered ALiBi distance, suppresses duplicate local/anchor candidates before softmax, uses hybrid full-span + recent-tail anchors with explicit `--sublinear_sinks` and `--sublinear_recent_anchors`, and includes optional pooled K/V landmark summaries behind `--sublinear_pooled_landmarks`. At the live 128/128/128 profile it keeps deep-past coverage while preserving recent anchors and the same VRAM-first key budget. See `sublinear_improved_snippet.py` for the minimal blocks and `sublinear_improved.py` for the coverage demo/standalone selector.
 
 License: Apache-2.0 (matching the upstream method).

nB300_agillm4_vram_dblock.py

@@ -926,6 +926,9 @@ DEFAULT_SUBLINEAR_WINDOW = _env_int("AGILLM_SUBLINEAR_WINDOW", 256)
 DEFAULT_SUBLINEAR_STRIDE = _env_int("AGILLM_SUBLINEAR_STRIDE", 64)
 DEFAULT_SUBLINEAR_MAX_ANCHORS = _env_int("AGILLM_SUBLINEAR_MAX_ANCHORS", 256)
 DEFAULT_SUBLINEAR_CHUNK = _env_int("AGILLM_SUBLINEAR_CHUNK", 128)
+DEFAULT_SUBLINEAR_SINKS = _env_int("AGILLM_SUBLINEAR_SINKS", 4)
+DEFAULT_SUBLINEAR_RECENT_ANCHORS = _env_int("AGILLM_SUBLINEAR_RECENT_ANCHORS", -1)  # -1 = half of max anchors
+DEFAULT_SUBLINEAR_POOLED_LANDMARKS = bool(_env_int("AGILLM_SUBLINEAR_POOLED_LANDMARKS", 0))
 DEFAULT_ANCHOR_MEMORY = bool(_env_int("AGILLM_ANCHOR_MEMORY", 0))
 DEFAULT_ANCHOR_STRIDE = _env_int("AGILLM_ANCHOR_STRIDE", 256)
 DEFAULT_ANCHOR_MAX = _env_int("AGILLM_ANCHOR_MAX", 2048)
@@ -1289,6 +1292,9 @@ class TuneableAttentionMHA(nn.Module):
         sublinear_stride: int = DEFAULT_SUBLINEAR_STRIDE,
         sublinear_max_anchors: int = DEFAULT_SUBLINEAR_MAX_ANCHORS,
         sublinear_chunk: int = DEFAULT_SUBLINEAR_CHUNK,
+        sublinear_sinks: int = DEFAULT_SUBLINEAR_SINKS,
+        sublinear_recent_anchors: int = DEFAULT_SUBLINEAR_RECENT_ANCHORS,
+        sublinear_pooled_landmarks: bool = DEFAULT_SUBLINEAR_POOLED_LANDMARKS,
     ):
         super().__init__()
         assert d % h == 0
@@ -1299,6 +1305,12 @@ class TuneableAttentionMHA(nn.Module):
         self.sublinear_stride = max(0, int(sublinear_stride))
         self.sublinear_max_anchors = max(0, int(sublinear_max_anchors))
         self.sublinear_chunk = max(1, int(sublinear_chunk))
+        self.sublinear_sinks = max(0, int(sublinear_sinks))
+        recent = int(sublinear_recent_anchors)
+        if recent < 0:
+            recent = self.sublinear_max_anchors // 2
+        self.sublinear_recent_anchors = min(max(0, recent), self.sublinear_max_anchors)
+        self.sublinear_pooled_landmarks = bool(sublinear_pooled_landmarks)
         # Exact n1 harvest: one fused QKV projection is mathematically the same
         # as three independent bias-free Linear(d, d) projections with their
         # weights stacked along out_features.
@@ -1395,6 +1407,29 @@ class TuneableAttentionMHA(nn.Module):
             return (idx // block) <= (q_pos[:, None] // block)
         raise ValueError(f"unknown structured attention mask kind: {kind}")
 
+    def _sublinear_anchor_positions(self, k_len: int, device):
+        anchor_start = self.sublinear_stride - 1
+        if self.sublinear_stride <= 0 or self.sublinear_max_anchors <= 0 or anchor_start >= k_len:
+            anchors = torch.empty(0, device=device, dtype=torch.long)
+        else:
+            all_anchors = torch.arange(anchor_start, k_len, self.sublinear_stride, device=device, dtype=torch.long)
+            if all_anchors.numel() <= self.sublinear_max_anchors:
+                anchors = all_anchors
+            else:
+                recent_budget = min(self.sublinear_recent_anchors, self.sublinear_max_anchors)
+                span_budget = max(0, self.sublinear_max_anchors - recent_budget)
+                parts = []
+                if span_budget > 0:
+                    span_sel = torch.linspace(0, all_anchors.numel() - 1, span_budget, device=device).round().long().unique()
+                    parts.append(all_anchors[span_sel])
+                if recent_budget > 0:
+                    parts.append(all_anchors[-recent_budget:])
+                anchors = torch.cat(parts).unique() if parts else torch.empty(0, device=device, dtype=torch.long)
+        if self.sublinear_sinks > 0 and k_len > 0:
+            sinks = torch.arange(min(self.sublinear_sinks, k_len), device=device, dtype=torch.long)
+            anchors = torch.cat([sinks, anchors]).unique() if anchors.numel() else sinks
+        return anchors
+
     def _sublinear_attention(self, q, k, v, attn_mask=None, rel_bias_tokens=None):
         """Local-window + landmark attention: O(N * (window + N/stride))."""
         bsz, heads, q_len, _ = q.shape
@@ -1407,25 +1442,20 @@ class TuneableAttentionMHA(nn.Module):
         if self.use_relpos and rel_bias_tokens is not None:
             slopes = _alibi_slopes(self.h).to(device=device, dtype=torch.float32)
 
-        anchor_start = self.sublinear_stride - 1
-        if self.sublinear_stride > 0 and self.sublinear_max_anchors > 0 and anchor_start < k_len:
-            anchors = torch.arange(
-                anchor_start,
-                k_len,
-                self.sublinear_stride,
-                device=device,
-                dtype=torch.long,
-            )
-            if anchors.numel() > self.sublinear_max_anchors:
-                # even-coverage: span the WHOLE past at fixed budget (not just recent tail)
-                _sel = torch.linspace(0, anchors.numel() - 1, self.sublinear_max_anchors, device=device).round().long().unique()
-                anchors = anchors[_sel]
-        else:
-            anchors = torch.empty(0, device=device, dtype=torch.long)
-        # attention sinks: always keep the first few tokens (StreamingLLM)
-        _sink = int(getattr(self, "sublinear_sinks", 4))
-        if _sink > 0 and k_len > 0:
-            anchors = torch.cat([torch.arange(min(_sink, k_len), device=device, dtype=torch.long), anchors]).unique()
+        anchors = self._sublinear_anchor_positions(k_len, device)
+        anchor_k = anchor_v = None
+        if anchors.numel() and self.sublinear_pooled_landmarks and self.sublinear_stride > 1:
+            # Optional pooled landmarks: each global anchor summarizes its stride segment.
+            # This is off by default because it adds cumsum work; enable after benchmarking.
+            ends = anchors + 1
+            starts = (ends - self.sublinear_stride).clamp_min(0)
+            zero_k = k.new_zeros(k.size(0), k.size(1), 1, k.size(3))
+            zero_v = v.new_zeros(v.size(0), v.size(1), 1, v.size(3))
+            prefix_k = torch.cat([zero_k, k.cumsum(dim=2)], dim=2)
+            prefix_v = torch.cat([zero_v, v.cumsum(dim=2)], dim=2)
+            denom = (ends - starts).to(dtype=k.dtype).view(1, 1, -1, 1).clamp_min(1)
+            anchor_k = (prefix_k[:, :, ends, :] - prefix_k[:, :, starts, :]) / denom
+            anchor_v = (prefix_v[:, :, ends, :] - prefix_v[:, :, starts, :]) / denom
 
         offsets = torch.arange(
             -self.sublinear_window,
@@ -1443,24 +1473,38 @@ class TuneableAttentionMHA(nn.Module):
             local_valid = (local_raw >= 0) & (local_raw < k_len)
             local_idx = local_raw.clamp(0, max(0, k_len - 1))
 
+            k_local = k[:, :, local_idx, :]
+            v_local = v[:, :, local_idx, :]
             if anchors.numel():
                 anchor_idx = anchors.view(1, -1).expand(cur, -1)
-                anchor_valid = torch.ones_like(anchor_idx, dtype=torch.bool)
+                local_lo = (q_pos - self.sublinear_window).clamp_min(0).view(-1, 1)
+                local_hi = (q_pos + self.sublinear_window).clamp_max(max(0, k_len - 1)).view(-1, 1)
+                # Drop anchor copies already present in the local window; duplicates bias softmax mass.
+                anchor_valid = (anchor_idx < local_lo) | (anchor_idx > local_hi)
                 idx = torch.cat([local_idx, anchor_idx], dim=1)
                 valid = torch.cat([local_valid, anchor_valid], dim=1)
+                if anchor_k is not None and anchor_v is not None:
+                    k_anchor = anchor_k.unsqueeze(2).expand(-1, -1, cur, -1, -1)
+                    v_anchor = anchor_v.unsqueeze(2).expand(-1, -1, cur, -1, -1)
+                else:
+                    k_anchor = k[:, :, anchor_idx, :]
+                    v_anchor = v[:, :, anchor_idx, :]
+                k_sel = torch.cat([k_local, k_anchor], dim=-2)
+                v_sel = torch.cat([v_local, v_anchor], dim=-2)
             else:
                 idx = local_idx
                 valid = local_valid
+                k_sel = k_local
+                v_sel = v_local
 
             structured_valid = self._structured_valid(attn_mask, q_pos, idx)
             if structured_valid is not None:
                 valid = valid & structured_valid
 
-            k_sel = k[:, :, idx, :]
             scores = (q[:, :, q_start:q_end, :].unsqueeze(-2) * k_sel).sum(dim=-1) * scale
 
             if slopes is not None:
-                dist = (idx.view(1, 1, cur, -1) - q_pos.view(1, 1, cur, 1)).clamp_min(0).to(torch.float32)
+                dist = (q_pos.view(1, 1, cur, 1) - idx.view(1, 1, cur, -1)).abs().to(torch.float32)
                 scores = scores + (-slopes * dist).to(scores.dtype)
 
             if torch.is_tensor(attn_mask) and attn_mask.size(-1) == k_len and attn_mask.size(-2) >= q_end:
@@ -1470,7 +1514,6 @@ class TuneableAttentionMHA(nn.Module):
 
             scores = scores.masked_fill(~valid.view(1, 1, cur, -1), float("-inf"))
             weights = torch.softmax(scores.float(), dim=-1).to(v.dtype)
-            v_sel = v[:, :, idx, :]
             outputs.append((weights.unsqueeze(-1) * v_sel).sum(dim=-2))
 
         return torch.cat(outputs, dim=2)
@@ -1544,6 +1587,9 @@ class Block(nn.Module):
         sublinear_stride: int = DEFAULT_SUBLINEAR_STRIDE,
         sublinear_max_anchors: int = DEFAULT_SUBLINEAR_MAX_ANCHORS,
         sublinear_chunk: int = DEFAULT_SUBLINEAR_CHUNK,
+        sublinear_sinks: int = DEFAULT_SUBLINEAR_SINKS,
+        sublinear_recent_anchors: int = DEFAULT_SUBLINEAR_RECENT_ANCHORS,
+        sublinear_pooled_landmarks: bool = DEFAULT_SUBLINEAR_POOLED_LANDMARKS,
     ):
         super().__init__()
         self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
@@ -1556,6 +1602,9 @@ class Block(nn.Module):
             sublinear_stride=sublinear_stride,
             sublinear_max_anchors=sublinear_max_anchors,
             sublinear_chunk=sublinear_chunk,
+            sublinear_sinks=sublinear_sinks,
+            sublinear_recent_anchors=sublinear_recent_anchors,
+            sublinear_pooled_landmarks=sublinear_pooled_landmarks,
         )
         self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.ReLU(), nn.Linear(4 * d, d))
 
@@ -1581,6 +1630,9 @@ class Encoder(nn.Module):
         sublinear_stride: int = DEFAULT_SUBLINEAR_STRIDE,
         sublinear_max_anchors: int = DEFAULT_SUBLINEAR_MAX_ANCHORS,
         sublinear_chunk: int = DEFAULT_SUBLINEAR_CHUNK,
+        sublinear_sinks: int = DEFAULT_SUBLINEAR_SINKS,
+        sublinear_recent_anchors: int = DEFAULT_SUBLINEAR_RECENT_ANCHORS,
+        sublinear_pooled_landmarks: bool = DEFAULT_SUBLINEAR_POOLED_LANDMARKS,
         anchor_memory: bool = DEFAULT_ANCHOR_MEMORY,
         anchor_stride: int = DEFAULT_ANCHOR_STRIDE,
         anchor_max: int = DEFAULT_ANCHOR_MAX,
@@ -1599,6 +1651,9 @@ class Encoder(nn.Module):
                 sublinear_stride=sublinear_stride,
                 sublinear_max_anchors=sublinear_max_anchors,
                 sublinear_chunk=sublinear_chunk,
+                sublinear_sinks=sublinear_sinks,
+                sublinear_recent_anchors=sublinear_recent_anchors,
+                sublinear_pooled_landmarks=sublinear_pooled_landmarks,
             )
             for _ in range(l)
         ])
@@ -1610,6 +1665,9 @@ class Encoder(nn.Module):
         self.sublinear_stride = sublinear_stride
         self.sublinear_max_anchors = sublinear_max_anchors
         self.sublinear_chunk = sublinear_chunk
+        self.sublinear_sinks = sublinear_sinks
+        self.sublinear_recent_anchors = sublinear_recent_anchors
+        self.sublinear_pooled_landmarks = bool(sublinear_pooled_landmarks)
         self.anchor_memory_enabled = bool(anchor_memory)
         self.anchor_stride = int(anchor_stride)
         self.anchor_max = int(anchor_max)
@@ -2487,7 +2545,9 @@ def train(args):
         "AGILLM-4 runtime: "
         f"attn_backend={args.attn_backend} grad_checkpoint={args.grad_checkpoint} "
         f"sublinear_window={args.sublinear_window} sublinear_stride={args.sublinear_stride} "
-        f"sublinear_max_anchors={args.sublinear_max_anchors} sublinear_chunk={args.sublinear_chunk}"
+        f"sublinear_max_anchors={args.sublinear_max_anchors} sublinear_chunk={args.sublinear_chunk} "
+        f"sublinear_sinks={args.sublinear_sinks} sublinear_recent_anchors={args.sublinear_recent_anchors} "
+        f"sublinear_pooled_landmarks={args.sublinear_pooled_landmarks}"
     )
     core = Encoder(
         cfg,
@@ -2498,6 +2558,9 @@ def train(args):
         sublinear_stride=args.sublinear_stride,
         sublinear_max_anchors=args.sublinear_max_anchors,
         sublinear_chunk=args.sublinear_chunk,
+        sublinear_sinks=args.sublinear_sinks,
+        sublinear_recent_anchors=args.sublinear_recent_anchors,
+        sublinear_pooled_landmarks=args.sublinear_pooled_landmarks,
         anchor_memory=getattr(args, "anchor_memory", DEFAULT_ANCHOR_MEMORY),
         anchor_stride=getattr(args, "anchor_stride", DEFAULT_ANCHOR_STRIDE),
         anchor_max=getattr(args, "anchor_max", DEFAULT_ANCHOR_MAX),
@@ -2717,6 +2780,9 @@ def infer(args):
         sublinear_stride=args.sublinear_stride,
         sublinear_max_anchors=args.sublinear_max_anchors,
         sublinear_chunk=args.sublinear_chunk,
+        sublinear_sinks=args.sublinear_sinks,
+        sublinear_recent_anchors=args.sublinear_recent_anchors,
+        sublinear_pooled_landmarks=args.sublinear_pooled_landmarks,
         anchor_memory=getattr(args, "anchor_memory", DEFAULT_ANCHOR_MEMORY),
         anchor_stride=getattr(args, "anchor_stride", DEFAULT_ANCHOR_STRIDE),
         anchor_max=getattr(args, "anchor_max", DEFAULT_ANCHOR_MAX),
@@ -2888,6 +2954,13 @@ def main():
                     help="For --attn_backend sublinear, cap landmark candidates per query chunk.")
     tr.add_argument("--sublinear_chunk", type=int, default=DEFAULT_SUBLINEAR_CHUNK,
                     help="For --attn_backend sublinear, query chunk size controlling peak gather memory.")
+    tr.add_argument("--sublinear_sinks", type=int, default=DEFAULT_SUBLINEAR_SINKS,
+                    help="For sublinear attention, always include this many first-token attention sinks.")
+    tr.add_argument("--sublinear_recent_anchors", type=int, default=DEFAULT_SUBLINEAR_RECENT_ANCHORS,
+                    help="For capped sublinear anchors, reserve this many anchors for the recent tail; -1 uses half.")
+    tr.add_argument("--sublinear_pooled_landmarks", action=argparse.BooleanOptionalAction,
+                    default=DEFAULT_SUBLINEAR_POOLED_LANDMARKS,
+                    help="Use stride-segment pooled K/V summaries for sublinear landmark anchors.")
     tr.add_argument("--no_structured_masks", action="store_true",
                     help="Disable structured causal/SAT masks for sublinear attention and fall back to dense masks.")
     tr.add_argument("--anchor_memory", action="store_true",
@@ -3009,6 +3082,10 @@ def main():
     inf.add_argument("--sublinear_stride", type=int, default=DEFAULT_SUBLINEAR_STRIDE)
     inf.add_argument("--sublinear_max_anchors", type=int, default=DEFAULT_SUBLINEAR_MAX_ANCHORS)
     inf.add_argument("--sublinear_chunk", type=int, default=DEFAULT_SUBLINEAR_CHUNK)
+    inf.add_argument("--sublinear_sinks", type=int, default=DEFAULT_SUBLINEAR_SINKS)
+    inf.add_argument("--sublinear_recent_anchors", type=int, default=DEFAULT_SUBLINEAR_RECENT_ANCHORS)
+    inf.add_argument("--sublinear_pooled_landmarks", action=argparse.BooleanOptionalAction,
+                     default=DEFAULT_SUBLINEAR_POOLED_LANDMARKS)
     inf.add_argument("--no_structured_masks", action="store_true")
     inf.add_argument("--nat_expand", type=int, default=2)
     inf.add_argument("--nat_passes", type=int, default=1)

relaunch_agillm4_dblock_sg2.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+# Relaunch AGILLM-4 dblock with SG2's tuned config + improved sublinear attention v2.
+set -Eeuo pipefail
+cd /workspace/agillm-4
+export TOKENIZERS_PARALLELISM=false
+export TOKENIZER_ID="${TOKENIZER_ID:-deepseek-ai/DeepSeek-V4-Pro}"
+export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512,expandable_segments:True
+export AGILLM_ATTN_BACKEND=sublinear
+[ -f /root/.cache/huggingface/token ] && { export HF_TOKEN="$(tr -d '\r\n' </root/.cache/huggingface/token)"; export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"; }
+SAVE_DIR=/workspace/agillm4_4090_ckpts
+CKPT="$(ls -1t "$SAVE_DIR"/pretrain_step*.pt 2>/dev/null | head -1)"
+exec >> /workspace/agillm4_floor_train.log 2>&1
+echo "RELAUNCH_AGILLM4_DBLOCK_SG2 $(date -u +%Y-%m-%dT%H:%M:%SZ) resume=$CKPT (improved sublinear v2: ALiBi distance + dedupe + hybrid anchors + sinks)"
+exec python -u nB300_agillm4.py train --preset agillm4_floor --resume "$CKPT" \
+  --dblock --dblock_blocks 4 --dblock_schedule loss_balanced --dblock_warmup_steps 16 \
+  --dblock_sigma_curriculum_steps 2000 --dblock_log_every 25 --dblock_objective_mode stochastic \
+  --dblock_ar_prob 0.85 --dblock_sat_prob 0.075 --dblock_nat_prob 0.075 \
+  --dblock_ar_loss_tokens 512 --dblock_sat_loss_tokens 0 --dblock_nat_loss_tokens 512 \
+  --tie_weights --batch_size 2 --block 1280 --amp --attn_backend sublinear \
+  --sublinear_window 128 --sublinear_stride 128 --sublinear_max_anchors 128 --sublinear_chunk 128 \
+  --sublinear_sinks 4 --sublinear_recent_anchors 64 --no-sublinear_pooled_landmarks \
+  --grad_checkpoint --optimizer paged_adamw8bit --sat_every 4 --nat_every 4 --nat_max_tokens 768 --nat_mask_ratio 0.5 \
+  --token_param_ratio 100 --save_dir "$SAVE_DIR" --save_every_sec 86400 --heartbeat_every_sec 300 \
+  --empty_cache_every_steps 0 --delta_every_steps 25000 --delta_max_keep 1 --max_ckpts 1

sublinear_improved.py

@@ -1,54 +1,58 @@
 """Improved sublinear-attention anchor selection for AGILLM-4.
 
-AGILLM-4's `sublinear` attention = local sliding window + strided "landmark"
-anchors. The original capped the anchor set with `anchors[-max_anchors:]`, which
-DROPS the entire deep past once N > max_anchors*stride (the trainer goes blind to
-everything older than the recent tail). This patch keeps whole-sequence coverage at
-the SAME key budget, plus StreamingLLM-style attention sinks.
-
-Drop-in replacement for the anchor-cap block inside MHA._sublinear_attention.
+V2 used by the live AGILLM-4 DBlock line:
+- fixes gathered ALiBi distance so past causal keys receive distance penalty
+- suppresses local/anchor duplicate candidates before softmax
+- uses hybrid full-span + recent-tail anchors under the same max-anchor budget
+- exposes first-token attention sinks as `--sublinear_sinks`
+- includes optional pooled landmark K/V summaries behind `--sublinear_pooled_landmarks`
 """
 import torch
 
-def select_anchors(k_len, stride, max_anchors, sinks, device):
-    """Even-coverage strided landmarks over the FULL past + attention sinks."""
+
+def select_hybrid_anchors(k_len, stride, max_anchors, sinks=4, recent_anchors=-1, device='cpu'):
+    """Full-span + recent-tail landmark positions, plus attention sinks."""
+    device = torch.device(device)
     start = stride - 1
-    if stride > 0 and max_anchors > 0 and start < k_len:
-        anchors = torch.arange(start, k_len, stride, device=device, dtype=torch.long)
-        if anchors.numel() > max_anchors:
-            # even-coverage subsample across the whole sequence (NOT the recent tail)
-            sel = torch.linspace(0, anchors.numel() - 1, max_anchors, device=device).round().long().unique()
-            anchors = anchors[sel]
-    else:
+    if stride <= 0 or max_anchors <= 0 or start >= k_len:
         anchors = torch.empty(0, device=device, dtype=torch.long)
-    if sinks > 0 and k_len > 0:                       # always keep the first few tokens
-        anchors = torch.cat([torch.arange(min(sinks, k_len), device=device, dtype=torch.long), anchors]).unique()
+    else:
+        all_anchors = torch.arange(start, k_len, stride, device=device, dtype=torch.long)
+        if all_anchors.numel() <= max_anchors:
+            anchors = all_anchors
+        else:
+            if recent_anchors < 0:
+                recent_anchors = max_anchors // 2
+            recent_budget = min(max(0, int(recent_anchors)), max_anchors)
+            span_budget = max(0, max_anchors - recent_budget)
+            parts = []
+            if span_budget > 0:
+                sel = torch.linspace(0, all_anchors.numel() - 1, span_budget, device=device).round().long().unique()
+                parts.append(all_anchors[sel])
+            if recent_budget > 0:
+                parts.append(all_anchors[-recent_budget:])
+            anchors = torch.cat(parts).unique() if parts else torch.empty(0, device=device, dtype=torch.long)
+    if sinks > 0 and k_len > 0:
+        sink_idx = torch.arange(min(int(sinks), k_len), device=device, dtype=torch.long)
+        anchors = torch.cat([sink_idx, anchors]).unique() if anchors.numel() else sink_idx
     return anchors
 
-# --- exact patch applied to nB300_agillm4.py MHA._sublinear_attention ---
-PATCH = '''
-# replace:
-#     if anchors.numel() > self.sublinear_max_anchors:
-#         anchors = anchors[-self.sublinear_max_anchors :]
-# with:
-    if anchors.numel() > self.sublinear_max_anchors:
-        _sel = torch.linspace(0, anchors.numel() - 1, self.sublinear_max_anchors, device=device).round().long().unique()
-        anchors = anchors[_sel]
-# ... and after the else branch add:
-_sink = int(getattr(self, "sublinear_sinks", 4))
-if _sink > 0 and k_len > 0:
-    anchors = torch.cat([torch.arange(min(_sink, k_len), device=device, dtype=torch.long), anchors]).unique()
-'''
 
-if __name__ == "__main__":
-    # Structural coverage demo at the live config, N beyond max_anchors*stride.
-    N, W, stride, maxA, sinks = 32768, 128, 128, 128, 4
-    i = N - 1
-    loc = set(range(i - W, i + 1))
-    allA = list(range(stride - 1, i + 1, stride))
-    OLD = sorted(loc | set(allA[-maxA:]))
-    sel = torch.linspace(0, len(allA) - 1, maxA).round().long().tolist()
-    NEW = sorted(loc | {allA[s] for s in sel} | set(range(sinks)))
-    print(f"N={N} (cap bites: {len(allA)} anchors > {maxA})")
-    print(f"OLD covers {min(OLD)}..{max(OLD)}  -> blind to 0..{min(OLD)-1}; first-half keys={sum(x<N//2 for x in OLD)}")
-    print(f"NEW covers {min(NEW)}..{max(NEW)}  -> full span; first-half keys={sum(x<N//2 for x in NEW)}")
+def local_anchor_valid(q_pos, anchors, window, k_len):
+    """False where an anchor is already inside that query's local window."""
+    anchor_idx = anchors.view(1, -1).expand(q_pos.numel(), -1)
+    local_lo = (q_pos - window).clamp_min(0).view(-1, 1)
+    local_hi = (q_pos + window).clamp_max(max(0, k_len - 1)).view(-1, 1)
+    return (anchor_idx < local_lo) | (anchor_idx > local_hi)
+
+
+if __name__ == '__main__':
+    N, window, stride, maxA, sinks, recent = 32768, 128, 128, 128, 4, 64
+    old_all = list(range(stride - 1, N, stride))
+    old = sorted(set(range(N - window - 1, N)) | set(old_all[-maxA:]))
+    new = select_hybrid_anchors(N, stride, maxA, sinks, recent).tolist()
+    print(f'N={N} stride={stride} maxA={maxA} sinks={sinks} recent={recent}')
+    print(f'OLD anchor coverage: {min(old)}..{max(old)} first_half={sum(x < N//2 for x in old)}')
+    print(f'NEW anchor coverage: {min(new)}..{max(new)} first_half={sum(x < N//2 for x in new)} recent_tail={sum(x >= N-8192 for x in new)}')
+    q = torch.tensor([N - 1])
+    print(f'duplicate-suppressed anchors for final query: {int(local_anchor_valid(q, torch.tensor(new), window, N).sum())}/{len(new)}')

sublinear_improved_snippet.py

@@ -1,35 +1,29 @@
-"""Improved AGILLM-4 sublinear attention anchor selection.
+"""Minimal AGILLM-4 sublinear attention V2 snippets.
 
-Drop this block into `_sublinear_attention` in place of the recent-tail anchor
-selection. It keeps the same local-window + anchor key budget shape, but avoids
-losing the deep past once `num_anchors > sublinear_max_anchors`.
+These are the core blocks now folded into `nB300_agillm4_vram_dblock.py`.
 """
 
-anchor_start = self.sublinear_stride - 1
-if self.sublinear_stride > 0 and self.sublinear_max_anchors > 0 and anchor_start < k_len:
-    anchors = torch.arange(
-        anchor_start,
-        k_len,
-        self.sublinear_stride,
-        device=device,
-        dtype=torch.long,
-    )
-    if anchors.numel() > self.sublinear_max_anchors:
-        # Span the whole sequence instead of keeping only the recent tail.
-        sel = torch.linspace(
-            0,
-            anchors.numel() - 1,
-            self.sublinear_max_anchors,
-            device=device,
-        ).round().long().unique()
-        anchors = anchors[sel]
-else:
-    anchors = torch.empty(0, device=device, dtype=torch.long)
+# Anchor selection: full-span + recent-tail + sinks.
+anchors = self._sublinear_anchor_positions(k_len, device)
 
-# StreamingLLM-style attention sinks: preserve the first tokens as stable global memory.
-sink = int(getattr(self, "sublinear_sinks", 4))
-if sink > 0 and k_len > 0:
-    anchors = torch.cat([
-        torch.arange(min(sink, k_len), device=device, dtype=torch.long),
-        anchors,
-    ]).unique()
+# Optional pooled landmarks are available behind --sublinear_pooled_landmarks.
+if anchors.numel() and self.sublinear_pooled_landmarks and self.sublinear_stride > 1:
+    ends = anchors + 1
+    starts = (ends - self.sublinear_stride).clamp_min(0)
+    zero_k = k.new_zeros(k.size(0), k.size(1), 1, k.size(3))
+    zero_v = v.new_zeros(v.size(0), v.size(1), 1, v.size(3))
+    prefix_k = torch.cat([zero_k, k.cumsum(dim=2)], dim=2)
+    prefix_v = torch.cat([zero_v, v.cumsum(dim=2)], dim=2)
+    denom = (ends - starts).to(dtype=k.dtype).view(1, 1, -1, 1).clamp_min(1)
+    anchor_k = (prefix_k[:, :, ends, :] - prefix_k[:, :, starts, :]) / denom
+    anchor_v = (prefix_v[:, :, ends, :] - prefix_v[:, :, starts, :]) / denom
+
+# Duplicate suppression: do not let an anchor double-count a key already in local attention.
+anchor_idx = anchors.view(1, -1).expand(cur, -1)
+local_lo = (q_pos - self.sublinear_window).clamp_min(0).view(-1, 1)
+local_hi = (q_pos + self.sublinear_window).clamp_max(max(0, k_len - 1)).view(-1, 1)
+anchor_valid = (anchor_idx < local_lo) | (anchor_idx > local_hi)
+
+# Gathered ALiBi distance: distance from query to selected key, not future-only clamp.
+dist = (q_pos.view(1, 1, cur, 1) - idx.view(1, 1, cur, -1)).abs().to(torch.float32)
+scores = scores + (-slopes * dist).to(scores.dtype)