Stop a failed CUDA graph commit from poisoning the encoder

[strayberry]

Proposed changes

Fixes #4326.

CommandEncoder::commit() exits before resetting its in-progress graph state when an exception is thrown — for example by the cache-thrashing check added in #2600, which throws from graph_cache_[graph_key].

from_nodes_, to_nodes_, the graph keys, node_map_ and graph_ all stay populated with the nodes of the failed graph. The next commit mixes them into a fresh graph and fails in cudaGraphAddDependencies with cudaErrorInvalidValue. After that first failure the encoder stays poisoned and every later evaluation keeps failing until the process is restarted.

Production logs show the transition within 29 ms:

[19:17:26.109] [error] exception: Cache thrashing is happening, please set the environment
                       variable MLX_CUDA_GRAPH_CACHE_SIZE to a larger value than 400 to fix
                       degraded performance.
[19:17:26.138] [error] exception: cudaGraphAddDependencies(
                           graph_, from_nodes_.data(), to_nodes_.data(), nullptr,
                           from_nodes_.size()) failed: invalid argument

The first error occurs once; the second then repeats for every new inference request — 870 occurrences in the captured service log — until restart.

This is not a request to change the thrashing detection. The problem is that a recoverable exception leaves CommandEncoder in an invalid state.

The fix

Split the body into commit_impl() and restore the graph state when it throws:

void CommandEncoder::commit() {
  try {
    commit_impl();
  } catch (...) {
    try {
      reset_graph_state_after_error();
    } catch (...) {
    }
    throw;
  }
}

reset_graph_state_after_error() restores exactly the state the successful path of commit_impl() restores. Three details it depends on:

The pending runtime error has to be cleared first. check_cuda_error() only formats and throws; it never calls cudaGetLastError(), so the error stays pending on the runtime. ~CudaHandle() skips its destroy while an error is pending, so every handle released after this point would leak, and recreating the graph below would fail because of an error that has already been reported. The cudaGraphExecUpdate failure path in commit_impl() already does the same thing.

The old graph has to be dropped even when destroying it fails. CudaHandle::operator= calls reset() first, so both the destroy and the create can throw, leaving graph_ holding the old handle. A graph left populated is worse than no graph at all: its nodes are still roots and run again on the next launch, while the fresh dependencies reference only the new nodes — so cudaGraphAddDependencies succeeds and the stale kernels re-execute silently against buffers whose temporaries were already released. CudaHandle::release() is added so the handle is given up even if cudaGraphDestroy fails.

The recovery must not replace the error being thrown. Wrapping it keeps the caller's original Cache thrashing is happening ... intact instead of surfacing an unrelated graph-creation failure.

Fields that cannot throw are cleared before the graph is touched, so a failure while recreating it still leaves the rest of the encoder empty.

The current evaluation still fails and returns the original error. The next one no longer inherits corrupted graph nodes and dependencies.

Test

tests/cuda_graph_recovery_tests.cpp, compiled when MLX_BUILD_CUDA is on.

The cache key is topological — one "K-" per node plus the dependency edges — and does not include shapes, so varying tensor sizes alone never produces a new key. The test instead lengthens the op chain on each iteration, and runs on a freshly created stream so that stream's CommandEncoder picks up the shrunken cache capacity at construction.

	succeeded	thrashing	corrupted
before	2	1	5
after	2	6	0

The 1 / 5 split reproduces the production signature exactly: the thrashing exception fires once, then every later commit degrades into cudaGraphAddDependencies.

Verified on RTX 4060 Ti (sm_89), CUDA 13.0, MSVC 19.44, Windows x64. Full suite: 3435 assertions, 0 failed — the four GGUF cases fail on this platform regardless of this change, since CMakeLists.txt forces MLX_BUILD_GGUF OFF under MSVC. Also exercised against a long-running inference workload with mixed model and LoRA shapes on sm_86 and sm_89 with the production settings (MLX_CUDA_GRAPH_CACHE_SIZE=1024, thrashing check disabled, CUDA graphs on), with no cache-thrashing or cudaGraphAddDependencies invalid argument errors.

Out of scope

Two related things this deliberately leaves alone:

• cache_misses_ in LRUCache is never reset, so once the threshold is crossed every subsequent miss keeps throwing. This change turns a permanently poisoned encoder into one failed evaluation per miss; requests that hit the cache recover, but a workload that keeps producing new graph topologies will keep seeing the thrashing error. Making the check one-shot seems like a separate decision for the author of Detect cache thrashing in LRUCache #2600.
• ConcurrentContext::~ConcurrentContext() calls CHECK_CUDA_ERROR(cudaGraphAddEmptyNode(...)) from a destructor, which terminates if it throws in the same sticky-error state that motivates this fix. Same class of problem in the same file, but unrelated to the commit path.

Note that increasing MLX_CUDA_GRAPH_CACHE_SIZE delays the first exception but does not make the commit path exception-safe.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Thanks!