Repository navigation
Add runner-owned CUDAGraphManager - #7180
sphinxkkkbc wants to merge 11 commits into
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
This PR appears to belong to: docs/design/module/ar_runtime.md. Module owners: @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123 @sphinxkkkbc, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Omni ReviewBot triage noteAutomated triage of commit
These are automated triage suggestions only — the final decision belongs to the maintainers. |
006e08a to
8dcfc28
Compare
108a99b to
e04ace4
Compare
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
e04ace4 to
2f81fff
Compare
|
@hsliuustc0106 @gcanlin @linyueqian PTAL, Thanks! |
hsliuustc0106
added
tts
linyueqian
left a comment
linyueqian
left a comment
There was a problem hiding this comment.
The design is sound and worth saying so first, because the division of labour here is harder to get right than the individual bugs. Keeping graph lifecycle in a runner-owned manager while the model only declares descriptors and supplies capture and replay routines is the right split. Failed captures degrade to eager per descriptor rather than per target, targets stay eager until every startup capture finishes so a partial failure cannot leave a half-bound model, and clone_output defaults to True, so a model has to deliberately opt out of cloning rather than remember to opt in. The tests reach the awkward cases: nested capture rejection, negative caching of failed descriptors, and concurrent lazy misses collapsing to one entry.
The PR does not currently start, though, and that is a one-word fix. arg_utils.py passes vocoder_cudagraph= where OmniModelConfig declares vocoder_cudagraph_config, and the omni-field validator rejects unknown kwargs. The key is passed unconditionally, so every model and every stage fails at config construction regardless of the feature, which is what the red lanes are. Both spellings are correct in their own layers and this line is the projection between them, so nothing else needs renaming.
Two blocking findings in the memory budget, which is written as a one-way accumulator. _captured_memory_bytes is incremented once and never refunded, on eviction or in clear(), and a descriptor rejected by the budget is neither registered nor negative-cached. Separately each is a slow degradation. Together they reinforce: eviction drifts the counter past the budget, and past the budget every lazy miss pays a full capture, warmups and stream sync included, before being discarded. Already-cached graphs keep replaying, so the symptom is that new descriptors quietly stop being captured while the cost of trying is paid on every request.
Both trigger under memory pressure, which is when a deployment least wants to pay for repeated capture, so this is worth fixing rather than tuning. Two cautions on the fix, because the symmetric-looking version is wrong. A budget rejection is not a permanent property of the descriptor the way a capture failure is, so blacklisting it defeats the LRU: once refunds work, something that does not fit now can fit later. And admission is checked before eviction, so refunding alone will not admit a replacement when the cache already sits at budget.
A related important finding is that the budget measures the wrong bytes. torch.accelerator.memory_allocated tracks live tensor allocations, while the graph is captured into the shared pool and its intermediates stay reserved rather than allocated, so captured_bytes mostly reflects the static buffers and not the graph. A routine with small persistent buffers and large temporaries passes a small budget while holding much more device memory. That is independent of the refund fix: consistent accounting of the wrong quantity is still the wrong quantity.
The other important one is that profile_cudagraph_memory returns 0 whenever the manager exists, so vocoder graphs are never reserved against the KV budget. The comment there flags lazy capture as an open question, but startup capture has the same hole and runs after KV allocation.
The two suggestions are lower confidence and I have marked them as such. The lock-free read of entries.keys() is latent rather than live, since no routine that ships here iterates available and execute_model is serialized, but both test routines do iterate and the concurrency was deliberately locked and tested, so it is worth closing while it is cheap. The other is that nothing sets max_memory_bytes in the manager tests, which is why the accounting findings survived; a test that called create_model_config() would likewise have caught the kwarg bug before CI did.
Validation: static review of all twelve files at 2f81fff6 against merge base eeee2579, following capture, registration, eviction and replay, the interface protocol, and the config plumbing from arg_utils.py through OmniModelConfig into the manager's reader. I did not execute the branch. I also considered and did not file the warmup running outside torch.inference_mode() while capture runs inside it, since capture_model and execute_model are both decorated with it and the inner context is redundant on both paths.
| stage_id=self.stage_id, | ||
| async_chunk=self.async_chunk, | ||
| session_mode=self.session_mode, | ||
| vocoder_cudagraph=self.vocoder_cudagraph, |
There was a problem hiding this comment.
[blocking] Wrong keyword name, and it takes down every engine start rather than just this feature.
from_vllm_model_config runs _validate_omni_fields, which raises ValueError: Unexpected omni kwarg: {key} for anything outside the declared dataclass fields (config/model.py:331-333). OmniModelConfig declares vocoder_cudagraph_config (config/model.py:130). Because this kwarg is passed unconditionally, a None value still trips it, so every model and every stage fails at config construction whether or not the YAML sets anything.
That is what the lanes are red on: general build 14852 hard-failed Simple · Engine&Entrypoints Test, Simple · Other Test, Engine&Model Executor Test and Entrypoints Test, and AMD 11575 also took down Engine Test and both TTS E2E steps.
The fix is one word, vocoder_cudagraph_config=self.vocoder_cudagraph. The two names are both correct in their own layers and the projection between them is exactly what this line was meant to perform: OmniStageModelConfig.vocoder_cudagraph (config/omni_config.py:447) carries the deploy-YAML spelling, and its comment on the line above already says it is projected to OmniModelConfig.vocoder_cudagraph_config by OmniEngineArgs. Nothing else needs renaming.
The reason this reached CI is that the new config tests assert on OmniStageModelConfig, which is the correct spelling for that class, so they never construct an OmniModelConfig and never reach this line. A test that calls create_model_config() would have caught it immediately.
Worth noting the validator earned its keep. from_vllm_model_config builds via object.__new__ and __dict__.update, so without that strict check the wrong name would have been set silently as an ad-hoc attribute, the manager's getattr(model_config, "vocoder_cudagraph_config", None) at vocoder_cudagraph_manager.py:182 would have returned None, and the feature would have been inert with no error anywhere.
There was a problem hiding this comment.
Resolved in 17220237b1b3ca6b41650ec5f966943fc0180ab5. Added regression test.
This was a cherry-pick error introduced while superseding #7042. It was not exposed by the tests in the previous head since they did not exercise the real Qwen3-TTS model configuration/construction path.
| return None | ||
| allocated_after = self._memory_allocated() | ||
| captured_bytes = max(0, allocated_after - allocated_before) | ||
| if self.max_memory_bytes is not None and self._captured_memory_bytes + captured_bytes > self.max_memory_bytes: |
There was a problem hiding this comment.
[blocking] A descriptor rejected by the budget is neither registered nor negative-cached, so with enable_lazy_capture it is re-captured on every request that resolves to it: allocate_buffers, cudagraph_num_of_warmups forward passes, a current_stream().synchronize() and a real graph capture, then discarded again. The capture-failure branch four lines up gets this right at line 332; this branch returns without it.
The failure mode is the inverse of what a memory budget is for: once memory is tight enough to start rejecting, affected requests become slower than they would have been with graphs disabled outright.
Startup makes it worse than a slow path. capture_and_bind runs inside _freeze_gc() (gpu_generation_model_runner.py:118), so a CUDAGraph dropped here is not collected until the freeze ends, and later startup descriptors can OOM against memory the manager already stopped counting.
One caution on the obvious fix: do not simply add these to failed_descriptors. Unlike a capture failure, a budget rejection is not a permanent property of the descriptor, and once eviction refunds the counter a descriptor that does not fit now can fit later. Permanent-fail only when a single graph is larger than the whole budget; otherwise evict until it fits and register it. Either way, do not capture and drop: register the graph or destroy it explicitly before returning.
| def _evict_lru_if_needed(self, managed: ManagedTarget) -> None: | ||
| assert managed.max_graphs is not None | ||
| while len(managed.entries) > managed.max_graphs: | ||
| managed.entries.popitem(last=False) |
There was a problem hiding this comment.
[blocking] Eviction never refunds _captured_memory_bytes, so the budget counter only rises while real usage stays bounded by max_graphs.
Line 344 is the counter's only write and it is an increment; clear() does not reset it either. max_graphs is set at line 521 to len(entries) + max_extra_graphs, so any deployment setting both max_extra_graphs and max_memory_bytes drifts upward until the check at line 336 rejects everything while actual memory sits well inside the budget. Already-cached graphs keep replaying, so the symptom is that new descriptors stop being captured rather than a hard failure.
It compounds with the finding above: past the budget, every lazy miss pays a full capture and is discarded, and the target never recovers short of a restart. A refund is not currently expressible, since VocoderCUDAGraphEntry records no charge. Store the charged bytes on the entry, subtract here and in clear(), and note that admission is checked before eviction, so refunding alone will not admit a replacement when the cache is already at budget.
| managed.failed_descriptors.add(descriptor) | ||
| return None | ||
| allocated_after = self._memory_allocated() | ||
| captured_bytes = max(0, allocated_after - allocated_before) |
There was a problem hiding this comment.
[important] This measures the wrong memory, so max_memory_bytes does not bound what its name promises.
_memory_allocated (line 200) uses torch.accelerator.memory_allocated, which tracks live tensor allocations in the caching allocator. The graph itself is captured into current_platform.get_global_graph_pool() (line 296), and intermediates freed during capture stay reserved in that pool rather than showing as allocated. By the time this delta is sampled, captured_bytes reflects roughly the allocate_buffers static tensors and not the graph's own footprint, which is usually the larger part.
A routine with small persistent buffers and large temporaries therefore passes a small budget while retaining substantially more device memory, which is the case the budget exists to prevent. This is independent of the refund fix above: storing and refunding the same delta keeps the accounting self-consistent but still counts the wrong bytes.
| # temporary profile capture would bind the same stable Targets and | ||
| # violate the one-manager lifecycle. Reservation for lazy capture | ||
| # remains an explicit design open question. | ||
| return 0 |
There was a problem hiding this comment.
[important] Returning 0 here means vocoder graph memory is never reserved against the KV-cache budget, and the comment only acknowledges the lazy-capture half of that.
Worker.determine_available_memory subtracts this estimate from the KV budget, and compile_or_warm_up_model calls capture_model() afterwards, so startup captures land in memory that has already been handed to the KV cache. Since this runner is a full GPUModelRunner subclass and still allocates KV, a generation stage can OOM inside capture_and_bind, or quietly over-commit when KV is large.
Reserving max_memory_bytes when it is set would be the cheap version, and it composes with fixing the accounting above.
|
|
||
| def runtime_callable(*args: Any, **kwargs: Any) -> Any: | ||
| routine.validate_runtime_inputs(args, kwargs) | ||
| resolution = routine.resolve_runtime(args, kwargs, entries.keys()) |
There was a problem hiding this comment.
[suggestion] entries.keys() is handed to resolve_runtime without _capture_lock, while lazy capture, eviction, and _touch_entry mutate that OrderedDict under it. entries is a MappingProxyType over the live dict, not a snapshot.
This is a suggestion rather than a finding because no routine that ships here iterates available: the PR provides only the abstract BaseVocoderCUDAGraphRoutine, and execute_model is serialized, so the runner path cannot hit it today. It is still latent in the API you have written and tested as concurrent. Both test routines do iterate (test_vocoder_cudagraph_manager.py:113-118 and the lazy-capture test), test_concurrent_lazy_misses_capture_one_entry calls one target from two threads, and _touch_entry's move_to_end can invalidate an iterator without changing the dict size, which is the version of this that is easiest to miss.
Snapshotting with tuple(managed.entries) under the lock costs nothing at these sizes. Worth noting the protocol also declares available: Set[VocoderCUDAGraphDescriptor] while a live KeysView is passed, so the contract and the argument disagree regardless of threading.
| assert manager.capture_attempts[("decode", 3)] == 0 | ||
|
|
||
|
|
||
| def test_lazy_entries_share_one_lru_with_startup_entries(monkeypatch) -> None: |
There was a problem hiding this comment.
[suggestion] No test in this file sets max_memory_bytes, so the entire budget path is unexercised: the rejection branch, the accounting increment, and their interaction with the eviction this test covers.
This test sets max_extra_graphs and asserts the shared LRU, which is the right shape. A sibling that also sets a small max_memory_bytes would have caught both accounting findings, since the drift only becomes visible where the two features meet. A test that calls create_model_config() would separately have caught the kwarg bug.
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
|
Updated Since the global CUDA Graph pool does not provide exact per-graph memory utilization, Also, as suggested by @hsliuustc0106 in this week's meeting, I renamed Target to Component to make the abstraction easier to understand. There are still two compatibility/design points to investigate:
As follow-up work in this PR, I will update the related design documentation and integrate Qwen3-TTS. The Qwen3-TTS migration was already exercised in #7042, but it should be refactored after the framework interface and YAML-side configuration are agreed upon. Any suggestions on the framework design and YAML-side configuration are appreciated, @linyueqian. I may also broaden the scope of this PR from vocoder-only CUDA Graph management to a more general downstream CUDA Graph framework. The framework was originally designed to support execution patterns such as VoxCPM2's unified and segmented graphs dispatch, so most of the underlying abstraction is not inherently vocoder-specific. I will work toward making that scope clearer in the next few commits while keeping the YAML-side changes as small as possible. I'll open follow-up issues if I think they're needed. |
linyueqian
left a comment
linyueqian
left a comment
There was a problem hiding this comment.
All seven round-one findings are addressed, and the two that mattered most are addressed properly. The vocoder_cudagraph_config rename matches the declared field at config/model.py:130, so engine start is no longer taken down for every deployment. Retiring max_memory_bytes outright was the right call rather than patching it: it removes the re-capture loop, the counter that only ever rose, and the fact that it was measuring allocator bytes rather than graph pool bytes, all in one move. The _available_descriptors snapshot under _capture_lock fixes the unlocked iteration cleanly.
What replaced them is where the remaining problems are, and both are in the new profiler and the eviction path it now exercises. profile_memory captures throwaway graphs into the persistent global pool and then discards them, which is the one thing the upstream profiler goes out of its way not to do, and making eviction destroy entries deterministically introduced a lifetime race against replay that the previous leak-it-to-GC behaviour happened to avoid. Neither is visible in the unit tests because they mock capture_entry, so no real graph or pool is involved.
This still cannot be approved on evidence even once those are resolved. No Buildkite lane has ever reported on 17220237: ready was already on the branch when you pushed, and pushing onto a branch that already carries ready fires nothing. The branch is also CONFLICTING against main. Main's general lane was red from 2026-09-08 to 2026-09-11 on an unrelated realtime event-name break and is green again as of build 15056, so a rebase now both resolves the conflict and gets a real four-lane result, in a way it would not have last week.
Validation: static review of the round-two delta from 2f81fff6 and of the full file at 17220237, the lock discipline across _evict_oldest, _destroy_entry, _available_descriptors and runtime_callable, and a read of profile_cudagraph_memory in the installed vllm 0.29.0 that docker/Dockerfile.ci targets, to check the claim that this follows the upstream approach. Fork head, so the branch was not executed. A four-model panel reviewed the same head and the findings below are theirs, each confirmed against the source before being raised.
| torch.inference_mode(), | ||
| torch.cuda.graph( | ||
| graph, | ||
| pool=current_platform.get_global_graph_pool(), |
There was a problem hiding this comment.
[blocking] Profiling captures into the persistent global graph pool and then discards those graphs, which is the specific thing the upstream profiler avoids.
capture_entry always captures with pool=current_platform.get_global_graph_pool(), and profile_memory calls it for up to two descriptors per component before destroying them in its finally and calling empty_cache(). So the throwaway profile graphs are created in, and released from, the pool that the real captures later reuse.
Upstream does the opposite, deliberately. In vllm 0.29.0, the version docker/Dockerfile.ci targets, GPUModelRunner.profile_cudagraph_memory allocates profiling_pool = current_platform.graph_pool_handle() and reassigns instance.graph_pool on every CUDAGraphWrapper and BreakableCUDAGraphWrapper for the duration, with the comment Use a temporary pool for profiling to avoid fragmentation in the main pool, restoring the originals afterwards. It also uses a temporary manager so the persistent one keeps no profiling-only graph state.
Two consequences follow. Pool memory released this way is not necessarily returned to the device, so Worker.determine_available_memory, which measures non-KV memory before calling this and then subtracts the returned estimate, can charge the sampled graphs twice: once as bytes still held, once in the estimate. And reopening a persistent pool whose use count was dropped to zero by the discard is exactly the hazard the upstream comment is guarding against, so there is a real risk of tripping the pool-reopen path on the next capture rather than merely fragmenting.
Pointing the capture at a throwaway graph_pool_handle() for the profiling phase, and restoring it after, matches upstream and removes both. Worth doing before the memory estimate is trusted for KV sizing.
There was a problem hiding this comment.
Resolved in a7987ea417416779fdff8ec98e09c4a4388511c2. Added a parameter to capture_entry and pass a temporary pool handle during profiling without affecting runtime capture, added a regression test to cover this case.
|
|
||
| @staticmethod | ||
| def _destroy_entry(entry: VocoderCUDAGraphEntry) -> None: | ||
| entry.graph.reset() |
There was a problem hiding this comment.
[blocking] _destroy_entry resets a graph without holding that entry's replay_lock, so eviction can destroy a graph another thread is replaying.
_evict_oldest pops the LRU entry and calls this under _capture_lock. The replay path takes a different lock: runtime_callable resolves an entry, then does with entry.replay_lock: and inside it calls routine.copy_runtime_inputs(args, kwargs, entry.buffers) and entry.graph.replay(). _capture_lock and replay_lock are unrelated, so nothing orders the two.
That makes the following interleaving reachable: thread A resolves an entry and enters its replay_lock, thread B takes a lazy-capture miss, _enforce_capacity evicts that same entry, and _destroy_entry calls entry.graph.reset() and sets entry.buffers = None while A is between copy_runtime_inputs and replay(). The mild outcome is an AttributeError on None buffers; the bad one is replaying a graph whose exec was just destroyed, which is a use-after-free on device memory rather than a clean failure.
This needs enable_lazy_capture plus eviction plus concurrent replay, but that is a supported configuration rather than an exotic one, and with max_extra_graphs=0 the first lazy miss already evicts. test_concurrent_lazy_misses_capture_one_entry establishes that concurrent replay is expected here, and round one avoided this only by accident: it dropped entries with popitem and left the exec to the garbage collector, so nothing was destroyed at a deterministic point.
Taking entry.replay_lock in _evict_oldest before destroying, or marking the entry dead and deferring the reset until its last replay releases, both close it.
| if not component_config.enabled or not component.descriptors: | ||
| continue | ||
| samples: list[int] = [] | ||
| for descriptor in component.descriptors[:2]: |
There was a problem hiding this comment.
[important] The extrapolation can under-count badly for exactly the workload this manager exists for.
Sampling component.descriptors[:2] takes the first two in declaration order, not the two largest, and the cost model is first_capture + per_graph * extra_graphs where per_graph is samples[1] or a 1 MiB floor. When a component has one usable sample, or when the second capture measures under 1 MiB, every remaining startup descriptor and every lazy graph is billed at 1 MiB regardless of shape. A vocoder component with many size buckets is the case where later descriptors are the expensive ones.
Because captures share the global pool, the measured delta for the second graph also reflects pool reuse rather than the marginal cost of a larger graph, so the per-graph figure is not a safe multiplier even when it is above the floor.
The failure is quiet: profile_cudagraph_memory under-reserves, KV cache takes the slack at startup, and the OOM surfaces later during capture_and_bind or on the first lazy capture of a large descriptor, far from the code that mis-estimated. Sampling the largest descriptor rather than the first, or taking the max of the samples instead of the second, would at least make the estimate an upper bound on the per-graph term.
There was a problem hiding this comment.
Resolved in a7987ea417416779fdff8ec98e09c4a4388511c2. Startup capture is now explicitly ordered from largest to smallest, and a regression test was added to cover this behavior.
| def runtime_callable(*args: Any, **kwargs: Any) -> Any: | ||
| routine.validate_runtime_inputs(args, kwargs) | ||
| resolution = routine.resolve_runtime(args, kwargs, self._available_descriptors(managed)) | ||
| entry = entries.get(resolution.descriptor) if resolution.descriptor is not None else None |
There was a problem hiding this comment.
[suggestion] One unlocked read of managed.entries is left. entries here is a MappingProxyType over the live OrderedDict captured at closure creation, and this get runs outside _capture_lock while lazy capture, move_to_end and popitem all mutate it under that lock.
Unlike the round-one keys() iteration this will not raise, since dict.get is atomic under the GIL and a stale None falls through to the miss handler, which is correct behaviour. Raising it only because the surrounding code now consistently snapshots under the lock, and leaving one direct read makes the discipline harder to verify later.
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
…coder-cudagraph-framework Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
|
Updated in Regarding the locking concern, I removed the locks and the related concurrency test, since there is currently no evidence of concurrent calls into the same runner/model-side CUDA Graph manager. One contract change is that the manager now assumes model-side execution will not concurrently replay different graphs that share the global CUDA Graph pool. UpdateGiven that individual LRU eviction is unsafe with a shared graph pool, I'm considering using the global graph pool when lazy capture is disabled, and giving every captured graph its own private pool handle when lazy capture is enabled. |
| if resolution.descriptor is not None: | ||
| self._descriptors[(component_id, resolution.descriptor.variant)] += 1 | ||
| key = (component_id, resolution.runtime_key.variant) | ||
| self._runtime_keys[key] += 1 |
There was a problem hiding this comment.
Bound runtime-key and descriptor statistics; retain aggregate counters when evicting detailed entries.
|
|
||
| entry = self.capture_entry(managed.component, descriptor) | ||
| if entry is None: | ||
| managed.failed_descriptors.add(descriptor) |
There was a problem hiding this comment.
Bound the failed-descriptor cache too; graph capacity limits never evict these entries.
| raise ValueError(f"Unknown config key(s) for vocoder Component {component_id}: {names}") | ||
| enabled = raw_component.get("enabled", True) | ||
| lazy = raw_component.get("enable_lazy_capture", False) | ||
| max_extra = raw_component.get("max_extra_graphs", 0) |
There was a problem hiding this comment.
Validate booleans and require a nonnegative integer for max_extra_graphs during prepare.
…ial-capture, add descriptor validation to handle special cases Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Omni ReviewBot: no human activity for 7 days@sphinxkkkbc this pull request has had no human commit, comment or review since 2026-09-16. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state. To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline. |
|
Status update — September 23 I’ll update this PR in the following follow-up steps:
|
sphinxkkkbc
changed the title
…ble policies, reduce change in yaml side, inherit upstream observability_config for statssink logging, remove lazy and lru combination in case of safety Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
sphinxkkkbc
changed the title
|
@linyueqian @Sy0307 @hsliuustc0106 PTAL, Thanks! |
Omni ReviewBot routing recordAssigned Strict on zcode (GLM-5.3-Flash) under experiment |
5 participants
Purpose
This PR introduces a runner-owned model-local CUDA Graph manager for graphable components inside generation models. It provides the framework and runner integration for both Model Runner V1 and V2; model-specific migrations remain follow-up work.
The ownership boundary is:
A model declares
supports_model_local_cudagraphand exposesModelLocalCUDAGraphComponentobjects. After all startup graphs have been captured, the manager binds an opaqueModelLocalGraphHandleto each active component. Model code does not hold the manager or access its graph entries.Capture timing is selected per component by the model:
PRECAPTURE,PRECAPTURE_LAZY, orPURE_LAZY. Runtime graphs are retained until manager cleanup; this PR does not perform runtime eviction.Known limitations
eager_call()when its graph-count limit is reached.max_extra_graphs: 0permits unbounded lazy capture, whose future memory use cannot be fully reserved during startup profiling.Design Overview
The runtime call path remains model-driven:
The model decides when and in what order components are called. The manager resolves each call against that component’s currently available descriptors. A unified component’s
eager_call()can still invoke graph-backed segmented components and interleave them with eager operations.Each component has one stable delegate:
Relationship to #4924
The core lifecycle remains:
Startup descriptors are captured largest first. Capturing every component before binding any handle matters for composition: capture-time calls through another component remain genuinely eager, while a runtime fallback from a unified component may use already-bound segmented graphs.
For each warmup and actual capture, the manager calls only:
The base
capture_context()callsprepare_for_capture()before the forward andafter_capture()infinally. Routines needing an additional scoped context can override it and compose withsuper().capture_context(...).Profiling captures use a temporary graph pool and are released afterward. The manager estimates graph memory from representative captures and logs aggregate memory around formal capture. Because graphs share a pool, an individual graph’s allocator delta is not a reliable refundable memory charge; this PR does not maintain per-entry memory accounting.
Key Refinements Beyond #4924
1. Component / Handle Binding
A model calls a stable component instead of calling a named manager API. The manager binds a narrow
ModelLocalGraphHandleafter capture; the handle exposes invocation and the descriptors currently backed by graphs, without exposing graph resources or lifecycle controls.Startup capture and binding are transactional. A capture or bind failure propagates, restores affected components to eager execution, and resets graphs captured during that attempt. Runtime lazy-capture failures also propagate; already retained entries are not evicted.
2. Multiple Graphable Components
Components have independent descriptor namespaces, capture modes, and runtime resolution. This allows model code to compose a larger graph path with smaller graph-backed paths and eager operations. The manager does not need to understand the model’s request-routing topology.
The model chooses whether a component is pre-captured, pre-captured with lazy extensions, or purely lazy. A pure-lazy component still declares representative descriptors for memory profiling.
3. Generation-Stage Configuration
A stage opts in with a direct component mapping under
model_local_cudagraph. The same field name is preserved from deploy YAML through the resolved model configuration:The model declares capture mode and descriptors. The framework-owned per-component option is
max_extra_graphs; model-defined options must be declared in that component’s supported configuration keys. Unknown component IDs and keys fail during preparation. Omitting the mapping leaves the existing runner behavior unchanged.Runtime statistics follow the upstream
observability_config.cudagraph_metricssetting. When enabled, aggregate counts are retained, detailed runtime-key and descriptor counts are bounded, and a summary is logged every 100 calls and at cleanup.4. Explicit Runtime Lifecycle Semantics
A graph hit copies inputs into static buffers, replays the graph, materializes the output, and clones tensor outputs by default. A coverage miss may capture a new descriptor if the component permits lazy capture and has capacity; otherwise it calls
eager_call().Runtime lazy capture uses a dedicated side stream with stream-ordering events. Startup capture uses vLLM’s
graph_capture()context and its non-default stream. Both paths use the global graph pool. Input-copy, capture, and replay errors propagate rather than being silently retried through eager execution.The V1 and MRV2 generation runners both prepare and clear the manager. The generation worker invokes temporary graph-memory profiling during memory planning and formal capture during startup. MRV2 keeps its existing eager outer dispatch so request routing continues to run for every call.
Sequence Diagram
sequenceDiagram participant Runner as Generation Runner participant Manager as ModelLocalCUDAGraphManager participant Model as Model participant Component participant Routine participant Graph as Graph Entry Runner->>Model: load Runner->>Manager: prepare(model) Manager->>Model: get_model_local_cudagraph_components() Runner->>Manager: profile_memory() Manager->>Graph: temporary captures, then release Runner->>Manager: capture_and_bind() loop all selected startup Descriptors Manager->>Routine: capture_context + forward_for_capture Manager->>Graph: retain capture in global pool end Manager->>Component: bind Handle after all captures succeed Model->>Component: call(runtime inputs) Component->>Manager: bound runtime callable Manager->>Routine: validate and resolve runtime alt retained graph matches Manager->>Graph: copy inputs and replay Manager-->>Model: materialized output else lazy capture is allowed and has capacity Manager->>Graph: capture and retain new entry Manager->>Graph: replay Manager-->>Model: materialized output else coverage miss Manager->>Routine: eager_call Routine-->>Model: fallback output end Runner->>Manager: shutdown / clear Manager->>Component: restore eager delegate Manager->>Graph: reset retained entriesTest Plan
Test Result
The focused suite above passes: 94 passed. Ruff checks pass for the changed code. Real CUDA capture and replay still require GPU validation as model-specific components are migrated.