tensoracle 0.5.0__tar.gz → 0.5.1__tar.gz

{tensoracle-0.5.0 → tensoracle-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: tensoracle
-Version: 0.5.0
+Version: 0.5.1
 Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
 Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
 Project-URL: repository, https://github.com/sync-dev-org/tensoracle
@@ -59,7 +59,7 @@ Choose the path that matches your work:
 
 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.
 
-Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate fp16 and plugins, and measure speed.
+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.
 
 ## Convert a PyTorch checkpoint
 
@@ -114,6 +114,8 @@ Use `InputShape.fixed((1, 64, 2))` for one input size; the graph may declare tha
 
 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`.
 
+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.
+
 ## Save and reopen an engine
 
 `open_engine` builds on the first call and reuses the saved engine on later calls, including calls from another process. It invokes `graph` only when building. Declare the expected input and output names, TensorRT types, and build-time shapes:

{tensoracle-0.5.0 → tensoracle-0.5.1}/README.md

@@ -29,7 +29,7 @@ Choose the path that matches your work:
 
 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.
 
-Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate fp16 and plugins, and measure speed.
+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.
 
 ## Convert a PyTorch checkpoint
 
@@ -84,6 +84,8 @@ Use `InputShape.fixed((1, 64, 2))` for one input size; the graph may declare tha
 
 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`.
 
+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.
+
 ## Save and reopen an engine
 
 `open_engine` builds on the first call and reuses the saved engine on later calls, including calls from another process. It invokes `graph` only when building. Declare the expected input and output names, TensorRT types, and build-time shapes:

{tensoracle-0.5.0 → tensoracle-0.5.1}/docs/index.md

@@ -2,7 +2,7 @@
 
 `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]`.
 
-Follow the [four-stage optimization workflow](optimization.md) to establish output quality, evaluate fp16 and plugins, and measure speed.
+Follow the [four-stage optimization workflow](optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
 
 ## Terms
 
@@ -161,7 +161,7 @@ Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutput
 
 ### Build a standard graph
 
-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.
+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.
 
 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).
 

{tensoracle-0.5.0 → tensoracle-0.5.1}/docs/optimization.md

@@ -10,9 +10,9 @@ Use the same representative inputs, shapes, and corresponding outputs throughout
 
 ## 2. Accept lower precision by quality
 
-- **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.
+- **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.
 - **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.
-- **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.
+- **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.
 
 ## 3. Replace a remaining bottleneck with a plugin
 

tensoracle-0.5.1/docs/precision.md

@@ -0,0 +1,154 @@
+# Build bf16, int8, and fp8 engines
+
+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.
+
+| Graph precision | Tensor placement in this example | Supported tensoracle GPU targets |
+| --- | --- | --- |
+| fp32 | fp32 matrix multiplication and bias, with TF32 disabled | sm86, sm89, sm120 |
+| bf16 | Cast features, matrix, and bias to bf16; cast the result to fp32 | sm86, sm89, sm120 |
+| int8 | Q/DQ the features and matrix with `trt.int8` and explicit scales | sm86, sm89, sm120 |
+| fp8 | Q/DQ the features and matrix with `trt.fp8` (E4M3) and explicit scales | sm89, sm120 |
+
+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.
+
+## Prepare weights and scales outside tensoracle
+
+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.
+
+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.
+
+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:
+
+```python
+from pathlib import Path
+
+import numpy as np
+from safetensors.numpy import save_file
+
+path = Path("linear.safetensors")
+features = np.random.default_rng(11).uniform(-1, 1, (16, 32)).astype(np.float32)
+rng = np.random.default_rng(29)
+matrix = rng.uniform(-0.25, 0.25, (32, 16)).astype(np.float32)
+bias = rng.uniform(-0.1, 0.1, (1, 16)).astype(np.float32)
+weights = {"matrix": matrix, "bias": bias}
+for precision, maximum in (("int8", 127), ("fp8", 448)):
+    weights[f"{precision}.input_scale"] = np.array(
+        [np.max(np.abs(features)) / maximum], dtype=np.float32
+    )
+    weights[f"{precision}.weight_scale"] = np.array(
+        [np.max(np.abs(matrix)) / maximum], dtype=np.float32
+    )
+save_file(weights, path)
+```
+
+[`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.
+
+## Define the typed graph
+
+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.
+
+```python
+import tensorrt as trt
+from safetensors.numpy import load_file
+
+
+def graph(path, precision):
+    def define(network):
+        weights = load_file(path)
+
+        def constant(array):
+            array = np.ascontiguousarray(array, dtype=np.float32)
+            return network.add_constant(array.shape, array).get_output(0)
+
+        features = network.add_input("features", trt.float32, (16, 32))
+        matrix = constant(weights["matrix"])
+        bias = constant(weights["bias"])
+        if precision == "bf16":
+            features = network.add_cast(features, trt.bfloat16).get_output(0)
+            matrix = network.add_cast(matrix, trt.bfloat16).get_output(0)
+            bias = network.add_cast(bias, trt.bfloat16).get_output(0)
+        elif precision in ("int8", "fp8"):
+            quantized_type = trt.int8 if precision == "int8" else trt.fp8
+
+            def qdq(tensor, key):
+                scale = constant(weights[f"{precision}.{key}_scale"])
+                quantized = network.add_quantize(tensor, scale, quantized_type).get_output(0)
+                return network.add_dequantize(quantized, scale, trt.float32).get_output(0)
+
+            features = qdq(features, "input")
+            matrix = qdq(matrix, "weight")
+        value = network.add_matrix_multiply(
+            features, trt.MatrixOperation.NONE, matrix, trt.MatrixOperation.NONE
+        ).get_output(0)
+        value = network.add_elementwise(value, bias, trt.ElementWiseOperation.SUM).get_output(0)
+        if precision == "bf16":
+            value = network.add_cast(value, trt.float32).get_output(0)
+        value.name = "result"
+        network.mark_output(value)
+
+    return define
+```
+
+## Build, save, and reuse each precision
+
+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.
+
+```python
+from tensoracle import (
+    BuilderSettings, EngineBuild, EngineIO, InputShape, TensorSpec, open_engine,
+)
+
+declaration = EngineBuild(
+    {"features": InputShape.fixed((16, 32))},
+    BuilderSettings(optimization_level=0),
+)
+io = EngineIO(
+    {"features": TensorSpec(trt.float32, (16, 32))},
+    {"result": TensorSpec(trt.float32, (16, 16))},
+)
+
+
+def open_linear(precision):
+    return open_engine(
+        graph(path, precision),
+        declaration,
+        name=f"linear-{precision}",
+        graph_version="1",
+        weights=(path,),
+        io=io,
+    )
+
+
+reference_handle = open_linear("fp32")
+candidate_handle = open_linear("bf16")  # Or "int8"; "fp8" on sm89/sm120.
+```
+
+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.
+
+## Compare outputs and choose limits
+
+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:
+
+```python
+from tensoracle import compare_outputs
+
+# candidate_output and fp32_output are the application's engine outputs.
+comparison = compare_outputs(candidate_output, fp32_output, atol=0.01, rtol=0)
+accepted = (
+    comparison.candidate_nonfinite_count == 0
+    and comparison.reference_nonfinite_count == 0
+    and comparison.outside_tolerance_count == 0
+    and comparison.relative_l2_error is not None
+    and comparison.relative_l2_error < 0.01
+)
+```
+
+These example limits belong to the small bf16 fixture. For the fixture above, a TensorRT 11.4 run on sm120 measured:
+
+| Candidate | Maximum absolute error | Relative L2 error | Test `atol` (`rtol=0`) | Test relative L2 limit |
+| --- | --- | --- | --- | --- |
+| bf16 | 0.00357 | 0.00294 | 0.01 | 0.01 |
+| int8 | 0.00710 | 0.00572 | 0.02 | 0.02 |
+| fp8 | 0.0501 | 0.0382 | 0.12 | 0.08 |
+
+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.

{tensoracle-0.5.0 → tensoracle-0.5.1}/mkdocs.yml

@@ -33,6 +33,7 @@ plugins:
 nav:
   - Home: index.md
   - Optimization workflow: optimization.md
+  - bf16, int8, and fp8 engines: precision.md
   - TensorRT SDK setup: tensorrt-sdk.md
   - API reference:
       - Engine and comparison: api/index.md

{tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_cache.py

@@ -236,7 +236,8 @@ def _write_weight_record(path: Path, record: dict) -> None:
                 pass
 
 
-def _gpu_name() -> str:
+def _gpu_identity() -> dict[str, str | int]:
+    """Describe device compatibility without binding the cache to a physical GPU."""
     library_name = ctypes.util.find_library("cuda") or ctypes.util.find_library("nvcuda")
     if library_name is None:
         raise RuntimeError("CUDA driver library was not found")
@@ -245,6 +246,8 @@ def _gpu_name() -> str:
     cuda.cuCtxGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)]
     cuda.cuDeviceGet.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int]
     cuda.cuDeviceGetName.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.c_int]
+    cuda.cuDeviceGetAttribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int]
+    cuda.cuDeviceTotalMem_v2.argtypes = [ctypes.POINTER(ctypes.c_size_t), ctypes.c_int]
     if cuda.cuInit(0) != 0:
         raise RuntimeError("CUDA driver initialization failed")
     device = ctypes.c_int()
@@ -256,7 +259,33 @@ def _gpu_name() -> str:
     name = ctypes.create_string_buffer(256)
     if cuda.cuDeviceGetName(name, len(name), device.value) != 0:
         raise RuntimeError("CUDA device name lookup failed")
-    return name.value.decode("utf-8")
+    properties: dict[str, str | int] = {"name": name.value.decode("utf-8")}
+    # TensorRT's documented compatibility checks, plus ECC and device memory:
+    # https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/engine-compatibility.html#compatibility-checks
+    # ECC can distinguish otherwise identical devices with different usable memory.
+    # Attribute numbers are CUdevice_attribute values from the CUDA driver API.
+    for field, attribute in (
+        ("compute_capability_major", 75),
+        ("compute_capability_minor", 76),
+        ("global_memory_bus_width", 37),
+        ("l2_cache_size", 38),
+        ("max_shared_memory_per_block", 8),
+        ("max_shared_memory_per_block_optin", 97),
+        ("max_shared_memory_per_multiprocessor", 81),
+        ("texture_alignment", 14),
+        ("multiprocessor_count", 16),
+        ("integrated", 18),
+        ("ecc_enabled", 32),
+    ):
+        value = ctypes.c_int()
+        if cuda.cuDeviceGetAttribute(ctypes.byref(value), attribute, device.value) != 0:
+            raise RuntimeError(f"CUDA device attribute lookup failed: {field}")
+        properties[field] = value.value
+    total_memory = ctypes.c_size_t()
+    if cuda.cuDeviceTotalMem_v2(ctypes.byref(total_memory), device.value) != 0:
+        raise RuntimeError("CUDA device total memory lookup failed")
+    properties["total_memory"] = total_memory.value
+    return properties
 
 
 def _conditions(
@@ -286,7 +315,7 @@ def _conditions(
             for name, shape in declaration.inputs.items()
         },
         "builder_settings": asdict(declaration.settings),
-        "gpu_name": _gpu_name(),
+        "gpu": _gpu_identity(),
         "os": platform.system(),
         "architecture": platform.machine(),
         "tensorrt_version": trt.__version__,

{tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.0'
-__version_tuple__ = version_tuple = (0, 5, 0)
+__version__ = version = '0.5.1'
+__version_tuple__ = version_tuple = (0, 5, 1)
 
 __commit_id__ = commit_id = None

tensoracle-0.5.1/tests/precision_runtime.py

@@ -0,0 +1,164 @@
+"""Build and execute the precision guide's linear graph, including after a restart."""
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import runtime as cudart
+from safetensors.numpy import load_file, save_file
+
+from tensoracle import (
+    BuilderSettings,
+    EngineBuild,
+    EngineIO,
+    InputShape,
+    TensorSpec,
+    open_engine,
+)
+
+
+def sample_input():
+    return np.random.default_rng(11).uniform(-1, 1, (16, 32)).astype(np.float32)
+
+
+def write_weights(path):
+    rng = np.random.default_rng(29)
+    matrix = rng.uniform(-0.25, 0.25, (32, 16)).astype(np.float32)
+    bias = rng.uniform(-0.1, 0.1, (1, 16)).astype(np.float32)
+    weights = {"matrix": matrix, "bias": bias}
+    # A small test fixture, not a model calibration algorithm or a package API.
+    for precision, maximum in (("int8", 127), ("fp8", 448)):
+        weights[f"{precision}.input_scale"] = np.array(
+            [np.max(np.abs(sample_input())) / maximum], dtype=np.float32
+        )
+        weights[f"{precision}.weight_scale"] = np.array(
+            [np.max(np.abs(matrix)) / maximum], dtype=np.float32
+        )
+    save_file(weights, path)
+
+
+def graph(path, precision):
+    def define(network):
+        weights = load_file(path)
+
+        def constant(array):
+            array = np.ascontiguousarray(array, dtype=np.float32)
+            return network.add_constant(array.shape, array).get_output(0)
+
+        features = network.add_input("features", trt.float32, (16, 32))
+        matrix = constant(weights["matrix"])
+        bias = constant(weights["bias"])
+        if precision == "bf16":
+            features = network.add_cast(features, trt.bfloat16).get_output(0)
+            matrix = network.add_cast(matrix, trt.bfloat16).get_output(0)
+            bias = network.add_cast(bias, trt.bfloat16).get_output(0)
+        elif precision in ("int8", "fp8"):
+            quantized_type = trt.int8 if precision == "int8" else trt.fp8
+
+            def qdq(tensor, key):
+                scale = constant(weights[f"{precision}.{key}_scale"])
+                quantized = network.add_quantize(tensor, scale, quantized_type).get_output(0)
+                return network.add_dequantize(quantized, scale, trt.float32).get_output(0)
+
+            features = qdq(features, "input")
+            matrix = qdq(matrix, "weight")
+        value = network.add_matrix_multiply(
+            features, trt.MatrixOperation.NONE, matrix, trt.MatrixOperation.NONE
+        ).get_output(0)
+        value = network.add_elementwise(value, bias, trt.ElementWiseOperation.SUM).get_output(0)
+        if precision == "bf16":
+            value = network.add_cast(value, trt.float32).get_output(0)
+        value.name = "result"
+        network.mark_output(value)
+
+    return define
+
+
+def open_linear(path, precision, *, reuse=False):
+    def forbidden(network):
+        raise AssertionError("a saved precision engine called its graph after a process restart")
+
+    return open_engine(
+        forbidden if reuse else graph(path, precision),
+        EngineBuild(
+            {"features": InputShape.fixed((16, 32))},
+            BuilderSettings(optimization_level=0),
+        ),
+        name=f"linear-{precision}",
+        graph_version="1",
+        weights=(path,),
+        io=EngineIO(
+            {"features": TensorSpec(trt.float32, (16, 32))},
+            {"result": TensorSpec(trt.float32, (16, 16))},
+        ),
+    )
+
+
+def _cuda(result):
+    status, *values = result
+    if status != cudart.cudaError_t.cudaSuccess:
+        raise RuntimeError(f"CUDA runtime returned {status}")
+    return values[0] if values else None
+
+
+def run(handle):
+    features = sample_input()
+    output = np.empty((16, 16), dtype=np.float32)
+    source = target = stream = None
+    with handle.create_execution() as execution:
+        assert execution.get_tensor_shape("result") == output.shape
+        assert execution.get_tensor_dtype("result") == trt.float32
+        try:
+            source = _cuda(cudart.cudaMalloc(features.nbytes))
+            target = _cuda(cudart.cudaMalloc(output.nbytes))
+            stream = _cuda(cudart.cudaStreamCreate())
+            _cuda(
+                cudart.cudaMemcpyAsync(
+                    source,
+                    features.ctypes.data,
+                    features.nbytes,
+                    cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
+                    stream,
+                )
+            )
+            execution.enqueue({"features": int(source), "result": int(target)}, stream=int(stream))
+            _cuda(
+                cudart.cudaMemcpyAsync(
+                    output.ctypes.data,
+                    target,
+                    output.nbytes,
+                    cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
+                    stream,
+                )
+            )
+        finally:
+            if stream is not None:
+                _cuda(cudart.cudaStreamSynchronize(stream))
+                _cuda(cudart.cudaStreamDestroy(stream))
+            if target is not None:
+                _cuda(cudart.cudaFree(target))
+            if source is not None:
+                _cuda(cudart.cudaFree(source))
+    return output
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("weights", type=Path)
+    parser.add_argument("precision", choices=("bf16", "int8", "fp8"))
+    parser.add_argument("outputs", type=Path)
+    args = parser.parse_args()
+    outputs = {}
+    paths = {}
+    for precision in ("fp32", args.precision):
+        handle = open_linear(args.weights, precision, reuse=True)
+        paths[precision] = str(handle.path)
+        outputs[precision] = run(handle)
+    np.savez(args.outputs, **outputs)
+    print(json.dumps(paths))
+
+
+if __name__ == "__main__":
+    main()

{tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_cache.py

@@ -95,7 +95,7 @@ def test_save_large_built_plan_keeps_other_threads_running(tmp_path, monkeypatch
 
     plan = b"\0" * (512 * 1024 * 1024)
     monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
-    monkeypatch.setattr(cache, "_gpu_name", lambda: "GPU A")
+    monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
     builder = SimpleNamespace(
         create_network=lambda flags: SimpleNamespace(num_inputs=0, num_outputs=0),
         create_builder_config=lambda: SimpleNamespace(
@@ -234,7 +234,7 @@ def test_saved_engine_deserialization_failure_guides_caller_to_check_plugins(tmp
 
     cache_root = tmp_path / "cache"
     monkeypatch.setenv("XDG_CACHE_HOME", str(cache_root))
-    monkeypatch.setattr(cache, "_gpu_name", lambda: "GPU A")
+    monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
     weights = tmp_path / "weights.safetensors"
     _weights(weights)
     declaration = tensoracle.EngineBuild({"features": tensoracle.InputShape.fixed((1, 3, 2))})
@@ -414,7 +414,7 @@ def test_build_failure_removes_placeholder_and_partial_files(tmp_path, monkeypat
     import tensoracle._cache as cache
 
     monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
-    monkeypatch.setattr(cache, "_gpu_name", lambda: "GPU A")
+    monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
     monkeypatch.setattr(
         cache, "_build_serialized_engine", lambda graph, declaration, plugins, logger: graph(None)
     )
@@ -474,7 +474,7 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
     import tensoracle._cache as cache
 
     monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
-    monkeypatch.setattr(cache, "_gpu_name", lambda: "GPU A")
+    monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
     weights = tmp_path / "weights.safetensors"
     _weights(weights)
     same_contents = tmp_path / "same.safetensors"
@@ -627,7 +627,7 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
     ]
     paths = {open_case(**variant).path for variant in variants}
     with monkeypatch.context() as patch:
-        patch.setattr(cache, "_gpu_name", lambda: "GPU B")
+        patch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU B"})
         paths.add(open_case().path)
     with monkeypatch.context() as patch:
         patch.setattr(cache.platform, "system", lambda: "Other OS")

{tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_distribution.py

@@ -130,30 +130,53 @@ def test_public_snapshot_includes_only_the_pages_workflow(publication_config, tm
     assert paths - {PAGES_WORKFLOW} <= excluded
 
 
+@pytest.mark.req("SP-TENSORACLE-024")
 def test_pages_workflow_builds_the_public_snapshot(publication_config):
-    """Only the public repository builds locked docs dependencies and deploys the site."""
+    """Only the public repository builds and deploys docs with limited permissions and no secrets."""
     workflow_path = ROOT / PAGES_WORKFLOW
     assert workflow_path.is_file()
-    workflow = yaml.safe_load(workflow_path.read_text(encoding="utf-8"))
+    workflow_source = workflow_path.read_text(encoding="utf-8")
+    assert "secrets." not in workflow_source
+    workflow = yaml.safe_load(workflow_source)
     triggers = workflow.get("on", workflow.get(True))
     assert triggers == {"push": {"branches": ["main"]}, "workflow_dispatch": None}
 
+    declared_permissions = set()
+    for scope in (workflow, *workflow["jobs"].values()):
+        permissions = scope.get("permissions", {})
+        assert isinstance(permissions, dict)
+        declared_permissions.update(permissions.items())
+    assert declared_permissions == {
+        ("contents", "read"),
+        ("pages", "write"),
+        ("id-token", "write"),
+    }
+
     steps = []
     for job in workflow["jobs"].values():
         assert job["if"] == "github.repository == 'sync-dev-org/tensoracle'"
         steps.extend(job["steps"])
     commands = [step["run"] for step in steps if "run" in step]
-    assert "uv sync --locked --only-group docs" in commands
+    assert commands == [
+        "uv sync --locked --only-group docs",
+        "uv run --no-sync mkdocs build --strict --config-file ../mkdocs.yml",
+    ]
     for job in workflow["jobs"].values():
         if any("run" in step for step in job["steps"]):
             assert job["defaults"]["run"]["working-directory"] == "python"
-    assert "uv run --no-sync mkdocs build --strict --config-file ../mkdocs.yml" in commands
     actions = {}
     for step in steps:
         if "uses" in step:
             action, version = step["uses"].split("@")
-            assert re.fullmatch(r"[0-9a-f]{40}|v\d+\.\d+\.\d+", version)
+            assert re.fullmatch(r"[0-9a-f]{40}", version)
             actions[action] = step
+    assert actions.keys() <= {
+        "actions/checkout",
+        "astral-sh/setup-uv",
+        "actions/configure-pages",
+        "actions/upload-pages-artifact",
+        "actions/deploy-pages",
+    }
     assert {
         "actions/configure-pages",
         "actions/upload-pages-artifact",
@@ -163,8 +186,9 @@ def test_pages_workflow_builds_the_public_snapshot(publication_config):
     deployment = next(job for job in workflow["jobs"].values() if "environment" in job)
     assert deployment["needs"] == "build"
     assert deployment["environment"]["name"] == "github-pages"
-    assert deployment["permissions"]["pages"] == "write"
-    assert deployment["permissions"]["id-token"] == "write"
+    deployment_permissions = deployment.get("permissions", workflow["permissions"])
+    assert deployment_permissions["pages"] == "write"
+    assert deployment_permissions["id-token"] == "write"
 
 
 def test_docs_dependencies_are_separate_from_gpu_development_packages(publication_config):

{tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_engine_io.py

@@ -50,7 +50,7 @@ def _cached_open(tmp_path, monkeypatch, engine, declaration, io):
     import tensoracle._cache as cache
 
     monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
-    monkeypatch.setattr(cache, "_gpu_name", lambda: "GPU A")
+    monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
     weights = tmp_path / "weights.safetensors"
     _weights(weights)
     conditions = cache._conditions("checked", "1", (weights,), (), declaration)