PaperlayerView exact revision on arXiv

Source-linked AI summary

Cache-Aware Conv3D Lowering Across Embedded World-Model Decoders

Jiaming Zhang, Wu Yang, Shuai Tao, Wulong Liu

arXiv:2609.31938v1cs.LGcs.PF

TL;DR

Embodied planners can use world-model rollouts to consider possible futures, but rollout latency constrains deliberation. This paper lowers supported causal Conv3D calls in fixed pretrained decoders to batched spatial Conv2D while preserving model weights and temporal-cache semantics. The approach accelerates the primary 25-frame Cosmos3-Edge workload and transfers to a distinct Wan2.1 VAE, while TensorRT remains faster in the evaluated steady-state comparison.

  • Problem

    World-model rollouts can support embodied planning, but their latency limits deliberation speed and the number of possible futures considered.

    Source: §1 · p. 1

  • Method

    The runtime lowering expresses supported cached causal Conv3D calls as batched Conv2D operations while preserving pretrained weights and temporal-cache semantics.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    The primary 25-frame Cosmos3-Edge workload achieves a 6.91× decoder speedup and 2.21× complete-generation speedup.

    Source: §8 · p. 14

  • Takeaways & Limitations

    The generic lowering captures 95.55% of the absolute eager-to-TensorRT latency reduction without per-module or per-state engine construction.

    Source: Our Contributions · p. 2

    Evaluation uses one Jetson AGX Orin and fresh-process runs; LingBot-World tests decoder-only transfer with synthetic latents, and other hardware, resolutions, schedulers, and software stacks remain untested.

    Source: §7.3 · p. 14

Abstract

from arXiv · show

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.

1 Introduction

Rollout latency limits how quickly embodied planners can deliberate and how many possible futures they can consider. This work tests a runtime-only transformation of fixed pretrained decoders, then evaluates its effects across complete generation, decoder workloads, and runtimes.

1 Introduction

  • World-model rollouts support planning through imagined observations or latent trajectories, and their latency limits deliberation speed and the number of futures considered.The work studies local execution on a 64-GB Jetson AGX Orin, whose unified-memory architecture has less compute capacity and bandwidth than datacenter GPUs; cloud or fog offloading adds networking, availability, and latency considerations.Source: §1 · pp. 1–2
  • The study asks whether runtime-only transformation can execute a fixed pretrained causal video decoder more efficiently without modifying the learned model.For Cosmos3-Edge, the evaluation separates encoder, decoder, and unchanged transformer-and-other work in complete image-to-video generation.Source: §1 · p. 2, Our Contributions · p. 2

Our Contributions

  • For the 17-frame Cosmos3-Edge decoder, full-coverage TensorRT has the lowest steady-state latency under aggressive fixed-model AOT specialization, while generic runtime lowering captures 95.55% of the absolute eager-to-TensorRT latency reduction.Source: Our Contributions · p. 2
  • The unchanged lowering transfers to LingBot-World’s distinct Wan2.1 VAE, while Cosmos3-Edge and Cosmos3-Nano share a byte-identical Wan2.2 VAE.The evaluation also reports numerical agreement, route coverage, completion, and memory behavior.Source: Our Contributions · p. 2
  • TensorRT achieves the lowest steady-state latency in the evaluated 17-frame Edge decoder comparison, but requires per-module and per-runtime-state AOT specialization.The generic runtime lowering captures 95.55% of the absolute eager-to-TensorRT latency reduction without that engine construction.Source: Our Contributions · p. 2

2 Related Work

Prior work spans predictive world models, video generation, spatiotemporal computation, and embedded runtimes. This paper targets operator-level decoder execution while preserving the behavior of fixed pretrained models.

World Models and Video Generation

  • World models support reasoning through imagined observations or latent trajectories, while visual foresight combines action-conditioned video prediction with model-predictive control.The study does not compare predictive accuracy, interaction quality, or policy performance.Source: §2 · p. 3

Spatiotemporal Computation and Decoder Acceleration

  • Unlike methods that redesign, prune, replace, or distill decoder components, this work changes only the runtime representation of supported calls in fixed pretrained decoders.It requires no additional training data, parameter updates, or changes to the decoder’s learned function.Source: §2 · p. 3
  • P3D and R(2+1)D factorize 3D kernels for network design and optimization, whereas this approach evaluates each existing causal Conv3D temporal tap as spatial Conv2D.Original weights and temporal indexing are maintained.Source: §2 · p. 3

Embedded Runtimes, Compilation, and Serving

  • The work addresses a complementary operator-level problem: preserving temporal-cache semantics while choosing a guarded runtime route inside the decoder stage.It evaluates support conditions, fallback behavior, route coverage, numerical agreement, memory use, complete-generation latency, and a full-coverage TensorRT reference.Source: §2 · p. 3

3 Deployment Characterization

The deployment characterization profiles Conv3D workloads in Cosmos3-Nano and LingBot-World, then relates them to the evaluated Edge decoder. The Nano and Edge VAEs are identical, while LingBot uses a distinct Wan2.1 VAE.

3.1 Cosmos3-Nano Profile and Edge Checkpoint Identity

  • Conv3D accounts for 139.91 of 145.82 seconds of Cosmos3-Nano VAE GPU time (95.95%), with expensive calls combining short temporal extents and large spatial feature maps.The profiled 17-frame decoder contains 169 causal Conv3D invocations; one implicit BF16 3D-convolution kernel family accounts for 75.6% of captured kernel time.Source: §3.1 · p. 4
  • Cosmos3-Edge’s VAE is byte-identical to Nano’s Wan2.2 VAE, so the Nano profile targets the same operator family, but its 95.95% attribution is specific to Nano execution.Edge is not independently profiled in Table 1.Source: §3.1 · p. 4, Table 1 · p. 4

3.2 Distinct Wan2.1 Transfer Profile

  • At 17 frames, LingBot executes 164 Conv3D calls across 28 signatures; their summed CUDA-event time is 45.801 seconds, or 97.8% of clean VAE-decode time.The instrumented decode takes 48.786 seconds, including approximately 1.97 seconds of measurement overhead.Source: §3.2 · p. 4
  • LingBot’s dominant calls use five-dimensional tensors, unit temporal stride, one convolution group, and temporal kernel size three, matching the lowering’s targeted execution contract.The two highest-cost signatures contribute 53.4% of attributed Conv3D time.Source: §3.2 · p. 4

4 Cache-Aware Runtime Lowering

The lowering reconstructs a causal Conv3D by applying spatial Conv2D to each valid temporal tap and accumulating the responses. Cache-aware interval construction preserves temporal alignment, while guarded routing falls back for unsupported calls.

4.1 Temporal-Tap Conv2D Formulation

  • Each unchanged Conv3D temporal tap uses its spatial weight slice in a Conv2D operation, and the responses are summed with bias added once.The original spatial stride, padding, and dilation are retained.Source: §4.1 · p. 5

4.2 Temporal Intervals

  • For each tap, the valid output and source intervals have equal length, allowing the source range to be folded into the batch for one Conv2D call.This avoids evaluating invalid temporal frames, preserves causal alignment, and supports temporal dilation through the tap offset.Source: §4.2 · p. 5

4.3 Guarded Runtime Procedure

  • The guarded procedure validates supported inputs, accumulates tap responses, adds bias once, restores output layout, and records fast-path calls; unsupported calls fall back to Conv3D.Strict mode terminates validation on an unsupported case or execution error; otherwise exceptions fall back.Source: §4.3 · p. 6

5 Experimental Methodology

The experiments compare the lowering across three decoder systems and evaluate it within complete Cosmos3-Edge generation, under controlled decoder workloads, and against TensorRT. They also validate numerical agreement and account for memory on a 64-GB Jetson AGX Orin.

5.1 Systems, Boundaries, and Provenance

  • Cosmos3-Edge is the primary deployment case, Cosmos3-Nano supplies profiling and text-to-video evaluation, and LingBot-World tests transfer to a distinct decoder architecture.The evaluated systems are summarized in Table 2; compression factors are listed spatially and temporally.Source: §5.1 · p. 7, Table 2 · p. 7

5.2 Complete-Pipeline Workloads

  • The Edge image-to-video workload fixes the image, prompt, resolution, denoising settings, seed, and BF16 precision, with 25 frames as primary and five as minimum.The Nano sweep varies frame count from 5 to 49 at 1280×720; its timer includes tokenization, text encoding, denoising, and decoding, but excludes output conversion and file writing.Source: §5.2 · pp. 7–8

5.3 Decoder Sweeps and Acceptance

  • Decoder sweeps test 5–49 frames with shared deterministic BF16 latents for Edge, while LingBot uses synthetic decoder-compatible latents.A point is accepted only with successful execution, a fresh correctly shaped nonempty output, no unexpected NaN/Inf, and consistent route counters; optimized runs must avoid unexplained fallbacks.Source: §5.3 · p. 8

5.4 Clean Production-Runtime Comparison

  • The clean 17-frame Edge comparison uses six balanced-order measurements per route on a fixed-frequency, otherwise idle Jetson, with fresh processes and cleared temporal state.TensorRT uses fixed-shape engines across runtime signatures; engine construction, deserialization, and context creation are excluded from steady-state timing and reported separately.Source: §5.4 · p. 8

5.5 Numerical Validation and Memory Accounting

  • Numerical validation compares routes on identical latents in BF16 and FP32, while paired generation captures finite-precision propagation through 50 unchanged denoising steps.Anonymous RSS and CUDA allocator metrics are overlapping memory views on unified-memory Jetson and are not summed.Source: §5.5 · p. 9

6 Results

The lowering substantially reduces decoder latency in complete Edge generation and across the tested decoder systems, with fast-path coverage throughout measured decoder boundaries. It agrees closely with baseline outputs, uses more memory, and is slower than fully specialized TensorRT, which requires extensive specialization.

6.1 Cosmos3-Edge Deployment

  • At 25 frames, decoder latency falls from 90.05 to 13.03 s (6.91×), and complete-generation latency falls from 180.68 to 81.94 s (2.21×).The five-frame workload shows 7.92× decoder acceleration and 1.85× complete-generation acceleration.Source: Table 3 · p. 9, §6.1 · p. 9

6.2 Decoder Sweeps and Architecture Transfer

  • Decoder acceleration spans 6.92×–7.93× for Edge, 4.74×–6.65× for Nano, and 2.70×–5.75× for LingBot, with no decoder fallbacks across the three systems.LingBot uses a distinct Wan2.1 VAE; Edge and Nano use byte-identical Wan2.2 VAEs through different generation modes and pipeline boundaries.Source: Table 4 · p. 10, §6.2 · p. 10

6.3 Numerical Agreement

  • Across the Edge frame sweep, same-latent BF16 cosine similarity is at least 0.99998, maximum absolute difference is at most 0.045, and no NaN/Inf mismatches occur.In FP32, maximum absolute difference is 1.2 × 10^-3 for Edge and cosine similarity is 0.99999998 for Nano and LingBot; TensorRT is also not bit-exact relative to eager BF16.Source: Table 5 · p. 11, §6.3 · p. 11

6.4 Memory Behavior and Completion

  • The lowering trades memory for latency: peak CUDA allocation is 7.31 GiB versus eager's 5.37 GiB in the clean 17-frame Edge comparison, while TensorRT reaches 14.57 GiB.All routes fit on the 64-GB Orin; Edge and LingBot also complete the 49-frame decoder configuration, and unified-memory process and CUDA metrics are not summed.Source: §6.4 · pp. 11–12

6.5 Production-Runtime Comparison

  • TensorRT is 1.36× faster than the lowering, which captures 95.55% of the eager-to-TensorRT latency reduction, while TensorRT requires 99 engine specializations across 33 runtime signatures.TensorRT also requires 1637.6 s of AOT compilation, 3.16 GB of serialized engines, and 7.58 s of cold setup; the lowering achieves most steady-state gain without per-module or per-state AOT engine construction.Source: §6.5 · p. 12

7 Discussion

The discussion shows that cache-aware lowering works across complete Cosmos3-Edge generation and transfers to the evaluated Wan2.1 decoder, while TensorRT remains faster at a cost in memory residency and specialization. The evidence is bounded by the tested hardware, workloads, and supported convolution cases, and numerical agreement does not establish perceptual or downstream-control equivalence.

7.1 What Edge Establishes and What Transfers

  • The lowering remains effective in the complete 50-step Edge pipeline; the unchanged transformer-and-other work means generation gains follow accelerated VAE phases, with larger gains at 25 frames where baseline VAE latency is greater.Edge and Nano share a byte-identical Wan2.2 VAE, whereas LingBot-World supplies transfer evidence through a distinct Wan2.1 VAE with different architecture and runtime shapes.Source: §7.1 · pp. 12–13

7.2 Performance, Memory, and Numerical Implications

  • 6.78 s versus 9.23 s: TensorRT is faster on the clean comparison, while the lowering captures 95.55% of the eager-to-TensorRT latency reduction without 99 module- and state-specific engines.The distinction is the specialization and memory residency required to reach the higher-throughput operating point.Source: §7.2 · p. 13
  • In the clean 17-frame comparison, the lowering uses 7.31 GiB peak CUDA allocation and 9.75 GiB process VmHWM, versus 14.57 GiB and 34.19 GiB for full-residency TensorRT; cached spatial weight slices account for about 1.06 GB.Overlapping temporary-buffer lifetimes prevent an additive peak-memory decomposition, though selective caching, buffer reuse, or streamed accumulation are possible opportunities.Source: §7.2 · p. 13
  • Selected Triton kernels improved isolated operations, yet an integrated 17-frame experiment increased VAE-decoding latency from 16.01 to 66.81 s, without an integrated kernel-time breakdown to identify the cause.The result cautions that operator-level gains need validation at the complete execution boundary.Source: §7.2 · p. 13
  • BF16 and FP32 evaluations support numerical agreement for tested inputs, but do not establish perceptual equivalence or invariance in downstream control.The smaller FP32 discrepancy is consistent with finite-precision execution and accumulation order rather than an observed cache-indexing error.Source: §7.2 · p. 13

7.3 Limitations

  • Edge and Nano complete-generation evaluations differ from LingBot-World's synthetic decoder-compatible latents; Nano and LingBot sweep points use one fresh-process measurement, and all experiments use one Jetson AGX Orin.Fresh-process execution does not represent sustained or concurrent serving; Edge decoder and generation sweeps use three repetitions, while the clean 17-frame comparison uses six balanced repetitions.Source: §7.3 · p. 14
  • TensorRT is 1.36× faster on the clean 17-frame workload, and the lowering is limited to supported five-dimensional inputs, unit temporal stride, and one convolution group.Twelve stride-two encoder calls remain native Conv3D, so zero-fallback claims apply to the decoder, not the complete pipeline; other hardware and workloads remain untested.Source: §7.3 · p. 14
  • The complete Edge numerical evaluation uses one image, prompt, seed, and scheduler; thermal throttling is not ruled out, and no reliable whole-board input-power rail was available for energy measurement.The evaluation also does not establish closed-loop task performance or downstream-control invariance.Source: §7.3 · p. 14

8 Conclusion

The paper presents cache-aware lowering for supported causal Conv3D calls, preserving pretrained weights and temporal-cache semantics while evaluating them as batched Conv2D operations. On the primary 25-frame Cosmos3-Edge workload, it reports 6.91× decoder and 2.21× complete-generation speedups, with a lighter runtime point than fully specialized TensorRT.

8 Conclusion

  • 6.91× decoder speedup and 2.21× complete-generation speedup on the primary 25-frame Cosmos3-Edge workload demonstrate the lowering's reported gains.The evaluated Wan-family decoders also showed complete fast-path coverage.Source: §8 · p. 14
  • The same lowering transfers to a distinct Wan2.1 VAE, indicating the optimization is not specific to a single decoder instance.The conclusion reports substantial decoder acceleration and complete fast-path coverage across evaluated Wan-family decoders.Source: §8 · p. 14
  • Fully specialized TensorRT achieves higher steady-state throughput, while the lowering offers a lighter runtime operating point without per-module or per-state engine construction.The comparison uses full-coverage TensorRT specialization.Source: §8 · p. 14
  • The runtime lowering evaluates supported causal Conv3D calls as batched Conv2D operations without changing pretrained weights or temporal-cache semantics.The conclusion also reports numerical and route validation supporting intended cache behavior and strong finite-precision agreement.Source: §8 · p. 14

Related papers

Loading 2609.31938v1…