Release generic MotionKernel artifact dispatch V1 - #26
Open
aryan5v wants to merge 69 commits into
Open
Conversation
Add a model-agnostic FastVideo launcher that executes one generation mode from a versioned MotionKernel workload manifest. Writes structured result JSON for native-versus-optimized end-to-end measurement without model-specific callables.
Prefer exact mode_env keys, keep dry-run from mutating os.environ, import torch once for timing, and load the launcher in tests without rewriting sys.path at import time.
Add a model-independent FX capture hook that runs alongside the existing worker-side optimization profile. Repeated module stacks are found structurally (any nn.ModuleList whose children share one class), so no architecture is named anywhere in the capture path. - forward hooks record only tensor signatures and call counters - symbolic tracing runs after the profiler window closes, so captured graphs never contaminate the exported timings - export gains optional regions / graph_breaks / unsupported keys plus a separately versioned capture block; readers that only understand rows are unaffected - capture stays off unless the optimization profile is requested and FASTVIDEO_OPTIMIZATION_PROFILE_CAPTURE_FX is set - trace and finalize failures are recorded as data; generation and the timing export continue - payloads are asserted metadata-only: no tensor values, weights, or prompts
Run a packaged optimization artifact in place of a repeated block's forward when, and only when, it provably matches -- otherwise run natively. Dispatch attaches to the same model-independent structure capture uses: children of an nn.ModuleList that share a class. No Wan-, LTX-, Cosmos- or Kandinsky-specific conditional exists anywhere in this path. Per stack and per observed input signature it runs the first call natively (which reveals the output signature), recomputes the module's graph fingerprint through the capture module so the value matches what the producer recorded, then selects a bundle whose fingerprint, tensor signatures and declared environment all match. The entry point is called as candidate(module, *args, **kwargs). Trust: executable code is loaded only from FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR, every declared file is re-hashed immediately before import, undeclared files in a bundle are a hard rejection, and bundles are imported under a private module namespace rather than sys.path. Fallback: a missing match, an unloadable bundle, an untraceable module or an exception from the candidate falls back to native execution and records a structured reason; a candidate that raised is demoted, not retried. The optional diagnostics report is metadata only. With FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR unset nothing is wrapped at all -- no forward is patched, no graph is traced, no artifact code is read. Tests: 28 new CPU tests with fake kernels and fake modules; 60 passed across the optimization suites. Verified end to end against a bundle packaged by MotionKernel's packager. Note: the pre-commit mypy hook fails with "FastVideo-v1-dispatch is not a valid Python package name" -- that is the checkout directory name and reproduces on untouched files. yapf, ruff and codespell pass.
[feat]: dispatch trusted graph artifacts with native fallback
…ent/v1-subgraph-dispatch
feat: execute graph-derived artifacts through generic dispatch
FastVideo's generic FX capture already found tensor metadata inside dataclass inputs, but torch.export flattens inputs and outputs through pytree and rejects any dataclass that is not a registered node. The dominant repeated block of models that pass a nested dataclass therefore fell back to symbolic tracing, which succeeds without example inputs and so produced a region with no executable IR and no graph break explaining the gap. Recursively collect the dataclass types reachable from the observed call and register them with torch.export.register_dataclass before export and dynamo. Traversal descends through dataclass fields, lists, tuples and mappings, is cycle-safe, and is bounded by depth and node budget. Only type(...) is read; field values are walked for structure alone and are never compared, formatted or exported. Export rejects an unregistered dataclass in the output as well, so the observed return structure contributes its types too -- types only, never the returned objects. Registration is skipped when the type is already a pytree node, checked against the live registry rather than a private cache, so repeated capture emits no duplicate-registration warning and a registration made by the model or another library is never overwritten. Traversal and registration failures raise DataclassRegistrationError and become the existing sanitized capture-failure metadata, which exports only a reason code and the exception class name. Registering dataclasses makes their string fields reachable as export placeholders, whose values sit in node.meta and would match the enum-shaped safe-string pattern. String-valued graph inputs are now refused, so caller data such as model ids or captions cannot reach the exported IR. Symbolic tracing needs no registration and is unchanged, so the auto fallback order is preserved. No model, architecture, class or field name is referenced.
Register structured dataclass inputs for export capture
* [debug]: log subgraph signature rejection metadata * [fix]: replay autocast during deferred graph capture * [test]: document canonical IR dtype format
…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).
…ent/v1-r4-dispatch-fix # Conflicts: # fastvideo/tests/optimization/test_fx_capture.py
[feat]: complete safe CUDA-graph artifact dispatch V1
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
aryan5v has reached the 50-credit limit for trial accounts. To continue receiving code reviews, upgrade your plan.
📝 WalkthroughWalkthroughThe PR adds workload-driven generation tooling and a generic optimization system. The system captures FX metadata, validates trusted artifacts, dispatches compatible graphs, supports subgraph replay, exports profiling data, integrates with pipelines, and adds tests and CI updates. ChangesOptimization infrastructure
Workload-driven generation launcher
Compatibility and validation fixes
Pre-commit workflow update
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Pipeline
participant FXCaptureSession
participant GraphDispatchSession
participant ArtifactRegistry
participant Artifact
Pipeline->>FXCaptureSession: profile forward call
FXCaptureSession->>Pipeline: record invocation metadata
Pipeline->>GraphDispatchSession: execute module
GraphDispatchSession->>ArtifactRegistry: find compatible candidate
ArtifactRegistry->>Artifact: verify and load bundle
Artifact-->>GraphDispatchSession: return callable
GraphDispatchSession->>Artifact: execute candidate
Artifact-->>GraphDispatchSession: return output or failure
GraphDispatchSession-->>Pipeline: return candidate output or native fallback
Possibly related PRs
🚥 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 |
There was a problem hiding this comment.
aryan5v has reached the 50-credit limit for trial accounts. To continue receiving code reviews, upgrade your plan.
![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 (6)
fastvideo/tests/encoders/test_reason1_chat_template.py (1)
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse an expected-behavior test module name.
Rename this module to
test_reason1_chat_template_normalization.py. The current name identifies the feature but not the expected behavior.🤖 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/encoders/test_reason1_chat_template.py` at line 1, Rename the test module from test_reason1_chat_template.py to test_reason1_chat_template_normalization.py so its name explicitly reflects the normalization behavior being tested.Source: Coding guidelines
fastvideo/optimization/subgraph.py (1)
144-166: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winMatch assertion targets by schema name, not by
str(node.target).
fastvideo/optimization/fx_capture.pylines 597-602 documents thatOpOverload.__str__is not stable across PyTorch distributions, and that NVIDIA builds may render it as a repr-like value whiletarget._schema.namestays canonical._strip_runtime_assertionscomparesstr(node.target)against"aten._assert_tensor_metadata.default", so on such a build the comparison never matches and no assertion is stripped.The failure is silent and costs only performance, which is the exact cost this function exists to remove. Prefer the schema name with the string form as a fallback.
♻️ Proposed change
+def _assertion_key(target: Any) -> str: + """Canonical identity for an FX call target, stable across torch builds.""" + schema_name = getattr(getattr(target, "_schema", None), "name", None) + overload = getattr(target, "_overloadname", None) or "default" + if isinstance(schema_name, str) and "::" in schema_name: + namespace, operation = schema_name.split("::", 1) + return f"{namespace}.{operation}.{overload}" + return str(target) + + def _strip_runtime_assertions(graph: fx.Graph) -> int: @@ - if str(node.target) not in _ASSERTION_TARGETS: + if _assertion_key(node.target) not in _ASSERTION_TARGETS: continue🤖 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 144 - 166, Update _strip_runtime_assertions to identify call_function targets using the canonical target._schema.name, with str(node.target) retained only as a fallback when schema metadata is unavailable. Compare the resulting name against _ASSERTION_TARGETS so assertion removal works consistently across PyTorch distributions.fastvideo/optimization/fx_capture.py (1)
1391-1400: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAlign the
get_forward_contextfailure handling with_observe.
_observeat lines 1039-1047 catchesExceptionwhen it reads the forward context, and treats a missing context as optional metadata. Here onlyAssertionErroris caught. Ifget_forward_contextraises any other exception type,capture_export_invocationpropagates it.GraphDispatchSession._resolvethen recordsgraph_identity_unavailableand demotes the scope to native for the whole run, even though the context is optional metadata.Use the same broad handling as
_observe.♻️ Proposed change
try: from fastvideo.forward_context import get_forward_context context = get_forward_context() observed_context: tuple[Any, Any] | None = ( context.current_timestep, context.attn_metadata, ) - except AssertionError: + except Exception: # noqa: BLE001 - optional metadata only observed_context = 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/optimization/fx_capture.py` around lines 1391 - 1400, Update the forward-context handling in capture_export_invocation to catch Exception, matching _observe’s behavior when reading get_forward_context. Preserve observed_context = None for any context-read failure so optional metadata does not propagate into GraphDispatchSession._resolve or demote the run to native.fastvideo/optimization/artifact.py (1)
778-791: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
sys.dont_write_bytecodeis a process-global toggle.The loader sets
sys.dont_write_bytecode = Truearoundcompile/execand restores it infinally.compileandexecdo not write bytecode caches, so the toggle only affects any other import that runs in a different thread during this window; that import silently skips its bytecode cache. The intent stated in the docstring — preventing a.pycinside the bundle — is already satisfied because no source loader is used.Consider dropping the toggle, or documenting that it is defensive against transitive imports performed by the artifact itself.
🤖 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/artifact.py` around lines 778 - 791, Remove the process-global sys.dont_write_bytecode toggle and its restoration from the artifact loading flow around compile/exec, since compile and exec already avoid writing bundle bytecode caches. Keep the existing exception cleanup and rejection behavior in the artifact import path unchanged.fastvideo/optimization/__init__.py (1)
4-5: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueImporting this package freezes two environment settings.
fastvideo.optimization.dispatchtransitively importssubgraphandtiming. Both read environment variables at module import time:_CUDA_GRAPHS_ENABLEDinsubgraph.pyline 209, andENABLED,SYNCHRONIZE,SHADOWintiming.pylines 44-51. Any process that importsfastvideo.optimizationbefore settingFASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHSorFASTVIDEO_OPTIMIZATION_ARTIFACT_TIMINGwill not observe those values.
FASTVIDEO_OPTIMIZATION_ARTIFACT_DIRis read lazily inattach_graph_dispatch, so the two groups behave differently. Document the ordering requirement, or read the CUDA-graph flag at session build time inrewrite_exported_subgraph.🤖 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/__init__.py` around lines 4 - 5, Update the optimization initialization flow to avoid freezing CUDA-graph and timing environment settings when fastvideo.optimization is imported. Prefer reading the CUDA-graph flag at session-build time within rewrite_exported_subgraph, and ensure timing settings are likewise evaluated after environment configuration; otherwise document the required environment-variable ordering at the relevant public initialization symbols.fastvideo/tests/optimization/test_dispatch.py (1)
1940-2001: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the shape-variant budget and for
_distributed_mode.Two behaviors that this cohort documents as safety bounds have no test here:
GraphDispatchSession.max_shape_variants. No test drives more distinct input signatures than the budget, so theshape_variant_budget_exhaustedpath and the growth of_decisionsare unverified. See the related finding onfastvideo/optimization/dispatch.pylines 253-261._distributed_mode. No test covers the multi-rankunspecifiedresult or the exception path. See the related finding onfastvideo/optimization/dispatch.pylines 596-614.Do you want me to generate these tests, or open an issue to track them?
🤖 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 1940 - 2001, Add tests covering GraphDispatchSession.max_shape_variants by driving more distinct input signatures than the configured budget, then assert the shape_variant_budget_exhausted decision and bounded _decisions growth. Also add _distributed_mode tests for the multi-rank unspecified result and its exception path, using the existing dispatch test helpers and preserving current behavior.
🤖 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 @.github/workflows/ci-precommit.yml:
- Around line 61-64: Update the pull_request_target pre-commit workflow so the
always_run check-filenames hook cannot scan inherited files despite the diff
arguments; make that hook explicitly diff-aware or disable it for this
invocation, while retaining complete-tree validation through the existing
workflow_call path.
In `@fastvideo/envs.py`:
- Around line 365-368: The optimization flags use inconsistent, case-sensitive
boolean parsing. Add one shared helper that strips and lowercases values,
treating "", "0", "false", "no", and "off" as disabled; use it for
FASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHS in fastvideo/envs.py (lines 365-368)
and replace bool(_SETTING) in timing.py (lines 44-51), ensuring SYNCHRONIZE and
SHADOW remain gated by the normalized ENABLED result.
In `@fastvideo/hooks/hooks.py`:
- Around line 74-86: Update run_with_forward to track each successfully
completed pre_forward hook and invoke the corresponding post_forward hooks
during exception unwinding, preserving hook order and re-raising the original
failure. Update the dispatch path that catches candidate failures so the native
fallback is invoked through run_with_forward rather than directly, and add a
regression test covering a raising candidate with matching pre-hook and
post-hook counts.
In `@fastvideo/optimization/dispatch.py`:
- Around line 283-317: Update the candidate exception path in the dispatch
method to run the matched ModuleHookManager post-forward cleanup before invoking
native(*args, **kwargs). Reuse the hook manager and hook context established for
this dispatch, ensure cleanup occurs after candidate failure and before native
fallback, and preserve the existing demotion, logging, and fallback behavior.
- Around line 596-614: Update _distributed_mode so the exception handler returns
"single" only when torch.distributed is unavailable; return "unspecified" for
failures from availability, initialization, or world-size checks, preserving the
documented fail-closed behavior for ambiguous multi-rank state.
- Around line 253-261: Update the dispatch decision caching around _dispatch and
_resolve so max_shape_variants bounds per-scope cached variants: track variant
counts directly rather than recomputing them from _decisions, and reuse one
shared shape-variant-budget-exhausted decision per scope keyed by (scope, "")
instead of storing a decision for every new shape_key. Ensure later signatures
reuse the shared over-budget decision while diagnostics and lookup costs remain
bounded.
In `@fastvideo/optimization/subgraph.py`:
- Around line 755-763: Update the per-manifest evaluation loop in
GraphDispatchSession._resolve to catch SubgraphRewriteError from
subgraph_signature_keys and reject only the current artifact with an appropriate
failure reason. Continue evaluating remaining candidates, preserving the
artifact isolation policy described by the existing artifact selection flow.
In `@fastvideo/pipelines/composed_pipeline_base.py`:
- Around line 240-243: Update the initialization flow around _dispatch_session
and attach_graph_dispatch to prevent enabling graph dispatch together with
enable_torch_compile unless explicit compatibility support is implemented and
covered. Prefer validating this configuration early and raising a clear error,
or otherwise ensure dispatch does not wrap compiled modules while preserving
existing behavior for either feature alone.
---
Nitpick comments:
In `@fastvideo/optimization/__init__.py`:
- Around line 4-5: Update the optimization initialization flow to avoid freezing
CUDA-graph and timing environment settings when fastvideo.optimization is
imported. Prefer reading the CUDA-graph flag at session-build time within
rewrite_exported_subgraph, and ensure timing settings are likewise evaluated
after environment configuration; otherwise document the required
environment-variable ordering at the relevant public initialization symbols.
In `@fastvideo/optimization/artifact.py`:
- Around line 778-791: Remove the process-global sys.dont_write_bytecode toggle
and its restoration from the artifact loading flow around compile/exec, since
compile and exec already avoid writing bundle bytecode caches. Keep the existing
exception cleanup and rejection behavior in the artifact import path unchanged.
In `@fastvideo/optimization/fx_capture.py`:
- Around line 1391-1400: Update the forward-context handling in
capture_export_invocation to catch Exception, matching _observe’s behavior when
reading get_forward_context. Preserve observed_context = None for any
context-read failure so optional metadata does not propagate into
GraphDispatchSession._resolve or demote the run to native.
In `@fastvideo/optimization/subgraph.py`:
- Around line 144-166: Update _strip_runtime_assertions to identify
call_function targets using the canonical target._schema.name, with
str(node.target) retained only as a fallback when schema metadata is
unavailable. Compare the resulting name against _ASSERTION_TARGETS so assertion
removal works consistently across PyTorch distributions.
In `@fastvideo/tests/encoders/test_reason1_chat_template.py`:
- Line 1: Rename the test module from test_reason1_chat_template.py to
test_reason1_chat_template_normalization.py so its name explicitly reflects the
normalization behavior being tested.
In `@fastvideo/tests/optimization/test_dispatch.py`:
- Around line 1940-2001: Add tests covering
GraphDispatchSession.max_shape_variants by driving more distinct input
signatures than the configured budget, then assert the
shape_variant_budget_exhausted decision and bounded _decisions growth. Also add
_distributed_mode tests for the multi-rank unspecified result and its exception
path, using the existing dispatch test helpers and preserving current behavior.
🪄 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: 663fa565-d039-49c2-ac5e-2e784648be35
📒 Files selected for processing (23)
.github/workflows/ci-precommit.ymlexamples/inference/optimizations/README.mdexamples/inference/optimizations/generation_launcher.pyfastvideo/envs.pyfastvideo/hooks/hooks.pyfastvideo/models/encoders/reason1.pyfastvideo/optimization/__init__.pyfastvideo/optimization/artifact.pyfastvideo/optimization/dispatch.pyfastvideo/optimization/fx_capture.pyfastvideo/optimization/identity.pyfastvideo/optimization/profiler.pyfastvideo/optimization/subgraph.pyfastvideo/optimization/timing.pyfastvideo/pipelines/composed_pipeline_base.pyfastvideo/tests/encoders/test_reason1_chat_template.pyfastvideo/tests/optimization/__init__.pyfastvideo/tests/optimization/test_dispatch.pyfastvideo/tests/optimization/test_fx_capture.pyfastvideo/tests/optimization/test_profiler.pytests/local_tests/optimizations/__init__.pytests/local_tests/optimizations/test_generation_launcher_workload.pytests/local_tests/optimizations/test_optimization_profiler.py
Comment on lines
+61
to
+64
| # A PR must not be forced to reformat unrelated inherited files. Lint | ||
| # only the base-to-head diff while retaining the base branch's trusted | ||
| # hook configuration. workflow_call still audits the complete tree. | ||
| extra_args: ${{ github.event_name == 'pull_request_target' && format('--from-ref {0} --to-ref HEAD --hook-stage manual', github.event.pull_request.base.sha) || '--all-files --hook-stage manual' }} |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
rg -n -C 8 'always_run:\s*true|git ls-files|repo:\s*local' \
.pre-commit-config.yamlRepository: aryan5v/FastVideo
Length of output: 1223
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== workflow file =="
fd -a 'ci-precommit\.yml$' . | sed 's#^\./##'
for f in $(fd 'ci-precommit\.yml$' .); do
echo "--- $f ---"
cat -n "$f"
done
echo
echo "== pre-commit action references =="
for f in *.yml .github/workflows/*.yml; do
[ -f "$f" ] || continue
rg -n -C 3 'pre-commit/action|id:\s*pre-commit|extra_args:|pre-commit run' "$f" || true
done
echo
echo "== config hook count/list =="
python3 - <<'PY'
from pathlib import Path
p=Path('.pre-commit-config.yaml')
for i,line in enumerate(p.read_text().splitlines(),1):
if line.strip().startswith('- id:'):
print(f"{i}: {line.strip()}")
PY
echo
echo "== action.yml from pre-commit/action v3.0.1 =="
set -euo pipefail
tmp="$(mktemp -d)"
git clone --depth 1 --branch v3.0.1 https://github.com/pre-commit/action.git "$tmp/action" 2>/dev/null || true
if [ -d "$tmp/action" ]; then
sed -n '1,220p' "$tmp/action/action.yml"
fiRepository: aryan5v/FastVideo
Length of output: 2369
Do not rely on PR diff arguments to skip always_run hooks.
pre-commit/action@v3.0.1 passes extra_args directly to pre-commit run. The check-filenames hook is local with always_run: true and uses git ls-files, so inherited files can still fail this PR-only workflow. Make the hook diff-aware, disable it for this invocation, or run full-tree checks separately.
🤖 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 @.github/workflows/ci-precommit.yml around lines 61 - 64, Update the
pull_request_target pre-commit workflow so the always_run check-filenames hook
cannot scan inherited files despite the diff arguments; make that hook
explicitly diff-aware or disable it for this invocation, while retaining
complete-tree validation through the existing workflow_call path.
Comment on lines
+365
to
+368
| # Replay rewritten subgraphs through CUDA graphs. Set to 0/false to force | ||
| # eager graph execution while retaining artifact dispatch. | ||
| "FASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHS": | ||
| lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHS", "1") not in {"0", "false", "False", ""}, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Inconsistent boolean parsing of the optimization environment flags. Both sites parse an operator-facing on/off flag with ad-hoc, case-sensitive logic instead of one shared rule. In each case a value an operator would reasonably use to turn the feature off leaves it on: CUDA_GRAPHS=FALSE keeps CUDA graph replay enabled, and TIMING=0 enables timing. Introduce one helper that normalizes the value with .strip().lower() and compares against {"", "0", "false", "no", "off"}, then use it at both sites.
fastvideo/envs.py#L365-L368: normalize theFASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHSvalue before the set membership test, soFALSE,no, andoffdisable replay.fastvideo/optimization/timing.py#L44-L51: replaceENABLED = bool(_SETTING)with the normalized off-value test, and gateSYNCHRONIZEandSHADOWon the resultingENABLED.
📍 Affects 2 files
fastvideo/envs.py#L365-L368(this comment)fastvideo/optimization/timing.py#L44-L51
🤖 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/envs.py` around lines 365 - 368, The optimization flags use
inconsistent, case-sensitive boolean parsing. Add one shared helper that strips
and lowercases values, treating "", "0", "false", "no", and "off" as disabled;
use it for FASTVIDEO_OPTIMIZATION_ARTIFACT_CUDA_GRAPHS in fastvideo/envs.py
(lines 365-368) and replace bool(_SETTING) in timing.py (lines 44-51), ensuring
SYNCHRONIZE and SHADOW remain gated by the normalized ENABLED result.
Comment on lines
+74
to
+86
| def run_with_forward( | ||
| self, | ||
| forward: Callable[..., Any], | ||
| *args: Any, | ||
| **kwargs: Any, | ||
| ) -> Any: | ||
| """Run ``forward`` through this module's installed hook lifecycle.""" | ||
| for hook in self.forward_hooks.values(): | ||
| args, kwargs = hook.pre_forward(self.module, *args, **kwargs) | ||
| output = forward(*args, **kwargs) | ||
| for hook in reversed(self.forward_hooks.values()): | ||
| output = hook.post_forward(self.module, output) | ||
| return output |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Unwind hooks when the forward callable fails.
run_with_forward skips post_forward when a pre-hook or forward raises. The dispatch caller catches candidate failures and invokes native outside this lifecycle. The failed candidate invocation can therefore leave offload or parameter-materialization state unbalanced.
Track successfully completed pre-hooks and unwind them on failure. Ensure the native fallback also runs through a balanced hook lifecycle. Add a regression test for a candidate that raises and verify matching pre-hook and post-hook counts.
🤖 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/hooks/hooks.py` around lines 74 - 86, Update run_with_forward to
track each successfully completed pre_forward hook and invoke the corresponding
post_forward hooks during exception unwinding, preserving hook order and
re-raising the original failure. Update the dispatch path that catches candidate
failures so the native fallback is invoked through run_with_forward rather than
directly, and add a regression test covering a raising candidate with matching
pre-hook and post-hook counts.
Comment on lines
+253
to
+261
| key = (wrapper.scope, shape_key) | ||
| decision = self._decisions.get(key) | ||
| if decision is None: | ||
| # First call for this signature: run native, learn the output | ||
| # layout, then decide once for every later call. | ||
| with timing.phase("dispatch.native_reference"): | ||
| output = native(*args, **kwargs) | ||
| self._decisions[key] = self._decide(wrapper, args, kwargs, output, input_metas) | ||
| return output |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
max_shape_variants does not bound the _decisions table.
_dispatch stores the returned _Decision for every new (scope, shape_key) at line 260, including the shape_variant_budget_exhausted decision produced by _resolve at lines 394-396. _variant_count then counts that entry too, so the count keeps rising past the budget and a new entry is added for every distinct input signature seen.
Two consequences for a model with dynamic input shapes:
_decisionsand thediagnostics()report grow without bound._variant_countisO(len(self._decisions))and runs on the first call of every new signature, so total cost is quadratic in the number of distinct signatures.
Cache one shared over-budget decision per scope instead of one per shape key, and track the variant count directly.
🐛 Proposed fix
@@ class GraphDispatchSession.__init__
self._decisions: dict[tuple[str, str], _Decision] = {}
+ self._variant_counts: dict[str, int] = defaultdict(int)
self._dropped_variants: dict[str, int] = defaultdict(int) def _variant_count(self, scope: str) -> int:
- return sum(1 for existing_scope, _ in self._decisions if existing_scope == scope)
+ return self._variant_counts[scope] def _resolve(self, ...):
scope = wrapper.scope
if self._variant_count(scope) >= self.max_shape_variants:
self._dropped_variants[scope] += 1
return _Decision(None, "shape_variant_budget_exhausted")
+ self._variant_counts[scope] += 1Then skip storing the entry when the budget is exhausted, so the table stays bounded:
- self._decisions[key] = self._decide(wrapper, args, kwargs, output, input_metas)
+ resolved = self._decide(wrapper, args, kwargs, output, input_metas)
+ if resolved.reason != "shape_variant_budget_exhausted":
+ self._decisions[key] = resolved
return outputNote that skipping the store makes every later call with that signature re-run _decide. If that is unacceptable, keep the store but cap it with a single shared over-budget entry keyed by (scope, "").
📝 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
| key = (wrapper.scope, shape_key) | |
| decision = self._decisions.get(key) | |
| if decision is None: | |
| # First call for this signature: run native, learn the output | |
| # layout, then decide once for every later call. | |
| with timing.phase("dispatch.native_reference"): | |
| output = native(*args, **kwargs) | |
| self._decisions[key] = self._decide(wrapper, args, kwargs, output, input_metas) | |
| return output | |
| key = (wrapper.scope, shape_key) | |
| decision = self._decisions.get(key) | |
| if decision is None: | |
| # First call for this signature: run native, learn the output | |
| # layout, then decide once for every later call. | |
| with timing.phase("dispatch.native_reference"): | |
| output = native(*args, **kwargs) | |
| resolved = self._decide(wrapper, args, kwargs, output, input_metas) | |
| if resolved.reason != "shape_variant_budget_exhausted": | |
| self._decisions[key] = resolved | |
| return output |
🤖 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 253 - 261, Update the
dispatch decision caching around _dispatch and _resolve so max_shape_variants
bounds per-scope cached variants: track variant counts directly rather than
recomputing them from _decisions, and reuse one shared
shape-variant-budget-exhausted decision per scope keyed by (scope, "") instead
of storing a decision for every new shape_key. Ensure later signatures reuse the
shared over-budget decision while diagnostics and lookup costs remain bounded.
Comment on lines
+283
to
+317
| try: | ||
| with timing.phase("dispatch.candidate_total"), self._materialized_candidate_parameters( | ||
| wrapper.parameter_manager): | ||
| hook_manager = ModuleHookManager.get_from(wrapper.module) | ||
| if hook_manager is None: | ||
| result = candidate_forward(*args, **kwargs) | ||
| else: | ||
| result = hook_manager.run_with_forward( | ||
| candidate_forward, | ||
| *args, | ||
| **kwargs, | ||
| ) | ||
| except Exception as exc: # noqa: BLE001 - untrusted candidate code | ||
| # Demote permanently: a candidate that raised once is not trusted | ||
| # to be retried thousands of times over the rest of the run. | ||
| decision.runtime_fallbacks += 1 | ||
| decision.candidate = None | ||
| decision.reason = f"{_REASON_RUNTIME_PREFIX}:{type(exc).__name__}" | ||
| # 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._safe_parameter_snapshot(wrapper)}, | ||
| sort_keys=True, | ||
| ), | ||
| exc_info=True, | ||
| ) | ||
| return native(*args, **kwargs) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect ModuleHookManager.run_with_forward exception handling.
set -euo pipefail
fd -i 'hooks.py' fastvideo/hooks --exec cat -n {}
ast-grep run --pattern 'def run_with_forward($$$):
$$$' --lang python fastvideo/hooks/hooks.pyRepository: aryan5v/FastVideo
Length of output: 6087
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "class LayerwiseOffloadHook|def pre_forward|def post_forward|materialized|_materialized|hooks\.py|hook_manager" fastvideo -S
for f in $(fd -i 'hooks.py' fastvideo); do
echo "=== $f ==="
sed -n '1,220p' "$f" | cat -n
doneRepository: aryan5v/FastVideo
Length of output: 9133
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== layerwise_offload.py ==="
cat -n fastvideo/hooks/layerwise_offload.py | sed -n '1,200p'
echo "=== dispatch _dispatch block ==="
cat -n fastvideo/optimization/dispatch.py | sed -n '243,335p'
echo "=== dispatcher hook attachment usages ==="
rg -n "append_forward_hook|replace_forward_hook|LayerwiseOffloadHook|_hook_manager|get_from_or_default" fastvideo -SRepository: aryan5v/FastVideo
Length of output: 17147
Run the hook post-forward cleanup on candidate runtime failures.
ModuleHookManager.run_with_forward() calls pre_forward, then forward, then post_forward; if candidate_forward raises, post_forward is skipped and the fallback calls native(*args, **kwargs) without re-entering the hook manager. This leaves a registered LayerwiseOffloadHook with its wait_and_replace_params() effects but none of its corresponding release_gpu_params() cleanup for that dispatch.
Run the matched reverse hook cleanup on this failure path before returning native.
🧰 Tools
🪛 ast-grep (0.45.0)
[info] 310-313: use jsonify instead of json.dumps for JSON output
Context: json.dumps(
{"at_failure": self._safe_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 283 - 317, Update the
candidate exception path in the dispatch method to run the matched
ModuleHookManager post-forward cleanup before invoking native(*args, **kwargs).
Reuse the hook manager and hook context established for this dispatch, ensure
cleanup occurs after candidate failure and before native fallback, and preserve
the existing demotion, logging, and fallback behavior.
Comment on lines
+596
to
+614
| def _distributed_mode() -> str: | ||
| """Report the sharding mode this process runs in. | ||
|
|
||
| A multi-rank run whose mode is not declared explicitly reports | ||
| ``unspecified``, which matches no artifact: an unsharded kernel silently | ||
| applied to a sharded module would be a correctness bug, so the ambiguous | ||
| case fails closed. | ||
| """ | ||
| declared = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE | ||
| if declared: | ||
| return declared | ||
| try: | ||
| import torch.distributed as distributed | ||
|
|
||
| if distributed.is_available() and distributed.is_initialized(): | ||
| return "single" if distributed.get_world_size() == 1 else "unspecified" | ||
| except Exception: # noqa: BLE001 - absence of distributed means single process | ||
| return "single" | ||
| return "single" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The broad except contradicts the documented fail-closed rule.
The docstring states that an ambiguous multi-rank run must report unspecified so no artifact matches. The handler at line 612 returns "single" for every exception, not only for a missing torch.distributed. If distributed.get_world_size() raises after is_initialized() returned True, this reports "single" and an unsharded kernel can be selected for a sharded module.
Narrow the handler: treat an absent torch.distributed as "single", and treat any other failure as "unspecified".
🐛 Proposed fix
declared = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE
if declared:
return declared
try:
import torch.distributed as distributed
-
- if distributed.is_available() and distributed.is_initialized():
- return "single" if distributed.get_world_size() == 1 else "unspecified"
- except Exception: # noqa: BLE001 - absence of distributed means single process
+ except ImportError:
+ # No distributed support built in: this is a single process.
return "single"
+ try:
+ if distributed.is_available() and distributed.is_initialized():
+ return "single" if distributed.get_world_size() == 1 else "unspecified"
+ except Exception: # noqa: BLE001 - an unreadable topology is ambiguous
+ # Ambiguous: matching no artifact is safer than assuming one rank.
+ return "unspecified"
return "single"📝 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 _distributed_mode() -> str: | |
| """Report the sharding mode this process runs in. | |
| A multi-rank run whose mode is not declared explicitly reports | |
| ``unspecified``, which matches no artifact: an unsharded kernel silently | |
| applied to a sharded module would be a correctness bug, so the ambiguous | |
| case fails closed. | |
| """ | |
| declared = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE | |
| if declared: | |
| return declared | |
| try: | |
| import torch.distributed as distributed | |
| if distributed.is_available() and distributed.is_initialized(): | |
| return "single" if distributed.get_world_size() == 1 else "unspecified" | |
| except Exception: # noqa: BLE001 - absence of distributed means single process | |
| return "single" | |
| return "single" | |
| def _distributed_mode() -> str: | |
| """Report the sharding mode this process runs in. | |
| A multi-rank run whose mode is not declared explicitly reports | |
| ``unspecified``, which matches no artifact: an unsharded kernel silently | |
| applied to a sharded module would be a correctness bug, so the ambiguous | |
| case fails closed. | |
| """ | |
| declared = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE | |
| if declared: | |
| return declared | |
| try: | |
| import torch.distributed as distributed | |
| except ImportError: | |
| # No distributed support built in: this is a single process. | |
| return "single" | |
| try: | |
| if distributed.is_available() and distributed.is_initialized(): | |
| return "single" if distributed.get_world_size() == 1 else "unspecified" | |
| except Exception: # noqa: BLE001 - an unreadable topology is ambiguous | |
| # Ambiguous: matching no artifact is safer than assuming one rank. | |
| return "unspecified" | |
| return "single" |
🤖 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 596 - 614, Update
_distributed_mode so the exception handler returns "single" only when
torch.distributed is unavailable; return "unspecified" for failures from
availability, initialization, or world-size checks, preserving the documented
fail-closed behavior for ambiguous multi-rank state.
Comment on lines
+755
to
+763
| def subgraph_signature_keys( | ||
| region: Mapping[str, Any], | ||
| manifest: ArtifactManifest, | ||
| ) -> tuple[tuple[tuple[Any, ...], ...], tuple[tuple[Any, ...], ...]]: | ||
| """Return live boundary input/output signatures for compatibility checks.""" | ||
| return ( | ||
| _meta_keys(region, manifest.boundary_refs), | ||
| _meta_keys(region, manifest.output_node_ids), | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
One malformed subgraph manifest can suppress every other candidate.
subgraph_signature_keys raises SubgraphRewriteError when the manifest's boundary_refs or output_node_ids do not resolve to IR metadata. GraphDispatchSession._resolve calls it at fastvideo/optimization/dispatch.py lines 442-445 inside the for manifest in candidates loop, and that loop has no per-manifest guard. The exception propagates to _decide, which records decision_failed:SubgraphRewriteError and returns a native fallback.
The remaining candidates for that signature are never evaluated, so one bad subgraph bundle disables a valid module bundle for the same scope. fastvideo/optimization/artifact.py lines 820-823 states the opposite policy: a failing bundle must never disable the ones that passed.
Wrap the per-manifest evaluation in dispatch so the failure becomes a rejection reason for that artifact only.
🐛 Proposed fix in `fastvideo/optimization/dispatch.py`
for manifest in candidates:
+ try:
if manifest.target_kind == "subgraph":
if export_region is None:
rejections.append(f"{manifest.artifact_id}:export_capture_missing")
continue
candidate_input_keys, candidate_output_keys = subgraph_signature_keys(
export_region,
manifest,
)
fingerprint = str(export_region.get("fingerprint", ""))
else:
...
+ except Exception as exc: # noqa: BLE001 - one bad bundle must not hide the rest
+ rejections.append(
+ f"{manifest.artifact_id}:signature_unavailable:{type(exc).__name__}")
+ continue🤖 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 755 - 763, Update the
per-manifest evaluation loop in GraphDispatchSession._resolve to catch
SubgraphRewriteError from subgraph_signature_keys and reject only the current
artifact with an appropriate failure reason. Continue evaluating remaining
candidates, preserving the artifact isolation policy described by the existing
artifact selection flow.
Comment on lines
240
to
+243
| self._trace_mgr = attach_activation_trace(self.modules.get("transformer")) | ||
| # Generic graph dispatch. Returns None unless a trusted artifact | ||
| # directory is configured, in which case nothing is wrapped at all. | ||
| self._dispatch_session = attach_graph_dispatch(self.modules) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect how GraphDispatchSession discovers repeated block stacks,
# to check compatibility with torch.compile-wrapped modules.
ast-grep run --pattern 'class GraphDispatchSession {
$$$
}' --lang python fastvideo/optimization/dispatch.py
rg -n -B3 -A15 'def attach_modules' fastvideo/optimization/dispatch.py
rg -n 'OptimizedModule|_orig_mod|torch.compile' fastvideo/optimization/dispatch.py fastvideo/optimization/fx_capture.pyRepository: aryan5v/FastVideo
Length of output: 1519
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '--- files ---\n'
git ls-files | rg '(^|/)composed_pipeline_base.py$|fastvideo/optimization/(dispatch|fx_capture).py$' || true
printf '\n--- composed pipeline relevant lines ---\n'
sed -n '1,80p;140,250p;528,540p' fastvideo/pipelines/composed_pipeline_base.py
printf '\n--- dispatch attach logic ---\n'
sed -n '1,260p' fastvideo/optimization/dispatch.py
printf '\n--- fx capture relevant classes/functions ---\n'
sed -n '430,590p' fastvideo/optimization/fx_capture.py
printf '\n--- searches ---\n'
rg -n "class GraphDispatchSession|OptimizedModule|_orig_mod|torch.compile|register_forward_pre_hook|register_forward_hook|default_capture_targets|register_module|submodules" fastvideo/optimization/dispatch.py fastvideo/optimization/fx_capture.py fastvideo/pipelines/composed_pipeline_base.py
printf '\n--- deterministic AST/text check for compiled module wrapping and attach timing ---\n'
python3 - <<'PY'
from pathlib import Path
import ast
path_base = Path("fastvideo/pipelines/composed_pipeline_base.py")
path_dispatch = Path("fastvideo/optimization/dispatch.py")
tree_base = ast.parse(path_base.read_text(), filename=str(path_base))
tree_dispatch = ast.parse(path_dispatch.read_text(), filename=str(path_dispatch))
# Find the init block names and relevant calls
for node in ast.walk(tree_base):
if isinstance(node, ast.FunctionDef) and node.name == "__init__":
print("__init__ line", node.lineno)
for call in ast.walk(node):
if isinstance(call, ast.Call):
fn = ast.unparse(call.func)
if fn in {"torch.compile", "attach_graph_dispatch", "attach_activation_trace"}:
print(call.func.lineno, "call:", fn)
if isinstance(call, ast.Assign):
if call.targets and hasattr(call.targets[0], "id") and call.targets[0].id == "_dispatch_session":
print("_dispatch_session assignment line", call.lineno, ast.unparse(call.value))
for node in ast.walk(tree_dispatch):
if isinstance(node, ast.FunctionDef) and node.name == "attach":
print("dispatch attach line", node.lineno)
src = ast.unparse(node)
for needle in ["_orig_mod", "compile", "compile", "nn.ModuleList", "register_forward_pre_hook", "register_forward_hook"]:
print(" contains", needle, bool(needle in src))
PYRepository: aryan5v/FastVideo
Length of output: 246
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- files ---'
git ls-files | rg '(^|/)composed_pipeline_base.py$|fastvideo/optimization/(dispatch|fx_capture).py$' || true
printf '%s\n' ''
printf '%s\n' '--- composed pipeline relevant lines ---'
sed -n '1,80p;140,250p;528,540p' fastvideo/pipelines/composed_pipeline_base.py
printf '%s\n' ''
printf '%s\n' '--- dispatch attach logic ---'
sed -n '1,260p' fastvideo/optimization/dispatch.py
printf '%s\n' ''
printf '%s\n' '--- fx capture relevant classes/functions ---'
sed -n '430,590p' fastvideo/optimization/fx_capture.py
printf '%s\n' ''
printf '%s\n' '--- searches ---'
rg -n "class GraphDispatchSession|OptimizedModule|_orig_mod|torch\.compile|register_forward_pre_hook|register_forward_hook|default_capture_targets|register_module|submodules" fastvideo/optimization/dispatch.py fastvideo/optimization/fx_capture.py fastvideo/pipelines/composed_pipeline_base.py
printf '%s\n' ''
printf '%s\n' '--- deterministic AST/text check ---'
python3 - <<'PY'
from pathlib import Path
import ast
path_base = Path("fastvideo/pipelines/composed_pipeline_base.py")
path_dispatch = Path("fastvideo/optimization/dispatch.py")
tree_base = ast.parse(path_base.read_text(), filename=str(path_base))
tree_dispatch = ast.parse(path_dispatch.read_text(), filename=str(path_dispatch))
for node in ast.walk(tree_base):
if isinstance(node, ast.FunctionDef) and node.name == "__init__":
print("__init__ line", node.lineno)
for call in ast.walk(node):
if isinstance(call, ast.Call):
fn = ast.unparse(call.func)
if fn in {"torch.compile", "attach_graph_dispatch", "attach_activation_trace"}:
print(call.func.lineno, "call:", fn)
if isinstance(call, ast.Assign):
if call.targets and hasattr(call.targets[0], "id") and call.targets[0].id == "_dispatch_session":
print("_dispatch_session assignment line", call.lineno, ast.unparse(call.value))
for node in ast.walk(tree_dispatch):
if isinstance(node, ast.FunctionDef) and node.name == "attach":
print("dispatch attach line", node.lineno)
src = ast.unparse(node)
for needle in ["_orig_mod", "compile", "compile", "nn.ModuleList", "register_forward_pre_hook", "register_forward_hook"]:
print(" contains", needle, bool(needle in src))
PYRepository: aryan5v/FastVideo
Length of output: 28352
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- composed_pipeline_base compile helpers ---'
sed -n '80,165p' fastvideo/pipelines/composed_pipeline_base.py
printf '%s\n' ''
printf '%s\n' '--- default_capture_targets implementation ---'
sed -n '860,970p' fastvideo/optimization/fx_capture.py
printf '%s\n' ''
printf '%s\n' '--- GraphDispatchSession call path ---'
sed -n '970,1190p' fastvideo/optimization/dispatch.py
printf '%s\n' ''
printf '%s\n' '--- dispatch attach modules call sites ---'
rg -n -B4 -A4 "attach_graph_dispatch|detach_graph_dispatch|GraphDispatchSession|attach_modules" fastvideo -g '*.py'Repository: aryan5v/FastVideo
Length of output: 47109
Document or guard the torch.compile + dispatch compatibility.
attach_graph_dispatch() runs after the compile path and wraps repeated nn.ModuleList children by assigning module.forward. The dispatch implementation does not unwrap torch.compile, and PyTorch’s Dynamo forwarding hooks can trigger graph breaks, so an artifact directory plus enable_torch_compile may skip dispatch or avoid compiling parts that hooks break. Treat this configuration as mutually exclusive or add explicit support/coverage before relying on both.
🤖 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/pipelines/composed_pipeline_base.py` around lines 240 - 243, Update
the initialization flow around _dispatch_session and attach_graph_dispatch to
prevent enabling graph dispatch together with enable_torch_compile unless
explicit compatibility support is implemented and covered. Prefer validating
this configuration early and raising a clear error, or otherwise ensure dispatch
does not wrap compiled modules while preserving existing behavior for either
feature alone.
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
Publishes the fork-only FastVideo side of the MotionKernel V1 technical preview:
Validation
This PR targets
aryan5v/FastVideoonly.Summary by CodeRabbit