cubecl-cuda 0.10.0: deterministic reproduction of the stale-page panics from #1384 (`Memory page N doesn't exist`, CUDA 13.3)

[Lunasciel]

PR #1384 (stale page index after ExclusiveMemoryPool::cleanup renumbers pages; the lookup panics on the dispatch worker) was closed as no longer reproducible.

I hit the same panics deterministically, every run, in a single process with no concurrent requests — a plain forward pass or greedy generation — where #1384 needed a batching server at concurrency ≥ 2.

8B model, during generation:

thread 'DSD-0-0' panicked at cubecl-cuda-0.10.0/src/compute/stream.rs:101:
called `Result::unwrap()` on an `Err` value:
couldn't find resource for that handle: Memory page 0 doesn't exist
WARN Task failed: Any { .. }

Sites (which of them fire is workload-dependent):

• src/compute/stream.rs:101 — handle_cursor → get_cursor(binding.memory.clone()).unwrap(): survivable (the cursor is cross-stream sync metadata), but fires ~3× per generated token; the unwind overhead drops throughput ~50–100×.
• src/compute/server.rs:701 — command.resource(binding).expect("Resource to exist."): fatal, exit 101.
• src/compute/server.rs:124 — command.reserve(size).unwrap() in initialize_memory: fatal; the reserve().unwrap() knock-on class that fix(memory): resolve stale page lookups instead of panicking #1384's A/B table counted 794×.

Environment:

• RTX 4090 (24 GB), driver 610.43.02, CUDA 13.3, cudarc 0.19.8,
• cubecl-cuda 0.10.0, burn 0.21.0, native bf16, fusion disabled,
• Linux (Arch-based).

Trigger: a dense-transformer forward (36 layers, vocab = 151936, bf16) creating many short-lived device bindings under real pool pressure — layer-streamed weights (upload → matmul → drop, hidden = 2560) or a ~16 GB resident model (hidden = 4096) computing a chunked full-vocab logits projection (chunks up to n = 65536).

Ruled out: autotune (cache-warm run panics identically), max_streams = 1 (no effect), fusion (disabled), bf16 math (surviving runs produce numerically correct output).

An isolated binding-churn program (same pattern, 1280 uploads) does not reproduce it — the stale index seems to need cleanup renumbering under real pool pressure (~16 GB resident on a 24 GB card).

[laurigates]

I reproduced this deterministically from a very different workload and traced the root cause. Short version up front: on current main the cubecl-runtime layer already fails gracefully here — the remaining crash is cubecl-cuda unwrapping that graceful error. So there are two separable pieces, and I have a ready branch for one of them.

Deterministic repro (single process, no concurrency)

Training a weight-only-int8 (QLoRA-style) LoRA on a ~12.8B-parameter MMDiT diffusion model — RTX 4090 (24 GB), burn 0.21 / cubecl-cuda 0.10, driver 580.126.18 / CUDA 13.0, Linux. The forward transiently dequantizes ~260 frozen base weights (upload → matmul → drop) under ~20 GB of pool pressure. That is the same "many short-lived bindings under real pool pressure" trigger @Lunasciel describes, reached from a workload with no batching server and no concurrency — so the trigger looks general to that allocation pattern, not specific to this model or to int8. Every run panics at cubecl-cuda/src/compute/server.rs:124 and stream.rs:101 with Memory page N doesn't exist.

Root cause (confirms @Lunasciel's and #1384's analysis)

ExclusiveMemoryPool::cleanup and SlicedPool::cleanup drain(..) the pages vector and renumber the survivors, writing the new index into each page's current slice.handle only. A binding whose descriptor has diverged from that handle — orphaned by a bind, but kept alive by the cross-stream cursor lookup — retains a stale page index. After enough frees that index no longer maps to the page it was minted for.

What main already does, and what it doesn't

MemoryManagement::find on main has the init == 0 guard and the id-mismatch check ("a stale location must surface as NotFound, never as another allocation's slice"). So a stale binding on main resolves to a graceful NotFound — it does not silently read or write another allocation's buffer. There is no data corruption on main; the id-check closes that.

