Skip to content

Make WanTransformer3DModel hidden states contiguous before the block loop#14236

Open
menglcai wants to merge 1 commit into
huggingface:mainfrom
menglcai:update_wan2.2_contiguous
Open

Make WanTransformer3DModel hidden states contiguous before the block loop#14236
menglcai wants to merge 1 commit into
huggingface:mainfrom
menglcai:update_wan2.2_contiguous

Conversation

@menglcai

Copy link
Copy Markdown
Contributor

What does this PR do?

Similar to #14186

The non-contiguous tensor originates from hidden_states.flatten(2).transpose(1, 2) produces a non-contiguous layout whose overhead accumulates across transformer blocks in WanTransformer3DModel.

This patch improves performance on both CUDA and ROCm, though the gain is more pronounced on ROCm.

Benchmarks (Wan-AI/Wan2.2-TI2V-5B-Diffuser, 1280x704, 50 steps, bf16):

Platform PyTorch Default +contiguous Speedup
NVIDIA RTX 5090 2.12.1+cu132 5.38 s/it 5.22 s/it +3.0%
AMD Radeon AI PRO R9700 (gfx1201) 2.12.0+rocm7.15.0a20260713 12.23 s/it 9.62 s/it +21.3%
AMD Radeon(TM) 8060S Graphics (gfx1151) 2.12.0+rocm7.15.0a20260711 115.73 s/it 79.30 s/it +31.5%

Before submitting

  • Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
    • Did you read the Coding with AI agents guide?
    • Did you run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • Did you read the contributor guideline?
  • Did you read our philosophy doc? (important for complex PRs)
  • Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@sayakpaul

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@github-actions github-actions Bot added size/S PR with diff < 50 LOC models labels Jul 19, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Hi @menglcai, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant