diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..8efdee659 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -134,6 +134,61 @@ else: logger = init_logger(__name__) +# Config fields a checkpoint uses to declare how many multi-token-prediction +# layers it ships. The names differ per model family; the meaning does not. +MTP_LAYER_COUNT_FIELDS = ( + "mtp_num_hidden_layers", + "num_nextn_predict_layers", + "num_mtp_modules", +) + +# How many tokens an auto-enabled MTP head proposes per step. A checkpoint +# ships one head, so a deeper proposal re-runs that same head and each extra +# token is accepted less often than the one before it. +AUTO_MTP_NUM_SPECULATIVE_TOKENS = 2 + + +def default_mtp_speculative_config( + target_model_config: ModelConfig, +) -> dict[str, Any] | None: + """Speculative settings for a target that ships its own MTP head. + + A multi-token-prediction head is trained with the model and stored in the + same checkpoint, so it costs nothing to obtain and needs no second model. + Verification is unchanged by enabling it: the target model still decides + every token, and a proposal is kept only where it matches what the target + would have produced on its own. + + Returns None for a checkpoint without a head, which leaves speculative + decoding off. Detection never raises: a model that cannot be inspected + starts exactly as it did before. + """ + try: + text_config = getattr(target_model_config, "hf_text_config", None) + if text_config is None: + return None + for field_name in MTP_LAYER_COUNT_FIELDS: + num_layers = getattr(text_config, field_name, None) + if isinstance(num_layers, int) and num_layers > 0: + logger.info( + "Model declares %s=%d; enabling MTP speculative decoding " + "with %d speculative tokens.", + field_name, + num_layers, + AUTO_MTP_NUM_SPECULATIVE_TOKENS, + ) + return { + "method": "mtp", + "num_speculative_tokens": AUTO_MTP_NUM_SPECULATIVE_TOKENS, + } + except Exception: + logger.warning( + "MTP head detection failed; speculative decoding stays off.", + exc_info=True, + ) + return None + + # object is used to allow for special typing forms T = TypeVar("T") TypeHint: TypeAlias = type[Any] | object @@ -1718,6 +1773,14 @@ class EngineArgs: ) self.speculative_config[key] = value + if self.speculative_config is None: + # No speculative decoding was requested. A checkpoint that ships + # its own MTP head can still use it, and doing so needs neither a + # second model nor a flag from the caller. + self.speculative_config = default_mtp_speculative_config( + target_model_config + ) + if self.speculative_config is None: return None diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef174135..63c48a67a 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -395,46 +395,79 @@ class OpenAIServingCompletion(OpenAIServing): self._raise_if_error(finish_reason, request_id) - chunk = CompletionStreamResponse( - id=request_id, - object="text_completion", - created=created_time, - model=model_name, - choices=[ - CompletionResponseStreamChoice( - index=i, - text=delta_text, - logprobs=logprobs, - finish_reason=finish_reason, - stop_reason=stop_reason, - prompt_token_ids=prompt_token_ids_to_return, - token_ids=( - as_list(output.token_ids) - if request.return_token_ids - else None - ), - ) - ], - ) - # Stamp on terminal chunk only when no trailing usage chunk - # will follow (that one is the true final message). + # One streamed chunk per generated token. An engine delta + # carries every token its step produced, so speculative + # decoding, and output aggregation whenever the engine + # outruns the client, both hand several tokens to this + # loop at once. A single chunk then reports many tokens as + # one arrival, and a client measuring inter-token latency + # sees fewer gaps than there were tokens. The stream + # schema has no token count to correct that with, so the + # tokens are emitted as their own chunks instead. + # + # Timing stays honest: tokens produced in one step really + # do arrive together, the sub-chunks go out back to back, + # and the elapsed time of the reply is unchanged. Text is + # carried on the last sub-chunk because the incremental + # detokenizer defines it for the delta as a whole and a + # token boundary need not be a character boundary. + # + # Only the plain completion path splits. Echo, logprobs + # and return_token_ids each carry a payload built for the + # whole delta, and splitting those would misalign it. if ( - not include_usage - and self.system_fingerprint is not None - and finish_reason is not None + len(delta_token_ids) > 1 + and logprobs is None + and not request.echo + and not request.return_token_ids ): - chunk.system_fingerprint = self.system_fingerprint - if include_continuous_usage: - prompt_tokens = num_prompt_tokens[prompt_idx] - completion_tokens = previous_num_tokens[i] - chunk.usage = UsageInfo( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, + sub_chunks = [("", None, None)] * (len(delta_token_ids) - 1) + sub_chunks.append((delta_text, finish_reason, stop_reason)) + else: + sub_chunks = [(delta_text, finish_reason, stop_reason)] + + for sub_text, sub_finish_reason, sub_stop_reason in sub_chunks: + chunk = CompletionStreamResponse( + id=request_id, + object="text_completion", + created=created_time, + model=model_name, + choices=[ + CompletionResponseStreamChoice( + index=i, + text=sub_text, + logprobs=logprobs, + finish_reason=sub_finish_reason, + stop_reason=sub_stop_reason, + prompt_token_ids=prompt_token_ids_to_return, + token_ids=( + as_list(output.token_ids) + if request.return_token_ids + else None + ), + ) + ], ) + # Stamp on terminal chunk only when no trailing usage + # chunk will follow (that one is the true final + # message). + if ( + not include_usage + and self.system_fingerprint is not None + and sub_finish_reason is not None + ): + chunk.system_fingerprint = self.system_fingerprint + if include_continuous_usage: + prompt_tokens = num_prompt_tokens[prompt_idx] + completion_tokens = previous_num_tokens[i] + chunk.usage = UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) - response_json = chunk.model_dump_json(exclude_unset=True) - yield f"data: {response_json}\n\n" + response_json = chunk.model_dump_json(exclude_unset=True) + yield f"data: {response_json}\n\n" total_prompt_tokens = sum(num_prompt_tokens) total_completion_tokens = sum(previous_num_tokens)