LIBERO is the benchmark I keep coming back to. It’s 130 robot manipulation tasks across five suites (spatial, object, goal, scene10, and scene90), each with human demonstrations and a language goal, and it’s the standard testbed for language-conditioned imitation learning. The reference implementation runs on robosuite with CPU MuJoCo, one environment per process. Ten 600-step episodes take 92.6 seconds that way, one environment at a time, which is 38.4 environment-steps per second.

I wanted that on a GPU. Not a datacenter GPU; a $700 one.

The reason is rollouts. Behavioral cloning needs them. You run the policy, watch where it fails, and generate labels from the failures. The counterfactual-regret work I wrote about in July needed 18,000 labels, and the CPU vectorized environment would have chewed on that for hours. The AMD Instinct blog used the same trick for XArm6 sim2real, on MI300X and MI355X: run hundreds of worlds at once on the GPU with MuJoCo Warp. I co-wrote that post; it was my project over the summer at AMD. I wanted to know if the consumer card could do it.

There is a second reason, and it is the one that surprised me. Generating the demonstrations for one LIBERO suite, 10 tasks, took me 8 to 9 hours on a 16-core/32-thread consumer CPU. All five suites is a multi-day job. The same work on the GPU is minutes, which turns dataset generation from something you do once into something you can iterate on. That matters if you are producing rollouts, counterfactual labels, or synthetic demos rather than just evaluating a policy.

So I ported all 130 LIBERO tasks to MuJoCo Warp, and as far as I can tell it’s the first MJX+Warp port of LIBERO. I trained and evaluated a BC transformer on it too, on an RX 9070 XT, which I believe is the first time that’s been done on RDNA 4 outside AMD. It works. Some of it doesn’t. I’ll get to both.

Everything is at github.com/poad42/libero_mjx.

The GPU

The RX 9070 XT is a $700 consumer card: RDNA 4, gfx1201, 16 GB of VRAM, 304 W. The FLOPS don’t matter much here; the 16 GB and the ROCm stack do. Warp compiles HIP kernels for gfx1201 at runtime, JAX sees the card as rocm:0, and PyTorch sees it as cuda. Same card renders the scene and runs the policy.

16 GB is the annoying part. With two 128x128 cameras per environment, the Warp renderer allocates ray-tracing buffers per world, and my eval runs out of memory at 100 environments. Fifty fit, and give 66.4 environment-steps per second. It’s a memory wall, not a compute wall; the physics kernels are half asleep at that batch size.

What the port actually is

Three pieces.

The 130 tasks come from 131 MJCF XMLs pulled out of robosuite. The XMLs are scene structure only; the meshes and textures stay in the LIBERO and robosuite packages and get resolved by path. I wrote a small resolver so the same XMLs load in a container, on my workstation, or on a cluster without rewriting 131 files by hand.

LiberoEnv wraps any of the 130 tasks behind one reset/step, batched with jax.vmap, with a flax.struct.PyTreeNode holding the mjx.Data. The controller is a JAX port of robosuite’s OSC_POSE: Jacobian at the gripper site, densify the mass matrix, solve for the torques that push the end-effector to a target pose. That part was mostly transcription.

WarpRenderer does the images, on the GPU, batched. That part was not transcription.

The Warp port

Warp is NVIDIA’s. The AMD HIP/ROCm port for Strix Halo was open-sourced by Tomas Thoresen, validated on gfx1151. It works great on RDNA 4 as well, which is what this whole post runs on. If you want Warp on an AMD GPU, that is the branch to clone: AMD-Ecosystem/warp, or Tomas’s own fork.

The Instinct blog from the intro is this same stack on Instinct MI300X and MI355X.

How fast it goes

The first working Warp path was slower than the CPU. I kept the OSC controller in JAX and ran 25 physics substeps per control step, and got 18.9 environment-steps per second at 10 environments. The JAX physics step launches dozens of small kernels, and at 10 worlds the launch overhead eats everything.

Cutting the substeps from 25 to 5 (a 0.01 s simulation timestep instead of 0.002 s) cut the step time by 2.9x and got me to 40.9 env-steps/s. Then I ported the OSC controller from JAX to Torch and called mujoco_warp.step() directly, which removed the JAX dispatch entirely: 51.7 env-steps/s, 1.35x the CPU, at 60% average success across three seeds.

The 10-environment numbers are launch-bound. When I profile physics alone and crank the batch, the card wakes up: 1,568 env-steps/s at 768 environments, and 17,137 at 2,048, where the solver finally becomes compute-bound. The July post quoted 9,838 steps/s across 1,024 environments, a 49x speedup over the CPU SubprocVectorEnv, which is what made those 18,000 labels take 7 minutes instead of hours.

For a lighter scene, the state-based PandaPickCube from MuJoCo Playground, the same card does 212,445 env-steps/s at 2,048 worlds, 2.53x the JAX backend. Brax PPO on that task runs 2.16x faster on Warp (64.7 s vs 140.0 s for 2 million timesteps).

The BC transformer

The policy is LIBERO’s BC transformer: a ResNet image encoder, a 4-layer temporal transformer (6 heads, embedding 64, sequence length 10), and a 5-mode Gaussian-mixture head. It takes two 128x128 cameras (agentview and robot0_eye_in_hand) plus joint and gripper state, and spits out a 7-DoF action.

Fifty epochs on spatial task 0 gets 50% success on CPU episodes. On Warp, with the renderer fixes below, the same checkpoint gets 42.5% across four seeds (50, 40, 30, 50). Without the fixes, it gets 0%.

Across the spatial suite the 50-epoch checkpoints average 51.5% (task 8 is 95%, task 2 is 0%). The object suite averages 56.5%.

The interesting number isn’t 42.5%. It’s that a policy trained on CPU pixels works at all on GPU pixels, which means the renderer is close enough that the policy recognizes what it’s looking at.

The renderer was the whole fight

The physics port was mechanical. The renderer was not. A BC policy is a fixed function of the pixels it was trained on, and the Warp ray tracer is not MuJoCo’s CPU renderer. Every difference in shading, orientation, or texture moves the policy off its training distribution. My first Warp eval scored 0%, and it took a while to figure out why.

The first one was dumb. The Warp renderer and MuJoCo’s OpenGL renderer both use a bottom-left origin; robosuite’s training images are top-left. The policy was looking at an upside-down scene. Flip the image and it works.

The second was a constant. The Warp megakernel kept 30% of the diffuse and specular light on shadowed geometry (visible = 0.3). MuJoCo zeroes it and gives shadowed geometry ambient only. Changing it to 0.0 dropped the Warp-vs-EGL RMSE from 34.4 to 22.2.

The third was brightness. The ray tracer is about 85% as bright as the CPU renderer, and not uniformly (0.856 at center, 0.929 at the edges), so it’s not a clean exposure scale. A 1.15x multiplier on the output moved average success from 30% to 42.5%. That one fix did more than the kernel changes combined. 1.10, 1.15, and 1.20 all gave the same result on the seed I checked, so the exact value isn’t critical.

The fourth was cube maps. Warp sampled every texture as 2D. MuJoCo supports cube textures stored as six faces, and a cube material sampled as 2D reads a horizontal line across the face strip. LIBERO’s table uses a cube texture for its wood grain, so the table rendered as a flat tan smear. Reproducing MuJoCo’s cube texgen, for both the 6:1 strip layout and the square texture MuJoCo repeats on all six faces, cut the eye-in-hand RMSE from 29.9 to 20.4 and brought the grain back.

CPU (EGL) Warp

After cube-map fix Before cube-map fix

Drag to compare. Eye-in-hand camera, spatial task 0. Before the fix the table is flat tan and the stove has no burner rings; after, the wood grain and rings are back. Eye-in-hand RMSE went from 29.9 to 20.4.

CPU (EGL) Warp

After cube-map fix Before cube-map fix

The same comparison for the agentview camera.

The fifth one cost me a day. The renderer fixes rewrite mujoco_warp’s source on disk, so they have to run before mujoco_warp is imported. I had them running after. The file changed under the compiler, and warp resolved the nested render megakernel against stale line numbers, which fails as “kernels cannot return values”. Patch first, import second.

The gap that’s left between Warp (42.5%) and CPU (50%) is the residual brightness mismatch, plus the ray tracer having no alpha blending for the transparent end-effector target geoms. Both are in the mujoco_warp render kernel, not in my code.

What I didn’t finish

Two things, and I want to be straight about them.

