Cache-aware method speeds up video decoding on edge devices
Cache-Aware Conv3D Lowering Across Embedded World-Model Decoders
Machine LearningPerformance
Summary
Running detailed 3D video decoding models on small devices like robots or drones is usually slow due to limited resources. The authors found a new way to change how these models work internally so they run faster without changing the original learned model. Their method breaks down 3D computations into more efficient 2D steps, which fits better with the device memory and speeds things up by a lot. This helps systems create video predictions faster and uses the same approach across different devices and models.
What this means in practice
- •For embedded system developers: Accelerate video-based world model decoding on embedded platforms like NVIDIA Jetson AGX Orin to reduce latency without retraining models.
- •For machine learning engineers: Implement cache-aware transformation of Conv3D layers to improve generative model runtime efficiency on resource-constrained hardware.
Authors
Jiaming Zhang, Wu Yang, Shuai Tao, Wulong Liu
Abstract
Generative world models can provide visual rollouts for embodied planning, yet their feasibility on edge devices depends not only on the learned model but also on how the execution runtime represents its operations. We introduce a cache-aware lowering that expresses supported causal Conv3D calls as batched spatial Conv2D operations while preserving pretrained weights, temporal-cache semantics, convolution parameters, bias placement, and output layout. Across the complete Cosmos3-Edge image-to-video pipeline on a 64-GB NVIDIA Jetson AGX Orin, the proposed route accelerates VAE decoding by approximately $7\times$ and reduces complete-generation latency by more than $2\times$, while repeated decoder evaluations maintain complete fast-path coverage without fallbacks. The unchanged lowering also improves Cosmos3-Nano and transfers to LingBot-World's architecturally distinct Wan2.1 VAE. A clean-device comparison against fully specialized TensorRT shows that TensorRT provides a further $1.36\times$ steady-state improvement, but requires substantially greater per-module and per-runtime-state AOT specialization. Same-latent BF16 and FP32 evaluations characterize the finite-precision differences introduced by the alternative execution order. Together, these results position cache-aware lowering as a lightweight runtime optimization that recovers most of the available decoder acceleration without modifying the learned models themselves.