From 485d5b4d1754120d4588c5bbd027d635a07d8dc0 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Tue, 28 Jul 2026 03:45:37 +0000 Subject: [PATCH 1/4] Add PyTorch input logprobs scoring --- docs/en/llm/api_server.md | 28 ++ docs/zh_cn/llm/api_server.md | 25 ++ lmdeploy/messages.py | 20 +- lmdeploy/metrics/stats.py | 15 +- lmdeploy/pytorch/engine/engine_loop.py | 58 +++- lmdeploy/pytorch/engine/inputs_maker.py | 48 ++- lmdeploy/pytorch/engine/model_agent/agent.py | 104 +++++- .../pytorch/engine/mp_engine/base_worker.py | 9 +- .../pytorch/engine/mp_engine/ray_engine.py | 34 +- lmdeploy/pytorch/engine/request.py | 3 + lmdeploy/pytorch/messages.py | 22 +- lmdeploy/pytorch/model_inputs.py | 6 + lmdeploy/pytorch/paging/block_trie.py | 33 +- lmdeploy/pytorch/strategies/ar/sequence.py | 6 + lmdeploy/serve/core/async_engine.py | 12 +- lmdeploy/serve/openai/api_server.py | 64 ++-- lmdeploy/serve/openai/protocol.py | 2 + lmdeploy/serve/openai/serving_generate.py | 36 ++- tests/pytorch/engine/test_inputs_maker.py | 71 ++++ tests/pytorch/engine/test_model_agent.py | 306 ++++++++++++++++++ tests/pytorch/engine/test_prefill_logprobs.py | 258 +++++++++++++++ tests/pytorch/engine/test_ray_mp_engine.py | 73 ++++- tests/pytorch/paging/test_block_trie.py | 153 +++++++++ tests/pytorch/test_model_inputs.py | 17 + .../engine/test_engine_instance_cleanup.py | 22 ++ .../serve/test_generate_logprobs.py | 254 +++++++++++++++ .../serve/test_session_cleanup.py | 97 +++++- tests/test_lmdeploy/test_messages.py | 29 ++ tests/test_lmdeploy/test_metrics_loggers.py | 22 +- 29 files changed, 1729 insertions(+), 98 deletions(-) create mode 100644 tests/pytorch/engine/test_prefill_logprobs.py create mode 100644 tests/test_lmdeploy/serve/test_generate_logprobs.py diff --git a/docs/en/llm/api_server.md b/docs/en/llm/api_server.md index 4f70d4d6fe..3806dc365e 100644 --- a/docs/en/llm/api_server.md +++ b/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. diff --git a/docs/zh_cn/llm/api_server.md b/docs/zh_cn/llm/api_server.md index 131e845a02..9f61bb9223 100644 --- a/docs/zh_cn/llm/api_server.md +++ b/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 客户端。 diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index e07a46ae5e..4ec070afae 100644 --- a/lmdeploy/messages.py +++ b/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 diff --git a/lmdeploy/metrics/stats.py b/lmdeploy/metrics/stats.py index 90b19d924f..1c01442014 100644 --- a/lmdeploy/metrics/stats.py +++ b/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: diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py index 8b39cab777..ade6b01218 100644 --- a/lmdeploy/pytorch/engine/engine_loop.py +++ b/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 diff --git a/lmdeploy/pytorch/engine/inputs_maker.py b/lmdeploy/pytorch/engine/inputs_maker.py index eac5506f2f..4378c25ba1 100644 --- a/lmdeploy/pytorch/engine/inputs_maker.py +++ b/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.""" diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py index 1a687e4347..3361c3077c 100644 --- a/lmdeploy/pytorch/engine/model_agent/agent.py +++ b/lmdeploy/pytorch/engine/model_agent/agent.py @@ -21,7 +21,7 @@ from lmdeploy.pytorch.distributed import DistContext, get_dist_manager from lmdeploy.pytorch.engine.cache_engine import CacheEngine, StateCacheEngine from lmdeploy.pytorch.engine.guided_process import GuidedDecodingManager -from lmdeploy.pytorch.engine.logits_process import FusedLogitsProcessor, SamplingInputs +from lmdeploy.pytorch.engine.logits_process import FusedLogitsProcessor, SamplingInputs, _torch_topk from lmdeploy.pytorch.memdecode import build_memdecode_agent from lmdeploy.pytorch.model_inputs import ModelInputs, ModelInputsDelta, step_ctx_manager from lmdeploy.pytorch.models.patch import BuildModelContext, add_adapters, build_patched_model, update_custom_module_map @@ -370,7 +370,7 @@ def _init_runtime_state(self): """Initialize request-local decode and chunk state.""" self.step_inputs = self.strategy_factory.build_step_inputs() self._prev_chunk_output: dict = None - # chunked-prefill ppl: last logit row of the previous chunk, used to score the cross-chunk boundary token + # Last logit row of the previous chunk, used to score cross-chunk prompt tokens. self._prev_chunk_last_logit: torch.Tensor | None = None def reset_runtime_state(self): @@ -485,25 +485,49 @@ def _postprocess_forward_output(self, output: dict, inputs: ModelInputs): output['hidden_states'] = hidden_states return output + def _get_input_logits(self, hidden_states: torch.Tensor, inputs: ModelInputs): + """Project the selected scoring rows with at most one lm-head call.""" + if not self._is_prefill_input_logprobs(inputs): + return None + logits_indices = inputs.logits_indices + flat_hidden = hidden_states.reshape(-1, hidden_states.shape[-1]) + selected_hidden = flat_hidden.index_select(0, logits_indices) + return self.get_logits(selected_hidden[None])[0] + + @staticmethod + def _is_prefill_input_logprobs(inputs: ModelInputs | None) -> bool: + """Whether this prefill step uses the V2 scoring-only input-logprob + path.""" + return (inputs is not None and not getattr(inputs, 'is_dummy', False) + and not getattr(inputs, 'is_decoding', False) + and getattr(inputs, 'logits_indices', None) is not None + and getattr(inputs, 'seq_logit_length', None) is not None) + async def _async_model_forward( self, inputs: ModelInputs, return_logits: bool, ): """Model forward.""" - if self.memdecode_agent is not None and return_logits: + memdecode_agent = getattr(self, 'memdecode_agent', None) + if memdecode_agent is not None and return_logits: raise RuntimeError('MemDecode does not support returned prompt logits yet.') ret = await self.async_forward(inputs) + if self._is_prefill_input_logprobs(inputs): + # This is the only lm-head projection for scoring: return before + # the ordinary all/sampling-logits projection below. + ret['logits'] = self._get_input_logits(ret['hidden_states'][0], inputs) + return ret if not return_logits: ret = self._postprocess_forward_output(ret, inputs) - if self.memdecode_agent is not None: + if memdecode_agent is not None: base_hidden_states = ret['hidden_states'] base_logits = self.get_logits(base_hidden_states) - return await self.memdecode_agent.fuse_with_base( + return await memdecode_agent.fuse_with_base( inputs=inputs, base_output=ret, base_logits=base_logits, @@ -516,6 +540,59 @@ async def _async_model_forward( ret['logits'] = logits return ret + def _get_outputs_with_logprobs( + self, + input_logits: torch.Tensor | None, + inputs: ModelInputs, + num_logprobs: int, + model_metas: list[dict[str, Any]] | None, + ): + """Build the scoring-only output without generation modifiers.""" + input_logprobs = None + if input_logits is not None: + source_indices = inputs.logits_indices + input_ids = inputs.input_ids.flatten() + if inputs.is_chunk: + prev_chunk_last_logit = None if inputs.is_first_chunk else self._prev_chunk_last_logit + if (source_indices.numel() > 0 and not inputs.is_last_chunk + and source_indices[-1].item() == input_ids.numel() - 1): + self._prev_chunk_last_logit = input_logits[-1:].clone() + input_logits = input_logits[:-1] + source_indices = source_indices[:-1] + else: + self._prev_chunk_last_logit = None + else: + prev_chunk_last_logit = None + + targets = input_ids.index_select(0, source_indices + 1) + if prev_chunk_last_logit is not None: + input_logits = torch.cat([prev_chunk_last_logit, input_logits], dim=0) + targets = torch.cat([input_ids[:1], targets]) + if self.misc_config.logprobs_mode == 'raw_logprobs': + input_logits = input_logits.log_softmax(dim=-1) + target_indices = targets[:, None] + vals = input_logits.gather(-1, target_indices) + if num_logprobs > 0: + topk_vals, topk_indices = _torch_topk(input_logits, num_logprobs, dim=-1) + vals = torch.cat([vals, topk_vals], dim=-1) + target_indices = torch.cat([target_indices, topk_indices], dim=-1) + input_logprobs = BatchedLogProbs(vals=vals, indices=target_indices.to(torch.int32)) + + batch_size = inputs.seq_length.numel() + device = inputs.input_ids.device + return BatchedOutputs( + next_token_ids=torch.zeros(batch_size, + dtype=inputs.input_ids.dtype, + device=device), + stopped=torch.ones(batch_size, dtype=torch.bool, device=device), + stop_pos=torch.full((batch_size,), + -1, + dtype=torch.long, + device=device), + model_metas=model_metas, + logprobs=input_logprobs, + ) + async def async_sampling_logits(self, logits: torch.Tensor, inputs: ModelInputs, extra_inputs: ExtraInputs, sampling_inputs: SamplingInputs): """Sampling logits.""" @@ -866,15 +943,30 @@ async def _async_step( f'is_last_chunk={inputs.is_last_chunk} ' f'dp_meta={inputs.dp_meta} ' f'is_decoding={inputs.is_decoding}') + prefill_input_logprobs = self._is_prefill_input_logprobs(inputs) output = await self._async_model_forward( inputs, return_logits=return_logits or return_ce_loss, - ) + ) if inputs.is_dummy and not self.spec_agent.is_enabled(): # skip dummy forward output return + # input-logprob mode is scoring-only: this branch emits compact + # logprob rows, and falling through would wrongly run sampling / + # sequence-update on a max_tokens=0 request. + if prefill_input_logprobs: + model_metas = output.get('model_metas') + if self.need_output: + output = self._get_outputs_with_logprobs(output.get('logits'), inputs, + sampling_inputs.max_num_logprobs, model_metas) + self._push_output(output) + if inputs.is_chunk and not inputs.is_last_chunk: + self._prev_chunk_output = {'model_metas': model_metas} + + return + logits = output['logits'][0] # [bs, seq, prob] -> [seq, prob] seq_length = output.get('seq_length', inputs.seq_length) last_logits = self._slice_outs(logits, seq_length) # [bs, 1, prob] -> [bs, prob] diff --git a/lmdeploy/pytorch/engine/mp_engine/base_worker.py b/lmdeploy/pytorch/engine/mp_engine/base_worker.py index d80700d9a6..10998d229f 100644 --- a/lmdeploy/pytorch/engine/mp_engine/base_worker.py +++ b/lmdeploy/pytorch/engine/mp_engine/base_worker.py @@ -154,7 +154,7 @@ def __init__(self): def get(self, stream_id): if stream_id not in self._output: - self._output[stream_id] = EngineOutput(status=None, token_ids=[], logprobs=[]) + self._output[stream_id] = EngineOutput(status=None, token_ids=[], logprobs=None) return self._output[stream_id] def add(self, stream_id, result): @@ -162,7 +162,10 @@ def add(self, stream_id, result): return output = self.get(stream_id) output.token_ids.extend(result.token_ids or []) - output.logprobs.extend(result.logprobs or []) + if result.logprobs is not None: + if output.logprobs is None: + output.logprobs = [] + output.logprobs.extend(result.logprobs) def pop(self, stream_id, result): if not isinstance(result, EngineOutput): @@ -171,7 +174,7 @@ def pop(self, stream_id, result): if output is None: return result result.token_ids = output.token_ids or [] - result.logprobs = output.logprobs or None + result.logprobs = output.logprobs return result def discard(self, stream_id): diff --git a/lmdeploy/pytorch/engine/mp_engine/ray_engine.py b/lmdeploy/pytorch/engine/mp_engine/ray_engine.py index 5b26aff2c2..087b39a49a 100644 --- a/lmdeploy/pytorch/engine/mp_engine/ray_engine.py +++ b/lmdeploy/pytorch/engine/mp_engine/ray_engine.py @@ -39,7 +39,8 @@ def __init__(self, async def _stream_task_wrapper(self, stream_id: int, init_event: asyncio.Event, func: str, *args, **kwargs): """Create a stream task.""" method = getattr(self, func) - event = self._stream_aiter[stream_id][0] + stream_out = self._stream_aiter[stream_id] + event = stream_out[0] # notify after add msg def _notify_add_msg(): @@ -50,18 +51,28 @@ def _notify_add_msg(): kwargs['notify_add_msg_func'] = _notify_add_msg result = EngineOutput(ResponseType.INTERNAL_ENGINE_ERROR, []) + has_yielded = False + failed = False try: generator = method(*args, **kwargs) async for result in generator: + has_yielded = True self._engine_output_gather.add(stream_id, result) - self._stream_aiter[stream_id][1] = (result, False) + stream_out[1] = result event.set() except Exception: + failed = True + result = EngineOutput(ResponseType.INTERNAL_ENGINE_ERROR, []) logger.exception(f'Stream task {stream_id} failed.') finally: if not init_event.is_set(): init_event.set() - self._stream_aiter[stream_id][1] = (result, True) + # Completion is a wake-up, not another copy of the last output. + # Keep an unconsumed result in place and mark the stream stopped; + # if it was already consumed, wake the poller with result=None. + stream_out[2] = True + if failed or (not has_yielded and stream_out[1] is None): + stream_out[1] = result event.set() async def create_stream_task(self, func, *args, **kwargs): @@ -69,7 +80,9 @@ async def create_stream_task(self, func, *args, **kwargs): stream_id = self._stream_id self._stream_id += 1 event_loop = asyncio.get_event_loop() - self._stream_aiter[stream_id] = [asyncio.Event(), None] + # [wake event, pending result, stopped]. The stopped flag is kept + # separately so the completion wake-up cannot replay the last result. + self._stream_aiter[stream_id] = [asyncio.Event(), None, False] init_event = asyncio.Event() task = event_loop.create_task(self._stream_task_wrapper(stream_id, init_event, func, *args, **kwargs)) self._stream_task[stream_id] = task @@ -84,13 +97,13 @@ async def get_stream_task_result(self, stream_id: int): return None, True event = stream_out[0] await event.wait() - stream_result = stream_out[1] + result = stream_out[1] + stopped = stream_out[2] + stream_out[1] = None event.clear() - if stream_result is None: - return None, True - result, stopped = stream_result - result = self._engine_output_gather.pop(stream_id, result) + if result is not None: + result = self._engine_output_gather.pop(stream_id, result) if stopped: self._stream_aiter.pop(stream_id, None) @@ -123,8 +136,7 @@ async def drop_stream_task(self, stream_id: int): logger.exception(f'Ray MP abandoned stream task failed during drop: stream_id={stream_id}.') if stream_out is not None: event = stream_out[0] - if stream_out[1] is None: - stream_out[1] = (None, True) + stream_out[2] = True event.set() self._stream_aiter.pop(stream_id, None) self._stream_task.pop(stream_id, None) diff --git a/lmdeploy/pytorch/engine/request.py b/lmdeploy/pytorch/engine/request.py index 1079900781..a39eb69fc9 100644 --- a/lmdeploy/pytorch/engine/request.py +++ b/lmdeploy/pytorch/engine/request.py @@ -34,6 +34,9 @@ class Response: err_msg: str = '' is_done: bool = False req_metrics: RequestMetrics = None + # Raw rows in InferOutput.logprobs format, accumulated across scoring + # prefill chunks until the terminal response. + _input_logprobs: list[tuple[list[float], list[int]]] | None = None @dataclass diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py index b1c5a59e68..4d3426a945 100644 --- a/lmdeploy/pytorch/messages.py +++ b/lmdeploy/pytorch/messages.py @@ -141,6 +141,7 @@ class SamplingParam: out_last_hidden_states: bool = False out_ce_loss: bool = False num_logprobs: int = -1 + logprob_start_len: int = -1 return_routed_experts: bool = False # ngram @@ -238,6 +239,7 @@ def from_gen_config(cls, gen_config: GenerationConfig): out_logits=(output_logits is not None), out_ce_loss=gen_config.return_ppl, num_logprobs=logprobs, + logprob_start_len=gen_config.logprob_start_len, return_routed_experts=gen_config.return_routed_experts, repetition_ngram_size=repetition_ngram_size, repetition_ngram_threshold=repetition_ngram_threshold, @@ -854,6 +856,14 @@ def num_token_ids(self): def num_valid_ids(self): return self._num_history_ids + self._num_token_ids + @property + def logprob_start_pos(self) -> int: + """Absolute source-token boundary for current-input scoring.""" + start = self.sampling_param.logprob_start_len + if self.sampling_param.num_logprobs < 0 or start < 0: + return -1 + return self.input_start_pos + start + @property def num_images(self): return self._num_images @@ -991,10 +1001,18 @@ def clamp_prefix_cache_match_step(self, step: int): clamped = next_step return clamped + def get_prefix_cache_max_candidate_step(self): + """Get the exclusive raw prefix-reuse limit before safety rewinds.""" + max_step = self.num_valid_ids - 1 + logprob_start = self.logprob_start_pos + if logprob_start >= 0: + max_step = min(max_step, logprob_start) + return max(0, max_step) + def get_prefix_cache_max_match_step(self): - """Get the deepest prefix step allowed for a cache hit.""" + """Get the deepest effective prefix step allowed for a cache hit.""" block_size = self.block_size - max_step = ((self.num_valid_ids - 1) // block_size) * block_size + max_step = self.clamp_prefix_cache_match_step(self.get_prefix_cache_max_candidate_step()) recompute_blocks = max(0, self.prefix_cache.match_recompute_blocks) if recompute_blocks > 0: max_step = max(0, max_step - recompute_blocks * block_size) diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index 39d68dd854..a4e5aa1d9b 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/lmdeploy/pytorch/model_inputs.py @@ -210,6 +210,10 @@ class ModelInputs: target_hidden_states: torch.Tensor | None = None target_position_ids: torch.Tensor | None = None target_inputs_embeds: torch.Tensor | None = None + # Hidden-state positions requiring lm-head projection. + logits_indices: torch.LongTensor | None = None + # Number of compact logprob rows emitted for each sequence. + seq_logit_length: torch.LongTensor | None = None is_chunk: bool = False is_first_chunk: bool = False is_last_chunk: bool = False @@ -240,6 +244,8 @@ def step(self, input_ids: torch.Tensor, step_seqlens: torch.Tensor = None): state_prefix_cache_dst_offsets=None, state_prefix_cache_save_src_offsets=None, state_prefix_cache_save_offsets=None, + logits_indices=None, + seq_logit_length=None, mrope_pos_ids=mrope_pos_ids, ) diff --git a/lmdeploy/pytorch/paging/block_trie.py b/lmdeploy/pytorch/paging/block_trie.py index 97f246b711..ee90a31236 100644 --- a/lmdeploy/pytorch/paging/block_trie.py +++ b/lmdeploy/pytorch/paging/block_trie.py @@ -1007,7 +1007,7 @@ def _verify_state_checkpoint_node(self, seq: SchedulerSequence, node: Node, inde return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.STALE_INDEX_ENTRY, reason='checkpoint adapter differs from lookup adapter') - max_step = ((seq.num_valid_ids - 1) // self.block_size) * self.block_size + max_step = seq.get_prefix_cache_max_candidate_step() if step > max_step: return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.REQUEST_MISMATCH, reason='checkpoint is longer than this request') @@ -1046,15 +1046,20 @@ def _match_state_checkpoint(self, seq: SchedulerSequence): init_num_matched = init_curr.num_matched recompute_blocks = max(0, seq.prefix_cache.match_recompute_blocks) + max_candidate_step = seq.get_prefix_cache_max_candidate_step() + normal_candidate_step = max(0, seq.num_valid_ids - 1) + input_scoring_caps_match = max_candidate_step < normal_candidate_step raw_match_step = -1 - if recompute_blocks == 0: - max_step = ((seq.num_valid_ids - 1) // self.block_size) * self.block_size + if recompute_blocks == 0 and not input_scoring_caps_match: + max_step = seq.get_prefix_cache_max_match_step() else: - # SSM checkpoint lookup is sparse, but MTP overlap needs the raw KV - # trie depth so the whole checkpoint-to-raw-hit span stays private. + # SSM checkpoint lookup is sparse, but both MTP overlap and input + # scoring need the raw KV trie depth so the whole + # checkpoint-to-raw-hit span stays private while it is recomputed. raw_match_step = self._find_raw_block_match_step(seq, init_curr) - max_step = max(init_num_matched, raw_match_step - recompute_blocks * self.block_size) - max_step = seq.clamp_prefix_cache_match_step(max_step) + capped_match_step = seq.clamp_prefix_cache_match_step(min(raw_match_step, max_candidate_step)) + max_step = max(init_num_matched, capped_match_step - recompute_blocks * self.block_size) + max_step = seq.clamp_prefix_cache_match_step(max_step) steps = self._state_checkpoint_steps.get(seq.adapter_name, ()) for step in sorted((step for step in steps if init_num_matched < step <= max_step), reverse=True): if seq.clamp_prefix_cache_match_step(step) != step: @@ -1086,8 +1091,7 @@ def _match_state_checkpoint(self, seq: SchedulerSequence): seq.prefix_cache.restore_state = node.state_idx seq.prefix_cache.restore_node = node seq.prefix_cache.last_shared_node = node - if recompute_blocks > 0: - self._set_private_recompute_range(seq, step, raw_match_step) + self._set_private_recompute_range(seq, step, raw_match_step) self._record_match_stats(seq, query_tokens=seq.num_all_ids - init_num_matched, hit_tokens=step - init_num_matched) @@ -1098,7 +1102,10 @@ def _match_state_checkpoint(self, seq: SchedulerSequence): return seq.prefix_cache.last_shared_node = init_curr - self._clear_private_recompute_range(seq) + if raw_match_step > init_num_matched: + self._set_private_recompute_range(seq, init_num_matched, raw_match_step) + else: + self._clear_private_recompute_range(seq) self._record_match_stats(seq, query_tokens=seq.num_all_ids - init_num_matched) if logger.isEnabledFor(logging.DEBUG): logger.debug(f'SSM prefix-cache miss: session_id={seq.session_id} seq_id={seq.seq_id} ' @@ -1180,9 +1187,9 @@ def __clamp_match_step(match_step: int): curr = init_curr num_matched = init_num_matched - max_match_step = seq.get_prefix_cache_max_match_step() raw_num_matched = num_matched - effective_num_matched = seq.clamp_prefix_cache_match_step(min(raw_num_matched, max_match_step)) + effective_num_matched = seq.clamp_prefix_cache_match_step( + min(raw_num_matched, seq.get_prefix_cache_max_match_step())) __clamp_match_step(effective_num_matched) self._set_private_recompute_range(seq, num_matched, raw_num_matched) @@ -1205,7 +1212,7 @@ def __clamp_match_step(match_step: int): if logger.isEnabledFor(logging.DEBUG): logger.debug(f'Prefix-cache match: session_id={seq.session_id} seq_id={seq.seq_id} ' f'init_step={init_num_matched} matched_step={num_matched} ' - f'candidate_step={raw_num_matched} ' + f'effective_step={effective_num_matched} raw_candidate_step={raw_num_matched} ' f'clamped={effective_num_matched != raw_num_matched}') def allocate(self, seq: SchedulerSequence): diff --git a/lmdeploy/pytorch/strategies/ar/sequence.py b/lmdeploy/pytorch/strategies/ar/sequence.py index 0d39b7e0ae..37e697e636 100644 --- a/lmdeploy/pytorch/strategies/ar/sequence.py +++ b/lmdeploy/pytorch/strategies/ar/sequence.py @@ -144,6 +144,12 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode all_routed_experts): if msg.status != MessageStatus.RUNNING: continue + if update_mode == UpdateTokenMode.PREFILL and msg.sampling_param.max_new_tokens == 0: + msg.update_token_ids(np.empty((0,), dtype=np.int64), + model_meta=model_meta, + mode=UpdateTokenMode.PREFILL) + msg.state.finish() + continue # fill token msg.update_token_ids(token, model_meta=model_meta, mode=update_mode, routed_experts=routed_experts) diff --git a/lmdeploy/serve/core/async_engine.py b/lmdeploy/serve/core/async_engine.py index 301622557d..4b5e093656 100644 --- a/lmdeploy/serve/core/async_engine.py +++ b/lmdeploy/serve/core/async_engine.py @@ -564,8 +564,8 @@ def remove_session_once(): gen_config = self._determine_gen_config(session, input_ids, gen_config=gen_config) input_len = len(input_ids) - - if gen_config.max_new_tokens == 0: + input_logprobs_requested = gen_config.logprob_start_len >= 0 + if gen_config.max_new_tokens == 0 and not input_logprobs_requested: logger.info(f'run out of tokens. session={session_id}.') metrics_processor.increase_failed_requests('error') remove_session_once() @@ -638,7 +638,6 @@ def is_error(status): logger.debug(f'[generate] session {session_id} started') hit_stop_token = 0 req_stats = RequestStats(prompt_tokens=input_len) # per-request stats - # We use this as default outputs in case the async_stream_infer of the Engine yields empty generator. outputs = EngineOutput(ResponseType.INTERNAL_ENGINE_ERROR, []) @@ -697,7 +696,8 @@ def is_error(status): finish_reason = 'abort' metrics_processor.increase_failed_requests('abort') else: - finish_reason = 'stop' if outputs.token_ids[-1] in stop_ids else 'length' + finish_reason = ('stop' if outputs.token_ids + and outputs.token_ids[-1] in stop_ids else 'length') metrics_processor.increase_succeeded_requests() # utf-8 char at the end means it's a potential unfinished byte sequence @@ -705,7 +705,9 @@ def is_error(status): # avoid returning the last response twice response = '' token_ids, logits, last_hidden_state, logprobs = [], None, None, None - if gen_config.include_stop_str_in_output and finish_reason == 'stop': + if input_logprobs_requested: + logprobs = outputs.logprobs + elif gen_config.include_stop_str_in_output and finish_reason == 'stop': # return the eos token id (MUST be in a list), eos string, eos token's logits and so on token_ids = outputs.token_ids[-1:] response = self.tokenizer.decode(token_ids, skip_special_tokens=False) diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py index 9fe802ca1f..2675dd9857 100644 --- a/lmdeploy/serve/openai/api_server.py +++ b/lmdeploy/serve/openai/api_server.py @@ -277,16 +277,15 @@ def _create_chat_completion_logprobs(tokenizer: PreTrainedTokenizerBase, return ChoiceLogprobs(content=content) -def _create_output_token_logprobs(token_ids: list[int] | None = None, - logprobs: list[dict[int, float]] | None = None): - """Create raw (logprob, token_id) pairs for output tokens.""" +def _create_token_logprobs(token_ids: list[int] | None = None, + logprobs: list[dict[int, float]] | None = None): + """Create raw (logprob, token_id) pairs.""" 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 @router.get('/health') @@ -593,7 +592,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: if request.logprobs and res.logprobs: logprobs = _create_chat_completion_logprobs(tokenizer, res.token_ids, res.logprobs) if request.return_logprob: - output_token_logprobs = _create_output_token_logprobs(res.token_ids, res.logprobs) + output_token_logprobs = _create_token_logprobs(res.token_ids, res.logprobs) if res.finish_reason and include_usage: final_usage = UsageInfo.build( prompt_tokens=res.input_token_len, @@ -712,7 +711,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: logprobs = _create_chat_completion_logprobs(tokenizer, final_token_ids, final_logprobs) output_token_logprobs = None if request.return_logprob and len(final_logprobs): - output_token_logprobs = _create_output_token_logprobs(final_token_ids, final_logprobs) + output_token_logprobs = _create_token_logprobs(final_token_ids, final_logprobs) assert final_res is not None choices = [] @@ -983,10 +982,15 @@ async def generate(request: GenerateReqInput, raw_request: Request = None): if error_check_ret is not None: return error_check_ret - session = VariableInterface.create_session(request.session_id) - prompt = request.prompt input_ids = request.input_ids + input_logprobs_requested = request.logprob_start_len >= 0 + input_logprob_token_ids = (input_ids[request.logprob_start_len + 1:] + if input_logprobs_requested + and input_ids is not None else None) + + session = VariableInterface.create_session(request.session_id) + image_data = request.image_data if image_data is not None: # convert to openai format @@ -1005,6 +1009,7 @@ async def generate(request: GenerateReqInput, raw_request: Request = None): gen_config = _build_serving_generation_config( request, logprobs=1 if request.return_logprob else None, + logprob_start_len=request.logprob_start_len, stop_words=request.stop, ) @@ -1018,11 +1023,18 @@ async def generate(request: GenerateReqInput, raw_request: Request = None): media_io_kwargs=request.media_io_kwargs, mm_processor_kwargs=request.mm_processor_kwargs) - def create_generate_response_json(res, text, output_ids, logprobs, finish_reason, routed_experts=None): + def create_generate_response_json(res, + text, + output_ids, + logprobs, + input_logprobs, + finish_reason, + routed_experts=None): # only output router experts in last chunk routed_experts = None if finish_reason is None else routed_experts meta = GenerateReqMetaOutput(finish_reason=dict(type=finish_reason) if finish_reason else None, - output_token_logprobs=logprobs or None, + output_token_logprobs=logprobs, + input_token_logprobs=input_logprobs, prompt_tokens=res.input_token_len, routed_experts=routed_experts, completion_tokens=res.generate_token_len) @@ -1035,14 +1047,17 @@ async def generate_stream_generator(): text = res.response or '' output_ids = res.token_ids routed_experts = res.routed_experts - logprobs = [] - if res.logprobs: - for tok, tok_logprobs in zip(res.token_ids, res.logprobs): - logprobs.append((tok_logprobs[tok], tok)) + if input_logprobs_requested: + logprobs = None + input_logprobs = _create_token_logprobs(input_logprob_token_ids, res.logprobs) + else: + logprobs = _create_token_logprobs(res.token_ids, res.logprobs) + input_logprobs = None response_json = create_generate_response_json(res, text, output_ids, logprobs, + input_logprobs, res.finish_reason, routed_experts=routed_experts) yield f'data: {response_json}\n\n' @@ -1058,6 +1073,7 @@ async def _inner_call(): text = '' output_ids = [] logprobs = [] + input_token_logprobs = None async with aclosing(_with_request_cleanup(result_generator, [result_generator], [session])) as generator: async for res in generator: if await raw_request.is_disconnected(): @@ -1066,16 +1082,20 @@ async def _inner_call(): return create_error_response(HTTPStatus.BAD_REQUEST, 'Client disconnected') text += res.response or '' output_ids.extend(res.token_ids or []) - logprobs.extend(res.logprobs or []) + if input_logprobs_requested: + cur_input_logprobs = _create_token_logprobs(input_logprob_token_ids, + res.logprobs) + if cur_input_logprobs is not None: + input_token_logprobs = cur_input_logprobs + elif res.logprobs: + logprobs.extend(res.logprobs) - output_token_logprobs = [] - if len(logprobs) and len(output_ids): - for tok, tok_logprobs in zip(output_ids, logprobs): - output_token_logprobs.append((tok_logprobs[tok], tok)) + output_token_logprobs = _create_token_logprobs(output_ids, logprobs) if logprobs else None nonlocal response meta = GenerateReqMetaOutput(finish_reason=dict(type=res.finish_reason) if res.finish_reason else None, output_token_logprobs=output_token_logprobs or None, + input_token_logprobs=input_token_logprobs, prompt_tokens=res.input_token_len, routed_experts=res.routed_experts, completion_tokens=res.generate_token_len) diff --git a/lmdeploy/serve/openai/protocol.py b/lmdeploy/serve/openai/protocol.py index 92a48d1b51..420e74a719 100644 --- a/lmdeploy/serve/openai/protocol.py +++ b/lmdeploy/serve/openai/protocol.py @@ -545,6 +545,7 @@ class GenerateReqInput(BaseModel): input_ids: list[int] | None = None image_data: ImageDataFormat | None = None return_logprob: bool | None = None + logprob_start_len: int = Field(default=-1, ge=-1) max_tokens: int = 128 stop: str | list[str] | None = None stop_token_ids: list[int] | None = None @@ -578,6 +579,7 @@ class GenerateReqMetaOutput(BaseModel): completion_tokens: int | None = None finish_reason: dict[str, Any] | None = None output_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) + input_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) routed_experts: list[list[list[int]]] | str | None = None # (num_token, num_layer, topk_expert) diff --git a/lmdeploy/serve/openai/serving_generate.py b/lmdeploy/serve/openai/serving_generate.py index 10e9dde597..2388d0c199 100644 --- a/lmdeploy/serve/openai/serving_generate.py +++ b/lmdeploy/serve/openai/serving_generate.py @@ -1,6 +1,9 @@ # Copyright (c) OpenMMLab. All rights reserved. from typing import TYPE_CHECKING +from lmdeploy.messages import PytorchEngineConfig +from lmdeploy.pytorch.disagg.config import EngineRole + from .protocol import GenerateReqInput if TYPE_CHECKING: @@ -10,14 +13,10 @@ def check_request(request: GenerateReqInput, server_context: 'VariableInterface') -> str: engine_config = server_context.get_engine_config() session_manager = server_context.get_session_manager() - try: - # Check logprobs settings - logprobs_mode = engine_config.logprobs_mode - return_logprob = request.return_logprob - if logprobs_mode is None and return_logprob: - return f'return_logprob({return_logprob}) requested but not enabled logprobs_mode in engine configuration.' - except AttributeError: - pass + logprobs_mode = getattr(engine_config, 'logprobs_mode', None) + return_logprob = request.return_logprob + if logprobs_mode is None and return_logprob: + return f'return_logprob({return_logprob}) requested but not enabled logprobs_mode in engine configuration.' if (request.prompt is not None) ^ (request.input_ids is None): return 'You must specify exactly one of prompt or input_ids' @@ -28,9 +27,26 @@ def check_request(request: GenerateReqInput, server_context: 'VariableInterface' if request.input_ids is not None and len(request.input_ids) == 0: return 'The input_ids must not be an empty list' - if request.max_tokens is not None and request.max_tokens <= 0: + if request.max_tokens is not None and request.max_tokens < 0: + return f'The max_tokens {request.max_tokens!r} must be non-negative.' + + # 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)}).') + elif request.max_tokens is not None and request.max_tokens == 0: return f'The max_tokens {request.max_tokens!r} must be a positive integer.' - if session_manager.has(request.session_id): return f'The session_id {request.session_id!r} is occupied.' diff --git a/tests/pytorch/engine/test_inputs_maker.py b/tests/pytorch/engine/test_inputs_maker.py index 228618297b..4ac19ce7a9 100644 --- a/tests/pytorch/engine/test_inputs_maker.py +++ b/tests/pytorch/engine/test_inputs_maker.py @@ -4,6 +4,7 @@ from types import SimpleNamespace import pytest +import torch import lmdeploy.pytorch.engine.inputs_maker as inputs_maker_module from lmdeploy.pytorch.disagg.config import EngineRole @@ -12,12 +13,51 @@ InputsMakerAsync, InputsMakerConfig, LongContextChunker, + _build_logits_indices, _compact_state_prefix_cache_restore_offsets, _compact_state_prefix_cache_save_offsets, ) from lmdeploy.pytorch.messages import MessageStatus +def _logprob_seq(history, input_start, input_end, logprob_start, all_ids): + return SimpleNamespace(num_history_ids=history, + input_start_pos=input_start, + input_end_pos=input_end, + logprob_start_pos=logprob_start, + all_ids=all_ids) + + +def test_logits_indices_keep_batch_alignment_and_exact_targets(): + disabled = _logprob_seq(2, 0, 5, -1, list(range(5))) + enabled = _logprob_seq(10, 8, 15, 11, list(range(20))) + empty = _logprob_seq(20, 17, 20, 20, list(range(24))) + + indices, row_counts = _build_logits_indices([disabled, enabled, empty], + [3, 5, 0]) + assert indices.tolist() == [4, 5, 6] + assert row_counts.tolist() == [0, 3, 0] + + +def test_logits_indices_shift_cross_chunk_target_to_next_chunk(): + first = _logprob_seq(0, 0, 10, 2, list(range(12))) + indices, row_counts = _build_logits_indices([first], [4]) + assert indices.tolist() == [2, 3] + assert row_counts.tolist() == [1] + + middle = _logprob_seq(4, 0, 10, 2, list(range(12))) + indices, row_counts = _build_logits_indices([middle], [4]) + assert indices.tolist() == [0, 1, 2, 3] + assert row_counts.tolist() == [4] + + # Final chunks exclude the last input hidden, which is reserved for + # generated/control-token sampling, but include the carried previous row. + final = _logprob_seq(8, 0, 10, 2, list(range(12))) + indices, row_counts = _build_logits_indices([final], [2]) + assert indices.tolist() == [0] + assert row_counts.tolist() == [2] + + @dataclass class _DummyMultiModal: start: int @@ -41,6 +81,7 @@ def __init__(self, self.return_routed_experts = False self.return_ce_loss = False self.status = MessageStatus.RUNNING + self.sampling_param = SimpleNamespace(max_new_tokens=512) def set_step(self, step: int): self.num_history_ids = step @@ -1102,6 +1143,36 @@ def test_normal_prefill_can_update_running_while_long_chunker_is_active(): assert maker.long_context_chunker.next_step == 0 +def test_scoring_prefill_never_enters_decode_running_set(): + seq = _DummySeq(history_ids=0, + token_ids=16, + all_multimodals={}, + input_multimodals={}) + seq.sampling_param.max_new_tokens = 0 + model_inputs = _fake_model_inputs() + model_inputs.logits_indices = torch.tensor([0]) + maker = _make_policy_maker(seq) + + maker.update_running_seqs([seq], model_inputs) + + assert maker.running_seqs == [] + + +def test_prefill_running_update_filters_by_decode_need_not_logits_metadata(): + long_seq = _DummySeq(history_ids=0, token_ids=1024, all_multimodals={}, input_multimodals={}) + scoring_seq = _DummySeq(history_ids=0, token_ids=16, all_multimodals={}, input_multimodals={}) + decode_seq = _DummySeq(history_ids=0, token_ids=16, all_multimodals={}, input_multimodals={}) + scoring_seq.sampling_param.max_new_tokens = 0 + decode_seq.sampling_param.max_new_tokens = 3 + model_inputs = _fake_model_inputs() + model_inputs.logits_indices = torch.tensor([0]) + maker = _make_policy_maker(long_seq) + + maker.update_running_seqs([scoring_seq, decode_seq], model_inputs) + + assert maker.running_seqs == [decode_seq] + + def test_last_long_context_chunk_runs_as_prefill_on_prefill_turn(): image = _DummyMultiModal(start=600, end=700) long_seq = _DummySeq(history_ids=512, diff --git a/tests/pytorch/engine/test_model_agent.py b/tests/pytorch/engine/test_model_agent.py index 859b93b557..ae1c05847d 100644 --- a/tests/pytorch/engine/test_model_agent.py +++ b/tests/pytorch/engine/test_model_agent.py @@ -7,6 +7,312 @@ import pytest import torch +from lmdeploy.pytorch.engine.logits_process import SamplingInputs +from lmdeploy.pytorch.engine.model_agent.agent import BaseModelAgent +from lmdeploy.pytorch.model_inputs import ModelInputs + + +def _input_logprob_model_inputs(input_ids, indices): + size = len(input_ids) + return ModelInputs(input_ids=torch.tensor([input_ids]), + seq_length=torch.tensor([size]), + history_lengths=torch.tensor([0]), + block_offsets=torch.zeros((1, 1), dtype=torch.long), + is_decoding=False, + num_ignored_history=torch.tensor([0]), + max_q_seqlen=size, + max_kv_seqlen=size, + sum_kv_seqlen=size, + logits_indices=None if indices is None else torch.tensor(indices, dtype=torch.long), + seq_logit_length=None if indices is None else torch.tensor([len(indices)])) + + +def test_get_input_logits_projects_only_selected_rows(): + vocab_size = 8 + hidden = torch.arange(18, dtype=torch.float32).reshape(1, 6, 3) + inputs = _input_logprob_model_inputs([0, 1, 2, 3, 4, 5], [0, 1, 3, 4]) + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.model_config = SimpleNamespace(vocab_size=vocab_size) + calls = [] + + def get_logits(values): + calls.append(values.shape[1]) + torch.testing.assert_close(values[0], hidden[0, [0, 1, 3, 4]]) + return torch.arange(values.shape[1] * vocab_size, dtype=torch.float32).reshape( + 1, values.shape[1], vocab_size) + + agent.get_logits = get_logits + output = agent._get_input_logits(hidden, inputs) + + assert calls == [4] + assert output.shape == (4, vocab_size) + + +def test_prefill_input_logprobs_helper_is_semantic(): + agent = BaseModelAgent.__new__(BaseModelAgent) + inputs = _input_logprob_model_inputs([0, 1, 2], [0, 1]) + assert agent._is_prefill_input_logprobs(inputs) + + missing_lengths = _input_logprob_model_inputs([0, 1, 2], [0, 1]) + missing_lengths.seq_logit_length = None + assert not agent._is_prefill_input_logprobs(missing_lengths) + + decode = _input_logprob_model_inputs([0, 1, 2], [0, 1]) + decode.is_decoding = True + assert not agent._is_prefill_input_logprobs(decode) + + dummy = _input_logprob_model_inputs([0, 1, 2], [0, 1]) + dummy.is_dummy = True + assert not agent._is_prefill_input_logprobs(dummy) + + ordinary = _input_logprob_model_inputs([0, 1, 2], None) + assert not agent._is_prefill_input_logprobs(ordinary) + + +@pytest.mark.parametrize('return_logits', [False, True]) +def test_async_forward_returns_only_input_logits_without_sampling_policy( + return_logits): + vocab_size = 8 + hidden = torch.arange(9, dtype=torch.float32).reshape(1, 3, 3) + inputs = _input_logprob_model_inputs([0, 1, 2], [1, 2]) + calls = [] + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.model_config = SimpleNamespace(vocab_size=vocab_size) + + async def async_forward(_inputs): + return {'hidden_states': [hidden]} + + def get_logits(values): + calls.append(values.shape[1]) + return torch.zeros((1, values.shape[1], vocab_size)) + + agent.async_forward = async_forward + agent.get_logits = get_logits + agent._postprocess_forward_output = lambda output, _inputs: output + agent.spec_agent = SimpleNamespace( + update_main_model_outputs=lambda *_: pytest.fail('scoring path must not enter speculative processing')) + + output = asyncio.run(agent._async_model_forward(inputs, return_logits=return_logits)) + + assert calls == [2] + assert 'input_logits' not in output + assert output['logits'].shape == (2, vocab_size) + + +def test_async_forward_requires_complete_prefill_input_logprob_metadata(): + vocab_size = 8 + hidden = torch.arange(15, dtype=torch.float32).reshape(1, 5, 3) + inputs = _input_logprob_model_inputs([0, 1, 2, 3, 4], [0, 1]) + inputs.seq_logit_length = None + calls = [] + agent = BaseModelAgent.__new__(BaseModelAgent) + + async def async_forward(_inputs): + return {'hidden_states': [hidden]} + + def get_logits(values): + calls.append(values.shape[1]) + return torch.zeros((1, values.shape[1], vocab_size)) + + agent.async_forward = async_forward + agent.get_logits = get_logits + agent._postprocess_forward_output = lambda output, _inputs: output + agent.spec_agent = SimpleNamespace( + is_enabled=lambda: False, + update_main_model_outputs=lambda output, _inputs: + (output.pop('hidden_states')[0], output)) + + output = asyncio.run(agent._async_model_forward(inputs, return_logits=False)) + + assert calls == [5] + assert 'input_logits' not in output + assert output['logits'].shape == (1, 5, vocab_size) + + +def test_async_forward_disabled_path_keeps_one_sampling_projection(): + vocab_size = 8 + hidden = torch.arange(15, dtype=torch.float32).reshape(1, 5, 3) + inputs = _input_logprob_model_inputs([0, 1, 2, 3, 4], None) + calls = [] + agent = BaseModelAgent.__new__(BaseModelAgent) + + async def async_forward(_inputs): + return {'hidden_states': [hidden]} + + def get_logits(values): + calls.append(values.shape[1]) + return torch.zeros((1, values.shape[1], vocab_size)) + + agent.async_forward = async_forward + agent.get_logits = get_logits + agent._postprocess_forward_output = lambda output, _inputs: output + agent.spec_agent = SimpleNamespace( + is_enabled=lambda: False, + update_main_model_outputs=lambda output, _inputs: (output.pop('hidden_states')[0], output)) + + output = asyncio.run(agent._async_model_forward(inputs, return_logits=False)) + + assert calls == [5] + assert 'input_logits' not in output + + +@pytest.mark.parametrize('mode', ['raw_logits', 'raw_logprobs']) +@pytest.mark.parametrize('num_logprobs', [0, 1, 2]) +def test_input_compaction_matches_shared_reference(mode, num_logprobs): + inputs = _input_logprob_model_inputs([0, 1, 2], [0, 1]) + input_logits = torch.tensor([[0.1, 0.2, 0.3, 0.4], [1.0, 0.5, -0.5, -1.0]]) + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.misc_config = SimpleNamespace(logprobs_mode=mode) + + outputs = agent._get_outputs_with_logprobs(input_logits, inputs, num_logprobs, model_metas=[{'meta': 1}]) + compact = outputs.logprobs + + reference = input_logits.log_softmax(-1) if mode == 'raw_logprobs' else input_logits + targets = torch.tensor([1, 2]) + torch.testing.assert_close(compact.vals[:, 0], reference.gather(-1, targets[:, None])[:, 0]) + assert compact.indices[:, 0].tolist() == targets.tolist() + assert compact.vals.shape == (2, num_logprobs + 1) + assert outputs.next_token_ids.tolist() == [0] + assert outputs.stopped.tolist() == [True] + assert outputs.stop_pos.tolist() == [-1] + assert outputs.model_metas == [{'meta': 1}] + + +def test_input_compaction_reuses_previous_chunk_last_logit_for_cross_chunk_target(): + first_inputs = _input_logprob_model_inputs([0, 1, 2], [0, 1, 2]) + first_inputs.is_chunk = True + first_inputs.is_first_chunk = True + first_inputs.is_last_chunk = False + first_logits = torch.arange(24, dtype=torch.float32).reshape(3, 8) + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.misc_config = SimpleNamespace(logprobs_mode='raw_logits') + agent._prev_chunk_last_logit = None + + first_outputs = agent._get_outputs_with_logprobs(first_logits, first_inputs, 0, model_metas=None) + assert first_outputs.logprobs.indices[:, 0].tolist() == [1, 2] + torch.testing.assert_close(agent._prev_chunk_last_logit, first_logits[-1:]) + + final_inputs = _input_logprob_model_inputs([3, 4, 5], [0, 1]) + final_inputs.is_chunk = True + final_inputs.is_first_chunk = False + final_inputs.is_last_chunk = True + final_logits = torch.arange(100, 116, dtype=torch.float32).reshape(2, 8) + + outputs = agent._get_outputs_with_logprobs(final_logits, final_inputs, 0, model_metas=None) + compact = outputs.logprobs + + expected_logits = torch.cat([first_logits[-1:], final_logits], dim=0) + assert compact.indices[:, 0].tolist() == [3, 4, 5] + torch.testing.assert_close(compact.vals[:, 0], expected_logits.gather(-1, compact.indices.long())[:, 0]) + assert agent._prev_chunk_last_logit is None + + +def test_input_compaction_accepts_empty_pre_boundary_chunk(): + inputs = _input_logprob_model_inputs([0, 1, 2], []) + inputs.is_chunk = True + inputs.is_first_chunk = True + inputs.is_last_chunk = False + input_logits = torch.empty((0, 8)) + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.misc_config = SimpleNamespace(logprobs_mode='raw_logits') + agent._prev_chunk_last_logit = None + + outputs = agent._get_outputs_with_logprobs(input_logits, inputs, 2, model_metas=None) + + assert outputs.logprobs.vals.shape == (0, 3) + assert outputs.logprobs.indices.shape == (0, 3) + assert agent._prev_chunk_last_logit is None + + +def test_dp_dummy_skips_input_compaction_without_sampling_inputs(): + inputs = _input_logprob_model_inputs([0], None) + inputs.is_dummy = True + generated = SimpleNamespace(next_token_ids=torch.tensor([0]), + output_token_ids=torch.tensor([[0]]), + logprobs=None) + + async def spec_sampling(_inputs, _extra, _sampling): + return generated + + async def spec_forward(_inputs, extra, _sampling): + return extra + + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.rank = 0 + agent.spec_agent = SimpleNamespace(is_enabled=lambda: True, + async_sampling_logits=spec_sampling, + async_model_forward=spec_forward) + + result = asyncio.run( + agent._step_postprocess_with_output(last_logits=torch.zeros((1, 8)), + logits=torch.zeros((1, 8)), + inputs=inputs, + sampling_inputs=None, + stopping_criteria=None, + model_metas=None, + need_broadcast_next=False)) + + assert result[-1] is generated.next_token_ids + + +def test_mtp_sampling_still_delegates_generated_rows(): + generated = SimpleNamespace(next_token_ids=torch.tensor([5]), + output_token_ids=torch.tensor([[5, 6]]), + logprobs='spec-generated') + + async def spec_sampling(_inputs, _extra, _sampling): + return generated + + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.spec_agent = SimpleNamespace(is_enabled=lambda: True, + async_sampling_logits=spec_sampling) + result = asyncio.run( + agent.async_sampling_logits(torch.zeros((1, 8)), + _input_logprob_model_inputs([0, 1, 2], [0, 1]), + SimpleNamespace(), SamplingInputs(max_num_logprobs=0))) + + assert result[0] is generated.next_token_ids + assert result[1] == 'spec-generated' + assert result[2] is generated.output_token_ids + + +def test_ordinary_non_final_chunk_keeps_upstream_generated_logprob_row(monkeypatch): + import lmdeploy.pytorch.engine.model_agent.agent as agent_module + + class _Processor: + + def __init__(self, *args, **kwargs): + pass + + async def __call__(self, logits): + return logits, logits + + def sampling(self, logits): + return torch.tensor([3]) + + async def accept_guided_tokens(self, token_ids): + pass + + def compute_logprobs(self, raw_logprobs, token_ids): + indices = token_ids[:, None] + return raw_logprobs.gather(-1, indices), indices.to(torch.int32) + + monkeypatch.setattr(agent_module, 'FusedLogitsProcessor', _Processor) + inputs = _input_logprob_model_inputs([0, 1, 2], [0, 1]) + inputs.is_chunk = True + inputs.is_last_chunk = False + agent = BaseModelAgent.__new__(BaseModelAgent) + agent.misc_config = SimpleNamespace(logprobs_mode='raw_logits') + agent.guided_decoding_manager = None + agent.spec_agent = SimpleNamespace(is_enabled=lambda: False) + agent.agent_strategy = SimpleNamespace(post_sampling=lambda _i, _l, ids, extra: (ids, extra)) + + result = asyncio.run( + agent.async_sampling_logits(torch.zeros((1, 8)), inputs, SimpleNamespace(), + SamplingInputs(max_num_logprobs=0))) + + assert result[1] is not None + @pytest.fixture def event_loop(): diff --git a/tests/pytorch/engine/test_prefill_logprobs.py b/tests/pytorch/engine/test_prefill_logprobs.py new file mode 100644 index 0000000000..3dcd696fb8 --- /dev/null +++ b/tests/pytorch/engine/test_prefill_logprobs.py @@ -0,0 +1,258 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from lmdeploy.messages import ResponseType +from lmdeploy.pytorch.engine.engine_loop import EngineLoop +from lmdeploy.pytorch.engine.model_agent.agent import BatchedLogProbs +from lmdeploy.pytorch.messages import ( + MessageStatus, + SamplingParam, + SchedulerSequence, + SequenceMeta, + UpdateTokenMode, +) +from lmdeploy.pytorch.strategies.ar.sequence import ( + ARSequenceStrategy, + SchedulerSequenceDefault, +) + + +def _compact(actual_ids, width=1): + if not actual_ids: + return BatchedLogProbs(vals=torch.empty((0, width)), + indices=torch.empty((0, width), + dtype=torch.int32)) + indices = torch.tensor([[token_id] * width for token_id in actual_ids], + dtype=torch.int32) + vals = torch.tensor([[-float(token_id)] * width + for token_id in actual_ids]) + return BatchedLogProbs(vals=vals, indices=indices) + + +def test_logprob_start_position_is_request_relative(): + seq = object.__new__(SchedulerSequence) + seq.input_start_pos = 17 + seq.sampling_param = SamplingParam(num_logprobs=0, + logprob_start_len=4) + assert seq.logprob_start_pos == 21 + seq.sampling_param = SamplingParam(num_logprobs=0, + logprob_start_len=-1) + assert seq.logprob_start_pos == -1 + + +def test_engine_loop_accumulates_mixed_rows_in_chunk_order(): + loop = EngineLoop.__new__(EngineLoop) + messages = [ + SimpleNamespace(logprob_start_pos=0, + resp=SimpleNamespace(_input_logprobs=None)), + SimpleNamespace(logprob_start_pos=0, + resp=SimpleNamespace(_input_logprobs=None)), + ] + inputs = SimpleNamespace( + is_decoding=False, + logits_indices=torch.tensor([0, 1, 3]), + seq_logit_length=torch.tensor([2, 1]), + ) + outputs = SimpleNamespace(logprobs=_compact([2, 3, 9])) + + loop._append_input_logprobs(outputs, messages, inputs) + + assert messages[0].resp._input_logprobs == [ + ([-2.0], [2]), + ([-3.0], [3]), + ] + assert messages[1].resp._input_logprobs == [([-9.0], [9])] + + empty_inputs = SimpleNamespace( + is_decoding=False, + logits_indices=torch.empty(0, dtype=torch.long), + seq_logit_length=torch.tensor([0, 0]), + ) + loop._append_input_logprobs(SimpleNamespace(logprobs=_compact([])), + messages, empty_inputs) + assert messages[0].resp._input_logprobs[-1] == ([-3.0], [3]) + + unrequested = SimpleNamespace( + logprob_start_pos=-1, + resp=SimpleNamespace(_input_logprobs=None), + ) + loop._append_input_logprobs( + SimpleNamespace(logprobs=_compact([])), + [unrequested], + SimpleNamespace(is_decoding=False, + logits_indices=torch.empty(0, dtype=torch.long), + seq_logit_length=torch.tensor([0])), + ) + assert unrequested.resp._input_logprobs is None + + +def test_engine_loop_rejects_malformed_compact_scoring_segments(): + messages = [ + SimpleNamespace(logprob_start_pos=0, + resp=SimpleNamespace(_input_logprobs=None)), + SimpleNamespace(logprob_start_pos=0, + resp=SimpleNamespace(_input_logprobs=None)), + ] + inputs = SimpleNamespace( + is_decoding=False, + logits_indices=torch.arange(2), + seq_logit_length=torch.tensor([1, 1]), + ) + with pytest.raises(RuntimeError, match='compact output mismatch'): + EngineLoop._append_input_logprobs( + SimpleNamespace(logprobs=_compact([0])), + messages, inputs) + + +@pytest.mark.parametrize(('actual_ids', 'expected'), [ + ([], None), + ([2, 3], [([-2.0], [2]), ([-3.0], [3])]), +]) +def test_final_scoring_prefill_uses_existing_logprob_carrier(actual_ids, + expected): + loop = EngineLoop.__new__(EngineLoop) + + def finish_running(**kwargs): + kwargs['running'][0].status = MessageStatus.STOPPED + + loop.seq_strategy = SimpleNamespace(update_running=finish_running) + loop.scheduler = SimpleNamespace( + block_trie=SimpleNamespace(cache_routed_experts=lambda running: None)) + loop.config = SimpleNamespace(num_speculative_tokens=None, + enable_metrics=False) + response = SimpleNamespace(data=None, + _input_logprobs=None, + is_done=False, + type=ResponseType.SUCCESS) + msg = SimpleNamespace( + session_id=7, + status=MessageStatus.RUNNING, + num_token_ids=3, + generated_ids=np.empty((0,), dtype=np.int64), + logprob_start_pos=0, + resp=response, + resp_cache=False, + sampling_param=SamplingParam(max_new_tokens=0, + num_logprobs=0, + logprob_start_len=0), + cached_tokens=0, + engine_events=[], + routed_experts=None, + return_ce_loss=False, + return_logits=False, + ) + inputs = SimpleNamespace( + is_decoding=False, + is_chunk=False, + is_last_chunk=False, + logits_indices=torch.arange(len(actual_ids)), + seq_logit_length=torch.tensor([len(actual_ids)]), + ) + batched = SimpleNamespace( + logprobs=_compact(actual_ids), + logits=None, + all_routed_experts=None, + ce_loss=None, + new_token_timestamp=0, + stop_pos=torch.tensor([-1]), + next_token_ids=torch.tensor([0]), + ) + + result = loop._make_infer_outputs(batched, [msg], inputs, delta=None)[7] + + assert result.finish + assert result.token_ids.tolist() == [] + assert result.logprobs == expected + + sent = [] + loop._send_resp = sent.append + loop._send_resps([result]) + if expected is None: + assert response.data is None + else: + assert response.data['logprobs'] == [ + dict(zip(indices, vals)) for vals, indices in expected + ] + assert sent == [result] + + +def test_cancelled_scoring_response_drops_partial_accumulator(monkeypatch): + """Cancellation must never expose rows accumulated by earlier chunks.""" + loop = EngineLoop.__new__(EngineLoop) + loop.req_manager = object() + sent = [] + monkeypatch.setattr( + 'lmdeploy.pytorch.engine.engine_loop.response_reqs', + lambda req_manager, resp, resp_type, data: sent.append( + (resp_type, data))) + response = SimpleNamespace( + data=None, + _input_logprobs=[([-0.5], [2])], + is_done=True, + type=ResponseType.CANCEL, + ) + out = SimpleNamespace( + resp=response, + finish=False, + token_ids=np.empty((0,), dtype=np.int64), + logits=None, + cache_block_ids=None, + req_metrics=None, + routed_experts=None, + ce_loss=None, + ) + + loop._send_resp(out) + + assert response._input_logprobs == [([-0.5], [2])] + assert sent == [(ResponseType.CANCEL, { + 'token_ids': out.token_ids, + 'logits': None, + 'cache_block_ids': None, + 'req_metrics': None, + 'routed_experts': None, + 'logprobs': None, + 'ce_loss': None, + })] + + +@pytest.mark.parametrize('sampling_param', [ + SamplingParam(max_new_tokens=0, logprob_start_len=-1), + SamplingParam(max_new_tokens=0, num_logprobs=0, logprob_start_len=0), +]) +def test_zero_token_prefill_finishes_without_committing_inert_token( + sampling_param): + strategy = ARSequenceStrategy() + session = SimpleNamespace( + seq_meta=SequenceMeta(block_size=16, strategy=strategy), + session_id=1, + ) + seq = SchedulerSequenceDefault( + seq_id=1, + session=session, + sampling_param=sampling_param, + ) + state = SimpleNamespace(status=MessageStatus.RUNNING) + state.finish = lambda: setattr(state, 'status', MessageStatus.STOPPED) + seq.set_state(state) + seq.update_token_ids(np.array([1, 2, 3]), mode=UpdateTokenMode.INPUTS) + outputs = SimpleNamespace( + next_token_ids=torch.tensor([99]), + stopped=torch.tensor([True]), + model_metas=[{'turn': 1}], + all_routed_experts=None, + ) + inputs = SimpleNamespace(seq_length=torch.tensor([3]), + is_decoding=False) + + strategy.update_running([seq], outputs, inputs, delta=None) + + assert seq.status == MessageStatus.STOPPED + assert seq.num_history_ids == 3 + assert seq.num_token_ids == 0 + assert seq.generated_ids.tolist() == [] + assert seq.model_meta == {'turn': 1} diff --git a/tests/pytorch/engine/test_ray_mp_engine.py b/tests/pytorch/engine/test_ray_mp_engine.py index 652dd3d368..baa1b89ffe 100644 --- a/tests/pytorch/engine/test_ray_mp_engine.py +++ b/tests/pytorch/engine/test_ray_mp_engine.py @@ -1,6 +1,8 @@ import asyncio from contextlib import suppress +import pytest + from lmdeploy.messages import EngineOutput, ResponseType from lmdeploy.pytorch.engine.mp_engine.base_worker import EngineOutputGather from lmdeploy.pytorch.engine.mp_engine.ray_engine import RayEngineWorker, RayMPEngine @@ -88,14 +90,15 @@ async def _async_test_ray_get_stream_task_result_after_drop_is_idempotent(): stream_id = 123 event = asyncio.Event() - stream_out = [event, None] + stream_out = [event, None, False] worker._stream_aiter[stream_id] = stream_out get_task = asyncio.create_task(worker.get_stream_task_result(stream_id)) await asyncio.sleep(0) worker._stream_aiter.pop(stream_id) - stream_out[1] = (EngineOutput(ResponseType.FINISH, [7]), True) + stream_out[1] = EngineOutput(ResponseType.FINISH, [7]) + stream_out[2] = True event.set() result, stopped = await asyncio.wait_for(get_task, timeout=1) @@ -106,3 +109,69 @@ async def _async_test_ray_get_stream_task_result_after_drop_is_idempotent(): def test_ray_get_stream_task_result_after_drop_is_idempotent(): asyncio.run(_async_test_ray_get_stream_task_result_after_drop_is_idempotent()) + + +async def _async_test_ray_completion_wakeup_does_not_replay_output(): + worker = RayEngineWorker.__new__(RayEngineWorker) + worker._stream_id = 0 + worker._stream_aiter = {} + worker._stream_task = {} + worker._engine_output_gather = EngineOutputGather() + allow_finish = asyncio.Event() + + async def stream(notify_add_msg_func=None): + notify_add_msg_func() + yield EngineOutput(ResponseType.FINISH, [], logprobs=[{2: -0.5}]) + await allow_finish.wait() + + worker.instance_async_stream_infer = stream + stream_id = await worker.create_stream_task( + 'instance_async_stream_infer') + + first, stopped = await worker.get_stream_task_result(stream_id) + assert stopped is False + assert first.logprobs == [{2: -0.5}] + + allow_finish.set() + second, stopped = await worker.get_stream_task_result(stream_id) + assert stopped is True + assert second is None + + +def test_ray_completion_wakeup_does_not_replay_output(): + asyncio.run(_async_test_ray_completion_wakeup_does_not_replay_output()) + + +async def _async_test_ray_completion_keeps_unconsumed_empty_carrier(): + worker = RayEngineWorker.__new__(RayEngineWorker) + worker._stream_id = 0 + worker._stream_aiter = {} + worker._stream_task = {} + worker._engine_output_gather = EngineOutputGather() + + async def stream(notify_add_msg_func=None): + notify_add_msg_func() + yield EngineOutput(ResponseType.FINISH, [], logprobs=[]) + + worker.instance_async_stream_infer = stream + stream_id = await worker.create_stream_task( + 'instance_async_stream_infer') + await worker._stream_task[stream_id] + + result, stopped = await worker.get_stream_task_result(stream_id) + assert stopped is True + assert result.logprobs == [] + assert stream_id not in worker._stream_aiter + + +def test_ray_completion_keeps_unconsumed_empty_carrier(): + asyncio.run(_async_test_ray_completion_keeps_unconsumed_empty_carrier()) + + +@pytest.mark.parametrize('carrier', [None, [], [{2: -0.5}]]) +def test_engine_output_gather_preserves_logprob_carrier_states(carrier): + gather = EngineOutputGather() + final = EngineOutput(ResponseType.FINISH, [], logprobs=carrier) + gather.add(1, final) + assert gather.pop(1, final).logprobs == carrier + assert 1 not in gather._output diff --git a/tests/pytorch/paging/test_block_trie.py b/tests/pytorch/paging/test_block_trie.py index 3ccb9f777b..1a73b49154 100644 --- a/tests/pytorch/paging/test_block_trie.py +++ b/tests/pytorch/paging/test_block_trie.py @@ -119,6 +119,47 @@ def _add_ready_ssm_checkpoint(self, scheduler, token_ids): assert scheduler.block_trie.commit_state_checkpoint_for_seq(seq) return seq, seq.prefix_cache.last_shared_node, state_idx + def test_logprob_prefix_match_is_capped_at_scoring_start(self, scheduler, block_mgr, block_trie): + block_size = scheduler.block_trie.block_size + token_ids = list(range(block_size * 3 + 1)) + + cached = scheduler.add_session(100).add_sequence(token_ids) + block_mgr.allocate(cached) + block_trie.allocate(cached) + cached_blocks = cached.logical_blocks.get_real_blocks().copy() + + param = SamplingParam(num_logprobs=0, logprob_start_len=block_size + 3) + scored = scheduler.add_session(101).add_sequence(token_ids, sampling_param=param) + assert scored.get_prefix_cache_max_candidate_step() == block_size + 3 + assert scored.get_prefix_cache_max_match_step() == block_size + block_trie.match(scored) + assert scored.num_history_ids == block_size + assert scored.prefix_cache.private_recompute_start_step == block_size + assert scored.prefix_cache.private_recompute_end_step == block_size * 3 + + block_mgr.allocate(scored) + private_blocks = scored.logical_blocks.get_real_blocks()[1:3].copy() + assert not np.array_equal(private_blocks, cached_blocks[1:3]) + block_trie.allocate(scored) + assert np.array_equal(scored.logical_blocks.get_real_blocks()[1:3], private_blocks) + + def test_logprob_start_minus_one_preserves_prefix_match_limit(self, scheduler, block_mgr, block_trie): + block_size = scheduler.block_trie.block_size + token_ids = list(range(block_size * 3 + 1)) + + cached = scheduler.add_session(102).add_sequence(token_ids) + block_mgr.allocate(cached) + block_trie.allocate(cached) + + default_seq = scheduler.add_session(103).add_sequence(token_ids) + enabled_generated_only = scheduler.add_session(104).add_sequence( + token_ids, sampling_param=SamplingParam(num_logprobs=0, logprob_start_len=-1)) + assert default_seq.get_prefix_cache_max_match_step() == enabled_generated_only.get_prefix_cache_max_match_step() + + block_trie.match(default_seq) + block_trie.match(enabled_generated_only) + assert default_seq.num_history_ids == enabled_generated_only.num_history_ids == block_size * 3 + def test_allocate(self, block_trie, block_mgr, scheduler): allocator = block_trie.allocator sess = scheduler.add_session(0) @@ -388,6 +429,118 @@ def test_ssm_match_recompute_overlap_extends_from_checkpoint_to_raw_hit(self, ss assert seq.prefix_cache.private_recompute_start_step == -1 assert seq.prefix_cache.last_shared_node.num_matched == block_size * 4 + def test_logprob_ssm_private_recompute_extends_through_raw_hit(self, ssm_scheduler): + block_trie = ssm_scheduler.block_trie + block_mgr = ssm_scheduler.block_manager + sess = ssm_scheduler.add_session(1) + block_size = sess.seq_meta.block_size + + checkpoint_tokens = [1] * block_size + checkpoint_seq = sess.add_sequence(checkpoint_tokens) + block_mgr.allocate(checkpoint_seq) + block_trie.allocate(checkpoint_seq) + state_idx = block_trie.reserve_state_checkpoint_for_seq(checkpoint_seq) + assert state_idx >= 0 + assert block_trie.commit_state_checkpoint_for_seq(checkpoint_seq) + + token_ids = checkpoint_tokens + [2] * block_size + [3] * block_size + [4] * block_size + [5] + cached = sess.add_sequence(token_ids) + block_mgr.allocate(cached) + block_trie.allocate(cached) + cached_blocks = cached.logical_blocks.get_real_blocks().copy() + + param = SamplingParam(num_logprobs=0, logprob_start_len=block_size * 2 + 3) + seq = sess.add_sequence(token_ids, sampling_param=param) + seq.prefix_cache.match_recompute_blocks = 1 + block_trie.match(seq) + + assert seq.num_history_ids == block_size + assert seq.prefix_cache.restore_state == state_idx + assert seq.prefix_cache.private_recompute_start_step == block_size + assert seq.prefix_cache.private_recompute_end_step == block_size * 4 + + block_mgr.allocate(seq) + private_blocks = seq.logical_blocks.get_real_blocks()[1:4].copy() + assert not np.array_equal(private_blocks, cached_blocks[1:4]) + block_trie.allocate(seq) + assert np.array_equal(seq.logical_blocks.get_real_blocks()[1:4], private_blocks) + + def test_logprob_ssm_private_recompute_without_mtp_overlap(self, ssm_scheduler): + block_trie = ssm_scheduler.block_trie + block_mgr = ssm_scheduler.block_manager + sess = ssm_scheduler.add_session(2) + block_size = sess.seq_meta.block_size + + checkpoint_tokens = [1] * block_size + checkpoint_seq = sess.add_sequence(checkpoint_tokens) + block_mgr.allocate(checkpoint_seq) + block_trie.allocate(checkpoint_seq) + state_idx = block_trie.reserve_state_checkpoint_for_seq(checkpoint_seq) + assert state_idx >= 0 + assert block_trie.commit_state_checkpoint_for_seq(checkpoint_seq) + + token_ids = checkpoint_tokens + [2] * block_size + [3] * block_size + [4] * block_size + [5] + cached = sess.add_sequence(token_ids) + block_mgr.allocate(cached) + block_trie.allocate(cached) + cached_blocks = cached.logical_blocks.get_real_blocks().copy() + + param = SamplingParam(num_logprobs=0, logprob_start_len=block_size * 2 + 3) + seq = sess.add_sequence(token_ids, sampling_param=param) + assert seq.prefix_cache.match_recompute_blocks == 0 + block_trie.match(seq) + + assert seq.num_history_ids == block_size + assert seq.prefix_cache.restore_state == state_idx + assert seq.prefix_cache.private_recompute_start_step == block_size + assert seq.prefix_cache.private_recompute_end_step == block_size * 4 + + block_mgr.allocate(seq) + private_blocks = seq.logical_blocks.get_real_blocks()[1:4].copy() + assert not np.array_equal(private_blocks, cached_blocks[1:4]) + block_trie.allocate(seq) + assert np.array_equal(seq.logical_blocks.get_real_blocks()[1:4], private_blocks) + + @pytest.mark.parametrize('use_scoring_cap', [False, True]) + def test_ssm_checkpoint_miss_keeps_raw_match_private(self, ssm_scheduler, use_scoring_cap): + block_trie = ssm_scheduler.block_trie + block_mgr = ssm_scheduler.block_manager + sess = ssm_scheduler.add_session(3 + int(use_scoring_cap)) + block_size = sess.seq_meta.block_size + token_ids = ([1] * block_size + [2] * block_size + [3] * block_size + + [4] * block_size + [5]) + + checkpoint_seq = sess.add_sequence(token_ids[:-1]) + block_mgr.allocate(checkpoint_seq) + block_trie.allocate(checkpoint_seq) + state_idx = block_trie.reserve_state_checkpoint_for_seq(checkpoint_seq) + assert state_idx >= 0 + assert block_trie.commit_state_checkpoint_for_seq(checkpoint_seq) + + cached = sess.add_sequence(token_ids) + block_mgr.allocate(cached) + block_trie.allocate(cached) + cached_blocks = cached.logical_blocks.get_real_blocks().copy() + + if use_scoring_cap: + param = SamplingParam(num_logprobs=0, logprob_start_len=block_size * 2 + 3) + seq = sess.add_sequence(token_ids, sampling_param=param) + else: + seq = sess.add_sequence(token_ids) + seq.prefix_cache.match_recompute_blocks = 1 + block_trie.match(seq) + + assert seq.num_history_ids == 0 + assert seq.prefix_cache.restore_state == -1 + assert seq.prefix_cache.private_recompute_start_step == 0 + assert seq.prefix_cache.private_recompute_end_step == block_size * 4 + + block_mgr.allocate(seq) + private_blocks = seq.logical_blocks.get_real_blocks()[:4].copy() + assert not np.array_equal(private_blocks, cached_blocks[:4]) + block_trie.allocate(seq) + assert np.array_equal(seq.logical_blocks.get_real_blocks()[:4], private_blocks) + def test_match_after_sequence_blocks_are_freed(self, block_trie, block_mgr, scheduler): sess = scheduler.add_session(0) block_size = sess.seq_meta.block_size diff --git a/tests/pytorch/test_model_inputs.py b/tests/pytorch/test_model_inputs.py index b947180eb6..f5ecad51fa 100644 --- a/tests/pytorch/test_model_inputs.py +++ b/tests/pytorch/test_model_inputs.py @@ -48,3 +48,20 @@ def test_step_context_global_is_decoding_uses_dp_global_state(): dp_meta=DPMeta(dp_is_decoding=False), ) assert not step_ctx.global_is_decoding() + + +def test_model_inputs_logprob_metadata_clone_device_and_step_lifecycle(): + inputs = _make_model_inputs(is_decoding=True) + inputs.logits_indices = torch.tensor([0]) + inputs.seq_logit_length = torch.tensor([1]) + + clone = inputs.clone() + assert clone.logits_indices.tolist() == [0] + assert clone.seq_logit_length.tolist() == [1] + moved = clone.to_device('cpu') + assert moved.logits_indices.tolist() == [0] + assert moved.seq_logit_length.device.type == 'cpu' + + stepped = inputs.step(torch.tensor([[4]])) + assert stepped.logits_indices is None + assert stepped.seq_logit_length is None diff --git a/tests/test_lmdeploy/pytorch/engine/test_engine_instance_cleanup.py b/tests/test_lmdeploy/pytorch/engine/test_engine_instance_cleanup.py index f7f9ed7ebf..35b14e3b36 100644 --- a/tests/test_lmdeploy/pytorch/engine/test_engine_instance_cleanup.py +++ b/tests/test_lmdeploy/pytorch/engine/test_engine_instance_cleanup.py @@ -76,6 +76,28 @@ def test_engine_instance_ends_session_after_finish(): assert (RequestType.END_SESSION, {'session_id': 7}) in req_sender.sent_async +def test_engine_instance_forwards_finish_logprob_carrier(): + carrier = [([0.0, -1.0], [11, 12]), ([-2.0, -3.0], [13, 14])] + req_sender = _FakeReqSender([ + _response(ResponseType.FINISH, + data={ + 'token_ids': np.array([], dtype=np.int64), + 'logprobs': carrier, + }), + ]) + instance = EngineInstance.__new__(EngineInstance) + instance.engine = _FakeEngine() + instance.req_sender = req_sender + instance.max_input_len = 1024 + instance._enable_transfer_obj_ref = False + + outputs = asyncio.run(_consume(instance.async_stream_infer(17, [1, 2]))) + + assert outputs[-1].status == ResponseType.FINISH + assert outputs[-1].token_ids == [] + assert outputs[-1].logprobs == carrier + + def test_engine_instance_ends_session_when_generator_is_closed(): req_sender = _FakeReqSender([ _response(ResponseType.SUCCESS, data={'token_ids': np.array([10])}), diff --git a/tests/test_lmdeploy/serve/test_generate_logprobs.py b/tests/test_lmdeploy/serve/test_generate_logprobs.py new file mode 100644 index 0000000000..ab8aae0970 --- /dev/null +++ b/tests/test_lmdeploy/serve/test_generate_logprobs.py @@ -0,0 +1,254 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import asyncio +import importlib +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from lmdeploy.messages import PytorchEngineConfig, TurbomindEngineConfig +from lmdeploy.pytorch.disagg.config import EngineRole +from lmdeploy.serve.openai.api_server import ( + VariableInterface, + _create_token_logprobs, + generate, +) +from lmdeploy.serve.openai.protocol import GenerateReqInput +from lmdeploy.serve.openai.serving_generate import check_request + +api_server_module = importlib.import_module('lmdeploy.serve.openai.api_server') + + +class _SessionManager: + + def has(self, session_id): + return False + + def get(self, *args, **kwargs): + raise AssertionError('empty scoring response must not create a session') + + +class _ServerContext: + + def __init__(self, config, speculative_config=None): + self.config = config + self.async_engine = SimpleNamespace(speculative_config=speculative_config) + + def get_engine_config(self): + return self.config + + def get_session_manager(self): + return _SessionManager() + + +def test_generate_input_logprob_validation(): + config = PytorchEngineConfig(logprobs_mode='raw_logprobs', + role=EngineRole.Hybrid) + context = _ServerContext(config) + valid = GenerateReqInput(input_ids=[1, 2], + return_logprob=True, + logprob_start_len=0, + max_tokens=0) + assert check_request(valid, context) == '' + + assert 'return_logprob=True' in check_request( + GenerateReqInput(input_ids=[1, 2], logprob_start_len=0), context) + assert 'last scorable source position' in check_request( + GenerateReqInput(input_ids=[1, 2], + return_logprob=True, + logprob_start_len=3, + max_tokens=0), context) + assert 'last scorable source position' in check_request( + GenerateReqInput(input_ids=[1, 2], + return_logprob=True, + logprob_start_len=1, + max_tokens=0), context) + assert 'max_tokens=0' in check_request( + GenerateReqInput(input_ids=[1, 2], + return_logprob=True, + logprob_start_len=0), context) + assert 'max_tokens=0' in check_request( + GenerateReqInput(input_ids=[1, 2], + return_logprob=True, + logprob_start_len=0, + max_tokens=5), context) + assert 'positive integer' in check_request( + GenerateReqInput(input_ids=[1, 2], max_tokens=0), context) + 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) == '' + + assert 'not enabled logprobs_mode' in check_request( + valid, _ServerContext(TurbomindEngineConfig())) + assert 'PyTorch hybrid' in check_request( + valid, _ServerContext(SimpleNamespace(logprobs_mode='raw_logprobs'))) + assert 'speculative decoding' in check_request( + valid, _ServerContext(config, speculative_config=SimpleNamespace(method='qwen3_5_mtp'))) + assert 'raw_logits or raw_logprobs' in check_request( + valid, + _ServerContext( + PytorchEngineConfig(logprobs_mode='processed_logprobs', + role=EngineRole.Hybrid))) + with pytest.raises(ValidationError): + GenerateReqInput(input_ids=[1], logprob_start_len=-2) + + +def test_generate_logprob_formatting_preserves_requested_empty(): + assert _create_token_logprobs(None, None) is None + assert _create_token_logprobs([], []) is None + assert _create_token_logprobs( + [4], [{4: -0.75, 8: -2.0}]) == [(-0.75, 4)] + with pytest.raises(ValueError): + _create_token_logprobs([4], []) + + assert _create_token_logprobs( + [5], [{5: -0.25}]) == [(-0.25, 5)] + + +class _RawRequest: + + async def is_disconnected(self): + return False + + +class _NonemptyEngine: + + def __init__(self, expect_image=False): + self.backend_config = PytorchEngineConfig( + logprobs_mode='raw_logprobs', role=EngineRole.Hybrid) + self.session_mgr = _SessionManager() + self.epoch = 1 + self.expect_image = expect_image + + async def generate(self, **kwargs): + if self.expect_image: + assert kwargs['input_ids'] is None + content = kwargs['messages'][0]['content'] + assert content[0] == {'type': 'text', 'text': [1, 2, 3]} + assert content[1] == { + 'type': 'image_url', + 'image_url': { + 'url': 'https://example.com/image.png' + } + } + else: + assert kwargs['input_ids'] == [1, 2, 3] + assert kwargs['gen_config'].max_new_tokens == 0 + assert kwargs['gen_config'].logprob_start_len == 0 + yield SimpleNamespace( + response='', + token_ids=[], + routed_experts=None, + logprobs=[{2: -0.25}, {3: -0.5}], + finish_reason='length', + input_token_len=3, + generate_token_len=0, + ) + + +async def _call_nonempty_route(monkeypatch, stream, image_data=None): + old_engine = VariableInterface.async_engine + VariableInterface.async_engine = _NonemptyEngine( + expect_image=image_data is not None) + + async def passthrough(generator, result_generators, sessions): + async for item in generator: + yield item + + monkeypatch.setattr(api_server_module, '_with_request_cleanup', passthrough) + monkeypatch.setattr(VariableInterface, 'create_session', + classmethod(lambda cls, session_id: SimpleNamespace())) + try: + request = GenerateReqInput(input_ids=[1, 2, 3], + image_data=image_data, + return_logprob=True, + logprob_start_len=0, + max_tokens=0, + stream=stream) + result = await generate(request, raw_request=_RawRequest()) + if stream: + return ''.join([chunk async for chunk in result.body_iterator]) + return result + finally: + VariableInterface.async_engine = old_engine + + +class _DecodeEngine: + + def __init__(self): + self.backend_config = PytorchEngineConfig( + logprobs_mode='raw_logprobs', role=EngineRole.Hybrid) + self.session_mgr = _SessionManager() + self.epoch = 1 + + async def generate(self, **kwargs): + assert kwargs['input_ids'] == [1, 2] + assert kwargs['gen_config'].max_new_tokens == 2 + assert kwargs['gen_config'].logprob_start_len == -1 + assert kwargs['gen_config'].return_routed_experts is True + yield SimpleNamespace( + response='x', + token_ids=[7], + routed_experts=[[[1, 2]]], + logprobs=[{7: -0.125, 8: -2.0}], + finish_reason='length', + input_token_len=2, + generate_token_len=1, + ) + + +def test_generate_keeps_decode_logprobs_and_routed_experts(monkeypatch): + old_engine = VariableInterface.async_engine + old_default = VariableInterface.default_gen_config + VariableInterface.async_engine = _DecodeEngine() + VariableInterface.default_gen_config = {} + monkeypatch.setattr(VariableInterface, 'create_session', + classmethod(lambda cls, session_id: SimpleNamespace())) + + try: + request = GenerateReqInput(input_ids=[1, 2], + max_tokens=2, + return_logprob=True, + return_routed_experts=True) + result = asyncio.run(generate(request, raw_request=_RawRequest())) + finally: + VariableInterface.async_engine = old_engine + VariableInterface.default_gen_config = old_default + + assert result.text == 'x' + assert result.output_ids == [7] + assert result.meta_info.output_token_logprobs == [(-0.125, 7)] + assert result.meta_info.input_token_logprobs is None + assert result.meta_info.routed_experts == [[[1, 2]]] + + +def test_generate_nonempty_input_logprobs_nonstream(monkeypatch): + result = asyncio.run(_call_nonempty_route(monkeypatch, stream=False)) + assert result.text == '' + assert result.output_ids == [] + assert result.meta_info.input_token_logprobs == [(-0.25, 2), (-0.5, 3)] + assert result.meta_info.output_token_logprobs is None + + +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 + + +def test_generate_nonempty_input_logprobs_stream(monkeypatch): + body = asyncio.run(_call_nonempty_route(monkeypatch, stream=True)) + assert body.count('"input_token_logprobs":[[-0.25,2],[-0.5,3]]') == 1 + assert body.count('\ndata: ') == 1 + assert body.endswith('data: [DONE]\n\n') diff --git a/tests/test_lmdeploy/serve/test_session_cleanup.py b/tests/test_lmdeploy/serve/test_session_cleanup.py index 8b8b1b0f3f..7a3d96417c 100644 --- a/tests/test_lmdeploy/serve/test_session_cleanup.py +++ b/tests/test_lmdeploy/serve/test_session_cleanup.py @@ -1,7 +1,16 @@ import asyncio from contextlib import aclosing, suppress +from types import SimpleNamespace -from lmdeploy.messages import GenerationConfig +import pytest + +from lmdeploy.messages import ( + EngineOutput, + GenerationConfig, + RequestMetrics, + ResponseType, +) +from lmdeploy.pytorch.disagg.config import EngineRole from lmdeploy.serve.core.exceptions import SafeRunException from lmdeploy.serve.managers import SessionManager @@ -257,3 +266,89 @@ async def _run_max_new_tokens_zero_cleans_up_session(): def test_max_new_tokens_zero_cleans_up_session(): asyncio.run(_run_max_new_tokens_zero_cleans_up_session()) + + +async def _run_input_logprob_outputs(engine_outputs, + logprob_start_len=0, + queued_metrics=None): + from lmdeploy.metrics.metrics_processor import metrics_processor + from lmdeploy.serve.core.async_engine import AsyncEngine + + queued_metrics = [] if queued_metrics is None else queued_metrics + + class _Handle: + + def async_stream_infer(self, *args, **kwargs): + + async def outputs(): + for output in engine_outputs: + yield output + + return outputs() + + async def async_cancel(self, session_id): + raise AssertionError('input-logprob terminal extraction should not cancel') + + async def async_end(self, session_id): + raise AssertionError('input-logprob terminal extraction should not end early') + + handle = _Handle() + engine = AsyncEngine.__new__(AsyncEngine) + engine.session_mgr = SessionManager() + engine.session_mgr.build_request_handle_pool( + SimpleNamespace(create_instance=lambda: handle), 1) + engine._determine_gen_config = ( + lambda session, input_ids, gen_config=None: gen_config) + engine._if_session_stale = lambda session, input_len: None + engine.backend = 'pytorch' + engine.backend_config = SimpleNamespace( + role=EngineRole.Hybrid, + logprobs_mode='raw_logprobs', + enable_prefix_caching=True, + ) + engine.hf_cfg = SimpleNamespace(vocab_size=10) + engine.num_spec_token = 0 + engine.request_logger = SimpleNamespace( + log_response=lambda *args, **kwargs: None) + + old_queue_update = metrics_processor.queue_update + metrics_processor.queue_update = queued_metrics.append + try: + session = engine.session_mgr.get(260609) + config = GenerationConfig(max_new_tokens=0, + logprobs=0, + logprob_start_len=logprob_start_len) + results = [ + out async for out in engine.generate( + None, + session, + input_ids=[1, 2, 3], + gen_config=config, + sequence_end=True, + ) + ] + return results, queued_metrics, engine.session_mgr + finally: + metrics_processor.queue_update = old_queue_update + + +@pytest.mark.parametrize('carrier', [ + [], + [{2: -0.5}, {3: -0.7}], +]) +def test_async_engine_reinterprets_terminal_logprobs_as_input_rows(carrier): + output = EngineOutput(ResponseType.FINISH, + [], + logprobs=carrier, + req_metrics=RequestMetrics()) + start = 2 if carrier == [] else 0 + results, queued_metrics, manager = asyncio.run( + _run_input_logprob_outputs([output], logprob_start_len=start)) + + assert len(queued_metrics) == 1 + assert manager.sessions == {} + assert len(results) == 1 + assert results[0].token_ids == [] + assert results[0].generate_token_len == 0 + assert results[0].finish_reason == 'length' + assert results[0].logprobs == carrier diff --git a/tests/test_lmdeploy/test_messages.py b/tests/test_lmdeploy/test_messages.py index 36e16e3155..bc62cd9eb4 100644 --- a/tests/test_lmdeploy/test_messages.py +++ b/tests/test_lmdeploy/test_messages.py @@ -3,6 +3,8 @@ from pydantic import ValidationError from lmdeploy import GenerationConfig, Tokenizer +from lmdeploy.messages import Response +from lmdeploy.pytorch.messages import SamplingParam from lmdeploy.serve.openai.protocol import ChatCompletionRequest from lmdeploy.utils import get_hf_gen_cfg @@ -13,6 +15,33 @@ def test_generation_config_repetition_ngram_clamped(): assert c.repetition_ngram_threshold == 0 +def test_input_logprobs_contract_and_sampling_conversion(): + assert GenerationConfig().logprob_start_len == -1 + + for start in (-1, 0, 3): + config = GenerationConfig(logprobs=0, logprob_start_len=start) + param = SamplingParam.from_gen_config(config) + assert param.num_logprobs == 0 + assert param.logprob_start_len == start + + # Preserve upstream generated-only validation behavior. + GenerationConfig(logprobs=-1) + with pytest.raises(ValueError, match='logprobs must be non-negative'): + GenerationConfig(logprobs=-1, logprob_start_len=0) + with pytest.raises(ValueError, match='greater than or equal to -1'): + GenerationConfig(logprobs=0, logprob_start_len=-2) + with pytest.raises(ValueError, match='logprobs must be non-negative'): + GenerationConfig(logprob_start_len=0) + + +@pytest.mark.parametrize('carrier', [[], [{2: -0.5}]]) +def test_response_extend_preserves_logprob_carrier(carrier): + response = Response('', 0, 3) + assert not hasattr(response, 'input_logprobs') + response.extend(Response('', 0, 3, logprobs=carrier)) + assert response.logprobs == carrier + + def test_chat_completion_request_repetition_ngram_ge_zero(): with pytest.raises(ValidationError): ChatCompletionRequest( diff --git a/tests/test_lmdeploy/test_metrics_loggers.py b/tests/test_lmdeploy/test_metrics_loggers.py index 69ba4db4e5..1e63475edd 100644 --- a/tests/test_lmdeploy/test_metrics_loggers.py +++ b/tests/test_lmdeploy/test_metrics_loggers.py @@ -2,8 +2,9 @@ import pytest +from lmdeploy.messages import EngineOutput, RequestMetrics, ResponseType from lmdeploy.metrics.loggers import PrometheusStatLogger -from lmdeploy.metrics.stats import SpeculativeDecodingStats +from lmdeploy.metrics.stats import IterationStats, RequestStats, SpeculativeDecodingStats prometheus_client = pytest.importorskip('prometheus_client') @@ -42,3 +43,22 @@ def test_prometheus_stat_logger_records_specdecode_metrics(): position_labels = labels | {'position': '2'} assert _get_sample_value('lmdeploy:spec_decode_num_accepted_tokens_per_pos_total', position_labels) == 0 assert _get_sample_value('lmdeploy:spec_decode_per_position_accept_rate', position_labels) == 0 + + +def test_zero_token_finish_records_request_reason(): + req_stats = RequestStats(prompt_tokens=4) + iteration_stats = IterationStats() + req_stats.scheduled_time = iteration_stats.iteration_timestamp - 2 + output = EngineOutput(ResponseType.FINISH, [], req_metrics=RequestMetrics()) + + iteration_stats.update_from_output(output, req_stats) + + assert iteration_stats.new_generation_tokens == 0 + assert req_stats.generation_tokens == 0 + assert req_stats.finish_reason == ResponseType.FINISH + assert req_stats.finish_time == iteration_stats.iteration_timestamp + assert req_stats.first_token_time == req_stats.finish_time + assert req_stats.prefill_time_interval == pytest.approx(2) + assert req_stats.decode_time_interval == 0 + assert iteration_stats.prompt_tokens == 4 + assert iteration_stats.ttft is None From 3a88511df10e011af081a269c53b727827aa474b Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Tue, 28 Jul 2026 06:36:42 +0000 Subject: [PATCH 2/4] Fix TurboMind generate logprob validation --- lmdeploy/serve/openai/serving_generate.py | 2 +- tests/test_lmdeploy/serve/test_generate_logprobs.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/lmdeploy/serve/openai/serving_generate.py b/lmdeploy/serve/openai/serving_generate.py index 2388d0c199..5fb39e6a60 100644 --- a/lmdeploy/serve/openai/serving_generate.py +++ b/lmdeploy/serve/openai/serving_generate.py @@ -15,7 +15,7 @@ def check_request(request: GenerateReqInput, server_context: 'VariableInterface' session_manager = server_context.get_session_manager() logprobs_mode = getattr(engine_config, 'logprobs_mode', None) return_logprob = request.return_logprob - if logprobs_mode is None and return_logprob: + if hasattr(engine_config, 'logprobs_mode') and logprobs_mode is None and return_logprob: return f'return_logprob({return_logprob}) requested but not enabled logprobs_mode in engine configuration.' if (request.prompt is not None) ^ (request.input_ids is None): diff --git a/tests/test_lmdeploy/serve/test_generate_logprobs.py b/tests/test_lmdeploy/serve/test_generate_logprobs.py index ab8aae0970..6b6042fcfe 100644 --- a/tests/test_lmdeploy/serve/test_generate_logprobs.py +++ b/tests/test_lmdeploy/serve/test_generate_logprobs.py @@ -86,7 +86,12 @@ def test_generate_input_logprob_validation(): logprob_start_len=0, max_tokens=0), context) == '' - assert 'not enabled logprobs_mode' in check_request( + assert check_request( + GenerateReqInput(prompt='Paris is the capital of', + max_tokens=2, + return_logprob=True), + _ServerContext(TurbomindEngineConfig())) == '' + assert 'PyTorch hybrid' in check_request( valid, _ServerContext(TurbomindEngineConfig())) assert 'PyTorch hybrid' in check_request( valid, _ServerContext(SimpleNamespace(logprobs_mode='raw_logprobs'))) From 699517f6f0f4642b9b1461e14341884fd0b1c3a4 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Tue, 28 Jul 2026 06:41:19 +0000 Subject: [PATCH 3/4] Cover PyTorch generate logprob validation --- tests/test_lmdeploy/serve/test_generate_logprobs.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_lmdeploy/serve/test_generate_logprobs.py b/tests/test_lmdeploy/serve/test_generate_logprobs.py index 6b6042fcfe..f8d9fdebd2 100644 --- a/tests/test_lmdeploy/serve/test_generate_logprobs.py +++ b/tests/test_lmdeploy/serve/test_generate_logprobs.py @@ -86,6 +86,11 @@ def test_generate_input_logprob_validation(): logprob_start_len=0, max_tokens=0), context) == '' + assert 'not enabled logprobs_mode' in check_request( + GenerateReqInput(prompt='Paris is the capital of', + max_tokens=2, + return_logprob=True), + _ServerContext(PytorchEngineConfig(role=EngineRole.Hybrid))) assert check_request( GenerateReqInput(prompt='Paris is the capital of', max_tokens=2, From ce8ca4cc9d8e35c761f4c2c02c5d824c5a4d28c9 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Thu, 30 Jul 2026 06:28:22 +0000 Subject: [PATCH 4/4] support topk logprobs --- lmdeploy/serve/openai/api_server.py | 51 ++++++++++++-- lmdeploy/serve/openai/protocol.py | 3 + lmdeploy/serve/openai/serving_generate.py | 5 ++ .../serve/test_generate_logprobs.py | 68 +++++++++++++++++-- 4 files changed, 117 insertions(+), 10 deletions(-) diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py index 2675dd9857..7fd965584e 100644 --- a/lmdeploy/serve/openai/api_server.py +++ b/lmdeploy/serve/openai/api_server.py @@ -288,6 +288,24 @@ def _create_token_logprobs(token_ids: list[int] | None = None, return token_logprobs or None +def _create_top_logprobs(token_ids: list[int] | None = None, + logprobs: list[dict[int, float]] | None = None, + top_logprobs_num: int = 0): + """Create raw top-logprob pairs.""" + if token_ids is None or logprobs is None or top_logprobs_num <= 0: + return None + + top_logprobs = [] + for tok, tok_logprobs in zip(token_ids, logprobs, strict=True): + items = [(prob, top_id) for top_id, prob in tok_logprobs.items()] + if len(items) > top_logprobs_num: + # Drop the extra selected/target row; expose model top-k only. + items = [(prob, top_id) for prob, top_id in items if top_id != tok] + items.sort(key=lambda x: x[0], reverse=True) + top_logprobs.append(items[:top_logprobs_num]) + return top_logprobs or None + + @router.get('/health') async def health() -> JSONResponse: """Health check.""" @@ -1006,9 +1024,10 @@ async def generate(request: GenerateReqInput, raw_request: Request = None): prompt = [dict(role='user', content=[text_input] + image_input)] input_ids = None + gen_logprobs = (request.top_logprobs_num or 1) if request.return_logprob else None gen_config = _build_serving_generation_config( request, - logprobs=1 if request.return_logprob else None, + logprobs=gen_logprobs, logprob_start_len=request.logprob_start_len, stop_words=request.stop, ) @@ -1028,6 +1047,8 @@ def create_generate_response_json(res, output_ids, logprobs, input_logprobs, + top_logprobs, + input_top_logprobs, finish_reason, routed_experts=None): # only output router experts in last chunk @@ -1035,6 +1056,8 @@ def create_generate_response_json(res, meta = GenerateReqMetaOutput(finish_reason=dict(type=finish_reason) if finish_reason else None, output_token_logprobs=logprobs, input_token_logprobs=input_logprobs, + output_top_logprobs=top_logprobs, + input_top_logprobs=input_top_logprobs, prompt_tokens=res.input_token_len, routed_experts=routed_experts, completion_tokens=res.generate_token_len) @@ -1050,14 +1073,23 @@ async def generate_stream_generator(): if input_logprobs_requested: logprobs = None input_logprobs = _create_token_logprobs(input_logprob_token_ids, res.logprobs) + top_logprobs = None + input_top_logprobs = (_create_top_logprobs(input_logprob_token_ids, res.logprobs, + request.top_logprobs_num) + if request.top_logprobs_num else None) else: logprobs = _create_token_logprobs(res.token_ids, res.logprobs) + top_logprobs = (_create_top_logprobs(res.token_ids, res.logprobs, request.top_logprobs_num) + if request.top_logprobs_num else None) input_logprobs = None + input_top_logprobs = None response_json = create_generate_response_json(res, text, output_ids, logprobs, input_logprobs, + top_logprobs, + input_top_logprobs, res.finish_reason, routed_experts=routed_experts) yield f'data: {response_json}\n\n' @@ -1073,7 +1105,9 @@ async def _inner_call(): text = '' output_ids = [] logprobs = [] + input_logprobs = [] input_token_logprobs = None + input_top_logprobs = None async with aclosing(_with_request_cleanup(result_generator, [result_generator], [session])) as generator: async for res in generator: if await raw_request.is_disconnected(): @@ -1083,19 +1117,26 @@ async def _inner_call(): text += res.response or '' output_ids.extend(res.token_ids or []) if input_logprobs_requested: - cur_input_logprobs = _create_token_logprobs(input_logprob_token_ids, - res.logprobs) - if cur_input_logprobs is not None: - input_token_logprobs = cur_input_logprobs + if res.logprobs: + input_logprobs = res.logprobs elif res.logprobs: logprobs.extend(res.logprobs) output_token_logprobs = _create_token_logprobs(output_ids, logprobs) if logprobs else None + output_top_logprobs = (_create_top_logprobs(output_ids, logprobs, request.top_logprobs_num) + if request.top_logprobs_num and logprobs else None) + if input_logprobs_requested and input_logprobs: + input_token_logprobs = _create_token_logprobs(input_logprob_token_ids, input_logprobs) + input_top_logprobs = (_create_top_logprobs(input_logprob_token_ids, input_logprobs, + request.top_logprobs_num) + if request.top_logprobs_num else None) nonlocal response meta = GenerateReqMetaOutput(finish_reason=dict(type=res.finish_reason) if res.finish_reason else None, output_token_logprobs=output_token_logprobs or None, input_token_logprobs=input_token_logprobs, + output_top_logprobs=output_top_logprobs, + input_top_logprobs=input_top_logprobs, prompt_tokens=res.input_token_len, routed_experts=res.routed_experts, completion_tokens=res.generate_token_len) diff --git a/lmdeploy/serve/openai/protocol.py b/lmdeploy/serve/openai/protocol.py index 420e74a719..be471263b3 100644 --- a/lmdeploy/serve/openai/protocol.py +++ b/lmdeploy/serve/openai/protocol.py @@ -545,6 +545,7 @@ class GenerateReqInput(BaseModel): input_ids: list[int] | None = None image_data: ImageDataFormat | None = None return_logprob: bool | None = None + top_logprobs_num: int | None = None logprob_start_len: int = Field(default=-1, ge=-1) max_tokens: int = 128 stop: str | list[str] | None = None @@ -580,6 +581,8 @@ class GenerateReqMetaOutput(BaseModel): finish_reason: dict[str, Any] | None = None output_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) input_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) + output_top_logprobs: list[list[tuple[float, int]]] | None = None # per-output-token top (logprob, token_id) + input_top_logprobs: list[list[tuple[float, int]]] | None = None # per-input-token top (logprob, token_id) routed_experts: list[list[list[int]]] | str | None = None # (num_token, num_layer, topk_expert) diff --git a/lmdeploy/serve/openai/serving_generate.py b/lmdeploy/serve/openai/serving_generate.py index 5fb39e6a60..3fdbbbb9a0 100644 --- a/lmdeploy/serve/openai/serving_generate.py +++ b/lmdeploy/serve/openai/serving_generate.py @@ -15,6 +15,11 @@ def check_request(request: GenerateReqInput, server_context: 'VariableInterface' session_manager = server_context.get_session_manager() logprobs_mode = getattr(engine_config, 'logprobs_mode', None) return_logprob = request.return_logprob + top_logprobs_num = request.top_logprobs_num or 0 + if top_logprobs_num < 0: + return f'The top_logprobs_num {request.top_logprobs_num!r} cannot be a negative integer.' + if top_logprobs_num > 0 and return_logprob is not True: + return 'top_logprobs_num requires return_logprob=True.' if hasattr(engine_config, 'logprobs_mode') and logprobs_mode is None and return_logprob: return f'return_logprob({return_logprob}) requested but not enabled logprobs_mode in engine configuration.' diff --git a/tests/test_lmdeploy/serve/test_generate_logprobs.py b/tests/test_lmdeploy/serve/test_generate_logprobs.py index f8d9fdebd2..aeae3c29f3 100644 --- a/tests/test_lmdeploy/serve/test_generate_logprobs.py +++ b/tests/test_lmdeploy/serve/test_generate_logprobs.py @@ -11,6 +11,7 @@ from lmdeploy.serve.openai.api_server import ( VariableInterface, _create_token_logprobs, + _create_top_logprobs, generate, ) from lmdeploy.serve.openai.protocol import GenerateReqInput @@ -74,6 +75,12 @@ def test_generate_input_logprob_validation(): max_tokens=5), context) assert 'positive integer' in check_request( GenerateReqInput(input_ids=[1, 2], max_tokens=0), context) + assert 'top_logprobs_num requires return_logprob=True' in check_request( + GenerateReqInput(input_ids=[1, 2], top_logprobs_num=2), context) + assert 'negative integer' in check_request( + GenerateReqInput(input_ids=[1, 2], + return_logprob=True, + top_logprobs_num=-1), context) assert check_request( GenerateReqInput(prompt='hi', return_logprob=True, @@ -121,6 +128,16 @@ def test_generate_logprob_formatting_preserves_requested_empty(): assert _create_token_logprobs( [5], [{5: -0.25}]) == [(-0.25, 5)] + assert _create_top_logprobs(None, None, 2) is None + assert _create_top_logprobs([], [], 2) is None + assert _create_top_logprobs([4], [{4: -0.75}], 0) is None + assert _create_top_logprobs( + [4], [{4: -0.75, 8: -2.0, 9: -3.0}], 3) == [[(-0.75, 4), (-2.0, 8), (-3.0, 9)]] + assert _create_top_logprobs( + [4], [{4: -0.75, 8: -2.0, 9: -3.0, 10: -4.0}], 3) == [[(-2.0, 8), (-3.0, 9), + (-4.0, 10)]] + with pytest.raises(ValueError): + _create_top_logprobs([4], [], 2) class _RawRequest: @@ -131,12 +148,13 @@ async def is_disconnected(self): class _NonemptyEngine: - def __init__(self, expect_image=False): + def __init__(self, expect_image=False, expected_logprobs=1): self.backend_config = PytorchEngineConfig( logprobs_mode='raw_logprobs', role=EngineRole.Hybrid) self.session_mgr = _SessionManager() self.epoch = 1 self.expect_image = expect_image + self.expected_logprobs = expected_logprobs async def generate(self, **kwargs): if self.expect_image: @@ -153,21 +171,23 @@ async def generate(self, **kwargs): assert kwargs['input_ids'] == [1, 2, 3] assert kwargs['gen_config'].max_new_tokens == 0 assert kwargs['gen_config'].logprob_start_len == 0 + assert kwargs['gen_config'].logprobs == self.expected_logprobs yield SimpleNamespace( response='', token_ids=[], routed_experts=None, - logprobs=[{2: -0.25}, {3: -0.5}], + logprobs=[{2: -0.25, 8: -2.0}, {3: -0.5, 9: -3.0}], finish_reason='length', input_token_len=3, generate_token_len=0, ) -async def _call_nonempty_route(monkeypatch, stream, image_data=None): +async def _call_nonempty_route(monkeypatch, stream, image_data=None, top_logprobs_num=None): old_engine = VariableInterface.async_engine VariableInterface.async_engine = _NonemptyEngine( - expect_image=image_data is not None) + expect_image=image_data is not None, + expected_logprobs=top_logprobs_num or 1) async def passthrough(generator, result_generators, sessions): async for item in generator: @@ -182,6 +202,7 @@ async def passthrough(generator, result_generators, sessions): return_logprob=True, logprob_start_len=0, max_tokens=0, + top_logprobs_num=top_logprobs_num, stream=stream) result = await generate(request, raw_request=_RawRequest()) if stream: @@ -193,16 +214,18 @@ async def passthrough(generator, result_generators, sessions): class _DecodeEngine: - def __init__(self): + def __init__(self, expected_logprobs=1): self.backend_config = PytorchEngineConfig( logprobs_mode='raw_logprobs', role=EngineRole.Hybrid) self.session_mgr = _SessionManager() self.epoch = 1 + self.expected_logprobs = expected_logprobs async def generate(self, **kwargs): assert kwargs['input_ids'] == [1, 2] assert kwargs['gen_config'].max_new_tokens == 2 assert kwargs['gen_config'].logprob_start_len == -1 + assert kwargs['gen_config'].logprobs == self.expected_logprobs assert kwargs['gen_config'].return_routed_experts is True yield SimpleNamespace( response='x', @@ -236,10 +259,36 @@ def test_generate_keeps_decode_logprobs_and_routed_experts(monkeypatch): assert result.text == 'x' assert result.output_ids == [7] assert result.meta_info.output_token_logprobs == [(-0.125, 7)] + assert result.meta_info.output_top_logprobs is None assert result.meta_info.input_token_logprobs is None assert result.meta_info.routed_experts == [[[1, 2]]] +def test_generate_keeps_decode_top_logprobs(monkeypatch): + old_engine = VariableInterface.async_engine + old_default = VariableInterface.default_gen_config + VariableInterface.async_engine = _DecodeEngine(expected_logprobs=2) + VariableInterface.default_gen_config = {} + monkeypatch.setattr(VariableInterface, 'create_session', + classmethod(lambda cls, session_id: SimpleNamespace())) + + try: + request = GenerateReqInput(input_ids=[1, 2], + max_tokens=2, + return_logprob=True, + top_logprobs_num=2, + return_routed_experts=True) + result = asyncio.run(generate(request, raw_request=_RawRequest())) + finally: + VariableInterface.async_engine = old_engine + VariableInterface.default_gen_config = old_default + + assert result.meta_info.output_token_logprobs == [(-0.125, 7)] + assert result.meta_info.output_top_logprobs == [[(-0.125, 7), (-2.0, 8)]] + assert result.meta_info.input_token_logprobs is None + assert result.meta_info.input_top_logprobs is None + + def test_generate_nonempty_input_logprobs_nonstream(monkeypatch): result = asyncio.run(_call_nonempty_route(monkeypatch, stream=False)) assert result.text == '' @@ -257,6 +306,15 @@ def test_generate_input_logprobs_allow_image_with_input_ids(monkeypatch): assert result.meta_info.output_token_logprobs is None +def test_generate_nonempty_input_top_logprobs_nonstream(monkeypatch): + result = asyncio.run( + _call_nonempty_route(monkeypatch, stream=False, top_logprobs_num=2)) + assert result.meta_info.input_token_logprobs == [(-0.25, 2), (-0.5, 3)] + assert result.meta_info.input_top_logprobs == [[(-0.25, 2), (-2.0, 8)], [(-0.5, 3), (-3.0, 9)]] + assert result.meta_info.output_token_logprobs is None + assert result.meta_info.output_top_logprobs is None + + def test_generate_nonempty_input_logprobs_stream(monkeypatch): body = asyncio.run(_call_nonempty_route(monkeypatch, stream=True)) assert body.count('"input_token_logprobs":[[-0.25,2],[-0.5,3]]') == 1