[PR #93] [MERGED] Implement a Memory Compiler #112

New issue

Closed

opened 2026-10-26 22:24:40 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:24:40 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-aimdo/pull/93
Author: @rattus128
Created: 8/24/2026
Status: ✅ Merged
Merged: 8/31/2026
Merged by: @rattus128

Base: master ← Head: prs/memory-compiler

📝 Commits (10+)

994bf8a Add native CUDA malloc graph compiler
ec46366 Expose malloc graph Python API
2fb6bf8 Test malloc graph allocation and nesting
258a4fe Test malloc graph replay contract
f8c3342 Test malloc graph ownership errors
25ae8cd Allow VBAR eviction during CUDA capture
c56e1b2 Evict VBARs for malloc graph pages
fa9ece9 Compile small CUDA allocations
bec233b Add malloc graph iteration API
43f21d7 Pause malloc graph recording

📊 Changes

55 files changed (+3323 additions, -12 deletions)

View changed files

📝 comfy_aimdo/control.py (+26 -0)
➕ comfy_aimdo/malloc_graph.py (+80 -0)
📝 src/control.h (+4 -0)
➕ src/malloc-graph.c (+1239 -0)
➕ src/malloc-rogue.c (+189 -0)
➕ src/malloc-rogue.h (+19 -0)
📝 src/model-vbar.c (+8 -2)
📝 src/plat.h (+6 -0)
📝 src/pyt-cu-plug-alloc-async.c (+39 -10)
➕ src/vmm-ref.c (+127 -0)
➕ src/vmm-ref.h (+39 -0)
➕ tests/malloc-graph-aliasing-cuda.py (+52 -0)
➕ tests/malloc-graph-assert-break-cuda.py (+24 -0)
➕ tests/malloc-graph-branch-join-cuda.py (+49 -0)
➕ tests/malloc-graph-branch-reuse-cuda.py (+26 -0)
➕ tests/malloc-graph-changed-size-cuda.py (+30 -0)
➕ tests/malloc-graph-cuda.py (+51 -0)
➕ tests/malloc-graph-deep-nesting-cuda.py (+36 -0)
➕ tests/malloc-graph-destructor-cuda.py (+26 -0)
➕ tests/malloc-graph-different-stream-cuda.py (+24 -0)

...and 35 more files

📄 Description

Implement a memory compiler.

The TLDR wins:

Absolute minimization of VRAM with 0 VRAM stuck in a reserved state on the VRAM peak
0 CUDA memory API calls during replay & 0 GPU traffic for memory ops - no racing the async allocator accounting mid step

This adds a compacting memory compiler for repeated execution of a model. When initialized and recording, it reserves a virtual address range, then allocates from that range and backs them with 8 MiB physical pages as needed. Freeing an allocation logically releases its physical pages for reuse but leaves the virtual allocation fully intact and still pointing to the physical pages. The allocator is then free to map those pages into other subsequent virtual segments of any size for other mismatched allocations, while the original range remains intact. At the end of recording, the virtual range remains fully mapped with its ranges aliasing physical memory. This means no GPU virtual address management is needed on replay. The replay is simply validation of the pattern and recall of the addresses as allocated (subject to graph break rules - see below).

Caches no unused allocations, keeping VRAM at the recorded peak.
Supports small allocations.
Supports nested and optional scopes.
Reports peak, virtual, and physical memory usage.
Evicts VBAR memory when additional physical pages are needed.

When an allocation trace diverges on later replay, the compiler reports a "graph break", reconstructs state from the scope snapshot, and records the new path as an alternate branch. Later executions matching that path replay normally; strict mode can instead reject breaks for validation for developers who wish to minimize memory layour churn, however graph breaks on execution are functional and performant. Graph breaks do not invalidate the whole graph (like the torch compiler) and only need to recompile the divergent section. The graph then becomes branching and both branches remain compiled concurrently, allowing flip-flop coding patterns etc.

More to come in the ComfyUI PR

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-aimdo/pull/93 **Author:** [@rattus128](https://github.com/rattus128) **Created:** 8/24/2026 **Status:** ✅ Merged **Merged:** 8/31/2026 **Merged by:** [@rattus128](https://github.com/rattus128) **Base:** `master` ← **Head:** `prs/memory-compiler` --- ### 📝 Commits (10+) - [`994bf8a`](https://github.com/Comfy-Org/comfy-aimdo/commit/994bf8a012ef48da3c4d63cea1c534d93f6be9bd) Add native CUDA malloc graph compiler - [`ec46366`](https://github.com/Comfy-Org/comfy-aimdo/commit/ec46366b1536173685fac96b17c7bfaa262201e1) Expose malloc graph Python API - [`2fb6bf8`](https://github.com/Comfy-Org/comfy-aimdo/commit/2fb6bf89812b84e0baac73c1cb14df36d1dc2085) Test malloc graph allocation and nesting - [`258a4fe`](https://github.com/Comfy-Org/comfy-aimdo/commit/258a4fe364a253477ccb9c2f593715004d6bfdf6) Test malloc graph replay contract - [`f8c3342`](https://github.com/Comfy-Org/comfy-aimdo/commit/f8c3342eeff641896e33a0ceaf2def033f3b866a) Test malloc graph ownership errors - [`25ae8cd`](https://github.com/Comfy-Org/comfy-aimdo/commit/25ae8cde540445e9d0b596b02bb4bd8bb1910e10) Allow VBAR eviction during CUDA capture - [`c56e1b2`](https://github.com/Comfy-Org/comfy-aimdo/commit/c56e1b2898ffa10ad97a7d2d0db1cd7d526544ec) Evict VBARs for malloc graph pages - [`fa9ece9`](https://github.com/Comfy-Org/comfy-aimdo/commit/fa9ece944da251d2b96d69264bbc8e4d1b6c325c) Compile small CUDA allocations - [`bec233b`](https://github.com/Comfy-Org/comfy-aimdo/commit/bec233b5b811276f8119f9322b7e020846ff89b0) Add malloc graph iteration API - [`43f21d7`](https://github.com/Comfy-Org/comfy-aimdo/commit/43f21d71331eaa8de892d4952fd50b46ce9ad913) Pause malloc graph recording ### 📊 Changes **55 files changed** (+3323 additions, -12 deletions) <details> <summary>View changed files</summary> 📝 `comfy_aimdo/control.py` (+26 -0) ➕ `comfy_aimdo/malloc_graph.py` (+80 -0) 📝 `src/control.h` (+4 -0) ➕ `src/malloc-graph.c` (+1239 -0) ➕ `src/malloc-rogue.c` (+189 -0) ➕ `src/malloc-rogue.h` (+19 -0) 📝 `src/model-vbar.c` (+8 -2) 📝 `src/plat.h` (+6 -0) 📝 `src/pyt-cu-plug-alloc-async.c` (+39 -10) ➕ `src/vmm-ref.c` (+127 -0) ➕ `src/vmm-ref.h` (+39 -0) ➕ `tests/malloc-graph-aliasing-cuda.py` (+52 -0) ➕ `tests/malloc-graph-assert-break-cuda.py` (+24 -0) ➕ `tests/malloc-graph-branch-join-cuda.py` (+49 -0) ➕ `tests/malloc-graph-branch-reuse-cuda.py` (+26 -0) ➕ `tests/malloc-graph-changed-size-cuda.py` (+30 -0) ➕ `tests/malloc-graph-cuda.py` (+51 -0) ➕ `tests/malloc-graph-deep-nesting-cuda.py` (+36 -0) ➕ `tests/malloc-graph-destructor-cuda.py` (+26 -0) ➕ `tests/malloc-graph-different-stream-cuda.py` (+24 -0) _...and 35 more files_ </details> ### 📄 Description Implement a memory compiler. The TLDR wins: - Absolute minimization of VRAM with 0 VRAM stuck in a reserved state on the VRAM peak - 0 CUDA memory API calls during replay & 0 GPU traffic for memory ops - no racing the async allocator accounting mid step This adds a compacting memory compiler for repeated execution of a model. When initialized and recording, it reserves a virtual address range, then allocates from that range and backs them with 8 MiB physical pages as needed. Freeing an allocation logically releases its physical pages for reuse but leaves the virtual allocation fully intact and still pointing to the physical pages. The allocator is then free to map those pages into other subsequent virtual segments of any size for other mismatched allocations, while the original range remains intact. At the end of recording, the virtual range remains fully mapped with its ranges aliasing physical memory. This means no GPU virtual address management is needed on replay. The replay is simply validation of the pattern and recall of the addresses as allocated (subject to graph break rules - see below). - Caches no unused allocations, keeping VRAM at the recorded peak. - Supports small allocations. - Supports nested and optional scopes. - Reports peak, virtual, and physical memory usage. - Evicts VBAR memory when additional physical pages are needed. When an allocation trace diverges on later replay, the compiler reports a "graph break", reconstructs state from the scope snapshot, and records the new path as an alternate branch. Later executions matching that path replay normally; strict mode can instead reject breaks for validation for developers who wish to minimize memory layour churn, however graph breaks on execution are functional and performant. Graph breaks do not invalidate the whole graph (like the torch compiler) and only need to recompile the divergent section. The graph then becomes branching and both branches remain compiled concurrently, allowing flip-flop coding patterns etc. More to come in the ComfyUI PR --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:24:40 +00:00

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-aimdo#112

No description provided.

Rows
Columns