[feat]: complete safe CUDA-graph artifact dispatch V1 - #25
Merged
Conversation
…om the hot path Run ltx-v1-overnight-20260801-r4-sol loaded four artifacts into vae.decoder.up_blocks.6.res_blocks at once, executed 56 candidate calls with zero runtime fallbacks, failed byte_equal parity, and regressed end-to-end from 3.2818s to 3.9410s (0.8327x). The evidence cannot attribute either the parity change or the latency to any individual artifact, because there was no way to enable one at a time short of staging a separate directory per trial. Add FASTVIDEO_OPTIMIZATION_ARTIFACT_ENABLE, a comma-separated allowlist of artifact IDs. Every bundle under the root is still discovered, verified and hashed, so a corrupt one is still reported; selection only narrows which verified manifests may be matched. Requesting an ID no bundle provides is an error rather than a silently empty run, which would otherwise read as "this artifact changes nothing". The registry summary records enabled_filter and excluded_ids so a trial's own dispatch.json says what was under test. Two hot-path costs, both paid per candidate call: - _dispatch built three _parameter_snapshot dictionaries per call -- walking module parameters, sorting hook names and listing up to 16 tensor shapes -- to diagnose FSDP parameter materialization. A run that never fails never reads them. The snapshot now happens on the failure path, where it is used. - _validate_runtime_inputs walked every node in the rewritten graph on every call to rediscover its placeholders. That graph is immutable for the lifetime of the dispatcher, so the placeholder contract is now frozen once at build time and the per-call check iterates only the placeholders. Neither changes what is accepted: the same shape/dtype contract is enforced, and the same failures raise with the same messages. pytest fastvideo/tests/optimization/ -> 93 passed (6 new).
The transformer artifact costs 3.104ms per call end-to-end while its kernel saves 124us, so gate 5 fails on framework overhead. Attributing that needs the cost broken down in situ -- same module, same live tensors, same stream -- not in a micro-harness that rebuilds the region separately. FASTVIDEO_OPTIMIZATION_ARTIFACT_TIMING enables per-phase accounting across shape-key construction, pytree flatten, input validation, rewritten-graph execution and unflatten. 'sync' synchronizes around each phase to attribute device time; 'shadow' additionally runs the native forward on every dispatched call and times it, which is the only honest comparison for whether the artifact path is cheaper than the one it replaced. Inert unless the variable is set: phase() returns a shared no-op context manager, so the hot path pays one module-level boolean check.
The dispatcher executes the rewritten export graph in place of the module's own forward, and shadow timing shows that path costs 11.57ms against the native forward's 8.18ms on the same inputs. An export graph is decomposed, so a single high-level call in the module can become many primitives here; identifying which ops the replay actually runs is the next question. FASTVIDEO_OPTIMIZATION_ARTIFACT_DUMP_GRAPH writes op names and counts only.
Gate 5 fails on dispatch cost, not arithmetic. Shadow timing on transformer.model.transformer_blocks measured the rewritten-graph replay at 11.57ms against the module's own forward at 8.18ms on identical inputs, and the op histogram explains it: the export graph is decomposed into 621 call_function nodes per call, each paying Python and PyTorch dispatcher cost. At ~5us per op that is the whole 3.39ms penalty, against an artifact saving of 124us. Attention is not the problem -- it survives capture as a single fastvideo._flash_attn_default_forward op. Two changes, both numerics-preserving: - Strip export's aten._assert_tensor_metadata nodes from the rewritten graph. 68 of the 621 compute nothing, and the dispatcher already validates the same metadata once per call via the placeholder contract. Only nodes with no users are removed. - Replay the graph from a torch.cuda.CUDAGraph capture. A CUDA graph runs the same kernels with the same parameters in the same order, so it is bitwise identical to the eager replay by construction. That matters here: a compiler backend could fuse or reassociate and put the workload's byte_equal parity at risk, which is the gate this work exists to keep. Capture is guarded rather than assumed. Non-CUDA or non-tensor inputs, a changed shape or dtype, a missing pool, a non-tensor output, or any capture failure raises CudaGraphUnavailable and the entry falls back to eager replay permanently. Declining only costs speed. Static input buffers and the graph memory pool are shared across a scope's blocks. Every repeated block is called with the same input signature, so one buffer set serves all of them; per-block buffers would have added roughly 19GB for this stack's 48 blocks and 27 placeholders, past the workload's 5% peak-memory allowance. Off via FASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHS=0. pytest fastvideo/tests/optimization/ -> 101 passed (8 new).
…back A capture failure raising an arbitrary exception escapes to the dispatcher's candidate-runtime handler, which demotes the artifact to native execution for the rest of the run. That reads a missing acceleration as a broken kernel: the candidate is fine, only the capture is unavailable. Wrap capture, drop any partial state, synchronize, and re-raise as CudaGraphUnavailable so the eager replay proceeds and the artifact keeps running.
CUDA graph replay cut 1.382ms per call on transformer.model.transformer_blocks but left 3.381ms, and copying the boundary tensors into static buffers is a large part of what remains: the 27 placeholders total roughly 565MB per call, so a copy-in is over 1.1GB of traffic before the graph runs. Most of those tensors do not move. Timesteps, positional embeddings and the text context are the same tensor object for all 48 blocks and all 8 steps of a generation; only the residual stream changes. The runner now records input addresses across warmup, captures reading the unchanged ones where they already live, and allocates static buffers only for the positions that move. Correctness rests on the pinned addresses still holding at replay time, so every pinned input's data_ptr is checked on every call. A tensor that moved makes the runner decline and the eager replay answer instead -- replaying would read whatever now occupies that address, which is the one failure this path must never have. pytest fastvideo/tests/optimization/ -> 104 passed (3 new).
Profiling showed subgraph.execute_cuda_graph entered 192 times on transformer.model.transformer_blocks -- 48 blocks x 3 warmups, plus 48 capture attempts -- while the eager path still ran all 1151 calls. Every capture failed, and the reason went only to logger.info, which the host's log level dropped. A silently declined fast path is indistinguishable from one that was never attempted. Reasons are now counted in the timing report and logged at warning level.
…lining Every CUDA graph capture on transformer.model.transformer_blocks was failing with 'runtime input 10 is bool, not a tensor', so the fast path never engaged once across 1151 calls -- the eager replay ran every time and the measured improvement was node variance, not the capture. Export flattens Python scalars and flags through as graph inputs. They hold no device memory and need no static buffer; the capture bakes them in as the constants they are. The graph is only valid while they keep that value, so every call checks type and equality and declines if either changed, because the captured kernels encode the old value. pytest fastvideo/tests/optimization/ -> 107 passed (2 new).
…nstant outputs Two failures stopped every capture on transformer.model.transformer_blocks: - Sharing one graph_pool_handle across the stack's 48 blocks tripped the caching allocator's 'use_count > 0 INTERNAL ASSERT FAILED'. Each graph keeps its output buffers alive for the life of the run, so the pool is never free of live allocations when the next capture begins. Each capture now gets its own pool. - The rewritten graph returns a bool alongside its tensors. Like a non-tensor input, it is a constant the graph reproduces identically on every replay, so it is returned as captured rather than copied out. Static input buffers stay shared across the scope; only pools are private.
A review of the capture path found defects that a byte_equal workload cannot carry. In order of severity: - The capture bakes in every parameter and buffer the graph reads through get_attr, and nothing re-checked them. FSDP2's reshard frees the all-gathered storage the capture recorded pointers into, so the next replay would read freed memory or another tensor; weight offload and dequantization caches move them the same way. Their identity is now recorded at capture and re-checked on every replay. - Non-tensor outputs were returned as captured on the assumption that 'not a Tensor' means 'immutable constant'. A list or dict leaf would have handed the caller the graph's own static buffers, which the next replay overwrites in place. Only None and scalars are treated as constants now; anything else declines. - Pinned inputs were validated by data_ptr alone. data_ptr folds in storage_offset but says nothing about strides, so a permuted view handed back the same allocator block passed every check while the captured kernels read elements in a different order. Pinning is now on data_ptr, shape, stride, dtype, device and storage_offset. - An aborted capture never called graph.reset(), pinning its private memory pool for the life of the run. - The constant comparison could raise a non-RuntimeError, escaping both handlers and permanently demoting a working artifact. Also: 'warming up' was signalled by comparing an exception message string, so rewording it would have silently turned every warmup into a permanent disable. It is a distinct exception type now. Most importantly, the argument that a capture must be bitwise identical to the eager replay is now checked rather than asserted. The graph contains one node export did not functionalize -- the artifact's own entry point -- so purity is an assumption about third-party code. After capture, the runner replays once, runs the eager graph, and refuses the capture unless every output is bitwise equal. pytest fastvideo/tests/optimization/ -> 111 passed.
Three findings from Greptile, all real: - A dead 'if True:' guard in the decline handler, left behind when the warmup signal moved from a message-string comparison to a distinct exception type. - The private memory pool was released on the unexpected-error path but not when _capture itself refused. That mattered most for the bitwise-verification step, which ran *after* self._graph was assigned: a rejected capture stayed published on the runner and its pool's lifetime depended on the caller dropping the runner and refcounting collecting it. The graph is now published only once verified, and both abort paths go through _release_capture. - timing's counters are process-global with no way to clear them, so two dispatch sessions in one process would produce a report attributable to neither. Added timing.reset(). pytest fastvideo/tests/optimization/ -> 113 passed (2 new).
📝 WalkthroughWalkthroughChangesThe optimization runtime now supports artifact-ID allowlisting, opt-in timing reports, autocast replay during deferred tracing, and validated CUDA graph capture with eager fallback. Optimization runtime
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Dispatch
participant CudaGraphRunner
participant CUDARuntime
participant EagerReplay
Dispatch->>CudaGraphRunner: Validate inputs and contracts
CudaGraphRunner->>CUDARuntime: Warm up and capture graph
CUDARuntime-->>CudaGraphRunner: Return capture result
CudaGraphRunner->>CUDARuntime: Replay static buffers
CUDARuntime-->>CudaGraphRunner: Return graph outputs
CudaGraphRunner-->>Dispatch: Reconstruct outputs
CudaGraphRunner->>EagerReplay: Use eager fallback when capture is refused
EagerReplay-->>Dispatch: Return eager outputs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (10)
fastvideo/tests/optimization/test_fx_capture.py (1)
496-520: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd CUDA autocast replay coverage.
The current test only exercises CPU
bfloat16autocast replay. Add a CUDA-gated regression infastvideo/tests/optimization/test_fx_capture.pythat captures CUDA autocast state, exercises replay throughfinalize, and documents the GPU assumption.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/tests/optimization/test_fx_capture.py` around lines 496 - 520, Add a CUDA-gated regression test alongside test_deferred_export_replays_observed_autocast_dtype that runs the same capture and finalize flow under CUDA autocast, verifies the captured executable IR preserves the expected CUDA autocast dtype, and confirms autocast state is restored after execution and finalize. Document the test’s GPU requirement and skip it when CUDA is unavailable.Source: Coding guidelines
fastvideo/optimization/dispatch.py (1)
284-289: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument that shadow mode can change results for stateful forwards.
The shadow run calls
native(*args, **kwargs)an extra time per dispatch. If the native forward consumes RNG, updates a cache, or mutates module state, the extra call changes the generation output. The timing docstring states only that the result is discarded and that cost roughly doubles.Add that constraint to the
SHADOWdocstring infastvideo/optimization/timing.py, so nobody enables it during a parity run.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/optimization/dispatch.py` around lines 284 - 289, Update the SHADOW documentation in the timing module to state that shadow mode performs an extra native forward and can change outputs for stateful forwards that consume RNG, update caches, or mutate module state. Explicitly warn against enabling SHADOW during parity or generation-result comparisons, while retaining the existing notes about discarded results and increased cost.fastvideo/optimization/timing.py (2)
105-107: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueBound
_notesor reconsider recording notes while timing is disabled.
note()records unconditionally, butsnapshot()andwrite_report()return data only whenENABLEDis true. In the default configuration the notes accumulate and nothing ever reads them. Key cardinality is bounded only by the distinct message strings callers build, andfastvideo/optimization/subgraph.pyformats a decline reason into the key.Add a size cap, so a caller that interpolates a variable value into a note cannot grow the dictionary without limit.
♻️ Proposed change
def note(message: str) -> None: """Count one structured observation. Always on, so it survives log config.""" - _notes[str(message)[:200]] += 1 + key = str(message)[:200] + if key in _notes or len(_notes) < 256: + _notes[key] += 1🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/optimization/timing.py` around lines 105 - 107, Bound the `_notes` dictionary updated by `note()` so variable message values cannot create unlimited distinct keys. Add a fixed maximum cardinality and ensure new note keys are ignored once the cap is reached, while continuing to increment existing keys and preserving the current message truncation behavior.
151-163: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider matching the atomic write used by
write_diagnostics.
GraphDispatchSession.write_diagnosticsinfastvideo/optimization/dispatch.pywrites to a.tmpfile and then callsreplace, so a reader never sees a partial report. This function writes in place. Both files are produced by the same teardown call, so a consistent strategy avoids a truncatedtiming.json.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/optimization/timing.py` around lines 151 - 163, The write_report function currently writes timing.json in place, allowing readers to observe partial output. Match GraphDispatchSession.write_diagnostics by writing the serialized snapshot to a temporary file in the destination directory, then atomically replacing the target with it; preserve the existing None-on-failure behavior and ensure temporary artifacts are cleaned up as appropriate.fastvideo/tests/optimization/test_dispatch.py (2)
1894-1907: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueIsolate the process-global timing state with a fixture.
The test mutates
timing's module-level counters directly.monkeypatchrestoresENABLED, but_totals,_counts, and_notesare left in whatever state the test produced.note()also ignoresENABLED, so it records during every other test in the session.Add an autouse or explicit fixture that calls
timing.reset()before and after, so a future test that asserts accumulated totals cannot become order-dependent.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/tests/optimization/test_dispatch.py` around lines 1894 - 1907, Add a fixture for the imported timing module that calls timing.reset() both before and after tests, and apply it to test_timing_state_can_be_reset_between_sessions (or make it autouse for the module). Keep the existing ENABLED monkeypatch and assertions, ensuring _totals, _counts, and _notes are cleared before and after each test.
1541-1586: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the real export assertion target in these tests.
_ASSERTION_TARGETSis hardcoded toaten._assert_tensor_metadata.default, but the tests only cover monkeypatched local function targets. Add one export-based or target-string assertion so a PyTorch export rename fails the suite instead of leaving the stripping logic as a no-op.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/tests/optimization/test_dispatch.py` around lines 1541 - 1586, Update the assertion-stripping tests around _assertion_graph and test_export_runtime_assertions_are_stripped to include a real torch export assertion target, or explicitly assert its expected target string. Ensure the test verifies that _ASSERTION_TARGETS contains the actual aten._assert_tensor_metadata.default target and that stripping removes it, so PyTorch target renames cannot silently make the logic a no-op.fastvideo/optimization/subgraph.py (4)
182-204: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe profile file is overwritten once per live block.
_dump_graph_profileis called frombuild, andbuildruns once for each repeated block in a stack. Every call writes the samedestinationpath, so only the last block's histogram survives. For a 48-block stack the earlier profiles are lost.Include the block identity in the filename, so each build produces its own record.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/optimization/subgraph.py` around lines 182 - 204, The _dump_graph_profile output path currently overwrites prior block profiles; update its filename construction to include the current block’s unique identity, using available manifest or build-context symbols, while preserving the configured destination directory and JSON output behavior.
950-961: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueA process-wide decline reason is logged once per block.
entry.cuda_graph = Nonedisables the accelerator for one block only. A reason that applies to the whole run, for example"CUDA is not available", is therefore re-derived and re-logged for every block in the stack. For a 48-block transformer that is 48 identical warnings.Use
logger.warning_once, whichfastvideo/logger.pyprovides, or hoist a run-wide refusal to the sharedgraph_scope.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/optimization/subgraph.py` around lines 950 - 961, The CudaGraphUnavailable warning in the graph replay exception path is emitted once per block instead of once per process. Update the logger.warning call in the CudaGraphUnavailable handler to use the existing logger.warning_once mechanism, preserving the current message and eager-replay behavior while still clearing entry.cuda_graph for the affected block.
209-214: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThree new environment settings bypass the
fastvideo/envs.pyregistry. This PR adds four environment settings for the optimization runtime. OnlyFASTVIDEO_OPTIMIZATION_ARTIFACT_ENABLEis registered centrally; the other three callos.getenvat module import, so they are undiscoverable throughenvs, fixed for the process lifetime, and not patchable from a test.
fastvideo/optimization/subgraph.py#L209-L214: moveFASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHSintoenvironment_variablesinfastvideo/envs.pyand read it throughenvswhere_Entryis built. Do the same forFASTVIDEO_OPTIMIZATION_ARTIFACT_DUMP_GRAPHat Line 182.fastvideo/optimization/timing.py#L41-L48: registerFASTVIDEO_OPTIMIZATION_ARTIFACT_TIMINGinfastvideo/envs.py. KeepENABLED,SYNCHRONIZE, andSHADOWas module constants derived from that value, so the hot path still pays only a boolean check.fastvideo/envs.py#L355-L361: add the three entries next to the existingFASTVIDEO_OPTIMIZATION_ARTIFACT_*block, with the same comment style.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/optimization/subgraph.py` around lines 209 - 214, Register FASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHS and FASTVIDEO_OPTIMIZATION_ARTIFACT_DUMP_GRAPH in fastvideo/envs.py alongside the existing optimization artifact entries, then update subgraph.py to read them through envs when building _Entry instead of os.getenv. Register FASTVIDEO_OPTIMIZATION_ARTIFACT_TIMING there as well and update timing.py so ENABLED, SYNCHRONIZE, and SHADOW remain module-level constants derived from the registered value. Apply these changes at fastvideo/optimization/subgraph.py lines 209-214, fastvideo/optimization/timing.py lines 41-48, and fastvideo/envs.py lines 355-361; preserve the existing comment style and boolean hot-path behavior.
461-472: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueTwo small hardening items in the runner lifecycle.
Line 466 uses a tuple assignment to read
_pending_graphand clear it in one statement. The evaluation order is correct, but the intent is hard to read, and thegetattrdefault is redundant because__init__always sets the attribute.Line 540 indexes
self._moving[index]directly._captureassigns every index to_pinned,_moving, or_constants, so the key is always present today. If that invariant ever breaks, theKeyErroris not aCudaGraphUnavailable, soGraphDispatchSession._dispatchdemotes the artifact permanently — the outcome Lines 498-508 exist to prevent.♻️ Proposed changes
def _release_capture(self) -> None: """Drop a capture and free the private memory pool it reserved.""" aborted, self._graph = self._graph, None self._static_outputs = () self._attributes = [] - self._pending_graph, pending = None, getattr(self, "_pending_graph", None) + pending = self._pending_graph + self._pending_graph = None for candidate in (aborted, pending):- buffer = self._moving[index] + buffer = self._moving.get(index) + if buffer is None: + raise CudaGraphUnavailable( + f"runtime input {index} has no captured role" + ) if live.shape != buffer.shape or live.dtype != buffer.dtype:Also applies to: 540-543
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@fastvideo/optimization/subgraph.py` around lines 461 - 472, In _release_capture, read self._pending_graph into pending and clear the attribute in separate statements, relying on its initialization rather than getattr. In the capture/replay path around self._moving[index], replace direct indexing with a guarded lookup that raises CudaGraphUnavailable when the index is absent, preserving GraphDispatchSession._dispatch’s recoverable failure handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@fastvideo/optimization/dispatch.py`:
- Around line 310-326: Guard the diagnostic snapshot inside the runtime-failure
handler around the fallback logic near the artifact execution path. Ensure any
exception from self._parameter_snapshot(wrapper) is caught or replaced with a
safe diagnostic value, then still emit the warning and always reach return
native(*args, **kwargs), preserving native fallback even when snapshot
collection fails.
- Around line 683-686: Update the diagnostics handling in detach_graph_dispatch
so an invalid output path cannot escape teardown: guard the timing report path
construction and write, including Path.with_name, with error handling while
preserving session.write_diagnostics behavior. Ensure values such as "." or "/"
are safely ignored or logged without preventing detach_graph_dispatch from
completing.
In `@fastvideo/optimization/subgraph.py`:
- Around line 55-74: Update _capture to reject any pinned input whose
_layout_identity returns None, matching the existing unreadable bound-attribute
handling. Ensure the capture fails before storing None in self._pinned, and keep
replay validation in the pinned-input path from treating two None layouts as a
valid match.
- Around line 815-818: Update the cache comment above the WeakKeyDictionary in
the graph-scope initialization to describe the value as an _Entry, including its
rewritten module, frozen placeholder contract, and cuda_graph; keep the cache
declaration and behavior unchanged.
In `@fastvideo/optimization/timing.py`:
- Around line 118-128: Call timing.reset() at the beginning of every dispatch
session, before any timing.phase(), timing.record(), or timing.note() calls
occur. Update the dispatch-session entry point rather than the timing.reset()
implementation, ensuring each model/session starts with cleared process-global
timing state while preserving existing report generation.
In `@fastvideo/tests/optimization/test_dispatch.py`:
- Around line 1853-1882: Extract the output-leaf classification logic into a
shared helper near _CudaGraphRunner, such as _classify_output_leaves, returning
constant leaves and raising CudaGraphUnavailable for mutable values. Update
_CudaGraphRunner._capture to use this helper, then change
test_a_mutable_output_is_refused_rather_than_aliased to call the helper directly
and remove the unused runner, fake graph, and monkeypatch setup.
- Around line 1470-1498: Update
test_placeholder_contract_is_derived_once_not_per_call to verify its caching
claim by monkeypatching _placeholder_contract and asserting dispatch invokes it
only once across repeated calls, following the pattern used by
test_dispatch_does_not_snapshot_parameters_on_the_success_path. If dispatch
cannot be exercised here, rename the test to reflect only contract derivation
and runtime validation.
- Around line 1717-1732: Update the type annotations in _CudaGraphRunner to
match the layout identities stored by the warmup logic: change _observed from
nested integer lists to entries containing the _layout_identity return type, and
change _pinned to map integers to that same layout-identity type. Preserve the
existing initialization and storage behavior while ensuring mypy accepts the
assignments.
---
Nitpick comments:
In `@fastvideo/optimization/dispatch.py`:
- Around line 284-289: Update the SHADOW documentation in the timing module to
state that shadow mode performs an extra native forward and can change outputs
for stateful forwards that consume RNG, update caches, or mutate module state.
Explicitly warn against enabling SHADOW during parity or generation-result
comparisons, while retaining the existing notes about discarded results and
increased cost.
In `@fastvideo/optimization/subgraph.py`:
- Around line 182-204: The _dump_graph_profile output path currently overwrites
prior block profiles; update its filename construction to include the current
block’s unique identity, using available manifest or build-context symbols,
while preserving the configured destination directory and JSON output behavior.
- Around line 950-961: The CudaGraphUnavailable warning in the graph replay
exception path is emitted once per block instead of once per process. Update the
logger.warning call in the CudaGraphUnavailable handler to use the existing
logger.warning_once mechanism, preserving the current message and eager-replay
behavior while still clearing entry.cuda_graph for the affected block.
- Around line 209-214: Register FASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHS and
FASTVIDEO_OPTIMIZATION_ARTIFACT_DUMP_GRAPH in fastvideo/envs.py alongside the
existing optimization artifact entries, then update subgraph.py to read them
through envs when building _Entry instead of os.getenv. Register
FASTVIDEO_OPTIMIZATION_ARTIFACT_TIMING there as well and update timing.py so
ENABLED, SYNCHRONIZE, and SHADOW remain module-level constants derived from the
registered value. Apply these changes at fastvideo/optimization/subgraph.py
lines 209-214, fastvideo/optimization/timing.py lines 41-48, and
fastvideo/envs.py lines 355-361; preserve the existing comment style and boolean
hot-path behavior.
- Around line 461-472: In _release_capture, read self._pending_graph into
pending and clear the attribute in separate statements, relying on its
initialization rather than getattr. In the capture/replay path around
self._moving[index], replace direct indexing with a guarded lookup that raises
CudaGraphUnavailable when the index is absent, preserving
GraphDispatchSession._dispatch’s recoverable failure handling.
In `@fastvideo/optimization/timing.py`:
- Around line 105-107: Bound the `_notes` dictionary updated by `note()` so
variable message values cannot create unlimited distinct keys. Add a fixed
maximum cardinality and ensure new note keys are ignored once the cap is
reached, while continuing to increment existing keys and preserving the current
message truncation behavior.
- Around line 151-163: The write_report function currently writes timing.json in
place, allowing readers to observe partial output. Match
GraphDispatchSession.write_diagnostics by writing the serialized snapshot to a
temporary file in the destination directory, then atomically replacing the
target with it; preserve the existing None-on-failure behavior and ensure
temporary artifacts are cleaned up as appropriate.
In `@fastvideo/tests/optimization/test_dispatch.py`:
- Around line 1894-1907: Add a fixture for the imported timing module that calls
timing.reset() both before and after tests, and apply it to
test_timing_state_can_be_reset_between_sessions (or make it autouse for the
module). Keep the existing ENABLED monkeypatch and assertions, ensuring _totals,
_counts, and _notes are cleared before and after each test.
- Around line 1541-1586: Update the assertion-stripping tests around
_assertion_graph and test_export_runtime_assertions_are_stripped to include a
real torch export assertion target, or explicitly assert its expected target
string. Ensure the test verifies that _ASSERTION_TARGETS contains the actual
aten._assert_tensor_metadata.default target and that stripping removes it, so
PyTorch target renames cannot silently make the logic a no-op.
In `@fastvideo/tests/optimization/test_fx_capture.py`:
- Around line 496-520: Add a CUDA-gated regression test alongside
test_deferred_export_replays_observed_autocast_dtype that runs the same capture
and finalize flow under CUDA autocast, verifies the captured executable IR
preserves the expected CUDA autocast dtype, and confirms autocast state is
restored after execution and finalize. Document the test’s GPU requirement and
skip it when CUDA is unavailable.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ec8684f9-824e-45ac-bab4-7e779eb5ec8f
📒 Files selected for processing (8)
fastvideo/envs.pyfastvideo/optimization/artifact.pyfastvideo/optimization/dispatch.pyfastvideo/optimization/fx_capture.pyfastvideo/optimization/subgraph.pyfastvideo/optimization/timing.pyfastvideo/tests/optimization/test_dispatch.pyfastvideo/tests/optimization/test_fx_capture.py
Comment on lines
+310
to
326
| # The FSDP-lifecycle snapshot is built here, on the failure path, | ||
| # rather than three times per successful call. It exists to | ||
| # diagnose parameter materialization; a run that never fails never | ||
| # needs it, and paying for it on every dispatch charged the | ||
| # candidate's measured saving for diagnostics it did not use. | ||
| logger.warning( | ||
| "Artifact %s failed at runtime for %s; falling back to native " | ||
| "execution for the rest of this run; materialization=%s", | ||
| decision.artifact_id, | ||
| wrapper.scope, | ||
| json.dumps(materialization, sort_keys=True), | ||
| json.dumps( | ||
| {"at_failure": self._parameter_snapshot(wrapper)}, | ||
| sort_keys=True, | ||
| ), | ||
| exc_info=True, | ||
| ) | ||
| return native(*args, **kwargs) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Guard the diagnostic snapshot so the fallback cannot be lost.
self._parameter_snapshot(wrapper) runs inside the except handler. If it raises, the new exception replaces the fallback and propagates to the caller, so return native(*args, **kwargs) at Line 326 never runs. Every other failure path in this module falls back to native execution instead of breaking generation.
Compute the snapshot defensively.
🛡️ Proposed fix
decision.reason = f"{_REASON_RUNTIME_PREFIX}:{type(exc).__name__}"
+ try:
+ snapshot = self._parameter_snapshot(wrapper)
+ except Exception: # noqa: BLE001 - diagnostics never break a run
+ snapshot = {"unavailable": True}
# The FSDP-lifecycle snapshot is built here, on the failure path,
# rather than three times per successful call. It exists to
# diagnose parameter materialization; a run that never fails never
# needs it, and paying for it on every dispatch charged the
# candidate's measured saving for diagnostics it did not use.
logger.warning(
"Artifact %s failed at runtime for %s; falling back to native "
"execution for the rest of this run; materialization=%s",
decision.artifact_id,
wrapper.scope,
- json.dumps(
- {"at_failure": self._parameter_snapshot(wrapper)},
- sort_keys=True,
- ),
+ json.dumps({"at_failure": snapshot}, sort_keys=True),
exc_info=True,
)🧰 Tools
🪛 ast-grep (0.45.0)
[info] 319-322: use jsonify instead of json.dumps for JSON output
Context: json.dumps(
{"at_failure": self._parameter_snapshot(wrapper)},
sort_keys=True,
)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/optimization/dispatch.py` around lines 310 - 326, Guard the
diagnostic snapshot inside the runtime-failure handler around the fallback logic
near the artifact execution path. Ensure any exception from
self._parameter_snapshot(wrapper) is caught or replaced with a safe diagnostic
value, then still emit the warning and always reach return native(*args,
**kwargs), preserving native fallback even when snapshot collection fails.
Comment on lines
+683
to
+686
| output = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS | ||
| if output: | ||
| session.write_diagnostics(output) | ||
| timing.write_report(Path(str(output)).with_name("timing.json")) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Protect teardown from an invalid diagnostics path.
Path(str(output)).with_name("timing.json") runs outside any error handling. with_name raises ValueError when the path has no name component, for example FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS="." or "/". The exception then escapes detach_graph_dispatch and breaks teardown. session.write_diagnostics handles the same input safely because it catches internally.
🛡️ Proposed fix
output = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS
if output:
session.write_diagnostics(output)
- timing.write_report(Path(str(output)).with_name("timing.json"))
+ try:
+ report = Path(str(output)).expanduser().with_name("timing.json")
+ except ValueError:
+ logger.warning("Cannot derive a timing report path from %s", output)
+ else:
+ timing.write_report(report)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| output = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS | |
| if output: | |
| session.write_diagnostics(output) | |
| timing.write_report(Path(str(output)).with_name("timing.json")) | |
| output = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS | |
| if output: | |
| session.write_diagnostics(output) | |
| try: | |
| report = Path(str(output)).expanduser().with_name("timing.json") | |
| except ValueError: | |
| logger.warning("Cannot derive a timing report path from %s", output) | |
| else: | |
| timing.write_report(report) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/optimization/dispatch.py` around lines 683 - 686, Update the
diagnostics handling in detach_graph_dispatch so an invalid output path cannot
escape teardown: guard the timing report path construction and write, including
Path.with_name, with error handling while preserving session.write_diagnostics
behavior. Ensure values such as "." or "/" are safely ignored or logged without
preventing detach_graph_dispatch from completing.
Comment on lines
+55
to
+74
| def _layout_identity(value: Any) -> tuple[Any, ...] | None: | ||
| """Everything a capture bakes in about one tensor's memory. | ||
|
|
||
| ``data_ptr()`` alone is not enough: it folds in ``storage_offset`` but says | ||
| nothing about strides, so a contiguous tensor and a permuted view handed | ||
| back the same allocator block compare equal while the captured kernels read | ||
| elements in a different order. Device and dtype are included for the same | ||
| reason -- the capture encodes all of it. | ||
| """ | ||
| try: | ||
| return ( | ||
| value.data_ptr(), | ||
| tuple(value.shape), | ||
| tuple(value.stride()), | ||
| str(value.dtype), | ||
| str(value.device), | ||
| value.storage_offset(), | ||
| ) | ||
| except Exception: # noqa: BLE001 - anything unreadable is not capturable | ||
| return None |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
An unreadable input layout passes the pinned-address check.
_layout_identity returns None when it cannot read a tensor. At Line 362 the capture stores that None into self._pinned[index] without rejecting it. At Line 535 the replay then compares _layout_identity(live) != pinned. If the live tensor is also unreadable, both sides are None, the comparison passes, and the graph replays against an address the runner never verified.
The capture already rejects this case for bound attributes at Lines 418-421. Apply the same rule to pinned inputs.
🛡️ Proposed fix in `_capture`
if stable[index]:
+ if pointers[index] is None:
+ raise CudaGraphUnavailable(
+ f"runtime input {index} has no readable layout"
+ )
capture_inputs.append(value)
self._pinned[index] = pointers[index]
continueAlso applies to: 530-539
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/optimization/subgraph.py` around lines 55 - 74, Update _capture to
reject any pinned input whose _layout_identity returns None, matching the
existing unreadable bound-attribute handling. Ensure the capture fails before
storing None in self._pinned, and keep replay validation in the pinned-input
path from treating two None layouts as a valid match.
Comment on lines
+815
to
+818
| # Value is (rewritten module, frozen placeholder contract): the contract is | ||
| # derived once per build so the per-call path never walks the graph. | ||
| cache: weakref.WeakKeyDictionary[nn.Module, _Entry] = weakref.WeakKeyDictionary() | ||
| graph_scope = _CudaGraphScope() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the stale cache comment.
The comment states the value is "(rewritten module, frozen placeholder contract)". The value is now an _Entry, which also carries cuda_graph.
✏️ Proposed fix
- # Value is (rewritten module, frozen placeholder contract): the contract is
- # derived once per build so the per-call path never walks the graph.
+ # Value is an _Entry holding the rewritten module, the frozen placeholder
+ # contract, and the optional CUDA graph runner. The contract is derived
+ # once per build so the per-call path never walks the graph.
cache: weakref.WeakKeyDictionary[nn.Module, _Entry] = weakref.WeakKeyDictionary()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # Value is (rewritten module, frozen placeholder contract): the contract is | |
| # derived once per build so the per-call path never walks the graph. | |
| cache: weakref.WeakKeyDictionary[nn.Module, _Entry] = weakref.WeakKeyDictionary() | |
| graph_scope = _CudaGraphScope() | |
| # Value is an _Entry holding the rewritten module, the frozen placeholder | |
| # contract, and the optional CUDA graph runner. The contract is derived | |
| # once per build so the per-call path never walks the graph. | |
| cache: weakref.WeakKeyDictionary[nn.Module, _Entry] = weakref.WeakKeyDictionary() | |
| graph_scope = _CudaGraphScope() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/optimization/subgraph.py` around lines 815 - 818, Update the cache
comment above the WeakKeyDictionary in the graph-scope initialization to
describe the value as an _Entry, including its rewritten module, frozen
placeholder contract, and cuda_graph; keep the cache declaration and behavior
unchanged.
Comment on lines
+118
to
+128
| def reset() -> None: | ||
| """Clear all accumulated timing state. | ||
|
|
||
| The counters are process-global, so two dispatch sessions in one process -- | ||
| serving two models in sequence, say -- would otherwise contribute to the | ||
| same totals and produce a report attributable to neither. Call this when | ||
| starting a session whose measurements must stand alone. | ||
| """ | ||
| _totals.clear() | ||
| _counts.clear() | ||
| _notes.clear() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find every call site of timing.reset outside the tests.
rg -n -C3 --type=py 'timing\.reset\s*\(|from fastvideo\.optimization import timing'Repository: aryan5v/FastVideo
Length of output: 155
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files =="
git ls-files | rg '(^|/)timing\.py$|dispatch.*\.py$|test.*timing.*\.py$' || true
echo
echo "== timing calls/import search case-insensitive exact =="
rg -n -C2 --type=py 'timing\.reset\s*\(|\.reset\s*\(\)|reset\s*\(\)' || true
rg -n -C2 --type=py 'from fastvideo\.optimization import timing|from fastvideo\.optimization\. import timing|fastvideo/optimization/timing' || true
echo
echo "== fastvideo/optimization/timing.py outline/contents =="
wc -l fastvideo/optimization/timing.py
sed -n '1,220p' fastvideo/optimization/timing.py
echo
echo "== dispatch imports and timing usage =="
rg -n -C2 --type=py 'timing|timing\.' fastvideo || trueRepository: aryan5v/FastVideo
Length of output: 394
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== timing.py relevant lines =="
sed -n '1,220p' fastvideo/optimization/timing.py
echo
echo "== dispatch.py timing-related context =="
rg -n -C3 'timing|write_report|record|reset|snapshot' fastvideo/optimization/dispatch.py || true
echo
echo "== timing calls across Python files (broad exact name) =="
rg -n -C2 --type=py '\.reset\s*\(|timing\.reset|record\s*\(|snapshot\s*\(|write_report\s*\(' fastvideo || true
echo
echo "== timing module references across all tracked text files =="
rg -n -C2 'timing\.record|timing\.snapshot|timing\.write_report|timing\.reset|reset\(\)' . || trueRepository: aryan5v/FastVideo
Length of output: 50373
Call timing.reset() before starting each dispatch session.
timing.phase(...), timing.record(...), and timing.note(...) all use process-global state. timing.write_report() is only at the diagnostics path, while production dispatch code never clears totals before accumulating new sessions. A multi-model single-process run can mix timing data unless timing.reset() is called at session start.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/optimization/timing.py` around lines 118 - 128, Call timing.reset()
at the beginning of every dispatch session, before any timing.phase(),
timing.record(), or timing.note() calls occur. Update the dispatch-session entry
point rather than the timing.reset() implementation, ensuring each model/session
starts with cleared process-global timing state while preserving existing report
generation.
Comment on lines
+1470
to
+1498
| def test_placeholder_contract_is_derived_once_not_per_call(): | ||
| """The per-call path must not walk the graph to rediscover placeholders.""" | ||
| import operator | ||
|
|
||
| from torch import fx | ||
|
|
||
| from fastvideo.optimization.subgraph import ( | ||
| _placeholder_contract, | ||
| _validate_runtime_inputs, | ||
| ) | ||
|
|
||
| graph = fx.Graph() | ||
| first = graph.placeholder("x") | ||
| first.meta["val"] = torch.zeros(2, 3) | ||
| second = graph.placeholder("y") | ||
| second.meta["val"] = torch.zeros(4) | ||
| for index in range(200): | ||
| graph.call_function(operator.add, args=(first, index)) | ||
| graph.output(first) | ||
|
|
||
| contract = _placeholder_contract(graph) | ||
| assert contract == (("x", ((2, 3), "torch.float32")), ("y", ((4,), "torch.float32"))) | ||
|
|
||
| _validate_runtime_inputs(contract, [torch.zeros(2, 3), torch.zeros(4)]) | ||
|
|
||
| with pytest.raises(SubgraphRewriteError, match="metadata changed"): | ||
| _validate_runtime_inputs(contract, [torch.zeros(2, 5), torch.zeros(4)]) | ||
| with pytest.raises(SubgraphRewriteError, match="input count differs"): | ||
| _validate_runtime_inputs(contract, [torch.zeros(2, 3)]) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
The test name claims more than the body verifies.
test_placeholder_contract_is_derived_once_not_per_call checks the contract value and the two _validate_runtime_inputs failure modes. It never checks that the per-call path avoids walking the graph. The 200 filler nodes at Lines 1486-1487 make the walk expensive but nothing measures or counts it.
Either rename the test to describe what it asserts, or add an assertion that dispatch derives the contract once. Counting calls to _placeholder_contract with monkeypatch would cover the caching claim, in the same style as test_dispatch_does_not_snapshot_parameters_on_the_success_path at Line 1501.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/tests/optimization/test_dispatch.py` around lines 1470 - 1498,
Update test_placeholder_contract_is_derived_once_not_per_call to verify its
caching claim by monkeypatching _placeholder_contract and asserting dispatch
invokes it only once across repeated calls, following the pattern used by
test_dispatch_does_not_snapshot_parameters_on_the_success_path. If dispatch
cannot be exercised here, rename the test to reflect only contract derivation
and runtime validation.
Comment on lines
+1717
to
+1732
| def test_warmup_records_input_addresses_for_the_stability_decision(): | ||
| from fastvideo.optimization.subgraph import ( | ||
| CudaGraphUnavailable, | ||
| _CudaGraphRunner, | ||
| _CudaGraphScope, | ||
| ) | ||
|
|
||
| runner = _CudaGraphRunner(runnable=None, scope=_CudaGraphScope()) | ||
| tensor = torch.zeros(4) | ||
| for _ in range(_CudaGraphRunner.WARMUP_ITERATIONS): | ||
| with pytest.raises(CudaGraphUnavailable, match="warming up"): | ||
| runner([tensor]) | ||
| from fastvideo.optimization.subgraph import _layout_identity | ||
|
|
||
| assert len(runner._observed) == _CudaGraphRunner.WARMUP_ITERATIONS | ||
| assert all(seen == [_layout_identity(tensor)] for seen in runner._observed) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
These assertions expose two wrong type annotations in the runner.
The test asserts that each entry of runner._observed is a _layout_identity tuple. In fastvideo/optimization/subgraph.py the annotations declare integers:
- Line 313:
self._observed: list[list[int | None]] = [] - Line 315:
self._pinned: dict[int, int] = {}
_layout_identity returns tuple[Any, ...] | None, and Line 362 stores that value into _pinned. The coding guidelines list mypy among the configured tools, so the annotations should match.
✏️ Proposed fix in `fastvideo/optimization/subgraph.py`
- self._observed: list[list[int | None]] = []
+ self._observed: list[list[tuple[Any, ...] | None]] = []
#: index -> captured address, for inputs read where they live.
- self._pinned: dict[int, int] = {}
+ self._pinned: dict[int, tuple[Any, ...] | None] = {}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_warmup_records_input_addresses_for_the_stability_decision(): | |
| from fastvideo.optimization.subgraph import ( | |
| CudaGraphUnavailable, | |
| _CudaGraphRunner, | |
| _CudaGraphScope, | |
| ) | |
| runner = _CudaGraphRunner(runnable=None, scope=_CudaGraphScope()) | |
| tensor = torch.zeros(4) | |
| for _ in range(_CudaGraphRunner.WARMUP_ITERATIONS): | |
| with pytest.raises(CudaGraphUnavailable, match="warming up"): | |
| runner([tensor]) | |
| from fastvideo.optimization.subgraph import _layout_identity | |
| assert len(runner._observed) == _CudaGraphRunner.WARMUP_ITERATIONS | |
| assert all(seen == [_layout_identity(tensor)] for seen in runner._observed) | |
| self._observed: list[list[tuple[Any, ...] | None]] = [] | |
| #: index -> captured address, for inputs read where they live. | |
| self._pinned: dict[int, tuple[Any, ...] | None] = {} |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/tests/optimization/test_dispatch.py` around lines 1717 - 1732,
Update the type annotations in _CudaGraphRunner to match the layout identities
stored by the warmup logic: change _observed from nested integer lists to
entries containing the _layout_identity return type, and change _pinned to map
integers to that same layout-identity type. Preserve the existing initialization
and storage behavior while ensuring mypy accepts the assignments.
Source: Coding guidelines
Comment on lines
+1853
to
+1882
| def test_a_mutable_output_is_refused_rather_than_aliased(monkeypatch): | ||
| """A list output would hand back the graph's own static buffers.""" | ||
| from fastvideo.optimization.subgraph import ( | ||
| CudaGraphUnavailable, | ||
| _CudaGraphRunner, | ||
| _CudaGraphScope, | ||
| ) | ||
|
|
||
| runner = _CudaGraphRunner(runnable=None, scope=_CudaGraphScope()) | ||
|
|
||
| class _FakeGraph: | ||
| def replay(self): | ||
| pass | ||
|
|
||
| def reset(self): | ||
| pass | ||
|
|
||
| # Drive just the output-classification branch of _capture. | ||
| leaves = ([torch.zeros(2)],) | ||
| with pytest.raises(CudaGraphUnavailable, match="mutable list"): | ||
| runner._output_constants = {} | ||
| for position, leaf in enumerate(leaves): | ||
| if isinstance(leaf, torch.Tensor): | ||
| continue | ||
| if leaf is None or isinstance(leaf, (bool, int, float, complex, str, bytes)): | ||
| continue | ||
| raise CudaGraphUnavailable( | ||
| f"output {position} is a mutable {type(leaf).__name__}; " | ||
| "returning it uncopied would alias the graph's static buffers" | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
This test asserts against a copy of the source, not against the runner.
Lines 1874-1882 re-implement the output-classification loop from _capture inside the test body and then assert that the copy raises. No production code runs. runner, _FakeGraph, and the monkeypatch fixture are created and never used.
The test therefore passes even if the real branch at fastvideo/optimization/subgraph.py Lines 393-402 is deleted. That branch is safety-critical: it prevents returning the graph's static buffers to the caller, which the next replay would overwrite.
Extract the classification into a helper that both _capture and this test call, then assert against the helper. A helper also keeps the test honest if the accepted-constant type list changes.
♻️ Suggested direction in `fastvideo/optimization/subgraph.py`
def _classify_output_leaves(leaves: tuple[Any, ...]) -> dict[int, Any]:
"""Return the non-tensor output constants, or decline a mutable leaf."""
import torch
constants: dict[int, Any] = {}
for position, leaf in enumerate(leaves):
if isinstance(leaf, torch.Tensor):
continue
if leaf is None or isinstance(
leaf, (bool, int, float, complex, str, bytes)
):
constants[position] = leaf
continue
raise CudaGraphUnavailable(
f"output {position} is a mutable {type(leaf).__name__}; "
"returning it uncopied would alias the graph's static buffers"
)
return constants_capture then calls self._output_constants = _classify_output_leaves(leaves), and the test calls _classify_output_leaves(([torch.zeros(2)],)).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastvideo/tests/optimization/test_dispatch.py` around lines 1853 - 1882,
Extract the output-leaf classification logic into a shared helper near
_CudaGraphRunner, such as _classify_output_leaves, returning constant leaves and
raising CudaGraphUnavailable for mutable values. Update
_CudaGraphRunner._capture to use this helper, then change
test_a_mutable_output_is_refused_rather_than_aliased to call the helper directly
and remove the unused runner, fake graph, and monkeypatch setup.
…ent/v1-r4-dispatch-fix # Conflicts: # fastvideo/tests/optimization/test_fx_capture.py
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
1 participant
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.
Summary
Completes the FastVideo runtime required to execute MotionKernel graph-derived artifacts safely and economically.
Why
Eager replay of a 621-node exported transformer subgraph cost more than the native block. CUDA-graph replay preserves the same kernels, parameters, and order while eliminating host dispatch overhead. Capture remains fail-closed and is accepted only after the artifact entry point and eager replay match bitwise.
Validation
pytest -q fastvideo/tests/optimization: 113 passedEvidence:
/mnt/nfs/vlm-aryan/ltx-v1-r4-targeted-fix-20260801-203751Summary by CodeRabbit
New Features
Bug Fixes
Tests
Greptile Summary
This PR completes the CUDA-graph replay layer for MotionKernel graph-derived artifacts. It adds
_CudaGraphRunner(per-block capture with bitwise verification, per-scope shared static input buffers, and graceful eager fallback),_CudaGraphScope(one buffer per input position shared across repeated blocks),_CudaGraphScope.buffer_forwith full layout/stride/device checking,_CudaGraphRunner._release_capturefor deterministic pool cleanup, and a newtiming.pymodule for opt-in per-phase host/device profiling with shadow mode.subgraph.py): warmup for allocator stability, capture on a private side stream, bitwise verification of replay vs. eager, attribute-address staleness checks on every replay, and explicit pool release on capture failure.artifact.py,envs.py): newFASTVIDEO_OPTIMIZATION_ARTIFACT_ENABLEenv var for per-artifact A/B isolation; all four previously undeclared env vars are now registered inenvs.py.dispatch.py,timing.py): parameter snapshot deferred to failure path only; timing phases instrumented around every hot-path branch; shadow mode for live native-vs-candidate comparison."Confidence Score: 5/5
Safe to merge. The CUDA-graph path is fail-closed at every stage, fallback to eager is tested on CPU, and bitwise verification guards the captured result before it is ever used.
All failure paths fall back to native execution without demoting the artifact permanently. The capture, verification, and replay logic is internally consistent, and the previous review findings (stride check in buffer_for, env var registration, timing.reset at session start) are all addressed. The one asymmetry noted — explicit pool release on capture-fail but not on post-capture validation-fail — is benign in CPython where refcount GC is immediate.
Files Needing Attention: fastvideo/optimization/subgraph.py — the post-capture CudaGraphUnavailable paths (arity mismatch, moved attribute, etc.) skip the explicit _release_capture() call that the capture-fail path uses for deterministic pool cleanup.
Important Files Changed
_CudaGraphRunner, including per-block buffer sharing via_CudaGraphScope, bitwise verification, attribute staleness checks, and graceful eager fallback; also adds_strip_runtime_assertionsand_placeholder_contractfor hot-path reduction. One asymmetry in explicit pool release between capture-fail and post-capture-fail paths.reset()is now called inattach_graph_dispatchaddressing the prior review concern.environment_variablesregistry, resolving the prior review finding.enabled_idsallowlisting toArtifactRegistry: all bundles are still verified; selection is a separate narrowing step. Missing-ID requests are surfaced as errors rather than silent no-ops.observed_autocastafter finalization.Flowchart
%%{init: {'theme': 'neutral'}}%% flowchart TD A[dispatch called] --> B{entry cached?} B -- No --> C[build runnable + _Entry with _CudaGraphRunner] C --> D B -- Yes --> D{cuda_graph is not None?} D -- No --> E[eager runnable] D -- Yes --> F[_CudaGraphRunner.__call__] F --> G{self._graph is None?} G -- warmups < WARMUP_ITERATIONS --> H[record addresses / raise CudaGraphWarmingUp] H --> I[fall through to eager] I --> E G -- warmups reached --> J[_capture: build inputs, warmup 3x, graph.capture, bitwise verify] J -- success --> K[self._graph set / next call replays] J -- CudaGraphUnavailable --> L[_release_capture / raise] J -- unexpected error --> M[_release_capture / synchronize / wrap in CudaGraphUnavailable] L --> N[entry.cuda_graph = None / eager forever] M --> N G -- self._graph set --> O{validate inputs and attributes} O -- mismatch --> P[raise CudaGraphUnavailable / NO explicit _release_capture] P --> N O -- OK --> Q[capture_stream.wait / graph.replay / current_stream.wait / clone outputs] Q --> R[return _unflatten] E --> RReviews (3): Last reviewed commit: "Merge remote-tracking branch 'fork/agent..." | Re-trigger Greptile