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

[NicolasRouquette]

Summary

Rebased onto the latest main (the Lean 4.32 refactor), as a single commit. 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>.

Review changes addressed

• Self-contained strict decimal parser (no strtoull) — the cap is parsed by a small
  hand-rolled loop that accepts exactly a non-empty digit string fitting in size_t; malformed or
  overflowing values are rejected with a warning and leave the cache unbounded, instead of being
  silently misread (strtoull would take 1MiB as a 1-byte cap and saturate overflow to
  ULLONG_MAX). This also removes the only strto* reference in the CUDA objects, so the branch
  links against the Lean toolchain's glibc with no compatibility layer — no dependency on feat(cuda): bundle + auto-link the glibc≥2.38 isoc23 link shim #18.
• Thread-safe one-time env init — the cap is read from the environment under pthread_once
  (torchlean_cuda_cache_byte_cap_init), so concurrent first callers cannot race on the parse and the
  value is published before any caller observes it. (The prior function-local static assigned after
  its declaration was a data race and depended on -fthreadsafe-statics.)
• Overflow-safe cap comparison — cache_bytes + incoming > cap is reformulated as
  incoming > cap || cache_bytes > cap - incoming, so the sum is never formed and cannot wrap
  size_t (and cap - incoming is only evaluated when incoming <= cap).
• Control subprocess pinned to 0 — the fork helper always sets TORCHLEAN_CUDA_CACHE_CAP_BYTES
  in the child's environment (capped = 1048576, control = 0), so neither child inherits a stray
  cap from the parent and the control run genuinely exercises unbounded growth.

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 the cap explicitly 0 the same workload caches the full 8 MiB;
• malformed — 1MiB is rejected and the cache stays unbounded (a prefix-parsing reader would
  cache nothing under a 1-byte cap);
• overflow — a value past the native word is rejected rather than truncated or saturated.

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.

Verification

CPU-stub lake build && lake exe nn_tests_suite: green, all four fork cases pass. Full CUDA build +
suite re-run on an RTX A4500 (CUDA 13.0), built deliberately without any glibc compatibility
object: links clean and all curated tests pass, with all four fork cases confirming the cap and the
rejection semantics on the real allocator.

[Robertboy18]

Hey Nicolas, thanks! This one still addresses a real gap: low-level or custom CUDA paths can retain an unbounded amount of released memory in the reuse cache.

We just merged a major Lean 4.32 refactor, so could you first rebase this onto the latest main? Before merging, could you also make three small changes:

• make the one-time environment initialization thread-safe;
• make the cap comparison overflow-safe;
• explicitly set the control subprocess cap to 0, so it cannot inherit a cap from the parent environment.

Once that is updated, we'll run the full CUDA stress suite and review it for merge. Thanks!

[NicolasRouquette]

Thanks! Rebased onto the latest main (the Lean 4.32 refactor) — it stands as a single commit; the only merge touch-ups were the test-suite renames (NN.Entrypoint.Tensor → NN.Tensor, tensorND! → tensorOfList!) and slotting the probe alongside the new BlockMask curated test. All three requested changes are in:

• Thread-safe one-time environment init. torchlean_cuda_cache_byte_cap() now reads the environment under pthread_once (a dedicated ..._init runs on a single thread while others block, and the parsed value is published before any caller observes it). The previous function-local static assigned after its declaration was a data race between concurrent first callers and also leaned on -fthreadsafe-statics; pthread_once removes both.

• Overflow-safe cap comparison. g_torchlean_cuda_cache_bytes + incoming > cap is reformulated as incoming > cap || g_torchlean_cuda_cache_bytes > cap - incoming. The sum is never formed, so it cannot wrap size_t; cap - incoming is only evaluated once incoming <= cap, so it never underflows.

• Control subprocess cap pinned to 0. The fork helper in runCacheCapTest now always sets TORCHLEAN_CUDA_CACHE_CAP_BYTES explicitly in the child's environment — the capped run to 1048576, the control run to 0 — so neither child can inherit a stray cap from the parent. Previously the control child left the variable unset and would silently observe a parent-set cap, masking the uncapped growth it exists to demonstrate.

CPU-stub lake build && lake exe nn_tests_suite is green, and I re-ran the full CUDA build + stress suite locally on an RTX A4500 (scripts/checks/check.sh --cuda): the fork test's capped and control cases both pass, with the cap confirmed as the binding constraint. Ready for your CUDA run whenever you like.

[NicolasRouquette]

Heads-up on ordering: on glibc ≥ 2.38 this PR does not link on its own — the new torchlean_cuda_cache_byte_cap_init parses TORCHLEAN_CUDA_CACHE_CAP_BYTES with strtoull, whose nvcc host pass emits an __isoc23_strtoull reference the older link-time glibc can't bind (ld.lld: undefined symbol: __isoc23_strtoull). #18 supplies that symbol (and fixes the same failure for any other strto* use in the CUDA sources), so please merge #18 first (or together). Verified together on Ubuntu 24.04 + CUDA 13.0: the cache-cap regression passes on-GPU (8 MiB of returns bounded to a 1 MiB cache; uncapped caches the full 8 MiB).

[Robertboy18]

Thanks, Nicolas — the cache bound, overflow-safe accounting, and subprocess test all look good. We’d like to merge this. Since #18 is needed for the strtoull compatibility issue, could you rebase this onto main after #18 is merged? Once the final CUDA tests are green on the rebased version, this should be ready to go.

[Robertboy18]

Thanks, Nicolas. The cache accounting and cap logic look good. I think we should avoid making this depend on the global __isoc23_strto* shim in #18, though. Could you replace the strtoull call with a small strict decimal parser here, rejecting malformed and overflowing values? That keeps this fix self-contained and avoids adding a broad glibc compatibility layer. With that change and the CUDA stress test green, I’d be happy to merge this.

[Robertboy18]

Thanks, Nicolas! This is merged now. The bounded CUDA cache and telemetry fit the direction of the runtime well, and I’m doing a final local cleanup pass so the behavior and configuration are documented clearly.