chore: import upstream snapshot with attribution

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)

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

@@ -0,0 +1,169 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0
+# ==============================================================================
+"""Z-Image 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,
+)
+
+
+def is_zimage_transformer(current_model: Any, call_kwargs: dict) -> bool:
+    return (
+        bcg_utils.transformer_class_name_matches(current_model, "zimage")
+        and "encoder_hidden_states" in call_kwargs
+        and "freqs_cis" in call_kwargs
+    )
+
+
+def _first_caption_tensor(encoder_hidden_states: Any) -> torch.Tensor | None:
+    tensor = bcg_utils.first_tensor(encoder_hidden_states)
+    if not torch.is_tensor(tensor):
+        return None
+    if tensor.dim() == 2:
+        return tensor
+    if tensor.dim() == 3:
+        return tensor[0]
+    return None
+
+
+def _caption_seq_len(tensor: torch.Tensor) -> int:
+    if tensor.dim() == 2:
+        return int(tensor.shape[0])
+    if tensor.dim() == 3:
+        return int(tensor.shape[1])
+    raise ValueError("Z-Image caption tensor must have rank 2 or 3")
+
+
+def _pad_caption(obj: Any, *, target: int) -> Any:
+    if torch.is_tensor(obj):
+        if obj.dim() == 2:
+            return bcg_utils.pad_tensor_dim(obj, 0, target)
+        if obj.dim() == 3:
+            return bcg_utils.pad_tensor_dim(obj, 1, target)
+        return obj
+    if isinstance(obj, list):
+        return [_pad_caption(item, target=target) for item in obj]
+    if isinstance(obj, tuple):
+        return tuple(_pad_caption(item, target=target) for item in obj)
+    return obj
+
+
+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 _build_caption_freqs(current_model: Any, *, target: int, device: torch.device):
+    rotary_emb = getattr(_unwrap_model(current_model), "rotary_emb", None)
+    if rotary_emb is None:
+        return None
+
+    axes = [
+        torch.arange(1, target + 1, dtype=torch.int32, device=device),
+        torch.zeros(target, dtype=torch.int32, device=device),
+        torch.zeros(target, dtype=torch.int32, device=device),
+    ]
+    cap_pos_ids = torch.stack(axes, dim=-1)
+    return rotary_emb(cap_pos_ids)
+
+
+def _pad_caption_freqs(freqs_cis: Any, current_model: Any, *, target: int) -> Any:
+    if not isinstance(freqs_cis, (tuple, list)) or len(freqs_cis) != 2:
+        return freqs_cis
+
+    cap_cache, image_cache = freqs_cis
+    cap_tensor = bcg_utils.first_tensor(cap_cache)
+    if torch.is_tensor(cap_tensor) and cap_tensor.dim() >= 1:
+        cap_freqs = _build_caption_freqs(
+            current_model, target=target, device=cap_tensor.device
+        )
+        if cap_freqs is not None:
+            cap_cache = cap_freqs
+
+    if isinstance(freqs_cis, tuple):
+        return (cap_cache, image_cache)
+    return [cap_cache, image_cache]
+
+
+def _caption_mask(
+    call_kwargs: dict, *, caption: torch.Tensor, seq: int, bucket: int
+) -> torch.Tensor:
+    mask = bcg_utils.first_tensor(call_kwargs.get("encoder_hidden_states_mask"))
+    if not torch.is_tensor(mask):
+        mask = bcg_utils.first_tensor(call_kwargs.get("encoder_attention_mask"))
+    if torch.is_tensor(mask):
+        if mask.dim() == 1:
+            mask = mask[:seq].unsqueeze(0)
+        elif mask.dim() >= 2:
+            mask = mask[:, :seq]
+        mask = mask.to(device=caption.device, dtype=torch.bool)
+    else:
+        batch = int(caption.shape[0]) if caption.dim() == 3 else 1
+        mask = torch.ones((batch, seq), device=caption.device, dtype=torch.bool)
+    return bcg_utils.pad_tensor_dim(mask, 1, bucket)
+
+
+def pad_zimage_prompt_kwargs(
+    call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
+) -> dict:
+    caption = _first_caption_tensor(call_kwargs.get("encoder_hidden_states"))
+    if caption is None:
+        return call_kwargs
+
+    seq = _caption_seq_len(caption)
+    cap_freq = None
+    freqs_cis = call_kwargs.get("freqs_cis")
+    if isinstance(freqs_cis, (tuple, list)) and len(freqs_cis) == 2:
+        cap_freq = bcg_utils.first_tensor(freqs_cis[0])
+    cap_freq_len = int(cap_freq.shape[0]) if torch.is_tensor(cap_freq) else seq
+
+    bucket = bcg_utils.select_text_bucket(max(seq, cap_freq_len), buckets)
+    if bucket is None:
+        return call_kwargs
+
+    out = {
+        key: value
+        for key, value in call_kwargs.items()
+        if key
+        in {
+            "hidden_states",
+            "timestep",
+            "guidance",
+            "encoder_hidden_states",
+            "encoder_attention_mask",
+            "encoder_hidden_states_mask",
+            "freqs_cis",
+            "image_seq_len_target",
+            "patch_size",
+            "f_patch_size",
+        }
+    }
+
+    if seq < bucket:
+        out["encoder_hidden_states"] = _pad_caption(
+            out["encoder_hidden_states"], target=bucket
+        )
+
+    caption_mask = _caption_mask(call_kwargs, caption=caption, seq=seq, bucket=bucket)
+    out["encoder_hidden_states_mask"] = caption_mask
+    out["caption_valid_lens"] = caption_mask.sum(dim=1).to(dtype=torch.long)
+    out["_use_caption_valid_mask"] = True
+    if out.get("encoder_attention_mask") is not None:
+        out["encoder_attention_mask"] = out["encoder_hidden_states_mask"]
+    out["freqs_cis"] = _pad_caption_freqs(
+        out.get("freqs_cis"), current_model, target=bucket
+    )
+    return out
+
+
+bcg_utils.register_prompt_padder(is_zimage_transformer, pad_zimage_prompt_kwargs)

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

@@ -0,0 +1,306 @@
+# 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.
+# ==============================================================================
+"""Utilities for breakable CUDA graph (BCG) prompt padding.
+
+These helpers bucket prompt-conditioning inputs by sequence length so diffusion
+DiT forward calls with different prompt lengths can reuse captured CUDA graphs.
+Model-specific padders can register custom handling under
+``breakable_cuda_graph.model_padders``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Prompt-conditioning kwarg keys, grouped by which dim carries the text length.
+PROMPT_MASK_KEYS = (
+    "encoder_attention_mask",
+    "encoder_hidden_states_mask",
+    "attention_mask",
+    "text_mask",
+    "prompt_attention_mask",
+    "negative_attention_mask",
+    "prompt_embeds_mask",
+    "negative_prompt_embeds_mask",
+)
+TEXT_DIM1_KEYS = (
+    "encoder_hidden_states",
+    "encoder_hidden_states_2",
+    "encoder_attention_mask",
+    "encoder_hidden_states_mask",
+    "attention_mask",
+    "text_mask",
+    "text_ids",
+    "text_pos_ids",
+    "txt_ids",
+    "prompt_embeds",
+    "negative_prompt_embeds",
+    "prompt_attention_mask",
+    "negative_attention_mask",
+    "prompt_embeds_mask",
+    "negative_prompt_embeds_mask",
+    "audio_encoder_hidden_states",
+    "audio_encoder_attention_mask",
+)
+TEXT_DIM0_KEYS = (
+    "txt_freqs_cis",
+    "text_freqs_cis",
+)
+TEXT_SEQ_LEN_KEYS = (
+    "txt_seq_lens",
+    "text_seq_lens",
+)
+
+
+def first_tensor(obj: Any) -> torch.Tensor | None:
+    """First tensor leaf found by depth-first traversal (dicts in sorted-key
+    order), or ``None``."""
+    if torch.is_tensor(obj):
+        return obj
+    if isinstance(obj, (list, tuple)):
+        for item in obj:
+            tensor = first_tensor(item)
+            if tensor is not None:
+                return tensor
+    if isinstance(obj, dict):
+        for key in sorted(obj):
+            tensor = first_tensor(obj[key])
+            if tensor is not None:
+                return tensor
+    return None
+
+
+def select_text_bucket(seq: int, buckets: tuple[int, ...]) -> int | None:
+    """Smallest bucket that fits ``seq``; ``None`` (and a warning) when ``seq``
+    exceeds the largest bucket so the caller runs that length eagerly."""
+    for bucket in buckets:
+        if seq <= bucket:
+            return bucket
+    logger.warning(
+        "[Diffusion BCG] text length %d exceeds max bucket %d; not padding "
+        "(this length captures its own graph). Raise --bcg-text-buckets.",
+        seq,
+        buckets[-1],
+    )
+    return None
+
+
+def pad_tensor_dim(tensor: Any, dim: int, target: int, value: float = 0) -> Any:
+    if not torch.is_tensor(tensor) or tensor.dim() <= dim:
+        return tensor
+    seq = tensor.shape[dim]
+    if seq >= target:
+        return tensor
+    pad = [0, 0] * tensor.dim()
+    pad_index = 2 * (tensor.dim() - dim - 1) + 1
+    pad[pad_index] = target - seq
+    return torch.nn.functional.pad(tensor, tuple(pad), value=value)
+
+
+def pad_nested_dim(
+    obj: Any,
+    *,
+    dim: int,
+    source: int,
+    target: int,
+    value: float = 0,
+) -> Any:
+    if torch.is_tensor(obj):
+        if obj.dim() > dim and obj.shape[dim] == source:
+            return pad_tensor_dim(obj, dim, target, value)
+        return obj
+    if isinstance(obj, list):
+        return [
+            pad_nested_dim(item, dim=dim, source=source, target=target, value=value)
+            for item in obj
+        ]
+    if isinstance(obj, tuple):
+        return tuple(
+            pad_nested_dim(item, dim=dim, source=source, target=target, value=value)
+            for item in obj
+        )
+    return obj
+
+
+def bucket_txt_seq_lens(txt_seq_lens: Any, bucket: int) -> Any:
+    if txt_seq_lens is None:
+        return txt_seq_lens
+    if torch.is_tensor(txt_seq_lens):
+        return torch.full_like(txt_seq_lens, bucket)
+    if isinstance(txt_seq_lens, list):
+        return [bucket_txt_seq_lens(seq_len, bucket) for seq_len in txt_seq_lens]
+    if isinstance(txt_seq_lens, tuple):
+        return tuple(bucket_txt_seq_lens(seq_len, bucket) for seq_len in txt_seq_lens)
+    if isinstance(txt_seq_lens, int):
+        return bucket
+    return txt_seq_lens
+
+
+def prompt_seq_and_dim(call_kwargs: dict) -> tuple[int, int] | None:
+    """Return ``(text_seq_len, seq_dim)`` inferred from the prompt embeddings or
+    a prompt mask, or ``None`` when no text conditioning is present."""
+    ehs_tensor = first_tensor(call_kwargs.get("encoder_hidden_states"))
+    if torch.is_tensor(ehs_tensor) and ehs_tensor.dim() >= 2:
+        if ehs_tensor.dim() == 2:
+            return int(ehs_tensor.shape[0]), 0
+        return int(ehs_tensor.shape[1]), 1
+
+    for key in PROMPT_MASK_KEYS:
+        tensor = first_tensor(call_kwargs.get(key))
+        if torch.is_tensor(tensor) and tensor.dim() >= 2:
+            if tensor.shape[0] == 1:
+                return int(tensor.shape[1]), 1
+            return int(tensor.shape[0]), 0
+    return None
+
+
+def pad_nested_text_dim(
+    obj: Any,
+    *,
+    source: int,
+    target: int,
+    preferred_dim: int,
+) -> Any:
+    if torch.is_tensor(obj):
+        if obj.dim() > preferred_dim and obj.shape[preferred_dim] == source:
+            return pad_tensor_dim(obj, preferred_dim, target)
+        for dim in (1, 0):
+            if dim != preferred_dim and obj.dim() > dim and obj.shape[dim] == source:
+                return pad_tensor_dim(obj, dim, target)
+        return obj
+    if isinstance(obj, list):
+        return [
+            pad_nested_text_dim(
+                item, source=source, target=target, preferred_dim=preferred_dim
+            )
+            for item in obj
+        ]
+    if isinstance(obj, tuple):
+        return tuple(
+            pad_nested_text_dim(
+                item, source=source, target=target, preferred_dim=preferred_dim
+            )
+            for item in obj
+        )
+    if isinstance(obj, dict):
+        return {
+            key: pad_nested_text_dim(
+                value, source=source, target=target, preferred_dim=preferred_dim
+            )
+            for key, value in obj.items()
+        }
+    return obj
+
+
+def bucket_text_seq_lens(obj: Any, *, target: int) -> Any:
+    if isinstance(obj, int) and not isinstance(obj, bool):
+        return target
+    if isinstance(obj, list):
+        return [bucket_text_seq_lens(item, target=target) for item in obj]
+    if isinstance(obj, tuple):
+        return tuple(bucket_text_seq_lens(item, target=target) for item in obj)
+    return obj
+
+
+def pad_masked_prompt_kwargs(call_kwargs: dict, buckets: tuple[int, ...]) -> dict:
+    """Generic, model-agnostic prompt padding for models that pass a prompt
+    attention mask alongside their text embeddings."""
+    seq_and_dim = prompt_seq_and_dim(call_kwargs)
+    if seq_and_dim is None:
+        return call_kwargs
+    seq, seq_dim = seq_and_dim
+    has_mask = any(
+        first_tensor(call_kwargs.get(key)) is not None for key in PROMPT_MASK_KEYS
+    )
+    if not has_mask:
+        return call_kwargs
+    bucket = select_text_bucket(seq, buckets)
+    if bucket is None or seq == bucket:
+        return call_kwargs
+
+    out = dict(call_kwargs)
+    for key in TEXT_DIM1_KEYS:
+        if key in out and out[key] is not None:
+            out[key] = pad_nested_text_dim(
+                out[key], source=seq, target=bucket, preferred_dim=seq_dim
+            )
+    for key in TEXT_DIM0_KEYS:
+        if key in out and out[key] is not None:
+            out[key] = pad_nested_dim(out[key], dim=0, source=seq, target=bucket)
+    for key in TEXT_SEQ_LEN_KEYS:
+        if key in out and out[key] is not None:
+            out[key] = bucket_text_seq_lens(out[key], target=bucket)
+    return out
+
+
+def transformer_class_name_matches(current_model: Any, needle: str) -> bool:
+    """True when ``current_model`` (or its ``module`` / ``_orig_mod`` wrapper)
+    is a transformer whose qualified class name contains ``needle``."""
+    candidates = [current_model]
+    for attr in ("module", "_orig_mod"):
+        wrapped = getattr(current_model, attr, None)
+        if wrapped is not None:
+            candidates.append(wrapped)
+    for candidate in candidates:
+        cls = type(candidate)
+        name = f"{cls.__module__}.{cls.__qualname__}".lower()
+        if needle in name:
+            return True
+    return False
+
+
+# --- Model-specific prompt-padder registry ------------------------------- #
+# Each model that needs custom prompt padding registers a (predicate, padder)
+# pair from its own module in ``model_specific_stages`` so the base denoising
+# stage stays model-agnostic. ``padder(call_kwargs, current_model, buckets)``
+# returns the padded kwargs.
+PromptPadder = Callable[[dict, Any, tuple], dict]
+_PROMPT_PADDERS: list[tuple[Callable[[Any, dict], bool], PromptPadder]] = []
+
+
+def register_prompt_padder(
+    predicate: Callable[[Any, dict], bool], padder: PromptPadder
+) -> None:
+    _PROMPT_PADDERS.append((predicate, padder))
+
+
+def select_prompt_padder(current_model: Any, call_kwargs: dict) -> PromptPadder | None:
+    """Return the registered model-specific padder for ``current_model``, or
+    ``None`` to fall back to :func:`pad_masked_prompt_kwargs`."""
+    _ensure_model_padders_registered()
+    for predicate, padder in _PROMPT_PADDERS:
+        if predicate(current_model, call_kwargs):
+            return padder
+    return None
+
+
+_model_padders_registered = False
+
+
+def _ensure_model_padders_registered() -> None:
+    """Import the model-specific padder modules once so they register."""
+    global _model_padders_registered
+    if _model_padders_registered:
+        return
+    _model_padders_registered = True
+    from sglang.multimodal_gen.runtime.breakable_cuda_graph.model_padders import (  # noqa: F401
+        ideogram,
+        qwen_image,
+        zimage,
+    )

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

