chore: import upstream snapshot with attribution

python/paddle/device/cuda/cuda_graphed_layer.py

@@ -0,0 +1,500 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+import logging
+import os
+from collections import deque
+from enum import Enum
+
+import paddle
+from paddle.base import log_helper
+
+from .graphs import CUDAGraph
+
+# CUDAGraphedLayer Debug tools
+enable_debug_print = bool(
+    int(os.getenv('PADDLE_DEBUG_ENABLE_CUDAGRAPH_LAYER_LOGGING', '0'))
+)
+debug_cudagraphedlayer_fallback_to_default = bool(
+    int(os.getenv('PADDLE_DEBUG_CUDAGRAPHEDLAYER_FALLBACK_TO_DEFAULT', '0'))
+)
+
+logger = log_helper.get_logger(
+    __name__, logging.INFO, fmt='[%(levelname)s] %(message)s'
+)
+
+
+def debug_print(x):
+    if not enable_debug_print:
+        return
+    logger.info(x)
+
+
+def print_tensor(
+    t,
+    name="Unnamed",
+    print_meta=True,
+    print_ptr=False,
+    print_hash=True,
+    hash=None,
+):
+    output = []
+    if name:
+        output.append(name)
+    if hash is None:
+        hash = lambda t: float((t.astype('float32') * 1000).sum())
+
+    if t is None:
+        debug_print(f"{name} is None")
+    elif isinstance(t, paddle.Tensor):
+        if print_meta:
+            output.append(f"shape = {t.shape}")
+            output.append(f"place = {t.place}")
+        if print_ptr:
+            output.append(f"ptr = {hex(t.data_ptr())}")
+        if print_hash:
+            output.append(f"hash = {hash(t)}")
+        debug_print(" | ".join(output))
+
+
+def printer(x, banner="printer"):
+    if not enable_debug_print:
+        return
+    debug_print(banner.center(100, "-"))
+    recursive_apply(print_tensor, x)
+
+
+# We need this function, for any kind of inputs with iterables
+# we recursively apply the function to the leave nodes
+def recursive_apply(function, input_var):
+    if isinstance(input_var, list):
+        return [recursive_apply(function, item) for item in input_var]
+    elif isinstance(input_var, tuple):
+        return tuple(recursive_apply(function, item) for item in input_var)
+    elif isinstance(input_var, dict):
+        return {
+            key: recursive_apply(function, value)
+            for key, value in input_var.items()
+        }
+    else:
+        return function(input_var)
+
+
+def detach_tensor(tensor):
+    # Detach an individual tensor and preserve its 'stop_gradient' property
+    if isinstance(tensor, paddle.Tensor):
+        detached_tensor = tensor.detach()
+        detached_tensor.stop_gradient = tensor.stop_gradient
+        return detached_tensor
+    return tensor
+
+
+# We try our best to flatten the input to list of tensors
+# example: args = ((t1,t2),(t3,(t4,t5))) -> [t1, t2, t3, t4, t5]
+def recursive_flatten(target):
+    ret = []
+
+    def append(arg):
+        if isinstance(arg, paddle.Tensor):
+            # [NOTE] sometimes unnecessary tensors, such as the constant `mask` tensor in the PP layer, is passed into subsequent layers.
+            # When a tensor is marked with `stop_gradient=True`, it indicates that it does not contribute to gradient calculations,
+            # suggesting it's unrelated to the main computational process.
+            # Therefore, I try to eliminate the copying of such tensors in the to optimize performance.
+            # if not arg.stop_gradient:
+            # [NOTE] However, `stop_gradient=True` propagation rules within the framework appear to be flawed, so directly eliminate stop_gradient may cause bug
+            ret.append(arg)
+
+    recursive_apply(append, target)
+    return ret
+
+
+# input any kind of args / kwargs structure, output list of tensor
+def recursive_flatten_args_kwargs(args, kwargs):
+    return [
+        *recursive_flatten(args),
+        *recursive_flatten(tuple(kwargs.values())),
+    ]
+
+
+detach = lambda x: recursive_apply(detach_tensor, x)
+
+
+def get_grad_tensor(x):
+    """Returns the gradient of a Paddle Tensor if it's a tensor; otherwise, returns the input."""
+    if isinstance(x, paddle.Tensor):
+        if x.stop_gradient:
+            return None
+        else:
+            return x.grad
+    return None
+
+
+# CUDA Graph with Static Input and Output
+class CUDAGraphWithStaticInputOutput:
+    def __init__(self, num_warmup_steps):
+        self.num_warmup_steps = num_warmup_steps
+        self.graph = CUDAGraph()
+
+        self.has_recorded = False
+
+        self.has_preserved_inputs = False
+        self.args_static = None
+        self.kwargs_static = None
+
+        # inputs is the recursively flattened args and kwargs
+        self.inputs_static = None
+        self.outputs_static = None
+
+    def preserve_or_copy(self, args, kwargs):
+        """
+        For the CUDA Graph, it is crucial that the buffer remains address-stable,
+        meaning that the buffer addresses for any inputs to the CUDA Graph should not change.
+        One solution to achieve this is to preserve all input tensors.
+
+        This function attempts to recursively flatten the input arguments and keyword arguments
+        to identify all tensors passed to the layer (though it may still miss some due to other implicit
+        ways inputs can be passed to a layer). It then preserves references to these input tensors
+        as `self.inputs_static` so that the buffer pointers can be reused later.
+
+        When this method is called subsequently, it copies the values back to the preserved input tensors
+        to ensure the buffers are reused.
+        """
+        if not self.has_preserved_inputs:
+            self.has_preserved_inputs = True
+            self.args_static = args
+            self.kwargs_static = kwargs
+            self.inputs_static = recursive_flatten_args_kwargs(
+                self.args_static, self.kwargs_static
+            )
+        else:
+            inputs = recursive_flatten_args_kwargs(args, kwargs)
+            for x_static, x in zip(self.inputs_static, inputs):
+                x_static.copy_(x, True)
+
+    def record(self, f, *args, **kwargs):
+        self.preserve_or_copy(args, kwargs)
+
+        self.graph.capture_begin()
+        self.outputs_static = f(*self.args_static, **self.kwargs_static)
+        self.graph.capture_end()
+        debug_print(
+            "[CUDAGraph] Record-Replay Start (Graph is replayed for the first time)"
+        )
+        self.graph.replay()
+
+        self.has_recorded = True
+
+        return self.outputs_static
+
+    def set_output_static(self, outputs_static):
+        self.outputs_static = outputs_static
+
+    def replay(self, *args, **kwargs):
+        if not self.has_recorded:
+            raise RuntimeError("Graph should be recorded first")
+
+        self.preserve_or_copy(args, kwargs)
+
+        debug_print("[CUDAGraph] Replay Start")
+        self.graph.replay()
+        return self.outputs_static
+
+    def save(self, name):
+        logging.info(f"save graph to {name}")
+        self.graph.print_to_dot_files(name)
+
+
+# CUDA Graph Layer Status Enumeration
+class CUDAGraphLayerStatus(Enum):
+    """Enum to represent the status of a CUDA Graph Layer."""
+
+    WARMUP = 1
+    RECORD = 2
+    CUDAGRAPH = 3
+
+
+class CUDAGraphForwardBackward:
+    def __init__(self, num_warmup_steps):
+        self.forward_graph = CUDAGraphWithStaticInputOutput(num_warmup_steps)
+        self.backward_graph = CUDAGraphWithStaticInputOutput(num_warmup_steps)
+        self.status = CUDAGraphLayerStatus.RECORD
+
+    def capture_end(self):
+        self.status = CUDAGraphLayerStatus.CUDAGRAPH
+
+    def is_record_step(self):
+        return self.status == CUDAGraphLayerStatus.RECORD
+
+    def is_cuda_graph_step(self):
+        return self.status == CUDAGraphLayerStatus.CUDAGRAPH
+
+
+class CUDAGraphContext:
+    """
+    Manages the context for CUDA graph execution in layers. This includes handling
+    the state of CUDA graph layers, managing forward and backward graphs, and
+    tracking the execution steps.
+    """
+
+    def __init__(self, layer, num_warmup_steps):
+        """
+        Initializes the CUDA graph context.
+        :param layer: The layer to be used in the CUDA graph.
+        :param num_warmup_steps: Number of warmup steps before recording starts.
+        """
+        self.layer = layer
+        self.num_warmup_steps = num_warmup_steps
+
+        # The state of context is in either WARMUP or CUDAGRAPH
+        self._step = 0
+        self.status = CUDAGraphLayerStatus.WARMUP
+
+        # Queue to support 1f1b/interleaved scheduler, assuming FIFO order
+        # data queue
+        self.data_queue = deque()
+
+        # graph queue
+        self.graph_queue = deque()
+
+    # Graph Operations
+    def get_graph(self):
+        if len(self.graph_queue) == 0:
+            return CUDAGraphForwardBackward(self.num_warmup_steps)
+        else:
+            return self.graph_queue.popleft()
+
+    def reuse_graph(self, g):
+        self.graph_queue.append(g)
+
+    # Tensor Queue Operations
+    def push_data(self, args):
+        self.data_queue.append(args)
+
+    def pop_data(self):
+        return self.data_queue.popleft()
+
+    # Finite State Machine of Layer State
+    def warmup_step(self):
+        self._step += 1
+        if self._step == self.num_warmup_steps:
+            self.status = CUDAGraphLayerStatus.CUDAGRAPH
+
+    def is_warmup_step(self):
+        return self.status == CUDAGraphLayerStatus.WARMUP
+
+    def is_cuda_graph_step(self):
+        return self.status == CUDAGraphLayerStatus.CUDAGRAPH
+
+
+def select_y_with_grad(ys, dys):
+    # [TODO] when there is multiple output tensor, we support only one y that allows backward
+    y, dy = None, None
+    if isinstance(ys, paddle.Tensor):
+        y, dy = ys, dys[0]
+    elif isinstance(ys, (list, tuple)):
+        for v, dv in zip(ys, dys):
+            if isinstance(v, paddle.Tensor) and (not v.stop_gradient):
+                y, dy = v, dv
+                break
+    assert isinstance(y, paddle.Tensor) and isinstance(dy, paddle.Tensor)
+    return y, dy
+
+
+# we get the output of the backward from the detached inputs after the backward is calculated
+# we save it to the graph itself
+def get_args_grad(inputs):
+    grad_inputs, detached_grad_inputs = inputs
+    args_grad = []
+    for x, detached_x in zip(grad_inputs, detached_grad_inputs):
+        # if required grad
+        if not x.stop_gradient:
+            if detached_x.grad is None:
+                # if input requires grad but we don't have grad, we just allocate some zeros
+                # x.stop_gradient = True
+                args_grad.append(paddle.zeros(detached_x.shape))
+                # args_grad.append(None)
+            else:
+                args_grad.append(detached_x.grad)
+        else:
+            args_grad.append(None)
+    return tuple(args_grad)
+
+
+class _CUDAGraphedLayer(paddle.autograd.PyLayer):
+    """
+    A custom layer that integrates CUDA Graph recording and execution into PaddlePaddle's autograd system.
+    It handles forward and backward operations differently based on the CUDA graph layer status.
+    """
+
+    @staticmethod
+    def forward(ctx, context, arg_tuple, *grad_inputs):
+        """
+        Handles the forward pass of the layer. It operates differently based on the
+        context's status: warmup, recording, or CUDA graph step.
+        """
+        args, kwargs = arg_tuple
+        # Detach all inputs from the computational graph
+        args = detach(args)
+        kwargs = detach(kwargs)
+
+        detached_grad_inputs = recursive_flatten_args_kwargs(args, kwargs)
+        inputs = (grad_inputs, detached_grad_inputs)
+
+        printer(detached_grad_inputs, "Forward input")
+        if (
+            context.is_warmup_step()
+            or debug_cudagraphedlayer_fallback_to_default
+        ):
+            debug_print("[CUDAGraph] Forward Step (Default)")
+
+            with paddle.enable_grad():
+                y = context.layer(*args, **kwargs)
+
+            context.push_data((CUDAGraphLayerStatus.WARMUP, None, inputs, y))
+        else:
+            graph = context.get_graph()
+            if graph.is_record_step():
+                # In record step, record the forward pass in CUDA graph
+                debug_print(f"[CUDAGraph] Forward Step (Record) id {id(graph)}")
+
+                def forward(*args, **kwargs):
+                    with paddle.enable_grad():
+                        return context.layer(*args, **kwargs)
+
+                y = graph.forward_graph.record(forward, *args, **kwargs)
+
+                context.push_data(
+                    (CUDAGraphLayerStatus.RECORD, graph, inputs, y)
+                )
+            else:
+                debug_print(f"[CUDAGraph] Forward Step (Graph) id {id(graph)}")
+                y = graph.forward_graph.replay(*args, **kwargs)
+
+                context.push_data(
+                    (CUDAGraphLayerStatus.CUDAGRAPH, graph, None, y)
+                )
+
+        debug_print("[CUDAGraph] Forward Step End")
+
+        ctx.save_for_backward(context)
+        printer(y, "Forward output")
+        return detach(y)
+
+    @staticmethod
+    def backward(ctx, *dys):
+        """
+        Handles the backward pass of the layer. Similar to forward, it handles
+        backward based on the context's status: warmup, record, or CUDAGraph.
+        """
+        (context,) = ctx.saved_tensor()
+
+        (status, graph, inputs, ys) = context.pop_data()
+        y, dy = select_y_with_grad(ys, dys)
+
+        printer((y, dy), "Backward input")
+
+        if status == CUDAGraphLayerStatus.WARMUP:
+            debug_print("[CUDAGraph] Backward Step (Default)")
+
+            # In warmup step, perform standard backward operation
+            y.backward(dy)
+            args_grad = get_args_grad(inputs)
+
+            context.warmup_step()
+        elif status == CUDAGraphLayerStatus.RECORD:
+            debug_print(f"[CUDAGraph] Backward Step (Record) id {id(graph)}")
+
+            # In record step, record the backward pass in CUDA graph
+            def backward(y, dy):
+                y.backward(dy)
+
+            graph.backward_graph.record(backward, y, dy)
+
+            # [NOTE] the get_args_grad should not put inside backward
+            # the args_grad should be calculated after graph is replayed
+            args_grad = get_args_grad(inputs)
+            graph.backward_graph.set_output_static(args_grad)
+            graph.capture_end()
+
+            context.reuse_graph(graph)
+        elif status == CUDAGraphLayerStatus.CUDAGRAPH:
+            debug_print(f"[CUDAGraph] Backward Step (Graph) id {id(graph)}")
+
+            # In CUDA graph step, replay the recorded graph for backward pass
+            args_grad = graph.backward_graph.replay(y, dy)
+            context.reuse_graph(graph)
+        else:
+            raise RuntimeError("Unknown cuda graph status")
+
+        debug_print("[CUDAGraph] Backward Step End")
+
+        printer(args_grad, "Backward output")
+        return args_grad
+
+
+class CUDAGraphedLayer(paddle.nn.Layer):
+    """
+    CUDAGraphedLayer: A PaddlePaddle Layer to convert an eager mode model to utilize CUDA Graphs.
+
+    CUDA Graphs provide a way to capture kernel-level operations of a model and play
+    them back efficiently, allowing for potential speedups in repetitive computations,
+    such as those during training iterations. This layer is a wrapper that enables
+    the usage of CUDA Graphs with PaddlePaddle models.
+
+    Overview:
+    - The layer encapsulates another layer (the model to be converted).
+    - During the first few (num_warmup_steps) iterations, the layer operates in
+      eager mode without any CUDA Graphs.
+    - After the warmup steps, the layer captures the forward and backward computations
+      and replays them using CUDA Graphs in subsequent iterations.
+
+    Usage:
+        model = Model()
+        graphed_model = CUDAGraphedLayer(model)
+
+    Parameters:
+    - layer (paddle.nn.Layer): The PaddlePaddle model/layer to be converted.
+    - num_warmup_steps (int): The number of iterations before the CUDA Graph
+      capture begins. Default is 3.
+
+    Notes:
+    - Restrictions:
+        * CPU-GPU Synchronization: Operations that synchronize the CPU with the GPU, like device to host transfers, are not allowed.
+        * CPU Work: Any operations on the CPU within the captured graph are not recorded.
+        * Memory Address (Pointer) Consistency: Replays consistently read from and write to identical virtual memory addresses.
+        * Dynamic Operations:
+            - Control Flow: Dynamic control flows, especially those based on CPU data like if/else statements, are prohibited.
+            - Tensor Shapes: Dynamic tensor shapes are not supported.
+
+    - Allowed Operations:
+        * CUDA RNG Operations: CUDA-based Random Number Generation operations are allowed.
+    """
+
+    def __init__(self, layer: paddle.nn.Layer, num_warmup_steps=3):
+        super().__init__()
+        self.context = CUDAGraphContext(layer, num_warmup_steps)
+        self.add_sublayer(f"Graphed {type(layer).__name__}", layer)
+
+    def forward(self, *args, **kwargs):
+        # We collect them into a list of tensor that required grad
+        grad_inputs = recursive_flatten_args_kwargs(args, kwargs)
+        return _CUDAGraphedLayer.apply(
+            self.context, (args, kwargs), *grad_inputs
+        )
+
+    def is_warmup_step(self):
+        return self.context.is_warmup_step()
+
+    def is_cuda_graph_step(self):
+        return self.context.is_cuda_graph_step()

