Metadata-Version: 2.4
Name: mm-libtorch-encoder
Version: 0.2.3
License-File: native/DEEPSEEK_LICENSE
License-File: native/PILLOW_LICENSE
License-File: native/NEMOTRON_LICENSE
Summary: Native Rust/C++ multimodal encoder for a matching Torch runtime
Requires-Python: >=3.11
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Requires-Dist: torch==2.14.1

# 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.

```python
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)
```

Pass `expected_tokens=[first_length, second_length]` to `generate()` to validate
frontend-planned token counts against the final embedding rows. Supply one positive
integer per image (a one-element iterable for a single bytes input). Invalid counts
or mismatched lengths raise `ValueError`; omitting the argument skips this check.

`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 available in the published
[0.2.1 wheels](https://pypi.org/project/mm-libtorch-encoder/0.2.1/) (0.2.0
wheels do not include it):

```python
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.

## Nemotron 3.5 Super VL images

Available from version **0.2.2** for BF16 and NVFP4 checkpoints; images only.

```python
with Encoder(
    "/models/nemotron35",
    "nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4",
    replicas_per_gpu=2,
) as encoder:
    tensor, metadata = encoder.generate(image_bytes)[0]
```

Keep `config.json` and `preprocessor_config.json` beside the encoder weights.
Local aliases can use `family="nemotron35"`. For `grid_thws=(1, h, w)`, the
embedding length and HTTP `expected_tokens` are `h * w / 4`, excluding text
delimiters. CUDA BF16/FP16 uses Torch Flash attention. Audio/video and extreme
aspect ratios exceeding the 256 MiB position workspace are rejected.

## Qwen3.5 images

Source builds support dense and MoE Qwen3.5 checkpoints with unquantized vision
weights under `model.visual.*`. Qwen3.5 support requires version 0.2.3 or newer.
Use the checkpoint's matching `preprocessor_config.json` and `config.json`:

```python
with Encoder("/models/qwen35", "Qwen/Qwen3.5-0.8B", family="qwen35") as encoder:
    tensor, metadata = encoder.generate(image_bytes)[0]
```

`grid_thws` retains the original `(1, h, w)` patch grid for the serving engine's
multimodal RoPE. The embedding length is `h * w / spatial_merge_size**2`.
Vision start/end tokens remain the tokenizer's responsibility; their embeddings
are not included. This implementation accepts images only and rejects videos.
Tiny upstream tower and processor fixtures exercise CPU parity; trained-weight
GPU parity has not yet been qualified.

## 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.

```python
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.

```python
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](https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/asynchronous-execution.html#stream-prioritization).
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](README.md#optional-output-cache) enables caching at
runtime. Both interfaces default to caching disabled.

```python
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:

```sh
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:

```sh
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](https://docs.pypi.org/trusted-publishers/creating-a-project-through-oidc/)
from your [PyPI publishing settings](https://pypi.org/manage/account/publishing/)
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:

```sh
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](https://pypi.org/manage/account/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](https://twine.readthedocs.io/en/stable/) in a separate environment:

```sh
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:

```sh
pip install --no-deps --only-binary=:all: mm-libtorch-encoder==0.2.3
```

## 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](https://github.com/basetenlabs/baseten/actions/runs/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.

