Enable CUDA Graph for DiT decode layers (#4)

* init

Signed-off-by: Isotr0py <mozf@inferact.ai>

* fix

Signed-off-by: Isotr0py <mozf@inferact.ai>

* test(diffusion): add Qwen-Image-2.1 unit tests

Signed-off-by: Isotr0py <mozf@inferact.ai>

* 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 <mozf@inferact.ai>

---------

Signed-off-by: Isotr0py <mozf@inferact.ai>

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"

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

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)