Skip to content

[Feat]: Support output input logprobs - #4793

Open
RunningLeon wants to merge 3 commits into
InternLM:mainfrom
RunningLeon:lmdeploy-logprobs-v2
Open

[Feat]: Support output input logprobs #4793
RunningLeon wants to merge 3 commits into
InternLM:mainfrom
RunningLeon:lmdeploy-logprobs-v2

Conversation

@RunningLeon

Copy link
Copy Markdown
Collaborator

Motivation

add PyTorch raw /generate input-token logprob scoring with logprob_start_len and scoring-only max_tokens=0

Usage

Start the server with --backend pytorch --logprobs-mode raw_logprobs, then
set return_logprob and logprob_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_logprobs is emitted once. For
N=logprob_start_len, rows align with input_ids[N + 1:], and every entry is
[value, token_id]; the boundary token itself is not emitted. At least one
token 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 not
available for TurboMind, DistServe, proxy /generate, or the
OpenAI/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

  1. Pre-commit or other linting tools are used to fix the potential lint issues.
  2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
  3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
  4. The documentation has been modified accordingly, like docstring or example tutorials.

Copilot AI review requested due to automatic review settings July 28, 2026 05:06

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_len across request/engine plumbing and emits meta_info.input_token_logprobs for scoring-only /generate requests.
  • 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.

Comment on lines 283 to +288
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
Comment thread docs/en/llm/api_server.md
Comment on lines +229 to +233
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.

Comment on lines +193 to +195
此直接接口要求使用 `input_ids`(不接受 prompt 或图像输入)、设置
`max_tokens=0` 并禁用推测解码。TurboMind、DistServe、代理 `/generate` 以及
OpenAI、Anthropic 和 Responses 协议暂不支持输入 logprob。

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_data is provided, but the documented contract for input scoring is input_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_len requests are documented as requiring input_ids and not supporting prompt/image input, but check_request currently allows logprob_start_len>=0 with prompt (i.e., input_ids=None) and with image_data. For prompt-based requests this silently sets up input-logprob mode without any input_ids to align rows against, and it contradicts the /generate input-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 including image_data, but the /generate input-logprob contract (docs + check logic) is input_ids-only and should reject prompt/image scoring requests. Once serving_generate.check_request enforces 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) == ''

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants