Cache-Aware Conv3D Lowering Across Embedded World-Model Decoders

Jiaming Zhang; Wu Yang; Shuai Tao; Wulong Liu

cs.LGSep 25, 2026

Cache-Aware Conv3D Lowering Across Embedded World-Model Decoders

Authors: Jiaming Zhang, Wu Yang, Shuai Tao, Wulong Liu

Organizations: University of Wisconsin–Madison, Madison, WI, USA · Beta Infinity, Beijing, China

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.

Figures & tables

Model	Rank	$C_{in}\!\to\!C_{out}$	$T_{in}\!\to\!T_{out}$	$H\!\times\!W$	$T_{c}$	Calls	Time / VAE share
Nano	1	256 $\to$ 256	4 $\to$ 4	360 $\times$ 640	2	15	26.08 s / 17.88%
Nano	2	512 $\to$ 512	4 $\to$ 4	180 $\times$ 320	2	15	25.59 s / 17.55%
Nano	3	1024 $\to$ 1024	2 $\to$ 2	90 $\times$ 160	2	15	12.76 s / 8.75%
LingBot	1	96 $\to$ 96	4 $\to$ 4	480 $\times$ 832	2	18	12.620 s / 27.0%
LingBot	2	192 $\to$ 192	4 $\to$ 4	240 $\times$ 416	2	18	12.343 s / 26.4%
LingBot	3	192 $\to$ 192	4 $\to$ 4	240 $\times$ 416	1	6	5.115 s / 10.9%

Table 1: Highest-cost causal Conv3D signatures. Cosmos3-Nano uses Nsight GPU time at

1280\times 720

; LingBot-World uses CUDA events at

480\times 832

. Edge uses the byte-identical Wan2.2 VAE but is not independently profiled.

Panel of Figure 1 — Figure 1: Cache-aware lowering of each temporal tap to one batched Conv2D.

Association	VAE	$z$	Compression	Relation	Boundary
Cosmos3-Edge	Wan2.2 TI2V	48	$16\times/4\times$	Primary deployment	Full I2V + decoder
Cosmos3-Nano	Wan2.2 TI2V	48	$16\times/4\times$	Byte-identical VAE	Full T2V + decoder
LingBot-World	Wan2.1	16	$8\times/4\times$	Distinct architecture	Decoder only

Table 2: Evaluated systems. Compression lists spatial/temporal factors.

Frames	Phase	Baseline (s)	Lowered (s)	Speedup
5	VAE encode	5.325	0.728	$7.32\times$
5	Transformer + other	21.516	21.581	$\approx 1.00\times$
5	VAE decode	18.798	2.374	$7.92\times$
5	Complete generation	45.639	24.681	$1.85\times$
25	VAE encode	25.561	4.024	$6.35\times$
25	Transformer + other	65.072	65.142	$\approx 1.00\times$

Table 3: Cosmos3-Edge complete image-to-video latency decomposition. Values are medians over three fresh-container runs per route.

T_{\mathrm{dit+other}}=T_{\mathrm{e2e}}-T_{\mathrm{encode}}-T_{\mathrm{decode}}

Model	Frames	C3D (s)	C2D (s)	Speedup	Coverage
Cosmos3-Edge	5	18.749	2.365	7.93 $\times$	67/67
Cosmos3-Edge	9	33.095	4.458	7.42 $\times$	101/101
Cosmos3-Edge	17	61.573	8.623	7.14 $\times$	169/169
Cosmos3-Edge	33	118.519	16.976	6.98 $\times$	305/305
Cosmos3-Edge	49	175.569	25.379	6.92 $\times$	441/441
Cosmos3-Nano	5	20.52	3.78	5.43 $\times$	67/67

Table 4: VAE-decoder BF16 latency. Cosmos3-Edge values are medians of

n=3

runs per route; Cosmos3-Nano and LingBot-World values have

n=1

. Coverage reports fast-path hits over WanCausalConv3d invocations within the measured decoder boundary, excluding pre-target activity. Cosmos3-Nano records 1,083/1,083 fast-path executions across the sweep.

Model	Full max	Mean abs.	RMSE	Cosine	NaN/Inf
Cosmos3-Edge	$4.49\mathrm{e}{-2}$	$1.84\mathrm{e}{-3}$	$2.75\mathrm{e}{-3}$	0.99998334	0/0
Cosmos3-Nano	$3.3203\mathrm{e}{-2}$	$1.8781\mathrm{e}{-3}$	$2.8215\mathrm{e}{-3}$	0.99998309	0/0
LingBot-World	$2.3440\mathrm{e}{-2}$	$1.5640\mathrm{e}{-3}$	$2.3230\mathrm{e}{-3}$	0.99998402	0/0

Table 5: Same-latent BF16 agreement at 17 output frames. Cosmos3-Edge and LingBot-World use synthetic decoder-compatible latents; Cosmos3-Nano uses a saved generated latent.

Method	Latency (s)	Speedup	Coverage
PyTorch eager	61.72	$1.00\times$	native
Ours	9.23	$6.69\times$	169/169
TensorRT	6.78	$9.10\times$	169/169

Table 6: Clean-device 17-frame decoder comparison on Jetson AGX Orin. Latency is the median over six runs; coverage denotes accelerated WanCausalConv3d calls.

Inference Acceleration

Inference acceleration reduces AI model latency, memory use, and serving cost while preserving output quality.

See all 18 papers

C $^3$ ache: Accelerating World Action Models with Cross Inference Chunk Cache
Jun 8, 2026Weisen Zhao, Lam Nguyen, Zhicong Lu +1World Model Learning Action-Conditioned World Models
George Mason University · University of Central Florida
LightVAE: Towards Compact and Efficient Video Autoencoders
Oct 10, 2026Shiqiao Gu, Yang Yong, Fengjun Zhong +5Variational Autoencoders Video Compression
SenseTime · Beihang University · Nanyang Technological University
WorldCache: Content-Aware Caching for Accelerated Video World Models
Mar 23, 2026Umair Nawaz, Ahmed Heakl, Ufaq Khan +3Diffusion Model Caching Diffusion Model Inference Acceleration
Mohamed bin Zayed University of Artificial Intelligence, UAE · Linköping University, Sweden

Jun 8, 2026J/K move · Enter open · S save

C $^3$ ache: Accelerating World Action Models with Cross Inference Chunk Cache

Weisen Zhao, Lam Nguyen, Zhicong Lu, Yuzhang Shang

George Mason University · University of Central Florida

World Model Learning Action-Conditioned World Models Efficient Inference for World Action Models Inference Acceleration World Action Models

Open paper AskPDF arXiv

World Action Models (WAMs) generalize better than standard Vision-Language-Action (VLA) policies to novel motions and environments, because a video-modeling objective lets them learn from abundant unlabeled video rather than scarce labeled robot demonstrations. This generalization is computationally expensive. To complete a task, a WAM runs over multiple inference chunks, and each chunk requires a costly denoising process. Existing acceleration methods reduce this cost by caching and reusing computation within a single chunk's denoising trajectory. Our empirical analysis reveals a substantial source of redundancy they overlook: redundancy across chunks. When a robot executes a smooth behavior, the residuals computed at a given denoising step are strongly correlated from one chunk to the next. We introduce C

^3

ache, a training-free method that caches and reuses these residuals across inference chunks at the same denoising step. Experiments on benchmarks with a Fast-WAM backbone show that C

^3

ache achieves up to a

2.5\times

speedup in total wall-clock inference time, with negligible degradation in task success rate.

Abstract

Inference Acceleration

Explore similar work

C3^33ache: Accelerating World Action Models with Cross Inference Chunk Cache

LightVAE: Towards Compact and Efficient Video Autoencoders

WorldCache: Content-Aware Caching for Accelerated Video World Models

C $^3$ ache: Accelerating World Action Models with Cross Inference Chunk Cache