Repository navigation
Allocate CPU output buffers outside torch's allocator - #20
Merged
EthanRosenthal merged 3 commits intoOct 6, 2026
Merged
Conversation
Decoded frame and sample buffers were allocated with at::empty, i.e. through torch's CPU allocator. On the aarch64 Linux wheels that allocator is mimalloc, which does not return a large block freed by a thread other than the one that allocated it: the pages are MADV_FREE'd but stay mapped and resident until the allocating thread exits. Data-loading pipelines decode on worker threads and drop the tensors elsewhere, so every decode left its whole output buffer (hundreds of MiB to GiBs) in resident memory. Allocate the buffers with posix_memalign instead and hand them to torch with at::from_blob plus a free() deleter, which runs on whichever thread drops the last reference; glibc munmaps large blocks immediately regardless of thread. tch-rs has no from_blob-with-deleter wrapper, so build.rs compiles a small C++ shim (csrc/) against the libtorch headers, located the same way torch-sys locates them. Tensor shape, dtype, strides and device are unchanged; CUDA outputs still go through torch. AVTENSOR_OUTPUT_ALLOCATOR=torch restores the previous behaviour. benchmarks/cross_thread_retention.py reproduces the problem: 12 decodes of a 20 s 720p clip (1.24 GiB each) on a 4-thread pool, freed on the main thread, left 14.9 GiB resident (all LazyFree) with at::empty on an aarch64 torch 2.11 build, and none with this change, at the same decode time.
Each tch release targets a single libtorch version (0.24 for torch 2.11, 0.26 for torch 2.13), and avtensor's source compiles against both. Accept the range in Cargo.toml so a build for torch 2.11 can select tch 0.24 with cargo update --precise instead of editing the manifest. Cargo.lock stays on 0.26, so the default build is unchanged.
|
| Commit | Scanned at | New | Resolved | Net |
|---|---|---|---|---|
fed29b8 |
2026-10-04 21:57 UTC | 0 | 0 | 0 |
d1cdd42 < |
2026-10-04 22:04 UTC | 0 | 0 | 0 |
Last scanned: d1cdd42 · 2026-10-04 22:04 UTC
|
| Commit | Scanned at | New | Resolved | Net |
|---|---|---|---|---|
fed29b8 |
2026-10-04 21:57 UTC | 0 | 0 | 0 |
d1cdd42 < |
2026-10-04 22:04 UTC | 0 | 0 | 0 |
Last scanned: d1cdd42 · 2026-10-04 22:04 UTC
RikHeijdens
approved these changes
Oct 5, 2026
There was a problem hiding this comment.
Gave it a quick look over. Overall this change seems reasonable to me. I really like that the allocator continues to be configurable.
I'm also curious how you found this bug? Was Claude able to diagnose and fix this issue in this manner on its own or did you have to handhold it along the way? If it mostly diagnosed and fixed this on its own then that's very impressive.
It might also be worthwhile to fix this issue upstream in libtorch / tch-rs since I reckon this bug may affect other use cases as well.
EthanRosenthal
added a commit
that referenced
this pull request
Oct 7, 2026
* Bump version to 0.5.2 Release plan_asset and richer probe_asset metadata (#18), exact in-window frames for end_time on B-frame streams (#19), and CPU output buffers allocated outside torch's allocator (#20). * Build the sdist against torch 2.13 The build-system torch pin was left at 2.11.0 when tch moved to 0.26.0 (torch 2.13) in #15, so isolated sdist builds installed a torch that doesn't match the locked tch. Align the pin with Cargo.lock and update the README requirements row to match.
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
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.
Problem
Decoded frame and sample buffers were allocated with
at::empty, i.e. torch's CPU allocator. On the aarch64 Linux torch wheels that allocator is mimalloc (embedded inlibc10.so), which doesn't return a large block freed by a thread other than the one that allocated it: the pages areMADV_FREE'd but stay mapped and resident until the allocating thread exits. Data-loading pipelines decode on worker threads and drop the tensors elsewhere, so every decode leaves its whole output buffer (hundreds of MiB to GiBs) resident.Change
posix_memalign(64-byte aligned, same as torch's CPU allocator) and handed to torch withat::from_blobplus afreedeleter. The deleter runs on whichever thread drops the last reference, and glibcmunmaps large blocks immediately on any thread.from_blob, sobuild.rscompiles a small C++ shim (csrc/tensor_from_blob.cpp) against the libtorch headers, located the same way torch-sys locates them (LIBTORCH_USE_PYTORCHandLIBTORCHmodes).at::empty. Shape, dtype, strides and device are unchanged.AVTENSOR_OUTPUT_ALLOCATOR=torchrestores the previous behaviour.Cargo.tomlacceptstch/torch-sys>=0.24, <0.27, so a build for torch 2.11 can select tch 0.24 withcargo update -p tch -p torch-sys --precise 0.24.0instead of editing the manifest.Cargo.lockstays on 0.26, so the default build is unchanged.Behavioural difference:
from_blobstorages aren't resizable, so.resize_()on a returned tensor now raises. Ops that allocate new storage (.clone(),.contiguous(),.pin_memory(),.to(...)) are unaffected.Measurements
12 decodes of a 20 s 720p clip (1.24 GiB output each) on a 4-thread pool, with the tensors dropped on the main thread: on aarch64 with torch 2.11,
at::emptyleft +14.85 GiB resident (allLazyFree); this change leaves +0.01 GiB, at the same decode time (1006 vs 1010 ms). On x86_64, where torch doesn't use mimalloc, both paths retain nothing and run at the same speed.Testing
src/alloc.rs: env parsing,empty-equivalence (shape/dtype/contiguity/alignment, writes through views), zero-size fallback, invalid shapes, and the deleter running exactly once when the last view drops on another thread.cargo fmt --check,cargo clippy,make testagainst torch 2.11 (tch 0.24, x86_64, FFmpeg 7.1.1): 100 passed, 3 ignored. torch 2.13 (tch 0.26, the lockfile default) is covered by CI.