sglang 0.4.10__tar.gz → 0.4.10.post2__tar.gz

{sglang-0.4.10/sglang.egg-info → sglang-0.4.10.post2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sglang
-Version: 0.4.10
+Version: 0.4.10.post2
 Summary: SGLang is yet another fast serving framework for large language models and vision language models.
 License:                                  Apache License
                                    Version 2.0, January 2004
@@ -250,7 +250,7 @@ Requires-Dist: transformers==4.54.1; extra == "runtime-common"
 Requires-Dist: timm==1.0.16; extra == "runtime-common"
 Requires-Dist: uvicorn; extra == "runtime-common"
 Requires-Dist: uvloop; extra == "runtime-common"
-Requires-Dist: xgrammar==0.1.21; extra == "runtime-common"
+Requires-Dist: xgrammar==0.1.22; extra == "runtime-common"
 Provides-Extra: srt
 Requires-Dist: sglang[runtime_common]; extra == "srt"
 Requires-Dist: sgl-kernel==0.2.8; extra == "srt"
@@ -301,6 +301,7 @@ Requires-Dist: matplotlib; extra == "test"
 Requires-Dist: pandas; extra == "test"
 Requires-Dist: peft; extra == "test"
 Requires-Dist: sentence_transformers; extra == "test"
+Requires-Dist: pytest; extra == "test"
 Provides-Extra: all
 Requires-Dist: sglang[srt]; extra == "all"
 Requires-Dist: sglang[openai]; extra == "all"

{sglang-0.4.10 → sglang-0.4.10.post2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sglang"
-version = "0.4.10"
+version = "0.4.10.post2"
 description = "SGLang is yet another fast serving framework for large language models and vision language models."
 readme = "README.md"
 requires-python = ">=3.8"
@@ -49,7 +49,7 @@ runtime_common = [
     "timm==1.0.16",
     "uvicorn",
     "uvloop",
-    "xgrammar==0.1.21",
+    "xgrammar==0.1.22",
 ]
 
 srt = [
@@ -108,6 +108,7 @@ test = [
     "pandas",
     "peft",
     "sentence_transformers",
+    "pytest",
 ]
 all = ["sglang[srt]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]", "sglang[torch_memory_saver]", "sglang[decord]"]
 all_hip = ["sglang[srt_hip]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]", "sglang[decord]"]
@@ -130,6 +131,7 @@ dev_cpu = ["sglang[all_cpu]", "sglang[test]"]
 "sglang" = [
     "srt/layers/moe/fused_moe_triton/configs/*/*.json",
     "srt/layers/quantization/configs/*.json",
+    "srt/mem_cache/storage/hf3fs/hf3fs_utils.cpp",
 ]
 
 [tool.setuptools.packages.find]

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/bench_offline_throughput.py

@@ -418,6 +418,26 @@ if __name__ == "__main__":
     ServerArgs.add_cli_args(parser)
     BenchArgs.add_cli_args(parser)
     args = parser.parse_args()
+
+    # handling ModelScope model downloads
+    if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() in ("true", "1"):
+        if os.path.exists(args.model_path):
+            print(f"Using local model path: {args.model_path}")
+        else:
+            try:
+                from modelscope import snapshot_download
+
+                print(f"Using ModelScope to download model: {args.model_path}")
+
+                # download the model and replace args.model_path
+                args.model_path = snapshot_download(
+                    args.model_path,
+                )
+                print(f"Model downloaded to: {args.model_path}")
+            except Exception as e:
+                print(f"ModelScope download failed: {str(e)}")
+                raise e
+
     server_args = ServerArgs.from_cli_args(args)
     bench_args = BenchArgs.from_cli_args(args)
 

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/compile_deep_gemm.py

@@ -17,6 +17,7 @@ import time
 
 import requests
 
+from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST
 from sglang.srt.entrypoints.http_server import launch_server
 from sglang.srt.managers.io_struct import GenerateReqInput
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
@@ -52,7 +53,9 @@ class CompileArgs:
 
 
 @warmup("compile-deep-gemm")
-async def warm_up_compile(tokenizer_manager: TokenizerManager):
+async def warm_up_compile(
+    disaggregation_mode: str, tokenizer_manager: TokenizerManager
+):
     print("\nGenerate warm up request for compiling DeepGEMM...\n")
     generate_req_input = GenerateReqInput(
         input_ids=[0, 1, 2, 3],
@@ -62,6 +65,10 @@ async def warm_up_compile(tokenizer_manager: TokenizerManager):
             "ignore_eos": True,
         },
     )
+    if disaggregation_mode != "null":
+        generate_req_input.bootstrap_room = 0
+        generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST
+
     await tokenizer_manager.generate_request(generate_req_input, None).__anext__()
 
 

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/global_config.py

@@ -30,7 +30,11 @@ class GlobalConfig:
         self.default_new_token_ratio_decay_steps = float(
             os.environ.get("SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS", 600)
         )
-
+        self.torch_empty_cache_interval = float(
+            os.environ.get(
+                "SGLANG_EMPTY_CACHE_INTERVAL", -1
+            )  # in seconds. Set if you observe high memory accumulation over a long serving period.
+        )
         # Runtime constants: others
         self.retract_decode_steps = 20
         self.flashinfer_workspace_size = os.environ.get(

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/srt/configs/model_config.py

@@ -112,6 +112,7 @@ class ModelConfig:
             mm_disabled_models = [
                 "Gemma3ForConditionalGeneration",
                 "Llama4ForConditionalGeneration",
+                "Step3VLForConditionalGeneration",
             ]
             if self.hf_config.architectures[0] in mm_disabled_models:
                 enable_multimodal = False

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/srt/conversation.py

@@ -954,20 +954,6 @@ register_conv_template(
     )
 )
 
-register_conv_template(
-    Conversation(
-        name="mimo-vl",
-        system_message="You are MiMo, an AI assistant developed by Xiaomi.",
-        system_template="<|im_start|>system\n{system_message}",
-        roles=("<|im_start|>user", "<|im_start|>assistant"),
-        sep="<|im_end|>\n",
-        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
-        stop_str=["<|im_end|>"],
-        image_token="<|vision_start|><|image_pad|><|vision_end|>",
-    )
-)
-
-
 register_conv_template(
     Conversation(
         name="qwen2-audio",
@@ -981,51 +967,11 @@ register_conv_template(
     )
 )
 
-register_conv_template(
-    Conversation(
-        name="llama_4_vision",
-        system_message="You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.",
-        system_template="<|header_start|>system<|header_end|>\n\n{system_message}<|eot|>",
-        roles=("user", "assistant"),
-        sep_style=SeparatorStyle.LLAMA4,
-        sep="",
-        stop_str="<|eot|>",
-        image_token="<|image|>",
-    )
-)
-
-register_conv_template(
-    Conversation(
-        name="step3-vl",
-        system_message="<｜begin▁of▁sentence｜>You are a helpful assistant",
-        system_template="{system_message}\n",
-        roles=(
-            "<|BOT|>user\n",
-            "<|BOT|>assistant\n<think>\n",
-        ),
-        sep="<|EOT|>",
-        sep_style=SeparatorStyle.NO_COLON_SINGLE,
-        stop_str="<|EOT|>",
-        image_token="<im_patch>",
-        # add_bos=True,
-    )
-)
-
 
 @register_conv_template_matching_function
 def match_internvl(model_path: str):
     if re.search(r"internvl", model_path, re.IGNORECASE):
         return "internvl-2-5"
-    if re.search(r"intern.*s1", model_path, re.IGNORECASE):
-        return "interns1"
-
-
-@register_conv_template_matching_function
-def match_llama_vision(model_path: str):
-    if re.search(r"llama.*3\.2.*vision", model_path, re.IGNORECASE):
-        return "llama_3_vision"
-    if re.search(r"llama.*4.*", model_path, re.IGNORECASE):
-        return "llama_4_vision"
 
 
 @register_conv_template_matching_function
@@ -1040,22 +986,6 @@ def match_vicuna(model_path: str):
         return "vicuna_v1.1"
 
 
-@register_conv_template_matching_function
-def match_llama2_chat(model_path: str):
-    if re.search(
-        r"llama-2.*chat|codellama.*instruct",
-        model_path,
-        re.IGNORECASE,
-    ):
-        return "llama-2"
-
-
-@register_conv_template_matching_function
-def match_mistral(model_path: str):
-    if re.search(r"pixtral|(mistral|mixtral).*instruct", model_path, re.IGNORECASE):
-        return "mistral"
-
-
 @register_conv_template_matching_function
 def match_deepseek_vl(model_path: str):
     if re.search(r"deepseek.*vl2", model_path, re.IGNORECASE):
@@ -1064,12 +994,6 @@ def match_deepseek_vl(model_path: str):
 
 @register_conv_template_matching_function
 def match_qwen_chat_ml(model_path: str):
-    if re.search(r"gme.*qwen.*vl", model_path, re.IGNORECASE):
-        return "gme-qwen2-vl"
-    if re.search(r"qwen.*vl", model_path, re.IGNORECASE):
-        return "qwen2-vl"
-    if re.search(r"qwen.*audio", model_path, re.IGNORECASE):
-        return "qwen2-audio"
     if re.search(
         r"llava-v1\.6-34b|llava-v1\.6-yi-34b|llava-next-video-34b|llava-onevision-qwen2",
         model_path,
@@ -1078,12 +1002,6 @@ def match_qwen_chat_ml(model_path: str):
         return "chatml-llava"
 
 
-@register_conv_template_matching_function
-def match_gemma3_instruct(model_path: str):
-    if re.search(r"gemma-3.*it", model_path, re.IGNORECASE):
-        return "gemma-it"
-
-
 @register_conv_template_matching_function
 def match_openbmb_minicpm(model_path: str):
     if re.search(r"minicpm-v", model_path, re.IGNORECASE):
@@ -1092,37 +1010,7 @@ def match_openbmb_minicpm(model_path: str):
         return "minicpmo"
 
 
-@register_conv_template_matching_function
-def match_moonshot_kimivl(model_path: str):
-    if re.search(r"kimi.*vl", model_path, re.IGNORECASE):
-        return "kimi-vl"
-
-
-@register_conv_template_matching_function
-def match_devstral(model_path: str):
-    if re.search(r"devstral", model_path, re.IGNORECASE):
-        return "devstral"
-
-
 @register_conv_template_matching_function
 def match_phi_4_mm(model_path: str):
     if "phi-4-multimodal" in model_path.lower():
         return "phi-4-mm"
-
-
-@register_conv_template_matching_function
-def match_vila(model_path: str):
-    if re.search(r"vila", model_path, re.IGNORECASE):
-        return "chatml"
-
-
-@register_conv_template_matching_function
-def match_mimo_vl(model_path: str):
-    if re.search(r"mimo.*vl", model_path, re.IGNORECASE):
-        return "mimo-vl"
-
-
-# @register_conv_template_matching_function
-# def match_step3(model_path: str):
-#     if re.search(r"step3", model_path, re.IGNORECASE):
-#         return "step3-vl"

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -88,6 +88,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
         self.extend_lens = [r.extend_input_len for r in reqs]
         self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
         self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
+        self.multimodal_inputs = [r.multimodal_inputs for r in reqs]
 
         # Build sampling info
         self.sampling_info = SamplingBatchInfo.from_schedule_batch(

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/srt/disaggregation/launch_lb.py

@@ -1,6 +1,8 @@
 import argparse
 import dataclasses
 
+from sglang.srt.disaggregation.mini_lb import PrefillConfig, run
+
 
 @dataclasses.dataclass
 class LBArgs:
@@ -18,7 +20,7 @@ class LBArgs:
         parser.add_argument(
             "--rust-lb",
             action="store_true",
-            help="Use Rust load balancer",
+            help="Deprecated, please use SGLang Router instead, this argument will have no effect.",
         )
         parser.add_argument(
             "--host",
@@ -115,25 +117,8 @@ def main():
     args = parser.parse_args()
     lb_args = LBArgs.from_cli_args(args)
 
-    if lb_args.rust_lb:
-        from sgl_pdlb._rust import LoadBalancer as RustLB
-
-        RustLB(
-            host=lb_args.host,
-            port=lb_args.port,
-            policy=lb_args.policy,
-            prefill_infos=lb_args.prefill_infos,
-            decode_infos=lb_args.decode_infos,
-            log_interval=lb_args.log_interval,
-            timeout=lb_args.timeout,
-        ).start()
-    else:
-        from sglang.srt.disaggregation.mini_lb import PrefillConfig, run
-
-        prefill_configs = [
-            PrefillConfig(url, port) for url, port in lb_args.prefill_infos
-        ]
-        run(prefill_configs, lb_args.decode_infos, lb_args.host, lb_args.port)
+    prefill_configs = [PrefillConfig(url, port) for url, port in lb_args.prefill_infos]
+    run(prefill_configs, lb_args.decode_infos, lb_args.host, lb_args.port)
 
 
 if __name__ == "__main__":

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/srt/disaggregation/mooncake/conn.py

@@ -37,6 +37,7 @@ from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
     format_tcp_address,
+    get_bool_env_var,
     get_free_port,
     get_int_env_var,
     get_ip,
@@ -198,6 +199,10 @@ class MooncakeKVManager(BaseKVManager):
             self.bootstrap_timeout = get_int_env_var(
                 "SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT", 300
             )
+
+            self.enable_custom_mem_pool = get_bool_env_var(
+                "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
+            )
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             self.heartbeat_failures = {}
             self.session_pool = defaultdict(requests.Session)
@@ -258,6 +263,26 @@ class MooncakeKVManager(BaseKVManager):
         socket.connect(endpoint)
         return socket
 
+    def _transfer_data(self, mooncake_session_id, transfer_blocks):
+        if not transfer_blocks:
+            return 0
+
+        # TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
+        if self.enable_custom_mem_pool:
+            # batch_transfer_sync has a higher chance to trigger an accuracy drop for MNNVL, fallback to transfer_sync temporarily
+            for src_addr, dst_addr, length in transfer_blocks:
+                status = self.engine.transfer_sync(
+                    mooncake_session_id, src_addr, dst_addr, length
+                )
+                if status != 0:
+                    return status
+            return 0
+        else:
+            src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
+            return self.engine.batch_transfer_sync(
+                mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
+            )
+
     def send_kvcache(
         self,
         mooncake_session_id: str,
@@ -283,17 +308,14 @@ class MooncakeKVManager(BaseKVManager):
 
         # Worker function for processing a single layer
         def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
+            transfer_blocks = []
             for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
                 src_addr = src_ptr + int(prefill_index[0]) * item_len
                 dst_addr = dst_ptr + int(decode_index[0]) * item_len
                 length = item_len * len(prefill_index)
+                transfer_blocks.append((src_addr, dst_addr, length))
 
-                status = self.engine.transfer_sync(
-                    mooncake_session_id, src_addr, dst_addr, length
-                )
-                if status != 0:
-                    return status
-            return 0
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
 
         futures = [
             executor.submit(
@@ -465,21 +487,17 @@ class MooncakeKVManager(BaseKVManager):
         dst_aux_ptrs: list[int],
         dst_aux_index: int,
     ):
-        src_addr_list = []
-        dst_addr_list = []
-        length_list = []
+        transfer_blocks = []
         prefill_aux_ptrs = self.kv_args.aux_data_ptrs
         prefill_aux_item_lens = self.kv_args.aux_item_lens
+
         for i, dst_aux_ptr in enumerate(dst_aux_ptrs):
             length = prefill_aux_item_lens[i]
             src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
             dst_addr = dst_aux_ptrs[i] + length * dst_aux_index
-            src_addr_list.append(src_addr)
-            dst_addr_list.append(dst_addr)
-            length_list.append(length)
-        return self.engine.batch_transfer_sync(
-            mooncake_session_id, src_addr_list, dst_addr_list, length_list
-        )
+            transfer_blocks.append((src_addr, dst_addr, length))
+
+        return self._transfer_data(mooncake_session_id, transfer_blocks)
 
     def sync_status_to_decode_endpoint(
         self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/srt/disaggregation/prefill.py

@@ -460,6 +460,7 @@ class SchedulerDisaggregationPrefillMixin:
 
         # We need to remove the sync in the following function for overlap schedule.
         self.set_next_batch_sampling_info_done(batch)
+        self.maybe_send_health_check_signal()
 
     def process_disagg_prefill_inflight_queue(
         self: Scheduler, rids_to_check: Optional[List[str]] = None

{sglang-0.4.10 → sglang-0.4.10.post2}/sglang/srt/distributed/device_communicators/pynccl.py

@@ -75,6 +75,7 @@ class PyNcclCommunicator:
         self.available = True
         self.disabled = False
 
+        self.nccl_version = self.nccl.ncclGetRawVersion()
         if self.rank == 0:
             logger.info("sglang is using nccl==%s", self.nccl.ncclGetVersion())
 
@@ -259,6 +260,12 @@ class PyNcclCommunicator:
             cudaStream_t(stream.cuda_stream),
         )
 
+    def register_comm_window_raw(self, ptr: int, size: int):
+        return self.nccl.ncclCommWindowRegister(self.comm, buffer_type(ptr), size, 1)
+
+    def deregister_comm_window(self, window):
+        return self.nccl.ncclCommWindowDeregister(self.comm, window)
+
     @contextmanager
     def change_state(
         self, enable: Optional[bool] = None, stream: Optional[torch.cuda.Stream] = None

sglang-0.4.10.post2/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -0,0 +1,133 @@
+import tempfile
+
+import torch
+from packaging import version
+from torch.cuda.memory import CUDAPluggableAllocator
+
+from sglang.srt.distributed.parallel_state import GroupCoordinator
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+
+nccl_allocator_source = """
+#include <nccl.h>
+extern "C" {
+
+void* nccl_alloc_plug(size_t size, int device, void* stream) {
+  void* ptr;
+  ncclResult_t err = ncclMemAlloc(&ptr, size);
+  return ptr;
+
+}
+
+void nccl_free_plug(void* ptr, size_t size, int device, void* stream) {
+  ncclResult_t err = ncclMemFree(ptr);
+}
+
+}
+"""
+
+_allocator = None
+_mem_pool = None
+_registered_base_addrs = set()
+_graph_pool_id = None
+
+
+def is_symmetric_memory_enabled():
+    return global_server_args_dict["enable_symm_mem"]
+
+
+def set_graph_pool_id(graph_pool_id):
+    global _graph_pool_id
+    _graph_pool_id = graph_pool_id
+
+
+def get_nccl_mem_pool():
+    global _allocator, _mem_pool
+    if _mem_pool is None:
+        out_dir = tempfile.gettempdir()
+        nccl_allocator_libname = "nccl_allocator"
+        torch.utils.cpp_extension.load_inline(
+            name=nccl_allocator_libname,
+            cpp_sources=nccl_allocator_source,
+            with_cuda=True,
+            extra_ldflags=["-lnccl"],
+            verbose=True,
+            is_python_module=False,
+            build_directory=out_dir,
+        )
+        _allocator = CUDAPluggableAllocator(
+            f"{out_dir}/{nccl_allocator_libname}.so",
+            "nccl_alloc_plug",
+            "nccl_free_plug",
+        ).allocator()
+        _mem_pool = torch.cuda.MemPool(_allocator)
+    return _mem_pool
+
+
+class use_symmetric_memory:
+    def __init__(self, group_coordinator: GroupCoordinator):
+        if not is_symmetric_memory_enabled():
+            self.group_coordinator = None
+            self._mem_pool_ctx = None
+            self.is_graph_capture = None
+            self.device = None
+            self.pre_2_8_0 = None
+        else:
+            self.group_coordinator = group_coordinator
+            self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool())
+            self.is_graph_capture = torch.cuda.is_current_stream_capturing()
+            self.device = torch.cuda.current_device()
+            self.pre_2_8_0 = version.parse(torch.__version__) < version.parse("2.8.0")
+
+    def __enter__(self):
+        if not is_symmetric_memory_enabled():
+            return self
+        assert (
+            self.group_coordinator.pynccl_comm is not None
+        ), f"Symmetric memory requires pynccl to be enabled in group '{self.group_coordinator.group_name}'"
+        assert (
+            self.group_coordinator.pynccl_comm.nccl_version >= 22703
+        ), "NCCL version 2.27.3 or higher is required for NCCL symmetric memory"
+        if self.is_graph_capture:
+            assert (
+                _graph_pool_id is not None
+            ), "graph_pool_id is not set under graph capture"
+            # Pause graph memory pool to use symmetric memory with cuda graph
+            if self.pre_2_8_0:
+                torch._C._cuda_endAllocateCurrentStreamToPool(
+                    self.device, _graph_pool_id
+                )
+            else:
+                torch._C._cuda_endAllocateToPool(self.device, _graph_pool_id)
+        self._mem_pool_ctx.__enter__()
+        return self
+
+    def tag(self, tensor: torch.Tensor):
+        if not is_symmetric_memory_enabled():
+            return
+        tensor.symmetric_memory = True
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if not is_symmetric_memory_enabled():
+            return
+        global _registered_base_addrs
+        self._mem_pool_ctx.__exit__(exc_type, exc_val, exc_tb)
+        for segment in get_nccl_mem_pool().snapshot():
+            if segment["address"] not in _registered_base_addrs:
+                if segment["stream"] == 0 and self.pre_2_8_0:
+                    # PyTorch version < 2.8.0 has a multi-thread MemPool bug
+                    # See https://github.com/pytorch/pytorch/issues/152861
+                    # Fixed at https://github.com/pytorch/pytorch/commit/f01e628e3b31852983ab30b25bf251f557ba9c0b
+                    # WAR is to skip allocations on the default stream since the forward_pass thread always runs on a custom stream
+                    continue
+                self.group_coordinator.pynccl_comm.register_comm_window_raw(
+                    segment["address"], segment["total_size"]
+                )
+                _registered_base_addrs.add(segment["address"])
+
+        if self.is_graph_capture:
+            if self.pre_2_8_0:
+                torch._C._cuda_beginAllocateToPool(self.device, _graph_pool_id)
+            else:
+                torch._C._cuda_beginAllocateCurrentThreadToPool(
+                    self.device, _graph_pool_id
+                )