FlashRT-PAI speeds up physical AI model inference with pipeline design

FlashRT-PAI: Fast Kernels Are Not Fast Pipelines for Physical AI Inference

Artificial Intelligence

Summary

Physical AI systems that use vision, language, and action models often struggle with slow responses when making decisions fast. The authors show that the main delay comes not from the actual computations but from how many small tasks are started and how data moves between them. They built FlashRT-PAI, a system that reduces these delays by combining model design with hardware pipelines and reusing intermediate data intelligently. This approach makes AI models run much faster on real hardware without changing the original AI models.

What this means in practice

  • For embedded ai developers: Deliver faster real-time inference for physical AI models on edge devices by integrating pipeline-optimized execution and state reuse.
  • For machine learning engineers: Accelerate PyTorch-based vision-language-action models by using FlashRT-PAI’s drop-in kernel replacement for improved latency on GPUs and embedded platforms.

Authors

Liang Su, David Chen, Jason Li, Kevin Guo, Teresa Lv

Abstract

Vision-language-action (VLA), video-action (VA), and world action models (WAMs) form the backbone of modern physical AI systems, yet their real-time single-batch inference remains a defining bottleneck. Latency is dominated not by individual kernel computation, but by system-level overhead: thousands of short kernel launches, intermediate data movement across operator boundaries, and redundant recomputation across consecutive decisions. General-purpose graph compilers cannot address these costs, as they optimize individual calls in isolation and lack visibility into cross-call state semantics and global pipeline structure. We present FLASHRT-PAI, a low-latency inference engine built around model-hardware pipeline co-design rather than isolated kernel optimization. FLASHRT-PAI combines three core mechanisms: deep kernel fusion with launch overlap eliminates inter-operator boundaries and launch gaps; an execution contract formalizes the pipeline's stage DAG, update rates, and named buffers to enable cross-decision state reuse and selective refresh; and execution-state capsules enable fast snapshot, restore, and fork of persistent pipeline state. The system also supports drop-in acceleration of unmodified PyTorch models via parity- and latency-gated kernel replacement. Evaluation covers six policies on RTX 5090 and four on Jetson AGX Thor. At matched 16-bit precision, FLASHRT-PAI is 1.3-3.7x faster than official compiled paths. Combining pipeline co-design with reduced precision delivers 3.38-12.66x speedup over official eager implementations. In LIBERO-10, four-bit pi0.5 achieves 94.04% success on Thor and 93.32% on RTX 5090, compared with 92.4% for the measured official FP16 reference. FLASHRT-PAI is open source, distributed via Hugging Face Kernel Hub, and its opt-in pi0.5 kernels have been adopted into the LeRobot framework.