[VLM] retire aborted disaggregated prefill results (#36988)

python/sglang/srt/disaggregation/prefill.py

@@ -729,6 +729,7 @@ class SchedulerDisaggregationPrefillMixin:
             result.indexer_topk_output = None
 
         logprob_pt = 0
+        aborted_reqs: List[Req] = []
         assert batch.spec_info is result.next_draft_input
         draft_input = result.next_draft_input
         draft_hidden_states_cpu = None
@@ -763,6 +764,13 @@ class SchedulerDisaggregationPrefillMixin:
             if req.inflight_middle_chunks <= 0:
                 req.time_stats.set_prefill_finished_time()
 
+                if is_aborted(req):
+                    if self._retire_aborted_prefill_result(req):
+                        req.time_stats.set_completion_time()
+                        aborted_reqs.append(req)
+                    advance_logprob_pt(i, req)
+                    continue
+
                 # Test hook: exercise the release/requeue retry path.
                 if req.pending_bootstrap and should_force_retry(req):
                     self.optimistic_release_and_requeue(req)
@@ -770,6 +778,24 @@ class SchedulerDisaggregationPrefillMixin:
                     continue
 
                 req.output_ids.append(next_token_id)
+                if req.grammar is not None:
+                    try:
+                        req.grammar.accept_token(next_token_id)
+                    except ValueError as e:
+                        error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
+                        prepare_abort(
+                            req,
+                            error_message,
+                            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                        )
+                    req.grammar.finished = req.finished()
+                    if is_aborted(req):
+                        if self._retire_aborted_prefill_result(req):
+                            req.time_stats.set_completion_time()
+                            aborted_reqs.append(req)
+                        advance_logprob_pt(i, req)
+                        continue
+
                 maybe_cache_unfinished_req(req, self.tree_cache)
                 self.disagg_prefill_inflight_queue.append(req)
                 if self.spec_algorithm.is_eagle() and draft_input is not None:
@@ -808,17 +834,6 @@ class SchedulerDisaggregationPrefillMixin:
                     self.send_kv_chunk(req, last_chunk=True)
                 req.time_stats.set_prefill_transfer_queue_entry_time()
 
-                if req.grammar is not None:
-                    try:
-                        req.grammar.accept_token(next_token_id)
-                    except ValueError as e:
-                        error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
-                        prepare_abort(
-                            req,
-                            error_message,
-                            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
-                        )
-                    req.grammar.finished = req.finished()
             else:
                 # being chunked reqs' prefill is not finished
                 req.inflight_middle_chunks -= 1
@@ -831,16 +846,18 @@ class SchedulerDisaggregationPrefillMixin:
                     req.extend_range is not None
                     and req.extend_range.end >= len(req.origin_input_ids)
                 )
-                if req.pending_bootstrap and not still_chunking:
+                # Abort is terminal. Do not requeue an aborted optimistic
-                    self.optimistic_release_and_requeue(req)
+                # request merely because bootstrap is still pending.
+                if is_aborted(req):
+                    if not still_chunking and self._retire_aborted_prefill_result(req):
+                        req.time_stats.set_completion_time()
+                        aborted_reqs.append(req)
                     advance_logprob_pt(i, req)
                     req.time_stats.set_last_chunked_prefill_finish_time()
                     continue
 
-                # Optimistic bootstrap can fail while this overlapped chunk is
+                if req.pending_bootstrap and not still_chunking:
-                # already running. Drop aborted chunks instead of sending KV.
+                    self.optimistic_release_and_requeue(req)
-                if is_aborted(req):
-                    self.clear_pending_chunk_send(req)
                     advance_logprob_pt(i, req)
                     req.time_stats.set_last_chunked_prefill_finish_time()
                     continue
@@ -876,6 +893,12 @@ class SchedulerDisaggregationPrefillMixin:
                 auxiliary_output_starts,
             )
 
+        if aborted_reqs:
+            self.output_streamer.stream_output(
+                aborted_reqs,
+                any(req.return_logprob for req in aborted_reqs),
+            )
+
         can_run_cuda_graph = result.can_run_cuda_graph
         self.metrics_reporter.report_prefill_stats(
             batch=batch,
@@ -1030,6 +1053,35 @@ class SchedulerDisaggregationPrefillMixin:
         """
         self.disagg_prefill_pending_chunk_rids.discard(req.rid)
 
+    def _retire_aborted_prefill_result(self: Scheduler, req: Req) -> bool:
+        """Release an aborted request when its last prefill result is safe."""
+        self.clear_pending_chunk_send(req)
+        owns_resources = (
+            req.kv.holds_kv or req.kv.holds_mamba or req.metadata_buffer_index >= 0
+        )
+        if not owns_resources:
+            # A bootstrap failure or earlier abort already retired it.
+            return False
+
+        sender = req.disagg_kv_sender
+        if sender is not None:
+            try:
+                sender.abort()
+            except Exception:
+                # Transport notification is best effort; local ownership must
+                # still be released or the next idle invariant check will fail.
+                logger.exception("Failed to notify KV sender of abort for %s", req.rid)
+
+        if req.to_finish is not None and not req.finished():
+            req.update_finish_state()
+        maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator)
+        req.pending_bootstrap = False
+        if self.enable_hicache_storage:
+            self.tree_cache.release_aborted_request(req.rid)
+        if req.kv.holds_kv or req.kv.holds_mamba:
+            release_kv_cache(req, self.tree_cache, is_insert=False)
+        return True
+
     def handle_bootstrap_failure(self: Scheduler, req: Req) -> None:
         self.clear_pending_chunk_send(req)
         error_message = (

test/registered/disaggregation/test_disaggregation_basic.py

@@ -1,13 +1,9 @@
 import asyncio
 import json
 import os
-import threading
-import time
 import unittest
-import uuid
 from concurrent.futures import ThreadPoolExecutor
 from types import SimpleNamespace
-from typing import Any
 
 import aiohttp
 import openai
@@ -28,7 +24,6 @@ from sglang.test.server_fixtures.disaggregation_fixture import (
 from sglang.test.test_utils import (
     DEFAULT_DRAFT_MODEL_EAGLE3,
     DEFAULT_MODEL_NAME_FOR_TEST,
-    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
     DEFAULT_TARGET_MODEL_EAGLE3,
 )
 
@@ -754,104 +749,5 @@ class TestDisaggregationPauseResumePrefillLeak(PDDisaggregationServerBase):
             )
 
 
-PD_CHUNKED_ABORT_EXTRA_ARGS = [
-    "--max-running-requests",
-    "4",
-    "--chunked-prefill-size",
-    "64",
-]
-_CHUNKED_ABORT_LONG_PROMPT = (
-    "The quick brown fox jumps over the lazy dog. "
-    "Pack my box with five dozen liquor jugs. "
-    "Sphinx of black quartz, judge my vow. "
-) * 900
-
-
-def _decode_response(response: requests.Response) -> Any:
-    try:
-        return response.json()
-    except ValueError:
-        return response.text
-
-
-def _is_abort_result(status_code: int, body: Any) -> bool:
-    if status_code == 200:
-        reason = (
-            body.get("meta_info", {}).get("finish_reason", {})
-            if isinstance(body, dict)
-            else {}
-        )
-        return isinstance(reason, dict) and reason.get("type") == "abort"
-
-    if status_code not in (500, 503):
-        return False
-
-    text = body if isinstance(body, str) else str(body)
-    return "abort" in text.lower()
-
-
-class TestDisaggChunkedPrefillAbort(PDDisaggregationServerBase):
-    @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
-        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
-        cls.extra_prefill_args = PD_CHUNKED_ABORT_EXTRA_ARGS
-        cls.extra_decode_args = PD_CHUNKED_ABORT_EXTRA_ARGS
-        cls.launch_all()
-
-    def _post_abort(self, rid: str):
-        for url in (self.prefill_url, self.decode_url):
-            requests.post(
-                url + "/abort_request",
-                json={"rid": rid, "abort_all": False},
-                timeout=10,
-            )
-
-    def test_abort_mid_chunked_prefill_by_rid(self):
-        rid = f"pd-chunked-prefill-abort-{uuid.uuid4().hex}"
-        result: dict[str, Any] = {}
-
-        def run_generate():
-            try:
-                response = requests.post(
-                    self.lb_url + "/generate",
-                    json={
-                        "rid": rid,
-                        "text": f"{rid}\n{_CHUNKED_ABORT_LONG_PROMPT}",
-                        "sampling_params": {
-                            "temperature": 0,
-                            "max_new_tokens": 4096,
-                            "ignore_eos": True,
-                        },
-                    },
-                    timeout=180,
-                )
-                result["status_code"] = response.status_code
-                result["body"] = _decode_response(response)
-            except requests.RequestException as exc:
-                result["exception"] = repr(exc)
-
-        thread = threading.Thread(target=run_generate)
-        thread.start()
-
-        time.sleep(1.0)
-        abort_deadline = time.monotonic() + 8
-        while thread.is_alive() and time.monotonic() < abort_deadline:
-            self._post_abort(rid)
-            time.sleep(0.2)
-
-        thread.join(timeout=60)
-        self.assertFalse(thread.is_alive(), "Chunked-prefill abort request hung")
-        self.assertNotIn("exception", result, result.get("exception"))
-        self.assertTrue(
-            _is_abort_result(result["status_code"], result["body"]),
-            f"Expected chunked-prefill request to abort, got {result}",
-        )
-
-        for url in (self.lb_url, self.prefill_url, self.decode_url):
-            health = requests.get(url + "/health", timeout=10)
-            self.assertEqual(health.status_code, 200, health.text)
-
-
 if __name__ == "__main__":
     unittest.main()

test/registered/disaggregation/test_disaggregation_chunked_prefill_abort.py

@@ -0,0 +1,117 @@
+import threading
+import time
+import unittest
+import uuid
+from typing import Any
+
+import requests
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="2-gpu-large")
+
+PD_CHUNKED_ABORT_EXTRA_ARGS = [
+    "--max-running-requests",
+    "4",
+    "--chunked-prefill-size",
+    "64",
+]
+_CHUNKED_ABORT_LONG_PROMPT = (
+    "The quick brown fox jumps over the lazy dog. "
+    "Pack my box with five dozen liquor jugs. "
+    "Sphinx of black quartz, judge my vow. "
+) * 900
+
+
+def _decode_response(response: requests.Response) -> Any:
+    try:
+        return response.json()
+    except ValueError:
+        return response.text
+
+
+def _is_abort_result(status_code: int, body: Any) -> bool:
+    if status_code == 200:
+        reason = (
+            body.get("meta_info", {}).get("finish_reason", {})
+            if isinstance(body, dict)
+            else {}
+        )
+        return isinstance(reason, dict) and reason.get("type") == "abort"
+
+    if status_code not in (500, 503):
+        return False
+
+    text = body if isinstance(body, str) else str(body)
+    return "abort" in text.lower()
+
+
+class TestDisaggChunkedPrefillAbort(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.extra_prefill_args = PD_CHUNKED_ABORT_EXTRA_ARGS
+        cls.extra_decode_args = PD_CHUNKED_ABORT_EXTRA_ARGS
+        cls.launch_all()
+
+    def _post_abort(self, rid: str):
+        for url in (self.prefill_url, self.decode_url):
+            requests.post(
+                url + "/abort_request",
+                json={"rid": rid, "abort_all": False},
+                timeout=10,
+            )
+
+    def test_abort_mid_chunked_prefill_by_rid(self):
+        rid = f"pd-chunked-prefill-abort-{uuid.uuid4().hex}"
+        result: dict[str, Any] = {}
+
+        def run_generate():
+            try:
+                response = requests.post(
+                    self.lb_url + "/generate",
+                    json={
+                        "rid": rid,
+                        "text": f"{rid}\n{_CHUNKED_ABORT_LONG_PROMPT}",
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 4096,
+                            "ignore_eos": True,
+                        },
+                    },
+                    timeout=180,
+                )
+                result["status_code"] = response.status_code
+                result["body"] = _decode_response(response)
+            except requests.RequestException as exc:
+                result["exception"] = repr(exc)
+
+        thread = threading.Thread(target=run_generate)
+        thread.start()
+
+        time.sleep(1.0)
+        abort_deadline = time.monotonic() + 8
+        while thread.is_alive() and time.monotonic() < abort_deadline:
+            self._post_abort(rid)
+            time.sleep(0.2)
+
+        thread.join(timeout=60)
+        self.assertFalse(thread.is_alive(), "Chunked-prefill abort request hung")
+        self.assertNotIn("exception", result, result.get("exception"))
+        self.assertTrue(
+            _is_abort_result(result["status_code"], result["body"]),
+            f"Expected chunked-prefill request to abort, got {result}",
+        )
+
+        for url in (self.lb_url, self.prefill_url, self.decode_url):
+            health = requests.get(url + "/health", timeout=10)
+            self.assertEqual(health.status_code, 200, health.text)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py

@@ -0,0 +1,208 @@
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
+from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
+from sglang.srt.managers.utils import GenerationBatchResult
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class _Req:
+    def __init__(self, *, inflight_middle_chunks: int, allocated: bool = True):
+        self.rid = "aborted-prefill"
+        self.inflight_middle_chunks = inflight_middle_chunks
+        self.kv = ReqKvInfo(
+            req_pool_idx=1 if allocated else None,
+            kv_allocated_len=1 if allocated else 0,
+            mamba_pool_idx=object() if allocated else None,
+        )
+        self.metadata_buffer_index = 7 if allocated else -1
+        self.pending_bootstrap = allocated
+        self.disagg_kv_sender = Mock()
+        self.to_finish = FINISH_ABORT() if allocated else None
+        self.finished_reason = None if allocated else FINISH_ABORT()
+        self.return_logprob = False
+        self.return_sampling_mask = False
+        self.grammar = None
+        self.output_ids = []
+        self.origin_input_ids = list(range(100))
+        self.extend_range = None
+        self.time_stats = SimpleNamespace(
+            set_prefill_finished_time=Mock(),
+            set_last_chunked_prefill_finish_time=Mock(),
+            set_completion_time=Mock(),
+        )
+
+    def finished(self):
+        return self.finished_reason is not None
+
+    def update_finish_state(self):
+        self.finished_reason = self.to_finish
+        self.to_finish = None
+
+
+class _Scheduler(SchedulerDisaggregationPrefillMixin):
+    def __init__(self):
+        self.batch_result_processor = SimpleNamespace(
+            snapshot_auxiliary_output_starts=Mock(return_value=[]),
+            move_logprobs_to_cpu=Mock(),
+            consume_auxiliary_output=Mock(),
+        )
+        self.spec_algorithm = SimpleNamespace(is_eagle=lambda: False)
+        self.tree_cache = Mock()
+        self.disagg_prefill_inflight_queue = []
+        self.disagg_prefill_pending_chunk_rids = {"aborted-prefill"}
+        self.send_kv_chunk = Mock()
+        self.output_streamer = Mock()
+        self.metrics_reporter = SimpleNamespace(report_prefill_stats=Mock())
+        self.req_to_metadata_buffer_idx_allocator = Mock()
+        self.enable_hicache_storage = True
+        self.chunked_req = None
+
+
+def _batch(req):
+    return SimpleNamespace(
+        reqs=[req],
+        spec_info=None,
+        prefill_stats=None,
+        dp_cooperation_info=None,
+    )
+
+
+def _result():
+    return GenerationBatchResult(next_token_ids=torch.tensor([11]))
+
+
+def _free_req(req, _tree_cache, *, is_insert):
+    assert is_insert is False
+    req.kv.req_pool_idx = None
+    req.kv.mark_kv_released()
+    req.kv.mamba_pool_idx = None
+
+
+@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req)
+@patch("sglang.srt.disaggregation.prefill.maybe_cache_unfinished_req")
+def test_aborted_final_result_releases_hybrid_cache(
+    maybe_cache_unfinished_req, release_kv_cache
+):
+    scheduler = _Scheduler()
+    req = _Req(inflight_middle_chunks=0)
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False)
+    maybe_cache_unfinished_req.assert_not_called()
+    req.disagg_kv_sender.abort.assert_called_once_with()
+    scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7)
+    scheduler.tree_cache.release_aborted_request.assert_called_once_with(req.rid)
+    scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
+    scheduler.send_kv_chunk.assert_not_called()
+    assert req.output_ids == []
+    assert req.finished()
+    assert req.metadata_buffer_index == -1
+    assert req.rid not in scheduler.disagg_prefill_pending_chunk_rids
+
+
+@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req)
+def test_aborted_middle_result_releases_after_last_chunk(release_kv_cache):
+    scheduler = _Scheduler()
+    req = _Req(inflight_middle_chunks=1)
+    req.extend_range = SimpleNamespace(end=50)
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    assert req.inflight_middle_chunks == 0
+    release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False)
+    scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
+
+
+@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req)
+def test_aborted_middle_result_waits_for_inflight_chunk(release_kv_cache):
+    scheduler = _Scheduler()
+    req = _Req(inflight_middle_chunks=1)
+    req.extend_range = SimpleNamespace(end=len(req.origin_input_ids))
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    release_kv_cache.assert_not_called()
+    scheduler.output_streamer.stream_output.assert_not_called()
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False)
+    scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
+
+
+@patch("sglang.srt.disaggregation.prefill.release_kv_cache")
+def test_delayed_result_ignores_already_retired_request(release_kv_cache):
+    scheduler = _Scheduler()
+    req = _Req(inflight_middle_chunks=0, allocated=False)
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    release_kv_cache.assert_not_called()
+    req.disagg_kv_sender.abort.assert_not_called()
+    scheduler.output_streamer.stream_output.assert_not_called()
+
+
+@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req)
+def test_sender_abort_failure_does_not_skip_local_cleanup(release_kv_cache):
+    scheduler = _Scheduler()
+    req = _Req(inflight_middle_chunks=0)
+    req.disagg_kv_sender.abort.side_effect = RuntimeError("transport is down")
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False)
+    scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7)
+    scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
+    assert req.finished()
+
+
+@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req)
+@patch("sglang.srt.disaggregation.prefill.maybe_cache_unfinished_req")
+def test_grammar_rejection_retires_prefill_before_transfer(
+    maybe_cache_unfinished_req, release_kv_cache
+):
+    scheduler = _Scheduler()
+    req = _Req(inflight_middle_chunks=0)
+    req.to_finish = None
+    req.grammar = Mock()
+    req.grammar.accept_token.side_effect = ValueError("invalid token")
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    req.grammar.accept_token.assert_called_once_with(11)
+    assert req.grammar.finished
+    assert req.finished()
+    release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False)
+    maybe_cache_unfinished_req.assert_not_called()
+    scheduler.send_kv_chunk.assert_not_called()
+    scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
+    assert scheduler.disagg_prefill_inflight_queue == []
+
+
+def test_aborted_result_releases_mamba_allocated_before_kv():
+    scheduler = _Scheduler()
+    scheduler.enable_hicache_storage = False
+    scheduler.tree_cache.supports_mamba.return_value = True
+    scheduler.tree_cache.req_to_token_pool.mamba_allocator.free = Mock()
+    req = _Req(inflight_middle_chunks=0, allocated=False)
+    req.kv.mamba_pool_idx = torch.tensor([3])
+    req.to_finish = FINISH_ABORT()
+    req.finished_reason = None
+
+    scheduler.process_batch_result_disagg_prefill(_batch(req), _result())
+
+    scheduler.tree_cache.req_to_token_pool.mamba_allocator.free.assert_called_once()
+    assert req.kv.mamba_pool_idx is None
+    scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/managers/test_generation_auxiliary_output.py

@@ -443,6 +443,8 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit():
     req = SimpleNamespace(
         output_ids=[],
         finished_len=None,
+        to_finish=None,
+        finished_reason=None,
         inflight_middle_chunks=0,
         pending_bootstrap=False,
         return_logprob=False,