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]

Can you move the clear()s to a separate helper and reuse it in commit_impl?

[strayberry]

Can you move the clear()s to a separate helper and reuse it in commit_impl?

Done in 5c594db — moved the clear()s into clear_graph_state() and call it from both commit_impl() and the error path, so the field list lives in one place.

Two notes on where I drew the line:

• The error path additionally zeroes node_count_ / bytes_in_graph_. commit_impl() already does that unconditionally further down, outside the use_cuda_graphs() branch, since both are also incremented when graphs are off — so folding them into the helper would have skipped them on the non-graph path.
• Recreating the graph stayed out of the helper too: the successful path can just assign a fresh CudaGraph, while the error path has to survive a failing cudaGraphDestroy first. Keeping the helper free of anything that can throw also makes it safe to call while unwinding.

One behaviour change worth flagging: the shared helper now also clears active_deps_, active_outputs_ and concurrent_nodes_ on the successful path. Those are consumed inside insert_graph_dependencies and are already empty by the time commit runs, so it is a no-op there — happy to drop them from the helper if you would rather keep the successful path exactly as it was.