feat: add backend-agnostic caching memory pool - #36
Open
nkAtchen wants to merge 6 commits into
Open
Conversation
`std::unordered_map` is node-based, so tracking a live block called `operator new` on every insert and `operator delete` on every erase -- a host heap round trip on the exact path a device allocator exists to keep short. `BlockTable` is an open-addressing table with linear probing and tombstones, held at a load factor of 1/2 in one flat array. Insert and erase touch no allocator at all once the array has grown. `ReleaseCached` also stops holding the lock across upstream frees: it now settles the stats and detaches the blocks under the lock, then frees outside it. `cudaFree` implicitly synchronizes the whole device, so holding a mutex across one stalls every other thread for the duration.
Two pieces the arena pool in the next commit needs, separated out because neither is specific to it. `NodeArena` is a bump allocator with a free list, for the fixed-size nodes a pool's own bookkeeping needs. An allocator that calls the host heap to describe its own state has a dependency it cannot control the latency of, and under a lock that latency is every thread's. `PointerTable` is the previous commit's open-addressing table generalized over the mapped type. `MemoryPool` keeps its own `BlockTable`, which maps to a fixed record type and stays private to that header; this one is the reusable form the arena needs, mapping a pointer to an arbitrary `Value`.
`MemoryPool` calls upstream once per cache miss. On a device that is
hundreds of microseconds, so a workload whose sizes keep missing pays
it over and over -- and a size-class cache misses structurally on the
shapes inference produces, because a KV cache slab that grows with
sequence length lands in a different class every step.
`ArenaMemoryPool` requests large backing stores and slices them, so
upstream calls scale with the high-water footprint rather than with the
allocation count. Each backing is an address-ordered chunk list, so
splitting and coalescing make a released block reusable at any size,
not only at the size class it was allocated with.
Three things make it fast enough for the win to survive contention:
- Coalescing is deferred. `Deallocate` files a chunk into its
exact-size bin and returns; merging runs in one pass only when a
request cannot be served from what is already indexed. A loop that
cycles through a handful of sizes never merges at all.
- The free index is in two parts: 128 exact-size fast bins with an
occupancy bitmap, plus a tree ordered by (size, address). An
occupied exact-fit bin is already the best fit, so the common case
never touches the tree.
- In front of both sits a per-thread cache of exact-size blocks. A
hit takes no lock. `Deallocate` parks into the releasing thread's
cache, so the next allocation of that size on that thread skips
both the acquisition and the best-fit lookup.
Growth follows a doubling ramp to a cap, past which further backings
are added at the cap size; a single request larger than the cap gets
its own oversize backing. The first non-oversize backing is resident so
a steady small workload keeps a warm arena. Everything else is subject
to automatic shrink, with two scans of hysteresis -- `cudaFree`
implicitly synchronizes the whole device, so thrashing it from the
allocation path costs more than holding the memory one more round.
The pool does not track streams: a block is reusable the instant
`Deallocate` returns, and coalescing means it may come back at a
different offset and size. Callers must ensure device-side access has
completed first. This is the same contract `MemoryPool` has, documented
here because coalescing makes violating it corrupt an unrelated
allocation rather than merely reuse a block early.
Neither pool replaces the other. The arena wins where upstream is
expensive or the shapes defeat size classes; the size-class pool wins
on large-block recycling and keeps a shorter critical section at low
thread counts. Both stay.
Neither pool is universally better, so the claim in the previous commit is only worth as much as the measurement behind it. This is that measurement. `perf_allocator_matrix` runs one set of workloads across every allocation strategy the current backend offers -- the backend allocator itself, `MemoryPool`, `ArenaMemoryPool`, and `cudaMallocAsync` where it exists -- through the real dispatch API. Arms share a surface, so each benchmark is written once and instantiated per arm. The workloads are chosen to be the ones a general pool benchmark cannot make: shapes that straddle the size-class boundary, large-block recycling at sizes that are not multiples of the large granularity, growth to a gigabyte-scale high-water mark, trim cost split into traversal and upstream frees, thread counts past 8, a layer-by-layer inference sequence, a multi-threaded one, latency quantiles, and a long random-lifetime fragmentation run with a large-block probe at the end. Two design notes that matter for reading the output. Timings are reported next to the exact upstream call counts that cause them, because on a device the call count *is* the timing and it carries no machine noise. And `cuda_async` is marked stream-ordered everywhere it appears: its release does not wait for pending device work, so it offers a weaker guarantee than the other three arms and its ratios are not a drop-in speedup. `scripts/compare_allocators.py` configures and runs both the host and device builds and pivots the rows into the three comparisons that have distinct answers -- direct vs arena, pool vs arena, vendor pool vs arena. Serialized rather than parallel because `generated/` is written into the source tree at configure time and its contents depend on which backends are enabled. The arena's config scale is part of each row's identity, so the reduced host config and the production device config are never averaged into one series. `arena_vs_pool` measures the same two designs in one process against its own upstream stubs, one of them a calibrated busy-wait standing in for a synchronous `cudaMalloc`. That is the only way to see the arena's amortization on a machine with no device, and it keeps both arms on one CPU at one thermal state.
The new arena-based memory pool has been introduced in previous commits. This removes the old implementation to avoid confusion. BREAKING CHANGE: memory_pool.h is no longer available. Users must migrate to arena_memory_pool.h.
nkAtchen
force-pushed
the
feat/runtime-memory-pool
branch
from
August 11, 2026 01:22
7e7850c to
c0a59c9
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.