Fix CUDA graph lifecycle across inference suspend and resume - #40
Draft
lmcafee-nvidia wants to merge 7 commits into
Draft
Fix CUDA graph lifecycle across inference suspend and resume#40lmcafee-nvidia wants to merge 7 commits into
lmcafee-nvidia wants to merge 7 commits into
Conversation
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>
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
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.
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
OFFLOADresume 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)
Scenario 2 (deleted graphs)
Scenario 3 (unified memory)
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, andPERSIST. 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
PERSISTRECOMPUTEOFFLOADIn 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:
5a72930867caf41fe770fefbd5e7065bbc731c1915519444exposed the UVM manual-backup assertion and a TP/PP fixture-size error15520282at the preceding production tip passedOFFLOAD,RECOMPUTE, and TP2/PP2 cases; the UVM test then exposed a test-side CPU block-index mismatch15521039completed successfully on an eight-H100 node:OFFLOAD,RECOMPUTE, static-UVM hybrid, andPERSISTTP2/PP24 passed, 0 failed, 0 skipped, 69 deselected(32rank-local passes total)context.unified_memory_level == 1on every rank, proving UVM stayed active rather than falling back to GPU memorygit diff --check, isort, Black, Ruff, and Python compilation checks passed