Add tiled VAE parallelism & extended model support - #14
Merged
Conversation
get_torch_distributed_backend raised NotImplementedError when neither CUDA nor MUSA was present, so every entry point that asks DistVAE which backend to use was unreachable on a CPU-only machine. Sharding is correctness-testable there, and gloo is the backend that gets it there, so offer that instead of refusing. The gloo marker was also unregistered, warning eight times per run and leaving no way to select or skip the multi-rank tests as a group. Register it, and put the spawn scaffolding those tests each carry a copy of in one place, so the per-family adapter tests to come do not add another. Co-authored-by: Cursor <cursoragent@cursor.com>
PatchGroupNorm summed each rank's variance about that rank's own mean, so the result was short of the variance the unsharded norm computes by exactly how far the patches sat from each other. A rank holding a brighter patch measured its deviations from a brighter middle, and nothing put that offset back. It also took torch.var at its default, which applies Bessel's correction, where nn.GroupNorm divides by the count. That one is small enough to hide under a loose tolerance in a single layer and accumulates through a decoder's worth. Both are why a sharded AutoencoderKL decode did not reproduce a single-rank one. The new test compares against nn.GroupNorm directly, over 4D and 5D inputs and both split axes, at a tolerance that would have caught either. Co-authored-by: Cursor <cursoragent@cursor.com>
DecoderAdapter is what every AutoencoderKL model shards through, xDiT's SD3 and Z-Image among them, and its only check lived in a torchrun script that needs NCCL and a GPU. WanDecoderAdapter had none at all. Both now decode against an unsharded reference over gloo at world size 1, 2 and 4. The harness these rest on gets its own test, since a comparison that quietly accepted anything would turn the whole suite green and mean nothing. Co-authored-by: Cursor <cursoragent@cursor.com>
Qwen-Image's decoder is a fork of Wan's: QwenImageCausalConv3d and QwenImageResidualBlock are the Wan classes under other names, down to the same causal padding tuple and the same RMS norms that need no sharding. The single difference that reaches the adapter is that its up blocks were not given Wan's first_chunk argument. So rather than a second copy of five adapters, each Wan adapter now names the block types it accepts and the child adapters it builds, and the Qwen ones declare their own. The names xDiT dispatches on are unchanged, and the existing Wan decode test pins the behaviour across the move. Unlocks --use_parallel_vae for Qwen-Image, Qwen-Image-Edit, Krea-2-Raw and Krea-2-Turbo. The blocks are imported optionally, so a diffusers too old to carry them still runs every other family's adapter. Co-authored-by: Cursor <cursoragent@cursor.com>
The direct path has two ways of padding a patch. For zeros it pads both sides of the halo-extended input and crops the halo off the output, letting the spurious padding at the interior edge land only on rows it then discards. For any other mode it instead drops the padding at the interior edge, since the halo already holds those rows, and the output is patch-sized as it stands. It then cropped the halo off that too, taking a row off each interior boundary and returning a patch shorter than the one it was given. Deferring to build_crop_slice settles it: that already recognises an output the size of the patch and leaves it alone, which is why the chunked path was right all along. Circular padding is now refused across ranks rather than answered wrongly. It reads from the opposite edge of the image, which is not on a neighbour, so no halo exchange can supply it. Nothing in scope uses it; a single rank still can. Co-authored-by: Cursor <cursoragent@cursor.com>
…ecoders Both pad their causal convolutions by replication, which left alone would have each rank repeat its own edge rows at a boundary that is not an edge of the image. Moving the spatial half of that padding into PatchConv3d fixes it, and the temporal half stays where it was, one-sided along an axis nobody splits. Where they differ is the two places sharding notices. HunyuanVideo normalises with GroupNorm, whose statistics span the split and so have to be summed across ranks, while 1.5 uses RMS and needs nothing. And HunyuanVideo's mid block flattens (F, H, W) into a sequence and builds the causal mask itself, both from the height in front of it, so its attention cannot be wrapped on its own the way every other family's can: hand it a patch and it gets a mask cut for a patch. Gathering around the whole mid block avoids reimplementing that forward and costs little, since it runs at the latent resolution. 1.5's attention takes the tensor whole and needs no such thing. Unlocks --use_parallel_vae for HunyuanVideo, HunyuanVideo 1.5 and its distilled variants. Co-authored-by: Cursor <cursoragent@cursor.com>
LTX-2 is the last of the four families xDiT could not decode in parallel, and the least work of them. Its causal convolution keeps spatial padding inside its nn.Conv3d rather than applying it around one, so swapping that convolution for a PatchConv3d is the whole of the change; the temporal padding repeats frames along an axis nobody splits. Its upsampler moves channels into space, one input position per output, and its norms reduce over channels, so neither needs to see rows a rank does not hold. Alone among these families its mid block has no attention, so nothing has to be gathered back together mid-decode. A residual block with inject_noise enabled is refused rather than sharded: the noise is drawn per rank and would not add up to what a single rank draws, so the decode could not match its reference. No shipped LTX-2 or LTX-2.3 config turns it on. The decoder takes a timestep embedding where the other families take a temporal cache, so the splitting and reassembling around the decode moves into _sharded_decode, which takes the call to make rather than assuming a signature. Co-authored-by: Cursor <cursoragent@cursor.com>
The QwenImage, HunyuanVideo and LTX-2 adapters name classes newer than this, which reads like the floor is stale. It is not: those classes are resolved through diffusers_blocks rather than imported, so only Wan's blocks have to exist for this package to import at all. Co-authored-by: Cursor <cursoragent@cursor.com>
Patchify padded the axis it splits up to a multiple of the rank count, and the adapters cropped the result afterwards. That is not the same computation as the one being reproduced: the pad is zeros only until the first convolution, after which it is the network's answer to zeros, and it reaches the rows that survive the crop through every receptive field and every attention that follows. Measured on a Qwen-Image decoder against a single-rank reference, 15 latent rows over 2 ranks moved 84% of the output pixels by up to 0.77, and 16 rows over 3 moved 93% by up to 1.0. Wan's decoder, which has shipped for longer, is the same. The 2D DecoderAdapter escaped it by splitting after its mid block rather than before, so nothing that goes through AutoencoderKL is affected. Bands are now cut in whole multiples of what the VAE narrows or widens the axis by, which keeps each one on the grid the strided convolutions step along, and uneven counts are absorbed by giving the first ranks one band more rather than by inventing rows. Bands therefore differ in size between ranks, which the two gathers could not do: dist.all_gather requires every rank to contribute the same shape, so they now pad for the length of the transfer and slice on the far side, where the padding cannot reach a convolution. That is also what WanZeroPadConv2d was refusing uneven bands over, and it no longer needs to. Each adapter's tests gain the case that was wrong. use_uniform_patch is gone from the adapters, having named the behaviour that has been removed. Co-authored-by: Cursor <cursoragent@cursor.com>
Qwen-Image's encoder is Wan 2.1's laid out flat and renamed, down to the resample that pads (0, 1, 0, 1) and strides 2 over it. So the adapter is Wan's with a different set of block classes, and the two now share a base that holds the skeleton they agree on: a causal convolution in, a run of down blocks, a mid block, a normalisation, a causal convolution out, and the split and gather around them. What a family supplies is which adapter fits which down block, whether its forward threads a temporal cache, and whether it ends on a norm that reduces over the axis being split. Wan's encoder adapter kept the two down block shapes it has to handle, 2.2's grouped stages and 2.1's flat list, and gains nothing else. Unlike before it refuses a down block it does not recognise rather than warning and leaving it unsharded, which produced latents no one had checked. WanEncoderAdapter had no test; both are now covered at one, two and four ranks, including an attention among the down blocks, a non-square image, a row count that does not divide by the rank count, and the chunked convolution path. Co-authored-by: Cursor <cursoragent@cursor.com>
…ncoders Both encoders are their own decoders read backwards, so the residual blocks, mid blocks and causal convolutions are already sharded by the adapters written for the decode side. What is new is the down block and the downsampler it holds. Neither downsampler needs more than its convolution sharded. HunyuanVideo's strides, and 1.5's folds each pair of rows and columns into channels, and both read one input position per output one, so a rank can do the rest to its own rows as long as it holds whole pairs of them. Cutting bands in whole multiples of what the encoder narrows by is what guarantees that, and is also why the assertion inside WanZeroPadConv2d that a band is even still holds. HunyuanVideo ends on a GroupNorm, which reduces over the axis being split and so is wrapped; 1.5 ends on an RMS norm, which reduces over channels and is left alone. Covered at one, two and four ranks, with and without the attention in HunyuanVideo's mid block, over a non-square image, a row count that does not divide by the rank count, and the chunked convolution path. Co-authored-by: Cursor <cursoragent@cursor.com>
The last of the causal families, and the least work: LTX-2's mid block has no attention, so nothing has to be gathered, and its spatial padding already sits inside the convolution rather than being applied around it. Its down block is the one place a checkpoint has a real choice. Three of the four downsample kinds fold space into channels after a stride-1 convolution, and the fourth is a bare strided convolution, which only a stage that does not widen can hold. Both routes are handled and both are tested, along with the reflection padding LTX-2 ships with and the zeros LTX-2.3 uses. Co-authored-by: Cursor <cursoragent@cursor.com>
Two tests per run were failing on EADDRINUSE, a different pair each time. The port fixture was picking from 29800-39799, which overlaps the range Linux hands out to outgoing connections, so a port a test was about to bind could be one the kernel had just given to something else. Picking from 20000-29999 puts every port below that range. The hash is crc32 rather than hash(), whose seed changes per interpreter, so a test now gets the same port every run and a failure can be reproduced. Neither rules out two tests colliding with each other, and no port can be held open for them because rank 0 has to bind the rendezvous socket itself, so run_distributed retries on a fresh port for that. Co-authored-by: Cursor <cursoragent@cursor.com>
AutoencoderKL's encoder, which Flux.2 builds as well, is the last of the seven this repo has a decoder adapter for to get one for the other half. Flux-Kontext, Qwen-Image-Edit and FLUX.2 all encode an image at the size they are about to generate at, which is the encode worth splitting. It splits over the down blocks alone. They carry the image at full size and are the reason to encode in parallel; the mid block and the normalisation after it run at a eighth of that, so the split is undone before them and they run whole on every rank. That is the 2D decoder adapter read backwards, which splits after its mid block rather than before, and it costs a mid block on every rank in exchange for needing nothing said about the attention inside it. Downsample2DAdapter is the part that reaches across the split. Told to pad by hand, as these encoders tell it, the downsampler pads (0, 1, 0, 1) in its own forward and then strides over the result, and a rank's bottom row is only padding if it is the bottom row of the whole image. The zero-padding convolution Wan's resample already used is exactly that pair, so it moves out of the Wan adapter to be shared, and the adapter runs the norm and the convolution itself in that case rather than delegating, so the pad is not applied twice. The scale factor is counted off the blocks rather than taken from the caller, because a 2D VAE does not record its ratio anywhere and a caller reaching for the usual 8 on a three-stage encoder would cut bands a later stage halves into a row the rank does not own. Also corrects two descriptions of the encoder split that still described the padding the ragged split replaced. Co-authored-by: Cursor <cursoragent@cursor.com>
What we tune here is a property of the adapter stack rather than of the weights: PatchGroupNorm issues the same collectives whether its input came from Flux.2 or from torch.randn. What has to be real is the shape of the work, and that lives in a VAE's config.json, so the architecture can be built with random weights in a second and measured on real GPUs without downloading anything. The counts are the point. An optimisation that removes an all_reduce shows up as an integer, not as a timing delta the size of the run-to-run noise on a consumer GPU, so a change like collapsing PatchGroupNorm's three collectives into one can land as an assertion rather than as a benchmark. Latency, peak memory and the single-rank equivalence check come along with it, the last being the invariant every change in here has to preserve. Attribution is by call site, read with sys._getframe rather than by walking the stack, which is cheap enough to leave enabled during a timed decode. This cannot speak for real activation distributions: random weights give a mean near zero and a variance near one, which is the easy case for any variance computation, so a change whose error depends on the mean being large relative to the spread still needs a real decode to sign off.
Three ways this measured something other than what a runner model does. Importing xfuser puts AITER's GroupNorm in torch.nn.GroupNorm's place, and both the adapter and xDiT's selection ask isinstance against whichever class is bound when they ask. Building the VAE before that import left it holding the class from before the swap, so a decoder made entirely of the blocks the 2D adapter wants was reported as fitting no adapter at all. xDiT then puts the stock class back when it validates --use_parallel_vae, because AITER's carries no num_channels for GroupNormAdapter to read; a VAE built here rather than by a runner model has to be walked through both steps by hand. Rank 0 alone computed the reference while the others went on to a barrier. That barrier was the first collective in the process, so it was also where the communicator got built, and the other three sat in init - not in the barrier - until the store gave up ten minutes later. The reference is now taken on every rank, which the matching seeds make free, and only at sizes where one card can hold the whole half; above that it is skipped with a note, since an unsharded decode at that size is the thing sharding exists to avoid. Selection now has to come from xDiT. Choosing an adapter here when xfuser was missing measured this file's opinion of which one fits, and reported the mismatch as an assertion from inside a half-replaced decoder. Agreement is reported against the reference's own scale. On random weights an absolute tolerance says little, and in bf16 one step at magnitude 1 is already 0.008.
PatchGroupNorm made three round trips per call and the first of them moved a single number: the rows this rank holds, all-reduced to the group's total so nelements can be worked out. That sum and the group sums are reductions over the same ranks and neither depends on the other, so they can share one tensor. The arithmetic is untouched - the same two totals come back, and a row count is exact in float32 - which is worth saying because the third reduction is not like this: it takes the squares about the mean the second one produces, and folding it in would mean changing the estimator. Measured at a 64x64 latent on four GPUs, this takes a Flux.2 decode from 75 all-reduces to 50, and the total from 163 collectives to 138. The payload was never the point: all 75 moved 0.01 MB between them, so what this buys is 25 fewer launches and 25 fewer points where the ranks have to meet.
exchange_halo offered its bottom halo to the next rank, then blocked in the receive from the previous one, and only once that had landed did it offer its top halo back the other way. The two directions do not depend on each other, so the second exchange was waiting on the first for no reason: two exposed round trips per sharded convolution where one would do. All four operations are now built as P2POps and issued together. NCCL groups a batch into one operation, which also settles a complaint it was making about the old shape - "An unbatched P2P op (send/recv) was called on this ProcessGroup with size 4. In lazy initialization mode, this will result in a new 2-rank NCCL communicator to be created" - once per op, and a Flux.2 decode at a 64x64 latent issues fifty-six of them. The receive-buffer caching is unchanged, just lifted into a helper now that both sides want it. The bench harness has to wrap distributed_c10d as well as the re-export from torch.distributed: P2POp validates the op it is handed against that module's own isend and irecv, so counting collectives would otherwise have turned this exchange into an invalid op.
use_uniform_patch built the patch boundaries from a rank's own patch size instead of gathering everyone's, on the assumption that all patches are the same size. Patchify cuts bands in whole multiples of the VAE's scale factor and gives the ranks that come first one extra band where the count does not divide, so that assumption is false for any row count that is not a multiple of the rank count, and the boundaries it derives are wrong for every rank past the remainder. Its cache made that worse by keying on the local patch size alone, which does not determine the layout it stands for. Both adapter entry points already passed False, so nothing reachable used it; what was left was a parameter threaded through thirteen files and one test that exercised the wrong path. The all_gather it was meant to avoid is worth avoiding, but it needs the row layout carried forward from Patchify rather than guessed at, which is a separate change. PatchAdaGroupNorm goes with it: nothing has referenced it, and it would raise on any modern torch, calling torch.tensor on a list of tensors and asking sum_to_size to reduce (N, C, H, W) to (N, groups).
Three gaps, all of which hid work we have never measured. --half encoder sharded the encoder and then timed vae.decode, so the encoder adapters have never been through RCCL here at all. The half under test now decides what it is handed and what is called, and the reference threshold is read in latent space either way so one number means the same thing for both. The table held the two 2D families. It now holds the five the adapters claim to support, with the shapes taken from each checkpoint's own vae/config.json rather than shrunk: Wan, Qwen-Image, HunyuanVideo, HunyuanVideo 1.5 and LTX-2. Four of those carry a frame axis, so --frames sizes it, and the compression ratios and latent width are stated in the table because they are readable off a built VAE under three different names depending on the class. Counts came from rank 0, which borders one neighbour where the middle ranks border two, and so sends and receives less of a halo than they do. Every collective is one they all wait on, so the report now carries the most any rank made alongside rank 0's, and the per-rank spread. smoke_families.py builds all seven on the meta device, which catches a config key the installed diffusers does not take before it costs a pod.
At stride 1 every term that mentions where a patch sits cancels out of the halo width, so the all_gather each convolution made to learn the other ranks patch sizes was buying an answer the kernel already gave. Skipping it removes 143 of the 288 collectives a Wan encode makes and 143 of 178 per-conv gathers on the decode. A strided conv still gathers: its halo turns on where in the global stride grid its patch begins, which is not local knowledge. The metadata tuple now carries global_start rather than the whole boundary list, because that is all three callers ever read from it, and it is None exactly when no one paid to find it out.
calc_bottom_halo_width asserts its way out of a patch narrower than the kernel reaches, so the [3,2,2]-with-a-7-kernel corner of the parametrisation had no gathered answer to agree with. Patchify refuses that split long before a convolution sees it.
…me port twice The parametrisation [2,-2] failed the gate as 'process 1 terminated with signal SIGABRT', with 'terminate called without an active exception' and both ranks reporting a clean Gloo connect. That is a rank tearing its context down while its peer still holds one, not a wrong answer, and the same case passed on an earlier run of the same commit. A barrier before the teardown makes the ranks leave in step. The port was also derived from hash(nodeid), which python salts per interpreter, so a failing test bound a different port on every run and could not be asked to fail again. crc32 keeps the per-test uniqueness the fixture was written for and makes it reproducible.
…order A comparison of parallel VAE against tiling against a narrowed window had no home: the harness sharded unconditionally and never tiled, so three of the four arms could only be had from a full model run. Tiling turns out to need nothing from a runner - vae_tiling reads a diffusers VAE and nothing else - so the runner sequence transplants whole, including the part that matters, which is reading the VAE's own tile area before --vae_tile_size narrows it. --no-parallel-vae gives the unsharded baseline the other arms are read against, and the docstring now says what the harness's peak VRAM is and is not: the VAE's own, never a run's.
The four-arm table wants crossing with resolution, and one pod per cell makes that ninety-odd pods for what is a few minutes of actual measurement: the wall clock is startup, the branch install, importing torch, and building the VAE, none of which the second cell needs to pay again. --grid-arms and --grid-shapes cross the two in one process instead. The VAE is still rebuilt per cell, because sharding and the batched decode both replace parts of it in place and unpicking that reliably is harder than paying for a fresh one from a fixed seed. What is reused is the reference, which depends only on the shape and is the expensive part. A cell that throws no longer takes the grid down with it. The ranks vote on each cell before any of them moves on, because a rank carrying into the next cell's collectives while the others unwind an exception would hang the pod rather than lose a row.
…fit from
A world_size 1 VAE group still ran every collective on the decode path. Measured on
flux2's decoder at 1 GPU: 52 of them per decode, about 5% of wall clock, all answering
questions the rank could answer alone.
PatchGroupNorm two all_reduce per norm, ~50 per decode. Summing one rank's numbers
across one rank returns them unchanged, so both are the identity.
gather_patches two all_gather, one for the sizes and one for the payload. The only
patch is the local one and the only size is its own.
_init_rank_mapping one all_gather_object per adapter constructed, since initialize()
clears the cache each time. A one-rank group's mapping is [this rank].
Arithmetic is untouched at every rank count: the guards skip reductions that are already
no-ops rather than computing anything differently, so the ws>1 paths are byte-identical
and ws==1 now returns what the collectives would have returned.
Note this does NOT make one rank use the chunked convolution path - PatchConv2d/3d still
short-circuit to a plain F.conv at world_size 1, ignoring block_size. Letting one rank
chunk without a halo is a separate memory win and a larger change.
Retain the tested local planner and documentation while joining the remote branch history that contains the same benchmark work. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Wrap Diffusers upsampling modules and decoder blocks in place so stale copied implementations and their unused imports can be removed. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep the suite focused on observable behavior by dropping source archaeology, one-time migration assertions, and superseded manual scripts. Co-authored-by: Cursor <cursoragent@cursor.com>
Use the tile height for row sharding so narrow rectangular windows are not rejected by their width. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep both latent axes available so row sharding, quality bounds, and benchmark area accounting each use the dimension they actually mean. Co-authored-by: Cursor <cursoragent@cursor.com>
Make tile marking and latent-row inspection part of the supported public boundary consumed by xDiT. Co-authored-by: Cursor <cursoragent@cursor.com>
Count rows and columns with their own strides so the ordering test also covers non-square grids correctly. Co-authored-by: Cursor <cursoragent@cursor.com>
Expose row count and processing factor so integrations can add operation-specific remediation without parsing error text. Co-authored-by: Cursor <cursoragent@cursor.com>
avjves
reviewed
Aug 14, 2026
There was a problem hiding this comment.
I'll be honest, this is a bit too big of a PR for me to review manually as I haven't touched the VAE code myself, so understanding the nuances of the changes would take a long time 😅 Hopefully @nsakkine has time to review soon.
I ran this through a couple of different models for automatic PRs, but most of the issues were invalid IMO. I left a few comments that might be relevant.
A couple of more high level questions I have:
- IIRC you said this tiling has a slight but not visible numeric change. Is it true? https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/parallelism/vae_parallelism.md In Omni, they also implement tile-based VAE, but claim it's bit-level identical to 1GPU. Is this a different way of doing it?
- Do you know if the AITER GroupNorm could now be used with this, or is it still orthogonal to this? Some parallel VAE implementations do allow for it, i.e they call the GroupNorm still, just with a smaller / different parameters.
Contributor
Author
This is equivalent to the 1 GPU version as well. Due to the nature of tiling, there is always the risk of artifacts (even in the 1 GPU case) primarily from 2 sources - poorly constructed seams (too small an overlap) and normalisation over the tile shifting the color (too small a tile).
I didn't specifically investigate this yet, but I do believe effort should be put into this as a follow up PR - It's an obvious area of improvement for DistVAE and needs to be properly benchmarked. |
nsakkine
reviewed
Aug 18, 2026
Collaborator
|
Since |
nsakkine
reviewed
Aug 18, 2026
nsakkine
reviewed
Aug 18, 2026
nsakkine
reviewed
Aug 18, 2026
Reuse the common chunk-bound calculation so asymmetric convolutions never produce inputs smaller than their kernels. Co-authored-by: Cursor <cursoragent@cursor.com>
Clarify how PatchGroupNorm aggregates shard statistics and that its forward path is inference-only. Co-authored-by: Cursor <cursoragent@cursor.com>
Document why Patchify clones narrow views to preserve rank-local storage and contiguity. Co-authored-by: Cursor <cursoragent@cursor.com>
Cap deterministic move and swap evaluation so first-time high-resolution layouts do not spend seconds hill-climbing negligible load differences. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
Author
Good point. Since Patchify now cuts bands on the VAE’s complete stride grid, strided-convolution boundaries remain aligned and their halo/crop metadata should be derivable locally. That could remove both the patch-size all_gather and the device-to-host synchronization. I’d prefer to handle this as a focused follow-up so I can properly test for any regressions across model families. |
Contributor
Author
|
@feifeibear I believe this PR is ready, could you please take a look when you have a chance? |
avjves
approved these changes
Aug 19, 2026
nsakkine
approved these changes
Aug 19, 2026
Fail stalled Gloo ranks within explicit deadlines and make minimum-dependency CI report the active test instead of running indefinitely. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
DistVAE's main branch currently provides row-sharded
AutoencoderKLdecoding and Wan encoding/decoding.Large decodes can exceed accelerator memory after denoising, including when row sharding is enabled. Diffusers tiling bounds the working set but processes tiles serially. This branch allows distributing complete tiles across VAE ranks, exchanges tile-edge data, and gathers decoded pieces for assembly.
The figure compares row sharding with whole-tile distribution. It shows where each mode communicates, how much work it repeats, and what sets peak activation memory.
Included
distvae.vaeAPI selects adapters, applies rectangular tile plans, installs tiled decode, and distributes whole tiles.AutoencoderKLencoding and support for VAEs used by most of the xDiT supported models.Benchmarks
The harness measures three paths: a complete unsharded decode on every rank, row-sharded adapted layers, and complete spatial tiles distributed across ranks. These synthetic-weight architecture benchmarks cover gfx950 and gfx1201 systems. Peak memory covers decoder execution only. Runs use
bfloat16.Decode paths on gfx950 at eight ranks
Values are latency / peak allocated accelerator memory. The selected tiled plan is the fastest measured plan whose peak is below row sharding. LTX-2 had no qualifying tiled plan, so its table entry is its fastest measured tiled plan. Shapes use output-pixel
H×W×frames; image shapes omit frames. Windows and overlaps use output-pixelH×W.kl(FLUX.1)flux2qwen_imagewanltx2input tensor must fit into 32-bit index mathhunyuan_videohunyuan_video_15The image VAEs are faster and lighter with the selected full-width strips. Wan favors row sharding for latency: 735.1 ms versus 825.6 ms, while the tiled plan lowers peak memory from 3255 MB to 3138 MB. LTX-2 favors row sharding on both metrics; every measured tiled plan used more memory. For both Hunyuan families, row sharding avoids the unsharded failure. Tiling lowers their peak further, but only HunyuanVideo 1.5 is faster than row sharding in this run.
Decode paths on gfx1201
These decoder-half results use four R9700 32GB cards. The tiled column contains the fastest tiled plan tested so far, regardless of memory. Bold marks the lowest latency in each row.
At four ranks:
kl(FLUX.1)flux2qwen_imagewanAt two ranks:
kl(FLUX.1)flux2qwen_imagewanInstrumented communication operations
The logger records tensor collective and point-to-point API calls during one VAE invocation. It does not count timing barriers or
all_gather_objectitself, and a batched point-to-point exchange contributes one call to the total. Tensor collectives used internally byall_gather_objectcontribute bytes without increasing the call total. The table usestotal_calls_max, the busiestrank's count.
kl(FLUX.1)flux2qwen_imageltx2hunyuan_video_15wanhunyuan_videoRow-sharding counts were unchanged across the tested shapes and rank counts for a given family. They follow the adapted architecture: convolutions exchange halos and distributed normalization reduces statistics. Whole-tile distribution normally exchanges tile-edge data and assembled pieces once. HunyuanVideo repeats those operations for each temporal chunk; its 129-frame case used 11 chunks and reported 22 calls.
Comparison with DistVAE beta5 (current main)
DistVAE
0.0.0beta5(current main) supports the plain 2d decoderPatchDecoder, and its Wan adapter can wrap a Wan decoder in place.The comparison installs those beta5 adapters directly and otherwise uses the same architectures, shapes, dtype, warmup, iterations, and communication logger.
gfx950
klklwanwangfx1201
klklklklwanwanwanwanFor
kl, beta5 issues 32 output gathers where the branch issues two; both transfer 28.3 MB through those gathers. The branch also reduces normalization reductions from 75 to 50 and replaces 28 separate send/receive pairs with 28 batched exchanges. For Wan, both versions issue 84 gathers. The branch replaces 693 separate send/receive pairs with 693 batched exchanges. These changes reduce distributed API calls without claiming a reduction in total network traffic.