[Feat]: Support output input logprobs

docs/en/llm/api_server.md

@@ -203,6 +203,34 @@ curl http://{server_ip}:{server_port}/v1/completions \
 }'
 ```
 
+- score input tokens with the direct PyTorch `/generate` endpoint
+
+Start the server with `--backend pytorch --logprobs-mode raw_logprobs`, then
+set `return_logprob` and `logprob_start_len`:
+
+```bash
+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.
+
 ## Launch multiple api servers
 
 Following are two steps to launch multiple api servers through torchrun. Just create a python script with the following codes.

docs/zh_cn/llm/api_server.md

@@ -169,6 +169,31 @@ for item in api_client.completions_v1(model=model_name, prompt='hi'):
 
 参考 [api_server_responses](./api_server_responses.md)。
 
+### 使用 `/generate` 计算输入 token 的 logprob
+
+使用 `--backend pytorch --logprobs-mode raw_logprobs` 启动服务，然后在直接
+`/generate` 请求中设置 `return_logprob` 和 `logprob_start_len`：
+
+```bash
+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` 只发送一次。对于
+`N=logprob_start_len`，返回行对应 `input_ids[N + 1:]`，每项为
+`[value, token_id]`。边界后必须至少有一个 token，否则端点会拒绝该请求。
+流式和非流式都只返回一个终止 payload，流式随后发送 `[DONE]`。
+
+此直接接口要求使用 `input_ids`（不接受 prompt 或图像输入）、设置
+`max_tokens=0` 并禁用推测解码。TurboMind、DistServe、代理 `/generate` 以及
+OpenAI、Anthropic 和 Responses 协议暂不支持输入 logprob。
+
 ### 使用 Java/Golang/Rust
 
 可以使用代码生成工具 [openapi-generator-cli](https://github.com/OpenAPITools/openapi-generator-cli) 将 `http://{server_ip}:{server_port}/openapi.json` 转成 java/rust/golang 客户端。

lmdeploy/messages.py

@@ -74,6 +74,10 @@ class GenerationConfig:
             around special tokens. The behavior of Fast tokenizers is to have
             this to False. This is setup to True in slow tokenizers.
         logprobs: Number of log probabilities to return per output token.
+        logprob_start_len: Source-token boundary in the current
+            model-processed input. Rows are returned for tokens after this
+            boundary. ``-1`` disables input logprobs while preserving
+            generated-token logprobs.
         response_format: Generate responses according to given formatting.
             Examples:
 
@@ -130,6 +134,7 @@ class GenerationConfig:
     skip_special_tokens: bool = True
     spaces_between_special_tokens: bool = True
     logprobs: int = None
+    logprob_start_len: int = -1
     response_format: dict | None = None
     logits_processors: list[LogitsProcessor] | None = None
     output_logits: Literal['all', 'generation'] = None
@@ -200,6 +205,10 @@ def __post_init__(self):
         assert self.temperature >= 0 and self.temperature <= 2  # [0,2]
         assert 0 <= self.min_p <= 1, \
             f'min_p should be in range [0, 1], but found {self.min_p}'
+        if self.logprob_start_len < -1:
+            raise ValueError('logprob_start_len must be greater than or equal to -1')
+        if self.logprob_start_len >= 0 and (self.logprobs is None or self.logprobs < 0):
+            raise ValueError('logprobs must be non-negative when logprob_start_len is non-negative')
         if self.repetition_ngram_size <= 0 or self.repetition_ngram_threshold <= 0:
             self.repetition_ngram_size = 0
             self.repetition_ngram_threshold = 0
@@ -584,7 +593,9 @@ class Response:
             stop point or a provided stop sequence, 'length' if the maximum
             number of tokens specified in the request was reached.
         token_ids: the output token ids.
-        logprobs: the top logprobs for each output position.
+        logprobs: the top logprobs for each output position. For a scoring-only
+            input-logprob request, this field carries the complete ordered
+            input-token rows on the terminal response.
         index: it refers to the position index of the input request batch.
     """
     text: str
@@ -649,7 +660,7 @@ def extend(self, other: 'Response') -> 'Response':
         self.index = other.index
         if other.token_ids:
             self.token_ids += other.token_ids
-        if other.logprobs:
+        if other.logprobs is not None:
             self.logprobs = self.logprobs or []
             self.logprobs += other.logprobs
         self.routed_experts = other.routed_experts
@@ -723,8 +734,9 @@ class EngineOutput:
     Args:
         status: the response type.
         token_ids: the newly generated token ids in each iteration.
-        logprobs: the top logprobs for each output
-            position.
+        logprobs: the top logprobs for each output position. For a scoring-only
+            input-logprob request, this internal field carries the complete
+            ordered input-token rows on its terminal output.
         cache_block_ids: send cache blocks back for migration in
             Disaggregated LLM Serving when Prefill Engine is Done.
         req_metrics: request metrics information

lmdeploy/metrics/stats.py

@@ -239,7 +239,17 @@ def update_from_output(self, outputs: EngineOutput, req_stats: RequestStats):
             return
 
         new_generation_tokens = len(outputs.token_ids)
+        if outputs.status != ResponseType.SUCCESS:
+            req_stats.finish_reason = outputs.status
+            req_stats.finish_time = self.iteration_timestamp
         if new_generation_tokens == 0:
+            if outputs.status != ResponseType.SUCCESS and req_stats.first_token_time == 0:
+                # A scoring-only request has no first generated token. Use the
+                # terminal timestamp as the internal prefill/decode boundary:
+                # full inference time is prefill and decode time is zero.
+                req_stats.first_token_time = req_stats.finish_time
+                req_stats.lastest_token_time = req_stats.finish_time
+                self.prompt_tokens = req_stats.prompt_tokens
             return
 
         self.new_generation_tokens = new_generation_tokens
@@ -256,11 +266,6 @@ def update_from_output(self, outputs: EngineOutput, req_stats: RequestStats):
         req_stats.lastest_token_time = outputs.req_metrics.token_timestamp
         req_stats.generation_tokens += new_generation_tokens
 
-        if outputs.status != ResponseType.SUCCESS:
-            req_stats.finish_reason = outputs.status
-            req_stats.finish_time = self.iteration_timestamp
-
-
 # modify from vllm
 @dataclass
 class SpeculativeDecodingStats:

lmdeploy/pytorch/engine/engine_loop.py

@@ -202,6 +202,31 @@ def reset_runtime_state(self):
         sessions."""
         self.inputs_maker.reset_runtime_state()
 
+    @staticmethod
+    def _append_input_logprobs(batched_outputs: 'BatchedOutputs', running: 'SeqList',
+                               model_inputs: 'ModelInputs | None'):
+        """Append compact prefill rows to their request-owned response."""
+        if model_inputs is None or model_inputs.is_decoding or model_inputs.logits_indices is None:
+            return
+        compact = batched_outputs.logprobs
+        row_counts = model_inputs.seq_logit_length.tolist()
+        num_rows = sum(row_counts)
+        if (compact is None or compact.vals.shape != compact.indices.shape
+                or num_rows != compact.vals.shape[0]):
+            raise RuntimeError('input-logprob compact output mismatch')
+
+        offset = 0
+        for msg, row_count in zip(running, row_counts, strict=True):
+            end = offset + row_count
+            if row_count > 0:
+                if msg.resp._input_logprobs is None:
+                    msg.resp._input_logprobs = []
+                msg.resp._input_logprobs.extend(
+                    zip(compact.vals[offset:end].tolist(),
+                        compact.indices[offset:end].tolist(),
+                        strict=True))
+            offset = end
+
     @staticmethod
     def _log_resps(outputs: list[InferOutput]):
         """Log resps."""
@@ -290,10 +315,8 @@ def __get_logit(msg, logits: torch.Tensor, seq_length: list[int], idx: int):
 
             return logit
 
-        def __get_logprobs(batched_outputs: 'BatchedOutputs'):
+        def __get_logprobs(logprobs, stop_positions, batch_size: int):
             """Get valid logprobs."""
-            batch_size = batched_outputs.stop_pos.size(0)
-            logprobs = batched_outputs.logprobs
             if logprobs is None:
                 return [None for _ in range(batch_size)]
             num_decode_tokens = logprobs.indices.shape[0] // batch_size
@@ -304,7 +327,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
                 start = idx * num_decode_tokens
                 end = (idx + 1) * num_decode_tokens
                 mask = logprobs.indices[start:end][:, 0] >= 0
-                stop_pos = batched_outputs.stop_pos[idx]
+                stop_pos = stop_positions[idx] if stop_positions is not None else -1
                 # only apply when stopped
                 if stop_pos > -1:
                     mask = torch.logical_and(mask, stop_pos >= range_tensor)
@@ -316,6 +339,13 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
         logits = batched_outputs.logits
         all_routed_experts = batched_outputs.all_routed_experts
         ce_loss = batched_outputs.ce_loss
+        logprobs = batched_outputs.logprobs
+        is_prefill_logprobs_mode = (model_inputs is not None and not model_inputs.is_decoding
+                                    and model_inputs.logits_indices is not None)
+        if is_prefill_logprobs_mode:
+            self._append_input_logprobs(batched_outputs, running, model_inputs)
+            # Packed scoring rows do not enter generated-row splitting below.
+            logprobs = None
 
         if model_inputs is not None and (model_inputs.is_chunk and not model_inputs.is_last_chunk):
             # chunk long context does not need to update seqs and outputs
@@ -327,9 +357,8 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
             return dict()
 
         new_token_timestamp = batched_outputs.new_token_timestamp
-        logprobs = batched_outputs.logprobs
 
-        all_logprobs = __get_logprobs(batched_outputs)
+        all_logprobs = __get_logprobs(logprobs, batched_outputs.stop_pos, len(running))
 
         seq_length = [seq.num_token_ids for seq in running]
         is_run = [seq.status == MessageStatus.RUNNING for seq in running]
@@ -346,12 +375,15 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
                 continue
             token_ids = msg.generated_ids
             finish = msg.status == MessageStatus.STOPPED or msg.status == MessageStatus.TO_BE_MIGRATED
-            if not finish and len(token_ids) == 0:
-                continue
-            resp_data = msg.resp.data
-            if resp_data is not None and len(resp_data.get('token_ids', [])) == len(token_ids):
-                # no new tokens
-                continue
+            input_logprobs = msg.resp._input_logprobs
+            if input_logprobs is None:
+                if not finish and len(token_ids) == 0:
+                    continue
+                resp_data = msg.resp.data
+                if (resp_data is not None and 'token_ids' in resp_data
+                        and len(resp_data['token_ids']) == len(token_ids)):
+                    # no new tokens
+                    continue
             session_id = msg.session_id
             if msg.resp_cache:
                 cache_block_ids = self.scheduler.block_manager.get_block_table(msg).tolist()
@@ -364,6 +396,8 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
             if logprobs is not None and num_logprobs > 0:
                 cur_logprobs = list([_dat[:num_logprobs + 1] for _dat in vals_indices]
                                     for vals_indices in all_logprobs[idx])
+            if input_logprobs is not None:
+                cur_logprobs = input_logprobs
 
             # get spec stats info
             spec_info = None

lmdeploy/pytorch/engine/inputs_maker.py

@@ -82,6 +82,43 @@ def _compact_state_prefix_cache_save_offsets(messages: list['SchedulerSequence']
     return tuple(src_offsets), tuple(dst_offsets)
 
 
+def _build_logits_indices(messages: list['SchedulerSequence'], query_lengths: list[int]):
+    """Build ordered source-hidden indices and per-sequence output rows.
+
+    Non-final long-prefill chunks still project their final source row, but
+    that row is emitted by the next chunk through ModelAgent's
+    ``_prev_chunk_last_logit`` carry.
+    """
+    if not any(msg.logprob_start_pos >= 0 for msg in messages):
+        return None, None
+
+    indices = []
+    row_counts = []
+    base = 0
+    for msg, query_len in zip(messages, query_lengths, strict=True):
+        chunk_start = int(msg.num_history_ids)
+        chunk_end = chunk_start + int(query_len)
+        logprob_start = int(msg.logprob_start_pos)
+        input_end = int(msg.input_end_pos)
+
+        range_start = max(chunk_start, logprob_start) if logprob_start >= 0 else chunk_end
+        range_end = min(chunk_end, input_end - 1)
+        num_projected_rows = max(0, range_end - range_start)
+        indices.extend(range(base + range_start - chunk_start, base + range_end - chunk_start))
+
+        has_prev_chunk_logit = (logprob_start >= 0 and int(msg.input_start_pos) < chunk_start
+                                and logprob_start < chunk_start and chunk_start < input_end)
+        stash_current_last_logit = (range_end == chunk_end and num_projected_rows > 0
+                                    and chunk_end < input_end)
+        if stash_current_last_logit:
+            assert len(messages) == 1, 'long-prefill cross-chunk logit carry requires a single request'
+        row_counts.append(num_projected_rows + int(has_prev_chunk_logit) - int(stash_current_last_logit))
+        base += int(query_len)
+
+    return (torch.tensor(indices, dtype=torch.long),
+            torch.tensor(row_counts, dtype=torch.long))
+
+
 @dataclass
 class InputsMakerConfig:
     """Input maker config.
@@ -882,6 +919,10 @@ def create_model_inputs(self, messages: 'SeqList', is_prefill: bool):
             sum_kv_seqlen=sum_kv_seqlen,
             model_metas=model_metas,
         )
+        if is_prefill:
+            logits_indices, seq_logit_length = _build_logits_indices(messages, [len(ids) for ids in token_ids])
+            model_inputs.logits_indices = logits_indices
+            model_inputs.seq_logit_length = seq_logit_length
 
         # adapters
         self._set_adapter_ids(model_inputs, messages)
@@ -964,6 +1005,9 @@ def create_model_inputs_long_context(self,
             model_metas=model_metas,
             is_chunk=True,
         )
+        logits_indices, seq_logit_length = _build_logits_indices([seq], [chunk_size])
+        model_inputs.logits_indices = logits_indices
+        model_inputs.seq_logit_length = seq_logit_length
 
         # adapters
         self._set_adapter_ids(model_inputs, [seq])
@@ -1125,7 +1169,9 @@ def update_running_seqs(self, running: 'SeqList', inputs: 'ModelInputs|None'):
         if is_decoding:
             self.running_seqs = running
         else:
-            self.running_seqs += running
+            for seq in running:
+                if seq.sampling_param.max_new_tokens > 0:
+                    self.running_seqs.append(seq)
 
     def deactivate_evict_seqs(self):
         """Deactivate and evict seqs."""