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

[perf] Assemble flat prompt top logprobs scheduler-side as numpy arrays (#32223)

Browse Source
This commit is contained in:
Sam Shleifer
2026-07-31 18:00:51 -07:00
committed by GitHub
parent ca07917c58
commit 58974ca16c
9 changed files with 593 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/detokenizer_manager.py
+3
View File
@@ -462,6 +462,9 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
output_token_logprobs_idx=recv_obj.output_token_logprobs_idx,
input_top_logprobs_val=recv_obj.input_top_logprobs_val,
input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
input_top_logprobs_val_flat=recv_obj.input_top_logprobs_val_flat,
input_top_logprobs_idx_flat=recv_obj.input_top_logprobs_idx_flat,
input_top_logprobs_flat_null_prefix=recv_obj.input_top_logprobs_flat_null_prefix,
output_top_logprobs_val=recv_obj.output_top_logprobs_val,
output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
input_token_ids_logprobs_val=recv_obj.input_token_ids_logprobs_val,
python/sglang/srt/managers/io_struct.py
+53
View File
@@ -38,6 +38,7 @@ from typing import (
List,
Literal,
Optional,
Tuple,
Type,
Union,
)
@@ -831,6 +832,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True):
stream: bool
# Whether to return sparse output-token support from top-k/top-p/min-p sampling.
return_sampling_mask: bool = False
# Assemble prompt top logprobs as flat arrays scheduler-side (see
# GenerateReqInput.return_flat_raw_top_logprobs). The b64 flag stays
# tokenizer-manager-side: the scheduler ships arrays either way.
return_flat_raw_top_logprobs: bool = False
# Whether to return hidden states
return_hidden_states: bool = False
@@ -1229,6 +1234,39 @@ CachedTokensDetails = Dict[str, Union[int, str]]
FinishReasonDict = Dict[str, Optional[Union[str, int, List[int]]]]
def build_flat_input_top_logprobs_arrays(
input_top_logprobs_val: List[Optional[List[float]]],
input_top_logprobs_idx: List[Optional[List[int]]],
top_logprobs_num: int,
) -> Tuple[np.ndarray, np.ndarray, int]:
"""Convert nested per-position prompt top logprob rows into the flat
arrays of the `return_flat_raw_top_logprobs` response format.
Returns (float32 values [rows, k], int32 token ids [rows, k],
null_prefix). The leading null rows are counted into null_prefix and
excluded from the arrays. Raises ValueError when the rows are not
representable by (shape, null_prefix): interior nulls or ragged k,
e.g. multi-item scoring.
"""
num_rows = len(input_top_logprobs_val)
null_prefix = 0
while null_prefix < num_rows and not input_top_logprobs_val[null_prefix]:
null_prefix += 1
val_rows = input_top_logprobs_val[null_prefix:]
idx_rows = input_top_logprobs_idx[null_prefix:]
k = len(val_rows[0]) if val_rows else top_logprobs_num
for offset, row in enumerate(val_rows):
if row is None or len(row) != k:
raise ValueError(
"return_flat_raw_top_logprobs requires rectangular top logprob "
f"rows with nulls only in the leading prefix; row {null_prefix + offset} "
f"has {None if row is None else len(row)} entries (expected {k})."
)
val_arr = np.asarray(val_rows, dtype=np.float32).reshape(len(val_rows), k)
idx_arr = np.asarray(idx_rows, dtype=np.int32).reshape(len(idx_rows), k)
return val_arr, idx_arr, null_prefix
class BatchTokenIDOutput(BaseBatchReq, kw_only=True):
# The finish reason
finished_reasons: List[Optional[FinishReasonDict]]
@@ -1319,6 +1357,15 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True):
spec_correct_drafts_histogram: Optional[List[List[int]]] = None
spec_cap_lens_histogram: Optional[List[List[int]]] = None
# Scheduler-side flat assembly of prompt top logprobs for requests with
# return_flat_raw_top_logprobs: float32 / int32 [rows, k] arrays plus the
# leading-null count (see build_flat_input_top_logprobs_arrays). For such
# requests the nested input_top_logprobs_val/idx entry is empty. None when
# no request in the batch uses the flat format.
input_top_logprobs_val_flat: Optional[List[Optional[np.ndarray]]] = None
input_top_logprobs_idx_flat: Optional[List[Optional[np.ndarray]]] = None
input_top_logprobs_flat_null_prefix: Optional[List[Optional[int]]] = None
class BatchStrOutput(BaseBatchReq, kw_only=True):
# The finish reason
@@ -1401,6 +1448,12 @@ class BatchStrOutput(BaseBatchReq, kw_only=True):
spec_correct_drafts_histogram: Optional[List[List[int]]] = None
spec_cap_lens_histogram: Optional[List[List[int]]] = None
# Detokenizer pass-through for the scheduler-side flat prompt top logprob
# arrays; see BatchTokenIDOutput.input_top_logprobs_val_flat.
input_top_logprobs_val_flat: Optional[List[Optional[np.ndarray]]] = None
input_top_logprobs_idx_flat: Optional[List[Optional[np.ndarray]]] = None
input_top_logprobs_flat_null_prefix: Optional[List[Optional[int]]] = None
class BatchEmbeddingOutput(BaseBatchReq, kw_only=True):
# The finish reason
python/sglang/srt/managers/multi_tokenizer_mixin.py
+18
View File
@@ -211,6 +211,15 @@ def _handle_output_by_index(output, i):
input_top_logprobs_idx=_extract_field_by_index(
output, "input_top_logprobs_idx", i, check_length=False
),
input_top_logprobs_val_flat=_extract_field_by_index(
output, "input_top_logprobs_val_flat", i, check_length=False
),
input_top_logprobs_idx_flat=_extract_field_by_index(
output, "input_top_logprobs_idx_flat", i, check_length=False
),
input_top_logprobs_flat_null_prefix=_extract_field_by_index(
output, "input_top_logprobs_flat_null_prefix", i, check_length=False
),
output_top_logprobs_val=_extract_field_by_index(
output, "output_top_logprobs_val", i, check_length=False
),
@@ -319,6 +328,15 @@ def _handle_output_by_index(output, i):
input_top_logprobs_idx=_extract_field_by_index(
output, "input_top_logprobs_idx", i, check_length=False
),
input_top_logprobs_val_flat=_extract_field_by_index(
output, "input_top_logprobs_val_flat", i, check_length=False
),
input_top_logprobs_idx_flat=_extract_field_by_index(
output, "input_top_logprobs_idx_flat", i, check_length=False
),
input_top_logprobs_flat_null_prefix=_extract_field_by_index(
output, "input_top_logprobs_flat_null_prefix", i, check_length=False
),
output_top_logprobs_val=_extract_field_by_index(
output, "output_top_logprobs_val", i, check_length=False
),
python/sglang/srt/managers/schedule_batch.py
+8
View File
@@ -687,6 +687,12 @@ class ReqLogprob:
input_token_logprobs_idx: Optional[List[int]] = None
input_top_logprobs_val: Optional[List[List[float]]] = None
input_top_logprobs_idx: Optional[List[List[int]]] = None
# Flat replacements for the rows above (see
# build_flat_input_top_logprobs_arrays); when set, the nested rows are
# emptied and the arrays ship instead.
input_top_logprobs_val_flat: Optional[np.ndarray] = None
input_top_logprobs_idx_flat: Optional[np.ndarray] = None
input_top_logprobs_flat_null_prefix: Optional[int] = None
input_token_ids_logprobs_val: Optional[List[List[float]]] = None
input_token_ids_logprobs_idx: Optional[List[List[int]]] = None
output_token_logprobs_val: Optional[list] = None
@@ -725,6 +731,7 @@ class Req(ReqDllmMixin):
dllm_config: Optional[DllmConfig] = None,
token_ids_logprob: List[int] = None,
return_sampling_mask: bool = False,
return_flat_raw_top_logprobs: bool = False,
stream: bool = False,
origin_input_ids_unpadded: Optional[array[int]] = None,
lora_id: Optional[str] = None,
@@ -943,6 +950,7 @@ class Req(ReqDllmMixin):
self.temp_scaled_logprobs = False
self.top_p_normalized_logprobs = False
self.return_sampling_mask = return_sampling_mask
self.return_flat_raw_top_logprobs = return_flat_raw_top_logprobs
# Logprobs (return values)
# True means the input logprob has been already sent to detokenizer.
python/sglang/srt/managers/scheduler.py
+1
View File
@@ -2242,6 +2242,7 @@ class Scheduler(
top_logprobs_num=recv_req.top_logprobs_num,
token_ids_logprob=recv_req.token_ids_logprob,
return_sampling_mask=recv_req.return_sampling_mask,
return_flat_raw_top_logprobs=recv_req.return_flat_raw_top_logprobs,
stream=recv_req.stream,
lora_id=recv_req.lora_id,
session_id=recv_req.session_id,
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py
+37
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import (
List,
@@ -10,6 +11,7 @@ import torch
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.io_struct import build_flat_input_top_logprobs_arrays
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.runtime_context import get_exec
from sglang.srt.server_args import (
@@ -17,6 +19,8 @@ from sglang.srt.server_args import (
ServerArgs,
)
logger = logging.getLogger(__name__)
@dataclass(kw_only=True, slots=True, frozen=True)
class SchedulerLogprobResultProcessor:
@@ -84,6 +88,35 @@ class SchedulerLogprobResultProcessor:
req.temp_input_top_logprobs_idx = None
req.temp_input_top_logprobs_val = None
def _flatten_input_top_logprobs(self, req: Req) -> None:
"""Replace the nested input top logprob rows with flat arrays for
requests that opted into return_flat_raw_top_logprobs, so the batch
output ships two ndarrays instead of num_positions * k python lists.
"""
if req.logprob.top_logprobs_num <= 0:
return
try:
(
req.logprob.input_top_logprobs_val_flat,
req.logprob.input_top_logprobs_idx_flat,
req.logprob.input_top_logprobs_flat_null_prefix,
) = build_flat_input_top_logprobs_arrays(
req.logprob.input_top_logprobs_val,
req.logprob.input_top_logprobs_idx,
req.logprob.top_logprobs_num,
)
except ValueError as e:
# Unrepresentable rows (e.g. multi-item scoring): keep the nested
# format, mirroring the tokenizer manager fallback.
logger.warning(
"Falling back to nested input top logprobs for rid=%s: %s",
req.rid,
e,
)
return
req.logprob.input_top_logprobs_val = []
req.logprob.input_top_logprobs_idx = []
def _process_input_token_ids_logprobs(self, req: Req) -> None:
"""Process input token IDs logprobs."""
if req.logprob.token_ids_logprob is None:
@@ -265,6 +298,10 @@ class SchedulerLogprobResultProcessor:
== relevant_tokens_len
)
# After the length checks: the flat arrays replace the nested rows.
if req.return_flat_raw_top_logprobs:
self._flatten_input_top_logprobs(req)
def add_logprob_return_values(
self,
i: int,
python/sglang/srt/managers/scheduler_components/output_streamer.py
+40
View File
@@ -312,6 +312,12 @@ class _GenerationStreamAccumulator:
output_token_logprobs_idx: Optional[list] = None
input_top_logprobs_val: Optional[list] = None
input_top_logprobs_idx: Optional[list] = None
# Per-request flat prompt top logprob arrays (return_flat_raw_top_logprobs);
# None entries for requests on the nested format.
input_top_logprobs_val_flat: Optional[list] = None
input_top_logprobs_idx_flat: Optional[list] = None
input_top_logprobs_flat_null_prefix: Optional[list] = None
has_input_top_logprobs_flat: bool = False
output_top_logprobs_val: Optional[list] = None
output_top_logprobs_idx: Optional[list] = None
input_token_ids_logprobs_val: Optional[list] = None
@@ -340,6 +346,9 @@ class _GenerationStreamAccumulator:
self.output_token_logprobs_idx = []
self.input_top_logprobs_val = []
self.input_top_logprobs_idx = []
self.input_top_logprobs_val_flat = []
self.input_top_logprobs_idx_flat = []
self.input_top_logprobs_flat_null_prefix = []
self.output_top_logprobs_val = []
self.output_top_logprobs_idx = []
self.input_token_ids_logprobs_val = []
@@ -464,6 +473,17 @@ class _GenerationStreamAccumulator:
)
self.input_top_logprobs_val.append(req.logprob.input_top_logprobs_val)
self.input_top_logprobs_idx.append(req.logprob.input_top_logprobs_idx)
self.input_top_logprobs_val_flat.append(
req.logprob.input_top_logprobs_val_flat
)
self.input_top_logprobs_idx_flat.append(
req.logprob.input_top_logprobs_idx_flat
)
self.input_top_logprobs_flat_null_prefix.append(
req.logprob.input_top_logprobs_flat_null_prefix
)
if req.logprob.input_top_logprobs_val_flat is not None:
self.has_input_top_logprobs_flat = True
self.input_token_ids_logprobs_val.append(
req.logprob.input_token_ids_logprobs_val
)
@@ -476,6 +496,9 @@ class _GenerationStreamAccumulator:
self.input_token_logprobs_idx.append([])
self.input_top_logprobs_val.append([])
self.input_top_logprobs_idx.append([])
self.input_top_logprobs_val_flat.append(None)
self.input_top_logprobs_idx_flat.append(None)
self.input_top_logprobs_flat_null_prefix.append(None)
self.input_token_ids_logprobs_val.append([])
self.input_token_ids_logprobs_idx.append([])
@@ -613,6 +636,23 @@ class _GenerationStreamAccumulator:
output_token_logprobs_idx=self.output_token_logprobs_idx,
input_top_logprobs_val=self.input_top_logprobs_val,
input_top_logprobs_idx=self.input_top_logprobs_idx,
# None on the common path so the wire payload is unchanged when no
# request in the batch uses the flat format.
input_top_logprobs_val_flat=(
self.input_top_logprobs_val_flat
if self.has_input_top_logprobs_flat
else None
),
input_top_logprobs_idx_flat=(
self.input_top_logprobs_idx_flat
if self.has_input_top_logprobs_flat
else None
),
input_top_logprobs_flat_null_prefix=(
self.input_top_logprobs_flat_null_prefix
if self.has_input_top_logprobs_flat
else None
),
output_top_logprobs_val=self.output_top_logprobs_val,
output_top_logprobs_idx=self.output_top_logprobs_idx,
input_token_ids_logprobs_val=self.input_token_ids_logprobs_val,
python/sglang/srt/managers/tokenizer_manager.py
+66 -21
View File
@@ -84,6 +84,7 @@ from sglang.srt.managers.io_struct import (
UpdateWeightFromDiskReqOutput,
async_sock_recv,
async_sock_send,
build_flat_input_top_logprobs_arrays,
sock_send,
unwrap_from_pickle,
)
@@ -233,6 +234,12 @@ class ReqState:
# prefill chunks arrive, so streaming decode chunks reuse the payload.
input_top_logprobs_flat_fields: Optional[Dict[str, Any]] = None
input_top_logprobs_flat_num_rows: int = -1
# Scheduler-assembled flat arrays (val float32 [rows, k], idx int32
# [rows, k], null_prefix), sent once at prefill completion. When present,
# the nested input_top_logprobs_val/idx above stay empty.
input_top_logprobs_scheduler_flat: Optional[Tuple[np.ndarray, np.ndarray, int]] = (
None
)
# For detokenized logprobs
input_token_logprobs: List[Any] = dataclasses.field(default_factory=list)
@@ -278,26 +285,38 @@ def _build_flat_input_top_logprobs_fields(
base64 contiguous little-endian binary; the dtype marker fields let the
widths change later without a wire break.
"""
num_rows = len(input_top_logprobs_val)
null_prefix = 0
while null_prefix < num_rows and not input_top_logprobs_val[null_prefix]:
null_prefix += 1
val_rows = input_top_logprobs_val[null_prefix:]
idx_rows = input_top_logprobs_idx[null_prefix:]
k = len(val_rows[0]) if val_rows else top_logprobs_num
for offset, row in enumerate(val_rows):
if row is None or len(row) != k:
# Not representable by (shape, null_prefix); e.g. multi-item scoring.
raise ValueError(
"return_flat_raw_top_logprobs requires rectangular top logprob "
f"rows with nulls only in the leading prefix; row {null_prefix + offset} "
f"has {None if row is None else len(row)} entries (expected {k})."
)
val_arr, idx_arr, null_prefix = build_flat_input_top_logprobs_arrays(
input_top_logprobs_val, input_top_logprobs_idx, top_logprobs_num
)
if return_b64:
return _build_flat_input_top_logprobs_fields_from_arrays(
val_arr, idx_arr, null_prefix, return_b64=True
)
# Flatten the original python rows so the JSON numbers keep their full
# (float64) precision, matching the pre-scheduler-flat output.
return {
"input_top_logprobs_val_flat": [
v for row in input_top_logprobs_val[null_prefix:] for v in row
],
"input_top_logprobs_idx_flat": [
i for row in input_top_logprobs_idx[null_prefix:] for i in row
],
"input_top_logprobs_shape": [val_arr.shape[0], val_arr.shape[1]],
"input_top_logprobs_null_prefix": null_prefix,
}
def _build_flat_input_top_logprobs_fields_from_arrays(
val_arr: np.ndarray,
idx_arr: np.ndarray,
null_prefix: int,
return_b64: bool = False,
) -> Dict[str, Any]:
"""Build the flat response fields from scheduler-assembled [rows, k]
arrays (see `_build_flat_input_top_logprobs_fields` for the field
semantics)."""
fields: Dict[str, Any] = {}
if return_b64:
val_arr = np.asarray(val_rows, dtype=np.float32)
idx_arr = np.asarray(idx_rows, dtype=np.int32)
fields["input_top_logprobs_val_flat_b64"] = pybase64.b64encode(
val_arr.tobytes()
).decode("utf-8")
@@ -307,9 +326,9 @@ def _build_flat_input_top_logprobs_fields(
fields["input_top_logprobs_val_flat_b64_dtype"] = "float32"
fields["input_top_logprobs_idx_flat_b64_dtype"] = "int32"
else:
fields["input_top_logprobs_val_flat"] = [v for row in val_rows for v in row]
fields["input_top_logprobs_idx_flat"] = [i for row in idx_rows for i in row]
fields["input_top_logprobs_shape"] = [len(val_rows), k]
fields["input_top_logprobs_val_flat"] = val_arr.reshape(-1).tolist()
fields["input_top_logprobs_idx_flat"] = idx_arr.reshape(-1).tolist()
fields["input_top_logprobs_shape"] = [val_arr.shape[0], val_arr.shape[1]]
fields["input_top_logprobs_null_prefix"] = null_prefix
return fields
@@ -1264,6 +1283,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
top_logprobs_num=obj.top_logprobs_num,
token_ids_logprob=obj.token_ids_logprob,
return_sampling_mask=obj.return_sampling_mask,
return_flat_raw_top_logprobs=obj.return_flat_raw_top_logprobs,
stream=obj.stream,
rid=obj.rid,
http_worker_ipc=obj.http_worker_ipc,
@@ -2297,7 +2317,23 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# Guarded by the caller's return_logprob check, so obj is a
# GenerateReqInput here.
use_flat = state.obj.return_flat_raw_top_logprobs
if use_flat:
if use_flat and state.input_top_logprobs_scheduler_flat is not None:
# The scheduler already assembled the flat arrays (sent once
# at prefill completion); encode them directly.
if state.input_top_logprobs_flat_fields is None:
val_arr, idx_arr, null_prefix = (
state.input_top_logprobs_scheduler_flat
)
state.input_top_logprobs_flat_fields = (
_build_flat_input_top_logprobs_fields_from_arrays(
val_arr,
idx_arr,
null_prefix,
return_b64=state.obj.return_flat_raw_top_logprobs_b64,
)
)
meta_info.update(state.input_top_logprobs_flat_fields)
elif use_flat:
# Flat replaces nested for the input side only.
if state.input_top_logprobs_flat_num_rows != len(
state.input_top_logprobs_val
@@ -2424,6 +2460,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
state.input_top_logprobs_idx.extend(
recv_obj.input_top_logprobs_idx[recv_obj_index]
)
if (
recv_obj.input_top_logprobs_val_flat is not None
and recv_obj.input_top_logprobs_val_flat[recv_obj_index] is not None
):
state.input_top_logprobs_scheduler_flat = (
recv_obj.input_top_logprobs_val_flat[recv_obj_index],
recv_obj.input_top_logprobs_idx_flat[recv_obj_index],
recv_obj.input_top_logprobs_flat_null_prefix[recv_obj_index],
)
state.output_top_logprobs_val.extend(
recv_obj.output_top_logprobs_val[recv_obj_index]
)