tensoracle 0.4.0__tar.gz → 0.5.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-0.4.0 → tensoracle-0.5.0}/.gitignore +1 -1
- {tensoracle-0.4.0 → tensoracle-0.5.0}/PKG-INFO +17 -6
- {tensoracle-0.4.0 → tensoracle-0.5.0}/README.md +16 -5
- tensoracle-0.5.0/build_backend.py +94 -0
- tensoracle-0.5.0/cpp/CMakeLists.txt +23 -0
- tensoracle-0.5.0/cpp/tests/core.cpp +31 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/docs/index.md +5 -3
- {tensoracle-0.4.0 → tensoracle-0.5.0}/docs/optimization.md +2 -2
- {tensoracle-0.4.0 → tensoracle-0.5.0}/docs/tensorrt-sdk.md +1 -1
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/README.md +3 -3
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -1
- {tensoracle-0.4.0 → tensoracle-0.5.0}/mkdocs.yml +2 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/pyproject.toml +23 -6
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/_build.py +1 -5
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/_cache.py +0 -1
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/_version.py +2 -2
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/build_memory_runtime.py +0 -1
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_build.py +15 -12
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_build_memory.py +3 -3
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_bundle.py +6 -6
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_cache.py +18 -30
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_comparison.py +8 -8
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_convert.py +9 -9
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_cpp_core.py +21 -54
- tensoracle-0.5.0/tests/test_distribution.py +445 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_engine_io.py +8 -11
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_example_plugin.py +31 -19
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_execution.py +20 -22
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_plugin_io.py +8 -8
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_plugins.py +8 -8
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_producer.py +10 -9
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_shape_tensor_outputs.py +18 -19
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/test_weight_hash.py +15 -15
- {tensoracle-0.4.0 → tensoracle-0.5.0}/uv.lock +12 -0
- tensoracle-0.4.0/tests/test_distribution.py +0 -155
- tensoracle-0.4.0/tests/test_public_mirror_spec.py +0 -324
- {tensoracle-0.4.0 → tensoracle-0.5.0}/LICENSE +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/cpp/src/cubin_plugin.cpp +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/docs/api/bundle.md +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/docs/api/index.md +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/docs/api/producer.md +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/CMakeLists.txt +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/bundle.json +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/kernels.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/plugin.cpp +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/examples/cubin_plugin/pyproject.toml +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/__init__.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/_comparison.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/_execution.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/_memory.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/_plugins.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/bundle.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/convert.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/src/tensoracle/producer.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/cubin_io/check.cpp +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/cubin_io/kernels.cu +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/cubin_io/plugin.cpp +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/example_runtime.py +0 -0
- {tensoracle-0.4.0 → tensoracle-0.5.0}/tests/plugin_io_runtime.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: tensoracle
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.5.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
|
|
@@ -44,7 +44,7 @@ Choose the path that matches your work:
|
|
|
44
44
|
| --- | --- | --- |
|
|
45
45
|
| Inference application | [Reopen a saved engine](#use-a-saved-engine-for-inference) and enqueue it with `handle.create_execution()`. | The application supplies device buffers and CUDA streams, for example through CuPy or `cuda-python`. |
|
|
46
46
|
| Standard graph developer | [Define and build a graph](#define-and-build-a-standard-graph), [save and reopen](#save-and-reopen-an-engine), and [compare outputs](#compare-output-arrays). | None for standard TensorRT layers. |
|
|
47
|
-
| Cubin plugin developer | [Build a cubin plugin](#build-a-cubin-plugin) and package it with the model. | `tensoracle[producer]`, CUDA toolkit, TensorRT SDK, CMake, and a C++17 compiler. |
|
|
47
|
+
| Cubin plugin developer | [Build a cubin plugin](#build-a-cubin-plugin) and package it with the model. | `tensoracle[producer]`, CUDA toolkit, TensorRT SDK, CMake 3.20 or later, and a C++17 compiler. |
|
|
48
48
|
|
|
49
49
|
| Requirement | Inference application | Standard graph developer | Cubin plugin developer |
|
|
50
50
|
| --- | --- | --- | --- |
|
|
@@ -55,7 +55,7 @@ Choose the path that matches your work:
|
|
|
55
55
|
| Triton 3.8.0, `ptxas` 12.9.86, `cuobjdump` 13.3 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
|
|
56
56
|
| CUDA toolkit | Not needed | Not needed | Native plugin build |
|
|
57
57
|
| TensorRT 11 SDK | Not needed | Not needed | Native plugin build |
|
|
58
|
-
| CMake and C++17 compiler | Not needed | Not needed | Native plugin build |
|
|
58
|
+
| CMake 3.20 or later and C++17 compiler | Not needed | Not needed | Native plugin build |
|
|
59
59
|
|
|
60
60
|
The supported targets are Linux x86_64 (including WSL2), TensorRT 11, and GPU generations sm86, sm89, and sm120. The producer and native build tools are separate from the `tensoracle` runtime dependencies. Engine building, engine reuse, output comparison, and cubin production do not require torch. Converting `.pth` checkpoints requires the `convert` extra described below.
|
|
61
61
|
|
|
@@ -104,7 +104,6 @@ def graph(network: trt.INetworkDefinition) -> None:
|
|
|
104
104
|
|
|
105
105
|
|
|
106
106
|
declaration = EngineBuild(
|
|
107
|
-
precision="fp32",
|
|
108
107
|
inputs={"features": InputShape.ranged((1, 1, 2), (1, 64, 2), (1, 448, 2))},
|
|
109
108
|
settings=BuilderSettings(workspace_limit=1 << 30, optimization_level=3),
|
|
110
109
|
)
|
|
@@ -113,7 +112,7 @@ plan: bytes = build_serialized_engine(graph, declaration)
|
|
|
113
112
|
|
|
114
113
|
Use `InputShape.fixed((1, 64, 2))` for one input size; the graph may declare that size directly or use `-1` for a dimension fixed by the profile. For a range, put `-1` in each varying dimension of the graph input and declare its minimum, optimum, and maximum sizes. Every graph input needs a matching entry in `EngineBuild.inputs`.
|
|
115
114
|
|
|
116
|
-
|
|
115
|
+
The engine's precision is determined by the tensor types that the graph function places in the network. For an fp16 graph, use `trt.float16` tensors and `np.float16` weights where needed; the graph can also contain fp32 tensors. `EngineBuild` declares input shapes and builder settings. TF32 is always disabled; `BuilderSettings(tf32=True)` raises `ValueError`. The function returns serialized plan bytes. Invalid declarations raise `ValueError`, and a failed TensorRT build raises `RuntimeError`.
|
|
117
116
|
|
|
118
117
|
## Save and reopen an engine
|
|
119
118
|
|
|
@@ -138,7 +137,7 @@ handle = open_engine(
|
|
|
138
137
|
engine = handle.engine # Keep handle.runtime alive while using the engine.
|
|
139
138
|
```
|
|
140
139
|
|
|
141
|
-
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version,
|
|
140
|
+
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. 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.
|
|
142
141
|
|
|
143
142
|
`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.
|
|
144
143
|
|
|
@@ -265,6 +264,18 @@ Install the producer extra when you need to compile cubins from your own kernels
|
|
|
265
264
|
pip install 'tensoracle[producer]'
|
|
266
265
|
```
|
|
267
266
|
|
|
267
|
+
Every repository that adopts tensoracle must use this top-level layout:
|
|
268
|
+
|
|
269
|
+
```text
|
|
270
|
+
cpp/
|
|
271
|
+
CMakeLists.txt
|
|
272
|
+
python/
|
|
273
|
+
pyproject.toml
|
|
274
|
+
src/
|
|
275
|
+
```
|
|
276
|
+
|
|
277
|
+
Keep C++ sources in `cpp/`, with a CMake project that builds and runs its tests through CTest without Python. Keep the Python project, its src-layout package, and its uv environment in `python/`; run the project's uv commands from that directory. Tensoracle itself follows this layout, and the installed C++ helper paths described below remain inside its Python package.
|
|
278
|
+
|
|
268
279
|
`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0**, **ptxas 12.9.86**, and **cuobjdump 13.3**. It finds the NVIDIA wheel executables by default; explicit `ptxas` and `cuobjdump` paths can override them. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
|
|
269
280
|
|
|
270
281
|
The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
|
|
@@ -14,7 +14,7 @@ Choose the path that matches your work:
|
|
|
14
14
|
| --- | --- | --- |
|
|
15
15
|
| Inference application | [Reopen a saved engine](#use-a-saved-engine-for-inference) and enqueue it with `handle.create_execution()`. | The application supplies device buffers and CUDA streams, for example through CuPy or `cuda-python`. |
|
|
16
16
|
| Standard graph developer | [Define and build a graph](#define-and-build-a-standard-graph), [save and reopen](#save-and-reopen-an-engine), and [compare outputs](#compare-output-arrays). | None for standard TensorRT layers. |
|
|
17
|
-
| Cubin plugin developer | [Build a cubin plugin](#build-a-cubin-plugin) and package it with the model. | `tensoracle[producer]`, CUDA toolkit, TensorRT SDK, CMake, and a C++17 compiler. |
|
|
17
|
+
| Cubin plugin developer | [Build a cubin plugin](#build-a-cubin-plugin) and package it with the model. | `tensoracle[producer]`, CUDA toolkit, TensorRT SDK, CMake 3.20 or later, and a C++17 compiler. |
|
|
18
18
|
|
|
19
19
|
| Requirement | Inference application | Standard graph developer | Cubin plugin developer |
|
|
20
20
|
| --- | --- | --- | --- |
|
|
@@ -25,7 +25,7 @@ Choose the path that matches your work:
|
|
|
25
25
|
| Triton 3.8.0, `ptxas` 12.9.86, `cuobjdump` 13.3 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
|
|
26
26
|
| CUDA toolkit | Not needed | Not needed | Native plugin build |
|
|
27
27
|
| TensorRT 11 SDK | Not needed | Not needed | Native plugin build |
|
|
28
|
-
| CMake and C++17 compiler | Not needed | Not needed | Native plugin build |
|
|
28
|
+
| CMake 3.20 or later and C++17 compiler | Not needed | Not needed | Native plugin build |
|
|
29
29
|
|
|
30
30
|
The supported targets are Linux x86_64 (including WSL2), TensorRT 11, and GPU generations sm86, sm89, and sm120. The producer and native build tools are separate from the `tensoracle` runtime dependencies. Engine building, engine reuse, output comparison, and cubin production do not require torch. Converting `.pth` checkpoints requires the `convert` extra described below.
|
|
31
31
|
|
|
@@ -74,7 +74,6 @@ def graph(network: trt.INetworkDefinition) -> None:
|
|
|
74
74
|
|
|
75
75
|
|
|
76
76
|
declaration = EngineBuild(
|
|
77
|
-
precision="fp32",
|
|
78
77
|
inputs={"features": InputShape.ranged((1, 1, 2), (1, 64, 2), (1, 448, 2))},
|
|
79
78
|
settings=BuilderSettings(workspace_limit=1 << 30, optimization_level=3),
|
|
80
79
|
)
|
|
@@ -83,7 +82,7 @@ plan: bytes = build_serialized_engine(graph, declaration)
|
|
|
83
82
|
|
|
84
83
|
Use `InputShape.fixed((1, 64, 2))` for one input size; the graph may declare that size directly or use `-1` for a dimension fixed by the profile. For a range, put `-1` in each varying dimension of the graph input and declare its minimum, optimum, and maximum sizes. Every graph input needs a matching entry in `EngineBuild.inputs`.
|
|
85
84
|
|
|
86
|
-
|
|
85
|
+
The engine's precision is determined by the tensor types that the graph function places in the network. For an fp16 graph, use `trt.float16` tensors and `np.float16` weights where needed; the graph can also contain fp32 tensors. `EngineBuild` declares input shapes and builder settings. TF32 is always disabled; `BuilderSettings(tf32=True)` raises `ValueError`. The function returns serialized plan bytes. Invalid declarations raise `ValueError`, and a failed TensorRT build raises `RuntimeError`.
|
|
87
86
|
|
|
88
87
|
## Save and reopen an engine
|
|
89
88
|
|
|
@@ -108,7 +107,7 @@ handle = open_engine(
|
|
|
108
107
|
engine = handle.engine # Keep handle.runtime alive while using the engine.
|
|
109
108
|
```
|
|
110
109
|
|
|
111
|
-
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. The cache identity includes the contents of all declared safetensors weight files and optional `plugins` files, the engine name, graph version,
|
|
110
|
+
Saved files live in the `tensoracle` directory under `XDG_CACHE_HOME`, or under `~/.cache` when that variable is unset. 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.
|
|
112
111
|
|
|
113
112
|
`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.
|
|
114
113
|
|
|
@@ -235,6 +234,18 @@ Install the producer extra when you need to compile cubins from your own kernels
|
|
|
235
234
|
pip install 'tensoracle[producer]'
|
|
236
235
|
```
|
|
237
236
|
|
|
237
|
+
Every repository that adopts tensoracle must use this top-level layout:
|
|
238
|
+
|
|
239
|
+
```text
|
|
240
|
+
cpp/
|
|
241
|
+
CMakeLists.txt
|
|
242
|
+
python/
|
|
243
|
+
pyproject.toml
|
|
244
|
+
src/
|
|
245
|
+
```
|
|
246
|
+
|
|
247
|
+
Keep C++ sources in `cpp/`, with a CMake project that builds and runs its tests through CTest without Python. Keep the Python project, its src-layout package, and its uv environment in `python/`; run the project's uv commands from that directory. Tensoracle itself follows this layout, and the installed C++ helper paths described below remain inside its Python package.
|
|
248
|
+
|
|
238
249
|
`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0**, **ptxas 12.9.86**, and **cuobjdump 13.3**. It finds the NVIDIA wheel executables by default; explicit `ptxas` and `cuobjdump` paths can override them. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
|
|
239
250
|
|
|
240
251
|
The wheel also contains `tensoracle/cpp/include`, `tensoracle/cpp/src`, and `tensoracle/cpp/cmake`. Pass the installed `tensoracle/cpp` directory to CMake as `TENSORACLE_CPP_ROOT` and include `TensoracleCubinPlugin.cmake`. Its `tensoracle_add_cubin_plugin` function compiles the common C++ source, the caller's plugin source, and the generated `bundle.cpp` into the caller's library. The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common source owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common source releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
|
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
"""Build from the Python project using the repository's canonical shared sources."""
|
|
2
|
+
|
|
3
|
+
import shutil
|
|
4
|
+
import tempfile
|
|
5
|
+
import tomllib
|
|
6
|
+
from contextlib import contextmanager
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
|
|
9
|
+
from hatchling import build as hatchling_build
|
|
10
|
+
from hatchling.builders.sdist import SdistBuilder
|
|
11
|
+
from hatchling.builders.wheel import WheelBuilder
|
|
12
|
+
|
|
13
|
+
get_requires_for_build_editable = hatchling_build.get_requires_for_build_editable
|
|
14
|
+
get_requires_for_build_sdist = hatchling_build.get_requires_for_build_sdist
|
|
15
|
+
get_requires_for_build_wheel = hatchling_build.get_requires_for_build_wheel
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@contextmanager
|
|
19
|
+
def _builder(builder_class):
|
|
20
|
+
project = Path(__file__).resolve().parent
|
|
21
|
+
config = tomllib.loads((project / "pyproject.toml").read_text(encoding="utf-8"))
|
|
22
|
+
# Source archives already contain the shared files next to pyproject.toml.
|
|
23
|
+
source_archive = (project / "PKG-INFO").is_file()
|
|
24
|
+
repository = project if source_archive else project.parent
|
|
25
|
+
shared = ("README.md", "LICENSE", "cpp", "docs", "examples", "mkdocs.yml", ".gitignore")
|
|
26
|
+
for name in shared:
|
|
27
|
+
if not (repository / name).exists():
|
|
28
|
+
raise FileNotFoundError(f"Required build source is missing: {repository / name}")
|
|
29
|
+
if source_archive:
|
|
30
|
+
yield builder_class(str(project), config=config)
|
|
31
|
+
return
|
|
32
|
+
|
|
33
|
+
with tempfile.TemporaryDirectory(prefix="tensoracle-build-") as directory:
|
|
34
|
+
staged = Path(directory)
|
|
35
|
+
for base, names in (
|
|
36
|
+
(project, ("pyproject.toml", "uv.lock", "build_backend.py", "src", "tests")),
|
|
37
|
+
(repository, shared),
|
|
38
|
+
):
|
|
39
|
+
for name in names:
|
|
40
|
+
source, destination = base / name, staged / name
|
|
41
|
+
if source.is_dir():
|
|
42
|
+
shutil.copytree(
|
|
43
|
+
source,
|
|
44
|
+
destination,
|
|
45
|
+
ignore=shutil.ignore_patterns("__pycache__", "*.pyc", ".pytest_cache"),
|
|
46
|
+
)
|
|
47
|
+
else:
|
|
48
|
+
shutil.copyfile(source, destination)
|
|
49
|
+
hatch = config["tool"]["hatch"]
|
|
50
|
+
hatch["version"]["raw-options"]["root"] = str(repository)
|
|
51
|
+
hatch["build"]["dev-mode-dirs"] = [str(project / "src")]
|
|
52
|
+
yield builder_class(str(staged), config=config)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def _build(builder_class, directory, version):
|
|
56
|
+
output = str(Path(directory).resolve())
|
|
57
|
+
with _builder(builder_class) as builder:
|
|
58
|
+
artifact = next(builder.build(directory=output, versions=[version]))
|
|
59
|
+
if version == "editable":
|
|
60
|
+
generated = Path(builder.root) / "src/tensoracle/_version.py"
|
|
61
|
+
target = Path(__file__).resolve().parent / "src/tensoracle/_version.py"
|
|
62
|
+
if generated != target:
|
|
63
|
+
shutil.copyfile(generated, target)
|
|
64
|
+
return Path(artifact).name
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def build_sdist(sdist_directory, config_settings=None):
|
|
68
|
+
return _build(SdistBuilder, sdist_directory, "standard")
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
|
|
72
|
+
return _build(WheelBuilder, wheel_directory, "standard")
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def build_editable(wheel_directory, config_settings=None, metadata_directory=None):
|
|
76
|
+
return _build(WheelBuilder, wheel_directory, "editable")
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None):
|
|
80
|
+
with _builder(WheelBuilder) as builder:
|
|
81
|
+
builder.metadata.validate_fields()
|
|
82
|
+
directory = Path(metadata_directory) / f"{builder.artifact_project_id}.dist-info"
|
|
83
|
+
directory.mkdir(parents=True, exist_ok=True)
|
|
84
|
+
metadata = builder.config.core_metadata_constructor(builder.metadata)
|
|
85
|
+
(directory / "METADATA").write_text(metadata, encoding="utf-8")
|
|
86
|
+
for relative in builder.metadata.core.license_files:
|
|
87
|
+
target = directory / "licenses" / relative
|
|
88
|
+
target.parent.mkdir(parents=True, exist_ok=True)
|
|
89
|
+
shutil.copyfile(Path(builder.root) / relative, target)
|
|
90
|
+
return directory.name
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
def prepare_metadata_for_build_editable(metadata_directory, config_settings=None):
|
|
94
|
+
return prepare_metadata_for_build_wheel(metadata_directory, config_settings)
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
cmake_minimum_required(VERSION 3.20)
|
|
2
|
+
project(tensoracle_cubin LANGUAGES CXX)
|
|
3
|
+
|
|
4
|
+
set(TENSORRT_ROOT "" CACHE PATH "TensorRT 11 SDK root containing include/NvInfer.h")
|
|
5
|
+
if(NOT EXISTS "${TENSORRT_ROOT}/include/NvInfer.h")
|
|
6
|
+
message(FATAL_ERROR "TensorRT SDK headers are missing; set TENSORRT_ROOT to the TensorRT 11 SDK")
|
|
7
|
+
endif()
|
|
8
|
+
find_package(CUDAToolkit REQUIRED)
|
|
9
|
+
|
|
10
|
+
add_library(tensoracle_cubin STATIC src/cubin_plugin.cpp)
|
|
11
|
+
target_compile_features(tensoracle_cubin PUBLIC cxx_std_17)
|
|
12
|
+
target_include_directories(tensoracle_cubin PUBLIC
|
|
13
|
+
"${CMAKE_CURRENT_SOURCE_DIR}/include" "${TENSORRT_ROOT}/include")
|
|
14
|
+
target_link_libraries(tensoracle_cubin PUBLIC CUDA::cuda_driver)
|
|
15
|
+
target_compile_options(tensoracle_cubin PRIVATE
|
|
16
|
+
$<$<CXX_COMPILER_ID:GNU,Clang>:-Wno-deprecated-declarations>)
|
|
17
|
+
|
|
18
|
+
include(CTest)
|
|
19
|
+
if(BUILD_TESTING)
|
|
20
|
+
add_executable(tensoracle_core_test tests/core.cpp)
|
|
21
|
+
target_link_libraries(tensoracle_core_test PRIVATE tensoracle_cubin)
|
|
22
|
+
add_test(NAME tensoracle_fields_and_sha256 COMMAND tensoracle_core_test)
|
|
23
|
+
endif()
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
// SP-TENSORACLE-004: retain typed fields in order and compute SHA-256.
|
|
2
|
+
#include "tensoracle/cubin_plugin.hpp"
|
|
3
|
+
#include <cstring>
|
|
4
|
+
using namespace tensoracle::cubin;
|
|
5
|
+
using namespace nvinfer1;
|
|
6
|
+
int main() {
|
|
7
|
+
if (sha256("abc", 3) !=
|
|
8
|
+
"ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") return 1;
|
|
9
|
+
float weight[2] = {1.25f, 2.5f};
|
|
10
|
+
int version = 7;
|
|
11
|
+
PluginField input[] = {
|
|
12
|
+
{"weight", weight, PluginFieldType::kFLOAT32, 2},
|
|
13
|
+
{"version", &version, PluginFieldType::kINT32, 1},
|
|
14
|
+
};
|
|
15
|
+
PluginFieldCollection source{2, input};
|
|
16
|
+
Fields fields;
|
|
17
|
+
if (!Fields::parse(&source, fields)) return 2;
|
|
18
|
+
auto const* saved = fields.collection();
|
|
19
|
+
if (saved->nbFields != 2 ||
|
|
20
|
+
std::strcmp(saved->fields[0].name, "weight") != 0 ||
|
|
21
|
+
std::strcmp(saved->fields[1].name, "version") != 0 ||
|
|
22
|
+
saved->fields[0].type != PluginFieldType::kFLOAT32 ||
|
|
23
|
+
saved->fields[0].length != 2 ||
|
|
24
|
+
std::memcmp(saved->fields[0].data, weight, sizeof(weight)) != 0) return 3;
|
|
25
|
+
weight[0] = 9.0f;
|
|
26
|
+
if (std::memcmp(saved->fields[0].data, weight, sizeof(weight)) == 0) return 4;
|
|
27
|
+
Fields copy = fields;
|
|
28
|
+
if (std::memcmp(copy.collection()->fields[0].data, saved->fields[0].data,
|
|
29
|
+
sizeof(weight)) != 0) return 5;
|
|
30
|
+
return 0;
|
|
31
|
+
}
|
|
@@ -16,7 +16,7 @@ Follow the [four-stage optimization workflow](optimization.md) to establish outp
|
|
|
16
16
|
| --- | --- | --- |
|
|
17
17
|
| Inference application | Use `open_engine` to reopen a matching saved engine, then enqueue it with `handle.create_execution()`. | The application supplies buffers and streams, for example through CuPy or `cuda-python`. |
|
|
18
18
|
| Standard graph developer | Define a TensorRT graph, build it with `build_serialized_engine` or `open_engine`, and measure outputs with `compare_outputs`. | None for standard TensorRT layers. |
|
|
19
|
-
| Cubin plugin developer | Produce cubins, build a native TensorRT plugin, and package it with the model. | `tensoracle[producer]`, CUDA toolkit, TensorRT SDK, CMake, and a C++17 compiler. |
|
|
19
|
+
| Cubin plugin developer | Produce cubins, build a native TensorRT plugin, and package it with the model. | `tensoracle[producer]`, CUDA toolkit, TensorRT SDK, CMake 3.20 or later, and a C++17 compiler. |
|
|
20
20
|
|
|
21
21
|
| Requirement | Inference application | Standard graph developer | Cubin plugin developer |
|
|
22
22
|
| --- | --- | --- | --- |
|
|
@@ -27,7 +27,7 @@ Follow the [four-stage optimization workflow](optimization.md) to establish outp
|
|
|
27
27
|
| Triton 3.8.0, `ptxas` 12.9.86, `cuobjdump` 13.3 | Not needed | Not needed | Installed with `tensoracle[producer]` for cubin production |
|
|
28
28
|
| CUDA toolkit | Not needed | Not needed | Native plugin build |
|
|
29
29
|
| TensorRT 11 SDK | Not needed | Not needed | Native plugin build |
|
|
30
|
-
| CMake and C++17 compiler | Not needed | Not needed | Native plugin build |
|
|
30
|
+
| CMake 3.20 or later and C++17 compiler | Not needed | Not needed | Native plugin build |
|
|
31
31
|
|
|
32
32
|
The supported platform is Linux x86_64 (including WSL2) with sm86, sm89, or sm120 GPUs and TensorRT 11. The producer and native build tools are separate from the `tensoracle` runtime dependencies. Engine building, engine reuse, output comparison, and cubin production do not require torch. Converting `.pth` checkpoints requires the `convert` extra.
|
|
33
33
|
|
|
@@ -161,12 +161,14 @@ Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutput
|
|
|
161
161
|
|
|
162
162
|
### Build a standard graph
|
|
163
163
|
|
|
164
|
-
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`. `open_engine` uses the same declaration plus a name, graph version, weight files, and `EngineIO` to save and reuse a checked engine.
|
|
164
|
+
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) and fp16. `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.
|
|
165
165
|
|
|
166
166
|
The application runs the engine and passes a CPU NumPy candidate output and a corresponding reference output to `compare_outputs`. The application chooses the acceptance limits. See [Compare output arrays](https://github.com/sync-dev-org/tensoracle#compare-output-arrays).
|
|
167
167
|
|
|
168
168
|
### Build a cubin plugin
|
|
169
169
|
|
|
170
|
+
Repositories that adopt tensoracle must follow the [repository layout in the plugin guide](https://github.com/sync-dev-org/tensoracle#build-a-cubin-plugin).
|
|
171
|
+
|
|
170
172
|
Install `tensoracle[producer]` when you need to compile cubins from your own kernels:
|
|
171
173
|
|
|
172
174
|
```bash
|
|
@@ -4,13 +4,13 @@ Use the same representative inputs, shapes, and corresponding outputs throughout
|
|
|
4
4
|
|
|
5
5
|
## 1. Establish the fp32 reference
|
|
6
6
|
|
|
7
|
-
- **tensoracle tools:** Build an fp32 TensorRT engine
|
|
7
|
+
- **tensoracle tools:** Build an fp32 TensorRT engine from a graph whose tensor types are fp32, using an `EngineBuild` declaration for input shapes and builder settings with `build_serialized_engine` or `open_engine`. TF32 is disabled for the engine. Run the engine in the application, then use `compare_outputs(candidate, reference, atol=..., rtol=...)` on corresponding CPU NumPy output arrays.
|
|
8
8
|
- **You provide:** Run the official PyTorch implementation in fp32 outside `tensoracle` and fix its outputs as the reference. Supply the TensorRT graph, weights, representative inputs, and an acceptance limit for each output. Keep the reference and candidate inputs and preprocessing identical.
|
|
9
9
|
- **Gate:** Reject non-finite outputs and any output that misses your limits. Advance only when every fp32 TensorRT output passes against its fixed PyTorch reference. Correct the graph or inputs before trying lower precision if this gate fails.
|
|
10
10
|
|
|
11
11
|
## 2. Accept lower precision by quality
|
|
12
12
|
|
|
13
|
-
- **tensoracle tools:** Build an fp16 TensorRT engine
|
|
13
|
+
- **tensoracle tools:** Build an fp16 TensorRT engine from a graph whose tensor types are fp16, using the same `EngineBuild` declaration for input shapes and builder settings. The graph function determines precision through its tensor types. When saving with `open_engine`, use a different `name` or `graph_version` from the fp32 engine. Use `compare_outputs` with each fp16 output as the candidate and its accepted fp32 TensorRT output as the reference. It reports non-finite counts and, for finite floating point arrays, maximum absolute error, relative L2 error, cosine similarity, and, when both tolerances are supplied, counts and fractions outside tolerance.
|
|
14
14
|
- **You provide:** Choose limits for those model-independent measurements. Compute any model-specific measurements separately, such as PSNR for images or application-specific time-series quality measures, and choose their limits. Run both engines on the same inputs.
|
|
15
15
|
- **Gate:** Accept the fp16 engine only when its outputs meet all applicable limits, including the model-specific ones. Otherwise keep the accepted fp32 engine or revise the candidate and repeat this stage.
|
|
16
16
|
|
|
@@ -102,7 +102,7 @@ Keep `TENSORRT_ROOT` set to the absolute path of the verified `sdk` directory. F
|
|
|
102
102
|
|
|
103
103
|
```bash
|
|
104
104
|
CMAKE_ARGS="-DTENSORRT_ROOT=$TENSORRT_ROOT" uv build examples/cubin_plugin --wheel
|
|
105
|
-
TENSORACLE_SDK_ROOT="$TENSORRT_ROOT" uv run pytest -q tests/test_example_plugin.py
|
|
105
|
+
TENSORACLE_SDK_ROOT="$TENSORRT_ROOT" uv run --project python pytest -q python/tests/test_example_plugin.py
|
|
106
106
|
```
|
|
107
107
|
|
|
108
108
|
The tests also need the CUDA toolkit; tests marked `gpu` require an NVIDIA GPU and driver. The SDK is a build input. The application still selects the TensorRT runtime and verifies its own plugin with that runtime.
|
|
@@ -9,7 +9,7 @@ The first plugin computes an fp32 32 by 16 dot product and then scales it in a s
|
|
|
9
9
|
From the tensoracle repository root, run:
|
|
10
10
|
|
|
11
11
|
```bash
|
|
12
|
-
uv run --extra producer python -c 'from pathlib import Path; from examples.cubin_plugin.kernels import kernel_specs; from tensoracle.producer import compile_bundle; compile_bundle(kernel_specs(), Path("examples/cubin_plugin/bundle"))'
|
|
12
|
+
uv run --project python --extra producer python -c 'from pathlib import Path; from examples.cubin_plugin.kernels import kernel_specs; from tensoracle.producer import compile_bundle; compile_bundle(kernel_specs(), Path("examples/cubin_plugin/bundle"))'
|
|
13
13
|
```
|
|
14
14
|
|
|
15
15
|
This installs `tensoracle[producer]` with Triton 3.8.0, ptxas 12.9.86, and cuobjdump 13.3, and recreates `bundle/` without using a GPU. `compile_bundle` finds ptxas and cuobjdump in their installed NVIDIA wheels. The producer checks the cubin ELF SM, public and hidden argument layout, register count, and zero scratch sizes before writing the manifest. Run the tests below with the producer extra alone so torch is absent; `compile_bundle` rejects an already imported torch.
|
|
@@ -17,7 +17,7 @@ This installs `tensoracle[producer]` with Triton 3.8.0, ptxas 12.9.86, and cuobj
|
|
|
17
17
|
Check that the regenerated files match the committed example without using a GPU or setting producer environment variables:
|
|
18
18
|
|
|
19
19
|
```bash
|
|
20
|
-
uv run --extra producer pytest -m "not gpu" tests/test_producer.py
|
|
20
|
+
uv run --project python --extra producer pytest -m "not gpu" python/tests/test_producer.py
|
|
21
21
|
```
|
|
22
22
|
|
|
23
23
|
Without the producer extra, the compilation tests skip with an explanation of the missing tools.
|
|
@@ -39,7 +39,7 @@ Install the tensoracle wheel and this example wheel into a runtime Python enviro
|
|
|
39
39
|
Run from the tensoracle repository with the SDK and CUDA toolkit paths set. The GPU tests build the example wheel, install it in a temporary runtime environment, and run dynamic shapes, direct builds, saved engine reuse, simultaneous execution in two distinct CUDA contexts, and an fp32 versus TF32 control. Each GPU subprocess has a 50 second limit.
|
|
40
40
|
|
|
41
41
|
```bash
|
|
42
|
-
TENSORACLE_SDK_ROOT="$TENSORRT_ROOT" uv run pytest -q tests/test_example_plugin.py
|
|
42
|
+
TENSORACLE_SDK_ROOT="$TENSORRT_ROOT" uv run --project python pytest -q python/tests/test_example_plugin.py
|
|
43
43
|
```
|
|
44
44
|
|
|
45
45
|
Record the GPU name, SM, TensorRT runtime version, test result, and wheel digest on sm86, sm89, and sm120 separately. The code and wheel contain all three cubin generations; successful execution on one generation does not establish the other two.
|
|
@@ -38,13 +38,16 @@ producer = [
|
|
|
38
38
|
|
|
39
39
|
[build-system]
|
|
40
40
|
requires = ["hatchling", "hatch-vcs"]
|
|
41
|
-
build-backend = "
|
|
41
|
+
build-backend = "build_backend"
|
|
42
|
+
backend-path = ["."]
|
|
42
43
|
|
|
43
44
|
[tool.hatch.version]
|
|
44
45
|
source = "vcs"
|
|
45
46
|
fallback-version = "0.0.0.dev0"
|
|
46
47
|
|
|
47
48
|
[tool.hatch.version.raw-options]
|
|
49
|
+
root = ".."
|
|
50
|
+
fallback_root = "."
|
|
48
51
|
version_scheme = "only-version"
|
|
49
52
|
local_scheme = "no-local-version"
|
|
50
53
|
|
|
@@ -52,20 +55,34 @@ local_scheme = "no-local-version"
|
|
|
52
55
|
version-file = "src/tensoracle/_version.py"
|
|
53
56
|
|
|
54
57
|
[tool.hatch.build.targets.sdist]
|
|
55
|
-
exclude = [
|
|
58
|
+
exclude = [
|
|
59
|
+
"/.nf/",
|
|
60
|
+
"/.github/",
|
|
61
|
+
"/AGENTS.md",
|
|
62
|
+
"/CLAUDE.md",
|
|
63
|
+
"/tools/",
|
|
64
|
+
"/noren.yaml",
|
|
65
|
+
"/noren-words.txt",
|
|
66
|
+
]
|
|
56
67
|
|
|
57
68
|
[tool.hatch.build.targets.wheel]
|
|
58
69
|
packages = ["src/tensoracle"]
|
|
59
70
|
|
|
60
71
|
[tool.hatch.build.targets.wheel.force-include]
|
|
61
|
-
cpp = "tensoracle/cpp"
|
|
72
|
+
"cpp/cmake" = "tensoracle/cpp/cmake"
|
|
73
|
+
"cpp/include" = "tensoracle/cpp/include"
|
|
74
|
+
"cpp/src" = "tensoracle/cpp/src"
|
|
62
75
|
|
|
63
76
|
[dependency-groups]
|
|
77
|
+
docs = [
|
|
78
|
+
"mkdocs>=1.6.1",
|
|
79
|
+
"mkdocs-material>=9.7.7",
|
|
80
|
+
"mkdocstrings[python]>=1.0.6",
|
|
81
|
+
]
|
|
64
82
|
dev = [
|
|
83
|
+
{ include-group = "docs" },
|
|
65
84
|
"cuda-python>=13.0",
|
|
66
85
|
"cupy-cuda13x>=14.2.0",
|
|
67
|
-
"mkdocs-material>=9.7.7",
|
|
68
|
-
"mkdocstrings[python]>=1.0.6",
|
|
69
86
|
"pytest>=8.0",
|
|
70
87
|
"ruff>=0.8",
|
|
71
88
|
]
|
|
@@ -75,7 +92,7 @@ testpaths = ["tests"]
|
|
|
75
92
|
addopts = ["--strict-markers"]
|
|
76
93
|
markers = [
|
|
77
94
|
"gpu: tests that build or load a TensorRT engine on a CUDA GPU",
|
|
78
|
-
"req(id):
|
|
95
|
+
"req(id): ID of the adopted spec item that the test verifies",
|
|
79
96
|
]
|
|
80
97
|
|
|
81
98
|
[tool.ruff]
|
|
@@ -4,7 +4,6 @@ import warnings
|
|
|
4
4
|
from collections.abc import Callable, Mapping, Sequence
|
|
5
5
|
from dataclasses import dataclass, field
|
|
6
6
|
from pathlib import Path
|
|
7
|
-
from typing import Literal
|
|
8
7
|
|
|
9
8
|
import tensorrt as trt
|
|
10
9
|
|
|
@@ -77,9 +76,8 @@ class BuilderSettings:
|
|
|
77
76
|
|
|
78
77
|
@dataclass(frozen=True)
|
|
79
78
|
class EngineBuild:
|
|
80
|
-
"""One engine's
|
|
79
|
+
"""One engine's named input shapes and builder settings."""
|
|
81
80
|
|
|
82
|
-
precision: Literal["fp32", "fp16"]
|
|
83
81
|
inputs: Mapping[str, InputShape]
|
|
84
82
|
settings: BuilderSettings = field(default_factory=BuilderSettings)
|
|
85
83
|
|
|
@@ -103,8 +101,6 @@ def _build_serialized_engine(
|
|
|
103
101
|
plugins: Sequence[Path] = (),
|
|
104
102
|
logger: trt.ILogger | None = None,
|
|
105
103
|
) -> tuple[bytes, tuple[str, ...]]:
|
|
106
|
-
if declaration.precision not in ("fp32", "fp16"):
|
|
107
|
-
raise ValueError(f"unsupported precision: {declaration.precision}")
|
|
108
104
|
if declaration.settings.tf32:
|
|
109
105
|
raise ValueError("TF32 cannot be enabled")
|
|
110
106
|
if declaration.settings.workspace_limit <= 0:
|
|
@@ -277,7 +277,6 @@ def _conditions(
|
|
|
277
277
|
"graph_version": graph_version,
|
|
278
278
|
"weights": [_weight_sha256(Path(path), force_rebuild=force_rebuild) for path in weights],
|
|
279
279
|
"plugins": [_file_sha256(Path(path)) for path in plugins],
|
|
280
|
-
"precision": declaration.precision,
|
|
281
280
|
"inputs": {
|
|
282
281
|
name: {
|
|
283
282
|
"minimum": list(shape.minimum),
|
|
@@ -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 = '0.
|
|
22
|
-
__version_tuple__ = version_tuple = (0,
|
|
21
|
+
__version__ = version = '0.5.0'
|
|
22
|
+
__version_tuple__ = version_tuple = (0, 5, 0)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -13,7 +13,14 @@ from safetensors.numpy import load_file, save_file
|
|
|
13
13
|
from tensoracle import BuilderSettings, EngineBuild, InputShape, build_serialized_engine
|
|
14
14
|
|
|
15
15
|
|
|
16
|
-
@pytest.mark.req("
|
|
16
|
+
@pytest.mark.req("SP-TENSORACLE-020")
|
|
17
|
+
def test_build_declaration_rejects_precision_keyword():
|
|
18
|
+
"""The build declaration rejects precision because the graph owns tensor types."""
|
|
19
|
+
with pytest.raises(TypeError, match="precision"):
|
|
20
|
+
EngineBuild(precision="fp32", inputs={"features": InputShape.fixed((1, 3, 2))})
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@pytest.mark.req("SP-TENSORACLE-005")
|
|
17
24
|
def test_package_requires_torch_only_for_checkpoint_conversion():
|
|
18
25
|
"""Only the conversion extra declares torch, leaving inference dependencies unchanged."""
|
|
19
26
|
dependencies = requires("tensoracle") or []
|
|
@@ -114,13 +121,12 @@ def _run(plan: bytes, value: np.ndarray) -> np.ndarray:
|
|
|
114
121
|
|
|
115
122
|
|
|
116
123
|
@pytest.mark.gpu
|
|
117
|
-
@pytest.mark.req("
|
|
124
|
+
@pytest.mark.req("SP-TENSORACLE-007")
|
|
118
125
|
def test_fixed_shape_build_runs_at_its_declared_size(tmp_path):
|
|
119
126
|
"""A fixed input shape produces a plan that computes the declared graph."""
|
|
120
127
|
path = tmp_path / "weights.safetensors"
|
|
121
128
|
_weights(path)
|
|
122
129
|
declaration = EngineBuild(
|
|
123
|
-
precision="fp32",
|
|
124
130
|
inputs={"features": InputShape.fixed((1, 3, 2))},
|
|
125
131
|
settings=BuilderSettings(optimization_level=0),
|
|
126
132
|
)
|
|
@@ -131,13 +137,12 @@ def test_fixed_shape_build_runs_at_its_declared_size(tmp_path):
|
|
|
131
137
|
|
|
132
138
|
|
|
133
139
|
@pytest.mark.gpu
|
|
134
|
-
@pytest.mark.req("
|
|
140
|
+
@pytest.mark.req("SP-TENSORACLE-007")
|
|
135
141
|
def test_range_build_runs_at_minimum_optimum_and_maximum_and_rejects_outside(tmp_path):
|
|
136
142
|
"""A ranged input accepts its three declared sizes and rejects an outside size."""
|
|
137
143
|
path = tmp_path / "weights.safetensors"
|
|
138
144
|
_weights(path)
|
|
139
145
|
declaration = EngineBuild(
|
|
140
|
-
precision="fp32",
|
|
141
146
|
inputs={"features": InputShape.ranged((1, 1, 2), (1, 3, 2), (1, 5, 2))},
|
|
142
147
|
settings=BuilderSettings(optimization_level=0),
|
|
143
148
|
)
|
|
@@ -154,13 +159,13 @@ def test_range_build_runs_at_minimum_optimum_and_maximum_and_rejects_outside(tmp
|
|
|
154
159
|
|
|
155
160
|
|
|
156
161
|
@pytest.mark.gpu
|
|
157
|
-
@pytest.mark.req("
|
|
162
|
+
@pytest.mark.req("SP-TENSORACLE-013")
|
|
163
|
+
@pytest.mark.req("SP-TENSORACLE-020")
|
|
158
164
|
def test_fp16_build_runs_a_graph_with_fp16_tensors(tmp_path):
|
|
159
|
-
"""
|
|
165
|
+
"""An engine built without a precision declaration runs the graph's fp16 tensors."""
|
|
160
166
|
path = tmp_path / "weights.safetensors"
|
|
161
167
|
_weights(path)
|
|
162
168
|
declaration = EngineBuild(
|
|
163
|
-
precision="fp16",
|
|
164
169
|
inputs={"features": InputShape.ranged((1, 1, 2), (1, 2, 2), (1, 3, 2))},
|
|
165
170
|
settings=BuilderSettings(optimization_level=0),
|
|
166
171
|
)
|
|
@@ -170,11 +175,10 @@ def test_fp16_build_runs_a_graph_with_fp16_tensors(tmp_path):
|
|
|
170
175
|
np.testing.assert_allclose(_run(plan, value), expected, rtol=2e-3, atol=2e-3)
|
|
171
176
|
|
|
172
177
|
|
|
173
|
-
@pytest.mark.req("
|
|
178
|
+
@pytest.mark.req("SP-TENSORACLE-013")
|
|
174
179
|
def test_tf32_request_is_rejected_before_graph_is_built():
|
|
175
180
|
"""An explicit request to enable TF32 fails without invoking the graph."""
|
|
176
181
|
declaration = EngineBuild(
|
|
177
|
-
precision="fp32",
|
|
178
182
|
inputs={"features": InputShape.fixed((1, 3, 2))},
|
|
179
183
|
settings=BuilderSettings(tf32=True),
|
|
180
184
|
)
|
|
@@ -182,7 +186,7 @@ def test_tf32_request_is_rejected_before_graph_is_built():
|
|
|
182
186
|
build_serialized_engine(lambda network: pytest.fail("graph was built"), declaration)
|
|
183
187
|
|
|
184
188
|
|
|
185
|
-
@pytest.mark.req("
|
|
189
|
+
@pytest.mark.req("SP-TENSORACLE-013")
|
|
186
190
|
def test_tf32_is_disabled_and_other_builder_settings_are_applied(monkeypatch):
|
|
187
191
|
"""The builder clears TF32 while applying the declared workspace and optimization level."""
|
|
188
192
|
import tensoracle._build as module
|
|
@@ -224,7 +228,6 @@ def test_tf32_is_disabled_and_other_builder_settings_are_applied(monkeypatch):
|
|
|
224
228
|
observed = []
|
|
225
229
|
monkeypatch.setattr(module.trt, "Builder", Builder)
|
|
226
230
|
declaration = EngineBuild(
|
|
227
|
-
precision="fp32",
|
|
228
231
|
inputs={"features": InputShape.fixed((1, 3, 2))},
|
|
229
232
|
settings=BuilderSettings(workspace_limit=4096, optimization_level=2),
|
|
230
233
|
)
|