Replay autocast state during deferred graph capture

[aryan5v]

Summary

• capture CPU/CUDA autocast enablement, dtype, and cache mode with each observed shape variant
• replay that state when deferred export/Dynamo capture runs after the profiled forward
• clear the bounded runtime metadata with all other live capture references
• add a CPU regression proving exported matmul metadata preserves the observed autocast dtype

GPU evidence

LTX diagnostic job 981 isolated the failure: discovery packaged the second subgraph boundary as bfloat16 while live FastVideo export observed float32. The existing artifact failed closed with input_signature_mismatch and zero candidate calls. This change fixes the generic capture-context omission; it contains no LTX/model-specific branch.

Validation

• 97 passed: fastvideo/tests/optimization + tests/local_tests/optimizations
• Ruff clean on changed Python files
• git diff --check clean

Greptile Summary

This PR fixes a dtype mismatch during deferred graph capture by snapshotting the ambient autocast state at profiling time and replaying it when finalize() triggers the actual export or Dynamo trace. Without the fix, mm and similar numerically sensitive operations acquired the ambient dtype at export time (outside the user's torch.autocast block) rather than the dtype they had during the real forward.

• Adds _AutocastState dataclass and _observe_autocast_state() to snapshot per-device enabled/dtype/cache state; stores the snapshot in _ShapeVariant.observed_autocast on first shape-key encounter.
• Adds _capture_autocast() context manager (used alongside _capture_forward_context in _trace) to replay the observed state via ExitStack; both cleanup paths in finalize() now clear observed_autocast alongside the other live references.
• Adds a regression test that asserts CPU bfloat16 autocast is faithfully reproduced in the exported IR dtype for aten.mm.default nodes.

Confidence Score: 5/5

Safe to merge. The autocast-replay logic is tightly scoped, correctly cleaned up in both the success and error paths, and directly validated by a new regression test.

The change is minimal and self-contained: it snapshots thread-local autocast state during the profiling forward and replays it via a properly scoped ExitStack during deferred export. Both cleanup paths in finalize already clear the field. The new test exercises the exact failure scenario described in the PR and confirms the exported IR dtype matches the observed state. No correctness issues were found.

Files Needing Attention: fastvideo/tests/optimization/test_fx_capture.py — the parametrized stable-metadata test is not updated to verify observed_autocast cleanup for symbolic and dynamo modes.

Important Files Changed

Filename	Overview
fastvideo/optimization/fx_capture.py	Introduces _AutocastState / _observe_autocast_state / _capture_autocast; wires observed_autocast into _ShapeVariant and both finalize cleanup paths; correct use of ExitStack for multi-device replay.
fastvideo/tests/optimization/test_fx_capture.py	Adds _AutocastBlock, test_deferred_export_replays_observed_autocast_dtype, and updates the invalid-tracer cleanup assertion; test correctly validates bfloat16 dtype in exported IR and checks observed_autocast is cleared post-finalize.

Sequence Diagram

sequenceDiagram
    participant User
    participant FXCaptureSession
    participant _ShapeVariant
    participant _observe_autocast_state
    participant _capture_autocast

    User->>FXCaptureSession: "model(inputs) inside torch.autocast(cpu, bf16)"
    FXCaptureSession->>_observe_autocast_state: "snapshot autocast state"
    _observe_autocast_state-->>FXCaptureSession: "_AutocastState(cpu enabled bf16, cuda disabled fp16)"
    FXCaptureSession->>_ShapeVariant: "store observed_autocast"

    Note over User,FXCaptureSession: User exits autocast context

    User->>FXCaptureSession: "session.finalize()"
    FXCaptureSession->>_capture_autocast: "enter with observed_autocast"
    _capture_autocast->>_capture_autocast: "ExitStack: enter autocast cpu bf16"
    _capture_autocast->>_capture_autocast: "ExitStack: enter autocast cuda disabled"
    FXCaptureSession->>FXCaptureSession: "torch.export / Dynamo trace sees bf16"
    _capture_autocast->>_capture_autocast: "ExitStack exits, autocast restored"
    FXCaptureSession->>_ShapeVariant: "observed_autocast = None"
    FXCaptureSession-->>User: "payload with bf16 IR nodes"

Loading

_{Reviews (2): Last reviewed commit: "[test]: document canonical IR dtype form..." | Re-trigger Greptile}

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 36805626-7287-4977-bdfd-399d8a301094

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}