@@ -0,0 +1,31 @@
+"""Replay-token tracking for diffusion BCG replays.
+
+The SRT BCG core does not stamp replays; the diffusion runner sets a fresh
+token around each graph replay so replay-local caches (e.g. varlen attention
+mask metadata in ``DynamicVarlenMaskMeta``) can be rebuilt once per replay
+while still being reused across the break points of that same replay.
+``get_current_replay_token`` returns ``None`` outside a replay (including
+during capture).
+"""
+
+import itertools
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+_current_replay_token_var: ContextVar[int | None] = ContextVar(
+    "mm_bcg_replay_token", default=None
+)
+_replay_token_counter = itertools.count(1)
+
+
+def get_current_replay_token() -> int | None:
+    return _current_replay_token_var.get()
+
+
+@contextmanager
+def replay_token_scope():
+    token = _current_replay_token_var.set(next(_replay_token_counter))
+    try:
+        yield
+    finally:
+        _current_replay_token_var.reset(token)

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

@@ -0,0 +1,472 @@
+# 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.
+# ==============================================================================
+"""Breakable CUDA graph (BCG) runner for diffusion DiT transformers.
+
+A runner wraps a callable ``nn.Module`` and turns it into an *eager runner* that
+transparently proxies every attribute to the wrapped module and, when called,
+replays a previously captured graph for the input signature — or runs the
+module eagerly when no graph was captured for that signature. Capture is an
+explicit, idempotent ``capture()`` call (driven at warmup) so that serving never
+triggers a fresh capture.
+
+This file is intentionally local to ``multimodal_gen``: diffusion reuses the
+low-level SRT BCG primitives, but the capture/replay runner owns diffusion DiT
+signature handling, static tensor buffers, prompt-bucket warmup, and fallback
+behavior.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph.replay_token import (
+    replay_token_scope,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    enable_breakable_cuda_graph,
+)
+
+# Log under the multimodal_gen namespace so the diffusion server's logging
+# config surfaces the "[Diffusion BCG] captured ..." lines.
+logger = logging.getLogger(__name__)
+
+
+def _env_int(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        logger.warning("[BCG] ignoring invalid integer %s=%r", name, raw)
+        return default
+
+
+def _env_float(name: str, default: float) -> float:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        return float(raw)
+    except ValueError:
+        logger.warning("[BCG] ignoring invalid float %s=%r", name, raw)
+        return default
+
+
+def _map_tensors(obj, fn):
+    """Rebuild ``obj`` applying ``fn`` to every tensor leaf, recursing into
+    list/tuple/dict containers; everything else passes through unchanged."""
+    if torch.is_tensor(obj):
+        return fn(obj)
+    if isinstance(obj, tuple):
+        return tuple(_map_tensors(o, fn) for o in obj)
+    if isinstance(obj, list):
+        return [_map_tensors(o, fn) for o in obj]
+    if isinstance(obj, dict):
+        return {k: _map_tensors(v, fn) for k, v in obj.items()}
+    return obj
+
+
+def _flatten_tensors(obj, out: list):
+    """Depth-first collect every tensor leaf into ``out`` (deterministic order:
+    dicts traversed in sorted-key order to match across calls)."""
+    if torch.is_tensor(obj):
+        out.append(obj)
+    elif isinstance(obj, (list, tuple)):
+        for o in obj:
+            _flatten_tensors(o, out)
+    elif isinstance(obj, dict):
+        for k in sorted(obj):
+            _flatten_tensors(obj[k], out)
+
+
+def _flatten_kwargs(kwargs: dict[str, Any]) -> list[torch.Tensor]:
+    out: list[torch.Tensor] = []
+    for name in sorted(kwargs):
+        _flatten_tensors(kwargs[name], out)
+    return out
+
+
+def _signature_leaf(obj: Any) -> Any:
+    if torch.is_tensor(obj):
+        return ("tensor", tuple(obj.shape), str(obj.dtype))
+    if isinstance(obj, tuple):
+        return ("tuple", tuple(_signature_leaf(o) for o in obj))
+    if isinstance(obj, list):
+        return ("list", tuple(_signature_leaf(o) for o in obj))
+    if isinstance(obj, dict):
+        return (
+            "dict",
+            tuple((k, _signature_leaf(obj[k])) for k in sorted(obj)),
+        )
+    if obj is None or isinstance(obj, (bool, int, float, str)):
+        return ("const", obj)
+    return ("object", type(obj).__module__, type(obj).__qualname__, id(obj))
+
+
+def _signature_kwargs(kwargs: dict[str, Any]) -> tuple:
+    return tuple((name, _signature_leaf(kwargs[name])) for name in sorted(kwargs))
+
+
+def _signature_summary_leaf(sig: Any, *, depth: int = 0) -> Any:
+    if not isinstance(sig, tuple) or not sig:
+        return sig
+
+    tag = sig[0]
+    if tag == "tensor":
+        return sig
+    if tag == "const":
+        value = sig[1]
+        if isinstance(value, str) and len(value) > 64:
+            value = value[:61] + "..."
+        return (tag, value)
+    if tag == "object":
+        return sig[:3]
+    if depth >= 2:
+        return (tag, "...")
+    if tag in ("tuple", "list"):
+        items = sig[1]
+        preview = tuple(
+            _signature_summary_leaf(item, depth=depth + 1) for item in items[:4]
+        )
+        if len(items) > 4:
+            preview += (("...", len(items) - 4),)
+        return (tag, len(items), preview)
+    if tag == "dict":
+        items = sig[1]
+        preview = tuple(
+            (key, _signature_summary_leaf(value, depth=depth + 1))
+            for key, value in items[:4]
+        )
+        if len(items) > 4:
+            preview += (("...", len(items) - 4),)
+        return (tag, len(items), preview)
+    return sig
+
+
+def _signature_summary(key: tuple) -> tuple:
+    return tuple((name, _signature_summary_leaf(value)) for name, value in key[:16]) + (
+        (("...", len(key) - 16),) if len(key) > 16 else ()
+    )
+
+
+def _clone_output(out: Any) -> Any:
+    if torch.is_tensor(out):
+        return out.clone()
+    if isinstance(out, tuple):
+        return tuple(_clone_output(o) for o in out)
+    if isinstance(out, list):
+        return [_clone_output(o) for o in out]
+    return out
+
+
+@dataclass
+class _CaptureEntry:
+    graph: BreakableCUDAGraph
+    # full captured kwargs with persistent static buffers at every tensor leaf
+    static_kwargs: dict[str, Any]
+    # the same static buffers, flattened in _flatten_kwargs order (replay copies
+    # live tensors into these positionally)
+    static_leaves: list[torch.Tensor]
+    output: Any
+    num_segments: int
+
+
+class _CaptureRejected(RuntimeError):
+    pass
+
+
+class BaseBreakableCudaGraphRunner:
+    """Eager runner around ``transformer`` with an explicit capture/replay API.
+
+    The capture/replay contract:
+
+    * :meth:`capture` captures a BCG graph for the given input signature, once
+      (idempotent). It is intended to be driven at warmup so that every
+      signature served later is already captured.
+    * :meth:`replay` copies live inputs into the captured static buffers and
+      replays the graph, returning a clone of the captured output.
+    * :meth:`__call__` is the *eager runner*: it replays when a graph exists for
+      the signature and otherwise runs ``transformer`` eagerly. It never
+      captures, so serving never pays a capture cost.
+
+    Any attribute not defined on the runner is proxied to ``transformer`` so the
+    runner can stand in for the wrapped module ("other functions directly
+    pass").
+    """
+
+    def __init__(
+        self,
+        transformer: nn.Module,
+        device: torch.device,
+        pool=None,
+    ) -> None:
+        self.transformer = transformer
+        self.device = device
+        self.device_module = torch.get_device_module(device)
+        # One shared mempool across all captured graphs/segments so per-block
+        # intermediates can be reclaimed and weak-ref'd safely.
+        self._pool = (
+            pool if pool is not None else self.device_module.graph_pool_handle()
+        )
+        self._capture_stream = self.device_module.Stream(device=device)
+        self.entries: dict[tuple, _CaptureEntry] = {}
+        # Signatures we have given up capturing (capture raised); run eager.
+        self._blocked: set[tuple] = set()
+        self._disabled_reason: str | None = None
+        self.max_entries = max(0, _env_int("SGLANG_DIFFUSION_BCG_MAX_ENTRIES", 32))
+        self.max_segments = max(0, _env_int("SGLANG_DIFFUSION_BCG_MAX_SEGMENTS", 128))
+
+    def __getattr__(self, name: str) -> Any:
+        # Only reached for attributes the runner itself does not define; proxy
+        # them to the wrapped transformer so callers can treat the runner as a
+        # transparent stand-in. Use __dict__ to avoid recursing through
+        # __getattr__ before ``transformer`` is assigned in __init__.
+        try:
+            transformer = self.__dict__["transformer"]
+        except KeyError as e:  # pragma: no cover - during/ before __init__
+            raise AttributeError(name) from e
+        return getattr(transformer, name)
+
+    # ------------------------------------------------------------------ #
+    # Public capture / replay API
+    # ------------------------------------------------------------------ #
+    @torch.no_grad()
+    def capture(self, **kwargs) -> bool:
+        """Capture a graph for ``kwargs``'s signature if not already captured.
+
+        Idempotent: returns ``True`` when a graph is available for the
+        signature afterwards (already captured or newly captured), ``False``
+        when capture is disabled/blocked or failed (the caller then runs eager).
+        """
+        if self._disabled_reason is not None:
+            return False
+        key = self._signature(kwargs)
+        if key in self._blocked:
+            return False
+        if key in self.entries:
+            return True
+        try:
+            entry = self._capture(kwargs, key)
+        except Exception as e:  # noqa: BLE001 — never break generation on capture
+            logger.warning(
+                "[Diffusion BCG] capture failed for signature %s (%s); "
+                "this signature will run eager.",
+                _signature_summary(key),
+                e,
+            )
+            self._blocked.add(key)
+            return False
+        self.entries[key] = entry
+        self._evict_entries_if_needed()
+        return True
+
+    def _should_capture_on_call(self, key: tuple) -> bool:
+        """Whether ``__call__`` may lazily capture an unseen signature.
+
+        Base runners only ever capture through the explicit :meth:`capture`
+        API, so this returns ``False``: serving never records a fresh graph.
+        Subclasses gate lazy capture on a warmup window (see the diffusion
+        runner) so warmup can capture by simply driving the forward as usual.
+        """
+        return False
+
+    @torch.no_grad()
+    def __call__(self, **kwargs) -> Any:
+        """Eager runner: replay a captured graph, else run ``transformer``.
+
+        While serving this never captures, so no new graph is recorded once
+        warmup is over. During the warmup window subclasses opt into lazy
+        capture via :meth:`_should_capture_on_call`.
+        """
+        if self._disabled_reason is not None:
+            return self.transformer(**kwargs)
+        key = self._signature(kwargs)
+        entry = self.entries.get(key)
+        if entry is None:
+            if not self._should_capture_on_call(key):
+                return self.transformer(**kwargs)
+            if not self.capture(**kwargs):
+                return self.transformer(**kwargs)
+            entry = self.entries[key]
+        return self.replay(entry, kwargs)
+
+    def replay(self, entry: _CaptureEntry, kwargs: dict[str, Any]) -> Any:
+        live_leaves = _flatten_kwargs(kwargs)
+        if len(live_leaves) != len(entry.static_leaves):
+            # Structure changed under a matching shape key — should not happen;
+            # fall back to eager rather than copy mismatched buffers.
+            return self.transformer(**kwargs)
+        for buf, live in zip(entry.static_leaves, live_leaves):
+            buf.copy_(live, non_blocking=True)
+        with replay_token_scope():
+            entry.graph.replay()
+        # Clone so the caller can hold the result across the next replay / the
+        # other CFG branch (which shares this static output buffer when shapes
+        # match). The clone is one cheap DtoD copy relative to the full DiT.
+        return _clone_output(entry.output)
+
+    # ------------------------------------------------------------------ #
+    # Internals
+    # ------------------------------------------------------------------ #
+    def _signature(self, kwargs: dict[str, Any]) -> tuple:
+        """Capture key for tensor leaves and non-tensor control values.
+
+        Tensor leaves are keyed by shape+dtype so their values can change per
+        replay. Non-tensor leaves are baked into the captured Python control
+        flow, so simple constants must be part of the key as well. Mutable
+        objects are keyed by identity to avoid replaying a graph whose eager
+        break points still reference a previous request's state object.
+        """
+        return _signature_kwargs(kwargs)
+
+    def _empty_cache(self) -> None:
+        empty_cache = getattr(self.device_module, "empty_cache", None)
+        if callable(empty_cache):
+            empty_cache()
+
+    @staticmethod
+    def _drop_entry(entry: _CaptureEntry) -> None:
+        entry.graph._break_fns.clear()
+        entry.graph._segments.clear()
+        entry.static_kwargs.clear()
+        entry.static_leaves.clear()
+        entry.output = None
+
+    def reset(self, *, disabled_reason: str | None = None) -> None:
+        for entry in self.entries.values():
+            self._drop_entry(entry)
+        self.entries.clear()
+        self._blocked.clear()
+        self._pool = None
+        self._empty_cache()
+        if disabled_reason is not None:
+            self._disabled_reason = disabled_reason
+
+    def _capture_limit_reason(self, entry: _CaptureEntry) -> str | None:
+        if self.max_segments and entry.num_segments > self.max_segments:
+            return (
+                f"captured {entry.num_segments} segments, above "
+                f"SGLANG_DIFFUSION_BCG_MAX_SEGMENTS={self.max_segments}"
+            )
+        return None
+
+    def _evict_entries_if_needed(self) -> None:
+        if not self.max_entries:
+            return
+        while len(self.entries) > self.max_entries:
+            evicted_key = next(iter(self.entries))
+            entry = self.entries.pop(evicted_key)
+            self._drop_entry(entry)
+            logger.info(
+                "[Diffusion BCG] evicted oldest capture for signature %s "
+                "(SGLANG_DIFFUSION_BCG_MAX_ENTRIES=%d)",
+                _signature_summary(evicted_key),
+                self.max_entries,
+            )
+        self._empty_cache()
+
+    def _capture(self, kwargs: dict[str, Any], key: tuple) -> _CaptureEntry:
+        if self._pool is None:
+            self._pool = self.device_module.graph_pool_handle()
+
+        # Persistent static buffers at every tensor leaf; bake non-tensors.
+        def _to_static(t: torch.Tensor) -> torch.Tensor:
+            # Static buffers live on the capture device. A CPU input (e.g. a
+            # scalar timestep/sigma or an index tensor built on the host)
+            # would otherwise force a CPU->CUDA copy inside the captured
+            # region, which is illegal; place its buffer on the device so the
+            # only host->device copy happens here, before capture, and replay
+            # is device-to-device.
+            if t.device.type == "cpu":
+                buf = torch.empty(t.shape, dtype=t.dtype, device=self.device)
+            else:
+                buf = torch.empty_like(t)
+            buf.copy_(t)
+            return buf
+
+        static_kwargs = {
+            name: _map_tensors(v, _to_static) for name, v in kwargs.items()
+        }
+        static_leaves = _flatten_kwargs(static_kwargs)
+
+        # Warm up on the capture stream so cuBLAS/cuDNN/Triton workspaces and
+        # any lazy JIT are materialized before capture (mirrors the LLM runner
+        # and torch.cuda.make_graphed_callables).
+        self.device_module.synchronize()
+        with self.device_module.stream(self._capture_stream):
+            for _ in range(2):
+                self.transformer(**static_kwargs)
+        self._capture_stream.synchronize()
+        self.device_module.synchronize()
+
+        graph = BreakableCUDAGraph()
+        with enable_breakable_cuda_graph():
+            with BreakableCUDAGraphCapture(
+                cuda_graph=graph, pool=self._pool, stream=self._capture_stream
+            ):
+                output = self.transformer(**static_kwargs)
+        self.device_module.synchronize()
+
+        logger.info(
+            "[Diffusion BCG] captured %d segment(s), %d tensor input(s) for "
+            "signature %s",
+            len(graph._segments),
+            len(static_leaves),
+            _signature_summary(key),
+        )
+        entry = _CaptureEntry(
+            graph=graph,
+            static_kwargs=static_kwargs,
+            static_leaves=static_leaves,
+            output=output,
+            num_segments=len(graph._segments),
+        )
+        limit_reason = self._capture_limit_reason(entry)
+        if limit_reason is not None:
+            self._drop_entry(entry)
+            self.reset(disabled_reason=limit_reason)
+            raise _CaptureRejected(
+                f"{limit_reason}; disabling this BCG runner and using eager"
+            )
+        return entry
+
+
+class DiffusionBreakableCudaGraphRunner(BaseBreakableCudaGraphRunner):
+    """Capture/replay a diffusion DiT ``transformer`` with BCG.
+
+    Unknown attributes proxy to the wrapped transformer, so the runner can
+    stand in for the module while only intercepting ``forward`` calls.
+    """
+
+    def _should_capture_on_call(self, key) -> bool:
+        try:
+            from sglang.multimodal_gen.runtime.managers.forward_context import (
+                get_forward_context,
+            )
+
+            forward_batch = get_forward_context().forward_batch
+        except Exception:
+            return False
+        return bool(getattr(forward_batch, "is_warmup", False))