ort-rocket

AI disclosure

Except for the prior work it builds on, ort-rocket was developed by AI, primarily Claude Code (Opus 4.8). Human involvement was mostly limited to setting project goals and providing hardware access. This is a side project for curiosity's sake and comes with no guarantee of quality, accuracy, or update frequency.

About ort-rocket

An ONNX Runtime execution provider for Rockchip NPUs (validated on the RK3588) via the mainline rocket DRM-accel driver. It offloads the encoder of transformer vision models, meaning ViT backbones and their conv necks, to the NPU through the standalone rocket-userspace driver library. Everything it does not claim falls back to ONNX Runtime's own CPU kernels: deformable-attention sampling, detection heads, pooling heads and other host ops.

It builds as a runtime-loadable plugin execution provider .so, which stock ONNX Runtime loads via register_execution_provider_library. The NPU then appears as a selectable device, exactly like ONNX Runtime's other execution providers. Any ONNX Runtime application can load it and run an unmodified .onnx model.

Five model families are recognized and offloaded today, identified by graph topology rather than exporter-generated tensor names:

• RF-DETR (nano + base): the DINOv2 ViT backbone and the CSP feature projector. The deformable-attention decoder and detection head stay on the CPU.
• CLIP and SigLIP image encoders: the plain global-attention ViT. The pooling head stays on the CPU.
• SAM (Segment Anything) ViT-Det image encoder: windowed and global attention with on-chip decomposed relative-position bias, plus the conv neck. The prompt encoder and mask decoder stay on the CPU.
• Depth Anything v2: the DINOv2 backbone. The DPT depth head stays on the CPU.

A graph outside that set runs entirely on ONNX Runtime's CPU kernels, unchanged. So does one whose recognizable topology has been rewritten by graph optimization, described below.

.
  rocket-userspace/   # the driver library (dependency): git clone https://github.com/gregordinary/rocket-userspace
  ort-rocket/         # this project: the ONNX Runtime execution provider

This is a separate project from the driver. rocket-userspace, the librocketnpu dependency, builds and runs on its own, and this provider links it. The full per-family faithfulness and performance detail, the quantized-model behavior, the complete ROCKET_ORT_* knob table, and the matcher / implementation notes are in API.md.

Performance

The value here is throughput and CPU offload rather than single-stream latency. On the RK3588 the matmul is DMA/dispatch-bound rather than MAC-bound, and the Cortex-A76 is a strong baseline. So offloading even an encoder that is 90-100% of a model's compute yields at most a modest single-stream speedup, and for one model a loss. The NPU earns its place by running several streams concurrently while the A76 cores handle the host tail.

All figures warm, RK3588 @ 600 MHz, EP versus an 8-thread ONNX Runtime CPU EP on ORT_ENABLE_ALL. Every offloaded encoder is numerically faithful to the CPU. [HW sweep] Reproduce the single-stream column with tools/bench_ep.py and the pool column with tools/bench_pool.py.

Model	Encoder offloaded	Single-stream (xCPU)	Pool (xsingle-stream)	Faithfulness vs CPU EP
SigLIP-B/16	plain ViT, d=768, 196 tok	1.27x	n/a	last_hidden cos 0.9999864
CLIP ViT-B/16	plain ViT, d=768, 197 tok	1.22x	n/a	last_hidden cos 0.9999955
RF-DETR base	DINOv2 ViT-S, d=384, windowed	1.05x	~1.6x	COCO mAP 0.564 = 0.564 (200-img)
RF-DETR nano	DINOv2 ViT-S, d=384, windowed	1.03x	~1.6x	COCO mAP 0.5049 vs 0.5051 (500-img)
SAM ViT-B	ViT-Det, d=768, 4096 tok	1.03-1.11x	~1.71x	enc cos 0.9999995, mask-IoU 0.9998
Depth Anything v2 Small	DINOv2 ViT-S, d=384, global 1370 tok	0.78x	n/a	depth cos 0.9999999

Which models win is predictable. The NPU runs the projection GEMMs, the large-K matmuls, about 2.7x more efficiently per-MAC than attention. Attention's score matmul contracts over the head dimension (64) and materializes an ntok x ntok score matrix per head. So a model's single-stream envelope tracks its effective attention-to-projection ratio, ntok / 4d, counting windowed attention at its window length, not the width alone:

Model	ntok	d	effective ntok/4d	attention	single-stream
SigLIP / CLIP B/16	~196	768	0.06	global, short	1.22-1.27x
RF-DETR base	1601	384	~0.07	windowed 101	1.05x
SAM ViT-B	4096	768	~0.10	windowed 196 + 4 global	1.05x
Depth Anything v2 Small	1370	384	0.89	global 1370	0.78x

Short-sequence wide encoders (CLIP/SigLIP) win, windowed encoders sit at parity, and a long global-attention encoder (Depth Anything Small) is the first to lose single-stream. Widening the encoder helps (d=768 beats d=384) by raising the projection share, but it cannot shrink an attention cost that only sequence length controls.

Throughput is the headline. A single encode is partly host-bound, in score readback, softmax, layout and the CPU tail. So the NPU's value is running one process per stream, and P=4 peaks the 3-core NPU.

The best host-core placement is model-dependent. RF-DETR scales best with the host threads left unpinned (ROCKET_ORT_PIN=0), and SAM with a per-process ROCKET_CPU_AFFINITY spread across distinct A76 cores. The NPU pool scales ~1.6x on RF-DETR and ~1.71x on SAM over a single NPU stream. It also frees the A76 cores a CPU-only run would saturate.

The CPU pools too, so the NPU-vs-CPU number measured against the CPU's own best pool is narrower. SAM is 1.44x there, and RF-DETR's margin is thinner, because the CPU pools well on that workload. Reproduce the vs-CPU-pool ratio with tools/bench_pool.py. The encoder-only families (CLIP/SigLIP/SAM) also have a ~100% offload share, so no host tail dilutes the encoder's edge.

The remaining lever is a tiled, fused global-attention pass, an FA-2-style running softmax that keeps the ntok x ntok scores off the host DMA path. It would lift both the single-stream number and the pool ceiling for the attention-bound models (SAM's four global layers, Depth Anything's every layer). It is the obvious next contribution.

What runs on the NPU

The EP claims one encoder subgraph per recognized family and hands everything else back to ONNX Runtime:

Offloaded, in each case to the NPU:

• The ViT encoder, meaning the patch-embed conv, LayerNorm, multi-head self-attention, FFN and GELU. It is reused across families.
• RF-DETR's CSP projector.
• SAM's windowed and global attention, with on-chip decomposed relative-position bias, and its conv neck.
• Depth Anything's four intermediate encoder taps.

Host, on ONNX Runtime's CPU kernels:

• Deformable-attention GridSample sampling and the DETR detection head (RF-DETR).
• The attention-pool and cls-token pooling heads (SigLIP/CLIP).
• The prompt encoder and mask decoder (SAM).
• The DPT reassemble and fusion head, and its align_corners resizes (Depth Anything).

None of these has an NPU route. GridSample and align_corners resize both need a gather the RK3588 NPU lacks, and each is a small share of its model.

Graph optimization must be off. The matcher keys on the exported op topology, so the session must run with ORT_DISABLE_ALL. ONNX Runtime's Gemm, Gelu and LayerNorm fusions engage from ORT_ENABLE_BASIC up. They rewrite the encoder block into a topology the matcher does not model, and the model then runs entirely on the CPU. A graph that is not one of the five families, or whose optimizations were left on, is left on the CPU, never mis-claimed.

Quantized models. The EP consumes int8 and int4 QDQ ONNX, ONNX Runtime's quantize_static format. By default it dequantizes the weights to fp16 and runs the same fp16 datapath, which is faithful and the fast path on this stack.

A native W8A8 int8 datapath (ROCKET_ORT_INT8=1) runs the encoder GEMMs as int8 x int8, and is ~1.8x slower. Mainline rocket has no on-chip int32 accumulation, so every int8 matmul reads its int32 result back to the host. It is a faithfulness and experimentation mode rather than a speed lever. Detail and the accuracy numbers are in API.md.

Architecture

A plugin execution provider, meaning an out-of-tree EP shared library:

• ONNX Runtime dlopens the library and calls the exported CreateEpFactories and ReleaseEpFactory entry points. The factory reports the devices it supports. ONNX Runtime then asks the EP, via GetCapability, which subgraphs it can run, and hands each claimed subgraph back through Compile. Whatever the EP does not claim runs on ONNX Runtime's own kernels.
• The provider calls ONNX Runtime only through the OrtApi and OrtEpApi function tables passed to CreateEpFactories. So the .so does not link libonnxruntime, and needs only the ONNX Runtime headers to build.
• It links librocketnpu, the rocket-userspace project, for the on-NPU operators.
• A matcher registry tries each family signature in turn during GetCapability and claims the first that matches. The offloaded subgraph is therefore identified by graph structure rather than by tensor names. The signature is the attention and MLP pattern, LayerScale, and the patch-embed and neck convs. Family-distinctive nodes complete it, such as SAM's relative-position Einsums or Depth Anything's four-tap DPT exit. A re-export that renames the interior still matches.

Requirements

• An RK3588 board on a mainline kernel carrying the rocket DRM-accel driver, with /dev/accel/accel0 present (lsmod | grep rocket).
• The sibling rocket-userspace driver library, cloned next to this repo or installed as a rocketnpu package.
• An ONNX Runtime build exposing the plugin-EP API (register_execution_provider_library).
• Privilege to open the accel node. Run with sudo -E (the -E preserves the ROCKET_* env knobs plain sudo strips).

Clock. The NPU boots at 200 MHz and the EP is correct there, and every figure above is at 600 MHz. Apply the patches/rocket clock patch and load the module with rocket_npu_clk_hz=600000000.

Build

Requires a C++17 compiler and CMake. The ONNX Runtime C/C++ API headers for the target ONNX Runtime version are vendored under third_party/onnxruntime/include. Override with -DORT_HEADER_DIR=<path> to build against a different release.

cmake -S . -B build
cmake --build build -j
# -> build/libonnxruntime_rocket.so

Use from an application

import onnxruntime as ort

ort.register_execution_provider_library("rocket", "/path/to/libonnxruntime_rocket.so")
dev = next(d for d in ort.get_ep_devices() if d.ep_name == "rocket")

so = ort.SessionOptions()
# The topology matcher keys on the raw exported ops, so disable ONNX Runtime's graph
# optimizations: its Gemm/SkipLayerNorm/Gelu fusions rewrite the encoder into a topology
# the matcher does not recognize (even ORT_ENABLE_BASIC fuses every MatMul+Add into Gemm).
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
so.add_provider_for_devices([dev], {})
sess = ort.InferenceSession("model.onnx", so)
outputs = sess.run(None, {"input": x})

The provider loads for any ONNX Runtime application and runs an unmodified .onnx model. The NPU offload activates for a recognized encoder, and any other graph runs on the CPU.

For a CLIP, SigLIP or Depth-Anything variant whose attention head dimension is not 64, set ROCKET_ORT_SIGLIP_HEADS to the head count. The EP reports the required value if it cannot infer it.

For multi-stream throughput, run one process per stream. The best host-core placement is model-dependent: RF-DETR unpinned via ROCKET_ORT_PIN=0, and SAM a per-process ROCKET_CPU_AFFINITY spread. See API.md for the per-model recipe and the full knob table.

The rocket NPU stack

This provider is one frontend of an open source stack for Rockchip NPUs, three userspace projects plus a set of optional kernel patches:

• rocket-userspace (librocketnpu): the userspace driver, matmul, and on-NPU op library. It is the dependency, and usable on its own.
• ort-rocket (this project): an ONNX Runtime execution provider for transformer vision models.
• ggml-rocket: a ggml backend .so for llama.cpp / whisper.cpp, linking the same driver.
• tflite-rocket: a TFLite external delegate for detection models, linking the same driver.
• patches (rocket/ scope): optional out-of-tree kernel-module patches for clock, voltage and IOMMU. They raise the NPU clock from its 200 MHz boot default to 600 MHz, and the figures here assume them.

License and credits

ort-rocket is GPL-3.0-or-later. It links the GPL-3 rocket-userspace driver library, whose NPU register headers are the GPL-3 reverse-engineering by Jasbir Matharu. It targets ONNX Runtime (MIT) as an out-of-tree plugin execution provider. The ONNX Runtime C/C++ API headers vendored under third_party/onnxruntime/ are MIT-licensed (Copyright Microsoft Corporation). It builds on the rocket-userspace driver for the NPU path.