From c356aadb0e5e5dfea498780479ec80ca96231902 Mon Sep 17 00:00:00 2001
From: sufubao <sufubao@sensetime.com>
Date: Sun, 5 Jul 2026 11:36:08 +0800
Subject: [PATCH] feat: gc.freeze hot processes and disable gc during cudagraph
 capture

---
 lightllm/common/basemodel/basemodel.py    |  8 +++++---
 lightllm/server/api_http.py               |  2 ++
 lightllm/server/detokenization/manager.py |  2 ++
 lightllm/server/router/manager.py         |  2 ++
 lightllm/utils/gc_utils.py                | 23 +++++++++++++++++++++++
 5 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 lightllm/utils/gc_utils.py

diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py
index 72d3f26cd..5f76c6824 100755
--- a/lightllm/common/basemodel/basemodel.py
+++ b/lightllm/common/basemodel/basemodel.py
@@ -24,6 +24,7 @@
 from lightllm.common.quantization import Quantcfg
 from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.gc_utils import freeze_gc, gc_frozen_and_disabled
 from lightllm.utils.dist_utils import get_dp_world_size
 from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type, get_added_mtp_kv_layer_num
 from lightllm.distributed.communication_op import dist_group_manager
@@ -139,11 +140,13 @@ def __init__(self, kvargs):
         self._autotune_warmup()
         self._full_att_decode_autotune()
         self._init_padded_req()
-        self._init_cudagraph()
-        self._init_prefill_cuda_graph()
+        with gc_frozen_and_disabled("cudagraph-capture"):
+            self._init_cudagraph()
+            self._init_prefill_cuda_graph()
         self._check_max_len_infer()
         torch.cuda.empty_cache()
         set_model_init_status(True)
+        freeze_gc("model-infer-worker")
         return
 
     def _init_config(self):
@@ -663,7 +666,6 @@ def _decode(
 
     @final
     def _context_forward(self, infer_state: InferStateInfo):
-
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"
diff --git a/lightllm/server/api_http.py b/lightllm/server/api_http.py
index 628a8f728..33dee4ecb 100755
--- a/lightllm/server/api_http.py
+++ b/lightllm/server/api_http.py
@@ -47,6 +47,7 @@
 from .api_lightllm import lightllm_get_score
 from lightllm.utils.envs_utils import get_env_start_args, get_lightllm_websocket_max_message_size
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.gc_utils import freeze_gc
 from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError
 from lightllm.server.metrics.manager import MetricClient
 from lightllm.utils.envs_utils import get_unique_server_name
@@ -542,4 +543,5 @@ async def startup_event():
     g_objs.set_args(get_env_start_args())
     loop.create_task(g_objs.httpserver_manager.handle_loop())
     logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}")
+    freeze_gc("httpserver")
     return
diff --git a/lightllm/server/detokenization/manager.py b/lightllm/server/detokenization/manager.py
index 58f932485..0d4c9ebbc 100644
--- a/lightllm/server/detokenization/manager.py
+++ b/lightllm/server/detokenization/manager.py
@@ -16,6 +16,7 @@
 import pickle
 import time
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.gc_utils import freeze_gc
 from lightllm.utils.envs_utils import get_unique_server_name
 from lightllm.utils.shm_port_args import get_shm_port_args
 
@@ -181,5 +182,6 @@ def start_detokenization_process(args, pipe_writer):
         pipe_writer.send(str(e))
         raise
     pipe_writer.send("init ok")
+    freeze_gc("detokenization")
     manager.handle_loop()
     return
diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py
index 01634d962..e35a76614 100644
--- a/lightllm/server/router/manager.py
+++ b/lightllm/server/router/manager.py
@@ -25,6 +25,7 @@
 from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient
 from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer
 from lightllm.utils.log_utils import init_logger, log_time_ready
+from lightllm.utils.gc_utils import freeze_gc
 from lightllm.utils.profiler import ProfilerCmd
 from lightllm.server.router.token_load import TokenLoad
 from lightllm.server.metrics.manager import MetricClient
@@ -508,5 +509,6 @@ def handle_exception(loop, context):
         raise
 
     pipe_writer.send("init ok")
+    freeze_gc("router")
     loop.run_until_complete(router.loop_for_fwd())
     return
diff --git a/lightllm/utils/gc_utils.py b/lightllm/utils/gc_utils.py
new file mode 100644
index 000000000..d726ebe79
--- /dev/null
+++ b/lightllm/utils/gc_utils.py
@@ -0,0 +1,23 @@
+import gc
+from contextlib import contextmanager
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def freeze_gc(tag: str = "") -> None:
+    gc.collect()
+    gc.freeze()
+    logger.info(f"gc.freeze done ({tag}): frozen={gc.get_freeze_count()}")
+
+
+@contextmanager
+def gc_frozen_and_disabled(tag: str = ""):
+    freeze_gc(tag)
+    was_enabled = gc.isenabled()
+    gc.disable()
+    try:
+        yield
+    finally:
+        if was_enabled:
+            gc.enable()