fix(cuda): make opportunistic cache cleanup transport-aware

[shuangwu]

Summary

Provide an opt-in policy for opportunistic CUDA cache cleanup, such as optional per-wave housekeeping. Do not suppress intentional memory release.

This revision removes the earlier broad rewrite of existing cleanup sites. Model loading, checkpoint resume, simulator shutdown, and other in-tree direct torch.cuda.empty_cache() calls are unchanged from upstream main.

Contract

• Applications use maybe_empty_cuda_cache() only when the flush is safe to omit. It returns whether PyTorch's flush was invoked, not the number of bytes released.
• NCCL/UCXX producer/consumer setup and generic prefetch startup suppress these optional helper requests before background/native work starts. The startup transition is serialized with an already-running helper flush.
• Optional suppression remains sticky after close or failed setup: an empty local queue does not prove global quiescence. It does not prevent intentional memory release.
• Intentional release retains its direct PyTorch call. Its owner must establish a safe boundary first; this PR neither implements a distributed quiescence protocol nor certifies existing release sites as safe.
• No PyTorch monkey patch, automatic delayed flush, force/reset flag, CUDA synchronization, or new cross-rank barrier is introduced by skipped optional requests.
• Custom transports bypassing built-in startup call suppress_opportunistic_cuda_cache_cleanup() before launching native/background work.

The practical prevention applies only once downstream optional per-wave callers adopt the helper or remove their flushes. No existing in-tree cleanup is guessed to be optional. This is intentionally not a blanket fix for every concurrent allocator-cleanup interaction.

Validation

• 248 focused cache-policy, NCCL/UCXX, transport strategy, prefetch and rollout-loop regression tests passed.
• Tests cover startup/flush races, concurrent optional requests, failed setup, fork lock handling, suppression before background startup, and preservation of intentional/direct cleanup. Regression checks retain existing model, checkpoint and simulator release call sites.
• Revised two-rank CPU/Gloo canary passed: 20 exact-payload transfers per rank while optional cleanup was skipped, followed by intentional direct cleanup after all transfers and distributed teardown completed.
• Revised two-rank GPU/NCCL canary passed: 20 exact-payload transfers per rank and more than 1,100 optional requests skipped per rank, followed by successful intentional direct cleanup after all transfers and distributed teardown. Both ranks and the job exited zero. Source hashes matched this revision.
• Repository formatting, lint and whitespace checks passed.

Portable reproduction: PYTHONPATH=. torchrun --standalone --nproc-per-node=2 tests/cuda_cache_cleanup_canary.py --backend nccl (or --backend gloo). Use a finite outer timeout. The canary tests real background communication and an explicit post-teardown boundary, not the original deadlock or a full training workload.

Scope and tradeoffs

Optional cleanup can leave allocator-reserved memory higher, but intentional memory handoff is no longer silently disabled. Applications must not mark a required release as optional when another allocator depends on it.

Independent of #747 (fault containment), #751 (owned teardown), and #754 (receive-memory budgets). Does not control allocator-internal reclamation, third-party engine cleanup, or native fault recovery. The same optional policy covers NCCL and UCXX; no live RDMA claim is made.

See docs/payload_transport/cuda_cache_cleanup.md for adoption requirements and safe-boundary obligations.

[shuangwu]

Closing after design review. The broad suppression policy changes legitimate memory-release behavior; the narrowed opt-in helper adds API and lifecycle state without improving on removing unnecessary per-wave flushes in the application. We will use that simpler mitigation instead. Transport fault containment (#747) and owned teardown (#751) remain independent and necessary. A future coordinated-cleanup proposal should be driven by a demonstrated need for memory handoff while transport remains active.