Skip to content
infinigencePublic

About

No description, website, or topics provided.

Resources

Stars

87 stars

Watchers

1 watching

Forks

Repository files navigation

apxinf-logo

ApxInf

ApxInf is a reimagined edge inference engine born of the agentic coding era, combining high performance, reliability, and energy efficiency across devices with an evolving agentic workflow that radically simplifies custom model development.

  • implemented in Rust, with backend-specific native dependencies
  • embodied AI is highest priority, VLA/WAM models on Jetson/DriveOS Thor/Orin
  • Agentically optimized CUDA Kernels

The first version of ApxInf ships with highly optimized PI-0.5 VLA model on Jetson Thor & Orin devices, and supports BF16, FP8 and INT8 precisions.

What's NEW!

  • [2026/10] 🔥 ApxInf supports Decision 3.0 mini and Qwen3.8-27B with image inference.
  • [2026/10] 🔥 ApxInf supoorts SmolVLA, deployable on Jetson Xavier / Thor / Orin / RTX 4090.
  • [2026/09] 🔥 ApxInf supports PI0-FAST, GR00T-N1.7, Qwen-Drive, Qwen3.8-27B.
  • [2026/09] 🔥 Initial open-source release of ApxInf repo, with PI0.5 support. Backends: Jetson Thor / Orin / RTX 4090.

Feature Matrix

VLA VLM Decision Hardware
  • PI0.5 ✅
  • PI0-FAST ✅
  • GR00T-N1.7 ✅
  • WALL-OSS-0.5 ✅
  • Qwen-Drive ✅
  • SmolVLA ✅
  • Qwen3.8-27B ✅
  • Qwen3-VL ✅
  • Decision 3.0 mini ✅
  • NVIDIA CUDA ✅
  • AMD ROCm
  • Apple Metal

Quick start

Make sure you have ApxInf built and installed, see Build ApxInf for instructions.

Native Apple Silicon text inference

The optional mlx feature runs Qwen3-0.6B and MiniCPM5-2B through native Rust models and the MLX 0.31.2 C++/Metal runtime. Native model composition lives in metal_qwen3 and metal_minicpm5, independent of the evolving multi-backend contract. It does not run a Python model. Qwen3 offers BF16 and mixed W8 selections; MiniCPM5 offers BF16 and an explicit DSpark drafter selection. Existing CUDA VLA families retain their own support scope.

export MLX_ROOT=/path/to/mlx-0.31.2
export DYLD_LIBRARY_PATH="$MLX_ROOT/lib"
export CARGO_TARGET_DIR=devlocal/apxinf-mlx/cargo-target
cargo build --release --features mlx
"$CARGO_TARGET_DIR/release/apxinf" generate \
  --model /path/to/Qwen3-0.6B --model-name metal_qwen3 --device metal --dtype bf16 \
  --model-variant bf16-compiled --greedy --max-tokens 160 \
  --chat-options '{"enable_thinking":false}' --prompt 'Explain gravity briefly.'

See the MLX architecture and supported profiles and native SDK setup for build requirements, MiniCPM/DSpark commands and supported options.

Run a d3-mini text decision

Download the complete vllm-sr/d3-mini package and build with CUDA. Construct D3_SYSTEM_PROMPT, D3_USER_PROMPT and token IDs with the model's d3_format.py.

cargo build --release --features cuda
./target/release/apxinf decide \
  --model /path/to/d3-mini \
  --system "$D3_SYSTEM_PROMPT" --prompt "$D3_USER_PROMPT" --candidates 3 \
  --autotune

The Python binding exposes the same pre-tokenized runtime:

import numpy as np
import apxinf_py

runner = apxinf_py.DecisionRunner.load("/path/to/d3-mini", autotune=True)
token_ids = np.asarray(formatted_d3_token_ids, dtype=np.uint32)
print(runner.decide(token_ids, 3))

Quickly benchmarking PI-0.5

Benchmarking PI-0.5 with randomly generated weights.

python scripts/bench_pi05.py --random-weights --model-variant bf16 --layer l1 \
  --views 2 --token-count 10 --action-horizon 10 --num-flow-steps 10 \
  --warmup 10 --samples 100 --autotune
python scripts/bench_pi05.py --random-weights --model-variant fp8_static --layer l1 \
  --views 2 --token-count 10 --action-horizon 10 --num-flow-steps 10 --autotune

Reported latency is P50 over 30 samples after 10 warm-up iterations (--warmup / --samples).

Run a policy through Python API

import numpy as np
from apxinf import AutoPolicy

policy = AutoPolicy.from_pretrained("<path-to-model>", model_variant="bf16")

observation = {
    key: np.zeros((256, 256, 3), np.uint8)
    for key in policy.metadata["image_keys"]
}
observation[policy.metadata["prompt_key"]] = "put both moka pots on the stove"
if state_key := policy.metadata.get("state_key"):
    observation[state_key] = np.zeros(policy.metadata["state_dim"], np.float32)

result = policy.infer(observation)

result["actions"]   # (H, policy.action_dim) float32, unnormalized
result["timing"]    # model_ms / total_ms
policy.close()

<path-to-model> is a checkpoint directory (model.safetensors, config.json, and that model's tokenizer/normalizer assets); none ships with this package. Resize, tokenization, normalization, and the flow sampler all run inside infer — pass raw frames.

Serve it with OpenPI compatible websocket server

python python/apxinf/examples/openpi_server.py \
  --model-dir <path-to-model> --model-variant bf16 --port 8000 \
  --image-keys observation/image,observation/wrist_image \
  --state-key observation/state

An unmodified openpi-client connects to it:

from openpi_client import websocket_client_policy

client = websocket_client_policy.WebsocketClientPolicy("127.0.0.1", 8000)
actions = client.infer(observation)["actions"]

Performance

d3-mini

Text decision model, BF16, batch 1. Public Jev Decision Index 0.3, 140,178 requests. See reproduction details.

Hardware Precision Public score Coverage P50 P95
Jetson AGX Thor BF16 54.91 100% 128 ms 1,165.7 ms

PI0.5

Two views, 224x224 NHWC uint8, 10 flow steps, H=10, batch 1. Latency is steady-state CUDA Graph replay P50.

Hardware Precision Latency Throughput
Jetson AGX Thor BF16 72.45 ms 13.8 Hz
Jetson AGX Thor FP8 41.16 ms 24.3 Hz
Jetson AGX Orin BF16 165.67 ms 6.0 Hz
RTX 4090 BF16 31.38 ms 31.9 Hz
RTX 4090 INT8 21.92 ms 45.6 Hz

With onestep action generation pruning.

Hardware Precision Latency Throughput
Jetson AGX Thor BF16 44.05 ms 22.7 Hz
Jetson AGX Thor FP8 26.32 ms 38.0 Hz
Jetson AGX Orin BF16 119.05 ms 8.4 Hz
RTX 4090 BF16 20.36 ms 49.1 Hz

LIBERO-10, 10 tasks x 50 episodes, H=10, replan=5, seed 7. PI0.5 reference is 92.4%.

Hardware Precision Trials Success Rate
Jetson AGX Thor BF16 500 464 92.8%
Jetson AGX Thor FP8 500 461 92.2%
Jetson AGX Orin BF16 500 460 92.0%

PI0-FAST

Two views, 224x224 NHWC uint8, batch 1.

Hardware Precision Prefix Per Token
Jetson AGX Thor BF16 33.1 ms 17.16 ms
Jetson AGX Thor FP8 31.0 ms 9.68 ms
Jetson AGX Orin BF16 117.3 ms 25.34 ms
RTX 4090 BF16 20.9 ms 5.24 ms

GR00T N1.7

One or two views, batch 1. Best recorded P50.

Hardware Precision 1-view P50 2-view P50
Jetson AGX Thor BF16 51.834 ms 54.216 ms
Jetson AGX Thor FP8 32.557 ms 35.436 ms
Jetson AGX Orin BF16 75.778 ms 84.864 ms
Jetson AGX Orin W8A8 56.711 ms 64.924 ms

Details.

Qwen3.8-27B-NVFP4

Text LLM, 48 Gated DeltaNet + 16 full-attention layers, mixed NVFP4/FP8. 2048-token prompt, 128 greedy tokens, batch 1, BF16 KV cache; steady state after one warm-up generation. See doc/qwen38-nvfp4.md for CLI inference.

Hardware Precision TTFT (2048) Decode Throughput
Jetson AGX Thor NVFP4/FP8 534 ms 78.1 ms/token 12.8 tok/s
cargo run -p apxinf-model --features cuda --release --example qwen38_bench -- \
    <path-to-Qwen3.8-27B-NVFP4> --prompt-len 2048 --max-new 128 --repeats 3

Qwen-Drive

Direct planning, BF16, batch 1, 10 flow steps and 12 input frames. Latency covers the end-to-end request from resident decoded images to the host trajectory.

Hardware Precision Latency Throughput PDM
Jetson AGX Thor SM110 BF16 482.60 ms 2.07 Hz 85.6786

Test setup, reproduction steps and accuracy results.

SmolVLA

Two views, 512x512 NHWC uint8, batch 1.

Hardware Precision Latency Throughput
Jetson AGX Thor BF16 29.7 ms 33.7 Hz
Jetson AGX Orin BF16 67.4 ms 14.8 Hz
Jetson Xavier NX FP16 312.38 ms 3.2 Hz
RTX 4090 BF16 13.47 ms 74.2 Hz

Port a new model with an agent

skills/model-port-workflow drives the whole sequence. Start with the documentation index and current module responsibilities. New VLA code separates Model forward computation, ModelRunner execution ownership, and weights; reuse the existing native binding and policy registries. The integration guide identifies current contracts and family-specific option limits.

Install it once, from the repository root:

# Claude Code
mkdir -p .claude/skills && ln -s ../../skills/model-port-workflow .claude/skills/

# Codex
mkdir -p ~/.agents/skills && ln -s "$(pwd)/skills/model-port-workflow" ~/.agents/skills/

Then invoke it with the model, the target, and the acceptance bar:

/model-port-workflow port GR00T N1.7 from <path-to-reference-implementation> to
ApxInf, Jetson Thor, BF16, parity against the reference within 1e-2

It works from the same guides a human would follow:

Build ApxInf

git clone <repo-url> && cd ApxInf
python3 -m venv .venv && source .venv/bin/activate
pip install maturin
CARGO_TARGET_DIR=target/wheel maturin build --release --features cuda --auditwheel skip -m crates/apxinf-py/Cargo.toml
pip install --force-reinstall target/wheel/wheels/apxinf_py-*.whl
pip install -e "python/apxinf[serving]"

Activate a venv or conda env before installing. The serving extra adds msgpack/websockets.

--features cuda is a Cargo feature, not a CUDA installation: it compiles the CUDA backend into the binding, and it is required — the PI0.5 runtime is only registered on CUDA devices, as is WallOSS.

The build queries the visible GPU for its compute capability and compiles the kernels for exactly that architecture, so build on the machine you deploy to; cross-compiling fails unless APXINF_CUDA_ARCH names the target (sm_87 Orin, sm_101 Thor-U, sm_110 Thor).

Confirm the binding imports and reaches the GPU:

python -c 'import apxinf_py; print(apxinf_py.__version__)'
python scripts/bench_pi05.py --random-weights --model-variant bf16 --layer l1 --samples 5

Needs a Linux host with an NVIDIA driver and a CUDA toolkit plus a stable Rust toolchain. If nvcc --version or cargo --version fails, set them up first:

Using ApxInf from Python

Three public layers, each wrapping the one before. Pick the outermost one that still leaves you the control you need.

L1 — native ModelRunner

You own resize, tokenization, noise, and unnormalization; the model takes already-resized frames and returns a normalized-domain chunk.

from apxinf import ModelRunner

model_runner = ModelRunner.load("pi05", "<path-to-model>/model.safetensors", model_variant="bf16")

# rgb: uint8 [views, H, W, 3] at model_runner.image_size; tokens: uint32; noise: float32
actions = model_runner.infer_rgb(rgb, "nhwc", token_ids, noise)   # (H, action_dim)
model_runner.action_horizon, model_runner.num_views, model_runner.image_size

L2 — policy

Adds the pre/post pipelines and reads the checkpoint's tokenizer and norm_stats, so it takes a raw observation dict and returns deployable actions — the Run a policy snippet. Beyond that call it exposes the serving contract, the pipelines, and the layer boundary:

policy = AutoPolicy.from_pretrained(
    "<path-to-model>",
    model_variant="bf16",
    action_dim=None,        # default: infer the model's full vector from checkpoint weights
)

policy.metadata             # model_type, action_horizon, image_keys, state_key, ...

result = policy.infer(observation)
result["normalized_actions"]  # what L1 returned, before trim + unnormalize

The pipelines are ordered named steps — image_stack -> tokenize in, trim -> unnormalize out — and every mutation returns a new one, so a custom step drops in as a value: policy.input_pipeline.replace("tokenize", MyTokenizeStep()). See the frontend README.

L3 — websocket server

Wraps an L2 policy in the OpenPI wire protocol. See OpenPI-compatible serving.

OpenPI-compatible serving

The server speaks OpenPI's websocket protocol, so an existing openpi-client robot stack connects without a code change — swap the endpoint and keep the observation dict you already send.

from apxinf import AutoPolicy
from apxinf.serving import WebsocketPolicyServer

policy = AutoPolicy.from_pretrained(
    "<path-to-model>",
    model_variant="bf16",
    image_keys=("observation/image", "observation/wrist_image"),
    state_key="observation/state",
)
WebsocketPolicyServer(policy, "0.0.0.0", 8000).serve_forever()

image_keys= / state_key= / prompt_key= say what your client sends. Omit image_keys and the policy falls back to the model's own view-slot names (base_0_rgb, ...), published as apxinf.CANONICAL_IMAGE_KEYS — a fallback, not a contract. state_key has no fallback at all, because the two failures are not alike: a wrong camera key raises on the first inference, a wrong state key is silent, so a policy that reads state refuses to be built without one.

The server keeps the checkpoint's native action width unless the user supplies --action-dim. It publishes the resolved wire contract in connect-time metadata so the client can assert it rather than guess.

Precisions

--model-variant selects the PI0.5 implementation; the serving command is otherwise unchanged.

python scripts/pi05_openpi_websocket_server.py \
  --model-dir <path-to-model> --model-variant bf16 --port 8000 \
  --image-keys observation/image,observation/wrist_image \
  --state-key observation/state
policy = AutoPolicy.from_pretrained("<path-to-model>", model_variant="bf16")
Precision Where Needs
bf16 every supported device; the default the checkpoint alone
fp8_static Thor only, where it is the fastest path — Orin has no FP8 Tensor Cores per-tensor activation scales
int8_dynamic W8A8, optimized for Orin (SM87) and Ada (SM89) the checkpoint alone

FP8 calibration

Pass the calibration generated for the deployment data explicitly; when omitted, ApxInf falls back to <path-to-model>/calibration.json:

python scripts/pi05_openpi_websocket_server.py \
  --model-dir <path-to-model> --model-variant fp8_static \
  --image-keys observation/image,observation/wrist_image \
  --state-key observation/state \
  --calibration <path-to-calibration.json> \
  --port 8000
policy = AutoPolicy.from_pretrained(
    "<path-to-model>",
    model_variant="fp8_static",
    calibration="<path-to-calibration.json>",
)

If the checkpoint does not contain calibration.json, generate one from representative Observations — a JSONL manifest, a directory of Observation NPZ files captured wherever the environment already lives, or the native LIBERO10 simulator driven in this process:

python3 scripts/calibrate_pi05.py --model-dir <path-to-model> \
  --manifest <path-to-observations.jsonl>

python3 scripts/calibrate_pi05.py --model-dir <path-to-model> \
  --input-dir <path-to-observations>

python3 scripts/calibrate_pi05.py --model-dir <path-to-model> \
  --libero-suite libero_10

The NPZ form makes the calibration input a reviewable artifact rather than a side effect of a rollout; --libero-suite drives MuJoCo here and needs the same dependencies as LIBERO evaluation.

See PI0.5 FP8 calibration for the Observation format, native LIBERO sampling, and output options.

LIBERO evaluation

Get the checkpoint

The published accuracy is pi05_libero_base, π0.5 fine-tuned on LIBERO — an arbitrary π0.5 checkpoint might not reproduce it.

pip install -U "huggingface_hub[cli]"
huggingface-cli download lerobot/pi05_libero_base --local-dir <path-to-model>
curl -fL https://storage.googleapis.com/openpi-assets/checkpoints/pi05_libero/assets/physical-intelligence/libero/norm_stats.json \
  -o <path-to-model>/norm_stats.json

The lerobot/pi05_libero_base checkpoint lost its normalization statistics during repository updates. To reproduce the officially reported performance, download OpenPI's LIBERO norm_stats.json separately as shown above and pass it explicitly with --norm-stats.

Run

The rollout needs LIBERO and MuJoCo:

python -c 'from libero.libero import benchmark'

If that fails, install LIBERO from source (pip install -r requirements.txt && pip install -e ., then export MUJOCO_GL=egl, or osmesa if the machine has no EGL). Its requirements.txt brings MuJoCo through robosuite and pins numpy==1.22.4, which predates Python 3.11 and so builds from source on a newer interpreter; --no-deps skips it.

--backend websocket additionally needs openpi-client, from an openpi checkout (pip install -e packages/openpi-client).

scripts/eval_libero.py builds the policy in-process — no server involved:

python scripts/eval_libero.py --backend in-process --model-dir <path-to-model> \
  --norm-stats <path-to-model>/norm_stats.json \
  --precision bf16 --action-horizon 10 \
  --suite libero_10 --tasks all --trials-per-task 50 \
  --results-jsonl <out-dir>/results.jsonl --summary-json <out-dir>/summary.json

That is the published protocol: all 10 LIBERO-10 tasks x 50 episodes at seed 7 (the default), 500 episodes in total.

Options

  • --suite picks the task suite, --tasks a comma list within it, and --trials-per-task the episode count; a smoke run is --tasks 0 --trials-per-task 1.
  • The model flags — --model-type, --norm-stats, --action-horizon, --action-dim, --discrete-state, FP8 --calibration — belong to --backend in-process alone.
  • --backend websocket --host <h> --port <p> evaluates a running server instead, on this machine or another. The model flags belong to the server there, and --precision only asserts what the server reports, so a mismatch fails at connect instead of skewing a run. Pass --norm-stats <path-to-model>/norm_stats.json to scripts/pi05_openpi_websocket_server.py when serving this checkpoint.
  • --image-keys / --state-key override the LIBERO wire keys. Both backends read them, so an in-process and a websocket run stay comparable.
  • Runs are resumable: completed task/trial rows in the JSONL ledger are skipped, and the summary reports success rate alongside per-segment latency.

π0-FAST

π0-FAST is a token VLA: instead of a continuous flow-matching head it autoregresses FAST discrete action tokens, and the policy layer detokenizes them (BPE then an orthonormal DCT) before unnormalizing. L1 therefore has its own entry point — the raw tokens — and L2 turns them into actions:

from apxinf import AutoPolicy

policy = AutoPolicy.from_pretrained("<path-to-model>", precision="bf16")

# L1: raw FAST action tokens, uint32 [max_action_tokens]
tokens = policy.model.infer_action_tokens_rgb(rgb, "nhwc", token_ids)

# L2: tokens -> detokenized, unnormalized actions
result = policy.infer(observation)        # observation: 2 cameras + state + prompt
result["action_tokens"]                   # the ids behind result["actions"]
result["normalized_actions"]              # before unnormalization

At L1 the runtime embeds exactly the ids it is handed, so token_ids must end with the PaliGemma BOS that triggers decoding — LeRobot concatenates it after the prompt before prefill, and Pi0FastPolicy appends it for you. Without that trailing BOS the model continues the prompt text (a caption) instead of emitting the Action: …| stream.

The prompt is assembled from the task and the discretized state, so the state is required and its width is a checkpoint property — policy.metadata["state_dim"] (8 for LIBERO's eef_pos + eef_axis_angle + gripper_qpos). The two tokenizers are read from the assets the checkpoint names (text_tokenizer_name, action_tokenizer_name) and resolved offline: pass a local path with tokenizer_path/fast_tokenizer_path, or set APXINF_PALIGEMMA_TOKENIZER / APXINF_FAST_TOKENIZER, or let the local Hugging Face cache answer.

Run

# the native binding then the frontend, as in "Build ApxInf" - no extra of its own
pip install --force-reinstall target/wheel/wheels/apxinf_py-*.whl
pip install -e python/apxinf
huggingface-cli download lerobot/pi0fast-libero-v044 --local-dir <path-to-model>

python scripts/eval_libero.py --backend in-process --model-dir <path-to-model> \
  --precision bf16 \
  --suite libero_10 --tasks all --trials-per-task 50 \
  --results-jsonl <out-dir>/results.jsonl --summary-json <out-dir>/summary.json

Unlike PI0.5 this checkpoint carries its own normalization statistics, so --norm-stats is neither needed nor accepted (its use is an error, not a silent no-op).

Benchmark

scripts/bench_pi05.py times the concentric serving shells so a regression can be attributed to the engine, the processors, or the transport.

python scripts/bench_pi05.py --model-dir <path-to-model> --model-variant bf16 --layer l1,l2
  • --layer selects any subset of l1 (bare model), l2 (full policy), l3 (websocket round trip). L3 attaches to a running server and needs no local weights.
  • --model-dir runs a real checkpoint at its native horizon; --random-weights runs the engine with no checkpoint on disk, and the shape knobs (--views, --image-size, --action-horizon, --num-flow-steps, --token-count) select the synthetic workload.
  • --calibration is FP8-only and synthetic-only; a checkpoint reads calibration.json from its own directory.
  • --action-horizon also applies to a checkpoint — the horizon is a sequence length, not a weight dimension — which is what makes a real checkpoint comparable to a synthetic run.
  • --warmup / --samples set the sampling protocol (default 10 and 30); --out writes the report as JSON.

Any registered model type works: AutoPolicy dispatches on the checkpoint's config.json, so the same command benchmarks the next model without a flag change.

For one-step evaluation, please refer to doc/run_warmstart_with_onestep.md

NVIDIA build environment

A complete CUDA toolkit is required: nvcc, CUDA headers and runtime, cuBLAS and cuBLASLt development libraries, and NVTX (libnvToolsExt on Jetson, libnvtx3interop on desktop CUDA). Also a C/C++ compiler, linker, ar, Git, pkg-config, and Python 3. The CUDA kernels, CUTLASS, and FlashAttention sources are vendored — no external checkout needed.

Install the driver and toolkit through the JetPack, DRIVE OS, or CUDA distribution for the machine, then check nvcc --version. If CUDA does not live at /usr/local/cuda, point CUDA_PATH at it.

Device Architecture Validated toolkit
Jetson AGX Thor sm_110 CUDA 13.0
Thor-U sm_101 CUDA 12.8
Jetson AGX Orin sm_87 CUDA 12.6, 13.2
RTX 4090 sm_89 CUDA 12.8

Rust toolchain

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
source "$HOME/.cargo/env"
rustup default stable

Built with Rust 1.95 and 1.96; no minimum supported version is declared.

Acknowledgement

The development of APXInf has been inspired by, and benefits from, the ideas and tooling of the broader open-source community. In particular, we would like to thank the teams and contributors behind FasterTransformer, TensorRT-LLM, llama.cpp, FlashAttention, FlashRT, vLLM, sgLang, and if we have inadvertently missed your project or contribution, please open an issue or a pull request so we can properly credit you.

License

Apache 2.0. Vendored third-party components retain their own licenses.

Community

Scan the QR Code to join our Wechat Group

wechat-group

About

No description, website, or topics provided.

Resources

Stars

87 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages