[Feat]: Support output input logprobs - #4793
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds support for scoring input-token logprobs via the direct /generate endpoint when using the PyTorch hybrid backend with --logprobs-mode raw_logits|raw_logprobs, enabled by a new logprob_start_len boundary and max_tokens=0 scoring-only requests.
Changes:
- Introduces
logprob_start_lenacross request/engine plumbing and emitsmeta_info.input_token_logprobsfor scoring-only/generaterequests. - Adds a scoring-only prefill path in the PyTorch engine (compact logit projection, accumulation across chunks, and correct finish/metrics handling).
- Updates prefix-cache matching / ray MP streaming logprob carrier handling and adds extensive unit tests + docs.
Reviewed changes
Copilot reviewed 29 out of 29 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_lmdeploy/test_metrics_loggers.py | Adds metrics coverage for zero-token FINISH/scoring-only request timing/finish_reason behavior. |
| tests/test_lmdeploy/test_messages.py | Verifies GenerationConfig.logprob_start_len contract + SamplingParam conversion and Response extension behavior. |
| tests/test_lmdeploy/serve/test_session_cleanup.py | Ensures scoring-only terminal handling does not cancel/end early and sessions are cleaned up. |
| tests/test_lmdeploy/serve/test_generate_logprobs.py | Adds /generate request validation + formatting tests and stream/non-stream behavior for input logprobs. |
| tests/test_lmdeploy/pytorch/engine/test_engine_instance_cleanup.py | Confirms FINISH logprob carriers are forwarded intact. |
| tests/pytorch/test_model_inputs.py | Validates ModelInputs new logprob metadata clone/to_device/step lifecycle. |
| tests/pytorch/paging/test_block_trie.py | Adds prefix-cache tests ensuring matching is capped appropriately for input scoring boundaries. |
| tests/pytorch/engine/test_ray_mp_engine.py | Tests ray MP stream completion behavior and logprob carrier preservation semantics. |
| tests/pytorch/engine/test_prefill_logprobs.py | Adds unit tests for compact input-logprob row accumulation and scoring-only prefill completion. |
| tests/pytorch/engine/test_model_agent.py | Tests scoring-only lm-head projection, compaction logic, and chunk boundary carry behavior. |
| tests/pytorch/engine/test_inputs_maker.py | Tests logits row index construction and ensures scoring-only requests don’t enter decode running sets. |
| lmdeploy/serve/openai/serving_generate.py | Adds /generate validation logic for input-logprob scoring-only requests (backend/mode/spec decode/max_tokens constraints). |
| lmdeploy/serve/openai/protocol.py | Adds logprob_start_len to request schema and input_token_logprobs to meta output. |
| lmdeploy/serve/openai/api_server.py | Emits input_token_logprobs for /generate and refactors token logprob formatting helper. |
| lmdeploy/serve/core/async_engine.py | Allows max_new_tokens=0 only for input scoring requests and preserves terminal logprob carrier. |
| lmdeploy/pytorch/strategies/ar/sequence.py | Finishes scoring-only prefill without committing an inert generated token. |
| lmdeploy/pytorch/paging/block_trie.py | Adjusts SSM/prefix-cache matching to respect input scoring caps and preserve privacy recompute ranges. |
| lmdeploy/pytorch/model_inputs.py | Adds logits_indices + seq_logit_length metadata for compact scoring projections. |
| lmdeploy/pytorch/messages.py | Adds SamplingParam.logprob_start_len and prefix-cache candidate/match step logic for scoring caps. |
| lmdeploy/pytorch/engine/request.py | Adds request-owned _input_logprobs accumulator for scoring-only chunked prefill. |
| lmdeploy/pytorch/engine/mp_engine/ray_engine.py | Prevents stream completion wake-up from replaying last output; separates stopped flag. |
| lmdeploy/pytorch/engine/mp_engine/base_worker.py | Preserves logprobs=None vs [] states and avoids forcing empty carriers. |
| lmdeploy/pytorch/engine/model_agent/agent.py | Implements scoring-only forward path with compact rows, chunk boundary carry, and bypasses sampling/updates. |
| lmdeploy/pytorch/engine/inputs_maker.py | Builds compact scoring projection indices/counts and filters scoring-only sequences from decode running sets. |
| lmdeploy/pytorch/engine/engine_loop.py | Appends compact input-logprob rows onto the response and routes terminal output to the input-logprob carrier. |
| lmdeploy/metrics/stats.py | Sets finish_reason/finish_time for non-success outputs and handles zero-token scoring-only timing boundaries. |
| lmdeploy/messages.py | Adds GenerationConfig.logprob_start_len validation and clarifies logprobs carrier semantics; preserves empty carriers in extend(). |
| docs/zh_cn/llm/api_server.md | Documents /generate input-token logprob scoring usage and constraints. |
| docs/en/llm/api_server.md | Documents /generate input-token logprob scoring usage and constraints. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if token_ids is None or logprobs is None: | ||
| return None | ||
|
|
||
| output_token_logprobs = [] | ||
| for tok, tok_logprobs in zip(token_ids, logprobs): | ||
| output_token_logprobs.append((tok_logprobs[tok], tok)) | ||
| return output_token_logprobs or None | ||
| token_logprobs = [(tok_logprobs[tok], tok) | ||
| for tok, tok_logprobs in zip(token_ids, logprobs, strict=True)] | ||
| return token_logprobs or None |
| This direct route requires `input_ids` (not prompt or image input), | ||
| `max_tokens=0`, and speculative decoding disabled. Input scoring is not | ||
| available for TurboMind, DistServe, proxy `/generate`, or the | ||
| OpenAI/Anthropic/Responses schemas. | ||
|
|
| 此直接接口要求使用 `input_ids`(不接受 prompt 或图像输入)、设置 | ||
| `max_tokens=0` 并禁用推测解码。TurboMind、DistServe、代理 `/generate` 以及 | ||
| OpenAI、Anthropic 和 Responses 协议暂不支持输入 logprob。 |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 29 out of 29 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (3)
tests/test_lmdeploy/serve/test_generate_logprobs.py:257
- This test asserts that input-logprob scoring works when
image_datais provided, but the documented contract for input scoring isinput_ids-only (no image input). After tightening request validation, this should instead assert that the request is rejected (or, if image+input_ids is intended to be supported, the docs/PR description should be updated accordingly).
def test_generate_input_logprobs_allow_image_with_input_ids(monkeypatch):
result = asyncio.run(
_call_nonempty_route(monkeypatch,
stream=False,
image_data='https://example.com/image.png'))
assert result.meta_info.input_token_logprobs == [(-0.25, 2), (-0.5, 3)]
assert result.meta_info.output_token_logprobs is None
lmdeploy/serve/openai/serving_generate.py:47
logprob_start_lenrequests are documented as requiringinput_idsand not supporting prompt/image input, butcheck_requestcurrently allowslogprob_start_len>=0withprompt(i.e.,input_ids=None) and withimage_data. For prompt-based requests this silently sets up input-logprob mode without anyinput_idsto align rows against, and it contradicts the/generateinput-logprob contract in the docs/PR description.
# Check for input-logprob specific requirements
if request.logprob_start_len >= 0:
if not isinstance(engine_config, PytorchEngineConfig) or engine_config.role != EngineRole.Hybrid:
return 'input logprobs are supported only by the PyTorch hybrid engine.'
if getattr(getattr(server_context, 'async_engine', None), 'speculative_config', None) is not None:
return 'logprob_start_len is not supported with speculative decoding.'
if return_logprob is not True:
return 'logprob_start_len requires return_logprob=True.'
if logprobs_mode not in ('raw_logits', 'raw_logprobs'):
return 'logprob_start_len requires raw_logits or raw_logprobs mode.'
if request.max_tokens != 0:
return 'logprob_start_len requires max_tokens=0.'
if request.input_ids is not None and request.logprob_start_len + 1 >= len(request.input_ids):
return (f'logprob_start_len({request.logprob_start_len}) exceeds '
f'the last scorable source position for input_ids length({len(request.input_ids)}).')
tests/test_lmdeploy/serve/test_generate_logprobs.py:87
- These assertions treat input-logprob scoring as valid for
prompt=...and for requests includingimage_data, but the/generateinput-logprob contract (docs + check logic) isinput_ids-only and should reject prompt/image scoring requests. Onceserving_generate.check_requestenforces this, these tests should assert the appropriate validation errors instead of expecting success.
This issue also appears on line 251 of the same file.
assert check_request(
GenerateReqInput(prompt='hi',
return_logprob=True,
logprob_start_len=0,
max_tokens=0), context) == ''
assert check_request(
GenerateReqInput(input_ids=[1, 2],
image_data='https://example.com/image.png',
return_logprob=True,
logprob_start_len=0,
max_tokens=0), context) == ''
Motivation
add PyTorch raw
/generateinput-token logprob scoring withlogprob_start_lenand scoring-onlymax_tokens=0Usage
Start the server with
--backend pytorch --logprobs-mode raw_logprobs, thenset
return_logprobandlogprob_start_len:curl http://{server_ip}:{server_port}/generate \ -H 'Content-Type: application/json' \ -d '{ "input_ids": [1234, 5678, 9012], "return_logprob": true, "logprob_start_len": 0, "max_tokens": 0 }'meta_info.input_token_logprobsis emitted once. ForN=logprob_start_len, rows align withinput_ids[N + 1:], and every entry is[value, token_id]; the boundary token itself is not emitted. At least onetoken must follow the boundary; otherwise the endpoint rejects the request. Both
streaming and non-streaming forms return one terminal payload; streaming then
emits
[DONE].This direct route requires
input_ids(not prompt or image input),max_tokens=0, and speculative decoding disabled. Input scoring is notavailable for TurboMind, DistServe, proxy
/generate, or theOpenAI/Anthropic/Responses schemas.
BC-breaking (Optional)
Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.
Use cases (Optional)
If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.
Checklist