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>
This commit is contained in:
2026-07-06 15:26:59 +08:00
parent 9a9084b3d6
commit 2cc0ce8684
5 changed files with 94 additions and 1 deletions
+3 -1
View File
@@ -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`.