tensoracle 1.0.1__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.
Files changed (86) hide show
  1. {tensoracle-1.0.1 → tensoracle-1.1.0}/PKG-INFO +67 -2
  2. {tensoracle-1.0.1 → tensoracle-1.1.0}/README.md +66 -1
  3. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/index.md +66 -1
  4. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_version.py +2 -2
  5. tensoracle-1.1.0/tests/test_cuda_graphs.py +191 -0
  6. {tensoracle-1.0.1 → tensoracle-1.1.0}/.gitignore +0 -0
  7. {tensoracle-1.0.1 → tensoracle-1.1.0}/LICENSE +0 -0
  8. {tensoracle-1.0.1 → tensoracle-1.1.0}/build_backend.py +0 -0
  9. {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/CMakeLists.txt +0 -0
  10. {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
  11. {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
  12. {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/src/cubin_plugin.cpp +0 -0
  13. {tensoracle-1.0.1 → tensoracle-1.1.0}/cpp/tests/core.cpp +0 -0
  14. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/api/bundle.md +0 -0
  15. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/api/index.md +0 -0
  16. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/api/producer.md +0 -0
  17. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/cpp-python-repositories.md +0 -0
  18. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/optimization.md +0 -0
  19. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/precision.md +0 -0
  20. {tensoracle-1.0.1 → tensoracle-1.1.0}/docs/tensorrt-sdk.md +0 -0
  21. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/CMakeLists.txt +0 -0
  22. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/README.md +0 -0
  23. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
  24. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.json +0 -0
  25. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
  26. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
  27. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
  28. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
  29. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
  30. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
  31. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
  32. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
  33. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
  34. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
  35. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
  36. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
  37. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
  38. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
  39. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
  40. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
  41. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
  42. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/kernels.py +0 -0
  43. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/plugin.cpp +0 -0
  44. {tensoracle-1.0.1 → tensoracle-1.1.0}/examples/cubin_plugin/pyproject.toml +0 -0
  45. {tensoracle-1.0.1 → tensoracle-1.1.0}/mkdocs.yml +0 -0
  46. {tensoracle-1.0.1 → tensoracle-1.1.0}/pyproject.toml +0 -0
  47. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/__init__.py +0 -0
  48. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_build.py +0 -0
  49. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_cache.py +0 -0
  50. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_comparison.py +0 -0
  51. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_execution.py +0 -0
  52. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_memory.py +0 -0
  53. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/_plugins.py +0 -0
  54. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/bundle.py +0 -0
  55. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/convert.py +0 -0
  56. {tensoracle-1.0.1 → tensoracle-1.1.0}/src/tensoracle/producer.py +0 -0
  57. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/build_memory_runtime.py +0 -0
  58. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/cubin_io/check.cpp +0 -0
  59. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/cubin_io/kernels.cu +0 -0
  60. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/cubin_io/plugin.cpp +0 -0
  61. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/example_runtime.py +0 -0
  62. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/native/CMakeLists.txt +0 -0
  63. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/native/export.h +0 -0
  64. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/native_build.py +0 -0
  65. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/plugin_io_runtime.py +0 -0
  66. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/precision_runtime.py +0 -0
  67. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_build.py +0 -0
  68. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_build_memory.py +0 -0
  69. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_bundle.py +0 -0
  70. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_cache.py +0 -0
  71. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_cmake_plugin.py +0 -0
  72. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_comparison.py +0 -0
  73. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_convert.py +0 -0
  74. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_cpp_core.py +0 -0
  75. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_distribution.py +0 -0
  76. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_engine_io.py +0 -0
  77. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_example_plugin.py +0 -0
  78. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_execution.py +0 -0
  79. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_gpu_cache.py +0 -0
  80. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_plugin_io.py +0 -0
  81. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_plugins.py +0 -0
  82. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_precision.py +0 -0
  83. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_producer.py +0 -0
  84. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_shape_tensor_outputs.py +0 -0
  85. {tensoracle-1.0.1 → tensoracle-1.1.0}/tests/test_weight_hash.py +0 -0
  86. {tensoracle-1.0.1 → 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.0.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
@@ -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. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
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.
@@ -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. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
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.
@@ -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. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
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.
@@ -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.0.1'
22
- __version_tuple__ = version_tuple = (1, 0, 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
+ )
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes