Repository navigation
perf(wan): 2.3x faster streaming Wan VAE decode - #425
Merged
Merged
Conversation
The streaming decoder (Cosmos3VAEDecoderARSubmodule) re-decoded the last 8 latents of the stream ahead of every window from an empty causal-conv cache, so each window paid for 11 latents to keep 3, and the result only approximated one decode of the whole stream. On one B200 (Wan 2.1 VAE, 3-latent chunk, bf16, compiled) that is 335 ms per chunk. - A live stream keeps the decoder's causal-conv cache in its request state (windowed_decode_warm_cache, on by default). A window decodes only its new latents, and the frames are those of one decode of the whole stream. The cache is freed after the last window or when the request goes away; sessions still store only latents. A stream without a cache (first window, resumed session, switch off) takes the old re-decode once. - The decode runs on channels-last weights and activations (vae_decode_channels_last, both decoder nodes), with the H/W zero padding done inside the conv, so no conv pays a layout conversion. - A window's new latents go through the decoder in one call (windowed_decode_latents_per_call=0; 1 restores per-latent calls). Peak activation memory grows with the latents per call. - cudnn.benchmark is set around the VAE decode only (vae_decode_cudnn_benchmark) and restored afterwards. - The wan22 node gets the channels-last and cuDNN switches. Compiling its decoder step is available (compile_vae_decode) but stays off. The shared pieces live in mstar/model/components/wan_vae.py. CPU tests cover the warm-cache stream against vae.decode for both decoder families (bit for bit in fp32, including a lone clip-start latent and a ragged tail), the first-window re-decode and cleanup paths, concurrent streams, and the layout switch.
Each warm stream keeps its decoder cache (1.76 GiB for the Wan 2.1 VAE, 1.33 GiB for Wan 2.2 at 832x480 in bf16) outside any engine resource pool. windowed_decode_max_warm_streams (default 0, no cap) bounds them per decoder copy: a stream that starts while that many streams hold a cache takes the stateless context re-decode for its whole life. The choice is made at the stream's first window, so a stream never switches paths. A debug log reports the live warm streams and their cache bytes when one starts or ends. The windowed configs carry the key commented out, and the windowed_decode_latents_per_call comment now states its memory trade-off: one decoder call per window is about 6% faster at 832x480 for a 21-27 GiB larger transient peak. The docstrings say what the first-window re-decode covers (a stream's first window, fresh or resumed) and that a stream stays on one decoder copy. CPU tests: with the cap reached a second stream stays stateless and matches the stateless path while the first stays warm, 0 means no cap, and a slot frees when a stream ends or is aborted.
arun-elr
force-pushed
the
perf/wan-vae-decode
branch
from
October 13, 2026 19:12
7a1b4cf to
60fbc5c
Compare
Measures the Wan VAE decode the way a streaming rollout uses it: a latent stream encoded from a natural clip, 3-latent chunks (12 frames at 480x832), 10 warm-up and 75 timed chunks replayed twice, median ms per chunk, one variant per fresh process with fresh compile caches, PSNR against an fp32 TF32-off reference. bench_wan_decoder.py runs the fp32 reference, stock diffusers, and M*'s streaming decoder node with this change's switches turned on one at a time (mstar_step0 is the stateless context re-decode, mstar_warm all switches on); run_all.sh runs that ladder for both VAEs and prints ms per chunk, peak memory and PSNR. The prototype harness, ablations, output-equivalence checks, profiles, charts, the vLLM-Omni / SGLang harnesses and all results are on the perf/wan-vae-decode-results branch.
arun-elr
force-pushed
the
perf/wan-vae-decode
branch
from
October 13, 2026 19:57
60fbc5c to
daa4414
Compare
One decoder call per latent is now the default for the streaming decode: the same decoder calls as vae.decode (byte-identical eagerly in the default layout), and the smallest transient peak. Decoding a window's new latents in one call (0) stays available. Served end to end on one B200 (cosmos3_edge.yaml, Wan 2.2 VAE, the DiT on the same GPU, 241 frames in 8 windows), 1 was as fast as 0 or faster everywhere (rollout 11.88 vs 12.01 s, first frame 1.38 vs 1.61 s), and at 8 concurrent streams 0 ran out of memory in 7 of them while 1 completed all 8 at 60.2 GiB. On the decoder alone (served 8-latent windows at 832x480, Wan 2.1 VAE) 1 takes 60.0 ms per 12 frames at an 8.2 GiB peak, 0 56.4 ms at 35.1 GiB. The config comment lists 1 / 2 / 3 / 0, the windowed configs carry the key commented out, the tests treat 0 as the opt-in, and the benchmark's mstar_warm follows the config defaults (mstar_warm_lpc0 is the opt-in).
arun-elr
changed the title
…aming decoder vae_decode_cudnn_benchmark now defaults to False: the timed algorithm pick can differ between launches, and it gains nothing on the Wan2.2 VAE. The wan22 node and the whole-clip Cosmos3 decoder behave as on main: the streaming node converts its own decoder-only VAE copy to channels-last instead of the shared instance, and only its decodes take cudnn.benchmark. List the new tests in test/cpu-core.txt (they skip without diffusers).
The cap check runs on the worker's GPU thread while its main thread can remove request states; iterating the live dict could raise.
Inductor drops diffusers' clone of each cached conv-input slice on channels-last activations, so a compiled call left every cache slot as a view of its conv's whole input buffer and a stream held 2-5x its logical cache. Cache slots are now written as compact copies (FeatCache), and compact_feat_cache catches any slot that still comes back as a view.
mstar/model/components/wan_vae.py imports diffusers at module level, and a top-level import of it made mstar.model.cosmos3 unimportable without diffusers (CI's CPU job collects test_model_kwargs_tensor_transport.py, which imports Cosmos3Model). Only the streaming decode paths load it now.
… warm path is bit-exact latents_per_call=0 costs ~10 GiB more peak on the served window after the compact cache (was ~21). The warm path is bit-identical to a whole-stream decode without compile at =1 (as the CPU tests assert), and matches it to bf16 rounding under compile.
arun-elr
marked this pull request as ready for review
October 14, 2026 01:43
… results branch The directory alone reproduces the headline; the results-branch link is for the full data.
merceod
approved these changes
Oct 14, 2026
merceod
left a comment
merceod
left a comment
Collaborator
There was a problem hiding this comment.
Read the change and the benchmark harness, ran the new tests (63 pass with diffusers 0.41), and measured it on H200s against the merge base. The numbers hold up and the output is right, so approving. Details:
- Exactness, real Wan 2.2 weights. The warm windowed decode (eager, default layout, one latent per call) is byte-identical to one stock
vae.decodeof the whole stream for windows of [1,3,3,3,2], [4,4,4] and [1,11] latents (45 frames at 480x832), and the per-stream cache is 1,425,216,000 bytes, the 1.33 GiB the description states. Channels-last eager differs by at most 4 of 255 levels (56.1 dB vs the whole-stream decode), compiled by at most 6 (54.6 dB), and main's stateless re-decode by up to 8 levels on the same random latents. - Decoder alone, in-tree harness, Wan 2.2, 3-latent chunks. Main's method 620 ms per chunk (620.5 ms with main's own code run through the harness), this PR's defaults 144.8 ms, so 4.3x on an H200 against 4.7x on your B200. Stock diffusers bf16 sits at 279 ms. PSNR vs fp32 is 54.5 dB here against 51.1 dB for the stateless path. Peak memory 5.48 GiB at the defaults, 9.61 with
latents_per_call=0. - Served Cosmos3-Edge, shipped config, 241 frames at 20 steps, 832x480, both trees back to back on one GPU. kv rollout 14.4 s vs 18.2 s (medians of 3, -21%), chained 12.9 s vs 17.5 s (-27%), time to first window 1.9 s on both, four concurrent streams 67 s vs 82 s per round (-19%), eight streams 134 s vs 165 s. Aborted streams (closed after 3 windows) leave nothing behind on either tree, and repeat runs are byte-identical within each tree. Reserved GPU memory on the branch grows with the peak number of concurrent streams (39.1 GB idle, 42.8 after 4, 49.7 after 8, which is 8 caches) and then stays flat through aborts and further rollouts, so the caches go back to the allocator. Main stays at 38.9 to 41.7 GB.
- 8 streams. One of three branch rounds lost 3 of 8 streams to the UniPC
assert self.this_order > 0in the DiT scheduler. The other two branch rounds and all three main rounds ran 8 of 8. That matches the pre-existing, load-dependent assertion the description lists; nothing here points at the decoder. - wan22 node. Same picture on both trees for a seeded 9-frame 256x256 clip at 4 steps, and within one server launch every request after the first is byte-identical on both trees (about 0.26 s a request on both once warm). Pixel parity between the trees cannot be measured end to end on this stack: the first request of a launch, and each new launch, render the same seeded sample a little differently (around 20 dB between renderings, the same on main), so the whole-clip decoder rests on the harness exactness check above.
- One observation about the served output, not about this PR. With the shipped config the kv windowed rollout drifts away from the conditioning image within the first two windows on this stack and is noise by frame 60, on main and on the branch alike (same latents feed both decoders, and the two trees' frames show the same content). The chained rollout keeps the scene but smears it more with every window. A non-windowed 121-frame clip stays on the scene, and a prompt that matches the image (a dashcam drive along a coastal road) changes nothing, so it is not the prompt. I would look at the windowed DiT path separately.
Notes:
- The 4.7x decoder headline is mostly work reduction: the 3-latent benchmark chunk decodes 11 latents on main and 3 here, so 3.4x of it is fewer frames decoded. The served 2.3x per window is the number that matters, and the description does say so lower down. Worth leading with it.
- The per-stream cache is 1.33 GiB only in bf16. An fp32 decode path doubles it, and the comments state 1.33 unconditionally.
- With
windowed_decode_latents_per_call=0the context re-decode is grouped too, so a resumed stream's first window (up to 13 latents) is one eager call, well above the documented peak for that opt-in. Grouping only steady windows would keep the documented peak honest. - An exception inside a window leaves a half-updated cache in the request state. Nothing retries today, so it is harmless, but a try/except that drops the cache would make it safe if that changes.
- None of the five new config keys is validated (a negative cap,
latents_per_callabove the window). - The two new test files are on the CPU core list but CI's core job installs no diffusers, so they never run there (42 tests). The description notes this as a follow-up; it is worth doing with this PR so the bit-exact tests guard the code they cover.
- The in-tree bench reproduces the 3-latent rows. The served-window row, the whole-stream PSNR and the byte-identical check are on the results branch only.
| else: | ||
| if st["ar_warm"]: | ||
| feat_cache = _wan_vae().new_feat_cache(self.vae) | ||
| pixels = self._decode_stream(stream_tail, feat_cache, clip_start=True) |
Collaborator
There was a problem hiding this comment.
At windowed_decode_latents_per_call=0 this first call groups the whole stream_tail, so a resumed stream's first window (8 latents of context plus the new ones, up to 13) goes through the decoder in one eager call, well above the peak the config comment quotes for =0. Grouping only the steady windows, or documenting the resumed-first-window peak, would keep that opt-in honest. Also worth a try/except that drops ar_feat_cache if a decode raises mid-window, so a half-updated cache cannot be continued if a retry ever lands.
Contributor
Author
There was a problem hiding this comment.
Both done:
=0grouping: only steady windows are grouped now. A stream's first window, fresh or resumed (context plus new latents, up to 13), runs one latent per call like the default, so the documented=0peak holds. A test checks the call sizes. The cost is one more compiled decoder graph on a=0server's first stream.- Failed window: if a warm window's decode raises,
ar_feat_cacheis dropped before the exception propagates.ar_tailandar_chunkshaven't advanced, so a retry re-decodes behind its context like a first window, then carries on warm. A test injects the failure.
| # re-decode for its whole life. The caches sit outside any engine resource | ||
| # pool, and the shipped configs set no max_concurrent_requests, so with no | ||
| # cap their total is bounded only by how many streams are admitted. | ||
| windowed_decode_max_warm_streams: int = 0 |
Collaborator
There was a problem hiding this comment.
Two small things on the new keys: nothing validates them (a negative cap behaves like no cap, latents_per_call above the window size silently groups everything), and the cap is a count that is blind to resolution, so the 1.33 GiB per stream it implies holds only at 832x480 in bf16 (an fp32 decode path doubles it; the comments state 1.33 unconditionally).
Contributor
Author
There was a problem hiding this comment.
Both fixed:
- Validation: the five keys are checked at load, after the yaml overrides. A
ValueErrornames the key when a switch isn't a boolean (a yaml"false"string used to read as true), when the cap orlatents_per_callisn't an integer >= 0, or whenlatents_per_callis larger than the default window (8 latents at 29 frames). Tests cover each case and a model load with a bad override. - Cache size: the comments now say 1.33 GiB holds at 832×480 in bf16, grows with frame area and doubles in fp32. The cap counts streams, not bytes, so size it from the per-stream cache at the served resolution and dtype. A byte budget would track resolution automatically; I've left that for later.
Collaborator
|
@arun-elr Overall looks good, quick question: could the Wan 2.2 model in the system also reuse the same decoder to get performance gains? |
1.33 GiB holds at 832x480 in bf16; an fp32 decode (cuDNN before 9.16, or COSMOS3_VAE_DECODE_FP32) doubles it, and the warm-stream cap counts streams, not bytes.
At windowed_decode_latents_per_call=0 a stream's first window grouped its whole decode input into one eager call: 7 latents for a fresh stream, up to 12 behind a resumed session's context, above the peak documented for =0. The first window now runs one latent per call, as the default does.
The cache is updated in place, so a failed decode left it half-advanced in the request state. Nothing retries a window today; if something does, it now re-decodes behind the window's context instead of continuing a broken cache.
The three switches must be booleans (a yaml "false" string was truthy), the warm-stream cap and latents per call integers >= 0 (a negative cap read as no cap), and latents per call no more than the default window's latent frames (more read as one call per window).
The Wan VAE decode tests on the core list import diffusers and were skipped in CI. They build tiny random models (no weights) and take about 40 s on CPU.
arun-elr
changed the title
Contributor
Author
@NSagan271 Thanks! Yes, the |
Measured on B200 serving cosmos3_edge.yaml at 832x480 in bf16: 48.50 GiB peak allocated at windowed_decode_latents_per_call=0 against 33.95 GiB at the default, not the ~10 GiB the comment gave. Comment-only.
Contributor
Author
|
Thanks @merceod, especially for re-measuring on H200s and the exactness check on real Wan 2.2 weights. Changes in this round:
Deferred:
|
Contributor
Author
|
@NSagan271 Follow-up with the B200 numbers. Decoder only (121 frames at 1280x704), the |
Contributor
Author
|
@NSagan271 the follow-up for the |
This was referenced Oct 15, 2026
Closed
merceod
pushed a commit
that referenced
this pull request
Oct 18, 2026
… the #425 decoder (#437) The wan22 VAE decoder node reuses the shared Wan VAE decoder from #425 for untiled decodes at the listed sizes (vae_decode_fast, default on; vae_decode_fast_sizes, default 1280x704): channels-last, compiled DecoderStep, 1.74x faster decode on an H200. Other sizes and tiled decodes switch the VAE back to its loaded layout and run stock vae.decode, byte-identical to the off switch with no compile. The compiled decode runs in Inductor's deterministic mode so every server process decodes the same latents to the same frames.
dghosef
added a commit
to dghosef/mstar
that referenced
this pull request
Oct 21, 2026
The fp32 Wan2.1 VAE decode was a third of a chunk. Channels-last weights and activations (the mstar-project#425 path, already wired as vae_decode_channels_last) plus cudnn.benchmark around the decode only: 547 -> 340 ms per 16-frame chunk at 480x832 on an H200, still fp32. Frames differ from the contiguous decode by conv accumulation order only (69.4 dB, max 2 uint8 levels). Served A/B, H200, 93 frames, 2 rounds x 2 reps ABBA with the DiT change: 9.60 -> 11.34 fps, first chunk 1.51 -> 1.30 s; both arms deterministic. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dghosef
added a commit
to dghosef/mstar
that referenced
this pull request
Oct 21, 2026
The fp32 Wan2.1 VAE decode was a third of a chunk. Channels-last weights and activations (the mstar-project#425 path, already wired as vae_decode_channels_last) plus cudnn.benchmark around the decode only: 547 -> 340 ms per 16-frame chunk at 480x832 on an H200, still fp32. Frames differ from the contiguous decode by conv accumulation order only (69.4 dB, max 2 uint8 levels). Served A/B, H200, 93 frames, 2 rounds x 2 reps ABBA with the DiT change: 9.60 -> 11.34 fps, first chunk 1.51 -> 1.30 s; both arms deterministic. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dghosef
added a commit
to dghosef/mstar
that referenced
this pull request
Oct 21, 2026
The fp32 Wan2.1 VAE decode was a third of a chunk. Channels-last weights and activations (the mstar-project#425 path, already wired as vae_decode_channels_last) plus cudnn.benchmark around the decode only: 547 -> 340 ms per 16-frame chunk at 480x832 on an H200, still fp32. Frames differ from the contiguous decode by conv accumulation order only (69.4 dB, max 2 uint8 levels). Served A/B, H200, 93 frames, 2 rounds x 2 reps ABBA with the DiT change: 9.60 -> 11.34 fps, first chunk 1.51 -> 1.30 s; both arms deterministic. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dghosef
added a commit
to dghosef/mstar
that referenced
this pull request
Oct 21, 2026
The fp32 Wan2.1 VAE decode was a third of a chunk. Channels-last weights and activations (the mstar-project#425 path, already wired as vae_decode_channels_last) plus cudnn.benchmark around the decode only: 547 -> 340 ms per 16-frame chunk at 480x832 on an H200, still fp32. Frames differ from the contiguous decode by conv accumulation order only (69.4 dB, max 2 uint8 levels). Served A/B, H200, 93 frames, 2 rounds x 2 reps ABBA with the DiT change: 9.60 -> 11.34 fps, first chunk 1.51 -> 1.30 s; both arms deterministic. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 23, 2026
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
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.
Summary
M*'s streaming Wan VAE decoder (
Cosmos3VAEDecoderARSubmodule) re-decodes 8 latents of context from an empty causal-conv cache before every window. This PR keeps each live stream's conv cache, so a window decodes only its new latents, and makes the decode cheaper (channels-last, padding inside the conv). Served, a window's decode is 2.3× faster (534 → 236 ms) and a rollout 16 to 23% faster.3-latent benchmark chunk; served kv windows have 8 new latents.
Cosmos3-Edge, shipped config, one B200 with the DiT co-located:
Decoder alone, ms per 3-latent chunk (480×832, bf16), same latents for every engine:
windowed_decode_latents_per_call=0+cudnn.benchmark)--vae-fast-path channels_last)=1, lowest memory)--quality extra-high)Most of the 4.7× against M* before is work not done: main decodes 11 latents per 3-latent chunk and this PR 3, which is 3.4×; the cheaper decode does the rest. The opt-in row takes 6.5% less time per chunk than vLLM-Omni's best on Wan 2.2 (12.8% on Wan 2.1); the defaults favour memory (5.5 GiB peak) and are within 2.1% (0.7%). All land at bf16-rounding quality; neither other engine streams Cosmos3.
Changes
windowed_decode_warm_cache(true)windowed_decode_max_warm_streams(0 = no cap)windowed_decode_latents_per_call(1)vae.decode's calls, 0 = a whole window)vae_decode_channels_last(true)vae_decode_cudnn_benchmark(false)cudnn.benchmarkThe whole-clip decoder and the
wan22node decode exactly as on main.Where to look
wan_vae.decode_causal: from an empty cache withclip_start, equal tovae.decodecall for call.wan_vae.DecoderStep.__call__: first two calls eager ("Rep"marker); new compiled signatures warm up on a cache copy.wan_vae.FeatCache/compact_feat_cache: compiled cache writes become compact copies (no pinned buffers).Cosmos3VAEDecoderARSubmodule.forward: warm vs stateless chosen once per stream (_warm_slot_free); cache freed after the last window or with the request.Cosmos3Model._build_decoder_only_vae: the channels-last copy; the whole-clip decoder keeps main's layout.Code notes
mstar/model/components/wan_vae.py(decode_causal,DecoderStep,to_channels_last,cudnn_benchmark);ChannelsLastCausalConv3ddoes the H/W zero padding inside the conv, with no padded copy.DecoderStepruns a stream's first two decoder calls eagerly (they still see diffusers'"Rep"cache markers; compiled, the second returns wrong frames on torch 2.12.1) and warms each new compiled signature on a copy of the cache, like the whole-clip decoder's existing_warmed_decode_shapes.FeatCachewrites cache slots as compact copies under compile, andcompact_feat_cachecatches any slot still returned as a view (see Memory).submodules.pyimports the helpers lazily, somstar.model.cosmos3still imports without diffusers.configs/cosmos3_edge.yamlandconfigs/cosmos3_nano_ar.yamlcarry the cap andwindowed_decode_latents_per_callas commented keys.vae_decode_cudnn_benchmarkgains nothing on Wan 2.2 (77.2 vs 77.8 ms) and frames can differ between launches; on an H200 it takes ~18% off a Wan 2.1 chunk.Proposed defaults (happy to change any of these)
windowed_decode_warm_cache: falserestores the stateless decode; withvae_decode_channels_last: falsetoo, today's decode and memory exactly.Correctness
Per-frame PSNR vs an exact fp32 decode, Wan 2.2 VAE, all 1,021 frames: M* before dips at chunk boundaries (worst 34.8 dB); this PR (in-tree node) tracks stock with no drift (all VAEs and references).
Exactness tiers
Whole 1,021-frame stream per configuration, fresh process, uint8 frames vs stock diffusers bf16 and an fp32 reference (TF32 off). These rows ran at
=0withcudnn.benchmarkon; the default=1issues the bit-exact tier's decoder calls and scores the same PSNR vs fp32 (55.40 / 54.42 dB, first 10 chunks). Full tables:RESULTS.md§8.=0,cudnn.benchmarkThe stateless re-decode's difference is the expected price of 8 latents of context (
windowed_decode_context_latentsexists to raise it); with context covering the whole stream it is exact too (test_warm_windows_are_the_whole_stream_decode).Videos
Top row should look identical; the ×10 heatmaps and the PSNR plot show where M* before differed (PSNR = image quality vs the exact decode, higher is closer).
Reference: fp32 (Cosmos3) or stock bf16 (Wan) whole-stream decode. Cosmos3 windows count from 1, Wan chunks from 0; a red tag marks M* before below 40 dB. Each video opens with a 4 s what-to-look-for card, plays at 0.5×, holds 2.5 s on M* before's worst frame, then replays at 1×.
Silent (neither model makes sound). Prompts, seeds, lossless frames: README.
Testing
ruffand 75 targeted + 277 related CPU tests pass (1 failure, also on main). CI's CPU job now installs diffusers, so the two new test files run there.Test list
ruff check .passes.pytest test/modular/test_wan_vae_decode.py mstar/model/cosmos3/tests/test_windowed_decode.py mstar/model/cosmos3/tests/test_windowed.py(torch 2.12.1, diffusers 0.39): 75 passed. They cover the warm stream bit-exact against onevae.decode(test_warm_cache_stream_is_the_whole_stream_decode,test_warm_windows_are_the_whole_stream_decode); cache release after the last window (test_cache_is_freed_after_the_last_window) and on abort throughcleanup_request(test_cache_goes_with_an_aborted_request); resumed sessions; the cap (overflow streams stay stateless, 0 is no cap, the slot frees when a stream ends or aborts, and the cap check survives concurrent request removal); the compact cache under a real Inductor compile on CPU (test_compiled_channels_last_stream_cache_owns_its_storage); the whole-clip decoder left as on main; andcudnn.benchmarkopt-in and scoped to the streaming decode, with the context manager restoring the flag when its body raises (test_cudnn_benchmark_is_scoped); a decode that raises mid-window dropping the cache (test_a_failed_window_drops_the_cache); a stream's first window running one latent per call at=0(test_resumed_first_window_is_not_grouped); and the config checks on the five new keys.pytest mstar/model/cosmos3/tests mstar/model/wan22 test/modular/test_cosmos3_edge_model.py test/modular/test_wan22_*.py test/modular/test_wan_vae_decode.py: 277 passed, 37 skipped, 1 failed. The failure,test_video_postprocess_falls_back_to_pyav, fails the same way on main.diffusers==0.39.0, so the two new test files ontest/cpu-core.txtrun there (54 tests, about 40 s on CPU, tiny random models, no weights).results/h200).Reproduce and data
bash run_all.sh(README); served rollouts, raw data, every table:results/INDEX.mdon the results branch (outside this diff).Commands
Decoder alone, from this branch. The input is a 6.5 s 4K natural video clip (hands drawing on a tablet), cropped to 832×480, ping-ponged to 1,021 frames and encoded with each VAE's own encoder;
prepare_latents.py --video PATHencodes any clip the same way (--randomwithout one). Timing doesn't depend on the latent values; PSNR does.Served rollouts, from the results branch (
e2e-cosmos3-edge):The 3-run tables come from
run_ablation.sh ladderon the results branch; the comparison fromexternal/.results/INDEX.mdlists each dataset (B200 main, H200 cross-check) with its GPU, software stack, code commit, run count and the table it feeds. More: samples, where the time goes, things that did not help.Details
End-to-end setup and tables
configs/cosmos3_edge.yamlas shipped (all nodes in one worker),nvidia/Cosmos3-Edge@344d602b, image-to-video/generatestreaming at 832×480, 241 frames, 20 steps, guidance 6.0, seed 0. kv: 8 windows of 8 latents; chained: 10 windows with 2 latents of overlap. before is main99f5108b; this PR is the PR at913295e4(compact cache,cudnn.benchmarkoff) at its defaults (=1) and with=0. Medians of 3 kv / 2 chained rollouts after a warm-up, one B200. The DiT takes M*'s fallback fromdensetoflashinferattention (thefa3-fwdwheel is sm90-only) in every variant. Setup and per-phase data: this PR, pre-fix. A second session on another B200 (e2e-cosmos3-edge-rerun) agrees: kv 14.26 → 11.93 s, chained 11.89 → 9.04 s.=1(default)=0=1=0=1(default)=0=1=0pre-fix is the PR at
60fbc5cd, before the compact-cache fix, withcudnn.benchmarkon.=1) is paid once.=0fits 8 streams, but it is still no faster end to end (11.94 vs 11.91 s; 17.5 fps either way) and was 0.14 s later to the first frame (0.22 to 0.23 s before the fix):DecoderStepruns a stream's first two decoder calls eagerly, and at=0the second was the first window's whole 7-latent decode. A stream's first window now runs one latent per call at=0too (032362aa); in a later B200 check that cut=0's first-frame lag behind the default from 0.15 to 0.01 s. So=1stays the default.=1, so every frame difference comes from the decoder. This PR at=1scores 53.46 dB (kv) and 52.33 dB (chained) against the fp32 whole-stream decode, the same as pre-fix; main scores 46.52 and 41.34 dB, with its error starting at window boundaries from window 3 on.Design
The stateless re-decode costs about 2× the useful work on the served 8-latent kv windows and 3.4× on the 3-latent benchmark chunks.
Warm path. A stream's first window runs the same context decode, but against a cache the node owns (
new_feat_cache,decode_causal): the last frames each causal conv saw, kept in the request's state asar_feat_cachenext to thear_tailandar_chunksthe node already keeps. Each later window decodes only its new latents and updates the cache in place.Lifecycle. Created on the stream's first window if the copy is below the cap (decided once per stream). Freed when the last window is decoded, or with the rest of the request state when the engine removes the request (
cleanup_requestfromEngine.remove_request: abort, failure, teardown). Tests cover the last-window and abort releases; a failed request goes through the conductor's same removal path as an abort. A decode that raises mid-window drops the half-updated cache at once, so a retried window re-decodes behind its context. Sessions are unchanged: theSessionStorekeeps only the 8-latent tail, and a resumed session's first window is the context re-decode behind it.Routing. Before and after this PR, the conductor binds each request to one copy of every component at admission (
_assign_worker_graphs_to_workers) and keeps it for the request's life. The decoder already relies on that forar_tail,ar_chunksandar_pixels; the cache adds no new routing requirement, and a stream's windows never move between copies.Per-change gains
One more switch per step, compiled, n = 3 fresh-process runs per row; green is the default configuration, orange the opt-in switches. Conv-internal padding comes with step 2. Eager rows, kernel counts and TFLOP/s:
RESULTS.md§1.=0, opt-in)cudnn.benchmark(opt-in)Step 2 is the default configuration. The ladder ran before the compact-cache fix (+1.7 ms per chunk); the defaults row in the Summary table was re-timed after it (
RESULTS.md§0).Full comparison table
Same latents, reference, protocol and timing, one B200; each engine's own streaming decoder, loader and documented flags. M* before and the opt-in row are the 3-run ladder above; the defaults row was re-timed after the compact-cache fix (2 runs,
RESULTS.md§0). Full table (defaults, compile, fp32 serving paths), kernel profiles and harnesses:RESULTS.md§7,external/.--vae-fast-path channels_last)--quality extra-high, bf16)=1, lowest memory, the same decoder calls asvae.decode)windowed_decode_latents_per_call=0+cudnn.benchmark, opt-in)† SGLang converts to uint8 from fp32; stock diffusers converted the same way scores the same, so its higher PSNR comes from the conversion, not the decode. Both other engines ran torch 2.13 (M*: 2.12.1).
Memory
max_concurrent_requests, so with no cap the total is bounded only by admission. At 8 streams (8 caches plus one ~4 GiB transient) the shipped Cosmos3-Edge config peaks at about 44 GiB allocated of the B200's 178 GiB.torch.compile. Inductor removes diffusers'clone()of each cached slice on channels-last activations, which left every slot a view of its conv's whole input buffer. Each slot is now written as a compact copy, so a stream holds exactly its logical cache, as vLLM-Omni and SGLang do. Output is byte-identical; the copy costs +1.7 / +2.2 ms per chunk (Wan 2.2 / Wan 2.1).windowed_decode_latents_per_callsets one decode's transient peak: on a served window=0is 74.3 vs 79.4 ms per 12 frames for about 10 GiB more peak, and end to end no faster (End-to-end above). Served windows before the fix:RESULTS.md§2b; the pre-fix sweep over 1, 2, 3 and 0 latents per call:lpc.csv.=1(default) /=0(end to end)=1/=0(§0)=1/=0(§0)Related, not in this PR
SessionStoredecodes the first window with no context.resume_sessionrequest to the copy holding its session.assert self.this_order > 0, under KV-page-pool pressure; 3 of 8 streams), also on the older pre-fix tree; not touched by this PR. The 8-stream memory peaks are unaffected.wan22decoder node); Cosmos3-edge PR 6: Cosmos3-Edge as a cosmos3 config variant, with streaming rollout and action loop #297 and Sliding Window Attention (AR video gen) + KV cache lifecycle management #198 (the streaming decoder this builds on). Similar techniques landed in SGLang in sgl-project/sglang#34125.Samples: still frames
Rows: stock diffusers, M* before, this PR, then ×10 difference against stock (black = identical). M* before matches stock only while the stream still fits its 8-latent context (chunks 0–2, frames 0–36, up to 1.5 s), then its error appears from chunk 3 (1.5 s) on (worst frame 27.0 dB at frame 49, 2.0 s, in chunk 4); this PR stays near black throughout (within bf16 rounding of stock), with no drift.
Side-by-side clips and per-step metrics: samples. "this PR" in the clips is the opt-in
=0+cudnn.benchmarkstep; the README lists the defaults too.Checklist
ruff check .passes