Free metadata_buffer_index after transfer finished (#5364)

14e8bd88 · Yongtong Wu · GitHub · adca585b · 14e8bd88
Unverified Commit 14e8bd88 authored Apr 14, 2025 by Yongtong Wu Committed by GitHub Apr 14, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

python/sglang/srt/disaggregation/prefill.py python/sglang/srt/disaggregation/prefill.py +5 -0

No files found.
--- 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)