Commit 29abd38
feat(annotation): add TTA annotation layer (export_as, lower_as, custom_plugin)
Adds torch_tensorrt.annotation (aliased as tta) — a zero-overhead annotation
layer that lets users tag regions of a PyTorch model for custom TensorRT
lowering, without modifying core torch_tensorrt internals beyond a small set
of generic extension hooks.
Core hooks added to torch_tensorrt (outside annotation/):
- _compiler.py: generic extension hook registries (register_compile_pass,
register_preserved_ep_attr, register_export_context, register_post_trace_hook)
+ EP attribute preservation across run_decompositions() + compile pass loop
- _compile.py: post-trace hook loop
- _settings.py: profiling_verbosity field in CompilationSettings
- _tracer.py: ExitStack wrapping for registered export context factories
- _ConversionContext.py: current_node field for converter access
- _TRTInterpreter.py: PREFER_AOT_PYTHON_PLUGINS flag, profiling_verbosity
routing, current_node tracking, layer metadata stamping
Annotation module (py/torch_tensorrt/annotation/):
- export_as: context manager to tag regions during tracing
- lower_as: context manager to specify custom TRT lowering (builtin/plugin/kernel)
- custom_plugin: Triton/CuTile/CuTeDSL AOT QDP plugin descriptors
- IR layer: region discovery, boundary validation, region views
- Full test suite (599 unit/integration tests + 25 pre-Blackwell tests)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>1 parent 2e26bfa commit 29abd38
79 files changed
Lines changed: 28236 additions & 11 deletions
File tree
- py/torch_tensorrt
- annotation
- _compile
- _custom_plugin
- _aot
- ir
- dynamo
- conversion
- tests/py/annotation
- integration
- unit
Some content is hidden
Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
320 | 320 | | |
321 | 321 | | |
322 | 322 | | |
| 323 | + | |
| 324 | + | |
| 325 | + | |
| 326 | + | |
| 327 | + | |
| 328 | + | |
323 | 329 | | |
324 | 330 | | |
325 | 331 | | |
| |||
Large diffs are not rendered by default.
0 commit comments