uncompiled attention: chunked over visible keys instead of eager flex

README.md

@@ -142,7 +142,7 @@ A broad, quality-filtered web mix, an educational slice, long documents for the
 
 The model uses custom code, so `trust_remote_code=True` is required.
 
-**Requirements:** PyTorch 2.5 or newer (attention runs on FlexAttention); tested with transformers 5.14. On an Ampere or newer GPU (compute capability 8.0+) with triton, attention is compiled; on CPU, older GPUs (T4, V100), or if compiling fails, it falls back to eager attention (same results, slower; set `SWA_NO_COMPILE=1` to force it). GPUs without bf16 (T4, V100) and CPUs: load with `dtype=torch.float32`. Generation uses a KV cache. Batches must not be padded: generate one prompt at a time, or batch only prompts of equal length.
+**Requirements:** PyTorch 2.5 or newer (attention runs on FlexAttention); tested with transformers 5.14. On an Ampere or newer GPU (compute capability 8.0+) with triton, attention is compiled; on CPU, older GPUs (T4, V100), or if compiling fails, it falls back to chunked attention: the same results, with sliding-window layers scoring only the 512 keys each query can see, so memory grows with length instead of length squared (set `SWA_NO_COMPILE=1` to force it). GPUs without bf16 (T4, V100) and CPUs: load with `dtype=torch.float32`. Generation uses a KV cache. Batches must not be padded: generate one prompt at a time, or batch only prompts of equal length.
 
 ```python
 import torch

modeling_swa.py

@@ -56,9 +56,9 @@ def _can_compile():  # compiled flex kernels need triton and a CUDA GPU of compu
 
 
 class SafeFlex:
-    """flex_attention compiled where it can be (fused, memory-light); on CPU, pre-Ampere GPUs (T4, V100), a missing
-    triton, or the first compile failure it runs eagerly for the rest of the process: the same attention, slower,
-    and materializing the full score matrix."""
+    """flex_attention compiled where it can be (fused, memory-light). On CPU, pre-Ampere GPUs (T4, V100), a missing
+    triton, or the first compile failure, `compiled` is None and the model uses chunked_attention instead (the call
+    that hit the failure finishes on eager flex_attention)."""
     def __init__(s): s.compiled = torch.compile(flex_attention) if _can_compile() else None
 
     def give_up(s, e):
@@ -96,6 +96,17 @@ def cached_attention(q, k, v, qpos, kpos, qdoc, kdoc, window, sink):
     return torch.einsum("bhtl,blhd->bthd", sc.softmax(-1)[..., :-1], v.float()).to(q.dtype)
 
 
+def chunked_attention(q, k, v, pos, docs, window, sink, chunk=512):
+    """The flex path without compiling: queries in chunks, each over only the keys it can see (the last `window`
+    positions for sliding-window layers, everything before it for full layers), so no T x T score matrix is built.
+    q (B, T, H, D); k, v (B, T, KV, D); pos (T,); docs (B, T). Same result as flex with the block masks."""
+    out = []
+    for i in range(0, q.shape[1], chunk):
+        j = min(i + chunk, q.shape[1]); a = max(0, i - window + 1) if window else 0
+        out.append(cached_attention(q[:, i:j], k[:, a:j], v[:, a:j], pos[i:j], pos[a:j], docs[:, i:j], docs[:, a:j], window, sink))
+    return torch.cat(out, 1)
+
+
 class SWACache:
     """KV cache for generation. Per layer: keys/values after norms, HoPE and value residual (sliding-window layers keep
     only the last `window` tokens, full layers keep all), their positions and document ids, and the last conv - 1 raw
@@ -129,12 +140,16 @@ class Block(nn.Module):
         if v1 is None: v1 = v                                                   # block 0 hands its values to every later block
         else: lam = s.lam.to(v.dtype); v = lam[0] * v + lam[1] * v1
         q, k = rope(s.qn(q), cos, sin).type_as(v), rope(s.kn(k), cos, sin).type_as(v)
-        if cache is None or "k" not in cache:                                   # training, eval, and the prompt pass
-            flex = flex_attention if torch.compiler.is_compiling() else _flex
+        if (cache is None or "k" not in cache) and masks is not None and (torch.compiler.is_compiling() or _flex.compiled is not None):
+            flex = flex_attention if torch.compiler.is_compiling() else _flex   # training, eval, and the prompt pass (fused)
             a, lse = flex(*(t.transpose(1, 2) for t in (q, k, v)), block_mask=masks[s.kind], enable_gqa=True, return_lse=True)
             a = a * torch.sigmoid(lse - s.sink.view(1, -1, 1)).unsqueeze(-1).type_as(a)   # softmax with sink: Z / (Z + e^sink)
             o = a.transpose(1, 2).reshape(B, T, c.h * c.hd)
             K, V, P, D = k, v, pos, docs
+        elif cache is None or "k" not in cache:                                 # the same, where flex cannot compile
+            P = pos if pos is not None else torch.arange(T, device=x.device)
+            o = chunked_attention(q, k, v, P, docs, c.window if s.kind == "swa" else 0, s.sink).reshape(B, T, c.h * c.hd)
+            K, V, D = k, v, docs
         else:                                                                   # decode: new queries over cached + new keys
             K, V = torch.cat([cache["k"], k], 1), torch.cat([cache["v"], v], 1)
             P, D = torch.cat([cache["pos"], pos]), torch.cat([cache["doc"], docs], 1)
@@ -164,7 +179,8 @@ class SWA(nn.Module):
         if docs is None:
             docs = (x == c.bos_token_id).cumsum(1)
             if off: docs = docs + cache.last_doc[:, None]                       # continue the cached document ids
-        masks = block_masks(c, docs) if not off else None                       # decode steps use cached_attention instead
+        fused = not off and (torch.compiler.is_compiling() or _flex.compiled is not None)
+        masks = block_masks(c, docs) if fused else None                         # decode steps and uncompiled runs use cached_attention
         x = s.emb(x)
         if torch.is_autocast_enabled(): x = x.to(torch.get_autocast_dtype("cuda"))
         inv = 1 / c.theta ** (torch.arange(0, c.hd, 2, device=x.device) / c.hd)   # rotary frequency per q/k dimension pair