# CUDA memory pool retention after ANS loading and MAPED merging - 2026-10-04

Paths and line numbers below are those of 2026-10-04, before the restructure.
`_compact/streamed.py` (with `retain`) is now `resident/cuda/counts.py`,
`io/_streamed.py` is `io/encoded.py`, `io/backends/cuda/precision.py` is
`resident/cuda/precision.py`, `io/_precision.py` is `io/precision.py`,
`io/_array_resident.py` is `io/arrays.py`, `io/_streamed_file.py` is
`io/qem.py`, and `backends/cuda/packed.py` and `backends/cuda/_ans.py` are
gone. [Downstream sync](downstream-sync.md) maps every module.

## Question

After MAPED merges seven ANS tilts on CUDA and closes them (`release_tilts=True`),
the process keeps about 8 GiB more than the 5.93 GiB merge, and `free_all_blocks()`
cannot return it. Right after loading, CuPy's default pool holds 10.55 GiB for
7.76 GiB of encoded tilts. Why, and can released residents return their memory
without changing a decoded value, slowing loading or the merge, or touching the
Apple MPS path? Can the complete MAPED run stay under 16 GB?

## Setup

- One NVIDIA RTX PRO 6000 Blackwell (96 GB), shared with other jobs. Python 3.14,
  CuPy 14.2.0, PyTorch 2.13.0.
- quantem.gpu `19c2b0a0` (base `fbfadf03`); MAPED `fac7b8ef` (merge regions of
  up to 4096 scan positions, 62 regions with `crop=True`).
- Seven native Arina tilts, 512 x 512 scan, 192 x 192 detector, uint32 files
  whose counts fit uint16. `io.load(io.discover(folder))` builds seven runtime
  ANS residents (7.76 GiB). `MAPED.merge` stores the merge with
  `io.load(MergedRegions, dtype="scaled_uint16")` (5.93 GiB, 62 regions).
- Process memory: `nvidia-smi --query-compute-apps` for the process, polled every
  50 ms (peaks shorter than that can be missed). CuPy: `MemoryPool.used_bytes()`
  and `total_bytes()`. Torch: `max_memory_allocated()` and `max_memory_reserved()`.
- Bit-exactness: SHA-256 over every array kept by every tilt and by the merged
  result, before and after, on the same MAPED commit.
- Timings: interleaved before/after processes; minima and medians from runs
  where the GPU showed 0 % utilization from other jobs.

## Cause

Both residents are `StreamedCounts` (`src/quantem/gpu/_compact/streamed.py`):
the tilts through `io/_streamed.py:175` (`source.append(raw)`), the merge
through `io/backends/cuda/precision.py:200` (`encode_ans`). Before this change,
`StreamedCounts.append` allocated the kept arrays from CuPy's default pool right
next to the scratch of the same chunk (line numbers before the change):

| Line | Array | Lifetime | Size per 16384-scan tilt chunk |
|---|---|---|---|
| 167 | `scratch` (encode) | scratch | 1.21 GB |
| 168 | `sizes`, `states` | scratch | 4.7 MB each |
| 169 | `models` | kept | 1.2 MB |
| 183 | `offsets` | kept | 4.7 MB |
| 186 | `payload` | kept | about 70 MB |
| 231, 237, 240 | `widths`, `starts`, `words` (`_index`) | kept | small |
| 116 | `valid` | kept | 36 KB |

The decoded counts (`raw`, 1.21 GB) and the decompressor buffers live in the same
pool. CuPy splits a request out of any free cached block and can only return a
block when every piece of it is free. A 70 MB payload or a 36 KB `valid` array
carved out of a freed 1.2 GB scratch block therefore keeps the whole block
reserved for the life of the resident; after `close()`, the released pieces
merge back but the block stays pinned by pieces of the next tilt or of the merge.
`load_h5_ans` freed the pool only at the start of each load, never at the end.

## Fix

1. `_compact/streamed.py`: new `retain(arrays)` copies the six kept arrays of a
   chunk into one exact allocation outside the pool
   (`cp.cuda.using_allocator()`), packed at 256-byte offsets and returned as
   views. The allocation is freed with its last view, so `close()` returns it
   at once. `append` and `index_encoded` keep `retain(...)`; `valid` and the
   per-device entropy tables are also allocated outside the pool. All
   intermediates stay pooled scratch.
2. `io/_streamed.py`: the CUDA H5 load returns its pooled scratch when it ends
   (an `ExitStack` callback that runs after the reader and corrector close).
3. `io/_precision.py`: a CUDA precision load returns cached blocks before
   converting and its own scratch after it.
4. `io/_array_resident.py`: the float32 path, which drops the count index after
   `append`, calls `retain` on the three code arrays so the dropped index bytes
   are not kept inside the shared allocation.

Decoded values and encoded bytes are unchanged; no MPS code path changed.

## Results

Memory (GiB, `nvidia-smi` for the process, same MAPED commit):

| Step | Before | After |
|---|---|---|
| After loading 7 tilts | 11.10 | 8.42 |
| Peak while loading | 11.10 | 12.03 to 12.08 |
| After loading, pool freed | 9.97 | 8.42 |
| CuPy pool after loading, used / held | 7.76 / 10.55 | 0.00 / 0.00 |
| Merge peak, tilts kept (cold / warm) | 16.94 / 17.04 | 17.08 / 17.19 |
| Route 2 peak, fresh process (3 runs) | 16.71 to 16.73 | 16.98 to 17.19 |
| After route 2 (tilts closed) | 15.08 | 6.77 |
| Plus pool freed, tilts dropped, both caches | 14.66 | 6.75 |
| CuPy pool after release, used / held | 5.93 / 13.87 | 0.00 / 0.00 |

