chore: import upstream snapshot with attribution

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