diff --git a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py index 018e8f35e..e3707516f 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py +++ b/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() diff --git a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py index 212df1a04..79ad54f1e 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py +++ b/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 diff --git a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py index b84c2912d..b351e052b 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py +++ b/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,