From 14e8bd889f27da2a9d116f6744fa48812348491b Mon Sep 17 00:00:00 2001
From: Yongtong Wu <914554688@qq.com>
Date: Mon, 14 Apr 2025 16:04:46 +0800
Subject: [PATCH] Free metadata_buffer_index after transfer finished (#5364)

---
 python/sglang/srt/disaggregation/prefill.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
index c6020767b..cdd1676a8 100644
--- a/python/sglang/srt/disaggregation/prefill.py
+++ b/python/sglang/srt/disaggregation/prefill.py
@@ -226,6 +226,11 @@ class SchedulerDisaggregationPrefillMixin:
             elif poll == KVPoll.Failed:
                 raise Exception("Transferring failed")
 
+        for req in done_reqs:
+            self.disagg_prefill_pending_queue.req_to_metadata_buffer_idx_allocator.free(
+                req.metadata_buffer_index
+            )
+
         # Stream requests which have finished transfer
         self.stream_output(done_reqs, False, None)