The loop isn’t closed. The Warp path collects data; it doesn’t train on that data in the loop. The 50-epoch policy trains offline on the LIBERO HDF5 demos in PyTorch, no simulator in the process. The Warp rollouts generate the counterfactual labels the repair head learns from and evaluate the trained policy, but the gradient step never sees the renderer. Rendering in the training loop is a different problem from rendering for evaluation, and I haven’t done it for LIBERO.

CPU rendering is broken on my reproducible base. I rebuilt the container on the TheRock AlmaLinux image, and its Mesa (23.1) is too old for gfx1201 EGL. The GPU Warp renderer doesn’t care, but the CPU (robosuite) renderer does, so eval_bc.py and the CPU half of render_comparison.py need a newer Mesa. I run those on a host that has one. The comparison image above came from there.

Where the in-loop version already exists

Rendering inside the training loop is not new, and I didn’t invent it. MuJoCo Playground does it with DeferredVisionWrapper: the env builds a Warp render context once (mjx.create_render_context) and, every step, calls mjx.refit_bvh, mjx.render, and mjx.get_rgb under jit and vmap. Physics, rendering, and the policy end up in one compiled program. PandaPickCubeCartesian is the reference env.

That path runs on one R9700, the 32 GB big brother of the 9070 XT, same RDNA 4 silicon. PandaPickCubeCartesian trained to 15,052,800 steps with the renderer in the loop, and the XArm6 pick-place vision runs reached 150,200,320 steps. The XArm6 run was a single card at 1,024 environments, and it did roughly 5,900 env-steps/s over 6.6 hours with rendering on the critical path. Those are PPO runs on the XArm6, not BC on LIBERO, but they show the pattern works on this hardware.

How this actually went

It wasn’t a straight line, and there were long gaps where I was working on other things. Roughly:

• May to June: the counterfactual-regret and flow-matching work. That’s where the first 10 LIBERO_SPATIAL tasks got ported, because I needed fast rollouts for gate labels.
• July 7: the CFR blog post.
• July 8 to 14: the real port. All 130 tasks, the parallel init, the Warp eval pipeline.
• July 29 to 31: a long push on performance and rendering. The Torch OSC port, the profiling, the shadow and brightness fixes, the CPU-vs-Warp comparison images.
• August to September: other projects. The port sat.
• October: I came back, rebuilt the container on public sources, and fixed the cube-map and patch-ordering bugs.

So the repo has bursts of work with gaps in between. That’s just how it went.

Making it reproducible

The port is public now, and the container builds from public sources only. TheRock manylinux base, ROCm from the pip SDK, PyTorch and JAX from the TheRock index, warp from the public cu-basil/warp fork, MuJoCo pinned on PyPI, and LIBERO from the public cu-basil/LIBERO fork. No private wheels, no local docker commit.

Two warp bugs I hit on the way, the LLVM 23 build failure and the nested-kernel source extraction, are upstream in cu-basil/warp now, so the container build applies no patches at all. A third is a runtime workaround in libero_mjx/warp_kernel_patch.py, and I will explain why it is not upstream yet. docker/verify_stack.py checks the GPU backend, the MJX-Warp wiring, and a reset-and-step on the real task XMLs.

The CCD kernel name drift

The example below surfaced a third warp bug. mujoco_warp’s CCD path calls wp.get_suggested_block_size, which loads the module at the default block_dim of 256, and then launches the same kernel at block_dim=64. For a nested module="unique" kernel, warp hashes the resolved static expressions, and adj.build() repopulates those when the module is compiled for the second block_dim. So the live kernel’s hash drifts away from the name the module was cached under, and ModuleExec.get_kernel_hooks looks up a mangled name the metadata does not have:

KeyError: 'ccd_kernel_builder__locals__ccd_kernel_<hash>_cuda_kernel_forward_smem_bytes'

It is intermittent, depending on hashing order, which is why it took a while to pin down. warp_kernel_patch.py wraps get_kernel_hooks so that when the current mangled name is missing it recovers the name the module was actually built with from the metadata. The match is exact (a single <key>_<8 hex> entry), so a module that legitimately holds several kernels under one key still fails loudly. It is a no-op where the hash does not drift, and it works on both the warp 1.13 and 1.17 APIs used here. It is a workaround, not the fix; the real fix belongs in warp, where a kernel’s hash should be stable across block_dim variants.

Try it

There’s a minimal example in the repo, examples/libero_quickstart.py. It loads one task, runs a batched rollout on Warp, renders both cameras on the GPU, and prints the throughput:

python examples/libero_quickstart.py --suite spatial --task-id 0 --n-envs 8 --steps 20

[quickstart] 100 init states from the LIBERO init file
[quickstart] spatial task 0, 8 envs, 128x128 cameras, impl=warp
  step   10      19.8 env-steps/s  success=0.00
  step   20      20.0 env-steps/s  success=0.00
[quickstart] 8 envs x 20 steps in 8.0s = 19.9 env-steps/s
[quickstart] success: 0/8 = 0%
[quickstart] wrote /tmp/frame.png (agentview | eye-in-hand)

Zero actions, so 0% success is the expected result; the point is that the env steps and the renderer work. At 8 environments the number is launch-bound, the same reason the 10-environment eval above is slow. --batch-bench sweeps the batch size and shows the curve:

  n_envs   env-steps/s   ms/step
      16          38.5    415.61
      64         151.9    421.27
     256         302.6    846.02

The --checkpoint flag runs a trained BC policy instead of zero actions. The 50-epoch spatial task 0 checkpoint (21.6 MB) ships in the repo, so you don’t have to train before trying it; it still needs that task’s demo HDF5 for the shape metadata. --save-image writes one agentview | eye-in-hand frame. Inside the container it’s:

./scripts/docker_run.sh python examples/libero_quickstart.py --suite spatial --task-id 0 --save-image /tmp/frame.png

Takeaways

• A $700 consumer GPU runs all 130 LIBERO tasks in parallel. Physics alone hits 17,137 env-steps/s on LIBERO at 2,048 worlds, and 212,445 on a lighter Playground scene. The full loop, with rendering and the BC policy, runs 51.7 env-steps/s at 10 environments.
• It is also a dataset generator. One suite of demonstrations took 8 to 9 hours on a 16-core/32-thread CPU; the full 130 tasks is a multi-day job. On the GPU it is minutes, so you can iterate instead of running it once.
• The renderer was the hard part, not the physics. Four fixes, and the brightness one alone was worth 12.5 points of success rate.
• The loop isn’t closed here. The Warp path collects data; training is offline. In-loop vision training works on this hardware (MuJoCo Playground, and the cluster’s 150M-step XArm6 runs), and porting it to LIBERO is the next thing I’d do.
• Everything builds from public sources. If you want to try it, the repo is poad42/libero_mjx.

Work with me

I’m looking for a role starting May 2027, in ML systems, GPU/edge deployment, or robotics infrastructure. If you’re hiring, or just want to talk about any of this, my email is in the footer. The code and the rest of my projects are public.

References

• Repairing a Frozen Visuomotor Policy With Counterfactual Regret Labels & Flow Matching (2026-07-07), where the port started.
• Unlocking Sim2Real for a Robotic Arm with RL Accelerated by AMD Instinct GPUs (AMD ROCm blog, 2026-09-29), the same Warp stack on MI300X and MI355X.
• Tomas Thoresen’s fork of Warp, where the Strix Halo (gfx1151) AMD port was developed: https://github.com/tomasthoresen/warp, and the public AMD fork: https://github.com/AMD-Ecosystem/warp
• MuJoCo Playground, and the DeferredVisionWrapper / PandaPickCubeCartesian vision path: https://github.com/google-deepmind/mujoco_playground
• MuJoCo Warp: https://github.com/google-deepmind/mujoco_warp
• LIBERO: https://github.com/Lifelong-Robot-Learning/LIBERO

This year I went further. Instead of fighting the black box, I opened it. I reverse-engineered the register commands the NPU actually executes, then built a compiler and runtime that talk to the hardware directly. No closed SDK. No vendor runtime. The result is OpenNPU, an MIT-licensed toolkit that runs GPT-2 and SigLIP from PyTorch, ONNX, and JAX, and that you can extend with your own operators.

This post is the story of that year. For the technical details, the wire format, the compiler pipeline, and the measured numbers, see the companion report: OpenNPU v1.0: an open compiler and runtime for the RK3588 NPU.

Last year we worked around the closed toolchain. This year we opened it.

Last year: the vision encoder hack

The RK3588 is a $100 chip with a 6 TOPS NPU. TOPS means trillions of operations per second, and 6 TOPS is a lot for the price. The catch is how you reach it. Rockchip ships the NPU behind two closed pieces of software: a compiler that turns a model into a binary file, and a runtime that loads that file and drives the hardware.

I wanted to run the SigLIP vision encoder from SmolVLM. The closed compiler refused. It threw an undocumented error, REGTASK Overflow (0xe010), and gave up. The model’s attention matrices were too big for the NPU’s small on-chip scratch memory, and the compiler had no way to handle them.

So I did the surgery myself. I sliced the attention into small tiles that fit the scratch memory, injected a tiny dummy operation to stop the compiler from fusing my tiles back into one giant block, and scaled the activations so the model’s wide value range survived quantization. It worked. The encoder ran 15x faster than the CPU, and the output matched the reference.

That post ended there. The model ran. But the approach had a ceiling.

The itch: a black box you cannot extend

The gray-box fix worked for one model, but it was hand surgery. Every new model needed the same treatment: find the constraint, tile around it, trick the compiler. The toolchain itself stayed opaque. I could not add an operator it did not know. I could not change how it scheduled work. I could not ship the result, because the whole thing depended on a closed compiler and a closed runtime.

The closed path hides the hardware behind two opaque layers. The open path talks to it directly.

The deeper problem was that reverse engineering the hardware limits was not the same as understanding the hardware. I knew the NPU had a small scratch memory. I did not know what bytes it actually executed. To build something reusable, I needed the format underneath.

Opening the box

The NPU is a computer. Like any computer, it runs instructions. Rockchip does not document those instructions, but the runtime has to issue them, and the runtime is a file on disk. I captured the bytes the runtime sent to the NPU and decoded the format.

The format turned out to be simple: the NPU is programmed by writing eight-byte register-command entries, and the runtime submits a chain of command blocks. Once the format was known, the closed runtime became optional. I could write the bytes myself and submit them through the public DRM ioctl interface. The report walks the format entry by entry.

Building the compiler

With the format in hand, I built a compiler. It takes a model graph, lowers it through a small IR, and emits the register commands. The whole thing is pure Python, a few hundred lines, with no dependency on LLVM or MLIR.

The compiler has one important trick. The NPU has no general matrix multiply instruction, so the compiler implements Y = X @ W as a one-by-one convolution, which the hardware calls a CNA descriptor. That single trick is what lets any transformer run: attention and the MLP layers are just chains of matrix multiplies. The report covers the descriptor and the pipeline in detail.

The framework: PyTorch, ONNX, and JAX

The compiler alone is not a toolkit. A toolkit is something you can call from the framework you already use. So OpenNPU exposes the same path through three front ends.

PyTorch gets the NPU as a device. You build a model with torch.nn, move it to the NPU, and run it. ONNX gets a runner that walks an exported graph and dispatches each node to the NPU or the CPU. JAX gets a real PJRT plugin, the same interface JAX uses for GPUs and TPUs, so jax.jit dispatches to the NPU and jax.devices() lists it as npu:0.

The JAX plugin is the most involved. It is pinned to a specific PJRT API version, because JAX reads the plugin’s function table at the offsets of the exact version it was compiled against. It parses the model bytecode JAX sends it, recognizes a small set of operations, and runs them on the NPU. The fast matmul path is exposed as a custom call that skips the parsing entirely.

Extending it

The point of a framework is that you can add to it. OpenNPU documents two routes for a new operator.

For a matrix-multiply-shaped operator, use the CNA descriptor path. The descriptor is parameterized, so a new shape or a new fused pattern is a configuration change, not a rewrite. For an elementwise or activation operator, generate a C template and let the plugin dispatch to it. Either way, the operator becomes callable from PyTorch and JAX once it is in the toolkit.

That is the difference from a one-off hack. Last year, adding an operator meant redoing the surgery. This year, it means adding a case to a table.

The results

GPT-2 decodes at 36 tok/s with KV caching. SigLIP runs at 1.1 img/s, or 2.7 img/s across three cores.

The GPT-2 decoder output matches PyTorch token for token. The SigLIP encoder matches the HuggingFace reference to cosine similarity 0.9999. Forty of the forty-one op and precision combinations pass on board. The report has the full comparison table and the honest hardware limits.

What it means

Reverse engineering the RK3588 NPU is not new. mtx512/rk3588-npu did it first, in early 2024, on the older 5.10 kernel, and the write-up is worth reading. Their approach was different: they derived the CNA and DPU register descriptors from the chip’s technical reference manual and ran matmul on the NPU’s internal SRAM. The register and submit formats changed in the newer 6.1 driver, so the descriptor here was re-derived on this kernel rather than ported; the older code does not run correctly on 6.1. The CNA descriptor matmul is one piece of that work, not the whole of it. The larger part was capturing the register-command templates and decoding the wire format, which is what the compiler and the forty op and precision combinations rest on. The contribution is the framework built on top of that: a compiler and runtime that work from PyTorch, ONNX, and JAX, that cover forty op and precision combinations, and that you can extend with your own operators.

The goal is a toolkit that makes the NPU easier to work on, not a claim about who did the reverse engineering first. The code is MIT licensed at github.com/poad42/opennpu_rk3588. The CNA descriptor matmul builds on prior work by mtx512/rk3588-npu, which is credited in the repository. If you have an RK3588 board, the examples are a good place to start.

This is the technical background. It covers the hardware contract, the wire format, the compiler pipeline, and the runtime.

The hardware contract

The RK3588 NPU is a 6 TOPS device with three cores. Rockchip exposes it through a DRM render node at /dev/dri/card1. Six ioctls cover the whole accelerator:

Three properties of this hardware dominate every design decision.

The NPU has no on-chip SRAM. The RK3576 NPU has on-die SRAM; the RK3588 does not. Every operation moves data DRAM to NPU to DRAM. That makes the accelerator DMA-bound. Measured throughput on the matmul path is about 1 GB/s, and the model’s fastest full-context forward sits at that wall.

Buffer addressing is a per-file-descriptor IOMMU pool. Each process gets its own virtual address window into the NPU’s physical memory. The runtime allocates buffers against that window, and some paths are position-sensitive: the scratch buffer must sit at the top of the pool, and a submit is rejected when scratch falls below a fixed address. One consequence is that the hardware runs one large matmul scratch at a time; the 3.5 MB scratch window fits one matmul’s working set, not three concurrent ones. The three cores are therefore data-parallel across independent requests, not model-parallel within a single matmul.

Third, the native numeric types are int8 and float16. There is no fp4. Every op and every weight layout has to fit one of those two.

The wire format

The NPU is programmed by writing register-command entries. Each entry is eight bytes: a 16-bit register, a 16-bit value, and a 32-bit tag.

One 8-byte entry, a 640-byte command block, and the two structs that drive a submit.

A 32-bit DMA address does not fit the 16-bit value field, so it is split across two fields. The decode is dma = (tag32 & 0xffff) << 16 | val16. A command block holds 80 entries (640 bytes). Of those, 69 are real register writes and one is the chain pointer. Entry 69 holds next_offset; a zero means the chain ends. The NPU follows the chain and stops at the end, so the runtime must submit exactly the number of tasks the chain forms. A job that claims more tasks than the chain runs makes the kernel wait out the full count.

Two structs drive a submit. The task struct is 40 bytes and carries flags, op_idx, enable_mask, int_mask, int_clear, int_status, regcfg_amount, regcfg_offset, regcmd_addr. The submit struct is 104 bytes and carries the task-object pointer, the IOMMU domain id, and the per-core configuration. The three cores are addressed as 0x1001, 0x2001, and 0x3001.

The compiler pipeline

The compiler in src/opennpu/ is a pure-Python MLIR-style flow. It does not link LLVM or MLIR; it is a small dataclass hierarchy that emits MLIR text and lowers it. That keeps the whole lowering readable in a few hundred lines and makes the project genuinely dependency-free.

From a model graph to the NPU: importer, IR, two lowering passes, register-command bytes, and the DRM runtime.

The path is:

ONNX importer -> IR -> TosaToNPULowering -> NPUToRegCmd -> runtime

• The importer reads an ONNX graph and builds the IR.
• TosaToNPULowering maps TOSA ops to rk3588_npu ops. The dialect defines over thirty op types: elementwise add, subtract, multiply, divide, relu, sigmoid, tanh, gelu, silu, exp, sqrt; linear matmul, gemm, batch matmul; conv2d, depthwise, transposed conv; pooling; norm; fused conv-relu and add-relu patterns; reshape, transpose, concat, slice, pad; reduce; and softmax and SDPA.
• NPUToRegCmd lowers each op to the register-command bytes through a codegen pass. It carries the hardware constants and the DMA encoding from the section above.
• The runtime allocates buffers, patches the addresses, and issues the submit.

Three codegen strategies exist and coexist. The capture-based path keeps captured reference templates and patches their DMA addresses at the known entry points. The pure-synthesize path generates the command bytes from the reversed constants, with no reference file at all. The MLIR-style path drives the same synthesizer through the dialect. All three produce the same regcmd bytes; the pure-synthesize and MLIR paths have no dependency on the closed toolkit.

The supported operator set was not assumed. It was taken from Rockchip’s RKNN Compiler Support Operator List (v1.5.2), which lists int8 and float16 as the native types.

Matmul through the convolution engine

The NPU has no general GEMM primitive. The compiler implements Y = X @ W as a one-by-one “direct convolution,” which the hardware calls a CNA descriptor.

Y = X @ W runs as a CNA descriptor. A CNA matmul uses five buffers.

A CNA matmul uses five buffers: weights, input, output, the task array, and the register-command array. Two constraints matter. When the inner dimension exceeds 768, the weights must be pre-tiled to fit the descriptor. And the descriptor supports int8, int16, and fp16; there is no fp4.

The matmul runtime in C, cna_matmul.c, keeps weights in a persistent buffer and streams only the input on each call. That cached-weight path is one half of the single-token speed. The KV cache is the other half: the decoder reads one token’s worth of memory instead of reprocessing the full context. Together they put single-token decode at 28 ms per token. The same file exposes attention on the NPU by loading K and V into descriptor slots, and layer norm and an erf approximation for GELU. It also exposes the three-core path through a core mask.

The framework front ends

Three front ends reach the same runtime.

ONNX uses onnx_runner.py, which walks the ONNX graph and dispatches each node to either the NPU or the CPU. PyTorch uses torch_npu.py, which presents the NPU as a device. JAX uses a real PJRT C plugin.

The JAX plugin, libpjrt_npu.so, is the most involved of the three. It is pinned to PJRT API version 0.112 and jaxlib 0.10.2. The pin matters: jaxlib reads the plugin’s function table at the offsets of the exact version it was compiled against, so the plugin must export that table layout. The header for 0.112 comes from the pinned XLA commit, and with it jax.devices() reports the NPU as npu:0.

The plugin does not link the MLIR/XLA compiler. Instead it parses the raw stablehlo bytecode that JAX sends to Client_Compile, using a small varint reader. It accepts a strict linear chain of recognized ops: dot_general becomes a matmul, tanh and maximum become the corresponding activation, and an add or multiply chain becomes the elementwise op. The fast matmul path is also exposed as a custom_call that bypasses HLO parsing entirely. That is the path examples/jax_matmul.py and the 28 ms-per-token decoder use.

Buffer lifecycle uses the same raw-DRM path as the Python runtime. The plugin opens the device and pre-allocates the weight cache before the first host-to-device copy, so the buffers sit at the top of the address window and one open file can be reused. Host to NPU copy goes through BufferFromHostBuffer and Buffer_ToHostBuffer.

The C kernels and the hybrid split

The compute-heavy work runs as C kernels compiled with unrolled loops. Two forwards exist.

Matmuls go to the NPU; attention, activations, and layer norm run on the CPU.

lm_forward.c runs GPT-2 and LLaMA-family models from one file, switched by flags for RoPE, SwiGLU, and RMSNorm. All matmuls go to the NPU through the cached CNA path. Attention, the activations, and layer norm run on OpenMP threads over the CPU’s A76 cores. vision_forward.c is the same idea for the SigLIP ViT: all 48 matmuls on the NPU, attention, GELU, and layer norm on the CPU.

The split is deliberate and measured, not ideological. Layer norm stays on the CPU because a submit costs about 0.5 ms and the CPU computes it in about 0.05 ms for a 768-wide activation. The NPU is memory-bound, so small operations do not amortize the cost of a submit. The compiler routes by cost.

Correctness and limits

The claims above are verified on an Orange Pi 5 Max (RK3588, Armbian, kernel 6.1.115-vendor-rk35xx).

The SigLIP encoder runs all 48 matmuls on the NPU. Its output matches the HuggingFace reference to cosine similarity 0.9999. The GPT-2 decoder output matches PyTorch token for token. The NPU runs at 98 percent load on the vision encoder. Forty of the forty-one op and precision combinations pass on board, covering the elementwise ops, the activations, matmul, softmax, concat, transpose, and layer norm across fp16, w8a8, and w16a16i.

The honest limits are the hardware’s. The NPU is DMA-bound around 1 GB/s. One matmul scratch fits the address window, so true model-parallel across the three cores is not possible. Prompt caching is off for the hybrid attention architecture. These bounds are measured, and they are the reason the runtime is a hybrid that puts the memory-bound work on the CPU.

How to use it and extend it

The repo runs a set of examples.

pip install -e .
python examples/gpt2_generate.py
python examples/vision_encoder.py --framework torch
python examples/vision_encoder.py --framework jax

Two documented routes exist for a bespoke op. For a matmul-shaped op, use the CNA descriptor path documented in the CNA guide. For an elementwise or activation op, generate a C template and let the plugin dispatch to it. Either op is callable from both PyTorch and JAX once the PJRT client is present.

Where the approach comes from

The direct-ioctl execution and the CNA descriptor matmul build on prior work. mtx512/rk3588-npu, a GPL v3 project by Jasbir Matharu, first demonstrated the descriptor approach in early 2024 on the older 5.10 kernel (write-up). Their approach differed: they derived the CNA and DPU register descriptors from the chip’s technical reference manual and ran matmul on the NPU’s internal SRAM, which avoids the DRAM roundtrip. The register and submit formats changed in the newer 6.1 driver, so OpenNPU re-derived the descriptor on this kernel rather than porting it; the older code does not run correctly on 6.1. The CNA descriptor matmul is one part of the reverse engineering, not the whole of it. The larger part was capturing the register-command templates and decoding the wire format, which the compiler and the op coverage rest on. OpenNPU credits that project and reimplements the descriptor independently under MIT. The MIT userland is separate from the GPL kernel driver, so the two do not mix.

You train a behavioral cloning policy. It works on the demos. You deploy it, and at some point a perturbation pushes it off track. The policy keeps producing actions that look reasonable. They’re just slightly off the manifold of successful executions, and there’s no restoring force.

The well-documented failure mode of pure BC (behavioral cloning) on manipulation is degenerate oscillation: the agent approaches the object, nudges it, backs off, re-approaches from the same angle, and loops. It never discovers it needs to go around.

The standard fixes have costs. Fine-tuning invalidates your safety case. Adding a planner requires a world model, and if your world model is wrong, the planner optimizes against a hallucination. Test-time adaptation reacts to the observation but ignores what the expert would’ve done.

My approach: don’t retrain the policy, don’t replace it, repair it in flight. Train a small recovery head, offline, that learns when the frozen policy is about to fail and how to push its action back toward expert behavior. The training signal is counterfactual regret labels. The transport mechanism is flow matching in DCT-compressed action space.

The repair manifold

Think of the policy as producing actions on a manifold in action space: the set of actions that lead to successful task completion. Expert demos live on this manifold. A well-trained BC policy produces actions near it. Under perturbation, the outputs slide off and there’s no restoring force.

A repair layer takes an off-manifold action & transports it back, conditioned on what the expert would’ve done from a similar state. Two questions, separable: when to intervene (a gate), and where to go (a decoder). Get either wrong and you either fail to repair or break working plans.

The gate is the whole game

The gate is load-bearing in the strongest sense. Flow repair applied unconditionally gets 0% success across every seed, every task, every PACE setting, and every model version I tested. The flow model transports BC actions toward expert anchors, but many anchors are bad. Unconditional flow is destructive. The gate is the only thing preventing it from destroying working plans.

A repair layer that fires too often is worse than one that fires too rarely. Every intervention is a chance to break a working plan. A gate that never fires is safe; you fall back to the frozen backbone. A gate that fires on every plan, even with a perfect decoder, adds noise & risk.

The first gate label was wrong

My first gate label was “any candidate has positive regret AND baseline fails.” That conflates “there exists a better action” with “the baseline would fail if left alone.” The gate fires whenever a better option exists, even when BC+MPPI (model predictive path integral control) would’ve succeeded. You get regressions on working plans.

Distance-mode labels via Warp

The conceptual fix is paired counterfactual rollouts: run both the baseline and the repair from the same frozen state, label TRUE only if the baseline fails AND the repair succeeds. In practice, long-horizon rollouts from mid-episode states diverge due to physics mismatch. Neither BC nor flow reaches success from most saved states. The binary success label is unusable.

I replaced it with a distance-mode gate label: “does the flow action get the object closer to the goal than the BC action after 8 steps?” This is a short-horizon proxy for “which action is better.” It doesn’t require the task to be solved. It’s computable in parallel.

I ported all 10 LIBERO_SPATIAL tasks to MuJoCo Warp for GPU-parallel rollouts. 1024 environments run at 9,838 steps/sec, a 49x speedup over the CPU SubprocVectorEnv. That’s 18K counterfactual labels in 7 minutes. The CPU pipeline took hours for the same volume.

The gate label is still precision-first. It says “flow is better than BC here,” not “flow is good.” The gate only fires when flow is demonstrably better than the alternative.

The gate threshold is as important as the gate label

The gate is a sigmoid classifier. Where you threshold it determines how often it fires. The default 0.5 is too aggressive.

Threshold 0.7 is 2.3x better than the default 0.5. Conservative thresholds are essential. Only apply flow when the gate confidence is high.

The result

Three-seed evaluation on LIBERO_SPATIAL task 0, 10 episodes per seed, 50-epoch BC baseline (50% success), PACE h=1:

The decoder: two approaches

Approach A: CFR (counterfactual regret) selection. MPPI proposes 256 chunks, a learned dynamics model rolls each forward, keep top-16. A transformer selects one. The world model is in the inference loop.

Approach B: flow-matching transport. Skip MPPI. Take the BC action, transport it toward an expert anchor via a learned velocity field in DCT (discrete cosine transform) space. No world model at inference.

The world model is the ceiling on A. I compared two dynamics models on multi-step rollout accuracy:

The 3-layer MLP is the best I have, and it’s still 46px off at 50 steps. The bigger residual is worse: it overfits short-horizon corrections & amplifies them over long rollouts. For world models, capacity is a liability. This pushed me toward Approach B.

Why DCT space

The flow head transports the BC action toward the anchor in DCT-edge coefficient space. The 8-step chunk keeps the first 2 timesteps raw (the edge: what you’re about to execute) and projects the rest onto 4 DCT-II basis coefficients. 8 raw values become 6 numbers.

Lower-dimensional: the velocity field has less to predict. Smooth by construction: you can’t produce a jerky chunk from 4 low-frequency coefficients. Exact edge: the immediate next action is kept raw, not smeared by DCT. When you have a strong prior, use it directly.

Both the gate & the flow head start conservative: negative bias init on the gate, zero-init on the velocity head. At initialization the repair does nothing. The model has to learn to deviate. A repair that never fires is safe; one that fires too often is dangerous.

PACE is structurally limited for single-action BC

The BC policy (BCTransformerPolicy) outputs 1 action per forward pass, not a chunk of 8. PACE h=8 tiles this single action for 8 consecutive steps. That’s equivalent to running the controller at 1/8 frequency. BC success drops from 50% (h=1) to 0% (h=8).

Stale-vs-fresh action L2 divergence: mean 0.557, max 2.83 (actions bounded [-1,1]). The stale action is completely wrong by the 90th percentile (2.03). PACE temporal coherence doesn’t apply when there’s no chunk to be coherent about. The BC produces a single action, and the “chunk” is just that action repeated. h=1 is the only viable setting. This is a structural mismatch between the BC architecture and PACE, not a tuning issue.

LIBERO results

PushT is 2D. LIBERO is 7-DoF manipulation from images. If the method is real, it survives the transfer. The portable pieces: anchor bank, DCT-edge encoding, flow velocity net, intervention gate. The one adaptation: flow only transports the continuous 6-DoF chunk. The gripper stays fixed from BC.

Three-seed evaluation on LIBERO_SPATIAL task 0, 50-epoch BC baseline (50% success), CaB with v2 hidden-state gate at threshold 0.7, PACE h=1:

CaB improves BC by +15% mean (v2, threshold 0.7). Flow alone gets 0%. The gate is the win. Flow matching is the transport mechanism, but the gate is what makes it safe.

Task 2 (10% BC baseline, the weakest task) doesn’t benefit. v1 CaB on task 2: 5% success (1/20), down from BC’s 10% (2/20). When BC is this weak, flow repair can’t help. The method needs a baseline that partially works.

The v2 model: hidden-state gate & anchor retrieval

The v1 model had a train/eval mismatch: training passed real anchor latents, eval passed anchor_latent=zeros(512). The v2 model fixes this with two changes.

GRU hidden state: a recurrent layer accumulates (latent, bc_action, gate history) across replan steps within an episode. The gate can detect “BC has been diverging for 5 steps” vs “BC is fine.” The hidden state gives the gate temporal context that a single-step classifier can’t.

Anchor retrieval: the v2 model retrieves the nearest anchor from the bank via cosine similarity on the BC latent, with coverage matching. This is the same at train and eval time. The v1 model’s zero-anchor eval was testing the gate with no anchor information.

Binary gate (BCE): simpler than the 3-class CaB label, with auto pos_weight for class imbalance.

The 20.9M dense model gives identical results to the 3.96M model. 18K labels give identical results to 11K labels. The bottleneck is gate label quality and gate threshold, not model capacity or data quantity. The hidden-state gate at 4.2M params outperforms all of them because it has temporal context, not because it’s bigger.

What didn’t work

Original gate label. “Positive regret AND baseline fails” produced over-intervention and regressions. The distance-mode label fixed it, but only after the Warp port enabled real physics rollouts.

More data with the wrong label. Scaled 50 to 200 episodes with the original noisy label. Offline F1 dropped 0.90 to 0.79. More data compounds label noise.

Bigger dynamics models are worse. The 6-layer residual gets 93px error at h=50 vs 46px for the 3-layer MLP. Capacity is a liability for world models.

Bigger flow models don’t help. 20.9M gives identical results to 3.96M. 18K labels give identical results to 11K. The bottleneck is gate quality, not decoder capacity.

Always-on flow repair. Flow with no gate gets 0% success on every seed, every task, every PACE setting. The gate is not optional infrastructure. It’s the only thing preventing flow from destroying working plans.

PACE h≥2. The BC policy outputs 1 action per forward pass. PACE h≥2 tiles this for h steps, reducing control frequency by h×. BC success drops from 50% to 0% at h=2. This is structural, not tunable.

The frozen policy knows when it’s failing

A raw linear probe on the BC’s 512-dim latent space hits F1 0.98 for failure detection. A sparse-autoencoder (SAE) drift detector gets F1 0.36 to 0.49. The raw latent beats the SAE by a factor of two.

The frozen policy’s own representations already know when it’s failing. You just have to read it out. The SAE direction was a detour. The BC latent also drives anchor retrieval (cosine similarity) & gate input in the v2 model. The 512-dim representation does triple duty: drift detection, anchor lookup, and gate context.

Status

The CaB +15% success rate is with a trained 50-epoch BC baseline and real Warp physics labels. The PushT +1 net strict is small. The gate threshold (0.7) is a hyperparameter that works for task 0 and might not generalize. Task 2 (10% BC) doesn’t benefit. 10 episodes per seed is a small sample. The next step is 50 episodes per seed across more tasks, and a per-task or adaptive threshold. If the signal survives those, there’s a paper. If not, the diagnosis (gate label is the lever, gate threshold is as important, flow alone is destructive, DCT flow is the right representation) is still the contribution.

Sharing work in progress. The next experiment might shift the picture. That’s the point.

CSCI 7000-011 — Current Topics in CS: Transformers for Robotics Instructor: Nikolaus Correll · CU Boulder · Spring 2026

Paper: “No Plan but Everything Under Control” — Mengers & Brock, ICRA 2025
My implementation: github.com/poad42/no_plan_everything_control (pure PyTorch, runnable without simulation)

The Problem Every Robotics Student Hits

You’re building a robot that needs to open a drawer. Sounds simple, right? But think about everything that has to happen:

1. The robot doesn’t know where the drawer handle is → it needs to look around first
2. Once it sees the handle, it needs to reach for it
3. It needs to grasp the handle
4. Then pull the drawer open

Traditional approaches say: write a planner. Define each sub-goal. Build a state machine that transitions between them. Handle failures at each stage. If something goes wrong mid-execution… re-plan.

This paper throws all of that away. Their claim: if you set up the math right, gradient descent alone will figure out the correct ordering of sub-goals. No planner. No state machine. No learned policy. Just calculus.

The Big Idea in One Sentence

Wire your sensors, actuators, and goal into a single differentiable graph. Take the gradient. Follow the steepest one. Repeat.

That’s it. The rest of the paper is explaining why this works and how to set up the graph correctly.

The Algorithm at a Glance

The entire algorithm is a loop:

1. Components update their state estimates from sensor data (Eq. 1)
2. Interconnections compute soft gates — p_visible, p_grasped, etc. (Eq. 2)
3. Goal cost is differentiated through all paths in the graph (Eq. 3)
4. The steepest gradient determines the next action (Eq. 4)
5. The action changes the world → new sensor data → repeat

There is no planner, no policy, no search. The feedback arrow closes the loop. Let’s look at each piece.

Building Block 1: Components (Eq. 1)

Think of a “component” as a little module that tracks one thing about the world. The robot’s end-effector position. The estimated location of the drawer handle. Whether the gripper is open or closed.

Each component updates its estimate at every timestep:

x_t = f(x_{t-1}; c_1, c_2, ..., c_N)

“My new estimate depends on my old estimate plus whatever information is flowing in from other components.”

Here’s what this looks like in practice — an EKF tracking the drawer handle position (components.py):

class EKFComponent(BaseComponent):
    """Extended Kalman Filter — tracks one quantity with uncertainty."""

    def update(self, priors):
        # Predict step
        x_pred, F = self._process_model(self._state, priors)
        Sigma_pred = F @ self._Sigma @ F.T + self._Q

        # Update step (if we have a measurement)
        if "measurement" in priors:
            z = priors["measurement"]
            z_pred, H = self._measurement_model(x_pred)
            S = H @ Sigma_pred @ H.T + self._R
            K = Sigma_pred @ H.T @ torch.linalg.solve(S, torch.eye(...))
            self._state = x_pred + K @ (z - z_pred)
            self._Sigma = (I - K @ H) @ Sigma_pred
        else:
            self._state = x_pred
            self._Sigma = Sigma_pred

Everything is PyTorch, so gradients flow through torch.autograd. Not just an estimator — a differentiable estimator.

Building Block 2: Active Interconnections (Eq. 2)

Components share information through differentiable gates that turn on or off depending on the current state:

c(x_1, x_2, ..., x_M)    — coupling changes based on state

Consider a visibility gate. Whether the wrist camera can see the drawer handle depends on where the end-effector is pointing. So p_visible is a soft sigmoid (interconnections.py):

class VisibilityGate(BaseInterconnection):
    def forward(self, ee_pose, object_pose):
        direction = object_pose[:3] - ee_pose[:3]
        cos_angle = F.cosine_similarity(direction, camera_forward, dim=0)
        return torch.sigmoid(self._sharpness * (cos_angle - self._threshold))

And p_grasped uses force-torque sensor readings:

class GraspGate(BaseInterconnection):
    def forward(self, distance, grip_force, hand_closed):
        p_close = torch.sigmoid(-self._dist_sharpness * (distance - self._dist_thresh))
        p_force = torch.sigmoid(self._force_sharpness * (grip_force - self._force_thresh))
        p_hand  = torch.sigmoid(self._hand_sharpness * (hand_closed - 0.5))
        return p_close * p_force * p_hand

This is the key insight. When p_visible = 0, the gradient path through vision is blocked. The algorithm naturally concludes: “I can’t optimize the grasp yet. The steepest gradient tells me to move the camera first.” Subgoal ordering falls out of the math — nobody programmed it.

Building Block 3: Steepest Gradient Action Selection (Eqs. 3–4)

The goal is a differentiable scalar cost. For opening a drawer 20 cm:

g(q) = (q − 0.20)²

The gradient propagates backward through all interconnections and components via the chain rule. Because different gates are open or closed, you get multiple gradient paths — each representing a different possible sub-goal.

The rule: pick the steepest one (gradient_descent.py):

# Compute all gradient paths
gradients = [path.gradient(action) for path in self._paths]

# Select steepest (Eq. 4)
norms = torch.stack([g.norm() for g in gradients])
best_idx = int(torch.argmax(norms).item())
grad_star = gradients[best_idx]

# Gradient descent (Eq. 3)
action_new = action - k * grad_star

That’s the entire control loop. No planner. No policy network.

Why It Works: The Locked Door Example

A point mass needs to reach (10, 0), but a wall at x=5 is locked. A button at (2, 4) opens it. The complete AICON implementation (demo_2d_locked_door.py):

class PointMassAICON:
    def step(self, lr=0.15):
        pos = self.agent_pos
        goal_dist = torch.norm(pos - self.goal_pos)
        dist_to_button = torch.norm(pos - self.button_pos)

        p_button_soft = torch.exp(-0.25 * (dist_to_button ** 2))
        door_current = door_prev + p_button * (1.0 - door_prev)

        factor = 1.0 + 10.0 * (1.0 - door_current)
        cost = goal_dist * factor + barrier * (1.0 - door_current)

        cost.backward()  # PyTorch does the rest

When the door is closed (door_current ≈ 0), factor = 11× — making the goal gradient weak relative to the button gradient. The agent:

1. Moves toward the button (steepest gradient when door is closed)
2. Touches the button (opens the door)
3. Heads straight to the goal (steepest gradient when door is open)

No planner told it to do this.

Scaling Up: Blocks World

AICON solves the classic Blocks World puzzle — rearranging block towers into a goal configuration — a benchmark that normally requires BFS or A*. The state is a matrix o[X,Y] (likelihood X is on Y) and vector c[X] (likelihood X is clear):

# Eq. 5: Is block X clear?
c[X] = 1 - (1/n) * sum(o[Y, X] for Y in blocks)

# Eq. 6: State update
o_new[X,Y] = o[X,Y] + c[X] * c[Y] * a_stack[X,Y] - c[X] * a_unstack[X,Y]

The gradient of the goal cost through these equations automatically decides whether to stack or unstack, and which blocks, and in what order (aicon_policy.py):

def _select_action(self, o, c, goal_cost):
    best_action, best_norm = ("stack", 0, 1), -1.0
    for X in range(n_blocks):
        for Y in range(n_blocks):
            if X == Y: continue
            if c[X] > 0.5 and c[Y] > 0.5:        # ∇_stack (Eq. 7)
                o_new = (o + c[X] * c[Y] * a_stack).clamp(0, 1)
                grad_norm = abs(goal_cost - goal_cost_fn(o_new))
                if grad_norm > best_norm:
                    best_action = ("stack", X, Y); best_norm = grad_norm
            if o[X,Y] > 0.5 and c[X] > 0.5:      # ∇_unstack (Eq. 8)
                o_new = (o - c[X] * a_unstack).clamp(0, 1)
                grad_norm = abs(goal_cost - goal_cost_fn(o_new))
                if grad_norm > best_norm:
                    best_action = ("unstack", X, Y); best_norm = grad_norm
    return best_action

Result: 100% solve rate on 130 randomly generated instances with 10–30 blocks, including tasks requiring 35+ steps. Matches optimal BFS in step count. No forward search.

Real-World: Drawer Manipulation

A Franka Panda with a wrist camera and force-torque sensor opens a physical drawer under uncertainty. Four components, two interconnections:

When Σ_drawer is large (high uncertainty), the steepest gradient path goes through reducing uncertainty — “move the camera to look around.” Once Σ_drawer shrinks, the gradient shifts to “reach and grasp.” After p_grasped flips to 1, it becomes “pull.” The robot sequences the entire task from one cost function: g(q) = (q − 0.20)².

Results vs. Planning Baselines

70 real-world trials: 7 conditions × 10 runs each. Baselines: state-space planner (fixed trajectory) and 67-dimensional belief-space planner (3 viewpoints + 2 kinematic explorations).

AICON handles both because it has no plan to become stale — the gradients recompute from live sensor feedback at every step. Under high uncertainty, it spontaneously invents triangulation — nobody programmed that behavior.

Ablations: Why the Gates Matter

Remove p_visible or p_grasped and performance craters. Sum all gradients instead of selecting the steepest and you get jerky, unstable motion. The steepest-gradient selector is crucial because competing subgoal gradients point in opposite directions and canceling them produces garbage.

My Take

What’s genuinely impressive:

• The math is minimal (sigmoid gates + chain rule + argmax) but the emergent behavior is sophisticated
• Zero-shot: no training, no reward shaping, no demonstrations
• Disturbance recovery is automatic — there’s no plan to invalidate, just gradients to recompute
• Matching BFS optimality in Blocks World with pure gradient descent is a striking result

Honest limitations:

• The graph structure is hand-designed per task. You need domain knowledge to pick the right components and gates
• Gradient path enumeration scales poorly. The drawer has ~12 paths; a kitchen scene would have thousands
• Everything must be differentiable. If your sensor model isn’t a PyTorch function, you’re stuck

Why it matters for this class:
Most current robotic AI either learns a policy end-to-end (which requires massive data) or uses a classical planner (which requires a precise world model). This paper suggests a third path: if you can express your world model differentiably, gradient descent handles task sequencing for free. That’s a structurally important insight, even if the current framework is limited to structured domains.

Code

Pure PyTorch, no simulator required for the demos:

git clone https://github.com/poad42/no_plan_everything_control
cd no_plan_everything_control

# Locked door (2D, instant, matplotlib)
python scripts/simple_demos/demo_2d_locked_door.py

# Blocks World solver
python scripts/run_blocks_world.py

Paper: arXiv:2503.01732
Supplementary videos: tu.berlin/robotics/papers/noplan

I disagree.

Call it: backporting FP8-style numerics experiments to the RTX 3090.

Not because Ampere magically does FP8 compute (it doesn’t), and not because this makes an RTX 3090 “faster” than Hopper (it won’t).

But because a lot of FP8 research and engineering is really about:

• how you store weights (bytes on the wire)
• when and where you expand them (decode)
• what scaling/quantization contract you enforce

You can explore a surprising amount of that on consumer Ampere, if you’re willing to treat FP8 as a storage format and map the math onto hardware that is available.

Quick note: if you see an acronym you don’t recognize, jump to the glossary.

Code: https://github.com/poad42/cuda-fp8-ampere

The plan

Ampere (sm_86) has extremely capable tensor cores, but it doesn’t have native FP8 tensor-core MMA. What it does have is a very fast path for INT8 tensor cores (IMMA / WMMA).

So the project becomes:

Keep weights stored as 1-byte FP8 bit patterns in VRAM, decode/scale/quantize on the fly, and use INT8 tensor cores for the matmul.

That’s the whole framing: democratize FP8 research by making the storage + numerics experimentable on hardware people actually have.

FP8-as-storage, in one paragraph

I am not trying to do “FP8 compute.” I’m trying to store weights in a compact FP8 format and only expand them when needed.

The VRAM part is simple: FP16/BF16 weights cost 2 bytes/weight, while FP8 weights cost 1 byte/weight. So for large weight matrices, storing FP8 can cut the resident weight footprint (and the bandwidth to stream it) by close to 2×.

In practice you also store scale factors (e.g. one FP16 scale per output channel), but that overhead is tiny compared to the full $N\times K$ weight matrix.

Conceptually:

1. Store weights as FP8 bytes (E4M3) — literally uint8 bit patterns.
2. Decode FP8 → FP16 on the fly using a 256-entry LUT.
3. Apply per-output-channel (per-column) scale.
4. Quantize to INT8 so the tensor cores can consume it.
5. Run IMMA (INT8×INT8→INT32 accumulate), then write FP16 output.

That’s the whole “FP8 without FP8 MMA” idea.

What’s actually new here (and what isn’t)

Three honesty bullets up front:

• This is not a claim that Ampere beats BF16/FP16 cuBLAS. In fact, for pure compute, cuBLAS is usually hard to beat.
• This is not full FP8 training. There’s no backward pass here.
• This project focuses on FP8(E4M3) storage. Extending to E5M2 is conceptually similar (another decode path), but I didn’t build it into this writeup.

So what is interesting?

Bit-level FP8 handling (LUT decode)

I store FP8 weights as raw uint8 bit patterns and decode them with a 256-entry LUT. Since there are only 256 possible FP8 bytes, decode is conceptually:

• u8 → fp16 via LUT[u8]

No __byte_perm tricks here — it’s mostly about making that decode cheap enough to hide behind the tensor-core pipe.

Scaling + quantization as a first-class contract

The weights aren’t “just FP8.” They’re FP8 bits + per-output-channel scale. The kernel makes that explicit: decode → apply scale → saturating quantize to int8 → IMMA.

Stochastic rounding (SR): important, but not implemented here

If you’re interested in FP8 training dynamics, stochastic rounding matters a lot. This project doesn’t implement SR (no backward pass), but if I were pushing this toward “training-like” experiments on older GPUs, SR would be near the top of the list.

Glossary (quick definitions)

• FP8(E4M3): an 8-bit float format. Great for storage, not great for high-accuracy math.
• MMA: matrix multiply-accumulate (the tensor core instruction family).
• IMMA / WMMA: NVIDIA’s tensor core path for int8 matrix multiply (instruction path / CUDA API).
• cuBLAS / cuBLASLt: NVIDIA’s GPU linear algebra libraries (GEMM).
• cp.async: an Ampere instruction to asynchronously copy from global memory to shared memory.
• l2pin: using “persisting L2” cache hints to keep hot tensors resident longer.
• Per-column scale: one scale factor per output channel; common in quantized inference.
• LUT decode: since there are only 256 FP8 bit patterns, decode can be a table lookup.

The pipeline, in one diagram

A (fp16/bf16)                B (uint8 fp8-e4m3 bits)         col_scales (u16 bits)
[M,K] row-major              [N,K] (represents KxN col-major)      [N]
      |                               |                               |
      |                               | (LUT in __constant__)        |
      |                               v                               |
      |                        fp8 -> fp16 decode                     |
      |                               |                               |
      |                               +-----------(per-column)--------+
      |                                           scale
      |                               |
      |                               v
      |                        fp16 -> int8 (sat)
      |                               |
      +--------------- int8 A --------+
                      (act quant)
                                      |
                                      v
                            WMMA/IMMA (int8) accumulate (int32)
                                      |
                                      v
                             D (fp16) written as [N,M]
                             (represents MxN col-major)

If you’ve never written CUDA kernels: that diagram is basically the whole story.

Baseline: PyTorch decode + matmul

Before writing any custom kernel, I wanted a baseline that matches the real workflow:

# weights stored as FP8 bytes
B_u8 = ...  # [N,K] uint8

# decode fp8 -> fp16 every iteration
B_fp16 = LUT[B_u8] * scales[:, None]

# compute in fp16 using standard matmul
out = A @ B_fp16.T

That’s the easiest FP8-as-storage implementation: store FP8 bytes, decode on demand, then use cuBLAS.

Two additional baselines are useful:

• Decode + matmul + downcast output to FP8: what the pipeline looks like if you want to store the output/activation in FP8.
• Matmul-only with fp16 weights cached: not apples-to-apples (you’re no longer storing FP8), but it’s a useful upper bound.

Baseline numbers (PyTorch)

Measured on RTX 3090 Ti (sm_86), CUDA-visible, shape $M=N=K=4096$.

Notes (important, and easy to misread):

• The “matmul only” baseline assumes fp16 weights are already resident. That defeats the FP8 VRAM savings.
• “Peak alloc” here is per-call peak allocated bytes; it does not include already-resident fp16 cached weights.

The naive decode+matmul being fast is not a paradox — cuBLAS is extremely optimized, and the decode step is embarrassingly parallel. My main motivation for the fused kernel is controlling memory traffic and keeping the pipeline “weight storage = FP8 bytes” end-to-end.

Fusing it into one kernel

Once you accept that IMMA wants int8 fragments, the kernel is a pipeline problem:

• Where does decode happen? (constant memory LUT vs texture vs global)
• Where does scaling happen? (apply scale in fp16, or bake it into an int8 conversion)
• Where does activation quant happen? (register path vs shared-memory staging)
• How do you feed tensor cores continuously? (avoid stalls from decode/scale/quant)

I ended up implementing variants as a way to test hypotheses.

Variants (experiments)

• v2: baseline fused path (FP8→INT8 JIT + IMMA). Keep it simple and measure.
• v2_i8lut: “what if I precompute a per-column FP8→INT8 table in shared memory?” (sounds clever; didn’t win).
• v3_act_f16: fused activation quantization, register path.
• v4_act_f16: cp.async staging for activations + shared-memory quantization, then IMMA.
• texscale: load per-column scales via TEX.
• l2pin: persisting-L2 hints for B/scales.

Kernel benchmark numbers

Measured via ./build/gpu_bench on RTX 3090 Ti (sm_86), driver 590.48.01, CUDA 13.1.

Shape: M=N=K=4096, --warmup 10 --iters 50.

Notes:

• *_l2pin can vary with driver/GPU state and other workloads.
• The int8gemm cuBLASLt number is not FP8-as-storage; it’s a ceiling for int8 TC GEMM on this machine.

Why do this at all?

After seeing the tables, the fair question is:

If naive Torch is already fast, why bother?

Because “fast” depends on what you’re measuring.

On pure matmul throughput, a highly tuned fp16/bf16 GEMM can absolutely win. This project is about a different constraint: weight storage and weight movement.

If your weights are truly stored in FP8 (1 byte/weight), then compared to fp16/bf16 (2 bytes/weight) you’re targeting up to 2× less weight traffic and 2× lower resident weight footprint. That’s a real lever for memory-bound inference workloads — even if you pay some extra compute to decode/scale/quantize.

Practically, the “democratizing FP8 research” win is:

• you can keep the storage format honest (FP8 bytes in VRAM)
• you can experiment with scaling/quantization contracts
• you can measure the cost of decode/quant instead of hiding it in a pre-processing step

So I view this as a tool for exploration: FP8-as-storage end-to-end, on hardware that doesn’t officially “support FP8.”

Try it yourself

Repo: https://github.com/poad42/cuda-fp8-ampere

Build:

git submodule update --init --recursive
cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build -j

Run tests:

cd build
ctest --output-on-failure

Run the kernel benches:

./build/gpu_bench --bench imma_fp8_jit_v2 --M 4096 --N 4096 --K 4096 --warmup 10 --iters 50
./build/gpu_bench --bench imma_fp8_jit_v4_act_f16 --M 4096 --N 4096 --K 4096 --warmup 10 --iters 50
./build/gpu_bench --bench imma_fp8_jit_v4_act_f16_texscale --M 4096 --N 4096 --K 4096 --warmup 10 --iters 50

Run the Torch baselines (including end-to-end downcast):

. .venv_torch_cuda312/bin/activate
python scripts/bench_torch_vs_fp8imma.py --M 4096 --N 4096 --K 4096 --kChunk 32 --report_mem --downcast_out_fp8

Next steps

If I had another weekend:

• Add a tiny numerical correctness harness (reference decode + GEMM with tolerances).
• Report a more honest memory metric: resident weights + peak workspace, not just per-call peak alloc.
• Try more realistic shapes (transformer-ish M, larger N, varying K) instead of only 4096³.

If you want to dig into the code, the repo contains:

• a CUDA kernel library (C++ API + C ABI)
• a benchmark harness (gpu_bench)
• a minimal PyTorch extension
• smoke tests (CTest + torch compile/import test)

If you look at the spec sheet for the Rockchip RK3588 (the chip inside the Orange Pi 5), it looks like a beast. It promises 6 TOPS of NPU performance. For $100, that’s a steal.

But if you try to run modern AI on it—specifically the Vision Encoder from SmolVLM—that promise falls apart.

The standard Computer Vision SDK (rknn-toolkit2) is optimized for older, predictable CNNs (like ResNet). When I fed it the SigLIP Vision Transformer used by SmolVLM, the driver choked. Even though the model is “smol,” the massive Attention matrices it generates triggered cryptic hex errors and refused to compile.

This left me with one option: running the model on the CPU. The result? A single image inference took ~30 seconds. The 6 TOPS accelerator sat idle while the CPU struggled.

I didn’t accept that. I decided to reverse-engineer the NPU to find out exactly why it was failing, and how to force it to run at full speed.

Context: Why do it the hard way? (First Principles)

A quick note for those following the ecosystem: You might see projects like QEngineering running the newer SmolVLM-v2 on Rockchip’s rknn-llm SDK.

That approach uses a specialized “black box” toolchain designed specifically for Transformers. Rockchip engineers have likely already implemented complex memory management inside that SDK to handle these models.

My project targets the original SmolVLM-v1, but more importantly, I built it on the legacy rknn-toolkit2 stack. Why hack the legacy stack? I wanted to take a “First Principles” approach. I didn’t want to use a black-box solver. I wanted to understand why the hardware was crashing on Attention layers and if I could find universal architectural patterns—like manual tiling and graph sharding—that could force any Transformer to run on any constrained edge accelerator.

The Detective Work: What is Error 0xe010?

Rockchip doesn’t publish a public Instruction Set Architecture (ISA). When I tried to compile the Attention layers, the driver kept spitting out an undocumented error: REGTASK Overflow (0xe010).

I hypothesized this was a memory overflow. Even though the model parameters are small (~96M), the intermediate activation matrices for a 1024-token sequence are huge (~25MB).

I wrote a script to generate synthetic ONNX graphs to probe the hardware limits:

• 8KB Tensor: Pass.
• 16KB Tensor: Pass.
• 32KB Tensor: Pass.
• 32.1KB Tensor: CRASH.

Discovery: The NPU has a hardware-enforced 32KB L1 SRAM Scratchpad for vector operations.

The standard compiler was trying to shove a 25MB Attention matrix into a 32KB slot.

The Fix: Nano-Tiling & The “Poison Pill”

To solve the 32KB limit, I wrote a “Nano-Tiling” algorithm in PyTorch. I manually sliced the massive 1024-token sequence into tiny 32x32 tiles that fit perfectly into the 32KB scratchpad.

But here is where it got messy. The rknn compiler is “smart.” It looked at my tiled graph, decided it was inefficient, and fused the operators back together into a single giant block… which immediately crashed the hardware again.

I had to trick the compiler. I needed a way to tell it: “Do not merge these nodes.”

I introduced a topological barrier I call the “Poison Pill.” I injected a dummy operation that looks mathematically significant to the dependency graph (preventing fusion) but is mathematically irrelevant to the model output.

# The "Poison Pill"
# 1. Take a slice (forcing a strided access)
slice_x = x[..., :1]

# 2. Apply a non-linear op (breaks compiler fusion heuristics)
# 3. Scale it down to near-zero so it doesn't affect the math
poison = torch.sigmoid(slice_x) * 1e-6 

# 4. Inject dependency
# The compiler sees 'out' depends on 'poison' and creates a barrier.
out = out + poison

By injecting this into the graph, I successfully forced the compiler to respect my tiling logic.

The “SigLIP Cliff”: Solving Accuracy Collapse

Getting it to run was step one. Getting it to be right was step two. When I first got the NPU running, the output was garbage. The cosine similarity compared to the original model was 0.02 (pure noise).

The culprit was the architecture of SigLIP. Unlike standard models, SigLIP has massive activation “spikes” (values around 300.0) sitting next to tiny visual signals (values around 0.05).

NPU quantization (INT8) works by mapping the range to -128/+127.

• If you zoom out to capture the 300.0, the 0.05 rounds down to 0. Signal lost.
• If you zoom in to capture the 0.05, the 300.0 overflows to infinity. Math crash.

I implemented a “Sandwich” Domain Shift:

1. CPU Pre-Scale: Multiply the input by 0.1. Now the max value is 30.0 (Safe for FP16).
2. NPU Execution: Run the heavy compute in this scaled-down “safe zone.”
3. CPU Post-Scale: Multiply the output by 10.0.

This simple trick restored the signal fidelity from 0.02 to 0.999 (effectively bit-exact).

The Architecture: Custom Runtime Scheduler

Finally, to bypass driver timeouts caused by the sheer number of tiles (thousands of tiny operations), I physically cut the model graph into 26 separate binary files (shards).

I wrote a custom User-Space Runtime in Python that acts as an orchestrator. It manually loads these shards onto the RK3588’s 3 separate NPU cores and fires them in a synchronized round-robin schedule (Core 0 -> Core 1 -> Core 2).

The Results

By ignoring the vendor’s “Unsupported” warnings and re-architecting the software to match the silicon’s physical reality, the results were drastic.

Conclusion

This project challenged the binary notion of “Supported Hardware.” The RK3588 didn’t support the SigLIP encoder out of the box on the standard SDK, but the silicon was always capable of it. It just needed an engineer to dig into the register overflow codes and manage the memory manually.

If you want to see the full code, including the tiling logic and the runtime orchestrator, check out the repo below.

View Source on GitHub