Skip to content

feat(cuda): bound the device block reuse cache with an env byte cap#10

Open
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:cuda-cache-byte-cap
Open

feat(cuda): bound the device block reuse cache with an env byte cap#10
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:cuda-cache-byte-cap

Conversation

@NicolasRouquette

@NicolasRouquette NicolasRouquette commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Summary

Rebased onto the updated main. Adds an opt-in byte cap on the CUDA buffer reuse cache (the
exact-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 therefore
accretes device memory that liveBytes does not even see — a returned block is accounted as freed
before it is cached.

TORCHLEAN_CUDA_CACHE_CAP_BYTES (0 = unbounded, the prior behaviour exactly) bounds it: 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.

What's in it

  • A running cache_bytes total under the cache mutex backs the cap decision.
  • New telemetry field AllocatorStats.cacheBytes (extern torchlean_cuda_allocator_cache_bytes; the
    CPU stub returns 0); AllocatorStats.format gains cache=<MiB>.

Test

runCacheCapTest (in nn_tests_suite, runs on both the CUDA build and the CPU stub). The cap is read
once natively, so the test forks the suite binary per configuration:

  • capped — a 1 MiB cap bounds an 8 MiB return workload to a 1 MiB cache;
  • control — with no cap the same workload caches the full 8 MiB.

On CUDA it 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.

No longer stacks on the arena PR

Rebased to stand alone on main. It previously reused a test helper from the arena PR; that helper
is 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's capped / control cases
pass on the stub). The CUDA-only binding-constraint assertion requires a GPU — run
scripts/checks/check.sh --cuda.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant