[docker] alleviate pd memory leakage (#1525)

Author: zhuzilin

docker/patch/latest/sglang.patch

@@ -16,32 +16,35 @@ index aa10cb08d..d41c31a09 100644
              self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN"
 
 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py
-index 51af67636..54716de5c 100644
+index 51af67636..661ea6fd6 100644
 --- a/python/sglang/srt/disaggregation/decode.py
 +++ b/python/sglang/srt/disaggregation/decode.py
-@@ -315,6 +315,13 @@ class DecodePreallocQueue:
+@@ -315,6 +315,16 @@ class DecodePreallocQueue:
          )
          return kv_manager
 
 +    def release_memory_occupation(self):
-+        if hasattr(self.kv_manager, "close"):
-+            self.kv_manager.close()
++        self.queue.clear()
++        self.retracted_queue.clear()
++        if hasattr(self.kv_manager, "deregister_buffer_to_engine"):
++            self.kv_manager.deregister_buffer_to_engine()
 +
 +    def resume_memory_occupation(self):
-+        self.kv_manager = self._init_kv_manager()
++        if hasattr(self.kv_manager, "register_buffer_to_engine"):
++            self.kv_manager.register_buffer_to_engine()
 +
      def add(self, req: Req, is_retracted: bool = False) -> None:
          """Add a request to the pending queue."""
          if self._check_if_req_exceed_kv_capacity(req):
 diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py
-index 32e8c0b69..df913da7b 100644
+index 32e8c0b69..dc93c5c5f 100644
 --- a/python/sglang/srt/disaggregation/mooncake/conn.py
 +++ b/python/sglang/srt/disaggregation/mooncake/conn.py
-@@ -1079,6 +1079,19 @@ class MooncakeKVManager(CommonKVManager):
-             f"Losing connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), {len(affected_rooms)} requests affected"
-         )
+@@ -253,6 +253,19 @@ class MooncakeKVManager(CommonKVManager):
+                 self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
+             )
 
-+    def close(self):
++    def deregister_buffer_to_engine(self):
 +        # Batch deregister KV data buffers
 +        if self.kv_args.kv_data_ptrs:
 +            self.engine.batch_deregister(self.kv_args.kv_data_ptrs)
@@ -54,23 +57,25 @@ index 32e8c0b69..df913da7b 100644
 +        if self.kv_args.state_data_ptrs:
 +            self.engine.batch_deregister(self.kv_args.state_data_ptrs)
 +
- 
- class MooncakeKVSender(CommonKVSender):
- 
+     def _transfer_data(self, mooncake_session_id, transfer_blocks):
+         if not transfer_blocks:
+             return 0
 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
-index a6eed743a..0124d8917 100644
+index a6eed743a..24a72ca70 100644
 --- a/python/sglang/srt/disaggregation/prefill.py
 +++ b/python/sglang/srt/disaggregation/prefill.py
-@@ -306,6 +306,13 @@ class PrefillBootstrapQueue:
+@@ -306,6 +306,15 @@ class PrefillBootstrapQueue:
          else:
              return bootstrapped_reqs, failed_reqs
 
 +    def release_memory_occupation(self):
-+        if hasattr(self.kv_manager, "close"):
-+            self.kv_manager.close()
++        self.queue.clear()
++        if hasattr(self.kv_manager, "deregister_buffer_to_engine"):
++            self.kv_manager.deregister_buffer_to_engine()
 +
 +    def resume_memory_occupation(self):
-+        self.kv_manager = self._init_kv_manager()
++        if hasattr(self.kv_manager, "register_buffer_to_engine"):
++            self.kv_manager.register_buffer_to_engine()
 +
 
  class SchedulerDisaggregationPrefillMixin:

docker/version.txt

@@ -1 +1 @@
-nightly-dev-20260129a
+nightly-dev-20260130a