Keep GPU arrays on the device with cuda_ipc

By default, a GPU array that crosses into or out of a served Tesseract takes a round trip through host memory. It is copied off the device, encoded into the request or response body, decoded on the other side, and copied back onto the GPU. For large arrays in a loop, such as an optimizer or a training step calling a GPU-resident Tesseract, that round trip can dominate the cost of each call.

This round trip can be avoided with the cuda_ipc GPU transport. Instead of the array bytes, the request or response carries a 64-byte CUDA IPC handle that the other process uses to map the same device memory. Shapes, dtypes, and schema validation stay in the JSON payload as usual.

Each call still pays a fixed cost of a few milliseconds for the HTTP request, validation, and handle exchange. For small arrays, below about 30,000 float32 elements (about 120 kB) in our measurements, the host round trip can be faster (see Performance trade-offs & optimization).

Requirements

  • Linux with an NVIDIA GPU. The transport is built on CUDA IPC.

  • Client and Tesseract on the same machine, seeing the same GPU. IPC handles are meaningless across hosts. The two processes may number the GPU differently, for example when a container sees only one of the host’s GPUs.

  • A shared IPC namespace. When the Tesseract runs in a container, it needs --ipc=host. The SDK adds this for you (see below).

  • GPU arrays on both sides. The client passes arrays that live on a CUDA device, and the Tesseract returns arrays that live on the device. (CPU arrays in the same payload are still encoded according to output_format.)

  • Matching versions. tesseract-core 1.15.0 or later on both sides, since earlier runtimes use a different wire format for cuda_ipc. For JAX and PyTorch, tesseract-jax 0.5.0 and tesseract-torch 0.2.0 or later.

Write a Tesseract that returns GPU arrays

Return outputs that live on the GPU, and adopt inputs with your framework before computing. The CuPy, JAX, and PyTorch examples compute s * a + b for two float32 vectors a and b:

def apply(inputs: InputSchema) -> OutputSchema:
    """Compute ``s * a + b`` on the GPU and return device memory."""
    a = cp.asarray(inputs.a)
    b = cp.asarray(inputs.b)
    result = inputs.s * a + b
    return OutputSchema(result=result)

cp.asarray keeps device arrays on the GPU and moves host arrays onto it.

def apply(inputs: InputSchema) -> OutputSchema:
    """Compute ``s * a + b`` on the GPU and return device memory."""
    a = as_jax_array(inputs.a)
    b = as_jax_array(inputs.b)
    result = inputs.s * a + b
    return OutputSchema(result=result)

as_jax_array (from tesseract_core.runtime.jax_recipes) calls jnp.from_dlpack on device arrays and jnp.asarray on everything else, so device arrays stay on the GPU and host arrays move onto the default device. Tesseracts built on the JAX recipe (tesseract init --recipe jax) use it for all their inputs already.

def apply(inputs: InputSchema) -> OutputSchema:
    """Compute ``s * a + b`` on the GPU and return device memory."""
    a = torch.as_tensor(inputs.a, device="cuda")
    b = torch.as_tensor(inputs.b, device="cuda")
    result = inputs.s * a + b
    return OutputSchema(result=result)

torch.as_tensor(x, device="cuda") and torch.from_dlpack(x) both adopt device arrays without a copy.

The runtime exports any output array that lives in GPU memory and exposes __cuda_array_interface__, which includes CuPy, PyTorch, and JAX arrays. Inputs that arrive by handle reach your endpoints as IpcDeviceArray objects, which expose both __cuda_array_interface__ and DLPack, so any of these frameworks can adopt them without a copy. Avoid np.asarray and anything else that goes through __array__. The device arrays support it for convenience, by copying to host.

Warning

Input dtypes must match the schema exactly. The host encodings cast an input to the dtype its schema declares, but cuda_ipc does not cast on the device and rejects a mismatched input instead.

Serve with GPU transport

Enable the transport where the Tesseract is served:

from tesseract_core import Tesseract

tess = Tesseract.from_image(
    "my-gpu-tesseract", gpus=["all"], gpu_transport="cuda_ipc"
)
with tess:
    ...

