Skip to content

refactor: store FX proxies directly on traced values - #45

Merged
frlai merged 4 commits into
feature/view-mutation-functionalizationfrom
refactor/direct-proxy-storage
Oct 29, 2026
Merged

frlai merged 4 commits into
feature/view-mutation-functionalizationfrom
refactor/direct-proxy-storage

Conversation

@frlai

@frlai frlai commented Oct 22, 2026 •

Copy link
Copy Markdown
Collaborator

PR #33 functionalizes recorded view mutations through PyTorch, but traced values still carry shared ProxyView cells. Store FX proxies directly on carriers and remove ProxyView, nested cell traversal, shared/new-view binding helpers, and the dual proxy= / view= factory interface.

Numerical writes retain stable FX tensor handles and record the existing mutation operations for functionalization. Input redeclarations use a small per-node binding table; retained aliases attach to declarations after a capture reset. Independent conversion copies, per-carrier ports, and weak whole-allocation Warp metadata retain their existing responsibilities.

New Warp outputs use a provisional FX Proxy shared directly by their actual aliases. Segment close resolves that proxy's producer once, without recording a tensor allocation or copy. Weak storage metadata captured before close is published on the eventual producer. Existing-storage writes still go through functionalization; finished nodes continue eager execution without creating pending producers.

Stacked on #33, targeting feature/view-mutation-functionalization at 914d27a (including PR #44). Implementation and experiments use the separate leapp-direct-proxy worktree.

Validation and branch update:

  • The resolved integration with target 914d27a passed 1,062 tests / 1,800 subtests, including examples and visualization, on Python 3.12 / Torch 2.14.0 / released Warp 1.18 / RTX 4090 with matching native ONNX/PT2 adapters.
  • After that full run, removed two migration-only unit checks: test_retained_alias_reattaches_after_declaration_reentry and test_new_warp_output_publishes_without_graph_allocation_or_copy. Production code is identical to the full-suite-tested integration. All pre-existing tests from the target branch remain.
  • Final branch checks: 64 tests / 6 subtests passed on the current CUDA stack; 63 focused view-mutation/Warp-boundary tests passed on minimum Torch 2.7.0 CPU. Ruff 0.17.1 formatting and diff checks pass.
  • Retained the ONNX/PT2 saved-artifact Warp publication regression, including conversion before/after explicit segment close and automatic segments, distinct capture samples, fresh-input replay, unchanged caller inputs, and rejection of synthetic allocation/copy operations. It guards the measured replay regression in the first prototype.
  • Earlier qualification at 1bf9364: 1,052 tests / 1,800 subtests, plus 26 additional ONNX/PT2 Warp allocation cases / 104 independently varied replay trials, weak-storage lifetime/cleanup and missing-covering-root guards, and five minimum-Torch allocation-metadata checks. These additional matrices were not rerun for the target-branch integration; historical full-version and TensorRT matrices from fix: preserve mutations through tensor views #33 are not claimed for this revision.

For the 65,536-element new-Warp-output probe, the final candidate's normalized FX and saved PT2 graph text are identical to baseline, with six normalized nodes. Warm native PT2 replay measured 44.25 vs 44.57 µs, with the same 524,288-byte additional Torch allocator peak. Two-pass capture medians were 144.20 vs 146.19 ms. These are local measurements, not general performance guarantees; publication introduces no additional replay operations.

Against the original d46edb7 baseline, the direct-proxy refactor removed 160 production physical lines, 31 Python AST statements, and six conditional/while branches. No replacement proxy-cell class or general view tree is introduced.

Published revision: 158979717e3f6f529885055670ef297ced71f4fe; exact source tree: 57025d9f33e0c9a1cbc5a97081d457d126cd4e4c. The branch includes the reviewed PR #44 import resolution and the agreed test cleanup.

frlai added 2 commits October 22, 2026 02:31
Signed-off-by: Frank Lai <frlai@nvidia.com>
Signed-off-by: Frank Lai <frlai@nvidia.com>
frlai added 2 commits October 29, 2026 14:53
Signed-off-by: Frank Lai <frlai@nvidia.com>
Resolve the tracing import conflict using direct proxy storage and the
ATen rewrite module from the target branch.

Signed-off-by: Frank Lai <frlai@nvidia.com>
@frlai
frlai merged commit e7a5476 into feature/view-mutation-functionalization Oct 29, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant