[diffusion] Enable breakable CUDA graph (BCG) for diffusion DiTs

python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py

@@ -58,26 +58,30 @@ def get_freqs_cis(self, batch, device, rotary_emb, dtype):
         return cos, sin
 
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
-        return {
+        kwargs = {
             "prior_token_id": batch.prior_token_id,
             "prior_token_drop": batch.prior_token_drop_cond,
             "crop_coords": batch.crop_coords,
             "target_size": batch.target_size,
-            "kv_caches": batch.kv_caches,
-            "kv_caches_mode": "read",
             "freqs_cis": self.get_freqs_cis(batch, device, rotary_emb, dtype),
         }
+        if getattr(batch, "prior_token_image_ids", None) is not None:
+            kwargs["kv_caches"] = batch.kv_caches
+            kwargs["kv_caches_mode"] = "read"
+        return kwargs
 
     def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
-        return {
+        kwargs = {
             "prior_token_id": batch.prior_token_id,
             "prior_token_drop": batch.prior_token_drop_uncond,
             "crop_coords": batch.crop_coords,
             "target_size": batch.target_size,
-            "kv_caches": batch.kv_caches,
-            "kv_caches_mode": "skip",
             "freqs_cis": self.get_freqs_cis(batch, device, rotary_emb, dtype),
         }
+        if getattr(batch, "prior_token_image_ids", None) is not None:
+            kwargs["kv_caches"] = batch.kv_caches
+            kwargs["kv_caches_mode"] = "skip"
+        return kwargs
 
     def get_decode_scale_and_shift(self, device, dtype, vae):
         latents_mean = (

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/__init__.py

@@ -0,0 +1 @@
+"""Diffusion breakable CUDA graph runtime helpers."""

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/__init__.py

@@ -0,0 +1 @@
+"""Model-specific prompt padders for diffusion breakable CUDA graph."""

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/ideogram.py

@@ -0,0 +1,131 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0
+# ==============================================================================
+"""Ideogram-4 breakable CUDA graph (BCG) prompt padding."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
+    prompt_padding as bcg_utils,
+)
+from sglang.multimodal_gen.runtime.layers.attention import DynamicVarlenMaskMeta
+
+_SEQUENCE_PADDING_INDICATOR = -1
+_OUTPUT_IMAGE_INDICATOR = 2
+_LLM_TOKEN_INDICATOR = 3
+_DYNAMIC_MASK_META_ATTR = "_sglang_bcg_ideogram_attn_mask_meta"
+
+
+def is_ideogram_transformer(current_model: Any, call_kwargs: dict) -> bool:
+    return (
+        bcg_utils.transformer_class_name_matches(current_model, "ideogram")
+        and "llm_features" in call_kwargs
+        and "x" in call_kwargs
+        and "indicator" in call_kwargs
+        and "position_ids" in call_kwargs
+    )
+
+
+def _unwrap_model(current_model: Any) -> Any:
+    for attr in ("module", "_orig_mod"):
+        wrapped = getattr(current_model, attr, None)
+        if wrapped is not None:
+            current_model = wrapped
+    return current_model
+
+
+def _dynamic_mask_meta(current_model: Any) -> DynamicVarlenMaskMeta:
+    model = _unwrap_model(current_model)
+    meta = getattr(model, _DYNAMIC_MASK_META_ATTR, None)
+    if not isinstance(meta, DynamicVarlenMaskMeta):
+        meta = DynamicVarlenMaskMeta()
+        setattr(model, _DYNAMIC_MASK_META_ATTR, meta)
+    return meta
+
+
+def _first_indicator(call_kwargs: dict) -> torch.Tensor | None:
+    indicator = bcg_utils.first_tensor(call_kwargs.get("indicator"))
+    if not torch.is_tensor(indicator) or indicator.dim() < 2:
+        return None
+    return indicator
+
+
+def _text_and_image_lengths(indicator: torch.Tensor) -> tuple[int, int] | None:
+    row = indicator[0]
+    if not torch.any(row == _LLM_TOKEN_INDICATOR):
+        return None
+    image_positions = (row == _OUTPUT_IMAGE_INDICATOR).nonzero(as_tuple=False)
+    if image_positions.numel() == 0:
+        return None
+    text_seq = int(image_positions[0].item())
+    if text_seq <= 0:
+        return None
+    image_seq = int(row.numel()) - text_seq
+    if image_seq <= 0:
+        return None
+    return text_seq, image_seq
+
+
+def _pad_total_dim(obj: Any, *, source: int, target: int, value: float = 0) -> Any:
+    return bcg_utils.pad_nested_dim(
+        obj, dim=1, source=source, target=target, value=value
+    )
+
+
+def pad_ideogram_prompt_kwargs(
+    call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
+) -> dict:
+    indicator = _first_indicator(call_kwargs)
+    if indicator is None:
+        return call_kwargs
+
+    lengths = _text_and_image_lengths(indicator)
+    if lengths is None:
+        return call_kwargs
+    text_seq, image_seq = lengths
+
+    bucket = bcg_utils.select_text_bucket(text_seq, buckets)
+    if bucket is None:
+        return call_kwargs
+
+    source_total = text_seq + image_seq
+    target_total = bucket + image_seq
+    out = dict(call_kwargs)
+
+    if source_total < target_total:
+        for key in ("llm_features", "x"):
+            if key in out and out[key] is not None:
+                out[key] = _pad_total_dim(
+                    out[key], source=source_total, target=target_total
+                )
+        if out.get("position_ids") is not None:
+            out["position_ids"] = _pad_total_dim(
+                out["position_ids"], source=source_total, target=target_total
+            )
+        if out.get("segment_ids") is not None:
+            out["segment_ids"] = _pad_total_dim(
+                out["segment_ids"],
+                source=source_total,
+                target=target_total,
+                value=_SEQUENCE_PADDING_INDICATOR,
+            )
+        if out.get("indicator") is not None:
+            out["indicator"] = _pad_total_dim(
+                out["indicator"], source=source_total, target=target_total
+            )
+        if out.get("attn_mask") is not None:
+            out["attn_mask"] = _pad_total_dim(
+                out["attn_mask"], source=source_total, target=target_total
+            )
+
+    if out.get("attn_mask") is not None:
+        out["attn_mask_meta"] = _dynamic_mask_meta(current_model)
+
+    return out
+
+
+bcg_utils.register_prompt_padder(is_ideogram_transformer, pad_ideogram_prompt_kwargs)

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/qwen_image.py

@@ -0,0 +1,102 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Qwen-Image breakable CUDA graph (BCG) prompt padding.
+
+Qwen-Image / Qwen-Image-Edit carry text length on dim 1 of
+``encoder_hidden_states`` and a separate ``freqs_cis`` text-rope cache plus
+``txt_seq_lens``; they may not pass an explicit prompt mask, so this padder
+synthesizes one. Registered with the base denoising stage's padder registry.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
+    prompt_padding as bcg_utils,
+)
+
+
+def is_qwen_transformer(current_model: Any, call_kwargs: dict) -> bool:
+    return (
+        bcg_utils.transformer_class_name_matches(current_model, "qwen")
+        and "txt_seq_lens" in call_kwargs
+        and "freqs_cis" in call_kwargs
+    )
+
+
+def pad_qwen_prompt_kwargs(
+    call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
+) -> dict:
+    ehs = call_kwargs.get("encoder_hidden_states")
+    ehs_tensor = bcg_utils.first_tensor(ehs)
+    if not torch.is_tensor(ehs_tensor) or ehs_tensor.dim() < 2:
+        return call_kwargs
+
+    seq = ehs_tensor.shape[1]
+    bucket = bcg_utils.select_text_bucket(seq, buckets)
+    if bucket is None:
+        return call_kwargs
+
+    out = dict(call_kwargs)
+    if seq < bucket:
+        out["encoder_hidden_states"] = bcg_utils.pad_nested_dim(
+            ehs, dim=1, source=seq, target=bucket
+        )
+        if (
+            "encoder_hidden_states_2" in out
+            and out["encoder_hidden_states_2"] is not None
+        ):
+            out["encoder_hidden_states_2"] = bcg_utils.pad_nested_dim(
+                out["encoder_hidden_states_2"], dim=1, source=seq, target=bucket
+            )
+
+    mask = out.get("encoder_hidden_states_mask")
+    if mask is None:
+        mask = torch.ones(
+            ehs_tensor.shape[:2],
+            device=ehs_tensor.device,
+            dtype=torch.bool,
+        )
+    if mask is not None:
+        out["encoder_hidden_states_mask"] = bcg_utils.pad_nested_dim(
+            mask, dim=1, source=seq, target=bucket
+        )
+
+    if "encoder_attention_mask" in out and out["encoder_attention_mask"] is not None:
+        out["encoder_attention_mask"] = bcg_utils.pad_nested_dim(
+            out["encoder_attention_mask"], dim=1, source=seq, target=bucket
+        )
+
+    freqs_cis = out.get("freqs_cis")
+    if isinstance(freqs_cis, tuple) and len(freqs_cis) == 2:
+        img_cache, txt_cache = freqs_cis
+        txt_cache = bcg_utils.pad_nested_dim(
+            txt_cache, dim=0, source=seq, target=bucket
+        )
+        out["freqs_cis"] = (img_cache, txt_cache)
+    elif isinstance(freqs_cis, list) and len(freqs_cis) == 2:
+        img_cache, txt_cache = freqs_cis
+        txt_cache = bcg_utils.pad_nested_dim(
+            txt_cache, dim=0, source=seq, target=bucket
+        )
+        out["freqs_cis"] = [img_cache, txt_cache]
+
+    out["txt_seq_lens"] = bcg_utils.bucket_txt_seq_lens(out.get("txt_seq_lens"), bucket)
+    return out
+
+
+bcg_utils.register_prompt_padder(is_qwen_transformer, pad_qwen_prompt_kwargs)