perf: in-place ops cut memory-read VRAM

[Borda]

Summary

Applies the VRAM-optimization pattern from "Optimizing PyTorch VRAM usage for mcbyte-tracker (16GB → under 8GB)" to Cutie's memory-read hot path: replacing chained non-in-place tensor arithmetic with in-place and fused ops, which the article found to be the actual VRAM win — not del + torch.cuda.empty_cache(), which that article measured as both slower and less stable than doing nothing.

get_similarity() in cutie/model/utils/memory_utils.py had the exact same anti-pattern as the article's baseline: several full-size B×N×HW intermediate tensors (a_sq, two_ab, b_sq) built via chained -a_sq + two_ab - b_sq, recomputed every frame against the full accumulated memory bank. A follow-up pattern search across cutie/ surfaced five more sites with the same shape, now fused or converted to in-place forms as well.

Test-first throughout: characterization tests (forward-numeric snapshots, backward-gradient snapshots, and CPU-vs-MPS parity) were written and confirmed green against the unmodified source before any production edit landed, per this repo's refactor discipline. See the companion commit test(cutie): guardrail tests before VRAM refactor.

Changes

• cutie/model/utils/memory_utils.py — get_similarity(): fuse the 2 * scalar multiply into two_ab.mul_(2), chain .sub_() for the a_sq/b_sq combination instead of allocating separate -a_sq + two_ab - b_sq intermediates. do_softmax() dense branch: .exp_() on the throwaway sub() output instead of a separately-allocated torch.exp() tensor.
• cutie/model/group_modules.py — MainToGroupDistributor (method='muladd'): x * g + g → torch.addcmul(g, x, g), one fused kernel instead of two allocations.
• cutie/model/channel_attn.py — CAResBlock.forward: x * w + downsample(r) → torch.addcmul(downsample(r), x, w).
• cutie/model/big_modules.py — KeyProjection shrinkage branch: d_proj(x) ** 2 + 1 → d_proj(x).pow(2).add_(1).
• cutie/model/cutie.py — encode_image/encode_mask normalization: (image - mean) / std → image.sub(mean).div_(std). sub() always copies, so the caller's frame tensor is never mutated — only the fresh copy is divided in place.
• cutie/inference/memory_manager.py — streaming object-value accumulation now uses .add_() in place instead of building a new_acc temp and reassigning both memory slices.
• scripts/bench_vram_mps.py (new) — standalone MPS peak-memory and wall-clock bench for the memory-read path. Not a pytest gate: MPS allocator readings are noisy run-to-run, so this is evidence to eyeball, not a CI assertion.

Why not every op went in-place

get_similarity/do_softmax are shared by the training path (cutie.py → get_affinity, needs autograd) and the inference path (memory_manager.py, no grad) — an in-place op that overwrites a tensor autograd needs for backward silently breaks training. One candidate op was reverted from in-place back to out-of-place after the backward-parity guardrail caught it: x_exp.div_(x_exp_sum) in do_softmax's dense branch raised RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation — Exp's own backward needs its output preserved, so that final normalization stays out-of-place (documented inline). The guardrail test, not assumption, decided every in-place/out-of-place call in this PR.

MPS bench (this MacBook, torch 2.13, N=32400 accumulated memory entries)

	Peak MPS memory	Wall-clock / iter
Before	2090.70 MiB	25.07 ms
After	1066.70 MiB	23.91 ms

~49% peak-memory cut, ~5% faster — consistent with the article's finding that structural in-place redesign outperforms aggressive cache clearing.

Test coverage

• tests/test_memory_utils.py — extended with 21 characterization cases: get_similarity (qe present/None, ms present/None, add_batch_dim, CK=1/N=1 boundaries), do_softmax (dense branch, top_k == N boundary, inplace identity, tie-break, large-magnitude stability), readout (T=1/T>1), plus forward and backward-gradient snapshots for every op on the training path.
• tests/test_model_ops_vram.py (new) — forward/backward parity for CAResBlock, MainToGroupDistributor(muladd), KeyProjection(need_s), and a caller-tensor-safety test for the cutie.py normalize idiom.
• tests/test_mps_parity.py (new) — CPU-vs-MPS numeric parity for get_similarity, do_softmax, and CAResBlock, actually executed on Apple GPU hardware (not skipped) — the first MPS coverage in this repo's test suite, which was previously CPU-only end to end.
• Full suite: 88 passed. ruff check cutie/ tests/ scripts/: clean.

Known limitation (not fixed here, out of scope)

do_softmax's top-k branch already breaks .backward() on main, regardless of the inplace flag — an unconditional values.exp_() on the torch.topk output. This is a pre-existing bug, unrelated to this refactor, pinned as-is by test_do_softmax_top_k_branch_backward_raises_inplace_version_error. No live impact today since the top-k path is only exercised inference-side (memory_manager.py, no grad); would need a separate fix if the top-k path is ever used with requires_grad.

[Copilot]

Pull request overview

This PR reduces VRAM usage in Cutie’s memory-read hot path by rewriting several chained tensor arithmetic patterns into fused and/or in-place forms (primarily in get_similarity() / do_softmax()), and adds guardrail tests (including first-run MPS execution) plus an optional MPS benchmarking script to validate performance and numerical parity.

Changes:

• Replace allocation-heavy arithmetic chains with in-place / fused ops in memory similarity + softmax and a few hot model blocks (addcmul, mul_, sub_, add_, div_, exp_).
• Add extensive characterization + backward guardrail tests for affected ops, plus CPU-vs-MPS numerical parity tests.
• Add a standalone MPS VRAM + wall-clock benchmark script for the memory-read path.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
cutie/model/utils/memory_utils.py	In-place/fused math in get_similarity() and dense do_softmax() to reduce intermediate tensor allocations.
cutie/model/group_modules.py	Use torch.addcmul in MainToGroupDistributor(method='muladd') to fuse multiply+add.
cutie/model/channel_attn.py	Use torch.addcmul in CAResBlock.forward residual path to fuse multiply+add.
cutie/model/big_modules.py	Use pow(2).add_(1) for shrinkage computation to reduce temporaries.
cutie/model/cutie.py	Normalize inputs via sub().div_() to keep a copy while reducing allocations.
cutie/inference/memory_manager.py	Switch streaming accumulation to in-place .add_() updates.
tests/test_memory_utils.py	Add reference-implementation checks, boundary cases, and backward guardrails for similarity/softmax/readout math.
tests/test_model_ops_vram.py	Add forward “golden” snapshots + backward finiteness tests for refactored/fused model ops and normalization caller-safety.
tests/test_mps_parity.py	Add CPU-vs-MPS parity tests (and a device-execution assertion) for key math ops.
scripts/bench_vram_mps.py	Add an optional standalone MPS memory/time bench for the memory-read path.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.