Beyond Prefill and Decode: Disaggregation Moves from Tokens to Frames

The distillation process serves two purposes. It reduces the number of denoising steps required for each generated block, and it trains the autoregressive student to remain stable as generation progresses. This helps mitigate error accumulation and exposure bias, which can otherwise lead to visual degradation during long autoregressive video generation.

In our implementation, the distilled model performs a one-step DiT rollout for each latent block. The VAE then immediately decodes that block into RGB frames before generation continues with the next block. Causal attention enables this incremental execution because the model no longer depends on future latent blocks.

This changes the systems architecture significantly. Instead of VAE decoding occurring primarily once at the end of video generation, DiT latent generation and VAE pixel decoding become recurring stages of the streaming pipeline. That repeated DiT-to-VAE handoff creates a natural boundary for disaggregation.

Different stages need different scaling strategies

The DiT can use replication and tensor parallelism. In this VAE implementation, however, the temporal cache is module-level state, so a worker decodes one session at a time rather than simply batching independent sessions together. The implementation accelerates VAE decoding through spatial sharding: dividing a latent by height across GPUs and exchanging the boundary rows needed for convolution.

These are different ways of using additional hardware. Separating the stages allows each to use the approach appropriate to its execution characteristics.

Scale the measured bottleneck-not every stage equally

Profiling shows that VAE decoding is the slower stage in this streaming pipeline. Spatially sharding VAE decoding across additional GPUs reduces that bottleneck and brings its execution rate closer to the DiT stage.

This motivates allocating additional GPU resources to VAE decoding while keeping DiT execution on a smaller GPU pool. The resource split follows the characteristics of each stage rather than requiring every stage to scale identically.

Separate memory budgets and overlap execution

DiT alone occupies approximately 25.9 GB of a 32 GB B70 GPU in the reported configuration. Separating the VAE and its activations creates more headroom for cache growth, encoding buffers, and other runtime allocations.

The split also enables a pipeline: DiT can generate latent block N+1 while VAE decodes block N. In steady state, the completion interval can approach the slower stage's execution time rather than the sum of both stages, subject to transfer and synchronization overhead. This improves the cadence of successive outputs; it does not remove the dependencies that determine first-output latency. ^[1]

Together, these characteristics provide the case for diffusion disaggregation: independent scaling, separate memory budgets and overlapping execution.