Matrix-Game-3: all multi-GPU runs fail with TypeError - attention() called with fa_version= instead of version=

[MMWilliams]

Summary

Every run with --ulysses_size > 1 crashes immediately in the first DiT block with a TypeError. This includes the configuration in the repo's own Matrix-Game-3/test.sh (which uses --ulysses_size 7/8), so the documented multi-GPU path does not run as shipped.

Tested at commit 71c3cd7f741311f8100f6cf9cde942b6c1378d11.

Error

File "Matrix-Game-3/wan/distributed/ulysses.py", line 63, in distributed_attention
    x = attention(
TypeError: attention() got an unexpected keyword argument 'fa_version'

Cause

wan/distributed/ulysses.py:63-70 calls attention() with fa_version=:

from ..modules.attention import attention
x = attention(
    q, k, v,
    k_lens=seq_lens,
    window_size=window_size,
    fa_version=fa_version,   # <-- wrong keyword
)

but the parameter in wan/modules/attention.py:176 is named version:

def attention(
    q, k, v,
    ...
    version=None,
):

Only this call site is affected. The other fa_version=fa_version call sites (sequence_parallel.py:307,364 and model.py:591,596,613,615,678,1015,1027) target functions that do declare fa_version, so they are correct.

Fix

One word, in wan/distributed/ulysses.py:69:

-        fa_version=fa_version,
+        version=fa_version,

Reproduce

torchrun --nproc_per_node=2 generate.py \
  --size 704*1280 --ckpt_dir Matrix-Game-3.0 \
  --image demo_images/001/image.png --prompt "..." \
  --ulysses_size 2 --dit_fsdp --t5_fsdp \
  --num_iterations 3 --num_inference_steps 3 --use_int8

Fails identically regardless of --fa_version, GPU model, or whether FlashAttention is installed — the TypeError is raised before any attention backend is selected.

Environment

2x RTX 5090 (sm_120), torch 2.11.0+cu128, Python 3.12, WSL2 Ubuntu 22.04. The bug is hardware-independent.

Related minor issues

Happy to open these separately if you'd prefer — noting them here since they came up while getting the model running on 2 GPUs:

1. WAN_VAE_SEGMENT_SIZE is undocumented. inference_pipeline.py:643 reads it (default 4), and it decides whether inference fits in 32 GB. At 4 the VAE decode exhausts VRAM at 704x1280 even with --use_int8, --t5_cpu and --lightvae_pruning_rate 0.5; at 1 the same run completes. A README note or a --vae_segment_size flag would help.

2. stream_decode swallows exceptions. wan/modules/vae2_2.py:1349 catches Exception, logs one line, and returns None. That None then surfaces ~200 lines away as AttributeError: 'NoneType' object has no attribute 'cpu' (inference_pipeline.py:652), hiding the real CUDA/OOM error. Adding traceback.format_exc() there would make debugging much easier.

Thanks for open-sourcing the weights and code — the one-word fix above got 2-GPU inference working end to end for us.