From ffb32a854835360ffc4e2e18130a8b3474670f42 Mon Sep 17 00:00:00 2001
From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com>
Date: Mon, 3 Nov 2025 05:51:27 -0800
Subject: [PATCH] Conditionally recapture cuda graph after model weight update
 from disk  (#12060)

---
 python/sglang/srt/managers/io_struct.py          | 2 ++
 python/sglang/srt/managers/tp_worker.py          | 2 +-
 python/sglang/srt/model_executor/model_runner.py | 4 ++++
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index 3e740621f..941b5a19d 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -1024,6 +1024,8 @@ class UpdateWeightFromDiskReqInput(BaseReq):
     torch_empty_cache: bool = False
     # Whether to keep the scheduler paused after weight update
     keep_pause: bool = False
+    # Whether to recapture cuda graph after weight udpdate
+    recapture_cuda_graph: bool = False
     # The trainer step id. Used to know which step's weights are used for sampling.
     token_step: int = 0
 
diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
index e8f318da0..97851260f 100644
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -100,7 +100,7 @@ class BaseTpWorker(ABC):
 
     def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
         success, message = self.model_runner.update_weights_from_disk(
-            recv_req.model_path, recv_req.load_format
+            recv_req.model_path, recv_req.load_format, recv_req.recapture_cuda_graph
         )
         return success, message
 
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index f48775fb3..963e2cd81 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -862,6 +862,7 @@ class ModelRunner:
         model_path: str,
         load_format: str,
         weight_name_filter: Optional[Callable[[str], bool]] = None,
+        recapture_cuda_graph: bool = False,
     ) -> tuple[bool, str]:
         """Update engine weights in-place from the disk."""
         logger.info(
@@ -917,6 +918,9 @@ class ModelRunner:
         self.server_args.load_format = load_format
         self.load_config = load_config
 
+        if recapture_cuda_graph and self.device == "cuda":
+            self.init_device_graphs()
+
         logger.info("Update weights end.")
         return True, "Succeeded to update model weights."