From 68e31a348507dc6371ac53f3c0aff893d579adc6 Mon Sep 17 00:00:00 2001
From: Zheng Wengang <zwg0606@gmail.com>
Date: Mon, 9 Feb 2026 11:34:29 +0800
Subject: [PATCH] [BugFix][PD]Fix metadata_buffer_index leak when aborted in PD
 (#17483)

---
 python/sglang/srt/disaggregation/prefill.py | 26 +++++++++++++++++++--
 python/sglang/srt/managers/scheduler.py     |  6 +++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
index aa86731fa..eadd04e21 100644
--- a/python/sglang/srt/disaggregation/prefill.py
+++ b/python/sglang/srt/disaggregation/prefill.py
@@ -62,6 +62,27 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def release_req_to_metadata_buffer(
+    req: Req, allocator: ReqToMetadataIdxAllocator
+) -> None:
+    """
+    Release the metadata buffer index allocated for a request in prefill disaggregation mode.
+
+    This function safely releases the metadata buffer index if it was allocated.
+
+    Args:
+        req: The request object that may have a metadata_buffer_index allocated
+        allocator: The ReqToMetadataIdxAllocator instance to free the index
+    """
+    if (
+        hasattr(req, "metadata_buffer_index")
+        and req.metadata_buffer_index is not None
+        and req.metadata_buffer_index >= 0
+    ):
+        allocator.free(req.metadata_buffer_index)
+        req.metadata_buffer_index = -1
+
+
 class PrefillBootstrapQueue:
     """
     Store the requests in bootstrapping
@@ -592,8 +613,9 @@ class SchedulerDisaggregationPrefillMixin:
         for req in done_reqs:
             req: Req
             req.add_latency(RequestStage.PREFILL_TRANSFER_KV_CACHE)
-            self.req_to_metadata_buffer_idx_allocator.free(req.metadata_buffer_index)
-            req.metadata_buffer_index = -1
+            release_req_to_metadata_buffer(
+                req, self.req_to_metadata_buffer_idx_allocator
+            )
             trace_slice(
                 RequestStage.PREFILL_TRANSFER_KV_CACHE, req.rid, thread_finish_flag=True
             )
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index 1cf897f2d..d1ee15275 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -47,6 +47,7 @@ from sglang.srt.disaggregation.encode_receiver import MMReceiverHTTP
 from sglang.srt.disaggregation.prefill import (
     PrefillBootstrapQueue,
     SchedulerDisaggregationPrefillMixin,
+    release_req_to_metadata_buffer,
 )
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
@@ -2726,6 +2727,11 @@ class Scheduler(
             # For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
             if self.disaggregation_mode == DisaggregationMode.DECODE:
                 release_kv_cache(req, self.tree_cache)
+            # For disaggregation prefill mode, free the metadata buffer index
+            if self.disaggregation_mode == DisaggregationMode.PREFILL:
+                release_req_to_metadata_buffer(
+                    req, self.req_to_metadata_buffer_idx_allocator
+                )
 
             # For mamba radix cache
             if (