Deploying VLA models on Jetson — what fits, what doesn't, and what TensorRT buys you

The inference problems I’ve spent the last year on — KV cache memory, all-reduce barriers, bandwidth-bound decode — are all about text. A prompt goes in, tokens come out. The world is a string.

Vision-Language-Action (VLA) models break that abstraction. A VLA takes a camera frame, a language instruction, and the robot’s current joint positions — and outputs an action chunk: the next 30 or 50 joint-angle deltas the robot should execute at 15–20 Hz. The world is a manipulator arm and a workspace, and the “output” has to arrive before the control loop deadline or the robot acts on stale data.

The deployment target I care about is Jetson — the same unified-memory platform I’ve been squeezing for LLM inference. Running a VLA on Jetson at 20 Hz is a harder constraint than running it in the cloud: the memory budget is fixed, there’s no PCIe, and the GPU is shared with the vision pipeline, the OS, and any co-resident processes. This post covers what I learned deploying VLA models with tether — the open-source edge-to-cloud deployment CLI from FastCrest — and specifically what the real constraints are, how TensorRT helps, and what you actually have to do to get a pi0-class model running on an Orin Nano.

Table of contents

What is a VLA model?

A Vision-Language-Action model is a robot policy that generalizes across tasks and environments. Traditional robot policies are narrow: train on pick-and-place in a specific workspace, deploy in that specific workspace, retrain for anything else. VLAs replace task-specific policies with a single large model — a VLM backbone (vision encoder + language model) that interprets a camera frame and a language instruction, combined with an action head that outputs robot joint targets.

The current generation of VLAs:

The weight count alone signals the problem. pi0.5 in FP32 is ~13 GB. SmolVLA is 1.6 GB. A Jetson Orin Nano has 8 GB of unified memory — shared between model weights, vision buffers, the OS, and everything else. The first engineering question isn’t “how fast” — it’s “does it fit?”

What fits on Jetson

The Jetson platform matrix is the starting point for any VLA deployment decision:

The Orin Nano boundary is hard. The Pi0 monolithic ONNX is 12.5 GB — it cannot be allocated in a 7.6 GB usable pool regardless of swapping, quantization, or how cleverly you tune the OS. SmolVLA at 1.6 GB fits comfortably and is the production VLA for Nano-class hardware. GR00T N1.6 at 4.4 GB fits with ~3 GB to spare for the OS and vision buffers.

This is the same constraint I documented in genie-ai-runtime. On Jetson Orin Nano 8 GB, unified memory is the single budget that governs everything — model weights, KV cache, vision buffers, the OS floor (~1.5 GB), and any co-resident processes like a voice pipeline or agent runtime. A 12.5 GB ONNX simply cannot exist on a 7.6 GB usable pool. The Nano target forces you toward smaller models, quantization, and distillation — not because of compute, but because of bytes.

So the question becomes: how do you get pi0-quality on a Nano? The answer is SnapFlow distillation.

SnapFlow distillation: getting pi0.5 down to 6.5 GB

Flow-matching VLAs like pi0/pi0.5 run a denoise loop at inference: 10 ODE steps from noise to action. Each step is a full forward pass through the action head. SnapFlow (arxiv 2604.05656) distills the 10-step teacher into a 1-step student — same architecture, trained to hit the teacher’s output directly.

Pi0.5 LIBERO-10 distillation results (v0.3.1, ~5-6 hours of training):

The student beat the teacher by 3.4 pp. This isn’t guaranteed — it’s a property of this particular distillation run on LIBERO — but it confirms that the 1-step student is not a quality-degraded approximation. It’s a different path through the same policy manifold, one that happens to be faster and more accurate on this benchmark.

The size drop is the Jetson story: 12.99 GB → 6.50 GB. At 6.5 GB, the FP16 student is below the 8 GB Nano’s total memory. The full-precision monolithic Pi0.5 export does not fit on an Orin Nano. The SnapFlow FP16 student does — barely, with ~1.5 GB remaining for the OS and inference buffers. That margin is tight enough that you need to be deliberate about co-resident processes (no full genie-ai-runtime stack alongside it, for example), but the model itself lands.

SnapFlow is also a 10× inference speedup on its own: 10 denoise steps per action → 1. On a Jetson where each denoise step is the bottleneck (not the VLM prefix), this matters more than any kernel-level optimization.

The TensorRT optimization path

Once the model fits in memory, TensorRT is the first optimization to apply. The path: export to ONNX, then serve with ONNX Runtime + TensorRT Execution Provider (ORT-TRT EP). ORT-TRT auto-builds an FP16 TRT engine on first serve and caches it — subsequent starts are ~1–2 s.

Reference numbers measured on A10G (sm_8.6), SmolVLA monolithic, batch=1:

5.55× speedup. What TensorRT does: it fuses the attention blocks (QK projection + softmax + V projection in one kernel), the LayerNorm and RMSNorm operations, and the denoise loop’s linear projections. For a VLA these are the dominant kernel patterns — large linear layers at batch=1. FP16 kernel fusion eliminates the memory round-trips that dominate at batch=1.

On Jetson Orin Nano, SmolVLA FP16 with ORT-TRT hits ~25 ms p50, ~20 ms with INT8. A 20 Hz control loop gives you 50 ms — SmolVLA clears it with 2× headroom. GR00T N1.6 (4-step DDIM, not 10-step) is faster per action; Pi0.5 SnapFlow 1-step at 6.5 GB is the boundary case.

The decomposed export splits the VLA into two ONNX sessions: vlm_prefix (vision + language → KV cache, fires once per instruction change) and expert_denoise (KV cache + state → action chunk, fires once per control step). CUDA Graphs on expert_denoise — the hot path — give an additional 3.03–3.76× speedup measured on Ampere. On Jetson AGX Orin (also Ampere, sm_8.7), CUDA Graph capture for expert_denoise succeeds and captures most of this gain.

The deployment gap: ONNX export correctness

Getting from a PyTorch checkpoint to a working ONNX on Jetson is not mechanical. The Pi0.5 monolithic exporter had 6 silent correctness bugs that survived in production ONNX for months until the v0.10.0 spine refactor caught them (issues #68–74):

• Wrong AdaRMSNorm gating — the adaptive modulation signal applied at the wrong layer boundary
• Wrong activation function — SiLU used where GELU-tanh was correct
• Wrong RMSNorm type — standard RMSNorm vs AdaRMSNorm mismatch in the DiT head

None of these caused shape errors. All of them caused wrong actions on real tasks. The verification harness now checks numerical parity against the PyTorch reference at export time:

Machine-precision parity. Anything looser is not a “good enough approximation” — it’s a policy that will behave differently on the robot. The verification step is mandatory, not optional, before you deploy on real hardware.

Engineering traps that bite every Jetson setup

Python version mismatch

lerobot (the training framework VLAs are trained with) requires Python 3.12. JetPack ships Python 3.10. The workflow is therefore split: ONNX export runs on a Python 3.12 desktop or cloud instance; the exported ONNX artifact is the portability layer that crosses onto the Jetson. You never run PyTorch or lerobot on the Jetson directly.

The practical implication: every model update requires an export step off-device. The ONNX is the artifact you version-control and deploy, not the PyTorch weights.

numpy<2 on aarch64

On JetPack aarch64, the NVIDIA Jetson AI Lab wheels require numpy<2. NumPy 2.x breaks ORT + torch on JetPack. Pin numpy<2 in the base environment before installing ORT — easy to miss, bites every fresh setup.

RecordWriter jitter in the control loop

Issue #48: the episode recorder’s flush() call was synchronous on the main event loop, introducing 50–100 ms jitter into a 100 Hz control loop. For a 20 Hz arm, a 50 ms flush delay is a full control step late. Fixed by moving the recorder to a background worker thread. Any I/O on the critical latency path of a real-time control loop needs async isolation — the same principle as avoiding synchronous cudaDeviceSynchronize() in the inference hot path.

TensorRT engine cache per JetPack version

The TRT engine is compiled for a specific GPU architecture and TRT version. An engine built on JetPack 6.0 does not load on JetPack 6.1. After any JetPack upgrade — including the JetPack 6.1 filesystem rebuild that IPv6/Thread support requires — the TRT engine cache must be invalidated and rebuilt. First-serve after a JetPack upgrade is ~2–5 minutes of recompile.

The optimization stack for Jetson

In order of applicability on Jetson:

1. SnapFlow distillation: 10-step → 1-step, pi0.5 FP16 from 13 GB → 6.5 GB — the prerequisite for Nano-class hardware, and a 10× denoise speedup
2. TRT EP (ORT-TRT): 5.55× baseline speedup from FP16 kernel fusion — turn on first, nearly free once ONNX is correct
3. CUDA Graphs on expert_denoise: 3× additional on the denoise hot path — verify capture succeeds on your specific JetPack and TRT version
4. INT8 quantization: ~20% additional latency reduction for models that survive accuracy validation at INT8
5. Episode KV cache: ~9× on repeated queries sharing a prefix — relevant for teleoperation tasks with stable instructions
6. Adaptive denoising early-stop: 58% latency savings on pi0, task-specific — profile before deploying
7. Triton fast-kernels (Pi0.5 only): 2.5× vs PyTorch compile on A100 — not yet available on Jetson, hardware-gated to sm ≥ 8.0 with full CUDA Graph support; falls back silently to ORT

The practical path for an Orin Nano targeting a 20 Hz arm: SmolVLA + ORT-TRT (item 2) + CUDA Graph on expert (item 3) lands at ~15–20 ms, clearing the 50 ms budget with 2.5× headroom. For an AGX Orin 32 GB running the SnapFlow Pi0.5 student: items 1–3 together should land at ~30–40 ms, clearing 20 Hz.

Why Jetson changes the inference problem

Running a VLA on Jetson is harder than running it on a cloud GPU in three ways that matter:

Unified memory is a shared budget, not dedicated VRAM. On a cloud A10G, the 24 GB VRAM is the VLA’s alone. On Jetson, the 8 GB unified pool is shared with the CPU, the OS, the vision capture pipeline, any LLM agent co-resident (genie-ai-runtime, genie-claw), and the TRT engine cache itself. The budget model I use for genie-ai-runtime applies directly here: OS (~1.5 GB) + VLA weights + TRT workspace + vision buffers = everything. There is no spill.

Hard latency deadlines, not throughput targets. A cloud serving engine queues requests and optimizes for throughput. A robot control loop fires at 20 Hz regardless — a missed deadline means the robot executes a stale action. p99 latency vs. the control loop period is the only metric that matters. At 20 Hz you have 50 ms. Any p99 overshoot is a real failure mode, not a percentile in a dashboard.

The action head is the hot path, not the VLM backbone. The flow-matching denoise loop runs at every control step (10 or 1 forward passes of the action head per action). The VLM prefix (vision encoder + language model) fires only when the instruction changes — which in normal robot operation is rare. On Jetson, where compute is limited, spending kernel budget on the VLM backbone and ignoring the denoise loop is optimizing the wrong thing. SnapFlow’s 10× denoise speedup is more valuable than any VLM kernel improvement on this hardware.

References

• tether — VLA deployment CLI
• SmolVLA (lerobot)
• pi0 (Physical Intelligence)
• GR00T N1.6 (NVIDIA)
• SnapFlow distillation (arxiv 2604.05656)
• genie-ai-runtime — Jetson LLM inference engine
• ONNX Runtime TensorRT Execution Provider