tensoracle 1.0.1__tar.gz → 1.2.0__tar.gz

{tensoracle-1.0.1 → tensoracle-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: tensoracle
-Version: 1.0.1
+Version: 1.2.0
 Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
 Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
 Project-URL: repository, https://github.com/sync-dev-org/tensoracle
@@ -24,8 +24,6 @@ Requires-Dist: tensorrt-cu13<12,>=11.0
 Provides-Extra: convert
 Requires-Dist: torch>=2.6; extra == 'convert'
 Provides-Extra: producer
-Requires-Dist: nvidia-cuda-cuobjdump<13.4,>=13.3.73; extra == 'producer'
-Requires-Dist: nvidia-cuda-nvcc-cu12==12.9.86; extra == 'producer'
 Requires-Dist: triton==3.8.0; extra == 'producer'
 Description-Content-Type: text/markdown
 
@@ -53,8 +51,8 @@ Choose the path that matches your work:
 | `tensorrt-cu13` 11.x | Open and run engines | Build and open engines | Open plugin engines |
 | NVIDIA GPU and driver | Open and run engines | Build and run engines | Run and test engines; cubin production itself can run without a GPU |
 | CuPy or `cuda-python` | The application chooses its buffer and stream library | The application chooses its buffer and stream library | `cuda-python` is used by the example runtime and tests, not by `tensoracle` |
-| Triton 3.8.0, `ptxas` 12.9.86, `cuobjdump` 13.3 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
-| CUDA toolkit | Not needed | Not needed | Native plugin build |
+| Triton 3.8.0 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
+| CUDA toolkit | Not needed | Not needed | Cubin production (`ptxas >=13.2,<14` and `cuobjdump`) and native plugin build |
 | TensorRT 11 SDK | Not needed | Not needed | Native plugin build |
 | CMake 3.20 or later and C++17 compiler | Not needed | Not needed | Native plugin build |
 
@@ -255,10 +253,75 @@ with handle.create_execution() as execution:
 
 The caller selects the CUDA device and context, orders inputs written on other streams before enqueue, and keeps buffers, streams, plugin preparation resources, and the execution unit alive until completion. Use one execution unit serially; do not call it concurrently from multiple threads. Confirm completion before changing shapes or streams, reusing buffers, or closing. `close()` and context exit release the TensorRT context before the handle, without waiting, and repeated close calls do nothing. Operations on a closed unit raise `RuntimeError` before touching TensorRT. The unit does not retain caller arrays or expose its TensorRT context.
 
-Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
+Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. Profile switching and caller-owned TensorRT workspace memory are outside this API.
 
 Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
 
+### Capture and replay with CUDA Graphs
+
+Tensoracle supports recording one `EngineExecution.enqueue` call with CUDA stream capture and replaying the resulting graph. The application owns graph creation, recapture, and destruction; tensoracle provides no capture or replay operations.
+
+Use a dedicated execution unit (`handle.create_execution()`) for each graph. After setting input shapes, enqueue once and wait for completion before starting capture, so TensorRT finishes deferred initialization. Capture and replay on a nondefault stream; the example uses a nonblocking stream. Keep synchronous CUDA calls and work on the legacy default stream out of capture.
+
+The graph fixes the input shapes, tensor addresses, and execution unit. Keep the execution unit, its handle, buffers, and plugin preparation resources alive until the graph is destroyed and outstanding work has finished. To use different shapes or addresses, capture a new graph for that combination or select another graph with its own execution unit. Input contents may change at the same addresses between completed replays.
+
+For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
+
+```python
+from contextlib import ExitStack
+
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import runtime as cudart
+
+with handle.create_execution() as execution:
+    execution.set_input_shapes({"features": (1, 4, 2)})
+    # Omit set_input_shapes for fixed inputs and single-point profiles.
+    features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
+    output = np.empty(
+        execution.get_tensor_shape("result"),
+        dtype=trt.nptype(execution.get_tensor_dtype("result")),
+    )
+    with ExitStack() as resources:
+        stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
+        resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
+        source = cuda(cudart.cudaMalloc(features.nbytes))
+        resources.callback(lambda: cuda(cudart.cudaFree(source)))
+        target = cuda(cudart.cudaMalloc(output.nbytes))
+        resources.callback(lambda: cuda(cudart.cudaFree(target)))
+        addresses = {"features": int(source), "result": int(target)}
+        try:
+            cuda(cudart.cudaMemcpyAsync(
+                source, features.ctypes.data, features.nbytes,
+                cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
+            ))
+            execution.enqueue(addresses, stream=int(stream))
+            cuda(cudart.cudaStreamSynchronize(stream))  # Warm up before capture.
+
+            cuda(cudart.cudaStreamBeginCapture(
+                stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
+            ))
+            try:
+                execution.enqueue(addresses, stream=int(stream))  # Exactly one call.
+            finally:
+                graph = cuda(cudart.cudaStreamEndCapture(stream))
+            resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
+            instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
+            resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
+
+            for _ in range(3):
+                cuda(cudart.cudaGraphLaunch(instance, stream))
+            cuda(cudart.cudaMemcpyAsync(
+                output.ctypes.data, target, output.nbytes,
+                cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
+            ))
+        finally:
+            # Wait before ExitStack destroys the graph and releases its buffers.
+            cuda(cudart.cudaStreamSynchronize(stream))
+```
+
+Engines with loops or conditionals, or layers or plugins that synchronize with the host during inference, cannot be captured. CUDA reports capture failures through `cudaErrorStreamCapture*` errors. Always end an attempted capture and release any graph resources before falling back to ordinary `enqueue`. See [TensorRT's CUDA Graph limitations](https://docs.nvidia.com/deeplearning/tensorrt/latest/performance/optimization.html#inference-with-cuda-graphs).
+
 ## Build a cubin plugin
 
 Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
@@ -281,7 +344,7 @@ python/
 
 Keep C++ sources in `cpp/`, with a CMake project that builds and runs its tests through CTest without Python. Keep the Python project, its src-layout package, and its uv environment in `python/`; run the project's uv commands from that directory. Tensoracle itself follows this layout, and the installed C++ helper paths described below remain inside its Python package.
 
-`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0**, **ptxas 12.9.86**, and **cuobjdump 13.3**. It finds the NVIDIA wheel executables by default; explicit `ptxas` and `cuobjdump` paths can override them. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
+`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0** with **ptxas >=13.2,<14** and **cuobjdump** from the host CUDA toolkit. It uses the toolkit directory named by `CUDA_HOME`, or `/usr/local/cuda` when `CUDA_HOME` is unset; explicit `ptxas` and `cuobjdump` paths override these defaults independently. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
 
 The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
 

{tensoracle-1.0.1 → tensoracle-1.2.0}/README.md

@@ -22,8 +22,8 @@ Choose the path that matches your work:
 | `tensorrt-cu13` 11.x | Open and run engines | Build and open engines | Open plugin engines |
 | NVIDIA GPU and driver | Open and run engines | Build and run engines | Run and test engines; cubin production itself can run without a GPU |
 | CuPy or `cuda-python` | The application chooses its buffer and stream library | The application chooses its buffer and stream library | `cuda-python` is used by the example runtime and tests, not by `tensoracle` |
-| Triton 3.8.0, `ptxas` 12.9.86, `cuobjdump` 13.3 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
-| CUDA toolkit | Not needed | Not needed | Native plugin build |
+| Triton 3.8.0 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
+| CUDA toolkit | Not needed | Not needed | Cubin production (`ptxas >=13.2,<14` and `cuobjdump`) and native plugin build |
 | TensorRT 11 SDK | Not needed | Not needed | Native plugin build |
 | CMake 3.20 or later and C++17 compiler | Not needed | Not needed | Native plugin build |
 
@@ -224,10 +224,75 @@ with handle.create_execution() as execution:
 
 The caller selects the CUDA device and context, orders inputs written on other streams before enqueue, and keeps buffers, streams, plugin preparation resources, and the execution unit alive until completion. Use one execution unit serially; do not call it concurrently from multiple threads. Confirm completion before changing shapes or streams, reusing buffers, or closing. `close()` and context exit release the TensorRT context before the handle, without waiting, and repeated close calls do nothing. Operations on a closed unit raise `RuntimeError` before touching TensorRT. The unit does not retain caller arrays or expose its TensorRT context.
 
-Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
+Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. Profile switching and caller-owned TensorRT workspace memory are outside this API.
 
 Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
 
+### Capture and replay with CUDA Graphs
+
+Tensoracle supports recording one `EngineExecution.enqueue` call with CUDA stream capture and replaying the resulting graph. The application owns graph creation, recapture, and destruction; tensoracle provides no capture or replay operations.
+
+Use a dedicated execution unit (`handle.create_execution()`) for each graph. After setting input shapes, enqueue once and wait for completion before starting capture, so TensorRT finishes deferred initialization. Capture and replay on a nondefault stream; the example uses a nonblocking stream. Keep synchronous CUDA calls and work on the legacy default stream out of capture.
+
+The graph fixes the input shapes, tensor addresses, and execution unit. Keep the execution unit, its handle, buffers, and plugin preparation resources alive until the graph is destroyed and outstanding work has finished. To use different shapes or addresses, capture a new graph for that combination or select another graph with its own execution unit. Input contents may change at the same addresses between completed replays.
+
+For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
+
+```python
+from contextlib import ExitStack
+
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import runtime as cudart
+
+with handle.create_execution() as execution:
+    execution.set_input_shapes({"features": (1, 4, 2)})
+    # Omit set_input_shapes for fixed inputs and single-point profiles.
+    features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
+    output = np.empty(
+        execution.get_tensor_shape("result"),
+        dtype=trt.nptype(execution.get_tensor_dtype("result")),
+    )
+    with ExitStack() as resources:
+        stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
+        resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
+        source = cuda(cudart.cudaMalloc(features.nbytes))
+        resources.callback(lambda: cuda(cudart.cudaFree(source)))
+        target = cuda(cudart.cudaMalloc(output.nbytes))
+        resources.callback(lambda: cuda(cudart.cudaFree(target)))
+        addresses = {"features": int(source), "result": int(target)}
+        try:
+            cuda(cudart.cudaMemcpyAsync(
+                source, features.ctypes.data, features.nbytes,
+                cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
+            ))
+            execution.enqueue(addresses, stream=int(stream))
+            cuda(cudart.cudaStreamSynchronize(stream))  # Warm up before capture.
+
+            cuda(cudart.cudaStreamBeginCapture(
+                stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
+            ))
+            try:
+                execution.enqueue(addresses, stream=int(stream))  # Exactly one call.
+            finally:
+                graph = cuda(cudart.cudaStreamEndCapture(stream))
+            resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
+            instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
+            resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
+
+            for _ in range(3):
+                cuda(cudart.cudaGraphLaunch(instance, stream))
+            cuda(cudart.cudaMemcpyAsync(
+                output.ctypes.data, target, output.nbytes,
+                cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
+            ))
+        finally:
+            # Wait before ExitStack destroys the graph and releases its buffers.
+            cuda(cudart.cudaStreamSynchronize(stream))
+```
+
+Engines with loops or conditionals, or layers or plugins that synchronize with the host during inference, cannot be captured. CUDA reports capture failures through `cudaErrorStreamCapture*` errors. Always end an attempted capture and release any graph resources before falling back to ordinary `enqueue`. See [TensorRT's CUDA Graph limitations](https://docs.nvidia.com/deeplearning/tensorrt/latest/performance/optimization.html#inference-with-cuda-graphs).
+
 ## Build a cubin plugin
 
 Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
@@ -250,7 +315,7 @@ python/
 
 Keep C++ sources in `cpp/`, with a CMake project that builds and runs its tests through CTest without Python. Keep the Python project, its src-layout package, and its uv environment in `python/`; run the project's uv commands from that directory. Tensoracle itself follows this layout, and the installed C++ helper paths described below remain inside its Python package.
 
-`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0**, **ptxas 12.9.86**, and **cuobjdump 13.3**. It finds the NVIDIA wheel executables by default; explicit `ptxas` and `cuobjdump` paths can override them. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
+`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0** with **ptxas >=13.2,<14** and **cuobjdump** from the host CUDA toolkit. It uses the toolkit directory named by `CUDA_HOME`, or `/usr/local/cuda` when `CUDA_HOME` is unset; explicit `ptxas` and `cuobjdump` paths override these defaults independently. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
 
 The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
 

tensoracle-1.2.0/docs/api/producer.md

@@ -0,0 +1,7 @@
+# Cubin producer API
+
+The following names are available from `tensoracle.producer`. Install `tensoracle[producer]` for Triton 3.8.0 and provide a host CUDA toolkit with `ptxas >=13.2,<14` and `cuobjdump`. `compile_bundle` uses the toolkit directory named by `CUDA_HOME`, or `/usr/local/cuda` when `CUDA_HOME` is unset. Explicit `ptxas` and `cuobjdump` paths override the defaults independently. A missing tool reports the searched path and how to set `CUDA_HOME` or pass an explicit path. Cubin production does not need a GPU. Normal imports, engine building, and engine reuse do not import the producer or Triton.
+
+::: tensoracle.producer.KernelSpec
+
+::: tensoracle.producer.compile_bundle

{tensoracle-1.0.1 → tensoracle-1.2.0}/docs/index.md

@@ -26,8 +26,8 @@ Follow [Build a C++ and Python repository](cpp-python-repositories.md) for sourc
 | `tensorrt-cu13` 11.x | Open and run engines | Build and open engines | Open plugin engines |
 | NVIDIA GPU and driver | Open and run engines | Build and run engines | Run and test engines; cubin production itself can run without a GPU |
 | CuPy or `cuda-python` | The application chooses its buffer and stream library | The application chooses its buffer and stream library | `cuda-python` is used by the example runtime and tests, not by `tensoracle` |
-| Triton 3.8.0, `ptxas` 12.9.86, `cuobjdump` 13.3 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
-| CUDA toolkit | Not needed | Not needed | Native plugin build |
+| Triton 3.8.0 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
+| CUDA toolkit | Not needed | Not needed | Cubin production (`ptxas >=13.2,<14` and `cuobjdump`) and native plugin build |
 | TensorRT 11 SDK | Not needed | Not needed | Native plugin build |
 | CMake 3.20 or later and C++17 compiler | Not needed | Not needed | Native plugin build |
 
@@ -155,12 +155,77 @@ with handle.create_execution() as execution:
 
 The caller selects the CUDA device and context, orders inputs written on other streams before enqueue, and keeps buffers, streams, plugin preparation resources, and the execution unit alive until completion. Use one execution unit serially; do not call it concurrently from multiple threads. Confirm completion before changing shapes or streams, reusing buffers, or closing. `close()` and context exit release the TensorRT context before the handle, without waiting, and repeated close calls do nothing. Operations on a closed unit raise `RuntimeError` before touching TensorRT. The unit does not retain caller arrays or expose its TensorRT context.
 
-Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
+Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. Profile switching and caller-owned TensorRT workspace memory are outside this API.
 
 Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
 
 `TensorSpec` accepts optional `location=trt.TensorLocation.DEVICE` and `format=trt.TensorFormat.LINEAR`. `open_engine` checks these only when declared, together with tensor names, modes, dtypes, build-time shapes, and profiles. An `EngineIOError` names the tensor, field, observed value, expected value, and saved file path. The execution unit does not accept a second declaration.
 
+#### Capture and replay with CUDA Graphs
+
+Tensoracle supports recording one `EngineExecution.enqueue` call with CUDA stream capture and replaying the resulting graph. The application owns graph creation, recapture, and destruction; tensoracle provides no capture or replay operations.
+
+Use a dedicated execution unit (`handle.create_execution()`) for each graph. After setting input shapes, enqueue once and wait for completion before starting capture, so TensorRT finishes deferred initialization. Capture and replay on a nondefault stream; the example uses a nonblocking stream. Keep synchronous CUDA calls and work on the legacy default stream out of capture.
+
+The graph fixes the input shapes, tensor addresses, and execution unit. Keep the execution unit, its handle, buffers, and plugin preparation resources alive until the graph is destroyed and outstanding work has finished. To use different shapes or addresses, capture a new graph for that combination or select another graph with its own execution unit. Input contents may change at the same addresses between completed replays.
+
+For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
+
+```python
+from contextlib import ExitStack
+
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import runtime as cudart
+
+with handle.create_execution() as execution:
+    execution.set_input_shapes({"features": (1, 4, 2)})
+    # Omit set_input_shapes for fixed inputs and single-point profiles.
+    features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
+    output = np.empty(
+        execution.get_tensor_shape("result"),
+        dtype=trt.nptype(execution.get_tensor_dtype("result")),
+    )
+    with ExitStack() as resources:
+        stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
+        resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
+        source = cuda(cudart.cudaMalloc(features.nbytes))
+        resources.callback(lambda: cuda(cudart.cudaFree(source)))
+        target = cuda(cudart.cudaMalloc(output.nbytes))
+        resources.callback(lambda: cuda(cudart.cudaFree(target)))
+        addresses = {"features": int(source), "result": int(target)}
+        try:
+            cuda(cudart.cudaMemcpyAsync(
+                source, features.ctypes.data, features.nbytes,
+                cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
+            ))
+            execution.enqueue(addresses, stream=int(stream))
+            cuda(cudart.cudaStreamSynchronize(stream))  # Warm up before capture.
+
+            cuda(cudart.cudaStreamBeginCapture(
+                stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
+            ))
+            try:
+                execution.enqueue(addresses, stream=int(stream))  # Exactly one call.
+            finally:
+                graph = cuda(cudart.cudaStreamEndCapture(stream))
+            resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
+            instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
+            resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
+
+            for _ in range(3):
+                cuda(cudart.cudaGraphLaunch(instance, stream))
+            cuda(cudart.cudaMemcpyAsync(
+                output.ctypes.data, target, output.nbytes,
+                cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
+            ))
+        finally:
+            # Wait before ExitStack destroys the graph and releases its buffers.
+            cuda(cudart.cudaStreamSynchronize(stream))
+```
+
+Engines with loops or conditionals, or layers or plugins that synchronize with the host during inference, cannot be captured. CUDA reports capture failures through `cudaErrorStreamCapture*` errors. Always end an attempted capture and release any graph resources before falling back to ordinary `enqueue`. See [TensorRT's CUDA Graph limitations](https://docs.nvidia.com/deeplearning/tensorrt/latest/performance/optimization.html#inference-with-cuda-graphs).
+
 ### Build a standard graph
 
 Define `graph(network)` with TensorRT's Python network API and load model weights from safetensors. Describe each graph input with `InputShape.fixed` or `InputShape.ranged`, and pass an `EngineBuild` declaration to `build_serialized_engine`. `EngineBuild` declares input shapes and builder settings; the graph function determines precision through the tensor types it places in the network. The supported precisions are fp32 (with TF32 disabled), fp16, bf16, int8, and fp8. The [bf16, int8, and fp8 guide](precision.md) shows tensor casts, explicit Q/DQ layers, and scales supplied through safetensors; fp8 targets sm89 and sm120. `open_engine` uses the same declaration plus a name, graph version, weight files, and `EngineIO` to save and reuse a checked engine. Use a different `name` or `graph_version` for engines with different tensor types, such as fp32 and fp16. The [graph example](https://github.com/sync-dev-org/tensoracle#define-and-build-a-standard-graph) shows the exact calls.
@@ -177,7 +242,7 @@ Install `tensoracle[producer]` when you need to compile cubins from your own ker
 pip install 'tensoracle[producer]'
 ```
 
-`tensoracle.producer.compile_bundle` uses the extra's Triton 3.8.0, `ptxas` 12.9.86, and `cuobjdump` 13.3 to produce cubins for sm86, sm89, and sm120. It finds the NVIDIA wheel executables by default and accepts explicit paths as overrides. Cubin production does not need a GPU or import torch; run it without importing torch first. Build the native plugin with a TensorRT 11 SDK, CUDA toolkit, CMake, and a C++17 compiler. Inspect the resulting bundle or wheel with `tensoracle.bundle`. The [cubin plugin example](https://github.com/sync-dev-org/tensoracle/tree/main/examples/cubin_plugin) covers rebuilding from a repository clone, packaging, and runtime use. Normal import, engine building, and engine reuse do not import the producer or Triton, and do not need the producer extra or native build tools.
+`tensoracle.producer.compile_bundle` uses the extra's Triton 3.8.0 with `ptxas >=13.2,<14` and `cuobjdump` from the host CUDA toolkit to produce cubins for sm86, sm89, and sm120. It uses the toolkit directory named by `CUDA_HOME`, or `/usr/local/cuda` when `CUDA_HOME` is unset, and accepts explicit paths as independent overrides. Cubin production does not need a GPU or import torch; run it without importing torch first. Build the native plugin with a TensorRT 11 SDK, CUDA toolkit, CMake, and a C++17 compiler. Inspect the resulting bundle or wheel with `tensoracle.bundle`. The [cubin plugin example](https://github.com/sync-dev-org/tensoracle/tree/main/examples/cubin_plugin) covers rebuilding from a repository clone, packaging, and runtime use. Normal import, engine building, and engine reuse do not import the producer or Triton, and do not need the producer extra or native build tools.
 
 Pass a `tensoracle::cubin::PluginIO` declaration as the fourth argument to `PluginBase(name, plugin_namespace, fields, io)`. The `inputs` vector determines the input count and declares each input's TensorRT `type` and `format` through `TensorSpec`; `output` declares the single output's type and format. `shape_input_index` selects the input whose shape the output copies. This example declares two fp16 / HWC8 inputs and an fp16 / HWC8 output with input 0's shape:
 

{tensoracle-1.0.1 → tensoracle-1.2.0}/examples/cubin_plugin/README.md

@@ -6,21 +6,21 @@ The first plugin computes an fp32 32 by 16 dot product and then scales it in a s
 
 ## Rebuild the cubins
 
-From the tensoracle repository root, run:
+Provide a CUDA toolkit with `ptxas >=13.2,<14` and `cuobjdump`. Set `CUDA_HOME` to its directory; when unset, the producer uses `/usr/local/cuda`. From the tensoracle repository root, run:
 
 ```bash
 uv run --project python --extra producer python -c 'from pathlib import Path; from examples.cubin_plugin.kernels import kernel_specs; from tensoracle.producer import compile_bundle; compile_bundle(kernel_specs(), Path("examples/cubin_plugin/bundle"))'
 ```
 
-This installs `tensoracle[producer]` with Triton 3.8.0, ptxas 12.9.86, and cuobjdump 13.3, and recreates `bundle/` without using a GPU. `compile_bundle` finds ptxas and cuobjdump in their installed NVIDIA wheels. The producer checks the cubin ELF SM, public and hidden argument layout, register count, and zero scratch sizes before writing the manifest. Run the tests below with the producer extra alone so torch is absent; `compile_bundle` rejects an already imported torch.
+This installs Triton 3.8.0 with `tensoracle[producer]` and recreates `bundle/` with the host toolkit tools without using a GPU. Explicit `ptxas` and `cuobjdump` paths passed to `compile_bundle` override the toolkit defaults independently. The producer checks the cubin ELF SM, public and hidden argument layout, register count, and zero scratch sizes before writing the manifest. Run the tests below with the producer extra alone so torch is absent; `compile_bundle` rejects an already imported torch.
 
-Check that the regenerated files match the committed example without using a GPU or setting producer environment variables:
+Check that the generated bundle has the expected files and passes inspection without using a GPU:
 
 ```bash
 uv run --project python --extra producer pytest -m "not gpu" python/tests/test_producer.py
 ```
 
-Without the producer extra, the compilation tests skip with an explanation of the missing tools.
+Without Triton or the toolkit tools, the compilation tests skip with an explanation of what is missing.
 
 ## Build and inspect the wheel