fix(wan2): staged VAE decode Stream(gpu, N) cross-thread error (#418) - #419
Merged
Conversation
…follow-up)
The Phase-2 staged pipeline (load_text_encoder → encode_text → load_dit →
denoise → load_vae → decode) raised
"RuntimeError: There is no Stream(gpu, 4) in current thread" at the VAE
decode mx.eval. Stages 1 (text encode) + 2 (DiT denoise) passed; stage 3
failed. The monolith generate() path was unaffected because it runs
encode/denoise/decode in one continuous executor call.
Root cause: MLX Metal streams are thread-local. The staged path round-trips
the denoised latent through the event-loop main thread between the denoise
executor call and the decode executor call. A *lazy* mx array (or a
[None]/[0] projection built off-thread) records the producing call's
auto-allocated internal Stream(gpu, N); evaluating it on a different thread
(or after that stream's table entry is gone) raises the error. Controlled
experiments confirmed: a latent built inside the decode executor decodes
fine; one built on the main thread fails; a concrete mx.array(numpy) is
portable.
Fix (2 files, +75/-17):
- wan2.py denoise(): build the 5D [None] batch projection and mx.eval it on
the executor thread before returning (was a main-thread lazy projection).
An eval'd array is portable across threads.
- wan2.py decode()/decode_tiled(): slice latent[0] inside _decode on the
executor thread, not on the main thread. Fixed t-axis index (2 for 5D,
1 for 4D).
- wan2.py load_vae(): get_executor("io") → get_executor("video") for
consistency with decode (io is a different thread pool).
- wan2.py _clear_mlx_cache(): 3 unload methods route mx.synchronize()/
mx.clear_cache() through the video executor thread (not main thread).
- stage.py run_denoise(): mx.eval(latents) before return.
- stage.py decode_wan_vae(): mx.eval(latent) at entry.
Verification:
- Real-model e2e (Wan2.1-T2V-1.3B, fusion-comfyui tests/e2e_wan2_staged.py):
full staged T2V PASSES, 57-65s, output (1,20,480,832,3), non-zero pixel
fraction 0.99. Sequential offload confirmed (mem 10.8GB→5.4GB→290MB).
- Unit (tests/unit/test_wan2_stage_api.py): 20/20 pass.
- ruff clean.
Co-Authored-By: Claude <noreply@anthropic.com>
dahai80
added a commit
that referenced
this pull request
Aug 8, 2026
* chore: bump version 0.8.11 → 0.8.12 Patch: GGUF load guard (#423), Wan2 staged VAE Stream fix (#419), DPO logprobs TypeError fix (#421). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * style: black-format training/reward.py (CI lint fix) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.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
1 participant
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.
Closes #418. Follow-up to #416 (the #410 stage API).
Summary
The Phase-2 staged pipeline raised
RuntimeError: There is no Stream(gpu, 4) in current threadat VAE decode. Root cause: MLX Metal streams are thread-local, and the staged path round-trips a lazy denoised latent through the event-loop main thread between the denoise and decode executor calls; evaluating a lazy array (or an off-thread[None]/[0]projection) on another thread fails.Fix (2 files, +75/-17)
denoise(): build the 5D[None]batch projection andmx.evalit on the executor thread before returning (was a main-thread lazy projection). An eval'd array is portable across threads.decode()/decode_tiled(): slicelatent[0]inside_decodeon the executor thread, not on the main thread. Fixed t-axis index (2 for 5D, 1 for 4D).load_vae():get_executor("io")→get_executor("video")(consistency with decode)._clear_mlx_cache(): 3 unload methods routemx.synchronize()/mx.clear_cache()through the video executor thread.run_denoise()/decode_wan_vae():mx.evalthe latent on the producing thread.The monolith
generate()path is unchanged in behavior (it shares one executor call; the new evals are no-ops for it).Verification
tests/e2e_wan2_staged.py): full staged T2V PASSES, 57-65s, output (1,20,480,832,3), non-zero pixel fraction 0.99. Sequential offload confirmed (mem 10.8GB→5.4GB→290MB).tests/unit/test_wan2_stage_api.py): 20/20 pass.See #418 for the controlled experiments that isolated the root cause.
🤖 Generated with Claude Code