Skip to main content

TensorTorrent logo

TensorTorrent

A heterogeneous PyTorch compiler and runtime for one machine with many CPUs, GPUs, and memory tiers.

CI status Latest version tag Python 3.10 to 3.13 Rust 1.85 or newer Apache-2.0 license

TensorTorrent exports a PyTorch model, partitions its graph, places regions across available compute, and runs the resulting schedule through a Rust data plane. Parameters can stream from slower storage and activations can spill when the model exceeds device or host memory.

Python compiles. Rust schedules. One immutable ExecutableArtifact describes the program.

[!IMPORTANT] Supported target: Linux with Python 3.10–3.13 and PyTorch 2.4 or newer. Validate every deployment machine with tensortorrent validate-hardware before serving production traffic.

Installation

pip install torch --index-url https://download.pytorch.org/whl/cpu   # or CUDA/ROCm from pytorch.org
pip install tensortorrent

Linux, Python 3.10–3.13, PyTorch ≥2.4. Install torch first if you need a specific build. Wheels: PyPI, Releases. Dev from source: uv + Rust 1.85+ (Quick start).

Quick start

git clone https://github.com/alhussein-jamil/TensorTorrent.git
cd TensorTorrent
make sync
make doctor

Compile a module and compare it with eager PyTorch:

import torch
import torch.nn as nn
import tensortorrent as tt  # import alias: tt

model = nn.Sequential(
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
).eval()
x = torch.randn(32, 256)

compiled = tt.compile(model, example_inputs=(x,))
torch.testing.assert_close(compiled(x), model(x), check_device=False)

compiled.save("artifact/")
reloaded = tt.load_compiled("artifact/")

Run uv run python examples/public_api_demo.py for hardware discovery, compile, and schedule output in one executable example.

What it handles

Area Implementation
PyTorch export and graph partitioning python/tensortorrent/compile
CPU, CUDA, ROCm, Intel XPU, and plugin discovery python/tensortorrent/backends
Native placement planner (Rayon subset + beam search) crates/tt-planner
Discrete-event schedule simulation (batch finalist ranking) crates/tt-runtime/simulator
Resource budget resolver (host memory, VRAM, CPU, disk) python/tensortorrent/hardware/budget.py
NUMA-aware host allocation and CPU budget enforcement crates/tt-backend-cpu
Scheduling, residency, transfer, stall watchdog, and cancellation crates/tt-runtime
Parameter streaming and activation spill crates/tt-storage
Atomic, checksummed artifact bundles python/tensortorrent/artifact_io.py
Concurrent request serving (HTTP, auth, metrics) python/tensortorrent/serve
Virtual accelerators for deterministic tests crates/tt-backend-virtual

The runtime supports NCCL, RCCL, oneCCL, Gloo, and explicit host-staged collective fallbacks where the installed hardware and libraries allow them.

Architecture

flowchart LR
    M[PyTorch module] --> E[Export / IR / profile]
    E --> N[Native Rust planner]
    N --> F[Top-K finalist schedules]
    F --> D[Parallel Rust DES]
    D --> W[Compile winner only]
    W --> R[Rust runtime]

TensorTorrent profiles the hardware, searches a large space of heterogeneous execution plans with a native parallel planner (multiple strong placements per competitive device subset), constructs bounded fair schedule variants (prefetch/staging), simulates those finalists with a Rust discrete-event model of compute/transfers/contention/memory, and selects the best feasible strategy before compiling only the winner and executing across CPUs/GPUs/storage. Planner parallelism is automatic and stays serial when the search is too small to benefit (subset-level Rayon for multi-device searches; intra-subset beam Rayon when parent×candidate fanout is large enough). Batch DES likewise stays serial for tiny schedule batches where thread-pool setup would dominate. Not every combinatorial plan is exhaustively simulated — the planner shortlists; DES ranks the strongest finalists.

The Python control plane owns export, normalization, partitioning, region compilation, public APIs, and diagnostics. The Rust data plane owns placement search, discrete-event simulation, the artifact, schedule, workers, residency, transfers, storage, cancellation, and telemetry. Torch compute regions may call back into Python; planning, scheduling, and data movement remain in Rust.

See the architecture guide for ownership boundaries and backend contracts for extension points.

Module composition

Compile a sequence as one graph to avoid opaque transfers between separately compiled artifacts:

compiled = tt.compile_modules(
    [encoder, projector, decoder],
    example_inputs=(x,),
    names=["encoder", "projector", "decoder"],
)

For branches, joins, structured arguments, or nested outputs, build a ModuleGraph from ModuleNode, GraphInput, and NodeOutput. Invalid names, forward references, and output paths are rejected before export.

Opt-in training

Compilation is inference-only by default. Set allow_training=True to use the same heterogeneous schedule with autograd:

config = tt.CompileConfig(allow_training=True)
compiled = tt.compile(model, example_inputs=(x,), config=config)

optimizer = torch.optim.Adam(compiled.parameters())
compiled.train()
optimizer.zero_grad()
loss = compiled(x).sum()
loss.backward()
optimizer.step()
compiled.eval()

Training cannot currently be combined with NVMe parameter streaming, activation spill budgets, or process workers. See the full product scope for intentional limits.

Does it actually work?

