Skip to content

Make WanTransformer3DModel hidden states contiguous before the block loop #14236

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
sayakpaul merged 2 commits into from
Jul 19, 2026
+3 −0
Merged

Make WanTransformer3DModel hidden states contiguous before the block loop #14236

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
105e4ae
flatten+transpose produces a non-contiguous tensor; make it contiguou…
menglcai Jul 19, 2026
30ff5a8
Merge branch 'main' into update_wan2.2_contiguous
sayakpaul Jul 19, 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
3 changes: 3 additions & 0 deletions src/diffusers/models/transformers/transformer_wan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -670,6 +670,9 @@ def forward(
hidden_states = self.patch_embedding(hidden_states)
hidden_states = hidden_states.flatten(2).transpose(1, 2)

# flatten+transpose produces a non-contiguous tensor; make it contiguous before the block loop.
hidden_states = hidden_states.contiguous()
Comment on lines +673 to +674

@sayakpaul sayakpaul Jul 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So the pattern is that we identify the first occurrence of ops that can render the tensor in a non-contiguous layout and call contiguous() on it?

@menglcai menglcai Jul 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, the pattern is to ensure the inputs (especially hidden_states) are contiguous before entering the transformer block loop (see also #14186, #14203). We observed that ops like flatten+transpose or pos_embed can leave tensors in a non-contiguous memory layout, and calling .contiguous() before the loop can improve performance. The benefit is most significant on ROCm.


# timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
if timestep.ndim == 2:
ts_seq_len = timestep.shape[1]
Expand Down
Loading