feat(cuda): bound the device block reuse cache with an env byte cap#10
Open
NicolasRouquette wants to merge 1 commit into
Open
feat(cuda): bound the device block reuse cache with an env byte cap#10NicolasRouquette wants to merge 1 commit into
NicolasRouquette wants to merge 1 commit into
Conversation
The reuse cache (return/take of exact-size device blocks) grew without bound: a dropped buffer is returned to the cache, and the cache is only emptied on a cudaMalloc-failure flush or at process exit. A loop over many distinct sizes therefore accretes device memory that `liveBytes` does not see -- a returned block is accounted as freed before it is cached. Add an opt-in byte cap, `TORCHLEAN_CUDA_CACHE_CAP_BYTES` (0 = unbounded, the prior behaviour exactly): a returned block that would grow the cache past the cap is freed immediately -- after waiting on its completion event, exactly as the flush path does, so an in-flight kernel never reads freed memory -- instead of being cached. A running `cache_bytes` total under the cache mutex backs the decision and is surfaced as a new `AllocatorStats.cacheBytes` telemetry field (0 in the CPU stub, which keeps no cache and so has nothing to cap). Regression test `runCacheCapTest` (nn_tests_suite, runs on both the CUDA build and the CPU stub): the cap is read once natively, so it forks the suite binary per configuration. With a 1 MiB cap an 8 MiB return workload is bounded to a 1 MiB cache; with no cap (control) the same workload caches the full 8 MiB. On CUDA the test also checks the cap is the binding constraint (the cache fills to within one block of it). On the stub cacheBytes stays 0, so the cap bound holds trivially. The test's scratch helper (`buildCacheScratch`) is self-contained, so this change stands alone on the updated main rather than stacking on the arena PR. CPU-stub `lake build && lake exe nn_tests_suite`: green.
62fe748 to
50e2001
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
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.
Summary
Rebased onto the updated
main. Adds an opt-in byte cap on the CUDA buffer reuse cache (theexact-size device-block free list behind
take_cached_block/return_cached_block).The cache grows without bound: a dropped buffer is returned to it, and it is only emptied on a
cudaMalloc-failure flush or at process exit. A loop over many distinct buffer sizes thereforeaccretes device memory that
liveBytesdoes not even see — a returned block is accounted as freedbefore it is cached.
TORCHLEAN_CUDA_CACHE_CAP_BYTES(0 = unbounded, the prior behaviour exactly) bounds it: a returnedblock that would grow the cache past the cap is freed immediately — after waiting on its completion
event, exactly as the flush path does, so an in-flight kernel never reads freed memory — instead of
being cached.
What's in it
cache_bytestotal under the cache mutex backs the cap decision.AllocatorStats.cacheBytes(externtorchlean_cuda_allocator_cache_bytes; theCPU stub returns
0);AllocatorStats.formatgainscache=<MiB>.Test
runCacheCapTest(innn_tests_suite, runs on both the CUDA build and the CPU stub). The cap is readonce natively, so the test forks the suite binary per configuration:
On CUDA it also checks the cap is the binding constraint (the cache fills to within one block of it).
On the stub
cacheBytesstays 0, so the cap bound holds trivially.No longer stacks on the arena PR
Rebased to stand alone on
main. It previously reused a test helper from the arena PR; that helperis now self-contained (
buildCacheScratch), so this PR is a single independent commit.Verification
CPU-stub
lake build && lake exe nn_tests_suite: green (the fork test'scapped/controlcasespass on the stub). The CUDA-only binding-constraint assertion requires a GPU — run
scripts/checks/check.sh --cuda.