Fix Qwen Image 2.1 compile recompiles

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,
                 )

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:
     """