fix(gpu): survive transient VRAM pressure — recover resident-table declines, close an R2 corruption race - #914
Open
ColoCarletti wants to merge 5 commits into
Open
fix(gpu): survive transient VRAM pressure — recover resident-table declines, close an R2 corruption race#914ColoCarletti wants to merge 5 commits into
ColoCarletti wants to merge 5 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.
ColoCarletti
force-pushed
the
gpu-vram-correctness-fixes
branch
from
August 7, 2026 18:23
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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).