diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 1a9ef250c519..11277f389b45 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -330,7 +330,7 @@ def _qwenimage21_prepare_qkv( rotary_emb: torch.Tensor | None, layer_cache: QwenImage21KVLayerCache | None, kv_cache_mode: str | None, - cache_write_slice: slice | None, + cache_write_len: int | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: """Shared QKV projection, norm, RoPE and KV-cache bookkeeping for both processors.""" query = attn.to_q(hidden_states) @@ -352,13 +352,13 @@ def _qwenimage21_prepare_qkv( key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) if layer_cache is not None: - if kv_cache_mode == "extract" and cache_write_slice is not None: + if kv_cache_mode == "extract" and cache_write_len is not None: # `clone()`, not `contiguous()`: at batch size 1 the prefix slice already counts as contiguous # (size-1 dims are ignored), so `contiguous()` returns the same view and the cache would pin the # whole prefill K/V for every step of the denoising loop. layer_cache.store( - key[:, cache_write_slice].clone(), - value[:, cache_write_slice].clone(), + key[:, :cache_write_len].clone(), + value[:, :cache_write_len].clone(), ) elif kv_cache_mode == "cached": cached_k, cached_v = layer_cache.get() @@ -400,12 +400,12 @@ def __call__( rotary_emb: torch.Tensor | None = None, layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, - cache_write_slice: slice | None = None, + cache_write_len: int | None = None, segments: list[tuple[int, int, bool]] | None = None, key_valid: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, seq_len_q = _qwenimage21_prepare_qkv( - attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice + attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_len ) seq_len_kv = key.shape[1] @@ -482,12 +482,12 @@ def __call__( rotary_emb: torch.Tensor | None = None, layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, - cache_write_slice: slice | None = None, + cache_write_len: int | None = None, segments: list[tuple[int, int, bool]] | None = None, key_valid: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, seq_len_q = _qwenimage21_prepare_qkv( - attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice + attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_len ) if segments is None: @@ -624,7 +624,7 @@ def forward( target_token_mask: torch.Tensor | None = None, layer_cache: QwenImage21KVLayerCache | None = None, kv_cache_mode: str | None = None, - cache_write_slice: slice | None = None, + cache_write_len: int | None = None, segments: list[tuple[int, int, bool]] | None = None, key_valid: torch.Tensor | None = None, ) -> torch.Tensor: @@ -637,7 +637,7 @@ def forward( rotary_emb=rotary_emb, layer_cache=layer_cache, kv_cache_mode=kv_cache_mode, - cache_write_slice=cache_write_slice, + cache_write_len=cache_write_len, segments=segments, key_valid=key_valid, ) @@ -974,7 +974,7 @@ def forward( 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 + cache_write_len = None block_segments, block_key_valid = None, None else: # prefill: the whole joint sequence. The block-causal structure goes down in whichever form the @@ -992,7 +992,7 @@ def forward( if all(isinstance(processor, QwenImage21FlexAttnProcessor) for processor in processors) else _qwenimage21_prefix_segments(image_ids, prefix_len) ) - cache_write_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None + cache_write_len = prefix_len if kv_cache_mode == "extract" else None block_key_valid = joint_key_valid for index_block, block in enumerate(self.transformer_blocks): @@ -1007,7 +1007,7 @@ def forward( modulation_mask, layer_cache, kv_cache_mode, - cache_write_slice, + cache_write_len, block_segments, block_key_valid, ) @@ -1020,7 +1020,7 @@ def forward( target_token_mask=modulation_mask, layer_cache=layer_cache, kv_cache_mode=kv_cache_mode, - cache_write_slice=cache_write_slice, + cache_write_len=cache_write_len, segments=block_segments, key_valid=block_key_valid, ) diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index e178b954bd1d..b212ddc9d696 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -98,6 +98,36 @@ def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str "img_mask": img_mask, } + def get_generation_inputs(self, index: int, device=torch_device) -> dict[str, torch.Tensor]: + """ + Inputs for the `index`-th generation of a multi-size session: prompt length, condition-image size and target + size all vary, so the prefix length (text + condition-image tokens) differs from one generation to the next. + """ + text_len = 4 + (index % 4) + condition_tokens = 4 * (1 + index // 4) + condition_hw = {4: (2, 2), 8: (2, 4), 12: (3, 4), 16: (4, 4)}[condition_tokens] + target_hw = [(2, 2), (2, 4), (4, 2), (4, 4)][index % 4] + target_tokens = target_hw[0] * target_hw[1] + + # The vision-language sequence is the text followed by one slot per 2x2 group of condition-image latents; the + # target image's slots are appended after it, so `img_mask` is longer than `encoder_hidden_states`. + vlm_len = text_len + condition_tokens // 4 + + hidden_states = randn_tensor((1, condition_tokens + target_tokens, 4), generator=self.generator, device=device) + encoder_hidden_states = randn_tensor((1, vlm_len, 8), generator=self.generator, device=device) + encoder_hidden_states_mask = torch.ones((1, vlm_len), device=device, dtype=torch.long) + img_mask = torch.zeros((1, vlm_len + target_tokens // 4), device=device, dtype=torch.bool) + img_mask[:, text_len:] = True + + return { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "encoder_hidden_states_mask": encoder_hidden_states_mask, + "timestep": torch.tensor([1.0], device=device), + "img_shapes": [[(1, *condition_hw), (1, *target_hw)]], + "img_mask": img_mask, + } + class TestQwenImage21Transformer(QwenImage21TransformerTesterConfig, ModelTesterMixin): @pytest.mark.skip( @@ -218,6 +248,47 @@ def test_attention_backend_applies_to_decode_only(self, processor_name): torch.testing.assert_close(outputs[0], outputs[1]) + @pytest.mark.parametrize("processor_name", ["QwenImage21FlexAttnProcessor", "QwenImage21AttnProcessor"]) + def test_compiled_blocks_do_not_recompile_per_generation(self, processor_name): + """ + Regression test for https://github.com/huggingface/diffusers/issues/15027. + + `prefix_len` (text + condition-image tokens) changes with every prompt length and condition-image size. + Values Dynamo guards by constant — the old `cache_write_slice.stop` — recompiled the repeated blocks for + every distinct prefix length, so a session of different-sized generations exhausted the recompile limit + (8) and `fullgraph=True` hard-failed. Plain `int` arguments become dynamic after their second distinct + value, so the compiled graph count must plateau instead of growing per generation. + """ + import diffusers.models.transformers.transformer_qwenimage21 as transformer_module + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device).eval() + model.set_attn_processor(getattr(transformer_module, processor_name)()) + + # `backend="eager"`: the guard behavior under test is backend- and device-independent, and it keeps the + # flex path's uncompiled CPU fallback (a dense fp32 score matrix) out of the picture. + model.compile_repeated_blocks(fullgraph=True, backend="eager") + + torch._dynamo.reset() + from torch._dynamo.utils import counters as dynamo_counters + + dynamo_counters.clear() + + graph_counts = [] + for index in range(12): + inputs = self.get_generation_inputs(index) + kv_cache = QwenImage21KVCache(init_dict["num_layers"]) + with torch.no_grad(): + model(**inputs, kv_cache=kv_cache, kv_cache_mode="extract", return_dict=False) + model(**inputs, kv_cache=kv_cache, kv_cache_mode="cached", return_dict=False) + graph_counts.append(dynamo_counters["stats"]["unique_graphs"]) + + # Before the fix this loop raises `torch._dynamo.exc.FailOnRecompileLimitHit` (recompile limit 8) part-way + # through: every new prefix length added a graph. The first two generations legitimately specialize + # (static shapes, then dynamic shapes); afterwards the count must plateau. + assert graph_counts[-1] - graph_counts[1] <= 4 + class TestQwenImage21BlockCausalMask: """