tensoracle 0.4.1__tar.gz → 0.5.1__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.1 → tensoracle-0.5.1}/.gitignore +1 -1
- {tensoracle-0.4.1 → tensoracle-0.5.1}/PKG-INFO +20 -7
- {tensoracle-0.4.1 → tensoracle-0.5.1}/README.md +19 -6
- tensoracle-0.5.1/build_backend.py +94 -0
- tensoracle-0.5.1/cpp/CMakeLists.txt +23 -0
- tensoracle-0.5.1/cpp/tests/core.cpp +31 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/docs/index.md +6 -4
- {tensoracle-0.4.1 → tensoracle-0.5.1}/docs/optimization.md +3 -3
- tensoracle-0.5.1/docs/precision.md +154 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/docs/tensorrt-sdk.md +1 -1
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/README.md +3 -3
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -1
- {tensoracle-0.4.1 → tensoracle-0.5.1}/mkdocs.yml +3 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/pyproject.toml +14 -6
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/_build.py +1 -5
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/_cache.py +32 -4
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/_version.py +2 -2
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/build_memory_runtime.py +0 -1
- tensoracle-0.5.1/tests/precision_runtime.py +164 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_build.py +9 -6
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_build_memory.py +1 -1
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_cache.py +10 -22
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_cpp_core.py +19 -52
- tensoracle-0.5.1/tests/test_distribution.py +469 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_engine_io.py +6 -9
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_example_plugin.py +14 -4
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_execution.py +0 -2
- tensoracle-0.5.1/tests/test_gpu_cache.py +204 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_plugin_io.py +4 -4
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_plugins.py +2 -2
- tensoracle-0.5.1/tests/test_precision.py +111 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_producer.py +3 -2
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_shape_tensor_outputs.py +1 -2
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_weight_hash.py +18 -18
- {tensoracle-0.4.1 → tensoracle-0.5.1}/uv.lock +12 -14
- tensoracle-0.4.1/tests/test_distribution.py +0 -164
- {tensoracle-0.4.1 → tensoracle-0.5.1}/LICENSE +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/cpp/src/cubin_plugin.cpp +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/docs/api/bundle.md +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/docs/api/index.md +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/docs/api/producer.md +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/CMakeLists.txt +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/bundle.json +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/kernels.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/plugin.cpp +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/examples/cubin_plugin/pyproject.toml +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/__init__.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/_comparison.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/_execution.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/_memory.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/_plugins.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/bundle.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/convert.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/src/tensoracle/producer.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/cubin_io/check.cpp +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/cubin_io/kernels.cu +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/cubin_io/plugin.cpp +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/example_runtime.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/plugin_io_runtime.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_bundle.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_comparison.py +0 -0
- {tensoracle-0.4.1 → tensoracle-0.5.1}/tests/test_convert.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: tensoracle
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.5.1
|
|
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,11 +55,11 @@ 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
|
|
|
62
|
-
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate
|
|
62
|
+
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
|
|
63
63
|
|
|
64
64
|
## Convert a PyTorch checkpoint
|
|
65
65
|
|
|
@@ -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,9 @@ 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`.
|
|
116
|
+
|
|
117
|
+
The same API builds bf16, int8, and fp8 graphs. Cast tensors to `trt.bfloat16` for bf16; for int8/fp8, place `add_quantize` / `add_dequantize` layers and load their scales from safetensors. The application determines scales with an external quantization library and chooses output quality limits with `compare_outputs`. fp8 targets sm89 and sm120; bf16 and int8 also target sm86. Follow [Build bf16, int8, and fp8 engines](https://github.com/sync-dev-org/tensoracle/blob/main/docs/precision.md) for a complete example, NumPy dtype constraints, and comparison against an fp32 engine. Use different engine names or graph versions for different precisions.
|
|
117
118
|
|
|
118
119
|
## Save and reopen an engine
|
|
119
120
|
|
|
@@ -138,7 +139,7 @@ handle = open_engine(
|
|
|
138
139
|
engine = handle.engine # Keep handle.runtime alive while using the engine.
|
|
139
140
|
```
|
|
140
141
|
|
|
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,
|
|
142
|
+
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
143
|
|
|
143
144
|
`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
145
|
|
|
@@ -265,6 +266,18 @@ Install the producer extra when you need to compile cubins from your own kernels
|
|
|
265
266
|
pip install 'tensoracle[producer]'
|
|
266
267
|
```
|
|
267
268
|
|
|
269
|
+
Every repository that adopts tensoracle must use this top-level layout:
|
|
270
|
+
|
|
271
|
+
```text
|
|
272
|
+
cpp/
|
|
273
|
+
CMakeLists.txt
|
|
274
|
+
python/
|
|
275
|
+
pyproject.toml
|
|
276
|
+
src/
|
|
277
|
+
```
|
|
278
|
+
|
|
279
|
+
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.
|
|
280
|
+
|
|
268
281
|
`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
282
|
|
|
270
283
|
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,11 +25,11 @@ 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
|
|
|
32
|
-
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate
|
|
32
|
+
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
|
|
33
33
|
|
|
34
34
|
## Convert a PyTorch checkpoint
|
|
35
35
|
|
|
@@ -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,9 @@ 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`.
|
|
86
|
+
|
|
87
|
+
The same API builds bf16, int8, and fp8 graphs. Cast tensors to `trt.bfloat16` for bf16; for int8/fp8, place `add_quantize` / `add_dequantize` layers and load their scales from safetensors. The application determines scales with an external quantization library and chooses output quality limits with `compare_outputs`. fp8 targets sm89 and sm120; bf16 and int8 also target sm86. Follow [Build bf16, int8, and fp8 engines](https://github.com/sync-dev-org/tensoracle/blob/main/docs/precision.md) for a complete example, NumPy dtype constraints, and comparison against an fp32 engine. Use different engine names or graph versions for different precisions.
|
|
87
88
|
|
|
88
89
|
## Save and reopen an engine
|
|
89
90
|
|
|
@@ -108,7 +109,7 @@ handle = open_engine(
|
|
|
108
109
|
engine = handle.engine # Keep handle.runtime alive while using the engine.
|
|
109
110
|
```
|
|
110
111
|
|
|
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,
|
|
112
|
+
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
113
|
|
|
113
114
|
`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
115
|
|
|
@@ -235,6 +236,18 @@ Install the producer extra when you need to compile cubins from your own kernels
|
|
|
235
236
|
pip install 'tensoracle[producer]'
|
|
236
237
|
```
|
|
237
238
|
|
|
239
|
+
Every repository that adopts tensoracle must use this top-level layout:
|
|
240
|
+
|
|
241
|
+
```text
|
|
242
|
+
cpp/
|
|
243
|
+
CMakeLists.txt
|
|
244
|
+
python/
|
|
245
|
+
pyproject.toml
|
|
246
|
+
src/
|
|
247
|
+
```
|
|
248
|
+
|
|
249
|
+
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.
|
|
250
|
+
|
|
238
251
|
`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
252
|
|
|
240
253
|
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
|
+
}
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
|
|
3
3
|
`tensoracle` builds TensorRT engines from model graphs supplied through TensorRT's Python network API. A model package owns its graph, safetensors weights, and output quality limits. `tensoracle` supplies engine building, saved engine reuse, execution, output comparison, and tools for cubin plugins. These workflows do not require torch. Only the optional `.pth` checkpoint converter requires torch, installed with `tensoracle[convert]`.
|
|
4
4
|
|
|
5
|
-
Follow the [four-stage optimization workflow](optimization.md) to establish output quality, evaluate
|
|
5
|
+
Follow the [four-stage optimization workflow](optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
|
|
6
6
|
|
|
7
7
|
## Terms
|
|
8
8
|
|
|
@@ -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), 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.
|
|
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,15 +4,15 @@ 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, bf16, int8, or fp8 TensorRT engine using the same `EngineBuild` declaration for input shapes and builder settings. The graph function determines precision through its tensor types. For int8 and fp8, the graph places explicit quantize/dequantize layers and reads their scales from safetensors; the application prepares those scales with an external quantization library. Follow [Build bf16, int8, and fp8 engines](precision.md) for a complete graph example and the GPU targets (fp8 requires sm89 or sm120). When saving with `open_engine`, use a different `name` or `graph_version` from the fp32 engine. Use `compare_outputs` with each lower precision 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
|
-
- **Gate:** Accept the
|
|
15
|
+
- **Gate:** Accept the lower precision 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
|
|
|
17
17
|
## 3. Replace a remaining bottleneck with a plugin
|
|
18
18
|
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
# Build bf16, int8, and fp8 engines
|
|
2
|
+
|
|
3
|
+
The graph function selects tensor types. Use the same `EngineBuild`, `open_engine`, and execution API for fp32, fp16, bf16, int8, and fp8. For bf16, cast tensors to `trt.bfloat16`. For int8 and fp8, place explicit quantize/dequantize (Q/DQ) layers and their scales in the graph. TensorRT 11's [`add_quantize` and `add_dequantize` operators](https://docs.nvidia.com/deeplearning/tensorrt/11.4.0/_static/operators/Quantize.html) accept the quantized type and dequantized type explicitly.
|
|
4
|
+
|
|
5
|
+
| Graph precision | Tensor placement in this example | Supported tensoracle GPU targets |
|
|
6
|
+
| --- | --- | --- |
|
|
7
|
+
| fp32 | fp32 matrix multiplication and bias, with TF32 disabled | sm86, sm89, sm120 |
|
|
8
|
+
| bf16 | Cast features, matrix, and bias to bf16; cast the result to fp32 | sm86, sm89, sm120 |
|
|
9
|
+
| int8 | Q/DQ the features and matrix with `trt.int8` and explicit scales | sm86, sm89, sm120 |
|
|
10
|
+
| fp8 | Q/DQ the features and matrix with `trt.fp8` (E4M3) and explicit scales | sm89, sm120 |
|
|
11
|
+
|
|
12
|
+
Do not build the fp8 graph on sm86. The graph keeps fp32 input and output buffers for every precision; the internal tensors determine its arithmetic. TensorRT chooses implementations for the typed graph, so an int8 or fp8 graph can retain higher precision operations such as the bias addition. This example does not establish a speed improvement.
|
|
13
|
+
|
|
14
|
+
## Prepare weights and scales outside tensoracle
|
|
15
|
+
|
|
16
|
+
For a real model, use an external quantization library, such as [NVIDIA Model Optimizer](https://github.com/NVIDIA/Model-Optimizer), to determine scales from representative data or quantization-aware training. The model package translates those results into its safetensors file and graph. tensoracle does not calculate scales, calibrate a model, or translate a quantizer's export format.
|
|
17
|
+
|
|
18
|
+
Store the fp32 weights and positive fp32 scales together in safetensors. In this example, `matrix` has shape `(32, 16)`, `bias` has shape `(1, 16)`, and each scale has shape `(1,)`. The keys `int8.input_scale`, `int8.weight_scale`, `fp8.input_scale`, and `fp8.weight_scale` distinguish the two quantization recipes. Both Q and DQ use the same scale for each tensor.
|
|
19
|
+
|
|
20
|
+
The following small fixture uses maximum absolute values to choose scales. Its inputs and weights match the GPU test; this scale calculation is only a demonstration for that fixture, not a calibration recipe for other models:
|
|
21
|
+
|
|
22
|
+
```python
|
|
23
|
+
from pathlib import Path
|
|
24
|
+
|
|
25
|
+
import numpy as np
|
|
26
|
+
from safetensors.numpy import save_file
|
|
27
|
+
|
|
28
|
+
path = Path("linear.safetensors")
|
|
29
|
+
features = np.random.default_rng(11).uniform(-1, 1, (16, 32)).astype(np.float32)
|
|
30
|
+
rng = np.random.default_rng(29)
|
|
31
|
+
matrix = rng.uniform(-0.25, 0.25, (32, 16)).astype(np.float32)
|
|
32
|
+
bias = rng.uniform(-0.1, 0.1, (1, 16)).astype(np.float32)
|
|
33
|
+
weights = {"matrix": matrix, "bias": bias}
|
|
34
|
+
for precision, maximum in (("int8", 127), ("fp8", 448)):
|
|
35
|
+
weights[f"{precision}.input_scale"] = np.array(
|
|
36
|
+
[np.max(np.abs(features)) / maximum], dtype=np.float32
|
|
37
|
+
)
|
|
38
|
+
weights[f"{precision}.weight_scale"] = np.array(
|
|
39
|
+
[np.max(np.abs(matrix)) / maximum], dtype=np.float32
|
|
40
|
+
)
|
|
41
|
+
save_file(weights, path)
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
[`safetensors.numpy` 0.8.0](https://github.com/huggingface/safetensors/blob/v0.8.0/bindings/python/py_src/safetensors/numpy.py) with NumPy 2.5 does not directly load bf16 or fp8 tensors into NumPy arrays: bf16 loading raises `TypeError`, and FP8 E4M3 loading raises `AttributeError`. A safetensors file can store those types, but its NumPy reader cannot represent them in this environment. Keep this example's weights and scales in fp32, then cast or quantize in TensorRT. If an existing checkpoint stores bf16 or fp8 weights, convert them with a reader that supports that dtype outside this NumPy graph path. Installing torch is not required for this example.
|
|
45
|
+
|
|
46
|
+
## Define the typed graph
|
|
47
|
+
|
|
48
|
+
This graph computes `features @ matrix + bias`. The bf16 branch casts all three operands before arithmetic and returns fp32. The int8/fp8 branch applies per-tensor Q/DQ to the features and matrix; its bias stays fp32. Keeping the stored matrix in fp32 lets TensorRT quantize it with the supplied scale during building.
|
|
49
|
+
|
|
50
|
+
```python
|
|
51
|
+
import tensorrt as trt
|
|
52
|
+
from safetensors.numpy import load_file
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def graph(path, precision):
|
|
56
|
+
def define(network):
|
|
57
|
+
weights = load_file(path)
|
|
58
|
+
|
|
59
|
+
def constant(array):
|
|
60
|
+
array = np.ascontiguousarray(array, dtype=np.float32)
|
|
61
|
+
return network.add_constant(array.shape, array).get_output(0)
|
|
62
|
+
|
|
63
|
+
features = network.add_input("features", trt.float32, (16, 32))
|
|
64
|
+
matrix = constant(weights["matrix"])
|
|
65
|
+
bias = constant(weights["bias"])
|
|
66
|
+
if precision == "bf16":
|
|
67
|
+
features = network.add_cast(features, trt.bfloat16).get_output(0)
|
|
68
|
+
matrix = network.add_cast(matrix, trt.bfloat16).get_output(0)
|
|
69
|
+
bias = network.add_cast(bias, trt.bfloat16).get_output(0)
|
|
70
|
+
elif precision in ("int8", "fp8"):
|
|
71
|
+
quantized_type = trt.int8 if precision == "int8" else trt.fp8
|
|
72
|
+
|
|
73
|
+
def qdq(tensor, key):
|
|
74
|
+
scale = constant(weights[f"{precision}.{key}_scale"])
|
|
75
|
+
quantized = network.add_quantize(tensor, scale, quantized_type).get_output(0)
|
|
76
|
+
return network.add_dequantize(quantized, scale, trt.float32).get_output(0)
|
|
77
|
+
|
|
78
|
+
features = qdq(features, "input")
|
|
79
|
+
matrix = qdq(matrix, "weight")
|
|
80
|
+
value = network.add_matrix_multiply(
|
|
81
|
+
features, trt.MatrixOperation.NONE, matrix, trt.MatrixOperation.NONE
|
|
82
|
+
).get_output(0)
|
|
83
|
+
value = network.add_elementwise(value, bias, trt.ElementWiseOperation.SUM).get_output(0)
|
|
84
|
+
if precision == "bf16":
|
|
85
|
+
value = network.add_cast(value, trt.float32).get_output(0)
|
|
86
|
+
value.name = "result"
|
|
87
|
+
network.mark_output(value)
|
|
88
|
+
|
|
89
|
+
return define
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
## Build, save, and reuse each precision
|
|
93
|
+
|
|
94
|
+
Use a distinct engine name or graph version whenever tensor types or the graph recipe change. The same weight file serves all four graphs here, so its digest cannot distinguish their precision. Scales are part of that file and therefore part of the saved engine identity.
|
|
95
|
+
|
|
96
|
+
```python
|
|
97
|
+
from tensoracle import (
|
|
98
|
+
BuilderSettings, EngineBuild, EngineIO, InputShape, TensorSpec, open_engine,
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
declaration = EngineBuild(
|
|
102
|
+
{"features": InputShape.fixed((16, 32))},
|
|
103
|
+
BuilderSettings(optimization_level=0),
|
|
104
|
+
)
|
|
105
|
+
io = EngineIO(
|
|
106
|
+
{"features": TensorSpec(trt.float32, (16, 32))},
|
|
107
|
+
{"result": TensorSpec(trt.float32, (16, 16))},
|
|
108
|
+
)
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
def open_linear(precision):
|
|
112
|
+
return open_engine(
|
|
113
|
+
graph(path, precision),
|
|
114
|
+
declaration,
|
|
115
|
+
name=f"linear-{precision}",
|
|
116
|
+
graph_version="1",
|
|
117
|
+
weights=(path,),
|
|
118
|
+
io=io,
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
reference_handle = open_linear("fp32")
|
|
123
|
+
candidate_handle = open_linear("bf16") # Or "int8"; "fp8" on sm89/sm120.
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
TF32 remains disabled by tensoracle, including for the fp32 reference. Engines are saved under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle`. Repeat the same call with unchanged build conditions in another process to reuse the plan without invoking the graph function. Keep the weight file and the engine names/versions stable across the restart. Follow [Execute an engine](index.md#execute-an-engine) to allocate fp32 input/output buffers and enqueue each handle with the same inputs.
|
|
127
|
+
|
|
128
|
+
## Compare outputs and choose limits
|
|
129
|
+
|
|
130
|
+
First establish the fp32 engine against fixed fp32 outputs from the official model implementation, as described in [stage 1](optimization.md#1-establish-the-fp32-reference). Then run the candidate and accepted fp32 engine with identical inputs. Pass their CPU NumPy outputs to `compare_outputs` and choose quality limits in the application:
|
|
131
|
+
|
|
132
|
+
```python
|
|
133
|
+
from tensoracle import compare_outputs
|
|
134
|
+
|
|
135
|
+
# candidate_output and fp32_output are the application's engine outputs.
|
|
136
|
+
comparison = compare_outputs(candidate_output, fp32_output, atol=0.01, rtol=0)
|
|
137
|
+
accepted = (
|
|
138
|
+
comparison.candidate_nonfinite_count == 0
|
|
139
|
+
and comparison.reference_nonfinite_count == 0
|
|
140
|
+
and comparison.outside_tolerance_count == 0
|
|
141
|
+
and comparison.relative_l2_error is not None
|
|
142
|
+
and comparison.relative_l2_error < 0.01
|
|
143
|
+
)
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
These example limits belong to the small bf16 fixture. For the fixture above, a TensorRT 11.4 run on sm120 measured:
|
|
147
|
+
|
|
148
|
+
| Candidate | Maximum absolute error | Relative L2 error | Test `atol` (`rtol=0`) | Test relative L2 limit |
|
|
149
|
+
| --- | --- | --- | --- | --- |
|
|
150
|
+
| bf16 | 0.00357 | 0.00294 | 0.01 | 0.01 |
|
|
151
|
+
| int8 | 0.00710 | 0.00572 | 0.02 | 0.02 |
|
|
152
|
+
| fp8 | 0.0501 | 0.0382 | 0.12 | 0.08 |
|
|
153
|
+
|
|
154
|
+
The test allows at least twice the observed errors for differences in GPU tactics. These are not tensoracle defaults or quality guarantees for other models. Choose representative data, per-output limits, and model-specific quality measurements for the real application. The GPU test also checks that the two engines occupy separate cache entries and reproduce their outputs in another process without calling either graph.
|
|
@@ -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.
|
|
@@ -17,6 +17,8 @@ plugins:
|
|
|
17
17
|
handlers:
|
|
18
18
|
python:
|
|
19
19
|
paths:
|
|
20
|
+
- python/src
|
|
21
|
+
# Standalone source archives keep src next to pyproject.toml.
|
|
20
22
|
- src
|
|
21
23
|
options:
|
|
22
24
|
docstring_style: google
|
|
@@ -31,6 +33,7 @@ plugins:
|
|
|
31
33
|
nav:
|
|
32
34
|
- Home: index.md
|
|
33
35
|
- Optimization workflow: optimization.md
|
|
36
|
+
- bf16, int8, and fp8 engines: precision.md
|
|
34
37
|
- TensorRT SDK setup: tensorrt-sdk.md
|
|
35
38
|
- API reference:
|
|
36
39
|
- Engine and comparison: api/index.md
|