On a single device TensorTorrent reaches eager parity at scale — matching or beating PyTorch on large MLPs and transformers. Eligible resident single-region graphs use the direct path by default. Measured resident CPU+accelerator branch plans can use the same low-overhead path after synchronized timing beats both schedule execution and full fusion (prefer_direct_path; override with TT_DIRECT_PATH=0/1). The product focus beyond that is multi-device placement, parameter streaming, and activation spill.

Measured tables and the same-device harness pin live in Benchmarks.

Resource budgets and guardrails

Every memory limit, CPU count, and disk quota flows through a single resolver that reads cgroup v2/v1 limits, live OS availability, and explicit config values — in that precedence order. Containers automatically see their cgroup limits, not host totals. The resolver provenance is shown by tensortorrent doctor.

See Resource budgets and guardrails for the full precedence chain, spill lifecycle, stall watchdog, and worked examples.

Development

make sync                 # create the environment and build the native extension
make check                # lint, types, Rust tests, Python tests, doctor
make audit                # cargo-audit (Rust) + pip-audit (Python)
make coverage             # run tests with coverage gate (Python 3.12)
make native-gate          # native extension smoke and execution checks
make hardware-test        # explicit: may consume most available VRAM or spill space

On a machine with a GPU, run everything that needs real hardware in one go:

bash tools/run_everything.sh     # tests + hardware suite + all benchmarks

It writes logs, JSON, and a SUMMARY.md to bench-results/<timestamp>/. Install the benchmark baselines first with uv sync --extra bench so the ONNX Runtime and Accelerate comparisons run instead of reporting as missing.

CI: PRs and pushes to main (Python 3.10 + 3.13, x86-64/ARM64). Hardware tests are opt-in. See CONTRIBUTING.md.

Repository map

python/tensortorrent/   Python control plane, public API, and serving
crates/tt-*/            Rust IR, runtime, memory, storage, backends, and FFI
tests/                  Unit, integration, end-to-end, property, and hardware tests
docs/                   Product, architecture, deployment, and reference guides
examples/               Small public API programs
bench/                  Runtime and planner comparisons
tools/                  Local quality and native-extension gates
deploy/                 Docker Compose and Kubernetes examples
Dockerfile              CPU-only production container
Dockerfile.cuda         CUDA GPU production container (validate on GPU host before use)

Documentation

Versions and releases

Versions follow Semantic Versioning; tags are vMAJOR.MINOR.PATCH. A tag builds wheels, a GitHub Release, and a PyPI publish — see docs/RELEASING.md.

License

Apache-2.0. See LICENSE.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

tensortorrent-0.2.8.tar.gz (393.9 kB view details)

Uploaded Source

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (1.5 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.17+ ARM64

tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (1.5 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.17+ ARM64

tensortorrent-0.2.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.11manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.10manylinux: glibc 2.17+ x86-64

File details

Details for the file tensortorrent-0.2.8.tar.gz.

File metadata

  • Download URL: tensortorrent-0.2.8.tar.gz
  • Upload date:
  • Size: 393.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for tensortorrent-0.2.8.tar.gz
Algorithm Hash digest
SHA256 acb380f9259a7c7bb506f174816a44dacad3111ab95e0e025c56275625df0694
MD5 50213f970475fbcd14cbc4d343dd19d7
BLAKE2b-256 3e2685944e196a6483f229f68dec1bd8575c7db90b937335e3b8ebb625b740db

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.8.tar.gz:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 f8af9cfca6ed4631f92a733a0419cf6f48f09842f840a9ae40d4be5ac77e7ac4
MD5 da18ae6c1c5294d14e2f41016e1d1362
BLAKE2b-256 94a7b5eecec63738eb29a4e0fbeac590830a79fb0d1669e89eb00e5006c0ec5b

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 d6f494707b85da1f50044bc87f4dd03b6ab17ee134fab7363c35c2740555aa92
MD5 240eaeef17621392b526a1da7130c039
BLAKE2b-256 a1247309ad061a9ec679f8c5229453749f4568bcea18016cf4598b2d77c5d7c6

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.8-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 204f4d793992e4b4c17035abc31ca3215230ec0d8bdced821331603c100b0f91
MD5 4d0a56cf10f70fae9911813ec9eddaf8
BLAKE2b-256 86d5bfcf316c606b0ff964e61faec69b77e732725925fb5ac574263c555b2dfa

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 6ea1c3b85701d2a552c7f8b03c243a153979d62404a61ebc56be77bebde7fabd
MD5 dec2e3052c9fa29b9d9bcb577a620b80
BLAKE2b-256 3a33c442f91dfdc42f79760bee933f977d22b1888a549751b04dbd642c2bf5c5

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.8-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 16231d996fafc97d56aa59ca92ecaebc47ac3bb1ed60feb2f2ec6ede3e298736
MD5 8de99a898ec75b9fda4fd748037030ba
BLAKE2b-256 9c5b7f44bffb33327058da99609e2acd36a3dc1b13f38d9715b8d873cbb0e905

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 e0eea17f2f54f40c2a0b79228868708a203cbc59503dbbb156774864f2a4a8e5
MD5 1ab6c8382c41616131a179b8c4bdfa47
BLAKE2b-256 e0340c75ee6f91cd886395b3319b79663d24a6b9ddf1883143380a04c81ffe1c

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page