From 8d5436b978796739a4ec64c817c9faa743cfe6ce Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Sun, 20 Sep 2026 19:31:03 +0800 Subject: [PATCH 1/2] refactor attention routes for qwen-image-2.1 --- diffsynth/models/qwen_image_21_dit.py | 171 ++++++++++++-------------- 1 file changed, 82 insertions(+), 89 deletions(-) diff --git a/diffsynth/models/qwen_image_21_dit.py b/diffsynth/models/qwen_image_21_dit.py index ae24a9f26..301293f1e 100644 --- a/diffsynth/models/qwen_image_21_dit.py +++ b/diffsynth/models/qwen_image_21_dit.py @@ -197,6 +197,26 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): ) +def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]: + # Maximal runs of equal image_ids, as (start, end, is_text). Derived once per forward rather than per block: + # tolist() is a device sync, and there is one processor call per layer. + prefix_ids = image_ids[:prefix_len].tolist() + segments = [] + start = 0 + for index in range(1, prefix_len + 1): + if index == prefix_len or prefix_ids[index] != prefix_ids[start]: + segments.append((start, index, prefix_ids[start] < 0)) + start = index + return segments + + +def _attention(query, key, value, attn_mask=None, use_flex=False): + return attention_forward( + query, key, value, q_pattern="b s n d", k_pattern="b s n d", v_pattern="b s n d", out_pattern="b s n d", + attn_mask=attn_mask, use_flex=use_flex, + ) + + class QwenImage21AttnProcessor: def __call__( self, @@ -206,7 +226,8 @@ def __call__( rotary_emb: torch.Tensor | None = None, kv_cache: dict[str, torch.Tensor] | None = None, cache_write_slice: slice | None = None, - is_causal: bool = False, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, ) -> torch.Tensor: query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) @@ -223,17 +244,21 @@ def __call__( query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) + decode = kv_cache is not None and "key" in kv_cache if kv_cache is not None: - if cache_write_slice is not None: - kv_cache["key"] = key[:, cache_write_slice].clone() - kv_cache["value"] = value[:, cache_write_slice].clone() - else: + if decode: key = torch.cat([kv_cache["key"], key], dim=1) value = torch.cat([kv_cache["value"], value], dim=1) + else: + kv_cache["key"] = key[:, cache_write_slice].clone() + kv_cache["value"] = value[:, cache_write_slice].clone() seq_len_q, seq_len_kv = query.shape[1], key.shape[1] - use_flex = FLEX_ATTN_AVAILABLE and isinstance(attention_mask, BlockMask) - if use_flex: + if decode: + # Fast decode path with kv cache + hidden_states = _attention(query, key, value, attn_mask=attention_mask) + elif FLEX_ATTN_AVAILABLE and isinstance(attention_mask, BlockMask): + # Flex attention route on the first step or without kv cache pad_q = int(math.ceil(seq_len_q / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_q pad_kv = int(math.ceil(seq_len_kv / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) - seq_len_kv if pad_q: @@ -241,20 +266,27 @@ def __call__( if pad_kv: key = F.pad(key.transpose(1, 3), (0, pad_kv)).transpose(1, 3) value = F.pad(value.transpose(1, 3), (0, pad_kv)).transpose(1, 3) - - hidden_states = attention_forward( - query, - key, - value, - q_pattern="b s n d", - k_pattern="b s n d", - v_pattern="b s n d", - out_pattern="b s n d", - attn_mask=attention_mask if not is_causal else None, - is_causal=is_causal, - use_flex=use_flex, - ) - hidden_states = hidden_states[:, :seq_len_q] + hidden_states = _attention(query, key, value, attn_mask=attention_mask, use_flex=True)[:, :seq_len_q] + else: + # Splited attention route on the first step or without kv cache + valid = None if key_valid is None else key_valid[:, None, None, :] + prefix_len = segments[-1][1] if segments else 0 + outputs = [] + for start, end, is_text in segments: + seg_mask = valid if valid is None else valid[..., :end] + if is_text: + seg_len = end - start + causal = torch.cat( + [ + torch.ones(seg_len, start, dtype=torch.bool, device=query.device), + torch.tril(torch.ones(seg_len, seg_len, dtype=torch.bool, device=query.device)), + ], + dim=1, + )[None, None] + seg_mask = causal if seg_mask is None else seg_mask & causal + outputs.append(_attention(query[:, start:end], key[:, :end], value[:, :end], attn_mask=seg_mask)) + outputs.append(_attention(query[:, prefix_len:], key, value, attn_mask=valid)) + hidden_states = torch.cat(outputs, dim=1) hidden_states = hidden_states.flatten(2, 3).type_as(query) hidden_states = attn.to_out[0](hidden_states) return attn.to_out[1](hidden_states) @@ -314,7 +346,8 @@ def forward( target_token_mask: torch.Tensor | None = None, kv_cache: dict[str, torch.Tensor] | None = None, cache_write_slice: slice | None = None, - is_causal: bool = False, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, ) -> torch.Tensor: mod1, mod2 = modulation.chunk(2, dim=-1) @@ -325,7 +358,8 @@ def forward( rotary_emb=rotary_emb, kv_cache=kv_cache, cache_write_slice=cache_write_slice, - is_causal=is_causal, + segments=segments, + key_valid=key_valid, ) hidden_states = hidden_states + img_gate1.tanh() * attn_output @@ -405,13 +439,11 @@ def __init__( axes_dims_rope: tuple[int, int, int] = (16, 56, 56), eps: float = 1e-6, causal_condition: bool = True, - causal_block: bool = True, ): super().__init__() self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.causal_condition = causal_condition - self.causal_block = causal_block self.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(axes_dims_rope)) self.time_text_embed = QwenImage21TimestepProjEmbeddings(embedding_dim=self.inner_dim) @@ -498,10 +530,10 @@ def forward( temb = self.time_text_embed(timestep, hidden_states) modulation = self.modulation(temb) - if kv_cache is not None and (not self.causal_condition or not self.causal_block): + if kv_cache is not None and not self.causal_condition: raise ValueError( - "kv_cache requires both `causal_condition=True` and `causal_block=True`; otherwise the condition " - "prefix depends on the changing target latent across denoising steps." + "kv_cache requires `causal_condition=True`; otherwise the condition prefix depends on the " + "changing target latent across denoising steps." ) if kv_cache is not None and len(kv_cache) != len(self.transformer_blocks): raise ValueError( @@ -520,76 +552,37 @@ def forward( prefix_len = int((~target_token_mask).sum()) is_decode = kv_cache is not None and len(kv_cache[0]) > 0 - use_flex = FLEX_ATTN_AVAILABLE and self.causal_block - + cache_write_slice = None if is_decode or kv_cache is None else slice(0, prefix_len) + segments = None if is_decode: joint_hidden_states = joint_hidden_states[:, prefix_len:] rotary_emb = rotary_emb[prefix_len:] modulation_mask = modulation_mask[prefix_len:] attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] - cache_write_slice = None - use_two_pass = False - elif use_flex: - cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None + elif FLEX_ATTN_AVAILABLE: attention_mask = build_qwenimage21_block_causal_mask( image_ids, joint_key_valid, batch_size, hidden_states.device ) - use_two_pass = False - elif self.causal_block: - use_two_pass = True - cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None - else: - cache_write_slice = slice(0, prefix_len) if kv_cache is not None else None - attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] - use_two_pass = False - - if use_two_pass: - prefix_hs = joint_hidden_states[:, :prefix_len] - target_hs = joint_hidden_states[:, prefix_len:] - prefix_rope = rotary_emb[:prefix_len] - target_rope = rotary_emb[prefix_len:] - prefix_mod_mask = modulation_mask[:prefix_len] if modulation_mask is not None else None - target_mod_mask = modulation_mask[prefix_len:] if modulation_mask is not None else None - for index_block, block in enumerate(self.transformer_blocks): - block_kv_cache = kv_cache[index_block] if kv_cache is not None else {} - prefix_cache = {} - prefix_hs = block( - hidden_states=prefix_hs, - modulation=modulation, - rotary_emb=prefix_rope, - attention_mask=None, - target_token_mask=prefix_mod_mask, - kv_cache=prefix_cache, - cache_write_slice=slice(0, prefix_len), - is_causal=True, - ) - block_kv_cache.update(prefix_cache) - target_hs = block( - hidden_states=target_hs, - modulation=modulation, - rotary_emb=target_rope, - attention_mask=None, - target_token_mask=target_mod_mask, - kv_cache=block_kv_cache, - cache_write_slice=None, - ) - - joint_hidden_states = torch.cat([prefix_hs, target_hs], dim=1) else: - for index_block, block in enumerate(self.transformer_blocks): - block_kv_cache = kv_cache[index_block] if kv_cache is not None else None - joint_hidden_states = gradient_checkpoint_forward( - block, - use_gradient_checkpointing, - use_gradient_checkpointing_offload, - hidden_states=joint_hidden_states, - modulation=modulation, - rotary_emb=rotary_emb, - attention_mask=attention_mask, - target_token_mask=modulation_mask, - kv_cache=block_kv_cache, - cache_write_slice=cache_write_slice, - ) + attention_mask = None + segments = _qwenimage21_prefix_segments(image_ids, prefix_len) + + for index_block, block in enumerate(self.transformer_blocks): + block_kv_cache = kv_cache[index_block] if kv_cache is not None else None + joint_hidden_states = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + hidden_states=joint_hidden_states, + modulation=modulation, + rotary_emb=rotary_emb, + attention_mask=attention_mask, + target_token_mask=modulation_mask, + kv_cache=block_kv_cache, + cache_write_slice=cache_write_slice, + segments=segments, + key_valid=joint_key_valid, + ) joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) return self.proj_out(joint_hidden_states) From 91508216c5766c37515941ef0be93cf16d3447c4 Mon Sep 17 00:00:00 2001 From: mi804 <1576993271@qq.com> Date: Sun, 20 Sep 2026 19:42:54 +0800 Subject: [PATCH 2/2] minor fix --- diffsynth/models/qwen_image_21_dit.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/diffsynth/models/qwen_image_21_dit.py b/diffsynth/models/qwen_image_21_dit.py index 301293f1e..5429e9fec 100644 --- a/diffsynth/models/qwen_image_21_dit.py +++ b/diffsynth/models/qwen_image_21_dit.py @@ -198,8 +198,6 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx): def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]: - # Maximal runs of equal image_ids, as (start, end, is_text). Derived once per forward rather than per block: - # tolist() is a device sync, and there is one processor call per layer. prefix_ids = image_ids[:prefix_len].tolist() segments = [] start = 0