diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c..6dd62a5 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -77,6 +77,11 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( ) +# Ceiling for the token-budget-derived cudagraph capture size. Capture costs +# time and memory per graph, so the budget is only honoured this far. +_MIXED_STEP_CUDAGRAPH_CAPTURE_LIMIT = 1024 + + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -1690,6 +1695,20 @@ class VllmConfig: max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) + # That bound counts one query position per running sequence, + # which is right for a pure-decode step. Under chunked prefill + # a step also carries prompt tokens, so its size is bounded by + # the token budget instead -- max_num_seqs 32 stops capture at + # 64 while real mixed steps run into the hundreds of tokens. + # Those steps exceed every captured graph and fall back to the + # eager path, paying full Python dispatch on every layer. + max_cudagraph_capture_size = max( + max_cudagraph_capture_size, + min( + self.scheduler_config.max_num_batched_tokens, + _MIXED_STEP_CUDAGRAPH_CAPTURE_LIMIT, + ), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f314..6508594 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1718,6 +1718,10 @@ class EngineArgs: ) self.speculative_config[key] = value + if self.speculative_config is None: + self.speculative_config = _checkpoint_mtp_speculative_config( + target_model_config + ) if self.speculative_config is None: return None @@ -2692,3 +2696,38 @@ def _raise_unsupported_error(feature_name: str): f"remove {feature_name} from your config." ) raise NotImplementedError(msg) + + +# Checkpoints that bundle a multi-token-prediction drafter in the same weights, +# so speculative decoding needs no second model and no extra download. +_CHECKPOINT_MTP_MODEL_TYPES: frozenset[str] = frozenset({"qwen3_5", "qwen3_5_moe"}) + +# How many tokens to draft per step. Request latency here is set by the number +# of sequential decode steps, not by the cost of one step: a short reply spends +# ~40 steps that each re-read the whole model. Drafting deeply collapses that +# count, and a rejected draft only costs the drafter's own (much smaller) pass. +# Shallow drafting is the worst of both worlds -- it pays the drafter on every +# step while barely reducing the step count. +_CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS: int = 5 + + +def _checkpoint_mtp_speculative_config(model_config: "ModelConfig") -> dict | None: + """Speculative config built from the checkpoint's own MTP head, or None. + + Only fills in a default: any explicit --speculative-config / --spec-* wins, + because create_speculative_config only calls this when nothing was set. + """ + try: + model_type = getattr(getattr(model_config, "hf_config", None), "model_type", None) + if model_type not in _CHECKPOINT_MTP_MODEL_TYPES: + return None + text_config = getattr(model_config, "hf_text_config", None) + n_mtp_layers = getattr(text_config, "mtp_num_hidden_layers", None) + if not isinstance(n_mtp_layers, int) or n_mtp_layers < 1: + return None + except Exception: + return None + return { + "method": "mtp", + "num_speculative_tokens": _CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS, + } diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef1741..c5214a4 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -395,46 +395,73 @@ 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). + # A RequestOutput can carry several new tokens at once: + # speculative decoding accepts a whole step's drafts, and + # the frontend coalesces outputs whenever it lags the + # engine. Emitting that as a single chunk hides token + # granularity from streaming clients and makes per-token + # arrival times unrecoverable, so split the delta into one + # chunk per token. Text, token accounting and finish_reason + # placement are all preserved. + delta_ids = as_list(output.token_ids) if ( - not include_usage - and self.system_fingerprint is not None - and finish_reason is not None + len(delta_ids) > 1 + and logprobs is None + and not (request.echo and prompt_token_ids_to_return) ): - 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, + pieces = _split_delta_by_token( + tokenizer, delta_ids, delta_text ) + else: + pieces = [(delta_text, delta_ids)] + + tokens_before = previous_num_tokens[i] - len(delta_ids) + last_piece = len(pieces) - 1 + for piece_idx, (piece_text, piece_ids) in enumerate(pieces): + is_last = piece_idx == last_piece + tokens_before += len(piece_ids) + chunk = CompletionStreamResponse( + id=request_id, + object="text_completion", + created=created_time, + model=model_name, + choices=[ + CompletionResponseStreamChoice( + index=i, + text=piece_text, + logprobs=logprobs if is_last else None, + finish_reason=finish_reason if is_last else None, + stop_reason=stop_reason if is_last else None, + prompt_token_ids=( + prompt_token_ids_to_return + if piece_idx == 0 + else None + ), + token_ids=( + piece_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 ( + is_last + and not include_usage + and self.system_fingerprint is not None + and finish_reason is not None + ): + chunk.system_fingerprint = self.system_fingerprint + if include_continuous_usage: + prompt_tokens = num_prompt_tokens[prompt_idx] + chunk.usage = UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=tokens_before, + total_tokens=prompt_tokens + tokens_before, + ) - 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) @@ -689,3 +716,39 @@ class OpenAIServingCompletion(OpenAIServing): tokens=out_tokens, top_logprobs=out_top_logprobs, ) + + + +def _split_delta_by_token( + tokenizer, + token_ids: list[int], + text: str, +) -> list[tuple[str, list[int]]]: + """Slice a multi-token delta into per-token ``(text, [token_id])`` pieces. + + Incremental detokenization only defines text for a whole group, so the + boundaries are recovered by decoding growing prefixes of the delta. They + are clamped into range and forced non-decreasing, and the final piece takes + whatever is left, so the pieces always concatenate back to exactly ``text`` + even when a boundary cannot be placed (multi-byte characters split across + tokens, unusual tokenizers). + """ + n = len(token_ids) + if tokenizer is None or n <= 1: + return [(text, list(token_ids))] + + pieces: list[tuple[str, list[int]]] = [] + cut = 0 + try: + for idx in range(1, n): + decoded = tokenizer.decode(token_ids[:idx], skip_special_tokens=False) + end = min(max(len(decoded), cut), len(text)) + pieces.append((text[cut:end], [token_ids[idx - 1]])) + cut = end + except Exception: + # Any tokenizer that cannot decode a bare prefix falls back to the + # original single-chunk behaviour rather than losing text. + return [(text, list(token_ids))] + + pieces.append((text[cut:], [token_ids[-1]])) + return pieces diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4ac8d49..f4d8ce0 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -321,8 +321,10 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] ] = { PlatformEnum.CUDA: [ - FlashInferFp8DeepGEMMDynamicBlockScaledKernel, + # DeepGEMM first: it covers every batch size with one kernel and needs + # no nvcc JIT, which a cold-started engine would pay for in full. DeepGemmFp8BlockScaledMMKernel, + FlashInferFp8DeepGEMMDynamicBlockScaledKernel, CutlassFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel,