Skip to content

DistVAE updates including tile-parallel VAE execution #750

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
avjves merged 83 commits into from
Aug 21, 2026
Merged

DistVAE updates including tile-parallel VAE execution #750

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
83 commits
Select commit Hold shift + click to select a range
cc115ed
fix(models): restore enable_tiling/enable_slicing dropped by capabili…
pds-amd Aug 3, 2026
57d2f1b
fix(causal-wan): declare the FP8 GEMMs the model already wires up
pds-amd Aug 4, 2026
1bafbd1
fix(runner): read boolean capabilities as flags, not degrees
pds-amd Aug 4, 2026
d0bab2d
fix(args): describe slicing in the --enable_slicing help
pds-amd Aug 4, 2026
425185a
fix(vae): refuse tiling and slicing flags the installed diffusers can…
pds-amd Aug 4, 2026
3595f88
feat(vae): apply tiling and slicing to every VAE a staged run decodes…
pds-amd Aug 4, 2026
b85ea72
feat(vae): add --vae_tile_size to narrow the VAE tile window
pds-amd Aug 4, 2026
2c061ff
feat(vae): explain a VAE decode that runs out of memory or fails to pad
pds-amd Aug 4, 2026
02efa40
feat(vae): one place that knows which DistVAE adapter a VAE needs
pds-amd Aug 4, 2026
d65e58c
feat(models): enable --use_parallel_vae where DistVAE already fits th…
pds-amd Aug 4, 2026
b51a652
fix(vae): read patch_size only where it means one factor
pds-amd Aug 4, 2026
57a6089
feat(models): enable --use_parallel_vae on the nine models DistVAE no…
pds-amd Aug 4, 2026
3f48b0b
fix(vae): choose an encoder adapter from the encoder's own blocks
pds-amd Aug 4, 2026
4ac0a0b
feat(models): enable --use_parallel_vae_encoder on the models that en…
pds-amd Aug 4, 2026
bea0dab
perf(vae): decode a tiled VAE's same-shaped tiles in one call
pds-amd Aug 5, 2026
7d6d63d
refactor(vae): keep the AITER GroupNorm revert with the code that nee…
pds-amd Aug 5, 2026
3910b64
fix(vae): keep a recounting feature cache the length sharding took fr…
pds-amd Aug 5, 2026
76dbf5b
test(sharding): call the function this module actually has
pds-amd Aug 5, 2026
c2ee7d1
test: skip the flash-attn tests where there is no flash-attn
pds-amd Aug 5, 2026
2ab1740
test: let tests/core import a sibling test module
pds-amd Aug 5, 2026
a1439d2
test(vae): give the core tests torch's GroupNorm, not AITER's
pds-amd Aug 5, 2026
1804ba0
fix(runner): import the two new wrappers where they are used, as the …
pds-amd Aug 5, 2026
48a62e9
test(envs): patch the checks get_device actually makes
pds-amd Aug 5, 2026
8d5b305
test(gilbert): expect the neighbours a row of points actually has
pds-amd Aug 5, 2026
c771574
docs(vae): say what a narrower tile window stops buying
pds-amd Aug 5, 2026
9feba6d
fix(vae): let a narrower tile window keep the memory it asked for
pds-amd Aug 5, 2026
5fa1292
docs(vae): state the batch budget in the units the flag is set in
pds-amd Aug 5, 2026
4467522
perf(vae): deal a tiled decode's tiles out to the ranks, rather than …
pds-amd Aug 6, 2026
1213589
perf(vae): deal the video VAEs' tiles out too, and fix a tolerance
pds-amd Aug 6, 2026
307b030
test(vae): give the patching Wan a config it can decode with
pds-amd Aug 6, 2026
888cc9d
perf(vae): give each rank a band of tile rows to blend, not tiles sca…
pds-amd Aug 6, 2026
dbe0037
test(vae): put torch GroupNorm back before a band test builds a VAE o…
pds-amd Aug 6, 2026
a0d0cc3
fix(vae): decide against bands from the window, not from the tiles a …
pds-amd Aug 6, 2026
549e260
Split tile runs by tile, not by row, so no rank waits on another
pds-amd Aug 6, 2026
16cc3de
Weigh a tile run by the latent it covers, not by the tiles in it
pds-amd Aug 6, 2026
6b4ef41
Import
pds-amd Aug 6, 2026
1ae04f9
Say run rather than band, now that a run is not a row
pds-amd Aug 6, 2026
1b644f6
Exchange the edges the runs need, not every edge there is
pds-amd Aug 6, 2026
eb395ef
Move a tile across where a run cannot be levelled any other way
pds-amd Aug 6, 2026
ae1c7de
Batch tiles only while a tile is too small to keep the device busy
pds-amd Aug 6, 2026
235c31d
Size the test thread pool from the CPU quota, and use RCCL where ther…
pds-amd Aug 6, 2026
71f450e
Floor the tile window rather than optimise below it, and drop tile ba…
pds-amd Aug 6, 2026
9dec9fe
Raise the useful tile-window floor from a quarter to a half
pds-amd Aug 6, 2026
5be3390
Hand round HunyuanVideo's tiles as well
pds-amd Aug 6, 2026
24a897a
Hand round LTX-2's tiles as well
pds-amd Aug 6, 2026
a381a56
Hand round HunyuanVideo 1.5's tiles as well
pds-amd Aug 7, 2026
fb71aa0
Give the tile grid a stride of its own
pds-amd Aug 7, 2026
c14409d
Ask the VAE whether it tiles, not only the command line
pds-amd Aug 7, 2026
fab5651
Check the tile against the ranks whatever set the window
pds-amd Aug 7, 2026
c89ef62
Note where the run decode has been checked bit-exact, and where it ha…
pds-amd Aug 7, 2026
09962d4
Convert every decoding VAE to channels last, not only the first stage's
pds-amd Aug 7, 2026
fdef2bc
Give the VAE tiling decisions one place, and the order between them a…
pds-amd Aug 7, 2026
f04f373
Fix the new setup tests fixtures, and stop the tile tests racing for …
pds-amd Aug 7, 2026
56110d5
Shard a VAE once, however many times a run is initialised
pds-amd Aug 7, 2026
19d3845
Refuse a tile window the stride loop would walk off, and read the ratio
pds-amd Aug 7, 2026
03519e7
Ask the tiling tests the questions they were only looking like they a…
pds-amd Aug 7, 2026
11251ec
Only compare the crop against the step where the two windows have a r…
pds-amd Aug 7, 2026
d999317
Write the tail both tiling loops end with once
pds-amd Aug 7, 2026
dfc1556
Move VAE orchestration into DistVAE
pds-amd Aug 10, 2026
9825979
Harden DistVAE integration tests
pds-amd Aug 10, 2026
09a85fc
Expose rectangular VAE tile controls
pds-amd Aug 10, 2026
05eac22
Remove scalar VAE tile sizing
pds-amd Aug 10, 2026
d0c79a3
Merge remote-tracking branch 'fork/feat/vae-tile-batching' into feat/…
pds-amd Aug 11, 2026
377b2f6
Plan VAE overlap for output shape
pds-amd Aug 11, 2026
9014ed5
Migrate VAE tiling to DistVAE beta9
pds-amd Aug 11, 2026
590ff43
Extract VAE orchestration into manager
pds-amd Aug 11, 2026
91db24b
Merge remote-tracking branch 'origin/main' into feat/vae-tile-batching
pds-amd Aug 11, 2026
e064a0b
Fix VAE rank setup and tile hints
pds-amd Aug 12, 2026
2731294
Fix VAE tiling configuration checks
pds-amd Aug 12, 2026
53dc28b
Inline VAE overlap argument help
pds-amd Aug 12, 2026
2663105
Refactor VAE lifecycle setup
pds-amd Aug 13, 2026
e7637de
Reduce DistVAE integration surface
pds-amd Aug 13, 2026
3a80944
Remove redundant VAE encoder comments
pds-amd Aug 13, 2026
f25235f
Cover rectangular VAE row sharding
pds-amd Aug 13, 2026
58cdec4
Fix parallel VAE process group selection
pds-amd Aug 13, 2026
db27d8a
Improve VAE row-sharding diagnostics
pds-amd Aug 13, 2026
10ce47f
fix: generalize VAE discovery to any number of pipeline stages
pds-amd Aug 17, 2026
58c9227
fix: require explicit --enable_tiling for VAE tile options
pds-amd Aug 17, 2026
a943f7e
fix: always enable VAE tiling for LTX-2.3 second-stage decode
pds-amd Aug 17, 2026
3d3ec03
fix: disable FP8 GEMMs by default for Causal Wan until properly tested
pds-amd Aug 17, 2026
a3bef17
fix: apply VAE tile overlap per run from input sample shape
pds-amd Aug 17, 2026
8d8021e
refactor: centralize runner preparation hook
pds-amd Aug 19, 2026
4b9cb40
chore: bump distvae dependency to >=0.1.0
pds-amd Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ def get_cuda_version():
"transformers>=4.39.1",
"sentencepiece>=0.1.99",
"beautifulsoup4>=4.12.3",
"distvae",
"distvae>=0.1.0",
"yunchang>=0.6.0",
"einops",
"diffusers>=0.33.0",
Expand Down
Loading