python/paddle/device/cuda/graphs.py

@@ -0,0 +1,276 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+from __future__ import annotations
+
+import os
+import warnings
+from typing import NoReturn, overload
+
+from paddle.base.core import (
+    CUDAPlace,
+    CustomPlace,
+    XPUPlace,
+    get_all_custom_device_type,
+    is_compiled_with_cuda,
+    is_compiled_with_custom_device,
+    is_compiled_with_rocm,
+    is_compiled_with_xpu,
+)
+
+
+def check_compiled_with_custom_device():
+    custom_device_flag = False
+    custom_devices_types = get_all_custom_device_type()
+    for device_type in custom_devices_types:
+        if is_compiled_with_custom_device(device_type):
+            custom_device_flag = True
+            break
+    return custom_device_flag
+
+
+if (
+    is_compiled_with_cuda()
+    or is_compiled_with_rocm()
+    or check_compiled_with_custom_device()
+    or is_compiled_with_xpu()
+):
+    from paddle.base.core import CUDAGraph as CoreCUDAGraph
+
+    def is_cuda_graph_supported():
+        return True
+else:
+    CoreCUDAGraph = None
+
+    def is_cuda_graph_supported():
+        return False
+
+
+def current_expected_place():
+    for device in get_all_custom_device_type():
+        selected_devices = os.getenv(f"FLAGS_selected_{device}s", "0").split(
+            ","
+        )
+        device_id = int(selected_devices[0])
+        return CustomPlace(device, device_id)
+    return None
+
+
+ALL_MODES = ["global", "thread_local", "relaxed"]
+cuda_graph_id = 0
+
+
+class CUDAGraph:
+    """
+    The native Paddle constructor takes ``place``, ``mode``, ``pool_id`` and
+    ``enable_replace``; the PyTorch-compatible ``keep_graph`` keyword is
+    accepted as well. ``capture_begin`` additionally accepts the PyTorch
+    keywords ``pool`` and ``capture_error_mode`` so the same instance can be
+    driven from either API style.
+    """
+
+    @overload
+    def __init__(self, keep_graph: bool, /) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        place: CUDAPlace | XPUPlace | CustomPlace | None = None,
+        mode: str = "thread_local",
+        pool_id: int | None = None,
+        enable_replace: bool = False,
+        *,
+        keep_graph: bool = False,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        place=None,
+        mode="thread_local",
+        pool_id=None,
+        enable_replace=False,
+        *,
+        keep_graph: bool = False,
+    ):
+        assert CoreCUDAGraph is not None, (
+            "CUDA Graph is only supported on PaddlePaddle compiled with NVIDIA GPU."
+        )
+
+        if isinstance(place, bool):
+            if keep_graph is not False:
+                raise TypeError(
+                    "keep_graph is specified both positionally and by keyword"
+                )
+            keep_graph = place
+            place = None
+
+        self._graph = None
+        if place is None and check_compiled_with_custom_device():
+            place = current_expected_place()
+        elif place is None:
+            if is_compiled_with_cuda():
+                device_id = int(os.environ.get('FLAGS_selected_gpus', 0))
+                place = CUDAPlace(device_id)
+            elif is_compiled_with_xpu():
+                device_id = int(os.environ.get('FLAGS_selected_xpus', 0))
+                place = XPUPlace(device_id)
+            else:
+                raise RuntimeError("Not Supported devices")
+
+        self._place = place
+        assert mode in ALL_MODES
+        self._mode = ALL_MODES.index(mode)
+        self._pool_id = pool_id
+        self._enable_replace = enable_replace
+        self._keep_graph = keep_graph
+        self._debug_mode = False
+
+    def capture_begin(
+        self, pool: int | None = None, capture_error_mode: str | None = None
+    ) -> None:
+        """Begin capturing CUDA work on the current stream.
+
+        Args:
+            pool (int, optional): A memory pool token from
+                :func:`paddle.cuda.graph_pool_handle` or another graph's
+                :meth:`pool`. When provided, this graph shares the indicated
+                memory pool. Overrides ``pool_id`` from the constructor.
+            capture_error_mode (str, optional): One of ``'global'``,
+                ``'thread_local'``, ``'relaxed'`` (see :data:`ALL_MODES`).
+                When ``None`` (default) the constructor's ``mode`` is used;
+                otherwise it overrides the constructor for this capture and a
+                :class:`UserWarning` is emitted to flag the precedence.
+                Invalid values raise :class:`ValueError`.
+        """
+        if pool is not None:
+            self._pool_id = pool
+        elif self._pool_id is None:
+            self._pool_id = CoreCUDAGraph.gen_new_memory_pool_id()
+
+        if capture_error_mode is None:
+            mode = self._mode
+        else:
+            if capture_error_mode not in ALL_MODES:
+                raise ValueError(
+                    f"capture_error_mode must be one of {ALL_MODES}, "
+                    f"but got {capture_error_mode!r}."
+                )
+            mode = ALL_MODES.index(capture_error_mode)
+            if mode != self._mode:
+                warnings.warn(
+                    f"capture_error_mode={capture_error_mode!r} differs from "
+                    f"the constructor mode={ALL_MODES[self._mode]!r}; the "
+                    f"explicit capture_error_mode takes precedence for this "
+                    f"capture.",
+                    stacklevel=2,
+                )
+
+        CoreCUDAGraph.begin_capture_with_pool_id(
+            self._place, mode, self._pool_id, self._enable_replace
+        )
+
+    def capture_end(self):
+        self._graph = CoreCUDAGraph.end_capture()
+
+    def _require_captured(self) -> None:
+        """Raise a clear error if no graph has been captured yet.
+
+        ``self._graph`` is only populated by :meth:`capture_end`; methods that
+        consume it (``replay`` / ``reset`` / ``debug_dump`` / ...) would
+        otherwise raise ``AttributeError`` on ``NoneType`` when called too
+        early. Centralizing the check produces a single, actionable message.
+        """
+        if self._graph is None:
+            raise RuntimeError(
+                "CUDAGraph has not been captured yet. "
+                "Call capture_begin/capture_end first."
+            )
+
+    def instantiate(self) -> CoreCUDAGraph:
+        """Return the instantiated core CUDA graph held by this wrapper.
+
+        Paddle builds the executable graph eagerly inside :meth:`capture_end`,
+        so by the time this is called the graph is already instantiated. It is
+        kept for source compatibility with ``torch.cuda.CUDAGraph.instantiate``
+        and returns the held core :class:`~paddle.base.core.CUDAGraph` produced
+        by :meth:`capture_end`.
+        """
+        self._require_captured()
+        return self._graph
+
+    def replay(self):
+        self._require_captured()
+        self._graph.replay()
+
+    def reset(self):
+        self._require_captured()
+        self._graph.reset()
+
+    def pool(self) -> int:
+        """Return an opaque integer token representing this graph's memory pool.
+
+        The token can be passed as the ``pool`` argument to another graph's
+        :meth:`capture_begin` (or to :class:`paddle.cuda.graph`) so the two
+        graphs share the same memory pool.
+        """
+        if self._pool_id is None:
+            self._pool_id = CoreCUDAGraph.gen_new_memory_pool_id()
+        return self._pool_id
+
+    def enable_debug_mode(self) -> None:
+        """Enable debug mode so that :meth:`debug_dump` is permitted."""
+        self._debug_mode = True
+
+    def debug_dump(self, debug_path) -> None:
+        """Dump the captured graph to ``debug_path`` for inspection.
+
+        :meth:`enable_debug_mode` must be called first.
+        """
+        if not self._debug_mode:
+            raise RuntimeError(
+                "debug_dump requires debug mode to be enabled first. "
+                "Call enable_debug_mode() before debug_dump()."
+            )
+        self._require_captured()
+        self.print_to_dot_files(debug_path)
+
+    def raw_cuda_graph(self) -> NoReturn:
+        """Paddle does not expose the raw ``cudaGraph_t`` handle."""
+        raise NotImplementedError(
+            "raw_cuda_graph is not yet supported in Paddle CUDAGraph. "
+            "The underlying cudaGraph_t handle is not exposed by the Python "
+            "binding."
+        )
+
+    def raw_cuda_graph_exec(self) -> NoReturn:
+        """Paddle does not expose the raw ``cudaGraphExec_t`` handle."""
+        raise NotImplementedError(
+            "raw_cuda_graph_exec is not yet supported in Paddle CUDAGraph. "
+            "The underlying cudaGraphExec_t handle is not exposed by the "
+            "Python binding."
+        )
+
+    def print_to_dot_files(self, dirname, flags=None):
+        if not isinstance(dirname, (str, bytes)):
+            dirname = dirname.name
+        os.makedirs(name=dirname, exist_ok=True)
+        assert os.path.isdir(dirname), (
+            f"The dirname {dirname} should be a directory"
+        )
+        if flags is None:
+            flags = 2047  # only all information. It can be any integer inside [1, 2048)
+        self._graph.print_to_dot_files(dirname, flags)
+
+    def replace_input_ptrs(self, old_ptrs, new_ptrs):
+        self._graph.replace_input_ptrs(old_ptrs, new_ptrs)

python/paddle/device/cuda/memory_analyzer.py

@@ -0,0 +1,540 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+from __future__ import annotations
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+
+# --- Constants ---
+KB = 1024
+MB = 1024 * 1024
+GB = 1024 * 1024 * 1024
+
+
+# --- Formatting Helpers ---
+def format_size(size_bytes):
+    if size_bytes == 0:
+        return "0 B"
+    if size_bytes < MB:
+        return f"{size_bytes / KB:.2f} KB"
+    if size_bytes < GB:
+        return f"{size_bytes / MB:.2f} MB"
+    return f"{size_bytes / GB:.2f} GB"
+
+
+def print_table(title, headers, rows):
+    if not rows:
+        return
+    # Calculate widths
+    col_widths = [len(str(h)) for h in headers]
+    for row in rows:
+        for i, cell in enumerate(row):
+            if i < len(col_widths):
+                col_widths[i] = max(col_widths[i], len(str(cell)))
+    col_widths = [w + 2 for w in col_widths]
+
+    # Build lines
+    row_fmt = "|" + "|".join([f"{{:^{w}}}" for w in col_widths]) + "|"
+    header_sep = "+" + "+".join(["=" * w for w in col_widths]) + "+"
+    inner_sep = "+" + "+".join(["-" * w for w in col_widths]) + "+"
+
+    print(f"\n### {title}")
+    print(header_sep)
+    print(
+        "|" + "|".join([f"{h:^{w}}" for h, w in zip(headers, col_widths)]) + "|"
+    )
+    print(header_sep)
+
+    for i, row in enumerate(rows):
+        print(row_fmt.format(*[str(c) for c in row]))
+        if (
+            title == "Block Size Distribution"
+            and (i + 1) % 2 == 0
+            and i != len(rows) - 1
+        ):
+            print(inner_sep)
+        elif title != "Block Size Distribution":
+            print(inner_sep)
+    if title == "Block Size Distribution":
+        print(header_sep)
+
+
+class MemoryAnalysisTool:
+    def __init__(self):
+        raise TypeError("Utility class should not be instantiated.")
+
+    @classmethod
+    def vmm_max_free_size(
+        self, device_id: int | None = None
+    ) -> tuple[int, int]:
+        name = 'paddle.device.cuda.vmm_max_free_size'
+        if not (core.is_compiled_with_cuda()):
+            raise ValueError(
+                f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
+            )
+        device_id = (
+            device_id
+            if device_id is not None
+            else core.get_cuda_current_device_id()
+        )
+        return core.vmm_max_free_size(device_id)
+
+    @classmethod
+    def vmm_free_block_info(
+        self,
+        device_id: int | None = None,
+    ) -> list[list[tuple[int, int]]]:
+        name = 'paddle.device.cuda.vmm_free_block_info'
+        if not (core.is_compiled_with_cuda()):
+            raise ValueError(
+                f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
+            )
+        device_id = (
+            device_id
+            if device_id is not None
+            else core.get_cuda_current_device_id()
+        )
+        return core.vmm_free_block_info(device_id)
+
+    @classmethod
+    def all_block_info(
+        self,
+        device_id: int | None = None,
+    ) -> list[list[tuple[int, int, bool]]]:
+        name = 'paddle.device.cuda.all_block_info'
+        if not (core.is_compiled_with_cuda()):
+            raise ValueError(
+                f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
+            )
+        device_id = (
+            device_id
+            if device_id is not None
+            else core.get_cuda_current_device_id()
+        )
+        info = core.all_block_info(device_id)
+        return [list(chunk) for chunk in info]
+
+    @classmethod
+    def vmm_all_block_info(
+        self,
+        device_id: int | None = None,
+    ) -> list[list[tuple[int, int, bool]]]:
+        return self.all_block_info(device_id)
+
+    @classmethod
+    def vmm_large_all_block_info(
+        self,
+        device_id: int | None = None,
+    ) -> list[list[tuple[int, int, bool]]]:
+        name = 'paddle.device.cuda.vmm_large_all_block_info'
+        if not (core.is_compiled_with_cuda()):
+            raise ValueError(
+                f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
+            )
+        device_id = (
+            device_id
+            if device_id is not None
+            else core.get_cuda_current_device_id()
+        )
+        return core.large_pool_block_info(device_id)
+
+    @classmethod
+    def vmm_small_all_block_info(
+        self,
+        device_id: int | None = None,
+    ) -> list[list[tuple[int, int, bool]]]:
+        name = 'paddle.device.cuda.vmm_small_all_block_info'
+        if not (core.is_compiled_with_cuda()):
+            raise ValueError(
+                f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
+            )
+        device_id = (
+            device_id
+            if device_id is not None
+            else core.get_cuda_current_device_id()
+        )
+        return core.small_pool_block_info(device_id)
+
+    @classmethod
+    def memory_summary(self, device_id: int | None = None) -> None:
+        device_id = (
+            device_id
+            if device_id is not None
+            else core.get_cuda_current_device_id()
+        )
+        nvidia_smi_AVAILABLE = False
+        try:
+            # import nvidia_smi, pip install nvidia-ml-py3
+            import nvidia_smi
+
+            nvidia_smi_AVAILABLE = True
+        except ImportError:
+            nvidia_smi_AVAILABLE = False
+
+        THRESHOLDS = [
+            1 * MB,
+            10 * MB,
+            50 * MB,
+            100 * MB,
+            200 * MB,
+            400 * MB,
+            600 * MB,
+            800 * MB,
+            1 * GB,
+            2 * GB,
+            3 * GB,
+        ]
+        RANGE_HEADERS = [
+            "[0B,1M)",
+            "[1M,10M)",
+            "[10M,50M)",
+            "[50M,100M)",
+            "[100M,200M)",
+            "[200M,400M)",
+            "[400M,600M)",
+            "[600M,800M)",
+            "[800M,1G)",
+            "[1G,2G)",
+            "[2G,3G)",
+            "[3G,+INF)",
+        ]
+
+        allocator_lists = self.all_block_info(device_id=device_id)
+        # --- Feature 1: Global Summary with NVML & Rates ---
+
+        # 1.1 Get Paddle Stats
+        mem_allocated = paddle.device.cuda.memory_allocated()
+        max_mem_allocated = paddle.device.cuda.max_memory_allocated()
+        mem_reserved = paddle.device.cuda.memory_reserved()
+        max_mem_reserved = paddle.device.cuda.max_memory_reserved()
+
+        # 1.2 Calculate Rates (Utilization of the Reserved Pool)
+        # Rate = How much of the reserved pool is actually holding tensor data?
+        max_alloc_rate = (
+            ((mem_reserved - max_mem_allocated) / mem_reserved)
+            if mem_reserved > 0
+            else 0.0
+        )
+
+        # 1.3 Get Physical Usage via nvidia_smi
+        phy_used_str = "N/A"
+        if nvidia_smi_AVAILABLE:
+            try:
+                nvidia_smi.nvmlInit()
+
+                handle = nvidia_smi.nvmlDeviceGetHandleByIndex(device_id)
+                info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
+                phy_used_str = format_size(info.used)
+                phy_total_str = format_size(info.total)
+                # nvidia_smi.nvmlShutdown() # Optional, depends on lifecycle
+            except Exception as e:
+                phy_used_str = "Err"
+                phy_total_str = "Err"
+        else:
+            print(
+                "Place install nvidia-smi to check real memory usage, pip install command: `pip install nvidia-ml-py3`"
+            )
+            phy_used_str = "No nvidia_smi"
+            phy_total_str = "No nvidia_smi"
+
+        global_headers = [
+            "Allocators",
+            "Allocated",
+            "Max Alloc",
+            "Reserved",
+            "Max Reserved",
+            "Max Frag Rate",
+            "Phy GPU Used / Total",
+        ]
+
+        global_rows = [
+            [
+                len(allocator_lists),
+                format_size(mem_allocated),
+                format_size(max_mem_allocated),
+                format_size(mem_reserved),
+                format_size(max_mem_reserved),
+                f"{max_alloc_rate:.2%}",
+                phy_used_str + ' / ' + phy_total_str,
+            ]
+        ]
+
+        print_table("Global Memory Snapshot", global_headers, global_rows)
+
+        # --- 2. Allocator Analysis ---
+        summary_rows = []
+        dist_rows = []
+
+        for idx, blocks in enumerate(allocator_lists):
+            allocator_name = f"Allocator_{idx}"
+
+            # A. Basic Counting
+            total_blocks = len(blocks)
+            free_blocks = 0
+            total_size = 0
+            free_size = 0
+            max_free_size = 0
+            max_used_size = 0
+            buckets = [[0, 0] for _ in range(len(RANGE_HEADERS))]
+
+            for size, addr, is_free in blocks:
+                total_size += size
+                if is_free:
+                    free_blocks += 1
+                    free_size += size
+                    max_free_size = max(max_free_size, size)
+                else:
+                    max_used_size = max(max_used_size, size)
+
+                # Bucket Mapping
+                b_idx = len(THRESHOLDS)
+                for i, t in enumerate(THRESHOLDS):
+                    if size < t:
+                        b_idx = i
+                        break
+                buckets[b_idx][0 if is_free else 1] += 1
+
+            used_blocks = total_blocks - free_blocks
+            used_size = total_size - free_size
+
+            # B. Summary Row (Total -> Used -> Free)
+            summary_rows.append(
+                [
+                    allocator_name,
+                    total_blocks,
+                    used_blocks,
+                    free_blocks,
+                    format_size(total_size),
+                    format_size(used_size),
+                    format_size(free_size),
+                    format_size(max_used_size),
+                    format_size(max_free_size),
+                ]
+            )
+
+            # D. Distribution Rows
+            dist_rows.append(
+                [allocator_name, "Free Blocks"] + [b[0] for b in buckets]
+            )
+            dist_rows.append(
+                [allocator_name, "Used Blocks"] + [b[1] for b in buckets]
+            )
+
+        # --- 3. Render Outputs ---
+        sum_headers = [
+            "ID",
+            "Tot Blks",
+            "Used Blks",
+            "Free Blks",
+            "Tot Size",
+            "Used Size",
+            "Free Size",
+            "Max Used",
+            "Max Free",
+        ]
+        print_table("Allocator Summary Statistics", sum_headers, summary_rows)
+
+        dist_headers = ["Allocator ID", "Block Type", *RANGE_HEADERS]
+        print_table("Block Size Distribution", dist_headers, dist_rows)
+
+    @classmethod
+    def allocate_record_table(self, data, output_filepath: str = ""):
+        if not data:
+            print("No data to display.")
+            return
+
+        print(f"Record data size: {len(data)}, start printing...")
+        headers = [
+            'Allocator_Instance',
+            'Is_Allocate',
+            'Seq_ID',
+            'Req_Size',
+            'Cur_Alloc',
+            'Cur_Rsrv',
+        ]
+        formatted_row = []
+        all_lines = []
+        all_lines.append("\t".join(headers))
+        for row in data:
+            formatted_row = [
+                str(row[0]),
+                "Allocate" if row[1] else "Free",
+                str(row[2]),
+                str(row[3]),
+                str(row[4]),
+                str(row[5]),
+            ]
+            line = "\t".join(formatted_row)
+            all_lines.append(line)
+
+        try:
+            with open(output_filepath, 'w', encoding='utf-8') as f:
+                f.write("\n".join(all_lines))
+            print(f"Data successfully written to: {output_filepath}")
+        except OSError as e:
+            print(f"Error writing to file {output_filepath}: {e}")
+
+    @classmethod
+    def allocate_record_plot(self, data, save_path: str = ""):
+        try:
+            import matplotlib.pyplot as plt
+            from matplotlib import ticker
+        except ImportError:
+            raise ImportError(
+                "matplotlib is required but not installed. Please install it using: pip install matplotlib"
+            )
+
+        if not data:
+            print("No data to plot.")
+            return
+        print(f"Record data size: {len(data)}, start plotting...")
+        data_np = np.array(data)
+        is_allocate = data_np[:, 1]
+        filter_mask = is_allocate == 1
+        data_np = data_np[filter_mask]
+        allocator_instance = data_np[:, 0]  # allocator_instance not used
+        ids = data_np[:, 2]
+        sizes = data_np[:, 3]
+        allocated = data_np[:, 4]
+        reserved = data_np[:, 5]
+
+        LOG_START_VALUE = 1
+        plt.style.use('seaborn-v0_8-whitegrid')
+        fig, (ax1, ax2) = plt.subplots(
+            2,
+            1,
+            sharex=True,
+            figsize=(16, 10),
+            dpi=120,
+            gridspec_kw={'height_ratios': [3, 1], 'hspace': 0},
+        )
+
+        # allocated event plot
+        ax1.plot(
+            ids, sizes, color='#2ca02c', linestyle='-', linewidth=1, alpha=0.3
+        )
+        ax1.scatter(
+            ids,
+            sizes,
+            color='#2ca02c',
+            s=60,
+            alpha=1.0,
+            edgecolors='white',
+            linewidth=0.5,
+            label='Request Size',
+            zorder=5,
+        )
+
+        ax1.set_ylabel(
+            'Request Size (Linear Scale)',
+            fontsize=12,
+            fontweight='bold',
+            labelpad=10,
+        )
+        ax1.set_title(
+            'Paddle GPU Memory Allocation Analysis',
+            fontsize=16,
+            fontweight='bold',
+            pad=20,
+        )
+
+        ax1.set_ylim(bottom=LOG_START_VALUE)
+        ax1.tick_params(axis='x', length=0)
+        plt.setp(ax1.get_xticklabels(), visible=False)
+
+        # memory allocated, reserved plot
+        ax2.plot(
+            ids,
+            reserved,
+            color='#d62728',
+            linestyle='--',
+            linewidth=1.5,
+            alpha=0.8,
+            label='Reserved (Pool)',
+        )
+        ax2.fill_between(ids, 0, reserved, color='#d62728', alpha=0.1)
+        ax2.plot(
+            ids,
+            allocated,
+            color='#1f77b4',
+            linestyle='-',
+            linewidth=2,
+            alpha=0.9,
+            label='Allocated (Used)',
+        )
+        ax2.fill_between(ids, 0, allocated, color='#1f77b4', alpha=0.15)
+
+        ax2.invert_yaxis()
+        ax2.set_ylim(reserved.max() * 3.0, LOG_START_VALUE)
+        # ax2.set_yscale('symlog', linthresh=1024 * 1024)
+        ax2.set_ylabel(
+            'Pool Status (Inverted)',
+            fontsize=11,
+            fontweight='bold',
+            labelpad=10,
+        )
+
+        ax2.set_xlabel('')
+        ax2.tick_params(axis='x', which='both', length=0)
+        plt.setp(ax2.get_xticklabels(), visible=False)
+
+        # y axis setting 0
+        def y_axis_formatter(x, pos):
+            val = abs(x)
+            if val <= LOG_START_VALUE * 1.5:
+                return '0'
+            return format_size(val).replace(" ", "")
+
+        formatter = ticker.FuncFormatter(y_axis_formatter)
+        ax1.yaxis.set_major_formatter(formatter)
+        ax2.yaxis.set_major_formatter(formatter)
+
+        for ax in [ax1, ax2]:
+            current_ticks = ax.get_yticks().tolist()
+            if LOG_START_VALUE not in current_ticks:
+                current_ticks.append(LOG_START_VALUE)
+            ax.set_yticks(sorted(current_ticks))
+
+        # axis setting
+        for ax in [ax1, ax2]:
+            for spine in ax.spines.values():
+                spine.set_edgecolor('black')
+                spine.set_linewidth(1.5)
+            ax.tick_params(
+                axis='both', which='major', colors='black', width=1.0, length=5
+            )
+
+        lines1, labels1 = ax1.get_legend_handles_labels()
+        lines2, labels2 = ax2.get_legend_handles_labels()
+
+        ax1.legend(
+            lines1 + lines2,
+            labels1 + labels2,
+            loc='upper right',
+            fontsize=10,
+            frameon=True,
+            facecolor='white',
+            framealpha=0.9,
+            edgecolor='black',
+            shadow=False,
+        )
+
+        plt.tight_layout()
+        plt.subplots_adjust(hspace=0.05)
+
+        plt.savefig(save_path)
+        plt.close()
+        print(f"Analysis plot saved to: {save_path}")

python/paddle/device/cuda/streams.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+from __future__ import annotations
+
+from paddle.base.core import CUDAEvent as Event, CUDAPlace, CUDAStream as Stream
+
+
+def create_stream(
+    device_id: CUDAPlace | int | None = None,
+    priority: int = 2,
+    device_type: str | None = None,  # Ignored for compatibility
+    blocking: bool = False,  # Ignored for compatibility
+):
+    """
+    Factory Function, used to create CUDA Stream
+    """
+    return Stream(device_id, priority)
+
+
+def create_event(
+    enable_timing: bool = False,
+    blocking: bool = False,
+    interprocess: bool = False,
+    device_type: str | None = None,
+    device_id: int = 0,
+):
+    """
+    Factory Function, used to create CUDA Event
+    """
+    return Event(enable_timing, blocking, interprocess)