The SDK runs the container with --ipc=host whenever the GPU transport is enabled, and raises a ValueError if no GPUs were requested.

from tesseract_core import Tesseract

tess = Tesseract.from_source("path/to/tesseract_api.py", gpu_transport="cuda_ipc")
with tess:
    ...

from_source serves the Tesseract in a subprocess on the same machine, without a container, so no extra GPU or IPC setup is needed.

$ tesseract serve my-gpu-tesseract --gpus all -e TESSERACT_GPU_TRANSPORT=cuda_ipc

As with the SDK, the container is started with --ipc=host, and the command fails if no GPUs were requested. Enabling the transport doesn’t make the server use it for every request. Each request asks for it, as the SDK and HTTP clients below do.

With from_image and from_source, gpu_transport can also be set through runtime_config={"gpu_transport": "cuda_ipc"}. If both are given, the keyword argument wins.

However you serve it, the Tesseract’s requirements need a CUDA-enabled build of its framework, such as jax[cuda13] or a CUDA build of PyTorch. The runtime finds the CUDA runtime library in those packages or in a system CUDA install.

With a JAX client and a JAX Tesseract on one GPU, each process preallocates 75% of the GPU’s memory by default, so the second one runs out of memory at startup. Disable preallocation (or lower XLA_PYTHON_CLIENT_MEM_FRACTION) in both processes before JAX initializes:

import os

# Before JAX initializes. A from_source Tesseract inherits it from this process.
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

# A container has to be passed it explicitly (`-e` with `tesseract serve`).
tess = Tesseract.from_image(
    "my-gpu-tesseract",
    gpus=["all"],
    gpu_transport="cuda_ipc",
    environment={"XLA_PYTHON_CLIENT_PREALLOCATE": "false"},
)

Warning

A Tesseract serving the GPU transport must run with a single worker (the default), and asking for more raises an error. That worker can serve concurrent requests.

A Tesseract created with gpu_transport="cuda_ipc" asks for the transport on every call. To check that it works between the two processes, call tess.resolve_gpu_transport(), which exchanges a small test array with the server and raises a RuntimeError with the reason if the exchange fails. Tesseract-JAX and Tesseract-Torch automatically run this check before their first call with GPU arrays.

Call it

When the inputs you pass to apply (or any other endpoint) live on the GPU, the client sends them by handle. GPU outputs come back as IpcDeviceArray objects, which any framework can adopt without a copy:

import cupy as cp

a = cp.arange(10_000_000, dtype=cp.float32)
b = cp.ones_like(a)

tess = Tesseract.from_image(
    "my-gpu-tesseract", gpus=["all"], gpu_transport="cuda_ipc"
)
with tess:
    out = tess.apply({"a": a, "b": b, "s": 3.0})["result"]

result = cp.asarray(out)      # stays on the GPU
host = out.copy_to_host()     # explicit copy to a NumPy array

For a Tesseract that someone else serves, Tesseract.from_url takes no transport argument. Ask the client which transport works, then request it with with_encoding:

tess = Tesseract.from_url("http://localhost:8000")
transport = tess.resolve_gpu_transport()  # "cuda_ipc" if it works from here, else "none"
out = tess.with_encoding(gpu_transport=transport).apply({"a": a, "b": b})

resolve_gpu_transport returns "none" with a warning when the server offers cuda_ipc but it doesn’t work from your process, for example because the server runs on another machine.

apply_tesseract uses the transport of the Tesseract you pass it. Under jax.jit on a CUDA backend, the call lowers to a native path that hands XLA’s device buffers to the Tesseract by handle and copies the results into XLA’s output buffers on the device. Gradients work as usual, since the VJP and JVP calls go through the same path:

import jax
import jax.numpy as jnp
from tesseract_core import Tesseract
from tesseract_jax import apply_tesseract

tess = Tesseract.from_image(
    "my-gpu-tesseract", gpus=["all"], gpu_transport="cuda_ipc"
)
with tess:

    def loss(a, b):
        out = apply_tesseract(tess, {"a": a, "b": b})
        return jnp.sum(out["result"])

    grad_a = jax.jit(jax.grad(loss))(a, b)

A Tesseract that asks for no transport, such as one from Tesseract.from_url, uses cuda_ipc if the server offers it and it works from your process, and copies through the host otherwise. To always copy through the host, pass tess.with_encoding(gpu_transport="none").

The native path is a compiled extension that ships with Tesseract-JAX, and it is used only when the call is compiled for a CUDA device. If the extension is unavailable, a Tesseract that asks for cuda_ipc raises an error instead of silently falling back to the host round trip, and one that asks for no transport falls back with a warning.

Warning

If JAX itself runs on the CPU, the GPU transport has no effect and the call takes the host round trip without an error. This happens with a CPU-only jaxlib, with JAX_PLATFORMS=cpu, or when JAX fails to load its CUDA plugin, in which case it falls back to the CPU with only a warning at startup. Check that jax.default_backend() returns "gpu" before relying on the GPU transport.

apply_tesseract from Tesseract-Torch uses the transport of the Tesseract you pass it. CUDA tensors are sent by handle, GPU outputs come back as CUDA tensors on their device, and autograd works as usual:

import torch
from tesseract_torch import apply_tesseract

a = torch.randn(1_000_000, device="cuda", requires_grad=True)
b = torch.randn(1_000_000, device="cuda")

tess = Tesseract.from_image(
    "my-gpu-tesseract", gpus=["all"], gpu_transport="cuda_ipc"
)
with tess:
    out = apply_tesseract(tess, {"a": a, "b": b})
    out["result"].sum().backward()

A Tesseract that asks for no transport, such as one from Tesseract.from_url, uses cuda_ipc if the server offers it and it works from your process, and copies through the host otherwise. Calls with only CPU tensors always copy through the host, and anything that comes through the host arrives on the CPU. To always copy through the host, pass tess.with_encoding(gpu_transport="none").

A client that speaks HTTP directly asks for the transport per request with a gpu_transport parameter on the Accept header, next to the usual CPU array encoding:

Accept: application/json+base64; gpu_transport=cuda_ipc

Without that parameter, or with gpu_transport=none, the response carries host copies even when the server has the transport enabled. A server without the transport enabled rejects cuda_ipc inputs with a 422 error.

The server keeps the GPU memory behind a response’s handles until the client says it has copied the arrays out. The SDK does this for you. A raw HTTP client sends the Tesseract-Exports-Done header on a later request, naming the IDs that earlier responses returned in their Tesseract-Exports header, and exports nobody claims are released after 5 minutes. See the array encodings reference for the details. Resolving a cuda_ipc handle requires calling into the CUDA runtime, so in practice this is for clients that already work with CUDA.

Troubleshooting

gpu_transport='cuda_ipc' requires GPU access, but no GPUs were requested. Pass gpus=["all"] (or specific GPU IDs) to from_image, or --gpus all to tesseract serve.

gpu_transport='cuda_ipc' was requested, but it does not work between this process and the Tesseract. Common causes are a Tesseract on another machine and a container started without access to the same GPU.

the native GPU FFI shim is unavailable. Your Tesseract-JAX install lacks its compiled extension. The published wheels for Linux on x86_64 and aarch64 include it, so reinstall from one of those. To compile it yourself, see building the GPU-direct shim from source.

The GPU transport is enabled, but calls are no faster. Check that your arrays actually live on the GPU on both sides. Common causes are inputs created on the CPU, a Tesseract that computes with NumPy, and adopting inputs with np.asarray. Setting TESSERACT_FORBID_DEVICE_HOST_COPY=1 in the client makes the SDK and a from_source Tesseract raise on any implicit copy of a GPU array to the host. With Tesseract-JAX, first check that jax.default_backend() returns "gpu" in the calling process. If JAX runs on the CPU, the native path is never used and nothing raises. Once JAX runs on the GPU, setting TESSERACT_JAX_DEBUG_CHECK_DEVICE_PTRS=1 makes any host pointer at the native boundary raise an error, which pinpoints where an array left the device.

Two JAX processes, and one fails with an out-of-memory error at startup. Each JAX process preallocates 75% of the GPU’s memory by default. Set XLA_PYTHON_CLIENT_PREALLOCATE=false in both, as described under Serve with GPU transport.