[PD-Disagg] Fix double free when prebuilt batch is aborted. (#18822)
This commit is contained in:
@@ -6,8 +6,6 @@ from typing import TYPE_CHECKING
|
|||||||
|
|
||||||
import torch
|
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.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
|
||||||
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
||||||
|
|
||||||
@@ -120,12 +118,15 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
|||||||
if req.grammar.current_token is None:
|
if req.grammar.current_token is None:
|
||||||
req.grammar.accept_token(req.output_ids[-1])
|
req.grammar.accept_token(req.output_ids[-1])
|
||||||
except ValueError as e:
|
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.
|
# 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.
|
# 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}"
|
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)
|
req.to_finish = FINISH_ABORT(
|
||||||
prepare_abort(
|
error_message, HTTPStatus.INTERNAL_SERVER_ERROR
|
||||||
req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
|
|
||||||
)
|
)
|
||||||
req.grammar.finished = req.finished()
|
req.grammar.finished = req.finished()
|
||||||
self.output_ids = torch.tensor(self.output_ids, device=self.device)
|
self.output_ids = torch.tensor(self.output_ids, device=self.device)
|
||||||
|
|||||||
Reference in New Issue
Block a user