fix(rollout): compact_logprobs + logprobs:0 + server-side text#247
fix(rollout): compact_logprobs + logprobs:0 + server-side text#247aoshen02 wants to merge 1 commit into
Conversation
Three rollout performance/correctness fixes: 1. Request logprobs:0 instead of logprobs:1 — the per-token logprob over-fetch added ~13.6% serialization overhead with no consumer. 2. Parse compact_logprobs response field when present, avoiding the nested logprobs.content[].logprob object traversal. 3. Prefer server-returned text over client-side tokenizer.decode(), offloading decode work to the engine and using asyncio.to_thread as fallback to avoid blocking the event loop. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates vime/rollout/vllm_rollout.py to support parsing compact_logprobs and text directly from the vLLM response, falling back to standard token_ids and logprobs parsing and asynchronous token decoding if they are not present. Additionally, the default logprobs in sampling parameters is changed from 1 to 0. The reviewer suggests wrapping the compact_logprobs parsing logic in a try-except block to prevent potential crashes from malformed server responses.
| clp = choice.get("compact_logprobs") | ||
| if isinstance(clp, list) and clp: | ||
| new_response_tokens = [int(pair[1]) for pair in clp] | ||
| new_response_log_probs = [float(pair[0]) for pair in clp] | ||
| else: | ||
| new_response_tokens = choice.get("token_ids") or [] | ||
| new_response_log_probs: list[float] = [] | ||
| lp = choice.get("logprobs") | ||
| if isinstance(lp, dict): | ||
| content_items = lp.get("content") or [] | ||
| new_response_log_probs = [ | ||
| float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items | ||
| ] | ||
| if not new_response_log_probs: | ||
| new_response_log_probs = [0.0] * len(new_response_tokens) |
There was a problem hiding this comment.
Parsing compact_logprobs directly without a try-except block is risky. If the server returns a malformed list, elements that are not pairs, or values that cannot be cast to int or float, it will raise an exception (such as IndexError, TypeError, or ValueError) and crash the entire rollout process. Wrapping the parsing logic in a try-except block and falling back to the standard token_ids and logprobs parsing ensures robustness.
clp = choice.get("compact_logprobs")
parsed_clp = False
if isinstance(clp, list) and clp:
try:
new_response_tokens = [int(pair[1]) for pair in clp]
new_response_log_probs = [float(pair[0]) for pair in clp]
parsed_clp = True
except (IndexError, TypeError, ValueError):
pass
if not parsed_clp:
new_response_tokens = choice.get("token_ids") or []
new_response_log_probs: list[float] = []
lp = choice.get("logprobs")
if isinstance(lp, dict):
content_items = lp.get("content") or []
new_response_log_probs = [
float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items
]
if not new_response_log_probs:
new_response_log_probs = [0.0] * len(new_response_tokens)
1 participant
Summary
logprobs: 1→logprobs: 0in_build_inference_sampling_params. The previous value caused vLLM to compute and serialize per-token log probabilities that no downstream consumer reads, adding ~13.6% serialization overhead (measured via bare-serve A/B bench).compact_logprobsresponse field (list of[logprob, token_id]pairs) when present, avoiding the deeply nestedlogprobs.content[i].logprobobject traversal.choice.textreturned by the engine over client-sidetokenizer.decode(). Falls back toasyncio.to_thread(decode)to avoid blocking the event loop when the server doesn't return text.Test plan
🤖 Generated with Claude Code