Chunk ViTMatte's global attention to cap its N^2 memory spike

ViTMatte's VitDet backbone runs 4 global attention blocks that materialize
the full [heads x N x N] map: ~19 GB transient at 2048x2048 (16384 tokens),
the pipeline's memory peak. transformers has no SDPA path for this
architecture (the decomposed rel-pos bias is added to raw scores), so
compute the same attention in query-row chunks instead: the bias
factorizes over query rows, making the chunked form mathematically exact
-- output verified BITWISE-identical (unit: fp32/bf16 x 3 sizes; full
pipeline: TestImage3 fp32 and 2048x2048 bf16, all byte-equal).

model.attn_query_chunk (default 2048, 0 = stock one-shot) engages only on
blocks seeing more tokens than the chunk size, so window blocks keep the
original path. Measured @2048x2048 bf16 (9700X): ViTMatte spike
19.8 -> 4.0 GB for ~15% more ViTMatte time; whole pipeline peak
21.4 -> 11.4 GB, warm 51 -> 53 s.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

README.md

@@ -99,7 +99,9 @@ file to tune the detailed parameters. Runtime choices stay on the command line:
 `--device` (default **CPU**; use `--device cuda` for GPU), `--precision` (default
 `fp32`; `bf16` speeds up all three models and halves the matting model's activation
 memory with visually identical alpha — needs bf16-capable hardware, falls back to
-fp32 elsewhere), `--seg-backend`
+fp32 elsewhere; large inputs additionally get ViTMatte's global attention computed
+in query chunks by default — exact, bitwise-identical, caps the memory spike at
+~4 GB instead of ~19 GB at 2048x2048, see `model.attn_query_chunk`), `--seg-backend`
 (default `birefnet`; `anime-seg` for anime characters), `--screen-color` (default:
 auto-detect the flat background), and `--trimap-mode`.
 

bgfilter/attn_chunk.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+"""Query-chunked attention for ViTMatte's VitDet backbone.
+
+ViTMatte's backbone runs 4 *global* attention blocks over every patch token.
+At 2048x2048 input that is 16384 tokens, and the stock transformers forward
+materializes the full [heads x N x N] attention map (plus a second copy when
+the decomposed relative-position bias is added, plus softmax) -- a ~19 GB
+transient that is the pipeline's memory peak. transformers has no SDPA path
+for this architecture (the decomposed bias must be added to the raw scores).
+
+Computing the same attention in query-row chunks is mathematically exact:
+each output row depends on its own scores row only, and the decomposed bias
+factorizes over query rows (``rel_h`` is indexed by query row, ``rel_w`` by
+query column), so both slice cleanly. Peak drops from O(N^2) to O(N * chunk).
+"""
+
+
+def patch_vitdet_attention(model, chunk: int) -> int:
+    """Replace VitDetAttention.forward with a query-chunked equivalent.
+
+    ``chunk`` is the number of query tokens per block; sequences no longer
+    than ``chunk`` (e.g. the 14x14 window blocks) keep the original one-shot
+    path. Returns the number of attention modules patched.
+    """
+    import types
+
+    import torch
+    from transformers.models.vitdet.modeling_vitdet import VitDetAttention, get_rel_pos
+
+    def chunked_forward(self, hidden_state, output_attentions=False):
+        batch_size, height, width, _ = hidden_state.shape
+        tokens = height * width
+        if output_attentions or tokens <= chunk:
+            return VitDetAttention.forward(self, hidden_state, output_attentions)
+
+        # Identical prologue to the stock forward.
+        qkv = (
+            self.qkv(hidden_state)
+            .reshape(batch_size, tokens, 3, self.num_heads, -1)
+            .permute(2, 0, 3, 1, 4)
+        )
+        queries, keys, values = qkv.reshape(3, batch_size * self.num_heads, tokens, -1).unbind(0)
+
+        if self.use_relative_position_embeddings:
+            rel_h = get_rel_pos(height, height, self.rel_pos_h)
+            rel_w = get_rel_pos(width, width, self.rel_pos_w)
+            r_q = queries.reshape(batch_size * self.num_heads, height, width, -1)
+
+        keys_t = keys.transpose(-2, -1)
+        rows = max(1, chunk // width)
+        out = torch.empty_like(queries)
+        for h0 in range(0, height, rows):
+            h1 = min(h0 + rows, height)
+            q_chunk = queries[:, h0 * width : h1 * width, :]
+            scores = (q_chunk * self.scale) @ keys_t
+            if self.use_relative_position_embeddings:
+                rq = r_q[:, h0:h1]
+                rel_height = torch.einsum("bhwc,hkc->bhwk", rq, rel_h[h0:h1])
+                rel_width = torch.einsum("bhwc,wkc->bhwk", rq, rel_w)
+                scores = (
+                    scores.view(batch_size * self.num_heads, h1 - h0, width, height, width)
+                    + rel_height[:, :, :, :, None]
+                    + rel_width[:, :, :, None, :]
+                ).view(batch_size * self.num_heads, (h1 - h0) * width, tokens)
+            out[:, h0 * width : h1 * width, :] = scores.softmax(dim=-1) @ values
+
+        hidden = out.view(batch_size, self.num_heads, height, width, -1)
+        hidden = hidden.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1)
+        return (self.proj(hidden),)
+
+    count = 0
+    for module in model.modules():
+        if isinstance(module, VitDetAttention):
+            module.forward = types.MethodType(chunked_forward, module)
+            count += 1
+    return count

bgfilter/settings.py

@@ -134,6 +134,12 @@ class ModelSettings:
     # hardware (any modern GPU, or a CPU with AVX512-BF16/AMX) — otherwise it
     # falls back to fp32 with a warning. Segmentation always stays fp32.
     precision: str = "fp32"
+    # Query-chunked global attention (exact math, bitwise-identical output):
+    # the VitDet backbone's 4 global blocks materialize an N^2 attention map —
+    # ~19 GB at 2048x2048. Blocks seeing more than this many tokens compute it
+    # in query chunks of this size instead, capping the transient at
+    # O(N * chunk): 19.8 -> 4.0 GB for ~15% more ViTMatte time. 0 disables.
+    attn_query_chunk: int = 2048
 
 
 @dataclass(frozen=True)

bgfilter/vitmatte_infer.py

@@ -25,6 +25,10 @@ class ViTMatteRunner:
         self.dtype = resolve_dtype(self.torch, self.device, settings.precision)
         self.model.to(device=self.device, dtype=self.dtype)
         self.model.eval()
+        if settings.attn_query_chunk > 0:
+            from .attn_chunk import patch_vitdet_attention
+
+            patch_vitdet_attention(self.model, settings.attn_query_chunk)
 
     def _resolve_device(self, requested: str):
         if requested == "cuda" and not self.torch.cuda.is_available():

configs/default.yaml

@@ -13,6 +13,10 @@ model:
   # AVX512-BF16 CPU), else falls back to fp32. --precision bf16 sets bf16 here
   # AND on the BiRefNet models below in one go.
   precision: fp32
+  # Query-chunked global attention (exact, bitwise-identical output): caps the
+  # VitDet global blocks' N^2 attention transient at O(N * chunk) — 19.8 -> 4.0
+  # GB at 2048x2048 for ~15% more ViTMatte time. 0 = stock one-shot attention.
+  attn_query_chunk: 2048
 
 segmentation:
   # enabled: true  = single-segmenter pipeline (one seg model drives the trimap)