tensoracle 1.0.0__tar.gz → 1.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {tensoracle-1.0.0 → tensoracle-1.1.0}/PKG-INFO +70 -5
- {tensoracle-1.0.0 → tensoracle-1.1.0}/README.md +68 -3
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/api/index.md +2 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/index.md +66 -1
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/optimization.md +1 -1
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/README.md +1 -1
- {tensoracle-1.0.0 → tensoracle-1.1.0}/pyproject.toml +1 -1
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_version.py +2 -2
- tensoracle-1.1.0/tests/test_cuda_graphs.py +191 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_distribution.py +1 -1
- {tensoracle-1.0.0 → tensoracle-1.1.0}/.gitignore +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/LICENSE +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/build_backend.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/CMakeLists.txt +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/src/cubin_plugin.cpp +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/tests/core.cpp +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/api/bundle.md +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/api/producer.md +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/cpp-python-repositories.md +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/precision.md +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/tensorrt-sdk.md +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/CMakeLists.txt +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.json +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/kernels.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/plugin.cpp +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/pyproject.toml +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/mkdocs.yml +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/__init__.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_build.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_cache.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_comparison.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_execution.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_memory.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_plugins.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/bundle.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/convert.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/producer.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/build_memory_runtime.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/cubin_io/check.cpp +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/cubin_io/kernels.cu +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/cubin_io/plugin.cpp +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/example_runtime.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/native/CMakeLists.txt +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/native/export.h +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/native_build.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/plugin_io_runtime.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/precision_runtime.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_build.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_build_memory.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_bundle.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_cache.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_cmake_plugin.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_comparison.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_convert.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_cpp_core.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_engine_io.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_example_plugin.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_execution.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_gpu_cache.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_plugin_io.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_plugins.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_precision.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_producer.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_shape_tensor_outputs.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_weight_hash.py +0 -0
- {tensoracle-1.0.0 → tensoracle-1.1.0}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: tensoracle
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.1.0
|
|
4
4
|
Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
|
|
5
5
|
Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
|
|
6
6
|
Project-URL: repository, https://github.com/sync-dev-org/tensoracle
|
|
@@ -8,7 +8,7 @@ Project-URL: issues, https://github.com/sync-dev-org/tensoracle/issues
|
|
|
8
8
|
Author-email: minamik <mia@sync.dev>
|
|
9
9
|
License-Expression: MIT
|
|
10
10
|
License-File: LICENSE
|
|
11
|
-
Classifier: Development Status ::
|
|
11
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
12
12
|
Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
|
|
13
13
|
Classifier: Intended Audience :: Developers
|
|
14
14
|
Classifier: License :: OSI Approved :: MIT License
|
|
@@ -142,7 +142,7 @@ handle = open_engine(
|
|
|
142
142
|
engine = handle.engine # Keep handle.runtime alive while using the engine.
|
|
143
143
|
```
|
|
144
144
|
|
|
145
|
-
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under the home directory's `.cache` when that variable is unset. The same rule applies on every OS; on Windows the home directory is `%USERPROFILE%`, so the default is `%USERPROFILE%\.cache\tensoracle`. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, input shapes, builder settings, GPU
|
|
145
|
+
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under the home directory's `.cache` when that variable is unset. The same rule applies on every OS; on Windows the home directory is `%USERPROFILE%`, so the default is `%USERPROFILE%\.cache\tensoracle`. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, input shapes, builder settings, the GPU's compatibility properties, OS, CPU architecture, and TensorRT version. The GPU is the device of the current CUDA context, or device 0 when no context is current. Its properties are the product name, compute capability, memory bus width, L2 cache size, shared memory limits, texture alignment, multiprocessor count, whether the GPU is integrated, ECC mode, and total memory; another device with the same properties reuses the saved engine. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Use a different `name` or `graph_version` for engines with different tensor types, such as fp32 and fp16. Include in `graph_version` any other caller value that changes the engine contents but is not part of the cache identity, such as a plugin manifest or a build recipe. Pass `force_rebuild=True` to replace a cached engine even when it is intact. tensoracle closes a saved engine file after reading it, so another process using an opened engine does not block replacement. On Windows, a process that keeps the saved file itself open prevents replacement or removal with `PermissionError`; after that process closes the file, a file left empty by the failed rebuild is rejected as `CorruptEngineError` and can be rebuilt with `force_rebuild=True`.
|
|
146
146
|
|
|
147
147
|
`plugins` names native TensorRT plugin libraries. Both `open_engine` and `build_serialized_engine` accept it. Each function checks that every library exports `getCreators`, loads it through TensorRT's registry, and retains it for the process before building the graph or deserializing an engine. Repeating the same library does not register it twice. `open_engine` also hashes the loaded library bytes into the engine identity; a changed library selects a different saved engine in a new process. A missing, unreadable, or incompatible library, or a library changed after loading in the same process, raises `PluginLibraryError` before a cached plan is opened. Keep the library and its cubin bundle together in the caller's wheel.
|
|
148
148
|
|
|
@@ -255,10 +255,75 @@ with handle.create_execution() as execution:
|
|
|
255
255
|
|
|
256
256
|
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.
|
|
257
257
|
|
|
258
|
-
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.
|
|
258
|
+
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.
|
|
259
259
|
|
|
260
260
|
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.
|
|
261
261
|
|
|
262
|
+
### Capture and replay with CUDA Graphs
|
|
263
|
+
|
|
264
|
+
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.
|
|
265
|
+
|
|
266
|
+
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.
|
|
267
|
+
|
|
268
|
+
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.
|
|
269
|
+
|
|
270
|
+
For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
|
|
271
|
+
|
|
272
|
+
```python
|
|
273
|
+
from contextlib import ExitStack
|
|
274
|
+
|
|
275
|
+
import numpy as np
|
|
276
|
+
import tensorrt as trt
|
|
277
|
+
from cuda.bindings import runtime as cudart
|
|
278
|
+
|
|
279
|
+
with handle.create_execution() as execution:
|
|
280
|
+
execution.set_input_shapes({"features": (1, 4, 2)})
|
|
281
|
+
# Omit set_input_shapes for fixed inputs and single-point profiles.
|
|
282
|
+
features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
|
|
283
|
+
output = np.empty(
|
|
284
|
+
execution.get_tensor_shape("result"),
|
|
285
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
286
|
+
)
|
|
287
|
+
with ExitStack() as resources:
|
|
288
|
+
stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
289
|
+
resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
|
|
290
|
+
source = cuda(cudart.cudaMalloc(features.nbytes))
|
|
291
|
+
resources.callback(lambda: cuda(cudart.cudaFree(source)))
|
|
292
|
+
target = cuda(cudart.cudaMalloc(output.nbytes))
|
|
293
|
+
resources.callback(lambda: cuda(cudart.cudaFree(target)))
|
|
294
|
+
addresses = {"features": int(source), "result": int(target)}
|
|
295
|
+
try:
|
|
296
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
297
|
+
source, features.ctypes.data, features.nbytes,
|
|
298
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
|
|
299
|
+
))
|
|
300
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
301
|
+
cuda(cudart.cudaStreamSynchronize(stream)) # Warm up before capture.
|
|
302
|
+
|
|
303
|
+
cuda(cudart.cudaStreamBeginCapture(
|
|
304
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
|
|
305
|
+
))
|
|
306
|
+
try:
|
|
307
|
+
execution.enqueue(addresses, stream=int(stream)) # Exactly one call.
|
|
308
|
+
finally:
|
|
309
|
+
graph = cuda(cudart.cudaStreamEndCapture(stream))
|
|
310
|
+
resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
|
|
311
|
+
instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
312
|
+
resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
313
|
+
|
|
314
|
+
for _ in range(3):
|
|
315
|
+
cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
316
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
317
|
+
output.ctypes.data, target, output.nbytes,
|
|
318
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
|
|
319
|
+
))
|
|
320
|
+
finally:
|
|
321
|
+
# Wait before ExitStack destroys the graph and releases its buffers.
|
|
322
|
+
cuda(cudart.cudaStreamSynchronize(stream))
|
|
323
|
+
```
|
|
324
|
+
|
|
325
|
+
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).
|
|
326
|
+
|
|
262
327
|
## Build a cubin plugin
|
|
263
328
|
|
|
264
329
|
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.
|
|
@@ -302,7 +367,7 @@ Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in th
|
|
|
302
367
|
|
|
303
368
|
For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
|
|
304
369
|
|
|
305
|
-
`tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records Triton
|
|
370
|
+
`tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records the Triton and ptxas versions and each kernel's source SHA-256, without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
|
|
306
371
|
|
|
307
372
|
## Tested environments
|
|
308
373
|
|
|
@@ -111,7 +111,7 @@ handle = open_engine(
|
|
|
111
111
|
engine = handle.engine # Keep handle.runtime alive while using the engine.
|
|
112
112
|
```
|
|
113
113
|
|
|
114
|
-
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under the home directory's `.cache` when that variable is unset. The same rule applies on every OS; on Windows the home directory is `%USERPROFILE%`, so the default is `%USERPROFILE%\.cache\tensoracle`. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, input shapes, builder settings, GPU
|
|
114
|
+
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under the home directory's `.cache` when that variable is unset. The same rule applies on every OS; on Windows the home directory is `%USERPROFILE%`, so the default is `%USERPROFILE%\.cache\tensoracle`. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version, input shapes, builder settings, the GPU's compatibility properties, OS, CPU architecture, and TensorRT version. The GPU is the device of the current CUDA context, or device 0 when no context is current. Its properties are the product name, compute capability, memory bus width, L2 cache size, shared memory limits, texture alignment, multiprocessor count, whether the GPU is integrated, ECC mode, and total memory; another device with the same properties reuses the saved engine. Raise `graph_version` when graph code changes; without a version change, the previous engine is reused. Use a different `name` or `graph_version` for engines with different tensor types, such as fp32 and fp16. Include in `graph_version` any other caller value that changes the engine contents but is not part of the cache identity, such as a plugin manifest or a build recipe. Pass `force_rebuild=True` to replace a cached engine even when it is intact. tensoracle closes a saved engine file after reading it, so another process using an opened engine does not block replacement. On Windows, a process that keeps the saved file itself open prevents replacement or removal with `PermissionError`; after that process closes the file, a file left empty by the failed rebuild is rejected as `CorruptEngineError` and can be rebuilt with `force_rebuild=True`.
|
|
115
115
|
|
|
116
116
|
`plugins` names native TensorRT plugin libraries. Both `open_engine` and `build_serialized_engine` accept it. Each function checks that every library exports `getCreators`, loads it through TensorRT's registry, and retains it for the process before building the graph or deserializing an engine. Repeating the same library does not register it twice. `open_engine` also hashes the loaded library bytes into the engine identity; a changed library selects a different saved engine in a new process. A missing, unreadable, or incompatible library, or a library changed after loading in the same process, raises `PluginLibraryError` before a cached plan is opened. Keep the library and its cubin bundle together in the caller's wheel.
|
|
117
117
|
|
|
@@ -224,10 +224,75 @@ with handle.create_execution() as execution:
|
|
|
224
224
|
|
|
225
225
|
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.
|
|
226
226
|
|
|
227
|
-
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.
|
|
227
|
+
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.
|
|
228
228
|
|
|
229
229
|
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.
|
|
230
230
|
|
|
231
|
+
### Capture and replay with CUDA Graphs
|
|
232
|
+
|
|
233
|
+
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.
|
|
234
|
+
|
|
235
|
+
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.
|
|
236
|
+
|
|
237
|
+
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.
|
|
238
|
+
|
|
239
|
+
For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
|
|
240
|
+
|
|
241
|
+
```python
|
|
242
|
+
from contextlib import ExitStack
|
|
243
|
+
|
|
244
|
+
import numpy as np
|
|
245
|
+
import tensorrt as trt
|
|
246
|
+
from cuda.bindings import runtime as cudart
|
|
247
|
+
|
|
248
|
+
with handle.create_execution() as execution:
|
|
249
|
+
execution.set_input_shapes({"features": (1, 4, 2)})
|
|
250
|
+
# Omit set_input_shapes for fixed inputs and single-point profiles.
|
|
251
|
+
features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
|
|
252
|
+
output = np.empty(
|
|
253
|
+
execution.get_tensor_shape("result"),
|
|
254
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
255
|
+
)
|
|
256
|
+
with ExitStack() as resources:
|
|
257
|
+
stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
258
|
+
resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
|
|
259
|
+
source = cuda(cudart.cudaMalloc(features.nbytes))
|
|
260
|
+
resources.callback(lambda: cuda(cudart.cudaFree(source)))
|
|
261
|
+
target = cuda(cudart.cudaMalloc(output.nbytes))
|
|
262
|
+
resources.callback(lambda: cuda(cudart.cudaFree(target)))
|
|
263
|
+
addresses = {"features": int(source), "result": int(target)}
|
|
264
|
+
try:
|
|
265
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
266
|
+
source, features.ctypes.data, features.nbytes,
|
|
267
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
|
|
268
|
+
))
|
|
269
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
270
|
+
cuda(cudart.cudaStreamSynchronize(stream)) # Warm up before capture.
|
|
271
|
+
|
|
272
|
+
cuda(cudart.cudaStreamBeginCapture(
|
|
273
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
|
|
274
|
+
))
|
|
275
|
+
try:
|
|
276
|
+
execution.enqueue(addresses, stream=int(stream)) # Exactly one call.
|
|
277
|
+
finally:
|
|
278
|
+
graph = cuda(cudart.cudaStreamEndCapture(stream))
|
|
279
|
+
resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
|
|
280
|
+
instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
281
|
+
resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
282
|
+
|
|
283
|
+
for _ in range(3):
|
|
284
|
+
cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
285
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
286
|
+
output.ctypes.data, target, output.nbytes,
|
|
287
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
|
|
288
|
+
))
|
|
289
|
+
finally:
|
|
290
|
+
# Wait before ExitStack destroys the graph and releases its buffers.
|
|
291
|
+
cuda(cudart.cudaStreamSynchronize(stream))
|
|
292
|
+
```
|
|
293
|
+
|
|
294
|
+
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).
|
|
295
|
+
|
|
231
296
|
## Build a cubin plugin
|
|
232
297
|
|
|
233
298
|
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.
|
|
@@ -271,7 +336,7 @@ Use `PluginBase(name, plugin_namespace, std::move(fields), residual_io())` in th
|
|
|
271
336
|
|
|
272
337
|
For configurations the declaration cannot express, such as multiple outputs or computed output shapes, developers can override the virtual V3 methods in a derived class or ask tensoracle's maintainers to extend the declaration.
|
|
273
338
|
|
|
274
|
-
`tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records Triton
|
|
339
|
+
`tensoracle.bundle.inspect_bundle(bundle_directory, plugin=library_path)` checks the manifest, each cubin's SHA-256 and actual ELF SM, the supported hidden arguments and zero scratch sizes, and the cubin bytes embedded in the library. `inspect_wheel(wheel_path, plugin="path/inside/wheel.so")` performs the same checks on a wheel without installing it. Their distinct exceptions identify a missing cubin, corrupt cubin, false SM claim, or plugin and bundle mismatch. The manifest records the Triton and ptxas versions and each kernel's source SHA-256, without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
|
|
275
340
|
|
|
276
341
|
## Tested environments
|
|
277
342
|
|
|
@@ -155,12 +155,77 @@ with handle.create_execution() as execution:
|
|
|
155
155
|
|
|
156
156
|
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.
|
|
157
157
|
|
|
158
|
-
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.
|
|
158
|
+
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.
|
|
159
159
|
|
|
160
160
|
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.
|
|
161
161
|
|
|
162
162
|
`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.
|
|
163
163
|
|
|
164
|
+
#### Capture and replay with CUDA Graphs
|
|
165
|
+
|
|
166
|
+
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.
|
|
167
|
+
|
|
168
|
+
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.
|
|
169
|
+
|
|
170
|
+
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.
|
|
171
|
+
|
|
172
|
+
For the `features` / `result` engine above, use the `cuda(...)` helper from the preceding example:
|
|
173
|
+
|
|
174
|
+
```python
|
|
175
|
+
from contextlib import ExitStack
|
|
176
|
+
|
|
177
|
+
import numpy as np
|
|
178
|
+
import tensorrt as trt
|
|
179
|
+
from cuda.bindings import runtime as cudart
|
|
180
|
+
|
|
181
|
+
with handle.create_execution() as execution:
|
|
182
|
+
execution.set_input_shapes({"features": (1, 4, 2)})
|
|
183
|
+
# Omit set_input_shapes for fixed inputs and single-point profiles.
|
|
184
|
+
features = np.arange(8, dtype=np.float32).reshape(1, 4, 2) / 4
|
|
185
|
+
output = np.empty(
|
|
186
|
+
execution.get_tensor_shape("result"),
|
|
187
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
188
|
+
)
|
|
189
|
+
with ExitStack() as resources:
|
|
190
|
+
stream = cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
191
|
+
resources.callback(lambda: cuda(cudart.cudaStreamDestroy(stream)))
|
|
192
|
+
source = cuda(cudart.cudaMalloc(features.nbytes))
|
|
193
|
+
resources.callback(lambda: cuda(cudart.cudaFree(source)))
|
|
194
|
+
target = cuda(cudart.cudaMalloc(output.nbytes))
|
|
195
|
+
resources.callback(lambda: cuda(cudart.cudaFree(target)))
|
|
196
|
+
addresses = {"features": int(source), "result": int(target)}
|
|
197
|
+
try:
|
|
198
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
199
|
+
source, features.ctypes.data, features.nbytes,
|
|
200
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream,
|
|
201
|
+
))
|
|
202
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
203
|
+
cuda(cudart.cudaStreamSynchronize(stream)) # Warm up before capture.
|
|
204
|
+
|
|
205
|
+
cuda(cudart.cudaStreamBeginCapture(
|
|
206
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
|
|
207
|
+
))
|
|
208
|
+
try:
|
|
209
|
+
execution.enqueue(addresses, stream=int(stream)) # Exactly one call.
|
|
210
|
+
finally:
|
|
211
|
+
graph = cuda(cudart.cudaStreamEndCapture(stream))
|
|
212
|
+
resources.callback(lambda: cuda(cudart.cudaGraphDestroy(graph)))
|
|
213
|
+
instance = cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
214
|
+
resources.callback(lambda: cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
215
|
+
|
|
216
|
+
for _ in range(3):
|
|
217
|
+
cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
218
|
+
cuda(cudart.cudaMemcpyAsync(
|
|
219
|
+
output.ctypes.data, target, output.nbytes,
|
|
220
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream,
|
|
221
|
+
))
|
|
222
|
+
finally:
|
|
223
|
+
# Wait before ExitStack destroys the graph and releases its buffers.
|
|
224
|
+
cuda(cudart.cudaStreamSynchronize(stream))
|
|
225
|
+
```
|
|
226
|
+
|
|
227
|
+
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).
|
|
228
|
+
|
|
164
229
|
### Build a standard graph
|
|
165
230
|
|
|
166
231
|
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.
|
|
@@ -20,7 +20,7 @@ First check whether a combination of standard TensorRT layers can preserve the o
|
|
|
20
20
|
|
|
21
21
|
- **tensoracle tools:** `tensoracle.producer.compile_bundle` generates cubins and a manifest for supported GPU generations. The shared C++ plugin helpers support a caller-built native plugin with declared tensor types, formats, and an input whose shape the output copies; `tensoracle.bundle.inspect_bundle` or `inspect_wheel` checks the cubins and their match with the plugin. Pass the plugin library through `plugins=` to `build_serialized_engine` or `open_engine`. Follow [Build a cubin plugin](index.md#build-a-cubin-plugin) for the `PluginIO` declaration, including fp16 / HWC8 and two-input examples, and the [cubin plugin example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md) for the generation, build, and inspection operations.
|
|
22
22
|
- **You provide:** Identify the bottleneck, write the kernel and model-specific plugin behavior, choose its supported shapes and types, build and package its native library, and run the resulting engines. An fp32 kernel must not use TF32.
|
|
23
|
-
- **Gate:** For a custom cubin plugin, the bundle and plugin inspections must pass. Reapply stage 1 to the replacement's fp32 engine against the same fixed PyTorch reference. If using
|
|
23
|
+
- **Gate:** For a custom cubin plugin, the bundle and plugin inspections must pass. Reapply stage 1 to the replacement's fp32 engine against the same fixed PyTorch reference. If using a lower precision engine, reapply stage 2 against an accepted fp32 engine with the same limits and model-specific measurements. Advance only when the replacement meets those quality gates; otherwise retain the accepted implementation.
|
|
24
24
|
|
|
25
25
|
## 4. Measure speed under isolated load
|
|
26
26
|
|
|
@@ -30,7 +30,7 @@ Follow [Prepare a TensorRT SDK](../../docs/tensorrt-sdk.md) to obtain and verify
|
|
|
30
30
|
CMAKE_ARGS="-DTENSORRT_ROOT=$TENSORRT_ROOT" uv build examples/cubin_plugin --wheel
|
|
31
31
|
```
|
|
32
32
|
|
|
33
|
-
The example's build backend and CUDA tooling belong to this project, not to tensoracle's runtime dependencies. Before installation, inspect the wheel with `tensoracle.bundle.inspect_wheel(wheel, plugin="example_cubin_plugin/libtensoracle_example_plugin.so")`. This compares the library's embedded digest and cubin bytes with the manifest and cubin files.
|
|
33
|
+
The example's build backend and CUDA tooling belong to this project, not to tensoracle's runtime dependencies. Before installation, inspect the wheel with `tensoracle.bundle.inspect_wheel(wheel, plugin="example_cubin_plugin/libtensoracle_example_plugin.so")`. For a wheel built on Windows, pass `plugin="example_cubin_plugin/tensoracle_example_plugin.dll"`. This compares the library's embedded digest and cubin bytes with the manifest and cubin files.
|
|
34
34
|
|
|
35
35
|
Install the tensoracle wheel and this example wheel into a runtime Python environment with TensorRT 11 and `cuda-python`, but without Triton, torch, or the TensorRT SDK. `open_engine` registers the library before calling `graph` or opening a saved engine. The tests use two separate processes and an `XDG_CACHE_HOME` directory to verify that the second process reuses the saved engine without building a graph.
|
|
36
36
|
|
|
@@ -7,7 +7,7 @@ readme = "README.md"
|
|
|
7
7
|
license = "MIT"
|
|
8
8
|
license-files = ["LICENSE"]
|
|
9
9
|
classifiers = [
|
|
10
|
-
"Development Status ::
|
|
10
|
+
"Development Status :: 5 - Production/Stable",
|
|
11
11
|
"Environment :: GPU :: NVIDIA CUDA :: 13",
|
|
12
12
|
"Intended Audience :: Developers",
|
|
13
13
|
"License :: OSI Approved :: MIT License",
|
|
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
|
|
|
18
18
|
commit_id: str | None
|
|
19
19
|
__commit_id__: str | None
|
|
20
20
|
|
|
21
|
-
__version__ = version = '1.
|
|
22
|
-
__version_tuple__ = version_tuple = (1,
|
|
21
|
+
__version__ = version = '1.1.0'
|
|
22
|
+
__version_tuple__ = version_tuple = (1, 1, 0)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -0,0 +1,191 @@
|
|
|
1
|
+
"""Capture one enqueue and replay it with caller-owned CUDA graphs."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from contextlib import ExitStack, contextmanager
|
|
5
|
+
|
|
6
|
+
import numpy as np
|
|
7
|
+
import pytest
|
|
8
|
+
import tensorrt as trt
|
|
9
|
+
from cuda.bindings import runtime as cudart
|
|
10
|
+
|
|
11
|
+
import tensoracle
|
|
12
|
+
from test_build import _check_cuda
|
|
13
|
+
from test_execution import _open_sum
|
|
14
|
+
|
|
15
|
+
pytest_plugins = ("test_example_plugin",)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@contextmanager
|
|
19
|
+
def _captured(handle, batches, *, shapes):
|
|
20
|
+
"""Keep one execution, its addresses, and its captured graph alive together."""
|
|
21
|
+
with handle.create_execution() as execution, ExitStack() as resources:
|
|
22
|
+
execution.set_input_shapes(shapes)
|
|
23
|
+
output = np.empty(
|
|
24
|
+
execution.get_tensor_shape("result"),
|
|
25
|
+
dtype=trt.nptype(execution.get_tensor_dtype("result")),
|
|
26
|
+
)
|
|
27
|
+
stream = _check_cuda(cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking))
|
|
28
|
+
resources.callback(lambda: _check_cuda(cudart.cudaStreamDestroy(stream)))
|
|
29
|
+
addresses = {}
|
|
30
|
+
for name, value in {**batches[0], "result": output}.items():
|
|
31
|
+
pointer = _check_cuda(cudart.cudaMalloc(value.nbytes))
|
|
32
|
+
resources.callback(lambda pointer=pointer: _check_cuda(cudart.cudaFree(pointer)))
|
|
33
|
+
addresses[name] = int(pointer)
|
|
34
|
+
|
|
35
|
+
def copy_inputs(batch):
|
|
36
|
+
for name, value in batch.items():
|
|
37
|
+
_check_cuda(
|
|
38
|
+
cudart.cudaMemcpyAsync(
|
|
39
|
+
addresses[name],
|
|
40
|
+
value.ctypes.data,
|
|
41
|
+
value.nbytes,
|
|
42
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
|
|
43
|
+
stream,
|
|
44
|
+
)
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
def read_output():
|
|
48
|
+
_check_cuda(
|
|
49
|
+
cudart.cudaMemcpyAsync(
|
|
50
|
+
output.ctypes.data,
|
|
51
|
+
addresses["result"],
|
|
52
|
+
output.nbytes,
|
|
53
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
|
|
54
|
+
stream,
|
|
55
|
+
)
|
|
56
|
+
)
|
|
57
|
+
_check_cuda(cudart.cudaStreamSynchronize(stream))
|
|
58
|
+
return output.copy()
|
|
59
|
+
|
|
60
|
+
try:
|
|
61
|
+
# The first ordinary enqueue flushes deferred shape and plugin initialization.
|
|
62
|
+
ordinary = []
|
|
63
|
+
for batch in batches:
|
|
64
|
+
copy_inputs(batch)
|
|
65
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
66
|
+
ordinary.append(read_output())
|
|
67
|
+
|
|
68
|
+
_check_cuda(
|
|
69
|
+
cudart.cudaStreamBeginCapture(
|
|
70
|
+
stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal
|
|
71
|
+
)
|
|
72
|
+
)
|
|
73
|
+
try:
|
|
74
|
+
execution.enqueue(addresses, stream=int(stream))
|
|
75
|
+
finally:
|
|
76
|
+
graph = _check_cuda(cudart.cudaStreamEndCapture(stream))
|
|
77
|
+
resources.callback(lambda: _check_cuda(cudart.cudaGraphDestroy(graph)))
|
|
78
|
+
instance = _check_cuda(cudart.cudaGraphInstantiate(graph, 0))
|
|
79
|
+
resources.callback(lambda: _check_cuda(cudart.cudaGraphExecDestroy(instance)))
|
|
80
|
+
|
|
81
|
+
def replay(index):
|
|
82
|
+
copy_inputs(batches[index])
|
|
83
|
+
# An empty graph or a launch writing to the wrong address must fail.
|
|
84
|
+
_check_cuda(
|
|
85
|
+
cudart.cudaMemsetAsync(addresses["result"], 0xFF, output.nbytes, stream)
|
|
86
|
+
)
|
|
87
|
+
_check_cuda(cudart.cudaGraphLaunch(instance, stream))
|
|
88
|
+
actual = read_output()
|
|
89
|
+
np.testing.assert_array_equal(actual, ordinary[index])
|
|
90
|
+
return actual
|
|
91
|
+
|
|
92
|
+
yield replay
|
|
93
|
+
finally:
|
|
94
|
+
# Finish work before graphs, buffers, streams, and the execution are released.
|
|
95
|
+
_check_cuda(cudart.cudaStreamSynchronize(stream))
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def _sum_batches(rows):
|
|
99
|
+
return [
|
|
100
|
+
{
|
|
101
|
+
"features": np.arange(rows * 2, dtype=np.float32).reshape(rows, 2) / 4 + offset,
|
|
102
|
+
"extra": np.full((rows, 2), -0.25 - offset / 2, dtype=np.float32),
|
|
103
|
+
}
|
|
104
|
+
for offset in (0, 2)
|
|
105
|
+
]
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
@pytest.mark.gpu
|
|
109
|
+
@pytest.mark.req("SP-TENSORACLE-028")
|
|
110
|
+
def test_standard_engine_graph_replay_matches_ordinary_enqueue(tmp_path, monkeypatch):
|
|
111
|
+
"""One captured enqueue on a nondefault stream reproduces ordinary sums for new input data."""
|
|
112
|
+
handle = _open_sum(tmp_path, monkeypatch, fixed=True)
|
|
113
|
+
batches = _sum_batches(3)
|
|
114
|
+
with _captured(handle, batches, shapes={}) as replay:
|
|
115
|
+
for index in (0, 1, 0):
|
|
116
|
+
np.testing.assert_array_equal(
|
|
117
|
+
replay(index), batches[index]["features"] + batches[index]["extra"]
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
@pytest.mark.gpu
|
|
122
|
+
@pytest.mark.req("SP-TENSORACLE-028")
|
|
123
|
+
def test_plugin_engine_graph_replay_matches_ordinary_enqueue(installed_example, tmp_path):
|
|
124
|
+
"""A captured cubin plugin runs repeatedly and matches ordinary enqueue with new inputs."""
|
|
125
|
+
library = installed_example["library"]
|
|
126
|
+
manifest = json.loads((library.parent / "bundle/bundle.json").read_text())
|
|
127
|
+
digest = np.frombuffer(manifest["bundle_sha256"].encode(), dtype=np.uint8)
|
|
128
|
+
bias = np.arange(16, dtype=np.float32) / 8
|
|
129
|
+
|
|
130
|
+
def graph(network):
|
|
131
|
+
creator = trt.get_plugin_registry().get_creator(
|
|
132
|
+
"TensoracleExampleBias", "1", "tensoracle.example"
|
|
133
|
+
)
|
|
134
|
+
assert creator is not None
|
|
135
|
+
fields = trt.PluginFieldCollection(
|
|
136
|
+
[
|
|
137
|
+
trt.PluginField("bundle_sha256", digest, trt.PluginFieldType.CHAR),
|
|
138
|
+
trt.PluginField("bias", bias, trt.PluginFieldType.FLOAT32),
|
|
139
|
+
]
|
|
140
|
+
)
|
|
141
|
+
plugin = creator.create_plugin("bias", fields, trt.TensorRTPhase.BUILD)
|
|
142
|
+
assert plugin is not None
|
|
143
|
+
features = network.add_input("features", trt.float32, (3, 16))
|
|
144
|
+
result = network.add_plugin_v3([features], [], plugin).get_output(0)
|
|
145
|
+
result.name = "result"
|
|
146
|
+
network.mark_output(result)
|
|
147
|
+
|
|
148
|
+
plan = tensoracle.build_serialized_engine(
|
|
149
|
+
graph,
|
|
150
|
+
tensoracle.EngineBuild(
|
|
151
|
+
{"features": tensoracle.InputShape.fixed((3, 16))},
|
|
152
|
+
tensoracle.BuilderSettings(optimization_level=0),
|
|
153
|
+
),
|
|
154
|
+
plugins=(library,),
|
|
155
|
+
)
|
|
156
|
+
logger = trt.Logger(trt.Logger.ERROR)
|
|
157
|
+
runtime = trt.Runtime(logger)
|
|
158
|
+
engine = runtime.deserialize_cuda_engine(plan)
|
|
159
|
+
assert engine is not None
|
|
160
|
+
handle = tensoracle.EngineHandle(engine, runtime, tmp_path / "bias.engine")
|
|
161
|
+
batches = [
|
|
162
|
+
{"features": np.arange(48, dtype=np.float32).reshape(3, 16) / 4 + offset}
|
|
163
|
+
for offset in (0, 2)
|
|
164
|
+
]
|
|
165
|
+
with _captured(handle, batches, shapes={}) as replay:
|
|
166
|
+
for index in (0, 1, 0):
|
|
167
|
+
np.testing.assert_array_equal(replay(index), batches[index]["features"] + bias)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
@pytest.mark.gpu
|
|
171
|
+
@pytest.mark.req("SP-TENSORACLE-028")
|
|
172
|
+
def test_ranged_engine_keeps_a_graph_and_execution_per_shape(tmp_path, monkeypatch):
|
|
173
|
+
"""Two shapes keep dedicated executions and graphs that replay correctly in alternation."""
|
|
174
|
+
handle = _open_sum(tmp_path, monkeypatch)
|
|
175
|
+
batches = [_sum_batches(rows) for rows in (2, 5)]
|
|
176
|
+
with ExitStack() as resources:
|
|
177
|
+
replays = [
|
|
178
|
+
resources.enter_context(
|
|
179
|
+
_captured(
|
|
180
|
+
handle,
|
|
181
|
+
values,
|
|
182
|
+
shapes={name: value.shape for name, value in values[0].items()},
|
|
183
|
+
)
|
|
184
|
+
)
|
|
185
|
+
for values in batches
|
|
186
|
+
]
|
|
187
|
+
for index in (0, 1, 0):
|
|
188
|
+
for replay, values in zip(replays, batches, strict=True):
|
|
189
|
+
np.testing.assert_array_equal(
|
|
190
|
+
replay(index), values[index]["features"] + values[index]["extra"]
|
|
191
|
+
)
|
|
@@ -410,7 +410,7 @@ def test_wheel_metadata_describes_the_public_package(distributions):
|
|
|
410
410
|
assert metadata.get_all("License-File") == ["LICENSE"]
|
|
411
411
|
assert metadata["Requires-Python"] == "<3.14,>=3.12"
|
|
412
412
|
assert set(metadata.get_all("Classifier", [])) == {
|
|
413
|
-
"Development Status ::
|
|
413
|
+
"Development Status :: 5 - Production/Stable",
|
|
414
414
|
"Environment :: GPU :: NVIDIA CUDA :: 13",
|
|
415
415
|
"Intended Audience :: Developers",
|
|
416
416
|
"License :: OSI Approved :: MIT License",
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py
RENAMED
|
File without changes
|
{tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|