Skip to content

Fix CUDA graph lifecycle across inference suspend and resume - #40

Draft
lmcafee-nvidia wants to merge 7 commits into
prefix-cache-core-stressfrom
prefix-cache-cuda-graph-lifecycle
Draft

Fix CUDA graph lifecycle across inference suspend and resume#40
lmcafee-nvidia wants to merge 7 commits into
prefix-cache-core-stressfrom
prefix-cache-cuda-graph-lifecycle

Conversation

@lmcafee-nvidia

@lmcafee-nvidia lmcafee-nvidia commented Aug 10, 2026

Copy link
Copy Markdown
Owner

What does this PR do?

This PR makes local CUDA graphs safe across inference suspend and resume. Graph deletion previously removed global records without invalidating every manager-local runner, and the engine could still select those stale graphs. It could also rebuild graphs immediately after OFFLOAD resume while retained requests still occupied the inference context; capture warmup uses dummy requests and resets that context, so rebuilding at that point could overwrite live state. The fix marks deleted graphs unavailable, clears local runners, recreates the deleted shared graph pool when capture resumes, executes eagerly until capture is safe, rebuilds only with an empty context, and preserves prefix-cache state and counters during warmup. It also skips redundant manual CPU backups when unified virtual memory already implements offload.

Example failure scenario(s)

Scenario 1 (offload resume)

  • request A is active, and request B is waiting.
  • the engine resumes after offloading its memory.
  • before the fix, it rebuilds graphs with 1 live request instead of 0.
  • graph warmup can reset the state for request A.
  • after the fix, request A finishes before graphs rebuild and request B starts.

Scenario 2 (deleted graphs)

  • the engine deletes the global graph pool, leaving 0 graphs.
  • at least 1 local runner still points to a deleted graph.
  • before the fix, the next batch can select that stale runner.
  • after the fix, every stale local runner is cleared.

Scenario 3 (unified memory)

  • unified memory level 1 already owns the offloaded memory.
  • before the fix, resume creates 1 extra CPU backup.
  • the two memory paths conflict and raise an error.
  • after the fix, resume does not create the extra backup.

How was this bug discovered?

Lifecycle coverage suspended the engine midway through real inference, deleted graph and memory state, and then resumed under OFFLOAD, RECOMPUTE, and PERSIST. The observed behavior included capture being attempted while retained work was present, capture warmup clearing prefix hashes and counters, stale manager-local graph objects surviving deletion, and an unexpected manual-backup assertion when static unified virtual memory was active. The expected invariants were that capture sees an empty context, warmup preserves prefix hashes and counters, deleted runners are unreachable, retained requests continue safely, and UVM does not take a second backup path. These exact lifecycle/state assertions—not token parity alone—were the discovery oracle; the focused runs then also checked successful request completion across GPT, hybrid, TP2, and PP2 paths.

Lifecycle contract

Cache-management mode CUDA-graph behavior after resume
PERSIST Keep the existing graph and retained inference state; no rebuild is needed.
RECOMPUTE Rebuild deleted graphs while the context is empty, before requests are re-admitted.
OFFLOAD Keep retained requests, run them eagerly until they drain, then rebuild before admitting new waiting work.

In every mode, graph capture requires an empty context, warmup preserves prefix state and counters, deletion clears both global and manager-local runners, and UVM-backed offload does not allocate redundant manual CPU backups.

Dependency and scope

Depends on NVIDIA#6418 and is based on prefix-cache-core-stress. Relative to that base, this PR changes four production files and two unit-test files, with +388/-2 unit-test lines. It contains no functional-test or recipe changes.

Validation

Exact signed tip: 5a72930867caf41fe770fefbd5e7065bbc731c19

  • red/intermediate DFW job 15519444 exposed the UVM manual-backup assertion and a TP/PP fixture-size error
  • DFW job 15520282 at the preceding production tip passed OFFLOAD, RECOMPUTE, and TP2/PP2 cases; the UVM test then exposed a test-side CPU block-index mismatch
  • final exact-tip DFW job 15521039 completed successfully on an eight-H100 node:
    • all four selected cases passed on ranks 0-7: OFFLOAD, RECOMPUTE, static-UVM hybrid, and PERSIST TP2/PP2
    • each rank reported 4 passed, 0 failed, 0 skipped, 69 deselected (32 rank-local passes total)
    • the static-UVM case kept context.unified_memory_level == 1 on every rank, proving UVM stayed active rather than falling back to GPU memory
    • exact diff-surface and git diff --check, isort, Black, Ruff, and Python compilation checks passed

Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant