Posted by hhuytho 4 hours ago
This repo has optimized runtimes for three VLAs on Jetson Thor: ABC-VLA, MolmoAct2, and pi-0.5. We used our agentic search to propose the custom fused kernels, and the runtime glue between engines. Changes are only accepted if they pass numerical parity and a behavioral test.
On MolmoAct2, a 30-action chunk went from 611 ms to 113 ms. Most of that comes from CUDA graphs and TensorRT. The last 1.5x was from the agentic search. No distillation or pruning was used.
On the recently published ABC-VLA model, the optimized solution is 32.4 ms compared to the TensorRT conversion at 63.4 ms. A significant part of this gain came from a lossless weight decoder, similar to DFloat11 and ZipServ. This was done automatically by the agents without any human guidance. Again, this is without any distillation or pruning. The plans are bf16/fp16, and no 8-bit or 4-bit quantization.
More information is in the README and blog posts. Please let us know if you have any problems reproducing on your hardware.
2 points | 0 comments