diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 81bdb722a..342296585 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -6,8 +6,6 @@ from typing import TYPE_CHECKING import torch -from sglang.srt.disaggregation.utils import prepare_abort -from sglang.srt.mem_cache.common import release_kv_cache from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo @@ -120,12 +118,15 @@ class ScheduleBatchDisaggregationDecodeMixin: if req.grammar.current_token is None: req.grammar.accept_token(req.output_ids[-1]) except ValueError as e: + from sglang.srt.managers.schedule_batch import FINISH_ABORT + # Grammar accept_token can raise ValueError if the token is not in the grammar. # This can happen if the grammar is not set correctly or the token is invalid. + # Use to_finish (not finished_reason) so that process_batch_result_prebuilt + # handles the release via check_finished -> release_kv_cache in one place. error_message = f"Grammar accept_token failed for req {req.rid} with token {req.output_ids[-1]}: {e}" - release_kv_cache(req, self.tree_cache) - prepare_abort( - req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR + req.to_finish = FINISH_ABORT( + error_message, HTTPStatus.INTERNAL_SERVER_ERROR ) req.grammar.finished = req.finished() self.output_ids = torch.tensor(self.output_ids, device=self.device)