From 4fc3cf8c6057ec6fbb1ada65012b14690d20b73f Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Fri, 18 Sep 2026 10:02:39 +0800 Subject: [PATCH] Enable CUDA Graph for DiT decode layers (#4) * init Signed-off-by: Isotr0py * fix Signed-off-by: Isotr0py * test(diffusion): add Qwen-Image-2.1 unit tests Signed-off-by: Isotr0py * feat(qwen_image_21): opt-in CUDA graph capture for KV-cache decode steps Decode steps (1..N) recompute only the target image's tokens at a fixed shape, but the per-request dict KV cache made addresses unstable, so the step could not be captured. Prefill now copies prefix K/V into persistent fixed-address buffers (keyed by branch/batch/prefix_len/target_tokens/ dtype/device/backend) and re-points the caller's cache dicts at them; decode replays a graph captured per key over the full 32-block forward. Eager fallback (logged) for non-CUDA, SP/ring/TP>1, HSDP, compiled or offload-hooked blocks, KV-cache quantization, dynamic LoRA wrappers, padded text masks, unknown keys, and capture failure. Off by default; opt in via OmniDiffusionConfig.enable_cuda_graph_decode or the QwenImage21Transformer2DModel constructor kwarg. E2E (GB200, T2I 1024x1024, 50 steps, CFG 4.0, CUDNN_ATTN): 171.2 -> 142.1 ms/step, 8.70s -> 7.24s total (1.20x), 44.9 dB PSNR vs eager output. * clean Signed-off-by: Isotr0py --------- Signed-off-by: Isotr0py --- vllm_omni/diffusion/data.py | 5 + .../models/qwen_image_21/decode_graph.py | 352 ++++++++++++++++++ .../qwen_image_21_transformer.py | 78 ++++ 3 files changed, 435 insertions(+) create mode 100644 vllm_omni/diffusion/models/qwen_image_21/decode_graph.py diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c30529520..3c44c7a8b 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -902,6 +902,11 @@ class OmniDiffusionConfig: # Compilation enforce_eager: bool = False + # Opt-in: capture fixed-shape KV-cache decode (denoising) steps into CUDA + # graphs. Currently implemented by Qwen-Image-2.1's transformer; unsupported + # configurations (SP/TP, ring, HSDP, dynamic LoRA, padded masks) log and + # fall back to eager decode. Other models ignore this flag. + enable_cuda_graph_decode: bool = False # Controls the generic compilation path used when a pipeline does not # provide its own setup_compile() implementation. diffusion_compile_granularity: str = "regional" diff --git a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py new file mode 100644 index 000000000..dcfdc2845 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA Graph capture for Qwen-Image-2.1 KV-cache decode steps. + +Decode (denoising steps 1..N) recomputes only the target image's tokens: the +shape is identical every step and the timestep-independent prefix K/V come from +the cross-step KV cache. That is the textbook CUDA-graph case, except the +default cache stores prefix K/V in freshly allocated tensors, so a captured +graph would bind addresses that the next request (or a step-mode cache merge) +no longer owns. + +This module keeps the prefix K/V in persistent, fixed-address buffers. Prefill +copies its K/V into them and re-points the caller's cache dict at the buffers; +decode then replays a captured graph that reads those buffers plus static input +buffers (target latents, timestep, RoPE frequencies). One graph is captured per +``(branch, batch, prefix_len, target_tokens, dtype, device, backend)`` key and +reused across requests: a new prefill with the same key rewrites the buffers in +place, which a captured graph reads correctly because it binds addresses, not +contents. + +Memory note: only the *prefix* K/V are persistent. The decode attention still +concatenates prefix K/V with the current step's target K/V inside the captured +region, and that transient tensor is served by the graph's memory pool, so the +persistent footprint is the same as the existing dict cache. + +Fallbacks (all logged, all eager): CUDA unavailable, 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. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import TYPE_CHECKING + +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import ( + QwenImage21Transformer2DModel, + ) + +logger = init_logger(__name__) + + +def _dynamic_lora_wrappers_present(module: torch.nn.Module) -> bool: + """True once ``DiffusionLoRAManager`` has wrapped layers under ``module``. + + A captured graph records the layers and branches that ran at capture time; + binding or rescaling an adapter afterwards changes the math without + changing the shapes the graph was keyed on, so graphs and dynamic LoRA + must not mix. Same reasoning as ``sensenova_u1.paged_decode``. + """ + try: + from vllm.lora.layers import BaseLayerWithLoRA + except ImportError: # pragma: no cover - depends on the wheel + return False + return any(isinstance(m, BaseLayerWithLoRA) for m in module.modules()) + + +class QwenImage21DecodeGraphEntry: + """Static state for one captured decode graph. + + Buffers are created outside inference mode so replay-side ``copy_`` is + legal regardless of the caller's grad context. ``graph`` is None until the + first decode step captures it; ``failed`` disables capture permanently for + this key after an exception. + """ + + def __init__( + self, + *, + branch: str, + batch_size: int, + prefix_len: int, + target_tokens: int, + in_channels: int, + dtype: torch.dtype, + device: torch.device, + kv_shapes: list[torch.Size], + kv_dtype: torch.dtype, + freqs: torch.Tensor, + ): + self.branch = branch + self.target_tokens = target_tokens + self.prefix_len = prefix_len + with torch.inference_mode(False): + self.hidden = torch.zeros(batch_size, target_tokens, in_channels, dtype=dtype, device=device) + self.timestep = torch.zeros(batch_size, dtype=dtype, device=device) + self.k = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes] + self.v = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes] + self.target_token_mask = torch.ones(target_tokens, dtype=torch.bool, device=device) + # RoPE frequencies are layout-dependent but step-independent. Clone + # into a normal (non-inference) tensor: the prefill runs under + # inference_mode, and the captured body takes views of this tensor + # outside it, which is forbidden for inference tensors. + self.freqs = freqs.clone() + # Decode needs no mask once padding is excluded (see the manager's + # prefill gate), so the captured body attends without metadata. + self.attn_metadata = None + # Per-block cache dicts with the same shape as the legacy protocol. + self.block_caches = [{branch: {"key": self.k[i], "value": self.v[i]}} for i in range(len(kv_shapes))] + self.graph: torch.cuda.CUDAGraph | None = None + self.output: torch.Tensor | None = None + self.failed = False + self.captures = 0 + + def capture(self, body) -> None: + """Warm up on a side stream, then capture one decode step. + + Warm-up runs the same computation it will record, allocating cuBLAS + workspaces and any lazy buffers outside the capture. Decode writes + nothing to the cache (prefix K/V are read-only), so warm-up is + idempotent. + """ + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side), torch.inference_mode(False), torch.no_grad(): + for _ in range(2): + body() + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + # The platform-wide pool is shared with every other captured path in + # the tree; entries only keep their *output* alive, whose tiny size + # makes reuse safe because the caller clones it before any other graph + # on the pool replays. + with torch.inference_mode(False), torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + self.output = body() + self.graph = graph + self.captures += 1 + + +class QwenImage21DecodeGraphManager: + """Owns the static KV buffers and captured graphs of one transformer.""" + + def __init__(self, model: QwenImage21Transformer2DModel, max_entries: int = 8): + self.model = model + self.max_entries = max_entries + self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry] = OrderedDict() + self._static_eligible: bool | None = None + + # ── eligibility ── + + def _check_static_eligibility(self) -> bool: + """One-time checks that cannot change after weights are loaded.""" + model = self.model + reason = None + param = next(model.parameters(), None) + if param is None or param.device.type != "cuda": + reason = f"model is not on CUDA (device={None if param is None else param.device})" + else: + parallel_config = getattr(model, "parallel_config", None) + if parallel_config is not None: + sp = getattr(parallel_config, "sequence_parallel_size", None) or 1 + if sp > 1 or getattr(parallel_config, "ring_degree", 1) > 1: + reason = "sequence/ring parallelism changes decode shapes per rank" + elif getattr(parallel_config, "use_hsdp", False): + reason = "HSDP wraps blocks with gather hooks outside the captured region" + if reason is None: + from vllm.distributed import get_tensor_model_parallel_world_size + + if get_tensor_model_parallel_world_size() > 1: + reason = "tensor parallelism runs collectives inside the captured region" + if reason is None: + for block in model.transformer_blocks: + if getattr(block.forward, "_torchdynamo_orig_callable", None) is not None: + reason = "transformer blocks are torch.compile'd" + break + if hasattr(block, "_hook_registry") or hasattr(block, "_omni_original_forward"): + reason = "transformer blocks carry offload/cache hooks" + break + if reason is None and model.transformer_blocks: + attn_layer = model.transformer_blocks[0].attn.attn + if getattr(attn_layer, "_kv_cache_dtype", None) is not None: + reason = "KV-cache quantization is enabled" + + if reason is not None: + logger.warning_once("Qwen-Image-2.1 CUDA graph decode disabled: %s. Falling back to eager decode.", reason) + return False + return True + + def eligible(self) -> bool: + if self._static_eligible is None: + self._static_eligible = self._check_static_eligibility() + return self._static_eligible + + def _backend_name(self) -> str: + backend = self.model.transformer_blocks[0].attn.attn.attn_backend + return backend.get_name() if backend is not None else "custom" + + # ── prefill registration ── + + def register_prefill( + self, + *, + kv_cache: list[dict[str, dict[str, torch.Tensor]]], + cache_branch: str, + prefix_len: int, + target_tokens: int, + target_freqs: torch.Tensor, + joint_key_valid: torch.Tensor | None, + dtype: torch.dtype, + ) -> None: + """Move a freshly prefilled prefix K/V into static buffers. + + The caller's ``kv_cache`` dicts are re-pointed at the static buffers, so + the existing decode path reads fixed addresses and a captured graph can + be replayed for later steps. Prefill itself always runs eagerly. + """ + if not self.eligible() or torch.compiler.is_compiling(): + return + if _dynamic_lora_wrappers_present(self.model): + if self.entries: + logger.warning_once( + "Qwen-Image-2.1 CUDA graph decode: LoRA wrappers appeared; " + "dropping captured graphs and falling back to eager decode." + ) + self.entries.clear() + return + if joint_key_valid is not None and not bool(joint_key_valid.all()): + # The masked decode path (varlen-unpack or 4D mask) branches on mask + # contents and cannot be captured; this request stays eager. + logger.warning_once( + "Qwen-Image-2.1 CUDA graph decode: padded text mask present; this request falls back to eager decode." + ) + return + + first = kv_cache[0].get(cache_branch) + if first is None or "key" not in first: + return + batch_size = first["key"].shape[0] + key = ( + cache_branch, + batch_size, + prefix_len, + target_tokens, + dtype, + first["key"].device.index, + self._backend_name(), + ) + + entry = self.entries.get(key) + if entry is None: + while len(self.entries) >= self.max_entries: + evicted_key, _ = self.entries.popitem(last=False) + logger.debug("Evicting decode graph entry %s (max_entries=%d)", evicted_key, self.max_entries) + entry = QwenImage21DecodeGraphEntry( + branch=cache_branch, + batch_size=batch_size, + prefix_len=prefix_len, + target_tokens=target_tokens, + in_channels=self.model.in_channels, + dtype=dtype, + device=first["key"].device, + kv_shapes=[block_cache[cache_branch]["key"].shape for block_cache in kv_cache], + kv_dtype=first["key"].dtype, + freqs=target_freqs, + ) + self.entries[key] = entry + else: + self.entries.move_to_end(key) + + for i, block_cache in enumerate(kv_cache): + parts = block_cache[cache_branch] + entry.k[i].copy_(parts["key"]) + entry.v[i].copy_(parts["value"]) + parts["key"] = entry.k[i] + parts["value"] = entry.v[i] + logger.debug( + "Registered decode graph entry key=%s (entries=%d, captured=%d)", + key, + len(self.entries), + sum(1 for e in self.entries.values() if e.graph is not None), + ) + + # ── decode replay ── + + def try_decode( + self, + *, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + kv_cache: list[dict[str, dict[str, torch.Tensor]]], + cache_branch: str, + target_tokens: int, + ) -> torch.Tensor | None: + """Replay the captured decode step, or return None to fall back to eager.""" + if not self.eligible() or torch.compiler.is_compiling(): + return None + first = kv_cache[0][cache_branch] + cached_key = first["key"] + if cached_key.device.type != "cuda": + # Registration is device-agnostic (unit tests exercise the buffer + # machinery on CPU), but capture/replay is CUDA-only: an empty + # graph over CPU ops would replay as a stale no-op. + return None + key = ( + cache_branch, + cached_key.shape[0], + cached_key.shape[1], + target_tokens, + hidden_states.dtype, + cached_key.device.index, + self._backend_name(), + ) + entry = self.entries.get(key) + if entry is None or entry.failed: + return None + self.entries.move_to_end(key) + + entry.hidden.copy_(hidden_states[:, -target_tokens:]) + entry.timestep.copy_(timestep) + # The caller's cache dicts normally already reference the entry's + # buffers (prefill re-pointed them). A step-mode batched decode instead + # hands over freshly merged tensors; copy those into the static buffers. + for i, block_cache in enumerate(kv_cache): + parts = block_cache[cache_branch] + if parts["key"] is not entry.k[i]: + entry.k[i].copy_(parts["key"]) + entry.v[i].copy_(parts["value"]) + parts["key"] = entry.k[i] + parts["value"] = entry.v[i] + + if entry.graph is None: + try: + entry.capture(lambda: self.model._decode_graph_forward(entry)) + except Exception as exc: + entry.failed = True + entry.graph = None + logger.warning( + "Qwen-Image-2.1 CUDA graph capture failed for key=%s: %s. " + "Falling back to eager decode for this shape.", + key, + exc, + ) + return None + logger.info("Captured Qwen-Image-2.1 decode graph for key=%s", key) + entry.graph.replay() + assert entry.output is not None + # Clone before returning: the output lives in the shared graph pool and + # the next replay (of this or another graph on the pool) overwrites it. + return entry.output.clone() + + def clear(self) -> None: + """Drop all entries, freeing buffers and captured graphs.""" + self.entries.clear() 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 57b98aad7..808926b50 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 @@ -32,6 +32,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelOutput, ) from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.models.qwen_image_21.decode_graph import QwenImage21DecodeGraphManager if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -716,6 +717,8 @@ class QwenImage21Transformer2DModel(CachedTransformer): causal_block: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = "transformer", + enable_cuda_graph_decode: bool | None = None, + cuda_graph_max_decode_graphs: int = 8, ): super().__init__() self.parallel_config = od_config.parallel_config @@ -786,6 +789,17 @@ 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) + # Opt-in CUDA Graph capture of the fixed-shape KV-cache decode steps. + # Default off: ``None`` defers to ``od_config.enable_cuda_graph_decode``. + if enable_cuda_graph_decode is None: + enable_cuda_graph_decode = bool(getattr(od_config, "enable_cuda_graph_decode", False)) + self.enable_cuda_graph_decode = enable_cuda_graph_decode + self._decode_graph_manager = ( + QwenImage21DecodeGraphManager(self, max_entries=cuda_graph_max_decode_graphs) + if enable_cuda_graph_decode + else None + ) + @staticmethod def build_token_metadata( image_pad_mask: torch.Tensor, img_shapes: list[tuple[int, int, int]] @@ -951,6 +965,19 @@ class QwenImage21Transformer2DModel(CachedTransformer): is_decode = kv_cache is not None and len(kv_cache) > 0 and "key" in kv_cache[0].get(cache_branch, {}) + if is_decode and self._decode_graph_manager is not None: + graph_output = self._decode_graph_manager.try_decode( + hidden_states=hidden_states, + timestep=timestep, + kv_cache=kv_cache, + cache_branch=cache_branch, + target_tokens=math.prod(layout[-1]), + ) + if graph_output is not None: + if not return_dict: + return (graph_output,) + return Transformer2DModelOutput(sample=graph_output) + if is_decode: txt_hidden_states = None else: @@ -1048,6 +1075,24 @@ class QwenImage21Transformer2DModel(CachedTransformer): cache_write_len=cache_write_len, ) + if ( + not is_decode + and cache_write_len is not None + and kv_cache is not None + and self._decode_graph_manager is not None + ): + # Move the freshly written prefix K/V into fixed-address buffers so + # decode steps can replay a captured CUDA graph against them. + self._decode_graph_manager.register_prefill( + kv_cache=kv_cache, + cache_branch=cache_branch, + prefix_len=cache_write_len, + target_tokens=target_freqs.shape[0], + target_freqs=target_freqs, + joint_key_valid=joint_key_valid, + dtype=hidden_states.dtype, + ) + if sp_active and not is_decode: # Only target tokens feed proj_out under SP: the prefix is replicated on every # rank and gathering it would duplicate it. Output is the (gathered) target image. @@ -1063,6 +1108,39 @@ class QwenImage21Transformer2DModel(CachedTransformer): return Transformer2DModelOutput(sample=output) + def _decode_graph_forward(self, entry) -> torch.Tensor: + """The exact decode-step computation, captured into a CUDA graph. + + Mirrors the eager decode path of ``forward``: target tokens only, + prefix K/V read from the entry's static buffers, no attention metadata + (the registration gate excludes padded masks), and an all-ones + modulation mask. All inputs live in the entry's fixed-address buffers. + """ + hidden_states = self.img_in(entry.hidden) + timestep = torch.cat([entry.timestep, entry.timestep.new_zeros(1)], dim=0) + temb = self.time_text_embed(timestep, hidden_states) + modulation = self.modulation(temb) + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states=hidden_states, + modulation=modulation, + freqs=entry.freqs, + target_token_mask=entry.target_token_mask, + attn_metadata=entry.attn_metadata, + sp_prefix_len=0, + sp_decode=False, + kv_cache=entry.block_caches[index_block], + cache_branch=entry.branch, + cache_write_len=None, + ) + hidden_states = self.norm_out(hidden_states, temb, entry.target_token_mask) + return self.proj_out(hidden_states) + + def release_captured_graphs(self) -> None: + """Drop captured decode graphs and their static buffers (e.g. sleep mode).""" + if self._decode_graph_manager is not None: + self._decode_graph_manager.clear() + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id)