Native Python integration

The python Cargo feature exposes the Rust/C++ encoder through PyO3. It shares preprocessing and model execution with the HTTP server, but does not start a server or compile its HTTP client, routing, and serialization dependencies. The first interface accepts encoded image bytes and returns owning Torch tensors.

from mm_libtorch_encoder import Encoder

with Encoder("/models/glm53", "zai-org/GLM-5.3-Flash", device="cuda:0") as encoder:
    results = encoder.generate([first_image_bytes, second_image_bytes])
    for embeddings, metadata in results:
        # Actual BF16 Torch tensor on cuda:0, sharing the native allocation.
        print(embeddings.shape, metadata.grid_thws, metadata.mm_hash)

generate() returns a list of ImageResult(tensor, metadata) pairs in input order, one per image. A single bytes object is also accepted. Each tensor owns its storage independently of the result tuple and encoder; keeping a tensor after deleting its result or closing the encoder is supported. The API does not copy embeddings, serialize them, or move them to CPU. Treat tensor contents as read-only; Torch does not enforce immutability.

ImageMetadata is a frozen dataclass with only immutable values: model, mm_hash, length, shape, dtype, device, strides (in elements), embedding_offsets, and grid_thws. It preserves crop boundaries and grid rows from the model output. Only the small CPU metadata tensors are converted to tuples. Metadata can be pickled independently; this API does not implement cross-process tensor transport or shared CPU/memfd storage. Images are executed sequentially, without automatic model batching.

The lower-level encode_image(image_bytes) interface remains available and returns the original dictionary of named Torch tensors, including the integer metadata tensors.

DeepSeek V4.1 Flash support is added in version 0.2.1 (the published 0.2.0 wheels do not include it):

with Encoder("/models/deepseek-v41", "deepseek-ai/DeepSeek-V4.1-Flash") as encoder:
    tensor, metadata = encoder.generate(image_bytes)[0]
    print(metadata.length, metadata.grid_thws)

Its tensor contains the full image span, including start/end and row newline embeddings. metadata.length is ceil(h / 3) * (ceil(w / 3) + 1) + 2, where (1, h, w) is the returned ViT patch grid. It supports images; audio/video are rejected. Local checkpoint aliases can select family="deepseek41" explicitly.

GPU replicas

Encoder(..., replicas_per_gpu=N) creates N persistent native workers on the selected GPU. The default is 1; the count must be an integer from 1 to 32. Larger counts are rejected before allocating queues or starting workers. Each worker owns a separate copy of the model and installs its own pooled CUDA stream before loading weights. Workers take jobs from one bounded queue. The optional Foyer cache is shared by all replicas, with one disk writer and one cache budget. encoder.replicas_per_gpu reports the configured count.

from concurrent.futures import ThreadPoolExecutor
from mm_libtorch_encoder import Encoder

with Encoder("/models/kimi-k3", "moonshotai/Kimi-K3", replicas_per_gpu=2) as encoder:
    with ThreadPoolExecutor(max_workers=2) as callers:
        outputs = list(callers.map(encoder.encode_image, image_bytes_list))

Calls from different Python threads can execute on different replicas; a single generate(images) call still processes its images sequentially and preserves their order. This does not batch requests together. Startup waits for all model loads and cleans up initialized workers if any load fails. close() drains accepted calls, joins every worker, then closes the shared cache. Each replica adds model weights and activation memory; throughput gains depend on the workload. The HTTP server exposes the same count as MM_ENCODER_REPLICAS_PER_GPU.

Sharing a GPU with prefill

Each worker executes at most one image at a time, so an instance can execute up to replicas_per_gpu images concurrently. Its CUDA stream requests stream_priority=1 by default. The pinned LibTorch implementation clamps every positive request to 0, the lowest supported priority, so this default has the same priority as the default Torch stream. It does not deprioritize encoding relative to default-priority work. Lower numbers mean higher priority; negative requests select higher-priority streams within the supported range. To favor prefill, run the engine's model work on a higher-priority stream.

import torch
from mm_libtorch_encoder import Encoder

encoder = Encoder("/models/glm53", "zai-org/GLM-5.3-Flash")  # Requests priority 1.
prefill_stream = torch.cuda.Stream(device="cuda:0", priority=-1)
print(encoder.stream_priority, prefill_stream.priority)  # Actual priorities after clamping.
# Integrate prefill_stream into the engine's own execution/graph lifecycle.
# Run encoder.generate(...) from a separate CPU thread for concurrent dispatch.

encoder.stream_priority and the native_ready log report the actual priority selected by Torch. The worker installs its stream before loading weights; ATen, attention, custom CUDA kernels and cuBLASLt use its current stream. Concurrent Python callers share the instance's queue. Create one Encoder per GPU and set replicas_per_gpu to control how many workers share that device.

Stream priority favors pending kernels; it does not preempt running kernels, reserve SMs, or isolate memory bandwidth. See the CUDA scheduling documentation. When the encoder is clamped to 0, it shares the default stream's priority. For stricter scheduling, the engine should defer calls to generate() while prefill is busy; avoid deferring an image that the current prefill needs. Existing engine CUDA graphs and auxiliary streams must be integrated by the engine, rather than assuming that an outer Python stream context changes them.

Results are fenced on the encoder stream before return and registered with the caller's consumer stream for allocator lifetime safety; the encoder does not add a device-wide synchronization. Set the consumer stream context around generate() when its GPU results will be consumed on a nondefault stream. CPU cache hits skip GPU work entirely. A consumer uploading cached CPU results still needs its normal copy/event dependencies.

Optional Foyer CPU and disk cache

The same Rust implementation is bundled into the standalone HTTP server; its MM_ENCODER_CACHE setting enables caching at runtime. Both interfaces default to caching disabled.

from mm_libtorch_encoder import CacheConfig, Encoder

with Encoder(
    "/models/glm53", "zai-org/GLM-5.3-Flash",
    cache=CacheConfig(
        memory_bytes=1024**3,
        disk_path="/cache/glm53-gpu0",
        disk_bytes=8 * 1024**3,
    ),
) as encoder:
    tensor, metadata = encoder.generate(image_bytes)[0]
    assert tensor.device.type == "cpu"
    print(encoder.cache_stats())

Caching is disabled by default. CacheConfig(memory_bytes=...) enables only the CPU tier; configure both disk_path and disk_bytes for Foyer's hybrid cache. With caching enabled, all results are CPU tensors, including misses. The first request transfers embeddings from GPU into owned CPU storage. Memory hits reuse that allocation without payload copies and skip preprocessing and GPU execution. Tensor views keep their backing storage alive after cache eviction or encoder shutdown. Clone tensors before modifying them: the cache shares their contents and Torch does not enforce immutability. The metadata remains frozen. Keep the CPU tensor alive until any asynchronous upload to a GPU has completed.

Disk writes serialize CPU bytes; disk reads allocate and reconstruct CPU tensors. The disk path is not zero-copy and this does not implement memfd/worker transport. Each disk directory allows one encoder writer at a time, enforced by a process lock. Use separate paths for simultaneous encoder instances. close() waits for Foyer's pending persistence work; cache read/close errors are counted and logged. Read errors fall back to recomputation. The cache is disposable, not a durable source of model results.

Keys include image content plus a namespace derived from the actual selected encoder weights, model/processor configuration, native extension binary, Torch/CUDA identity, GPU architecture/name and matmul precision. Hashing weights adds one encoder-weight read at cache-enabled startup. Keep checkpoint files and execution precision settings unchanged while the encoder is running.

memory_bytes limits byte-weighted cache retention, not total process RSS. Outstanding tensors can keep evicted storage alive; live_cpu_bytes_process reports live cached tensor payload across all encoders, including those external references. Torch allocator retention, metadata, disk I/O buffers and transient disk decoding allocations are additional memory. Request admission must bound the number/size of outstanding results.

max_entry_bytes defaults to 256 MiB (up to 1 GiB). Larger results are returned normally but bypass caching. Disk storage needs at least four blocks; each block is max_entry_bytes + 1 MiB, rounded up to 4 KiB. One block-sized I/O buffer pool is configured. Counters expose memory_hits, disk_hits, misses, bypasses, errors, and live_cpu_bytes_process.

Checkpoint files must already be present. The host application handles fetching, request admission, and batching. Each encoder owns its configured worker replicas; Python releases the GIL while waiting. Create one instance per GPU inside the serving process after process creation. Instances reject calls after fork. close() drains accepted calls and releases model ownership. Returned tensors retain their storage after the encoder closes.

The producer uses a private CUDA stream and fences execution before exporting DLPack tensors. Allocations are recorded on the caller's current CUDA stream; consume returned tensors on that stream, or synchronize before releasing them when using another stream. Embeddings stay on the GPU, with no HTTP, MessagePack, base64, or GPU-to-CPU output transfer. This first interface is synchronous; asynchronous stream handoff is a future optimization. Encoding and the serving engine still compete for GPU compute and memory, so reserve space for encoder weights and outputs when sizing the engine's KV cache.

Wheel and runtime boundary

The crate and Python package share a release version. Cargo.toml owns the version; Maturin derives the wheel version through project.dynamic in pyproject.toml. Update the crate version and its Cargo.lock entry together for each release. Read the installed wheel version with importlib.metadata.version("mm-libtorch-encoder").

The wheel builder reads the installed Torch SDK, CUDA version and C++ ABI instead of pinning a single engine runtime. Each wheel records its Torch version/source revision/build configuration, CUDA version, C++ ABI, machine architecture and libc version. Build and final engine stages must retain compatible system C++ libraries; the intended recipe builds inside the engine's own SDK environment. The Python adapter verifies this manifest before importing native code. CPython 3.11+ uses the abi3 extension. Exact engine images still require qualification; matching metadata alone does not establish model correctness or performance.

Component	Packaging
Rust and our C++ bridge/model code	Compiled into the extension
Custom CUDA DSO and restricted media libraries	Bundled in the wheel with private repaired library names
Torch/ATen, CUDA runtime and NVIDIA math libraries	Dynamically linked to the matching host installation
NVIDIA driver	Supplied by the host
Model weights	Supplied on disk, outside the wheel

A wheel can bundle shared libraries: static linking is not required for a single pip-installable artifact. Static linking a second Torch runtime would not provide process isolation for its dispatcher, allocator, and CUDA dependencies. Sharing the host runtime is the intended in-process contract. Images with a different Torch/CUDA stack require a separately built and qualified matching wheel.

Build inside each engine's SDK environment with Rust 1.93+, a C++20 compiler, the restricted media libraries from the Dockerfile, and an NVCC toolkit matching that engine's torch.version.cuda. Runtime-only images may need these build tools in a separate builder stage. CUDA 12 split NVIDIA wheel headers and CUDA 13's unified layout are both discovered by the builder:

python -m pip install maturin==1.12.6 patchelf==0.19.1.0 auditwheel==6.6.0 cmake==3.31.6
export CUDA_HOME=/path/to/matching/cuda
export FFMPEG_DIR=/path/to/restricted/media
python python/build_wheel.py

The build script selects the installed Torch SDK, builds the extension, and repairs media/kernel dependencies into dist/<runtime-key>/. Compiler caches are also separated by runtime key so changing Torch at the same SDK path cannot reuse incompatible compiled objects. Torch and NVIDIA libraries are explicitly excluded. The artifact keeps a Linux platform tag because its runtime contract extends beyond manylinux's system-library baseline. Install the resulting wheel into the qualified engine image using the local wheel selector:

python python/install_wheel.py /path/to/wheelhouse

The selector inspects embedded manifests and requires exactly one matching wheel, then installs with --no-deps. It neither downloads code nor replaces Torch. Collect runtime-specific dist/ subdirectories in one wheelhouse for SGLang, vLLM and TRT-LLM images; each image selects its own implementation at build time. A missing match requires compiling and qualifying a wheel in that engine's SDK build environment. Use separate source checkouts for concurrent runtime builds, because the builder writes a generated package manifest. Runtime startup only loads the selected extension; it does not compile anything.

Media decoding retains the server's size bounds and decoder restrictions. It runs in the host process with that process's privileges; this integration does not add a security isolation boundary.

Validation so far

A repaired wheel was installed in the matching Python 3.12/Torch 2.14.1 CUDA 13 SDK container and exercised on an H100 with trained GLM5.3 weights. The smoke covers BF16 GPU output, repeatability with two Python caller threads, malformed image recovery, delayed consumer-stream work, and output ownership after close. Installing that CUDA 13 wheel in the Torch 2.11/CUDA 12.9 image was rejected before native import; building a matching wheel for the older image passed. All 18 existing server library tests passed after the shared-module extraction.

This establishes the packaging and ownership path, not production throughput or a new accuracy/performance comparison for every model family. The smoke is runnable as python python/tests/smoke.py /models/glm53 after installing the wheel.

Compiled PyPI releases

Manual runs of mm-standalone-encoder.yml also build compiled AMD64 and ARM64 wheels against the pinned Torch 2.14.1/CUDA 13 SDK. The encoder-wheel-amd64 and encoder-wheel-arm64 artifacts contain only .whl files; no source distribution is built or uploaded. The Rust/C++/CUDA implementation and media libraries are compiled into the wheel alongside the Python interface. Torch, CUDA libraries, the driver, and model weights remain external.

These public wheels use CPython 3.11+ abi3 and audited manylinux_2_36 tags. They require glibc 2.36 or newer and torch==2.14.1; startup additionally checks the exact Torch source/build, CUDA 13 runtime, C++ ABI, and CPU architecture. Unlike exact engine wheels, the public manifest allows newer glibc versions. Other engine runtimes still need their own qualified wheels.

Publish from the release workflow

To enable publishing, configure a PyPI Trusted Publisher for project mm-libtorch-encoder, GitHub owner basetenlabs, repository baseten, workflow mm-standalone-encoder.yml, and environment pypi. For the first release of a new project, create a pending publisher from your PyPI publishing settings using those same fields. The first successful upload creates the project. For an existing project, add the publisher in that project's Publishing settings. CI authenticates with GitHub OIDC; no PyPI API token is needed in GitHub secrets. Then dispatch the same image/wheel release workflow from your local machine:

gh workflow run mm-standalone-encoder.yml --repo basetenlabs/baseten --ref master \
  -F publish_image=true -F publish_pypi=true

The publish job uploads only the compiled wheel artifacts after both architecture builds and standalone GPU qualification pass. Publishing is disabled by default. Bump Cargo.toml and Cargo.lock together before each new PyPI version; PyPI does not allow replacing existing release files.

Publish built wheels from a local machine

You can also upload the built wheels directly from a laptop. Download both architecture wheels into one directory, such as ./encoder-wheelhouse/. Uploading does not import the extension or compile code, so the publishing machine does not need Linux, Torch, CUDA, Rust, or the model weights.

Create a PyPI API token authorized to publish mm-libtorch-encoder. The first upload of a new project requires an account-scoped token; once the project exists, a project-scoped token can be used. Run Twine in a separate environment:

python3 -m venv /tmp/mm-encoder-publish
/tmp/mm-encoder-publish/bin/python -m pip install twine==6.2.0
/tmp/mm-encoder-publish/bin/python -m twine check --strict ./encoder-wheelhouse/*.whl
/tmp/mm-encoder-publish/bin/python -m twine upload --username __token__ ./encoder-wheelhouse/*.whl

Twine prompts for the API token as the password. These commands upload only .whl files, including both AMD64 and ARM64; no source distribution is uploaded. On Windows, use the publishing environment's Scripts/python.exe path instead of bin/python. A repeated upload of an existing filename is rejected; bump the package version and rebuild for the next release.

Install in a serving image

Install into a serving image that already has the matching CUDA Torch runtime:

pip install --no-deps --only-binary=:all: mm-libtorch-encoder==0.2.0

Engine-image qualification

Manually dispatch .github/workflows/mm-standalone-encoder.yml with qualify_python=true on the native encoder branch to qualify the SGLang, vLLM and TRT-LLM images pinned by the repository's versions/<engine>.env files. The matrix resolves image tags to immutable digests, builds a wheel with Dockerfile.python, and checks loading in the original engine image. Each artifact includes that digest and its wheel. After all three wheels build, a serialized B200 job installs each wheel in its original engine image, imports that engine, and runs the trained GLM smoke. It downloads only the pinned vision checkpoint subset and uses the shared node GPU mutex. This tests native inference with engine imports in the same process; full serving-engine model execution and production throughput still require target-deployment qualification.

CI run 38857534601 passed wheel builds, original-image imports, the B200 trained GLM smoke for all three engines, and the standalone server checks. The artifacts record the exact base-image digests and runtime manifests; wheels are approximately 5.8 MB and contain neither Torch nor CUDA runtime libraries.

Engine image	Bundled Torch	CUDA	Build, import, GPU smoke
SGLang nightly a6a8c560	2.14.1+cu130	13.0	Passed
vLLM v0.26.0	2.11.0+cu130	13.0	Passed
TRT-LLM 9442083b7d	2.14.0a0+4fdf77b940.nv26.08	13.4	Passed

These results cover Linux x86_64. ARM64 builds have not been qualified. The GPU smoke imports each engine in the encoder process but does not load an LLM or exercise that engine's scheduler, CUDA graphs, or custom memory allocator.

Local H100 validation also covered a Torch 2.11.0/CUDA 12.9 build. On three GLM5.3 image geometries (224x224, 319x257, 672x448), two outputs matched the Torch 2.14.1/CUDA 13.0 build bitwise; the third had cosine similarity 0.99999958, relative L2 error 0.000916, and maximum absolute difference 0.0004883. These are small synthetic-image checks, not a replacement for the trained-model corpus.