What is left is the crash: cubecl-cuda turns that NotFound into a dispatch-thread abort by unwrapping it — stream.rs:101 .unwrap(), server.rs:124 .unwrap(), server.rs:701 .expect(). That unwrap is what a user hits today. Removing the renumbering (below) does not by itself remove that unwrap; the two are separable.

Two separable pieces (matching @nathanielsimard's "fix the root cause, then relax how we fail" in #1384)

1. Root-cause removal — ready. Stop renumbering page indices on cleanup, so a stale index can never arise in the first place. Candidate keeps pages in stable slots: a deallocated page becomes a None tombstone (reused by later allocations before the vector grows) instead of drain-and-renumber, in both pools; the SlicedPool max_pages cap counts live pages only. Branch: laurigates/cubecl@fix/pool-page-renumber-stale-index. It builds on main and the cubecl-runtime suite shows no new failures versus unmodified main HEAD (the 9 memory_manage failures I saw are present on stock main too — unrelated bytes_reserved/capacity accounting). Caveat on testing: that suite is the runtime abstraction layer; I have not run the CUDA-backend tests (no second NVIDIA host here). My CUDA-path evidence is the downstream repro above.
2. Relaxed failure — needs your design input. With the renumbering gone, the cubecl-cuda .unwrap()/.expect() sites could recover from a NotFound instead of aborting the dispatch thread. I did not touch this — how you'd want a stream to recover a missing cursor is your call.

I found all this working on a downstream pure-Rust burn trainer, so I'd rather hand you the analysis and a branch than drop a PR you didn't ask for. Concretely: I'm ready to open a PR for piece 1 now if you'll take that approach — should I? And is piece 2 something you'd want in the same PR or a separate discussion?

[laurigates]

Update — I traced the mechanism further and tried two fixes against the deterministic repro (int4 QLoRA forward/backward on a ~12.8B MMDiT, RTX 4090, single process).

A safe partial fix (clears the survivable cursor panic)

On top of a tombstone pool (stable page indices — free a page by leaving a None slot instead of drain-and-renumber, so a live page's index never moves), two changes make the survivable site recover instead of aborting the dispatch thread:

• MemoryManagement::find: replace assert_eq!(slice.handle.descriptor(), binding.descriptor()) with a graceful NotFound when the binding diverged.
• cubecl-cuda handle_cursor: .unwrap() → .unwrap_or(u64::MAX) (the cursor only gates a cross-stream wait; u64::MAX errs toward adding the wait — never skips one).

This is memory-safe rather than error-hiding because the tombstone pool guarantees find returns either the right live page or a None slot — never the wrong survivor's slice, which I think addresses @nathanielsimard's concern on #1384 about masking. It reliably clears the stream.rs cursor-lookup panic in my repro. (Happy to open a PR upstream if useful.)

What it does NOT fix, and the refined root cause

The fatal server.rs sites (reserve in initialize_memory, and resource for a kernel data buffer) still fire under real pressure. The mechanism:

1. a transient T is bound (dequant of a frozen weight), a matmul (x, T) is queued,
2. T's host handle drops → its page is is_free,
3. a later alloc OOMs → the reclaim (cleanup(true)) tombstones T's page,
4. the queued matmul is finally dispatched → find(T) → tombstoned/reused slot → panic.

Because bindings resolve lazily at dispatch, the reclaim races the command queue. I tested a cuStreamSynchronize before the reclaim (i.e. drain GPU work, then tombstone) — it did not help: a stream sync drains submitted work but not cubecl's command queue of kernels not yet dispatched to the stream, so the reclaim still tombstones a page a queued kernel will later look up.

Suggested directions

• Command-queue drain before reclaim — flush/dispatch all pending commands so every queued binding is resolved before cleanup(true) frees any page.
• Cursor-aware cleanup — pages already carry slice.cursor; only free a is_free page whose cursor <= completed_cursor. Least-invasive; needs the completed cursor exposed to cleanup.

Note the live working set fits well under 24 GB here (the ~24 GB peak is the transient ratchet, since SlicedPool::cleanup only reclaims on the explicit path), so a correct reclaim should fit — this reads as reclaim-correctness, not a genuine OOM. Glad to test candidate fixes against the repro.