H3 Memory Probe (ledger + allocator trace, alpha)

A VRAM instrument for the H3 diffusion model, from the pack's VRAM Lab work. Long de-ropes are memory monsters - the README's numbers say a de-rope at d_max 4 on an 8–12 second clip is roughly 200k packed tokens, and the stock H3 block materializes its fused QKV and SwiGLU tensors for the whole sequence (8.6 and 15.4 GiB at that length), which OOMs 24 GB cards and, at 1376×768, even a 96 GB one. When you're fighting that, guessing at the culprit wastes more GPU hours than the node costs. H3 MemoryProbe tells you exactly where the bytes go.

How it works

It wraps the diffusion model and records a per-block, per-phase ledger of PyTorch's allocator counters at every phase boundary - start, block i's kv/attn/mlp, end - when it's fed by H3 Streamed Blocks upstream. Stock (unstreamed) blocks contribute start/end only. That gives you a JSONL ledger (ledger.jsonl) showing which block and which phase was holding what, no device sync and negligible cost. If the ledger alone isn't enough, record_history_forwards turns on the caching allocator's full alloc/free trace with Python stacks for that many model forwards, then dumps snapshot.pickle and a standalone trace.html - a hoverable timeline where you hover a band to see the stack that allocated it. That's the torch memory_viz style trace, rendered locally.

Inputs, all required:

• model - the model to instrument; it passes through, so it sits inline in the model wire.
• tag (default probe) - run label; files land in out_dir/<tag>_<time>/.
• ledger (on) - the per-forward JSONL. Keep it on for the cheap numbers.
• record_history_forwards (0) - 0 = off. A full H3 forward at 200k tokens is ~20k allocator events, a few percent of runtime while recording, and it's a lot of data - turn it on only when you actually need the trace.
• max_entries (300000) - ring size for the allocator trace.
• out_dir (default output/h3_memprobe) - relative to the ComfyUI working directory. Not /tmp.

Outputs: model (the wrapped model, continues the wire) and a report string.

When to reach for it

The README's small-card story is the use case: streaming the DiT in token chunks on a fenced-down card, with --reserve-vram and --fast-disk, took a 702-frame pass that OOMed a 96 GB card and rendered it on a 16 GB card in a 32 GB machine. If you're tuning that kind of stack - or just got an OOM and want evidence instead of vibes - this is the node that turns "why did it die" into a ledger.

Install

cd ComfyUI/custom_nodes
git clone https://github.com/matlowai/ComfyUI-MAINodes

or ComfyUI Manager → ComfyUI-MAINodes, restart. Category MAINodes/VRAM Lab. No extra deps.

Common gotchas

The ledger is only as good as what runs through the patched model - it says per-block/per-phase with H3 Streamed Blocks upstream, so don't expect granular block phase splits if you're running stock blocks. The trace mode is genuinely heavy; leaving record_history_forwards above 0 in a graph you run daily will slow every render and grow your disk. Turn it on, capture, turn it off. The out_dir guidance mirrors the pack's general rule: avoid /tmp, it's a RAM disk on most Linux installs and you want the probe data to survive.