[WIP]Feature/warp tracing - #7
Open
frlai wants to merge 61 commits into
Open
Conversation
frlai
force-pushed
the
feature/warp_tracing
branch
from
June 26, 2026 17:56
67d95f1 to
b83ce24
Compare
frlai
force-pushed
the
feature/warp_tracing
branch
from
July 8, 2026 15:35
ffbc2a2 to
17f8bc8
Compare
frlai
force-pushed
the
feature/warp_tracing
branch
from
July 20, 2026 23:58
5537959 to
cb014d7
Compare
…alls to warp functions
…matic i/o detection for node based operators
…utomatic used output detection and masking
…be registered as part of onnx and torchscript via the script path, export to onnx seems to work
…doesn't support jit-trace, onnx-torchscript or jit-script.
…aph capture into the onnx file
updated input signature to use Tensor[] inputs, str runtime_metadata, Tensor bundle. all metadata will be stored inside runtime_metadata. this includes the mask and i/o configs todo: need to prune runtime_metadata for uneeded fields.
Squashed from two_pass_apic_capture: - implements two pass tracing - force ptx sm 75 - updated to use single pass, refactored code ownership
…on and broadened support for automatic reentry validation
…ion inside the warp op
…stricter hierarchy, fixed bug with seg fault on reentry because loss of buffer creation ownership still need to fix bug with incorrect number of i/o at the boundry
…orch reuse cuda memory so torch memory creation functions sill trip the cupti oracle
…dry run, derrived values when node is not tracing are no longer promoted
… we now never demote and update a port in in the traced_data instead to track tagged tensors
…s itself across node boundries and survives most copy attempts
A Warp view of a Torch buffer created before tracing began had no conversion to intercept, so the runner writing it had no consumers and pruning deleted it while eager returned the mutated buffer. - input_tensors indexes each declared input by layout and adopts an existing identical-layout alias, so two names over one buffer collapse to one port instead of founding two roots - re-declaring a live carrier of this node rebinds that object in place rather than wrapping it, which keeps a promoted wp.array on the current pass's root across the Warp discovery and capture passes - the declared-input/output alias check runs after build_graph_module and only faults on an input that survived pruning, since trimming an unused declaration is ordinary - WarpSegment.knows_array replaces two disagreeing rules for whether a segment already holds an array
An adapter holding a persistent Torch/Warp pair reads its own attribute after a kernel writes the buffer, with no conversion for the tracer to intercept. Wrapping left that object plain, so everything downstream of the segment was untraced and output_tensors rejected it. TracedTensor already subclasses torch.Tensor and _promote_plain_tensor already served the copy_ and __setitem__ destination paths, so declaration now routes there and such code traces unchanged. NumPy stays wrapped, since a raw ndarray cannot be class-swapped. - restore __hash__ on TracedTensor: the elementwise __eq__ nulls the inherited hash, which breaks every identity set a carrier reaches, nn.Module member traversal included once a registered buffer is promoted - exempt a placeholder-rooted carrier from the cross-node "not a registered output" error, because one external tensor declared on two nodes now arrives at the second still carrying the first's placeholder - the dry-run carriers test asserted the wrapping contract and now asserts the in-place one
…ctor reframes patching as a simulated function dispatcher and the uses similar traced_data semantics as torch and numpy
Allow traceable function replacements while centralizing atomic patch installation and restoration across all backends.
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
Adds experimental NVIDIA Warp tracing/export support to LEAPP, including explicit
annotate.warp_op(...)regions, automatic Warp segment detection, APIC capture replay, and ONNX/ExportedProgram export support through aleapp::warp_runnercustom operator.Details
leapp::warp_runnerand lowers it tocom.nvidia.warp::WrpRunnerfor ONNX exports.InferenceManagersmoke-test guidance.Notes
onnxruntime-gpu, CUDA EP availability, andLEAPP_WARP_ONNX_CUSTOM_OP_LIBRARYpointing at a custom op built against the same ONNX Runtime version in the environment.