Skip to content

[BUG][CUDA] CommandEncoder remains corrupted after graph-cache thrashing exception #4326

Closed

Description

@strayberry

Describe the bug

On the CUDA backend, when the CUDA graph LRU cache throws the cache-thrashing exception introduced by #2600, CommandEncoder::commit() exits before resetting its in-progress graph state.

The next request then reuses stale graph nodes/dependency vectors and fails in cudaGraphAddDependencies with cudaErrorInvalidValue (invalid argument). After that first failure, the encoder remains poisoned and subsequent otherwise valid requests continue to fail until the process is restarted.

This is not a request to remove the cache-thrashing detection. The issue is that a recoverable exception from graph_cache_[graph_key] leaves CommandEncoder in an invalid state.

Observed sequence:

Cache thrashing exception
  -> CommandEncoder::commit() exits before graph state reset
  -> from_nodes_ / to_nodes_ / graph keys / graph_ remain populated
  -> next request mixes stale node handles with the next graph
  -> cudaGraphAddDependencies(...) returns invalid argument
  -> subsequent requests keep failing

Production logs show the transition within 29 ms:

[2026-08-17 19:17:26.109] [error] [LlamaWorkerTiming] 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.

[2026-08-17 19:17:26.138] [error] [LlamaWorkerTiming] exception:
cudaGraphAddDependencies(
    graph_,
    from_nodes_.data(),
    to_nodes_.data(),
    nullptr,
    from_nodes_.size()
) failed: invalid argument

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

To Reproduce

The production reproduction is a long-lived CUDA inference worker processing many graph shapes and LoRA/task variants:

  1. Run with CUDA graphs enabled and the defaults:

    MLX_CUDA_GRAPH_CACHE_SIZE=400
    MLX_ENABLE_CACHE_THRASHING_CHECK=1
    
  2. Generate more than 2 * capacity CUDA graph cache misses.

  3. Catch the resulting Cache thrashing is happening exception and keep the worker alive.

  4. Submit another valid inference request.

  5. The next request fails in cudaGraphAddDependencies with invalid argument, and later requests continue to fail.

A small regression can reduce the cache capacity to 1, construct at least three distinct CUDA graph keys, catch the expected thrashing exception, and then evaluate one known-valid graph. The final evaluation should succeed; currently it should expose the stale encoder state.

Relevant source paths:

  • mlx/backend/cuda/lru_cache.h: LRUCache::emplace() throws after the miss threshold.
  • mlx/backend/cuda/device.cpp: CommandEncoder::commit() accesses graph_cache_[graph_key], while graph-state cleanup happens only later on the success path.

The issue is present in current main commit 47a1bd613ebb67ed03da764c6ffedb8df9939997.

Expected behavior

After any exception during CommandEncoder::commit(), the encoder should return to a valid empty state. Catching the cache-thrashing exception should fail only the affected evaluation; it should not corrupt all subsequent evaluations.

At minimum, the error path should reset/recreate:

from_nodes_
to_nodes_
graph_deps_key_
graph_nodes_key_
node_map_
active_deps_
active_outputs_
concurrent_nodes_
graph_
is_graph_updatable_
node_count_
bytes_in_graph_

One possible approach is to split the current implementation into commit_impl() and wrap it:

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

A regression test should intentionally trigger an exception during graph-cache insertion and verify that the following graph evaluation succeeds.

Desktop (please complete the following information):

  • OS Version: Windows, x86_64
  • CUDA: CUDA 13 API path (the failing call includes the CUDA 13 edgeData argument)
  • GPU backend: CUDA
  • MLX source commit: 47a1bd613ebb67ed03da764c6ffedb8df9939997
  • Workload: long-lived C++ inference service, one MLX worker, multiple model/LoRA task shapes

Additional context

  • Detect cache thrashing in LRUCache #2600 intentionally converts prolonged cache thrashing into an explicit exception; it does not restore CommandEncoder state after that exception.
  • [CUDA] Update calls to cudaMemAdvise and cudaGraphAddDependencies for CUDA 13  #2525 is already present: the build uses the CUDA 13 cudaGraphAddDependencies signature.
  • Increasing MLX_CUDA_GRAPH_CACHE_SIZE delays the first exception but does not make the commit path exception-safe.
  • Disabling MLX_ENABLE_CACHE_THRASHING_CHECK avoids this particular trigger, but other exceptions inside commit() could leave the same stale state.
  • Restarting the process clears the poisoned encoder and temporarily restores inference.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions