chore: import upstream snapshot with attribution

python/sglang/srt/model_executor/runner_utils/__init__.py

@@ -0,0 +1,28 @@
+"""Low-level utilities used by the CUDA graph runners.
+
+Mirror of cuda_graph_backend_utils/ for runner-side state — buffer
+dataclasses, process-global capture flags, the speculative-shared
+graph memory pool, and the DeepEP capture/replay adapter. Runners in
+cuda_graph_runner/ import from here; nothing here should import
+back into cuda_graph_runner/.
+"""
+
+from sglang.srt.model_executor.runner_utils.buffers import (  # noqa: F401
+    DecodeInputBuffers,
+    PrefillInputBuffers,
+    _grouped_foreach_copy_,
+)
+from sglang.srt.model_executor.runner_utils.capture_mode import (  # noqa: F401
+    _set_capture_lora_variant,
+    compile_in_capture_mode,
+    get_capture_lora_variant,
+    get_is_capture_mode,
+    model_capture_mode,
+)
+from sglang.srt.model_executor.runner_utils.deepep_adapter import (  # noqa: F401
+    DeepEPCudaGraphRunnerAdapter,
+)
+from sglang.srt.model_executor.runner_utils.pool import (  # noqa: F401
+    get_global_graph_memory_pool,
+    set_global_graph_memory_pool,
+)

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -0,0 +1,435 @@
+# 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.
+# ==============================================================================
+"""Static-buffer dataclasses used by the CUDA graph runners.
+
+DecodeInputBuffers backs the decode-phase capture/replay path.
+PrefillInputBuffers backs the prefill-phase capture/replay path.
+
+Both subclass ForwardInputBuffers so that buffer-pool sharing works
+the same way as for non-cuda-graph forward paths.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_batch_info import (
+    ForwardBatch,
+    NgramEmbeddingInfo,
+    PPProxyTensors,
+    compute_local_num_token_non_padded,
+)
+from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+
+_has_foreach_copy = hasattr(torch, "_foreach_copy_")
+
+
+def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+    """Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs."""
+
+    def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+        if _has_foreach_copy:
+            torch._foreach_copy_(dsts, srcs)
+        else:
+            for dst, src in zip(dsts, srcs):
+                dst.copy_(src)
+
+    groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {}
+    for dst, src in zip(dsts, srcs):
+        key = (dst.dtype, src.dtype)
+        if key not in groups:
+            groups[key] = ([], [])
+        groups[key][0].append(dst)
+        groups[key][1].append(src)
+    for group_dsts, group_srcs in groups.values():
+        foreach_copy(group_dsts, group_srcs)
+
+
+@dataclass
+class DecodeInputBuffers(ForwardInputBuffers):
+    input_ids: torch.Tensor
+    input_embeds: torch.Tensor
+    req_pool_indices: torch.Tensor
+    seq_lens: torch.Tensor
+    seq_lens_cpu: torch.Tensor
+    out_cache_loc: torch.Tensor
+    positions: torch.Tensor
+    mrope_positions: torch.Tensor
+    num_token_non_padded: torch.Tensor
+    custom_mask: torch.Tensor
+    next_token_logits_buffer: torch.Tensor
+    mamba_track_indices: Optional[torch.Tensor]
+    mamba_track_mask: Optional[torch.Tensor]
+    global_num_tokens_gpu: torch.Tensor
+    global_num_tokens_for_logprob_gpu: torch.Tensor
+    encoder_lens: Optional[torch.Tensor]
+    pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
+    ngram_embedding_info: Optional[NgramEmbeddingInfo]
+    rids_int: Optional[torch.Tensor]
+    bootstrap_room_ids_int: Optional[torch.Tensor]
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        device: torch.device,
+        max_bs: int,
+        max_num_token: int,
+        hidden_size: int,
+        next_token_logits_buffer: torch.Tensor,
+        dtype: torch.dtype,
+        dp_size: int,
+        pp_size: int,
+        is_encoder_decoder: bool,
+        require_mlp_tp_gather: bool,
+        seq_len_fill_value: int,
+        encoder_len_fill_value: int,
+        num_tokens_per_bs: int,
+        cache_loc_dtype: torch.dtype,
+        enable_mamba_track: bool,
+        ne_token_table: Optional[torch.Tensor] = None,
+        hc_hidden_size: Optional[int] = None,
+        pp_proxy_topk_size: Optional[int] = None,
+    ) -> DecodeInputBuffers:
+        with torch.device(device):
+            input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
+            input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
+            req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
+            seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int64)
+            out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
+            positions = torch.zeros((max_num_token,), dtype=torch.int64)
+            mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
+            num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
+            custom_mask = torch.ones(
+                (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
+                dtype=torch.bool,
+            )
+            mamba_track_indices = (
+                torch.zeros((max_bs,), dtype=torch.int64)
+                if enable_mamba_track
+                else None
+            )
+            mamba_track_mask = (
+                torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
+            )
+
+            if pp_size > 1:
+                is_mhc = hc_hidden_size is not None
+                hs = hc_hidden_size if is_mhc else hidden_size
+                pp_proxy_tensors = {
+                    "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
+                }
+                if not is_mhc:
+                    pp_proxy_tensors["residual"] = torch.zeros(
+                        (max_bs, hidden_size), dtype=dtype
+                    )
+                if pp_proxy_topk_size is not None:
+                    pp_proxy_tensors["topk_indices"] = torch.zeros(
+                        (max_num_token, pp_proxy_topk_size), dtype=torch.int32
+                    )
+            else:
+                pp_proxy_tensors = None
+
+            if is_encoder_decoder:
+                encoder_lens = torch.full(
+                    (max_bs,), encoder_len_fill_value, dtype=torch.int32
+                )
+            else:
+                encoder_lens = None
+
+            if require_mlp_tp_gather:
+                global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
+                global_num_tokens_for_logprob_gpu = torch.zeros(
+                    (dp_size,), dtype=torch.int32
+                )
+            else:
+                global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+                global_num_tokens_for_logprob_gpu = torch.zeros((1,), dtype=torch.int32)
+
+            ngram_embedding_info = (
+                NgramEmbeddingInfo(
+                    token_table=ne_token_table,
+                    column_starts=torch.zeros([max_bs], dtype=torch.int32),
+                    req_lens=torch.ones([max_bs], dtype=torch.int32),
+                    out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
+                    out_req_lens=torch.ones([max_bs], dtype=torch.int32),
+                    skip_token_table_update=torch.zeros([max_bs], dtype=torch.bool),
+                )
+                if ne_token_table is not None
+                else None
+            )
+
+            if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
+                rids_int = torch.zeros((max_bs,), dtype=torch.int64)
+                bootstrap_room_ids_int = torch.full((max_bs,), -1, dtype=torch.int64)
+            else:
+                rids_int = None
+                bootstrap_room_ids_int = None
+
+        seq_lens_cpu = torch.full(
+            (max_bs,),
+            seq_len_fill_value,
+            dtype=torch.int64,
+            device="cpu",
+        )
+
+        return cls(
+            input_ids=input_ids,
+            input_embeds=input_embeds,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            out_cache_loc=out_cache_loc,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            num_token_non_padded=num_token_non_padded,
+            custom_mask=custom_mask,
+            next_token_logits_buffer=next_token_logits_buffer,
+            mamba_track_indices=mamba_track_indices,
+            mamba_track_mask=mamba_track_mask,
+            encoder_lens=encoder_lens,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+            pp_proxy_tensors=pp_proxy_tensors,
+            ngram_embedding_info=ngram_embedding_info,
+            rids_int=rids_int,
+            bootstrap_room_ids_int=bootstrap_room_ids_int,
+        )
+
+    def populate_from_forward_batch(
+        self,
+        *,
+        forward_batch: ForwardBatch,
+        raw_bs: int,
+        raw_num_token: int,
+        bs: int,
+        seq_len_fill_value: int,
+        require_gathered_buffer: bool,
+        num_tokens_per_bs: int,
+        dsa_enable_prefill_cp: bool,
+        enable_num_token_non_padded_flag: bool,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+    ):
+        if bs != raw_bs:
+            self.seq_lens.fill_(seq_len_fill_value)
+            self.out_cache_loc.zero_()
+            if self.mamba_track_indices is not None:
+                self.mamba_track_indices.zero_()
+            if self.mamba_track_mask is not None:
+                self.mamba_track_mask.fill_(False)
+
+        # Build batched copy lists for all GPU tensors.
+        dsts = [
+            self.input_ids[:raw_num_token],
+            self.req_pool_indices[:raw_bs],
+            self.seq_lens[:raw_bs],
+            self.out_cache_loc[:raw_num_token],
+            self.positions[:raw_num_token],
+        ]
+        srcs = [
+            forward_batch.input_ids,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            forward_batch.out_cache_loc,
+            forward_batch.positions,
+        ]
+
+        if self.ngram_embedding_info is not None:
+            ngram_embedding_info = forward_batch.ngram_embedding_info
+            self.ngram_embedding_info.column_starts[:raw_bs].copy_(
+                ngram_embedding_info.column_starts
+            )
+            self.ngram_embedding_info.req_lens[:raw_bs].copy_(
+                ngram_embedding_info.req_lens
+            )
+
+        if (
+            self.mamba_track_indices is not None
+            and forward_batch.mamba_track_indices is not None
+        ):
+            dsts.append(self.mamba_track_indices[:raw_bs])
+            srcs.append(forward_batch.mamba_track_indices)
+        if (
+            self.mamba_track_mask is not None
+            and forward_batch.mamba_track_mask is not None
+        ):
+            dsts.append(self.mamba_track_mask[:raw_bs])
+            srcs.append(forward_batch.mamba_track_mask)
+
+        if self.encoder_lens is not None and forward_batch.encoder_lens is not None:
+            dsts.append(self.encoder_lens[:raw_bs])
+            srcs.append(forward_batch.encoder_lens)
+
+        if forward_batch.mrope_positions is not None:
+            dsts.append(self.mrope_positions[:, :raw_num_token])
+            srcs.append(forward_batch.mrope_positions)
+
+        if self.rids_int is not None and forward_batch.rids_int is not None:
+            dsts.append(self.rids_int[:raw_bs])
+            srcs.append(forward_batch.rids_int)
+        if (
+            self.bootstrap_room_ids_int is not None
+            and forward_batch.bootstrap_room_ids_int is not None
+        ):
+            dsts.append(self.bootstrap_room_ids_int[:raw_bs])
+            srcs.append(forward_batch.bootstrap_room_ids_int)
+
+        if require_gathered_buffer:
+            self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs)
+            self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)
+
+        if enable_num_token_non_padded_flag:
+            if require_gathered_buffer and not dsa_enable_prefill_cp:
+                num_tokens_per_dp = bs * num_tokens_per_bs
+                local = compute_local_num_token_non_padded(
+                    global_num_token_non_padded=forward_batch.num_token_non_padded,
+                    num_tokens_per_dp=num_tokens_per_dp,
+                )
+                dsts.append(self.num_token_non_padded)
+                srcs.append(local)
+            else:
+                dsts.append(self.num_token_non_padded)
+                srcs.append(forward_batch.num_token_non_padded)
+
+        # Pipeline-parallel proxy tensors.
+        if pp_proxy_tensors is not None and self.pp_proxy_tensors is not None:
+            for key, buf in self.pp_proxy_tensors.items():
+                src = pp_proxy_tensors.tensors[key]
+                dim = src.shape[0]
+                dsts.append(buf[:dim])
+                srcs.append(src)
+
+        # Batch all GPU copies, grouped by dtype pair.
+        _grouped_foreach_copy_(dsts, srcs)
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                self.seq_lens_cpu.fill_(seq_len_fill_value)
+            self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+
+
+@dataclass
+class PrefillInputBuffers(ForwardInputBuffers):
+    input_ids: torch.Tensor
+    out_cache_loc: torch.Tensor
+    num_token_non_padded: torch.Tensor
+    mamba_track_indices: Optional[torch.Tensor]
+    mamba_track_mask: Optional[torch.Tensor]
+    mamba_track_seqlens: Optional[torch.Tensor]
+    positions: torch.Tensor
+    input_embeds: Optional[torch.Tensor]
+    mrope_positions: Optional[torch.Tensor]
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        device: torch.device,
+        max_bs: int,
+        max_num_tokens: int,
+        cache_loc_dtype: torch.dtype,
+        is_multimodal: bool,
+        hidden_size: int,
+        dtype: torch.dtype,
+        enable_mamba_track: bool,
+    ) -> PrefillInputBuffers:
+        with torch.device(device):
+            input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
+            out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype)
+            num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
+            mamba_track_indices = (
+                torch.zeros((max_bs,), dtype=torch.int64)
+                if enable_mamba_track
+                else None
+            )
+            mamba_track_mask = (
+                torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
+            )
+            mamba_track_seqlens = (
+                torch.zeros((max_bs,), dtype=torch.int32)
+                if enable_mamba_track
+                else None
+            )
+            positions = torch.zeros((max_num_tokens,), dtype=torch.int64)
+
+            if is_multimodal:
+                input_embeds = torch.zeros((max_num_tokens, hidden_size), dtype=dtype)
+                mrope_positions = torch.zeros((3, max_num_tokens), dtype=torch.int64)
+            else:
+                input_embeds = None
+                mrope_positions = None
+
+        return cls(
+            input_ids=input_ids,
+            out_cache_loc=out_cache_loc,
+            num_token_non_padded=num_token_non_padded,
+            mamba_track_indices=mamba_track_indices,
+            mamba_track_mask=mamba_track_mask,
+            mamba_track_seqlens=mamba_track_seqlens,
+            positions=positions,
+            input_embeds=input_embeds,
+            mrope_positions=mrope_positions,
+        )
+
+    def populate_from_forward_batch(
+        self,
+        *,
+        forward_batch: ForwardBatch,
+        raw_num_tokens: int,
+        static_num_tokens: int,
+        is_multimodal: bool,
+    ) -> None:
+        """Copy serving-batch values into static buffers and zero out
+        the padding region between raw_num_tokens and
+        static_num_tokens.
+        """
+        if static_num_tokens != raw_num_tokens:
+            self.out_cache_loc.zero_()
+            self.input_ids[raw_num_tokens:static_num_tokens].zero_()
+            self.positions[raw_num_tokens:static_num_tokens].zero_()
+            if is_multimodal:
+                self.input_embeds[raw_num_tokens:static_num_tokens].zero_()
+            if forward_batch.mrope_positions is not None:
+                self.mrope_positions[:, raw_num_tokens:static_num_tokens].zero_()
+
+        bs = forward_batch.batch_size
+
+        self.input_ids[:raw_num_tokens].copy_(forward_batch.input_ids)
+        self.positions[:raw_num_tokens].copy_(forward_batch.positions)
+        self.out_cache_loc[:raw_num_tokens].copy_(forward_batch.out_cache_loc)
+
+        if (
+            self.mamba_track_indices is not None
+            and forward_batch.mamba_track_indices is not None
+        ):
+            self.mamba_track_indices[:bs].copy_(forward_batch.mamba_track_indices)
+        if (
+            self.mamba_track_mask is not None
+            and forward_batch.mamba_track_mask is not None
+        ):
+            self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask)
+        if (
+            self.mamba_track_seqlens is not None
+            and forward_batch.mamba_track_seqlens is not None
+        ):
+            self.mamba_track_seqlens[:bs].copy_(forward_batch.mamba_track_seqlens)
+
+        if forward_batch.mrope_positions is not None:
+            self.mrope_positions[:, :raw_num_tokens].copy_(
+                forward_batch.mrope_positions
+            )

python/sglang/srt/model_executor/runner_utils/capture_mode.py

@@ -0,0 +1,78 @@
+# 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.
+# ==============================================================================
+"""Process-global capture-mode flags shared by the decode runner and the
+speculative-draft runners. Read by model code that needs to take a
+capture-time branch (e.g. lora dual-graph capture decides per-batch
+which variant to use).
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Optional
+
+import torch
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
+
+# Detect whether the current forward pass is in capture mode.
+is_capture_mode = False
+
+# When capturing dual MoE backends, tracks which variant is being captured.
+# None = not dual, "lora" = capturing lora variant, "nolora" = capturing nolora variant.
+_capture_lora_variant: Optional[str] = None
+
+
+def get_is_capture_mode() -> bool:
+    return is_capture_mode or is_in_breakable_cuda_graph()
+
+
+def compile_in_capture_mode(func):
+    """Decorator: wrap func with torch.compile only when defined
+    inside model capture mode; passthrough otherwise.
+
+    Used by model code (e.g. DeepSeek-V4) to opt nested helpers into
+    torch.compile during cuda-graph capture without paying the
+    compilation cost in the eager forward path.
+    """
+    if is_capture_mode:
+        return torch.compile(func)
+    return func
+
+
+def get_capture_lora_variant() -> Optional[str]:
+    """Return the lora variant being captured, or None if not in dual capture."""
+    return _capture_lora_variant
+
+
+def _set_capture_lora_variant(variant: Optional[str]) -> None:
+    global _capture_lora_variant
+    _capture_lora_variant = variant
+
+
+@contextmanager
+def model_capture_mode():
+    global is_capture_mode
+    from sglang.srt.runtime_context import get_flags
+
+    # Disable dispose_tensor() during capture: freeing mid-capture records data_ptr()==0 into the graph.
+    is_capture_mode = True
+    get_flags().capture.disable_dispose_tensor = True
+    try:
+        yield
+    finally:
+        is_capture_mode = False
+        get_flags().capture.disable_dispose_tensor = False

python/sglang/srt/model_executor/runner_utils/deepep_adapter.py

@@ -0,0 +1,42 @@
+# 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.
+# ==============================================================================
+"""DeepEP capture/replay adapter — records the dispatch mode used during
+capture and re-applies it during replay so DeepEP all-to-all has
+consistent expert routing across the captured graph.
+"""
+
+from __future__ import annotations
+
+from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
+from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend
+
+
+class DeepEPCudaGraphRunnerAdapter:
+    def __init__(self) -> None:
+        # Record DeepEP mode used during capture to ensure replay consistency.
+        self._captured_deepep_mode = None
+
+    def capture(self, is_extend_in_batch: bool) -> None:
+        if not get_moe_a2a_backend().is_deepep():
+            return
+        self._captured_deepep_mode = get_deepep_mode().resolve(
+            is_extend_in_batch=is_extend_in_batch
+        )
+        DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)
+
+    def replay(self) -> None:
+        if not get_moe_a2a_backend().is_deepep():
+            return
+        assert self._captured_deepep_mode is not None
+        DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -0,0 +1,40 @@
+# 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.
+# ==============================================================================
+"""Process-wide CUDA graph memory pool shared across the prefill and
+decode graph backends. The two phases never replay concurrently, so
+sharing one pool reserves only the larger phase's capture footprint.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from sglang.srt.runtime_context import get_resources
+
+
+def get_global_graph_memory_pool() -> Optional[Any]:
+    return get_resources().graph_memory_pool
+
+
+def set_global_graph_memory_pool(val: Any) -> None:
+    get_resources().graph_memory_pool = val
+
+
+def get_or_create_global_graph_memory_pool(device_module: Any) -> Any:
+    """Return the shared graph memory pool, creating it on first use so
+    later backends reuse the same handle."""
+    resources = get_resources()
+    if resources.graph_memory_pool is None:
+        resources.graph_memory_pool = device_module.graph_pool_handle()
+    return resources.graph_memory_pool