[router][bugfix] Fix input_logprobs handling with None value and logprob_start_len = -1 (#11113)

This commit is contained in:
Chang Su
2025-09-30 16:09:40 -07:00
committed by GitHub
parent fb367acfcb
commit 8ce830a8b0
7 changed files with 237 additions and 134 deletions

View File

@@ -415,7 +415,11 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
mm_inputs=None, # TODO: implement mm support
sampling_params=sampling_params,
return_logprob=grpc_req.return_logprob,
logprob_start_len=grpc_req.logprob_start_len or -1,
logprob_start_len=(
grpc_req.logprob_start_len
if grpc_req.logprob_start_len is not None
else -1
),
top_logprobs_num=grpc_req.top_logprobs_num or 0,
stream=grpc_req.stream or False,
lora_id=grpc_req.lora_id if grpc_req.lora_id else None,
@@ -486,10 +490,10 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
ignore_eos=grpc_params.ignore_eos,
)
def _convert_logprobs_to_proto(
def _convert_output_logprobs_to_proto(
self, logprobs_data: Dict
) -> Optional[sglang_scheduler_pb2.LogProbs]:
"""Convert logprobs dict to proto LogProbs format (transport RAW data only)."""
) -> Optional[sglang_scheduler_pb2.OutputLogProbs]:
"""Convert output logprobs dict to proto (no None values, plain floats)."""
if not logprobs_data:
return None
@@ -509,8 +513,47 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
)
)
return sglang_scheduler_pb2.LogProbs(
token_logprobs=token_logprobs_val,
return sglang_scheduler_pb2.OutputLogProbs(
token_logprobs=token_logprobs_val, # Plain float array
token_ids=token_logprobs_idx,
top_logprobs=top_logprobs_proto,
)
def _convert_input_logprobs_to_proto(
self, logprobs_data: Dict
) -> Optional[sglang_scheduler_pb2.InputLogProbs]:
"""Convert input logprobs dict to proto (first token is None, wrapped in InputTokenLogProb)."""
if not logprobs_data:
return None
token_logprobs_val = logprobs_data.get("token_logprobs_val", [])
token_logprobs_idx = logprobs_data.get("token_logprobs_idx", [])
top_logprobs_val = logprobs_data.get("top_logprobs_val", [])
top_logprobs_idx = logprobs_data.get("top_logprobs_idx", [])
# Wrap values in InputTokenLogProb (None for first token, value for others)
token_logprobs_wrapped = [
(
sglang_scheduler_pb2.InputTokenLogProb()
if x is None
else sglang_scheduler_pb2.InputTokenLogProb(value=x)
)
for x in token_logprobs_val
]
# Build TopLogProbs entries
top_logprobs_proto = []
if top_logprobs_val and top_logprobs_idx:
for val_list, idx_list in zip(top_logprobs_val, top_logprobs_idx):
top_logprobs_proto.append(
sglang_scheduler_pb2.TopLogProbs(
values=val_list,
token_ids=idx_list,
)
)
return sglang_scheduler_pb2.InputLogProbs(
token_logprobs=token_logprobs_wrapped,
token_ids=token_logprobs_idx,
top_logprobs=top_logprobs_proto,
)
@@ -522,12 +565,12 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
meta_info = output.get("meta_info", {})
# Convert output logprobs if present
output_logprobs_proto = self._convert_logprobs_to_proto(
output_logprobs_proto = self._convert_output_logprobs_to_proto(
output.get("output_logprobs")
)
# Convert input logprobs if present (only in first chunk)
input_logprobs_proto = self._convert_logprobs_to_proto(
input_logprobs_proto = self._convert_input_logprobs_to_proto(
output.get("input_logprobs")
)
@@ -576,12 +619,12 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
matched_stop_kwargs["matched_stop_str"] = matched
# Convert output logprobs if present
output_logprobs_proto = self._convert_logprobs_to_proto(
output_logprobs_proto = self._convert_output_logprobs_to_proto(
output.get("output_logprobs")
)
# Convert input logprobs if present
input_logprobs_proto = self._convert_logprobs_to_proto(
input_logprobs_proto = self._convert_input_logprobs_to_proto(
output.get("input_logprobs")
)