Time (seconds):

| Step | Before | After |
|---|---|---|
| Load 7 tilts, warm page cache, quiet GPU, minimum | 3.49 | 3.56 |
| Load 7 tilts, same, typical | 3.49 to 3.58 | 3.56 to 3.63 |
| Encode and index inside the load | 0.66 to 0.68 | 0.66 to 0.67 |
| Warm merge, 10 runs, median | 2.05 | 2.09 |
| Route 2 in a fresh process (includes compiling), median of 3 | 3.83 | 3.81 |

Bit-exactness: the SHA-256 of all seven tilts' kept arrays and of the merged
result's kept arrays are identical before and after; merged RMSE 0.005432755805088518
in both.

What sets the merge peak after the fix (route 2, 62 regions):

| Part | GiB |
|---|---|
| Seven tilts plus the complete merge | 13.69 |
| Torch reserved by the merge kernel (allocated peak 1.60) | 1.90 |
| CuPy conversion scratch, pool high-water | 0.80 |
| CUDA context after compiling | about 0.7 |
| Rounding of exact allocations (112 tilt chunks, 62 merge chunks) | about 0.2 |

MAPED region size, measured with the fix only (route 2 in a fresh process,
one run each, includes compiling):

| Region frames | Regions | Time (s) | Peak (GiB) | After (GiB) | Torch reserved (GiB) | Merged RMSE |
|---|---|---|---|---|---|---|
| 4096 (MAPED default) | 62 | 3.84 | 16.98 | 6.77 | 1.90 | 0.0054328 |
| 2048 | 124 | 4.09 | 16.46 | 6.86 | 1.07 | 0.0053741 |
| 1024 | 248 | 6.35 | 15.51 | 7.00 | 0.68 | 0.0053380 |

Allocation granularity measured on this GPU (200 exact allocations each):

| Requested | Device memory per allocation |
|---|---|
| 0.6 MB | 0.69 MB |
| 1.2 MB | 2.09 MB |
| 4.7 MB | 6.29 MB |
| 47 MB | 48.2 MB |
| 70 MB | 71.3 MB |

## Conclusion

Released tilts now return their memory: after route 2 the process holds the
merge plus the context (6.77 GiB) instead of 15.08 GiB, and after loading it
holds the tilts (8.42 GiB) instead of 11.10 GiB. Output is bit-identical. Cost:
about 0.06 s on a 3.5 s load and about 0.04 s on a 2.05 s warm merge (one exact
allocation and copy per chunk, plus scratch allocated again for each tilt now
that it is returned).

The run peak does not drop; it rises by 0.1 to 0.5 GiB. Before, the merge
scratch and payloads were carved out of the 2.8 GiB of load scratch the pool
could never return; now that scratch is returned and the merge allocates its
own. The peak is 13.69 GiB of live data plus the merge working set in two
allocators (torch 1.90, CuPy 0.80) plus the context. Staying under 16 GB needs
a smaller MAPED merge working set: 1024-frame regions reach 15.51 GiB (merge
about 1.65 times slower, RMSE changes because each region keeps its own
calibration), 2048-frame regions 16.46 GiB.

The load peak rises from 11.10 to 12.03 GiB: each tilt's 16384-scan staging
(decoded counts 1.21 GB, encode scratch 1.21 GB, decompressor buffers) now sits
beside the tilts already loaded instead of being partly reused for their
payloads. It stays below the merge peak.

## Rejected ideas

- **One exact allocation per array** (`using_allocator()` around each kept
  `cp.empty`): every allocation above about 1 MiB is rounded up to whole 2 MiB
  pages. With six arrays per chunk this cost 0.35 GiB after loading (8.77 GiB)
  and the warm merge peaked at 18.04 GiB. One allocation per chunk fixes it.
- **A dedicated `cp.cuda.MemoryPool()` per resident**, as
  `backends/cuda/packed.py` does: a fresh pool still requests each allocation
  from the device exactly (same rounding), destroying the pool object frees
  memory that views may still use, and release must also free the pool.
- **Only freeing the pool at the end of each load**: kept arrays are still
  carved out of scratch blocks, so `close()` still cannot return them.
- **Keeping cached scratch for the merge to reuse** (no free before converting):
  the 1.13 GiB block left by the detector summaries stays beside the merge;
  warm merge peak 17.41 to 17.61 GiB.
- **Writing payload and index straight into the exact allocation** instead of
  copying: needs every size before the compact and pack kernels, so `_index`
  would be split in two; the copy costs about 0.05 ms per chunk on a quiet GPU.
- **`cudaMemcpyAsync` instead of an elementwise copy**: equal on a quiet GPU
  (0.225 and 0.227 ms for 70 MB including the allocation), slower under load.
- **Encoding each merge region in smaller chunks** to shrink the encode scratch:
  more chunks means more rounding, and it changes the saved chunk layout.

## Not covered

- Exact allocations bypass CuPy's free-and-retry on out-of-memory; on a nearly
  full GPU a kept array can fail where the pool would have freed cached blocks
  first.
- Other CUDA residents still allocate from the default pool:
  `CudaANSResidentCounts` and `CudaPackedResidentCounts` (`backends/cuda/_ans.py`,
  used by saved `.qem` reopening and by global-scale precision), and
  the arena in `io/_streamed_file.py`.
- The MPS path was not changed and was not rerun here.