minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Refactor batch_result_processor into per-step prefill/decode helpers (#25709)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-19 09:17:17 +08:00
committed by GitHub
parent 7e7cb969e9
commit 0e198f0f4f
1 changed files with 291 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+291 -173
View File
@@ -7,6 +7,7 @@ from typing import (
Callable,
List,
Optional,
Tuple,
Union,
)
@@ -202,27 +203,7 @@ class SchedulerBatchResultProcessor:
# Move next_token_ids and logprobs to cpu
next_token_ids = next_token_ids.tolist()
if batch.return_logprob:
if logits_output.next_token_logprobs is not None:
logits_output.next_token_logprobs = (
logits_output.next_token_logprobs.tolist()
)
if logits_output.input_token_logprobs is not None:
logits_output.input_token_logprobs = tuple(
logits_output.input_token_logprobs.tolist()
)
if logits_output.next_token_top_logprobs_val:
logits_output.next_token_top_logprobs_val = [
v.tolist() for v in logits_output.next_token_top_logprobs_val
]
logits_output.next_token_top_logprobs_idx = [
x.tolist() for x in logits_output.next_token_top_logprobs_idx
]
if logits_output.next_token_token_ids_logprobs_val:
logits_output.next_token_token_ids_logprobs_val = [
v.tolist()
for v in logits_output.next_token_token_ids_logprobs_val
]
self._move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
hidden_state_offset = 0
@@ -256,58 +237,30 @@ class SchedulerBatchResultProcessor:
self._maybe_collect_customized_info(i, req, logits_output)
if batch.return_logprob:
assert extend_logprob_start_len_per_req is not None
assert extend_input_len_per_req is not None
extend_logprob_start_len = extend_logprob_start_len_per_req[i]
extend_input_len = extend_input_len_per_req[i]
num_input_logprobs = (
self.logprob_result_processor.calculate_num_input_logprobs(
req,
extend_input_len,
extend_logprob_start_len,
)
logprob_pt = self._apply_prefill_logprobs(
req=req,
i=i,
logits_output=logits_output,
extend_input_len_per_req=extend_input_len_per_req,
extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
next_token_ids=next_token_ids,
logprob_pt=logprob_pt,
)
if req.return_logprob:
self.logprob_result_processor.add_logprob_return_values(
i,
req,
logprob_pt,
next_token_ids,
num_input_logprobs,
logits_output,
)
logprob_pt += num_input_logprobs
if (
req.return_hidden_states
and logits_output.hidden_states is not None
):
req.hidden_states.append(
logits_output.hidden_states[
hidden_state_offset : (
hidden_state_offset := hidden_state_offset
+ len(req.origin_input_ids)
)
]
.cpu()
.clone()
.tolist()
hidden_state_offset = self._append_prefill_hidden_states(
req=req,
logits_output=logits_output,
hidden_state_offset=hidden_state_offset,
)
if req.grammar is not None:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
try:
req.grammar.accept_token(next_token_id)
except ValueError as e:
# 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.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
)
self.abort_request(AbortReq(rid=req.rid))
req.grammar.finished = req.finished()
self._apply_prefill_grammar(
req=req, next_token_id=next_token_id
)
else:
# being chunked reqs' prefill is not finished
@@ -319,25 +272,14 @@ class SchedulerBatchResultProcessor:
# Incrementally update input logprobs.
if batch.return_logprob:
extend_logprob_start_len = extend_logprob_start_len_per_req[i]
extend_input_len = extend_input_len_per_req[i]
if extend_logprob_start_len < extend_input_len:
# Update input logprobs.
num_input_logprobs = self.logprob_result_processor.calculate_num_input_logprobs(
req,
extend_input_len,
extend_logprob_start_len,
)
if req.return_logprob:
self.logprob_result_processor.add_input_logprob_return_values(
i,
req,
logits_output,
logprob_pt,
num_input_logprobs,
last_prefill_chunk=False,
)
logprob_pt += num_input_logprobs
logprob_pt = self._apply_chunked_prefill_logprobs(
req=req,
i=i,
logits_output=logits_output,
extend_input_len_per_req=extend_input_len_per_req,
extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
logprob_pt=logprob_pt,
)
req.time_stats.set_last_chunked_prefill_finish_time()
@@ -345,26 +287,9 @@ class SchedulerBatchResultProcessor:
if result.copy_done is not None:
result.copy_done.synchronize()
is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
embeddings = result.embeddings
embeddings = self._convert_embeddings(result=result)
phs = result.pooled_hidden_states
if is_sparse:
batch_ids, token_ids = embeddings.indices()
values = embeddings.values()
embeddings = [{} for _ in range(embeddings.size(0))]
for i in range(batch_ids.shape[0]):
embeddings[batch_ids[i].item()][token_ids[i].item()] = values[
i
].item()
else:
if isinstance(embeddings, torch.Tensor):
embeddings = embeddings.tolist()
else:
embeddings = [tensor.tolist() for tensor in embeddings]
if phs is not None:
if isinstance(phs, list):
phs = [t.cpu().detach() for t in phs]
@@ -407,6 +332,152 @@ class SchedulerBatchResultProcessor:
dp_cooperation_info=batch.dp_cooperation_info,
)
def _convert_embeddings(self, *, result: EmbeddingBatchResult) -> list:
is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
embeddings = result.embeddings
if is_sparse:
batch_ids, token_ids = embeddings.indices()
values = embeddings.values()
embeddings = [{} for _ in range(embeddings.size(0))]
for i in range(batch_ids.shape[0]):
embeddings[batch_ids[i].item()][token_ids[i].item()] = values[i].item()
else:
if isinstance(embeddings, torch.Tensor):
embeddings = embeddings.tolist()
else:
embeddings = [tensor.tolist() for tensor in embeddings]
return embeddings
def _move_logprobs_to_cpu(
self,
*,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
) -> None:
if batch.return_logprob:
if logits_output.next_token_logprobs is not None:
logits_output.next_token_logprobs = (
logits_output.next_token_logprobs.tolist()
)
if logits_output.input_token_logprobs is not None:
logits_output.input_token_logprobs = tuple(
logits_output.input_token_logprobs.tolist()
)
if logits_output.next_token_top_logprobs_val:
logits_output.next_token_top_logprobs_val = [
v.tolist() for v in logits_output.next_token_top_logprobs_val
]
logits_output.next_token_top_logprobs_idx = [
x.tolist() for x in logits_output.next_token_top_logprobs_idx
]
if logits_output.next_token_token_ids_logprobs_val:
logits_output.next_token_token_ids_logprobs_val = [
v.tolist() for v in logits_output.next_token_token_ids_logprobs_val
]
def _apply_prefill_logprobs(
self,
*,
req: Req,
i: int,
logits_output: LogitsProcessorOutput,
extend_input_len_per_req: Optional[List[int]],
extend_logprob_start_len_per_req: Optional[List[int]],
next_token_ids: List[int],
logprob_pt: int,
) -> int:
assert extend_logprob_start_len_per_req is not None
assert extend_input_len_per_req is not None
extend_logprob_start_len = extend_logprob_start_len_per_req[i]
extend_input_len = extend_input_len_per_req[i]
num_input_logprobs = self.logprob_result_processor.calculate_num_input_logprobs(
req,
extend_input_len,
extend_logprob_start_len,
)
if req.return_logprob:
self.logprob_result_processor.add_logprob_return_values(
i,
req,
logprob_pt,
next_token_ids,
num_input_logprobs,
logits_output,
)
logprob_pt += num_input_logprobs
return logprob_pt
def _append_prefill_hidden_states(
self,
*,
req: Req,
logits_output: LogitsProcessorOutput,
hidden_state_offset: int,
) -> int:
req.hidden_states.append(
logits_output.hidden_states[
hidden_state_offset : (
hidden_state_offset := hidden_state_offset
+ len(req.origin_input_ids)
)
]
.cpu()
.clone()
.tolist()
)
return hidden_state_offset
def _apply_prefill_grammar(self, *, req: Req, next_token_id: int) -> None:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
try:
req.grammar.accept_token(next_token_id)
except ValueError as e:
# 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.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
)
self.abort_request(AbortReq(rid=req.rid))
req.grammar.finished = req.finished()
def _apply_chunked_prefill_logprobs(
self,
*,
req: Req,
i: int,
logits_output: LogitsProcessorOutput,
extend_input_len_per_req: Optional[List[int]],
extend_logprob_start_len_per_req: Optional[List[int]],
logprob_pt: int,
) -> int:
extend_logprob_start_len = extend_logprob_start_len_per_req[i]
extend_input_len = extend_input_len_per_req[i]
if extend_logprob_start_len < extend_input_len:
# Update input logprobs.
num_input_logprobs = (
self.logprob_result_processor.calculate_num_input_logprobs(
req,
extend_input_len,
extend_logprob_start_len,
)
)
if req.return_logprob:
self.logprob_result_processor.add_input_logprob_return_values(
i,
req,
logits_output,
logprob_pt,
num_input_logprobs,
last_prefill_chunk=False,
)
logprob_pt += num_input_logprobs
return logprob_pt
def _resolve_spec_overlap_tokens(
self,
result: GenerationBatchResult,
@@ -488,31 +559,12 @@ class SchedulerBatchResultProcessor:
result.can_run_cuda_graph,
)
if batch.spec_algorithm.is_none() or batch.is_spec_v2:
if batch.is_spec_v2:
next_token_ids = self._resolve_spec_overlap_tokens(result, batch)
elif isinstance(next_token_ids, list):
pass # MLX path: already a list[int], skip torch round-trip
else:
next_token_ids = next_token_ids.tolist()
if batch.return_logprob:
next_token_logprobs = logits_output.next_token_logprobs.tolist()
if logits_output.next_token_top_logprobs_val:
logits_output.next_token_top_logprobs_val = [
v.tolist() for v in logits_output.next_token_top_logprobs_val
]
logits_output.next_token_top_logprobs_idx = [
x.tolist() for x in logits_output.next_token_top_logprobs_idx
]
if logits_output.next_token_token_ids_logprobs_val:
logits_output.next_token_token_ids_logprobs_val = [
v.tolist()
for v in logits_output.next_token_token_ids_logprobs_val
]
# else: Spec V1 — output_ids, check_finished, grammar, and reasoning tokens
# are already handled in the verify phase (eagle_info.py / ngram_info.py).
next_token_ids, next_token_logprobs = self._normalize_decode_outputs(
batch=batch,
result=result,
logits_output=logits_output,
next_token_ids=next_token_ids,
)
self.metrics_reporter.num_generated_tokens += len(batch.reqs)
if not batch.spec_algorithm.is_none():
@@ -571,36 +623,14 @@ class SchedulerBatchResultProcessor:
self._handle_finished_req(req, i, logits_output)
if req.return_logprob:
# Spec v1 handles logprobs inside its own worker.
# Normalize: non-spec has 1 token, spec v2 has multiple.
if batch.is_spec_v2:
accepted_logprobs = next_token_logprobs[i]
accepted_ids = next_token_id
max_accept = len(accepted_logprobs)
else:
accepted_logprobs = [next_token_logprobs[i]]
accepted_ids = [next_token_id]
max_accept = 1
for j, tok_id in enumerate(accepted_ids):
req.output_token_logprobs_val.append(accepted_logprobs[j])
req.output_token_logprobs_idx.append(tok_id)
if req.top_logprobs_num > 0:
flat_idx = i * max_accept + j
req.output_top_logprobs_val.append(
logits_output.next_token_top_logprobs_val[flat_idx]
)
req.output_top_logprobs_idx.append(
logits_output.next_token_top_logprobs_idx[flat_idx]
)
if req.token_ids_logprob is not None:
flat_idx = i * max_accept + j
req.output_token_ids_logprobs_val.append(
logits_output.next_token_token_ids_logprobs_val[flat_idx]
)
req.output_token_ids_logprobs_idx.append(
logits_output.next_token_token_ids_logprobs_idx[flat_idx]
)
self._apply_decode_logprobs(
req=req,
i=i,
batch=batch,
next_token_id=next_token_id,
next_token_logprobs=next_token_logprobs,
logits_output=logits_output,
)
if req.return_hidden_states and logits_output.hidden_states is not None:
req.hidden_states.append(
@@ -608,23 +638,9 @@ class SchedulerBatchResultProcessor:
)
if req.grammar is not None:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
try:
if batch.spec_algorithm.is_none():
# Normal decode: single token
req.grammar.accept_token(next_token_id)
elif batch.is_spec_v2:
# Speculative decode: next_token_id is a list of accepted tokens
for token_id in next_token_id:
req.grammar.accept_token(token_id)
except ValueError as e:
# 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.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
)
self.abort_request(AbortReq(rid=req.rid))
req.grammar.finished = req.finished()
self._apply_decode_grammar(
req=req, next_token_id=next_token_id, batch=batch
)
self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
self.token_to_kv_pool_allocator.free_group_end()
@@ -638,6 +654,108 @@ class SchedulerBatchResultProcessor:
num_correct_drafts=result.num_correct_drafts,
)
def _normalize_decode_outputs(
self,
*,
batch: ScheduleBatch,
result: GenerationBatchResult,
logits_output: LogitsProcessorOutput,
next_token_ids: Union[torch.Tensor, List[int]],
) -> Tuple[Union[List[int], List[List[int]]], Optional[List[float]]]:
next_token_logprobs = None
if batch.spec_algorithm.is_none() or batch.is_spec_v2:
if batch.is_spec_v2:
next_token_ids = self._resolve_spec_overlap_tokens(result, batch)
elif isinstance(next_token_ids, list):
pass # MLX path: already a list[int], skip torch round-trip
else:
next_token_ids = next_token_ids.tolist()
if batch.return_logprob:
next_token_logprobs = logits_output.next_token_logprobs.tolist()
if logits_output.next_token_top_logprobs_val:
logits_output.next_token_top_logprobs_val = [
v.tolist() for v in logits_output.next_token_top_logprobs_val
]
logits_output.next_token_top_logprobs_idx = [
x.tolist() for x in logits_output.next_token_top_logprobs_idx
]
if logits_output.next_token_token_ids_logprobs_val:
logits_output.next_token_token_ids_logprobs_val = [
v.tolist()
for v in logits_output.next_token_token_ids_logprobs_val
]
# else: Spec V1 — output_ids, check_finished, grammar, and reasoning tokens
# are already handled in the verify phase (eagle_info.py / ngram_info.py).
return next_token_ids, next_token_logprobs
def _apply_decode_logprobs(
self,
*,
req: Req,
i: int,
batch: ScheduleBatch,
next_token_id: Union[int, List[int]],
next_token_logprobs: list,
logits_output: LogitsProcessorOutput,
) -> None:
# Spec v1 handles logprobs inside its own worker.
# Normalize: non-spec has 1 token, spec v2 has multiple.
if batch.is_spec_v2:
accepted_logprobs = next_token_logprobs[i]
accepted_ids = next_token_id
max_accept = len(accepted_logprobs)
else:
accepted_logprobs = [next_token_logprobs[i]]
accepted_ids = [next_token_id]
max_accept = 1
for j, tok_id in enumerate(accepted_ids):
req.output_token_logprobs_val.append(accepted_logprobs[j])
req.output_token_logprobs_idx.append(tok_id)
if req.top_logprobs_num > 0:
flat_idx = i * max_accept + j
req.output_top_logprobs_val.append(
logits_output.next_token_top_logprobs_val[flat_idx]
)
req.output_top_logprobs_idx.append(
logits_output.next_token_top_logprobs_idx[flat_idx]
)
if req.token_ids_logprob is not None:
flat_idx = i * max_accept + j
req.output_token_ids_logprobs_val.append(
logits_output.next_token_token_ids_logprobs_val[flat_idx]
)
req.output_token_ids_logprobs_idx.append(
logits_output.next_token_token_ids_logprobs_idx[flat_idx]
)
def _apply_decode_grammar(
self,
*,
req: Req,
next_token_id: Union[int, List[int]],
batch: ScheduleBatch,
) -> None:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
try:
if batch.spec_algorithm.is_none():
# Normal decode: single token
req.grammar.accept_token(next_token_id)
elif batch.is_spec_v2:
# Speculative decode: next_token_id is a list of accepted tokens
for token_id in next_token_id:
req.grammar.accept_token(token_id)
except ValueError as e:
# 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.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
)
self.abort_request(AbortReq(rid=req.rid))
req.grammar.finished = req.finished()
def _handle_finished_req(
self,
req: Req,