diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 6003a4c55..3360c858c 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -693,7 +693,6 @@ class SchedulerDisaggregationPrefillMixin: 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}" - release_kv_cache(req, self.tree_cache) prepare_abort( req, error_message, @@ -812,7 +811,8 @@ class SchedulerDisaggregationPrefillMixin: undone_reqs.append(req) elif poll == KVPoll.Success: # transfer done release_kv_cache(req, self.tree_cache) # unlock the tree - req.finished_reason = FINISH_LENGTH(length=0) + if not isinstance(req.finished_reason, FINISH_ABORT): + req.finished_reason = FINISH_LENGTH(length=0) # FIXME: clean up req's data in transfer engine req.disagg_kv_sender.clear() done_reqs.append(req) @@ -884,7 +884,10 @@ class SchedulerDisaggregationPrefillMixin: logger.warning(error_message) req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) release_kv_cache(req, self.tree_cache) # unlock the tree - prepare_abort(req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR) + if not isinstance(req.finished_reason, FINISH_ABORT): + prepare_abort( + req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR + ) if self.metrics_reporter.enable_metrics: self.metrics_collector.increment_transfer_failed_reqs() return exc