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.
Files changed (86) hide show
  1. {tensoracle-1.0.0 → tensoracle-1.1.0}/PKG-INFO +70 -5
  2. {tensoracle-1.0.0 → tensoracle-1.1.0}/README.md +68 -3
  3. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/api/index.md +2 -0
  4. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/index.md +66 -1
  5. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/optimization.md +1 -1
  6. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/README.md +1 -1
  7. {tensoracle-1.0.0 → tensoracle-1.1.0}/pyproject.toml +1 -1
  8. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_version.py +2 -2
  9. tensoracle-1.1.0/tests/test_cuda_graphs.py +191 -0
  10. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_distribution.py +1 -1
  11. {tensoracle-1.0.0 → tensoracle-1.1.0}/.gitignore +0 -0
  12. {tensoracle-1.0.0 → tensoracle-1.1.0}/LICENSE +0 -0
  13. {tensoracle-1.0.0 → tensoracle-1.1.0}/build_backend.py +0 -0
  14. {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/CMakeLists.txt +0 -0
  15. {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
  16. {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
  17. {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/src/cubin_plugin.cpp +0 -0
  18. {tensoracle-1.0.0 → tensoracle-1.1.0}/cpp/tests/core.cpp +0 -0
  19. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/api/bundle.md +0 -0
  20. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/api/producer.md +0 -0
  21. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/cpp-python-repositories.md +0 -0
  22. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/precision.md +0 -0
  23. {tensoracle-1.0.0 → tensoracle-1.1.0}/docs/tensorrt-sdk.md +0 -0
  24. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/CMakeLists.txt +0 -0
  25. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
  26. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/bundle.json +0 -0
  27. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
  28. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
  29. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
  30. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
  31. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
  32. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
  33. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
  34. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
  35. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
  36. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
  37. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
  38. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
  39. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
  40. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
  41. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
  42. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
  43. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
  44. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/kernels.py +0 -0
  45. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/plugin.cpp +0 -0
  46. {tensoracle-1.0.0 → tensoracle-1.1.0}/examples/cubin_plugin/pyproject.toml +0 -0
  47. {tensoracle-1.0.0 → tensoracle-1.1.0}/mkdocs.yml +0 -0
  48. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/__init__.py +0 -0
  49. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_build.py +0 -0
  50. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_cache.py +0 -0
  51. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_comparison.py +0 -0
  52. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_execution.py +0 -0
  53. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_memory.py +0 -0
  54. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/_plugins.py +0 -0
  55. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/bundle.py +0 -0
  56. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/convert.py +0 -0
  57. {tensoracle-1.0.0 → tensoracle-1.1.0}/src/tensoracle/producer.py +0 -0
  58. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/build_memory_runtime.py +0 -0
  59. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/cubin_io/check.cpp +0 -0
  60. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/cubin_io/kernels.cu +0 -0
  61. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/cubin_io/plugin.cpp +0 -0
  62. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/example_runtime.py +0 -0
  63. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/native/CMakeLists.txt +0 -0
  64. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/native/export.h +0 -0
  65. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/native_build.py +0 -0
  66. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/plugin_io_runtime.py +0 -0
  67. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/precision_runtime.py +0 -0
  68. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_build.py +0 -0
  69. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_build_memory.py +0 -0
  70. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_bundle.py +0 -0
  71. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_cache.py +0 -0
  72. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_cmake_plugin.py +0 -0
  73. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_comparison.py +0 -0
  74. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_convert.py +0 -0
  75. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_cpp_core.py +0 -0
  76. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_engine_io.py +0 -0
  77. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_example_plugin.py +0 -0
  78. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_execution.py +0 -0
  79. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_gpu_cache.py +0 -0
  80. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_plugin_io.py +0 -0
  81. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_plugins.py +0 -0
  82. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_precision.py +0 -0
  83. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_producer.py +0 -0
  84. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_shape_tensor_outputs.py +0 -0
  85. {tensoracle-1.0.0 → tensoracle-1.1.0}/tests/test_weight_hash.py +0 -0
  86. {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.0.0
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 :: 4 - Beta
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 product name, OS, CPU architecture, and TensorRT version. 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`.
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. 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.
@@ -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, ptxas, and kernel source digests without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
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 product name, OS, CPU architecture, and TensorRT version. 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`.
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. 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.
@@ -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, ptxas, and kernel source digests without a producer host or OS path. The TensorRT versions tested with the included example are listed below.
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
 
@@ -12,6 +12,8 @@ The following names are available directly from `tensoracle`.
12
12
 
13
13
  ::: tensoracle._build.build_serialized_engine
14
14
 
15
+ ::: tensoracle._build.ShapeTensorOutputWarning
16
+
15
17
  ## Saved engines
16
18
 
17
19
  ::: tensoracle._cache.TensorSpec
@@ -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.
@@ -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 fp16, 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.
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 :: 4 - Beta",
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.0.0'
22
- __version_tuple__ = version_tuple = (1, 0, 0)
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 :: 4 - Beta",
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