Skip to content

fix Qwen-Image series context parallel #12970

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 8 commits into from
Jan 15, 2026
+5 −4
Merged

fix Qwen-Image series context parallel #12970

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0052f58
fix qwen-image cp
DefTruth Jan 13, 2026
a5957b1
Merge branch 'main' into fix-qwen-cp
DefTruth Jan 14, 2026
caf595a
relax attn_mask limit for cp
DefTruth Jan 14, 2026
22db1fd
CP plan compatible with zero_cond_t
DefTruth Jan 14, 2026
6314b77
move modulate_index plan to top level
DefTruth Jan 14, 2026
928ccdc
Merge branch 'main' into fix-qwen-cp
DefTruth Jan 14, 2026
846bdcc
Merge branch 'main' into fix-qwen-cp
DefTruth Jan 15, 2026
5f5eb4b
Merge branch 'main' into fix-qwen-cp
DefTruth Jan 15, 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
2 changes: 0 additions & 2 deletions src/diffusers/models/attention_dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1573,8 +1573,6 @@ def _templated_context_parallel_attention(
backward_op,
_parallel_config: Optional["ParallelConfig"] = None,
):
if attn_mask is not None:
raise ValueError("Attention mask is not yet supported for templated attention.")
if is_causal:
raise ValueError("Causal attention is not yet supported for templated attention.")
if enable_gqa:
Expand Down
7 changes: 5 additions & 2 deletions src/diffusers/models/transformers/transformer_qwenimage.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -761,11 +761,14 @@ class QwenImageTransformer2DModel(
_no_split_modules = ["QwenImageTransformerBlock"]
_skip_layerwise_casting_patterns = ["pos_embed", "norm"]
_repeated_blocks = ["QwenImageTransformerBlock"]
# Make CP plan compatible with https://github.com/huggingface/diffusers/pull/12702
_cp_plan = {
"": {
"transformer_blocks.0": {
"hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
"encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
"encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
},
"transformer_blocks.*": {
"modulate_index": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
},
"pos_embed": {
0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
Expand Down
Loading