Add Wan2.2: text/image-to-video (issue #126) - #165
Conversation
NSagan271
left a comment
There was a problem hiding this comment.
Sorry for taking so long to review this. Overall, implementation looks clean and my comments are minor. Once this is rebased off of main I can also do some spot testing.
|
|
||
| # [B, S, H, D] -> [B, H, S, D] SDPA -> back (the reference's | ||
| # dispatch_attention_fn native path; bidirectional, no mask). | ||
| out = F.scaled_dot_product_attention( |
There was a problem hiding this comment.
Would it be possible to use flash attention here, or does this attention computation need to be in fp32?
There was a problem hiding this comment.
It's bf16 SDPA, not fp32, and under torch.compile it already dispatches to fused kernels. An explicit FlashAttention backend for the DiT is the follow-up in the benchmark comment above: it pays about 15% at the 720p tier and nothing at dense, with sdpa staying as the verified bit-exact default.
There was a problem hiding this comment.
Makes sense, yeah agreed to have it as a follow-up. Can you put adding flash attention to Wan2.2 as a github issue before merging?
…ar-project#126) Adds the bottom layer of Wan2.2-TI2V-5B support: the config dataclass, the native video DiT (30 layers, 3072 hidden, batched-CFG forward), the inline UniPC solver port (bh2, order 2, flow prediction), and the mstar weight loader that maps the checkpoint onto them. The DiT is native rather than a diffusers wrapper so it can be built on meta device and materialized straight onto the GPU through mstar's loader (no CPU staging copy of 5B weights), and so the serving dtypes are ours to pin: bf16 weights with fp32 islands (time embedder, scale-shift tables, norm affines) that the reference keeps in fp32 and that the numerics depend on. UniPC is ported inline instead of held in a scheduler object because the solver state has to travel with the request (see the model commit). Architecture values are facts of the Wan-AI/Wan2.2-TI2V-5B-Diffusers checkpoint and are hardcoded so that constructing the model never touches the network. The loader checks them against the checkpoint at load time and refuses to serve on any drift, and it requires a bijection between checkpoint keys and native parameters: an unexpected or unloaded key is an error, not a warning.
… (issue mstar-project#126) The four NodeSubmodules the engine executes. The dit submodule owns one UniPC step: it seeds the latents from the request's seed at iteration 0, runs the batched-CFG forward, advances the solver, and carries (latents, step index, UniPC history, last_sample) on the loop's own edges. UMT5-XXL and the Wan2.2-VAE are thin HF/diffusers wrappers — there is nothing to gain from reimplementing them and a great deal of numerics to lose. Two dtype seams worth knowing about: * the VAE decode dtype is chosen from the live cuDNN version, because cuDNN ships fast bf16 conv3d for this decode only from 9.16 and older builds are much slower in bf16 than in fp32; * the decoder holds its own VAE instance in the checkpoint dtype while the I2V encoder holds a bf16 one, so neither has to re-cast the full VAE when the two dtypes disagree. All four nodes run eager. The transformer alone compiles cleanly, but the engine compiles the whole node, and Inductor then fuses the UniPC step's sigma arithmetic into a Triton kernel and passes the deliberately CPU-resident sigma as a device pointer. That sigma is what keeps the solver bit-exact against the reference, so numerics win and the node stays eager. The VAE decode is tiled unconditionally: the untiled conv3d workspace OOMs a 32 GiB card at the dense tier. Tiling is not free (it decodes ~2.5x the pixels it needs); choosing it on free VRAM instead of always is a possible follow-up.
…sue mstar-project#126) The Model ABC implementation and the registry/config entries that make "wan22" servable. Four stateless nodes (text_encoder, vae_encoder, dit, vae_decoder) and four graph walks; T2V runs encode_text -> video_gen, I2V inserts encode_image and swaps in a walk whose dit additionally consumes the persisted first-frame latent. Output is muxed to mp4 by PyAV, at the request's fps (a playback rate, not a generation knob). The denoise loop is dynamic: Loop.max_iters is the config ceiling and the per-request num_inference_steps stops it via check_stop, so steps is a request knob rather than a server one. Why UniPC state rides the loop edges instead of living in a scheduler object, which is the one deliberate structural difference from cosmos3: the solver's history buffer and corrector state are per-request, so putting them on the request's own edges keeps requests independent, keeps the node stateless, and lets the loop resume on whichever rank picks it up. It also means the dual-DiT A14B variants can reuse the loop as-is later. Request validation lives in process_prompt, on the DATA WORKER, because that is the seam where a ValueError becomes a 400 for the one offending request. The same raise at the conductor, or on the result path, is swallowed and the client hangs until timeout, so those checks are backstops only. Validated here: output modality, prompt presence, a declared-but-missing image, and all four geometry knobs — * height/width must be multiples of 32 (VAE downsample 16 x DiT patch 2). 720x1280 is impossible for this model (720/32 = 22.5); the 720p-class tier is 704x1280. Unaligned sizes previously reached the DiT as a mid-forward shape mismatch, which faults the WORKER rather than failing the request — observed as a non-terminating error loop that hung the client, blocked the queue and grew the log to gigabytes. * num_frames must be 4k+1 (the VAE compresses time by 4 around an anchor frame). Non-positive values made the latent time extent <= 0 and crashed the worker on a negative tensor dim; unaligned ones did not fail at all, which is worse — 32 frames requested was silently floored to 29 returned. * fps must be positive. postprocess also rejects it, but postprocess runs after the whole video is generated and its raise is swallowed, so the client burned a full generation and then hung. Each rejection names the rule and the nearest valid values.
Registers wan22 on /v1/videos/generations. The adapter splits the OpenAI-shaped
"size" ("WxH") into the model's width/height kwargs, passes seed/num_frames/fps
as first-class fields and the rest (guidance_scale, num_inference_steps,
negative_prompt) through extra_body, and turns an "image" into an
image-to-video request.
"video" conditioning is rejected outright rather than silently ignored:
Wan2.2 has no video-conditioned mode, and a request that asks for one is
asking for something it will not get.
The native /generate route needs no adapter and serves the same model.
…order (issue mstar-project#126) Three suites and the recorder that makes the third one reproducible. * test_wan22_model.py — dummy mode (no weights, no GPU, no network): graph walks, loop-carried state, walk transitions, check_stop boundaries, mp4 postprocess, and the full request-rejection matrix (bad size, bad num_frames, bad fps, wrong output modality, promptless, image declared but absent). * test_wan22_native_dit_equivalence.py — the native DiT against the diffusers module it replaces: 825/825 parameters bitwise identical under a remap asserted to be a true bijection (injective AND onto — equal counts alone would not catch two checkpoint keys colliding onto one native key), and bitwise-identical per-layer outputs. This is what licenses reimplementing the DiT at all. * test_wan22_reference_equivalence.py — the whole serving path against a recorded diffusers oracle: text embeddings, the 50-step trajectory, the inline UniPC port in lockstep with the reference scheduler, CFG batched vs sequential, final latents, decode PSNR, and tiled vs untiled decode. record_oracle.py records that oracle, and it ships because without it the suite cannot be run on a new machine. Its docstring carries the rule the thresholds depend on: the oracle must be recorded under the serving process's numerics flag (set_float32_matmul_precision("high")) and on the same torch build as the server under test. Record it without the flag and the reference itself moves — by 1.3e-3 at step 0, growing to 4.2 by step 49. An oracle is therefore only valid for the machine it was recorded on, and it is not distributable, so there is no baked-in path: point WAN22_ORACLE_DIR at what you recorded. Unset, the suite skips with instructions rather than silently comparing against someone else's numerics. Thresholds are derived from a matched run and are NOT to be loosened to make a run pass; a regression there is a finding.
…ect#126) A phase-instrumented t2v/i2v benchmark for wan22, with vllm-omni and SGLang-Diffusion as baselines and stock diffusers as the compute floor. The per-phase breakdown reuses mstar's own request profiler (--log-stats-file) rather than adding instrumentation: wan22's graph has one node per phase and the dit node runs once per UniPC step, so dit.total/exec_count IS the mean per-step wall. Per-step p50/p95 needs a distribution the profiler discards, so it comes from an opt-in step log, enabled by an env var, off by default and numerically inert when off. The measurement rules encoded here were each paid for with a wrong number: * Cross-system LATENCY requires one system at a time, alone on one GPU. Serving each system on its own GPU concurrently is valid for throughput and invalid for latency: cards settle at different clocks under power cap (a 6.7% inter-GPU spread was measured), which is larger than the effect being measured. reproduce.sh gains a preflight that tears down the process TREE and then VERIFIES the card is empty — orphaned workers survive a port-based kill, keep their CUDA contexts, and are then counted as the next run's peak. * Peak VRAM is labelled with WHICH metric it is. nvidia-smi memory.used is whole-GPU and is not a working set; torch.cuda.max_memory_allocated is process-scoped but unreadable from a client, so only the in-process diffusers baseline reports it. The two differ by GiB on one identical run. * torch/cuDNN are probed from the SERVER's interpreter (--server-python), not the client's. Baselines run in their own venvs, so a client-stamped torch on a baseline row records a stack that never ran the work. * The baselines' poll interval is 20ms, not 250ms: it biases their e2e upward by half its value, and 125ms is the same order as the differences measured. Every row carries the engine's timing_boundary so two e2e numbers are never differenced without it. * Per-phase columns at concurrency>1 are batch-shared, not per-request (three distinct requests report bit-identical node timings), and are withdrawn. * An i2v cell on a baseline is refused rather than run: those clients send no conditioning frame, so the cell would be a t2v run labelled i2v. reproduce.sh's serve commands encode two baseline gotchas the obvious recipe gets wrong: * vllm-omni: `vllm serve --omni` serves no video routes and /v1/videos/sync 405s; the real path is a different entrypoint and an async job API. It must run vllm 0.18.0 — 0.20.0+ are cu13-only wheels — which pins torch 2.10 on the baseline side (disclosed, since it differs from mstar's). * SGLang: most of its documented diffusion flags are absent from the shipped build and kill the server on unrecognized arguments; its client is a dispatch stub because SGLang-Diffusion is blocked upstream. The measured prices for the deferred work (compile, CUDA graphs, VRAM-conditional tiling) live in the PR description, cited to the recorded artifacts. The layout matches benchmark/cosmos3/: the measurement engine, its profile parser, and reproduce.sh, whose header carries the run order and the rules above.
Adds wan22 to the model tables in docs/models.rst and README.md, and a docs/models.rst section covering the supported variant and checkpoint, both routes with working curl examples, the generation knobs, the size and frame-count rules, and the inline-UniPC design note. Also registers wan22 in the CLI's DEFAULT_CONFIGS, without which "mstar serve wan22" fails despite configs/wan22.yaml existing.
…doc (issue mstar-project#126) Review follow-ups from PR mstar-project#165, the mechanical batch: * dit.py: max_period is an argument of _sinusoidal_timestep_embedding (default 10000, the only value used); the RoPE CPU tables build at init instead of first use — same float64->fp32 path, values bit-identical, only the build time moves. The table factories pass device="cpu" explicitly because init now runs under the weight loader's meta-device context, where an ambient-device arange would build meta tensors. * submodules.py: the text-embed zero-padding comment now states WHY the output pads to text_max_seq_len (the reference does, the DiT cross-attends the full unmasked stream, and the equivalence gates pin it); the compile-region comment is shortened; conv3d decode-dtype resolution carries a TODO to share a helper with cosmos3. * submodules.py: the vae-None gate in Wan22VaeDecoderSubmodule.__init__ is deleted — no live path constructs the decoder without a VAE (dummy mode returns from _create_submodule before construction), and a comment now records that invariant. * test_wan22_reference_equivalence.py: the compiled-vs-eager bound is recalibrated to 4.8 = worst observed x ~1.3 (observed trajectory max: 2.66 RTX 5090, 2.997 H100, 3.675 H200); the header now documents the oracle protocol — record from the suite's own venv with NVIDIA_TF32_OVERRIDE=0 and CUBLAS_WORKSPACE_CONFIG=:4096:8 set for recorder and suite alike (per-process TF32 GEMM algo selection, cross-venv kernel dispatch, and allocator-history-dependent cuBLAS workspace are otherwise unstable — observed on H200 / CUDA 12.9). * benchmark/vllm_omni_instructions.md deleted (stale, pre-wan22); the still-live vllm-omni venv pin recipe moved into benchmark/wan22/reproduce.sh's header, which no longer points at the deleted file.
…embed skip (issue mstar-project#126) Three review follow-ups from PR mstar-project#165 that move work to the prepare_inputs seam: * vae_encoder: the uint8 -> float [0, 1] conversion moves from forward into prepare_inputs — the per-request preprocessing seam — so a future batched forward sees dtype-uniform inputs. Same ops in the same order; numerics unchanged. * dit: async scheduling ON (the cosmos3 pattern). The conductor may dispatch a speculative iteration N after check_stop fired at N - 1; prepare_inputs reads the looped time_index and vetoes the overshoot by returning None, which both engines turn into a skipped forward on current main. The skip is logged per occurrence. check_stop's boundary math is unchanged; its docstring now describes the overshoot instead of denying it. * text_encoder: with CFG off (guidance_scale <= 1.0) the negative prompt's UMT5 forward is skipped entirely; text_embeds_neg becomes a 1-element placeholder that only keeps the persisted edge populated. The dit's batch-1 no-CFG path never reads it. Dummy coverage: the walk-structure test pins enable_async_scheduling=True, new tests pin the overshoot veto (None at k == N, real inputs at k == N - 1) and the negative-skip decision at guidance 5.0 / 1.0 / 0.0.
11e0255 to
d7e35df
Compare
Adds an sglang engine to the wan22 video benchmark: multipart POST /v1/videos -> poll -> GET /content, the same async job API shape as vllm-omni, with its own TIMING_BOUNDARY entry. The old refusal is removed: the R1 blocker (cu12x wheels shipping a CUDA-13-linked deep_gemm) is avoided by installing sglang from its cu129 wheel index, and sglang 0.5.x accepts the 480p dense tier, so the client was written against a live server and produced rows on an H200 before this commit. flow_shift is not an sglang request field; its Wan2.2-TI2V-5B pipeline config pins 5.0 server-side, which matches the grid pin and is recorded per request in the server's sampling-params echo. No trtllm client ships: TensorRT-LLM's VisualGen exists only in CUDA-13-only releases (first at v1.3.0rc12), which cannot run on a CUDA 12.x driver, so it produced zero rows and a client would be dead code by this harness's own rule.
d7e35df to
7f94cc0
Compare
|
Thanks for the review! Rebased onto current main. Everything is re-verified on an H200, including the bit-exact comparisons against diffusers. I re-measured performance on the H200, now with SGLang-Diffusion following the same protocol as the
Per-step times for vllm-omni and SGLang are derived from each server's own logged denoise rate; peak VRAM is whole-GPU for every column. mstar is fastest at the dense tier and second at the 720p headline tier, where SGLang's FlashAttention kernel outpaces our compiled SDPA attention on long sequences: on its For the attention kernel I can open a small follow-up PR: an optional |
There was a problem hiding this comment.
Looks good to me, and I also ran some spot checks / looked at the outputs and it seems correct. @Gaurav-Shah05 will leave it to you to merge, and can you also make a short github issue for adding a flash attention backend?
|
Thanks! I will merge this, and I have opened up the issue #208 |
Wan2.2-TI2V-5B: text/image-to-video (issue #126)
Adds
wan22- Wan2.2-TI2V-5B (Wan-AI/Wan2.2-TI2V-5B-Diffusers) - as a served model:text-to-video and image-to-video on
/v1/videos/generationsand/generate.What's in it:
built on the meta device (shapes only, no memory), then checkpoint weights load straight
into GPU memory with no CPU staging - bf16, with fp32 kept exactly where the reference
numerics need it.
and risks numerics.
Wan2.2 also ships larger 14B variants ("A14B") that run two DiTs per video, one for the
early noisy steps and one for the late refinement steps. This PR serves only the 5B; the
config rejects the 14B variants rather than half-supporting them.
The graph
Four stateless nodes (no KV cache):
text_encoder,vae_encoder,dit,vae_decoder.Three walks:
encode_textruns UMT5 once and persists the embeddings;encode_image(I2Vonly) encodes the conditioning frame;
video_genloops the dit node for the requestednumber of steps, then decodes. The denoise loop is a named loop, so each request sets its
own step count through
check_stop, and everything the solver needs between steps travelson the loop's own edges.
That last part is the one deliberate departure from cosmos3, which holds its scheduler as a
per-request Python object. Solver state is part of the request, so wan22 carries it on the
request's graph edges next to the latents. The dit node stays stateless, requests cannot
leak into each other, and the state follows the request to whichever rank runs the next
step - which is what will let the 14B variants swap their two DiTs mid-loop without
touching the loop. The port matches the reference scheduler exactly, verified
step-for-step in the tests.
What this PR optimizes, and what it doesn't
The inner DiT region (patchify -> blocks -> head) is torch.compiled while the
UniPC solver stays eager, since its CPU-resident sigma is what the solver's bit-exactness
depends on; the two CFG branches run as one batched forward instead of two calls; the
meta-device weight loading described above; the VAE decodes in fp32 or bf16 depending on
which the installed cuDNN runs faster; VAE decode is VRAM-gated - untiled when free memory
at decode time comfortably clears the estimated untiled peak, tiled otherwise, and the
tiled path is what still makes a 32 GB card work; frames cross the worker boundary as
uint8; and one resolved-settings line is logged per request, so an operator reads back the
knobs the server resolved instead of asserting them (the vllm-style arg dump).
Not done here, priced below: CUDA-graph capture, batching across requests, CFG parallelism
(the two branches share one GPU today; running them on separate ranks is expressible in the
graph but no such config ships), and multi-rank placement in general - the shipped config
is single-GPU.
Correctness
mstar's full path - tokenize, UMT5, per-token timesteps, native DiT, inline UniPC - matches
stock diffusers bit for bit at every one of the 50 denoise steps, under matched numerics.
The native DiT is verified layer by layer against the diffusers module (identical outputs
at every one of the 30 blocks), and the weight loader maps each of the 825 checkpoint
weights to exactly one native weight, with the tests failing if any is missing, duplicated,
or left over. Decode PSNR against the reference video is 35.0 dB on H100.
One place our computation differs from the reference: diffusers computes the guided and
unguided predictions in two model calls; we compute both in one batched call. On H100 the
results are identical. On the RTX 5090 dev GPU they differ by that card's rounding noise,
which the test bounds allow for.
The tests compare against recordings made from stock diffusers. Recordings depend on the
GPU and torch build, so they cannot be shipped;
test/wan22/record_oracle.pyregeneratesthem in a few minutes, and without one the suite skips and prints instructions.
Performance
1x H100, every system alone on the same GPU, run one after another, 10 runs per cell.
Dense tier, 480x832 x 33 frames, concurrency 1:
mstar is the fastest of the three at every tier measured: 1.4x vllm-omni's best previous
numbers and 1.76-1.83x stock diffusers at the dense tier. The two optimizations account for
the gain additively (both off: 8.89 s). Throughput plateaus at ~0.22 req/s from concurrency
2 - one request saturates the GPU.
Pending
Issue #126 checklist
Known limitations
construction (the conditioning image is inserted where and when the reference inserts it)
and by inspecting outputs, but has no recorded reference trajectory yet.
cleanup dependency; ASCII prompts are identical).
Request validation
Height and width must be multiples of 32 (the VAE shrinks each side 16x, then the DiT
patches 2x2), num_frames must be 4k+1 (the VAE compresses time 4x around an anchor frame),
and fps must be positive. All three fail fast as a 400 naming the rule and the nearest
valid values.
Reproduce