Address review: tuple image input, NPU index_select rope, per-key decode graph ownership

- Normalize tuple image inputs in pre_process_func (empty tuple -> T2I path)
- Use torch.index_select in QwenImage21Rope for NPU compatibility (#6320 style)
- Refuse concurrent aliasing of shared decode-graph prefix KV entries: track a
  single live owner per key, non-owners fall back to eager decode

vllm_omni/diffusion/models/qwen_image_21/decode_graph.py

@@ -14,6 +14,12 @@ fixed-address graph buffers when the active prefix tensors change. Graphs
 are keyed by branch, batch, prefix length, complete image layout,
 dtype, device and backend, so equal token counts with different RoPE layouts stay separate.
 
+Each key admits a single live owner, tracked by a weak reference to the
+registering request's prefix K tensor. A second in-flight request whose cache
+aliases the same key is refused graph decode and stays eager until the
+owner's cache is released, so shared static buffers are never overwritten
+under a concurrent owner.
+
 Memory note: graph entries own an additional copy of the prefix K/V. Decode
 attention concatenates that prefix with target K/V inside the captured region;
 transient tensors use the graph memory pool.
@@ -22,8 +28,8 @@ Fallbacks (all logged, all eager): CUDA unavailable, model-level CPU offload
 without persistent weight staging, sequence/ring/tensor parallelism,
 HSDP/offload hooks, torch.compile'd blocks, KV-cache quantization, dynamic
 LoRA wrappers, padded text masks (the masked attention path branches on
-mask contents, which cannot be captured), unknown graph key at decode, and
-capture failure.
+mask contents, which cannot be captured), unknown graph key at decode, a
+second live request aliasing an owned graph key, and capture failure.
 """
 
 from __future__ import annotations
@@ -160,6 +166,10 @@ class QwenImage21DecodeGraphManager:
         self.max_entries = max_entries
         self.model_level_offload = model_level_offload
         self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry | None] = OrderedDict()
+        # Key -> weak reference to the prefix K tensor of the request that owns
+        # the entry's static buffers. Dead references mean the owning request
+        # finished and its cache was released, so the key is free to adopt.
+        self._owners: dict[tuple, weakref.ReferenceType[torch.Tensor]] = {}
         self._static_eligible: bool | None = None
 
     # ── eligibility ──
@@ -251,7 +261,12 @@ class QwenImage21DecodeGraphManager:
         joint_key_valid: torch.Tensor | None,
         dtype: torch.dtype,
     ) -> None:
-        """Prepare graph buffers without changing the request's prefix K/V."""
+        """Prepare graph buffers without changing the request's prefix K/V.
+
+        The first in-flight request to register a key owns its static buffers;
+        registrations for the same key from another live request are refused
+        so their decode stays eager until the owner completes.
+        """
         if not self.eligible() or torch.compiler.is_compiling():
             return
         if _dynamic_lora_wrappers_present(self.model):
@@ -297,6 +312,7 @@ class QwenImage21DecodeGraphManager:
         if key not in self.entries:
             while len(self.entries) >= self.max_entries:
                 evicted_key, evicted_entry = self.entries.popitem(last=False)
+                self._owners.pop(evicted_key, None)
                 del evicted_entry
                 logger.debug("Evicting decode graph entry %s (max_entries=%d)", evicted_key, self.max_entries)
             allocation_error = None
@@ -330,6 +346,20 @@ class QwenImage21DecodeGraphManager:
 
         entry = self.entries[key]
         if entry is not None:
+            owner = self._owners.get(key)
+            owner_prefix = owner() if owner is not None else None
+            if owner_prefix is not None and owner_prefix is not first["key"]:
+                # Another in-flight request still owns this entry's static
+                # buffers; refreshing them would overwrite that request's
+                # prefix mid-generation. This request decodes eagerly until
+                # the owner's cache is released and the weak reference dies.
+                logger.warning_once(
+                    "Qwen-Image-2.1 CUDA graph decode: entry key=%s is owned by another in-flight "
+                    "request; this request falls back to eager decode.",
+                    key,
+                )
+                return
+            self._owners[key] = weakref.ref(first["key"])
             entry.prefix_sources = []
 
         logger.debug(
@@ -383,6 +413,14 @@ class QwenImage21DecodeGraphManager:
             return None
         self.entries.move_to_end(key)
 
+        owner = self._owners.get(key)
+        owner_prefix = owner() if owner is not None else None
+        if owner_prefix is not None and owner_prefix is not cached_key:
+            # The entry is owned by a different in-flight request (see
+            # register_prefill); this request stays on the eager decode path
+            # until the owner's cache is released.
+            return None
+
         entry.hidden.copy_(hidden_states[:, -entry.target_tokens :])
         entry.timestep.copy_(timestep)
         # Graph buffers are scratch space; request caches must never alias them.
@@ -415,3 +453,4 @@ class QwenImage21DecodeGraphManager:
     def clear(self) -> None:
         """Drop all entries, freeing buffers and captured graphs."""
         self.entries.clear()
+        self._owners.clear()

vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py

@@ -196,6 +196,8 @@ def get_qwen_image_21_pre_process_func(
         if "additional_information" not in prompt:
             prompt["additional_information"] = {}
 
+        if isinstance(raw_image, tuple):
+            raw_image = list(raw_image)
         if raw_image is None or (isinstance(raw_image, list) and len(raw_image) == 0):
             request.batch_compatibility_key = ("qwen_image_21", ())
             request.prompt = prompt

vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py

@@ -385,7 +385,16 @@ class QwenImage21Rope(nn.Module):
         height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device=device)
         width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device=device)
 
-        return torch.cat([self.freqs[0][frame_index], self.freqs[1][height_index], self.freqs[2][width_index]], dim=-1)
+        height_index = torch.where(height_index < 0, height_index + self.freqs[1].shape[0], height_index)
+        width_index = torch.where(width_index < 0, width_index + self.freqs[2].shape[0], width_index)
+        return torch.cat(
+            [
+                self.freqs[0].index_select(0, frame_index),
+                self.freqs[1].index_select(0, height_index),
+                self.freqs[2].index_select(0, width_index),
+            ],
+            dim=-1,
+        )
 
 
 def _enable_pattern_ignored_layers(quant_config: QuantizationConfig | None) -> QuantizationConfig | None:
@@ -898,12 +907,14 @@ class QwenImage21Transformer2DModel(CachedTransformer):
         # Module boundary where _sp_plan shards the target tokens + their RoPE freqs.
         self.sequence_prepare = QwenImage21SequencePrepare(self.img_in, self.pos_embed)
 
-        # Use decode graphs unless the caller requests eager execution.
-        # Model-level (sequential) offload registers its swap hook on this
-        # top-level module rather than on individual blocks; graphs stay
-        # eligible only when that hook keeps weights on persistent staging
-        # storage (see QwenImage21DecodeGraphManager._offload_reason).
-        self.enable_cuda_graph_decode = not od_config.enforce_eager
+        # Use decode graphs unless the caller disables them or requests eager
+        # execution. Model-level (sequential) offload registers its swap hook
+        # on this top-level module rather than on individual blocks; graphs
+        # stay eligible only when that hook keeps weights on persistent
+        # staging storage (see QwenImage21DecodeGraphManager._offload_reason).
+        self.enable_cuda_graph_decode = (
+            getattr(od_config, "enable_cuda_graph_decode", True) and not od_config.enforce_eager
+        )
         self._decode_graph_manager = (
             QwenImage21DecodeGraphManager(
                 self,