chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/vla/denoise_cuda_graph.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import torch
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.vla.prefix_cache import (
+    PrefixContext,
+    VLADensePrefixCache,
+)
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class VLADenoiseGraphSignature:
+    batch_size: int
+    prefix_len: int
+    action_horizon: int
+    action_dim: int
+    dtype: str
+    parallel_layout: str
+
+
+@dataclass
+class _CapturedDenoiseGraph:
+    graph: torch.cuda.CUDAGraph
+    static_prefix_context: PrefixContext
+    static_x_t: torch.Tensor
+    static_timestep: torch.Tensor
+    static_output: torch.Tensor
+    current_context_id: int | None = None
+    current_context_digest: str | None = None
+
+
+def _clone_past_key_values(past_key_values: Any) -> Any:
+    return VLADensePrefixCache(
+        tuple(
+            (keys.detach().clone(), values.detach().clone(), sliding_window)
+            for keys, values, sliding_window in past_key_values
+        )
+    )
+
+
+def _copy_past_key_values_(dst: Any, src: Any) -> None:
+    for (dst_keys, dst_values, _), (src_keys, src_values, _) in zip(
+        dst, src, strict=True
+    ):
+        dst_keys.copy_(src_keys)
+        dst_values.copy_(src_values)
+
+
+def _clone_prefix_context(prefix_context: PrefixContext) -> PrefixContext:
+    return PrefixContext(
+        past_key_values=_clone_past_key_values(prefix_context.past_key_values),
+        prefix_pad_masks=prefix_context.prefix_pad_masks.detach().clone(),
+        prefix_len=prefix_context.prefix_len,
+        layout=dict(prefix_context.layout),
+        cache_key_digest=prefix_context.cache_key_digest,
+    )
+
+
+def _copy_prefix_context_(dst: PrefixContext, src: PrefixContext) -> None:
+    dst.prefix_pad_masks.copy_(src.prefix_pad_masks)
+    _copy_past_key_values_(dst.past_key_values, src.past_key_values)
+    dst.cache_key_digest = src.cache_key_digest
+
+
+class VLADenoiseGraphRunner:
+    """Full CUDA graph runner for one VLA action-denoise step.
+
+    Each signature owns fixed input and output buffers. This does not use
+    diffusion BCG and does not capture prefix encoding or token decode.
+    """
+
+    def __init__(self, enabled: bool = True):
+        self.enabled = enabled
+        self._captured: dict[VLADenoiseGraphSignature, _CapturedDenoiseGraph] = {}
+        self._disabled_signatures: set[VLADenoiseGraphSignature] = set()
+        self._capture_stream: torch.cuda.Stream | None = None
+        self._graph_pool: Any = None
+
+    def _sync_context_if_needed(
+        self,
+        captured: _CapturedDenoiseGraph,
+        prefix_context: PrefixContext,
+    ) -> None:
+        context_id = id(prefix_context.past_key_values)
+        context_digest = prefix_context.cache_key_digest
+        if (
+            context_digest is not None
+            and captured.current_context_digest == context_digest
+        ):
+            captured.current_context_id = context_id
+            return
+        if captured.current_context_id == context_id:
+            return
+        _copy_prefix_context_(captured.static_prefix_context, prefix_context)
+        captured.current_context_id = context_id
+        captured.current_context_digest = context_digest
+
+    def _capture(
+        self,
+        signature: VLADenoiseGraphSignature,
+        step_fn: Callable[..., torch.Tensor],
+        prefix_context: PrefixContext,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> _CapturedDenoiseGraph:
+        static_prefix_context = _clone_prefix_context(prefix_context)
+        static_x_t = x_t.detach().clone()
+        static_timestep = timestep.detach().clone()
+
+        device_module = torch.get_device_module(x_t.device)
+        if self._capture_stream is None:
+            self._capture_stream = device_module.Stream(device=x_t.device)
+        if self._graph_pool is None:
+            self._graph_pool = get_or_create_global_graph_memory_pool(device_module)
+            set_graph_pool_id(self._graph_pool)
+
+        # warm up lazy kernels and workspaces before capture
+        device_module.synchronize()
+        with device_module.stream(self._capture_stream), torch.inference_mode():
+            step_fn(
+                static_prefix_context,
+                static_x_t,
+                static_timestep,
+            )
+        self._capture_stream.synchronize()
+
+        graph = torch.cuda.CUDAGraph()
+        with (
+            device_module.graph(
+                cuda_graph=graph,
+                pool=self._graph_pool,
+                stream=self._capture_stream,
+            ),
+            torch.inference_mode(),
+        ):
+            static_output = step_fn(
+                static_prefix_context,
+                static_x_t,
+                static_timestep,
+            )
+        self._capture_stream.synchronize()
+
+        captured = _CapturedDenoiseGraph(
+            graph=graph,
+            static_prefix_context=static_prefix_context,
+            static_x_t=static_x_t,
+            static_timestep=static_timestep,
+            static_output=static_output,
+            current_context_id=id(prefix_context.past_key_values),
+            current_context_digest=prefix_context.cache_key_digest,
+        )
+        self._captured[signature] = captured
+        logger.info(
+            "Captured VLA denoise CUDA graph: batch=%d prefix=%d action=%dx%d "
+            "dtype=%s",
+            signature.batch_size,
+            signature.prefix_len,
+            signature.action_horizon,
+            signature.action_dim,
+            signature.dtype,
+        )
+        return captured
+
+    def capture_or_run(
+        self,
+        signature: VLADenoiseGraphSignature,
+        step_fn: Callable[..., torch.Tensor],
+        prefix_context: PrefixContext,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self.enabled or signature in self._disabled_signatures:
+            return step_fn(prefix_context, x_t, timestep)
+
+        if x_t.device.type != "cuda":
+            return step_fn(prefix_context, x_t, timestep)
+
+        captured = self._captured.get(signature)
+        try:
+            if captured is None:
+                captured = self._capture(
+                    signature, step_fn, prefix_context, x_t, timestep
+                )
+                captured.graph.replay()
+            else:
+                self._sync_context_if_needed(captured, prefix_context)
+                captured.static_x_t.copy_(x_t)
+                captured.static_timestep.copy_(timestep)
+                captured.graph.replay()
+            return captured.static_output
+        except Exception:
+            self._disabled_signatures.add(signature)
+            self._captured.pop(signature, None)
+            logger.warning(
+                "VLA denoise CUDA graph disabled for signature %s",
+                signature,
+                exc_info=True,
+            )
+            return step_fn(prefix_context, x_t, timestep)