fix(gpu): survive transient VRAM pressure — recover resident-table declines, close an R2 corruption race - #914
fix(gpu): survive transient VRAM pressure — recover resident-table declines, close an R2 corruption race#914ColoCarletti wants to merge 11 commits into
Conversation
… on an R2 miss The device-only gate is a static predicate over a dynamic dispatch: it cannot mirror every reason the device R2 path might decline (parts count, kernel eligibility, transient errors, shapes a new workload brings), and each miss was a hard abort that deadlocked the epoch pipeline — DECODE on the synthetic workload, then a second table on the real-block bench. Instead of excluding tables one by one, treat the resident handles as the source of truth: on a miss, download the main/aux LDEs back to host, clear the device-only flag, and continue on the host path. Slower for that table, never wrong; the abort remains only when the handles themselves cannot serve the data. gpu_device_only_downgrades() counts recoveries so a persistently-missing condition still gets mirrored into the gate.
…declines A transient CUDA OOM on the resident aux LDE was a hard prove failure: the resident build leaves no host aux trace to fall back to. A device drain releases the concurrent VRAM peaks, so one retry usually keeps the table fully resident; if it still declines, download the resident aux trace (and the main LDE when the table is device-only) and continue host-backed. The drain before dropping the resident buffer also keeps kernels enqueued by the failed attempt from reading pool memory reused by a concurrent table.
…ption race Concurrent device R2 windows under VRAM pressure can transiently produce a fully wrong H for one or two tables while every input stays correct (rerunning the same chain on the same resident inputs matches the host), yielding a proof that fails the composition check. Serializing only the constraint-eval + decompose window across tables eliminates it; commits and host arms stay parallel, and the windows overlap rarely enough that the lock is near-free. LAMBDA_VM_GPU_SERIALIZE_R2=0 lifts the lock to bisect further or once the underlying race is found.
The device R2 path only exists for the d=2 quotient split; a table with any other composition bound (DECODE proves with num_parts == 1) would skip it entirely and hard-abort on its device-only trace.
b5e3e52 to
8f62d7b
Compare
A mixed state (one commit fell back to CPU while the other stayed device-only) left the recovery refusing to proceed: it treated a missing device handle as fatal even when that side already had a valid host copy. Only the missing side is downloaded now, and the R3 host-arm guards check the buffer they are about to read instead of the table-wide flag.
download_ext3_columns came along in a cherry-pick but its only consumer (the cross-check post-mortem) ships separately; dead code under the cuda feature.
|
/ai-review |
Codex Code Review
|
AI ReviewPR #914 · 3 changed files Findings
Status column reflects the verdict from the verifier: deepseek-verifier (openrouter/deepseek/deepseek-v4-pro). AI-004: R2 serialization guard uses .unwrap() and could poison on panic
Claim
Evidence Line 69: Suggested fix Use AI-005: New GPU_DEVICE_ONLY_DOWNGRADES counter not reset in reset_all_gpu_call_counters
Claim The PR adds a new public counter Evidence Line 1439 declares Suggested fix Add AI-006: GPU_DEVICE_ONLY_DOWNGRADES counter double-counts per table
Claim The GPU_DEVICE_ONLY_DOWNGRADES counter increments in both materialize_lde_trace_host() and materialize_aux_trace_host(), but a single table downgrade can trigger both (main + aux), counting one table as two downgrades. Evidence materialize_lde_trace_host() calls GPU_DEVICE_ONLY_DOWNGRADES.fetch_add(1) at line 1480. materialize_aux_trace_host() calls the same at line 1547. In prover.rs aux stage (line 3470-3485), when aux materialize succeeds and device_only was true, it also calls download_main_lde_row_major which triggers materialize_lde_trace_host logic, causing both counters to increment for the same table. Suggested fix Either rename counter to track sides not tables, or use a per-table flag to count only once per table. Simplest: increment only in materialize_lde_trace_host since that handles both sides, and remove increment from materialize_aux_trace_host. AI-011: materialize_lde_trace_host aux slab download skips the length validation that the main path performs
Claim The aux arm of materialize_lde_trace_host indexes Evidence download_main_lde_row_major (line 1553) checks Suggested fix After AI-014: materialize_aux_trace_host mutates trace.aux_table/num_aux_column before the final device-sync check, then returns false on sync failure
Claim materialize_aux_trace_host assigns Evidence Lines 1600-1611: Suggested fix Perform AI-015: set_num_rows doc comment was hijacked by the inserted set_host_data
Claim The pre-existing doc comment describing Evidence Lines 535-538 are the original Suggested fix Move the Reviewer Lanes
Verification Lanes
Native Codex and Claude reviews run separately and post their own comments. They are not included in this structured provenance report. Discarded candidates (9) — rejected by the verifier
Raw lane outputs, candidates, final issues, and model metrics are uploaded as workflow artifacts. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| /// `main_data.len()` — the caller supplies it from the device handle's | ||
| /// `lde_size` instead. | ||
| #[cfg(feature = "cuda")] | ||
| /// Install downloaded host buffers on a device-only table and clear the |
There was a problem hiding this comment.
Medium — misplaced insertion: the new fn was pasted between set_num_rows's doc comment / #[cfg(feature = "cuda")] and its signature. Two consequences:
set_host_datanow carries the "Override the LDE row count…" doc (and a duplicated#[cfg]).set_num_rowsloses its#[cfg(feature = "cuda")]and is compiled unconditionally on non-cuda builds.
It still compiles, but the gating is accidental. Move the new function below set_num_rows (or above the /// Override the LDE row count block) and drop the duplicated attribute.
| // decompose, where H is born) eliminates a transient whole-buffer | ||
| // H corruption seen under concurrent R2 windows on VRAM pressure. | ||
| // The commit and every host arm run outside the lock. | ||
| let _r2_serial_guard = crate::gpu_lde::r2_serialize_guard(); |
There was a problem hiding this comment.
Medium (perf) — the host fallback runs inside the global lock. The guard covers the whole block, so the None => arm (download_comp_h_to_field D2H of the full H, then Self::decompose_and_extend_d2, a host iFFT + two LDEs over lde_size) is serialized across all tables. That arm is exactly the path taken under VRAM pressure — the case this PR targets — so the "windows rarely overlap, lock is near-free" measurement won't hold there.
Suggest keeping only the device window under the lock, e.g. capture h_dev/the decompose result inside the block, drop(_r2_serial_guard) (or end the block) before the download + host decompose.
Minor, same spot: the guard is taken before number_of_parts == 2 is evaluated, so tables that never enter the device R2 path (num_parts == 1) still queue on the lock. Cheap to reorder the check ahead of the acquire.
| static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); | ||
| if *ENABLED.get_or_init(|| !std::env::var("LAMBDA_VM_GPU_SERIALIZE_R2").is_ok_and(|v| v == "0")) | ||
| { | ||
| Some(LOCK.lock().unwrap()) |
There was a problem hiding this comment.
Low — lock poisoning masks the first failure. The critical section contains assert!s and expects (and per-table prove threads are joined via thread::scope). If one table panics while holding this lock, every other table then panics with a PoisonError from this unwrap(), burying the original cause under N unrelated panics. Since the guarded state is (), poisoning carries no information:
| Some(LOCK.lock().unwrap()) | |
| Some(LOCK.lock().unwrap_or_else(|e| e.into_inner())) |
| /// [`materialize_lde_trace_host`]). Nonzero values mean the device-only gate | ||
| /// admitted a table some dispatch later declined — correct but slower, and | ||
| /// worth mirroring the missing condition into the gate. | ||
| pub(crate) static GPU_DEVICE_ONLY_DOWNGRADES: AtomicU64 = AtomicU64::new(0); |
There was a problem hiding this comment.
Low — counter is write-only. GPU_DEVICE_ONLY_DOWNGRADES isn't added to reset_all_gpu_call_counters() (every other counter is), and gpu_device_only_downgrades() has no caller, so as it stands it's dead API.
Two small follow-ups make it earn its keep: add it to the reset list, and assert gpu_device_only_downgrades() == 0 in prover/tests/cuda_path_integration.rs::gpu_device_only_residency_fires_and_verifies — that test already resets the counters and is the natural place to catch a gate/dispatch drift that silently starts downgrading every table (correct proof, lost residency win, no signal today).
| // empty together for this table. | ||
| #[cfg(feature = "cuda")] | ||
| let device_only = Self::device_only_for(*air, domain); | ||
| let mut device_only = Self::device_only_for(*air, domain) |
There was a problem hiding this comment.
Low — stale docs around this now-asymmetric gate. The aux gate is no longer the same value as the main commit's, which invalidates two comments:
- the line just above ("Same gate as the Round 1 main commit"), and
device_only_for's doc (~L1014): "Derived purely from the AIR + domain so the round-1 main-commit and aux-commit closures compute the identical value" — the aux side now also ANDs in the presence of the main device handle.
Worth restating the new contract: aux may be more conservative than main (never less), which is what keeps the mixed GPU-aux/CPU-main state out.
Also stale after this PR: try_expand_leaf_and_tree_ext3_row_major_keep_dev's eprintln comment in gpu_lde.rs ("This path has no CPU fallback … so the caller hard-aborts") — the caller now retries and downgrades.
|
Review — GPU VRAM-pressure recovery + R2 serialization Read the three changed files against the surrounding device paths ( What I checked and believe is right
Findings (inline)
Minor, not blocking
The |
The fallback arm (download H, host iFFT + LDEs) executed under the serialization lock, so under VRAM pressure — exactly when that arm runs — it serialized every other table's device window behind pure CPU work. The lock now covers only the device eval + decompose + the H download; the host decompose and every host arm run outside it. The lock is also acquired only for d=2 tables (the others never enter the device path).
set_host_data had been inserted between set_num_rows' doc and its signature, stealing its doc comment and un-gating it from the cuda feature; the serialization lock now recovers from poisoning instead of cascading PoisonErrors over the original panic; the downgrade counter joins reset_all_gpu_call_counters and the device-only residency test asserts it stays at zero on the happy path; stale comments about the aux gate mirroring the main gate rewritten with the actual contract.
| static ENABLED: OnceLock<bool> = OnceLock::new(); | ||
| static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); | ||
| if *ENABLED.get_or_init(|| !std::env::var("LAMBDA_VM_GPU_SERIALIZE_R2").is_ok_and(|v| v == "0")) | ||
| { | ||
| // The guarded state is (), so a panic while holding the lock carries | ||
| // no information — recover instead of burying the original panic | ||
| // under a cascade of PoisonErrors from every other table. | ||
| Some(LOCK.lock().unwrap_or_else(|e| e.into_inner())) | ||
| } else { | ||
| None | ||
| } |
There was a problem hiding this comment.
Consulting the env var is not that expensive, and unless we set it inside the code it won't change during the lifetime of the binary:
| static ENABLED: OnceLock<bool> = OnceLock::new(); | |
| static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); | |
| if *ENABLED.get_or_init(|| !std::env::var("LAMBDA_VM_GPU_SERIALIZE_R2").is_ok_and(|v| v == "0")) | |
| { | |
| // The guarded state is (), so a panic while holding the lock carries | |
| // no information — recover instead of burying the original panic | |
| // under a cascade of PoisonErrors from every other table. | |
| Some(LOCK.lock().unwrap_or_else(|e| e.into_inner())) | |
| } else { | |
| None | |
| } | |
| static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); | |
| if std::env::var("LAMBDA_VM_GPU_SERIALIZE_R2") != Ok("0") { | |
| // The guarded state is (), so a panic while holding the lock carries | |
| // no information — recover instead of burying the original panic | |
| // under a cascade of PoisonErrors from every other table. | |
| Some(LOCK.lock().unwrap_or_else(|e| e.into_inner())) | |
| } else { | |
| None | |
| } |
…ansposes (#921) * fix(gpu): harden the downgrade download recovery Three fixes on the device-only downgrade path, all in the graceful degradation function whose whole point is to avoid a hard abort. - The aux branch of `materialize_lde_trace_host` sliced the downloaded slabs without checking their length, so a short download would panic inside the recovery instead of degrading. Both sibling download paths already validate (`download_main_lde_row_major` checks `col_major.len() != m * lde`, `materialize_aux_trace_host` checks `raw.len() != rows * cols * 3`); this adds the matching check. - Restore the `len/capacity % 3` guard the other two ext3 `from_raw_parts` sites carry, spelled `is_multiple_of` because clippy's `manual_is_multiple_of` rejects the older form here. - The failure error claimed "host aux trace is empty" on a path where that is false: when the aux download succeeded and the follow-up main-LDE download failed, the host aux trace had just been populated. Track which recovery step failed and name it. Control flow unchanged. * perf(gpu): parallelize the downgrade recovery transposes Both conversions in the recovery path were single-threaded nested loops over the full LDE: the col-major -> row-major main transpose in `download_main_lde_row_major`, and the de-interleaved-slabs -> row-major interleaved aux conversion in `materialize_lde_trace_host`. For MEMW at LDE 2^20 those are a 411 MB and a 327 MB buffer respectively, walked with a strided access on one core. Both now follow the existing idiom in `trace.rs` ("Parallel col-major -> row-major transpose"): parallelize over OUTPUT row chunks with `par_chunks_exact_mut`, so every element is still written exactly once and no unsafe is involved. The index math is unchanged -- chunk `r` of width `m` is `row_major[r * m + c]`, and chunk `r` of width `m * 3` sub-chunked by 3 is `interleaved[(r * m + c) * 3 + k]` -- because the layout was verified against the kernels. Gated on the `parallel` feature with the sequential loop kept for builds without it, and skipped when `m == 0` since `chunks_exact_mut(0)` panics. These loops run on a scheduler driver thread holding no locks, so rayon is safe here, unlike the pinned-staging unpack in math-cuda.
…antics (#920) This branch turned two of the device-only hard-aborts into downloads that recover and continue host-backed, but the surrounding docs still describe the old contract: "every host read hard-aborts", "the prove aborts loudly", "a mis-gate panics one of the guards". Rewrite those to say what the code now does — R2 and the R1 resident-aux commit recover and bump GPU_DEVICE_ONLY_DOWNGRADES, R3/R4 still abort, and the R3 guards check the individual buffer so mixed states are legal. Also correct the R2 lock comment (it serializes submission, not execution, for device-only tables), note that the numeric gate is not the complete predicate on its own, broaden the downgrade counter's doc to cover resident-aux declines on tables that were never device-only, and drop the false "only" from materialize_lde_trace_host's failure list. Comments, doc comments, two assertion message strings and one doc-comment run command (--test-threads=1, matching the Makefile target). No behavior changes.
diegokingston
left a comment
diegokingston
left a comment
There was a problem hiding this comment.
We should check whether it makes sense to land 911 before
4 participants
Main today fails ~1/3 of real-block proves at epoch 2^22 on a 32 GB GPU (reproduced: 2/6 runs on 8b88a8d), and under sustained VRAM pressure can silently emit an invalid proof (~1/15 runs). This PR fixes the three failure modes without touching the perf envelope.
Hard failures under transient CUDA OOM. The resident-aux LDE and the device R2 path had no fallback for a device-only table: the aux decline was a hard prove error ("resident aux LDE failed; host aux trace is empty") and an R2 device miss was a hard abort that kills the prover thread. Now the aux decline drains the device and retries once (the transient peak is usually gone, so the table stays resident); if it still declines, the resident aux trace and the main LDE are downloaded and the table continues host-backed. The R2 miss recovers the same way (materialize the resident LDEs, continue on the host evaluator). Both paths log the downgrade with the table name.
Silent invalid proofs. Concurrent device R2 windows under VRAM pressure can transiently produce a fully wrong composition H for one or two tables while every resident input stays correct (rerunning the same chain matches the host recompute), yielding a proof that fails the verifier's composition check. Serializing only the constraint-eval + decompose window across tables eliminates it (30/30 clean cycles vs ~5 failures/30 without); commits and host arms stay parallel and the windows overlap rarely enough that the lock is near-free (~0% on real block, ~1-2% on small synthetic workloads). LAMBDA_VM_GPU_SERIALIZE_R2=0 lifts the lock. The underlying race is still being hunted; the lock is the safe default until then.
Device-only gate tightening. The device R2 path only exists for the d=2 composition split; tables with any other bound (DECODE proves with num_parts == 1) now never enter device-only mode instead of aborting downstream. Mixed states (GPU aux commit with a CPU main commit) are also excluded.
Diagnostics (env-gated, zero cost when off). LAMBDA_VM_GPU_XCHECK=1 runs the verifier's composition consistency check inside the prover per table (~µs) and, on a failure, a post-mortem that recomputes each device stage on host, reports the corruption shape and reruns the device chain to distinguish a transient race from corrupted resident inputs. LAMBDA_VM_GPU_FORCE_DOWNGRADE=1 exercises the recovery paths end to end. These are how the fixes above were found and validated.
Validation on a 5090 (32 GB): 30/30 real-block e22 prove+verify cycles under reduced VRAM budget with zero corruption and all transient OOMs recovered (vs main failing 2/6 and the pre-fix branch failing ~5/30); full GPU test suite green (898 tests).