feat(cuda): scope eager device-memory reclamation to the tape (withStep/withForward)

[NicolasRouquette]

Summary

Reworked in response to the review here: instead of a public device arena that frees buffers by
allocation epoch — which can free a device buffer while its Lean Buffer is still reachable
(use-after-free if a value or alias is missing from keep) — scoped device-memory reclamation now
lives inside the eager session/tape ownership model, keyed on tape membership.

Two exception-safe combinators on EagerSession bracket an eager phase and reclaim the tape at its
boundary:

• withStep — one training step. resetTape opens it; the body runs forward/backward/optimizer;
  on exit every tape-owned temporary is released and the tape snapshot retired, so the device working
  set stays flat across steps.
• withForward — one inference/eval forward pass. It transfers the body's result off the tape
  (e.g. getValue, which downloads to host) before releasing the forward temporaries.

The three open-coded CUDA optimizer-step reclamation sites in Trainer now use withStep.

Why this is safe (addresses the arena's use-after-free)

Reclamation is keyed on tape membership, not allocation epoch. The only buffers a scope frees are
tape node values and workspace (cleanup) — which are runtime-private: callers hold
TensorRef/Param handles, never a raw Buffer. So a scope can never invalidate a buffer the caller
still holds. Persistent parameter mirrors are preserved by the paramsByLeaf role check.

withStep is also correct on the error path: if the body raises before the optimizer re-mirrors
the parameters, the trainable leaf buffers may still be the live mirrors, so it falls back to the
conservative release (releaseCudaTapeNonParamValues, which preserves every parameter value) before
re-raising — it never frees a live mirror.

This removes the class of hazard that blocked the public withCudaArena API; the arena implementation
(torchlean_cuda_arena.h, the epoch registry, the buffer arena_reg field, the UAF detector) is
dropped.

Scope vs. the original problem

The refactored main already bounds the training path (each optimizer step releases the tape);
this generalizes that into a reusable, exception-safe scope and gives the inference/eval path the
same guarantee. The one case the arena uniquely served — a raw foldl of Buffer → Buffer with no
tape and no IO — is intentionally out of scope: real models go through the tape, and for a genuine
in-op reduction the existing releaseThen/releaseManyThen (threaded through the carried result) is
the safe tool.

Tests

NN/Tests/Runtime/Cuda/Stress.lean (in nn_tests_suite, wired into the stress suite):

• runLeakBoundStress (folded from the closed leak-bound regression): drives a small eager Buffer
  workload through the explicit release discipline for two equal blocks of steps and asserts, via
  Buffer.allocatorStats, that live allocations and live bytes are flat across step count and return
  to baseline. Runs on the CPU stub.
• runEagerScopeStress: checks withForward end-to-end on a CPU session (every build), and — on a
  real CUDA session — that the device working set is flat across forward scopes and returns to
  baseline. The CUDA loop is skipped on the CPU-stub build, which rejects .cuda sessions.

Verification

CPU-stub lake build && lake exe nn_tests_suite && lake lint: green. The CUDA-specific assertions
(runEagerScopeStress Part 2) require a GPU — run scripts/checks/check.sh --cuda.

[Robertboy18]

Thanks, Nicholas. The memory-pressure problem here is real, but I don't think we can merge the arena API in its current form. It can free storage while the corresponding Lean Buffer is still reachable, so missing a value or alias from keep creates a use-after-free. The optional detector does not make that safe as a public API.

With the new runtime refactor, I think this should instead be implemented inside the eager session/tape ownership model. Temporary buffers should belong to a step or tape node, workspace should be released through WithWorkspace and node cleanup, and surviving outputs should be transferred automatically. That way no reachable public buffer is invalidated, and nested or concurrent sessions do not share a global arena.

If you're interested in revising it in that direction, we can keep discussing the design here.