diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..dcf35614d 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,29 @@ def _raise_unsupported_error(feature_name: str): f"remove {feature_name} from your config." ) raise NotImplementedError(msg) + + +# Checkpoints whose text config carries a bundled MTP drafter that vLLM can load +# from the same weights (method "mtp", draft = target model). +_CHECKPOINT_MTP_MODEL_TYPES: frozenset[str] = frozenset({"qwen3_5", "qwen3_5_moe"}) +_CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS: int = 5 + + +def _checkpoint_mtp_speculative_config(model_config: ModelConfig) -> dict | None: + """Default speculative config from the checkpoint's own MTP head, or None. + + Only applies when the user gave no speculative config at all. The drafter + weights live in the target checkpoint, so nothing extra is downloaded. + """ + hf_config = getattr(model_config, "hf_config", None) + model_type = getattr(hf_config, "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 + return { + "method": "mtp", + "num_speculative_tokens": _CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS, + } diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 911421029..e1eccb8fc 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -33,6 +33,9 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionStreamResponse, ChatMessage, ) +from vllm.entrypoints.openai.completion.serving import ( + _split_delta_by_token, +) from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, ErrorResponse, @@ -674,21 +677,33 @@ class OpenAIServingChat(OpenAIServing): delta=True, ) - if output.finish_reason is None: - # Send token-by-token response for each request.n - choice_data = ChatCompletionResponseStreamChoice( - index=i, - delta=delta_message, - logprobs=logprobs, - finish_reason=None, - token_ids=( - as_list(output.token_ids) - if request.return_token_ids - else None - ), - ) + delta_ids = as_list(output.token_ids) - # if the model is finished generating + # One choice chunk per generated token. A RequestOutput + # can carry several new tokens (speculative decoding + # accepts a step's drafts at once), which hides token + # granularity from streaming consumers; split content-only + # deltas into per-token chunks, mirroring the legacy + # completions endpoint. Skipped when logprobs are requested + # or when a structured-output parser owns the delta. + if ( + len(delta_ids) > 1 + and logprobs is None + and parser is None + and delta_message is not None + and delta_message.content is not None + ): + pieces = [ + (DeltaMessage(content=piece_text), piece_ids) + for piece_text, piece_ids in _split_delta_by_token( + tokenizer, delta_ids, delta_message.content + ) + ] + else: + pieces = [(delta_message, delta_ids)] + + if output.finish_reason is None: + finish_reason_ = None else: # check for error finish reason and abort streaming # finish_reason='error' indicates a retryable error @@ -705,51 +720,70 @@ class OpenAIServingChat(OpenAIServing): finish_reason_ = ( output.finish_reason if output.finish_reason else "stop" ) - choice_data = ChatCompletionResponseStreamChoice( - index=i, - delta=delta_message, - logprobs=logprobs, - finish_reason=finish_reason_, - stop_reason=output.stop_reason, - token_ids=( - as_list(output.token_ids) - if request.return_token_ids - else None - ), - ) - finish_reason_sent[i] = True + tokens_before = previous_num_tokens[i] - len(delta_ids) + n_pieces = len(pieces) + for piece_idx, (piece_message, piece_ids) in enumerate(pieces): + is_last = piece_idx == n_pieces - 1 + tokens_before += len(piece_ids) + + if output.finish_reason is not None and is_last: + choice_data = ChatCompletionResponseStreamChoice( + index=i, + delta=piece_message, + logprobs=logprobs, + finish_reason=finish_reason_, + stop_reason=output.stop_reason, + token_ids=( + piece_ids if request.return_token_ids else None + ), + ) + else: + choice_data = ChatCompletionResponseStreamChoice( + index=i, + delta=piece_message, + logprobs=logprobs, + finish_reason=None, + token_ids=( + piece_ids if request.return_token_ids else None + ), + ) - choice_data = maybe_filter_parallel_tool_calls(choice_data, request) - chunk = ChatCompletionStreamResponse( - id=request_id, - object=chunk_object_type, - created=created_time, - choices=[choice_data], - model=model_name, - ) - # Stamp the fingerprint on terminal chunks only (those with - # finish_reason set). When ``include_usage`` is on, the - # trailing usage chunk below overrides this as the true - # final message. - if ( - not include_usage - and self.system_fingerprint is not None - and choice_data.finish_reason is not None - ): - chunk.system_fingerprint = self.system_fingerprint - - # handle usage stats if requested & if continuous - if include_continuous_usage: - completion_tokens = previous_num_tokens[i] - chunk.usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=num_prompt_tokens + completion_tokens, + choice_data = maybe_filter_parallel_tool_calls( + choice_data, request + ) + chunk = ChatCompletionStreamResponse( + id=request_id, + object=chunk_object_type, + created=created_time, + choices=[choice_data], + model=model_name, ) + # Stamp the fingerprint on terminal chunks only (those + # with finish_reason set). When ``include_usage`` is + # on, the trailing usage chunk below overrides this as + # the true final message. + if ( + not include_usage + and self.system_fingerprint is not None + and choice_data.finish_reason is not None + ): + chunk.system_fingerprint = self.system_fingerprint - data = chunk.model_dump_json(exclude_unset=True) - yield f"data: {data}\n\n" + # handle usage stats if requested & if continuous + if include_continuous_usage: + completion_tokens = tokens_before + chunk.usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=num_prompt_tokens + completion_tokens, + ) + + data = chunk.model_dump_json(exclude_unset=True) + yield f"data: {data}\n\n" + + if output.finish_reason is not None: + finish_reason_sent[i] = True # once the final token is handled, if stream_options.include_usage # is sent, send the usage diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef174135..661cc4b43 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). + # One choice chunk per generated token. A RequestOutput can + # carry several new tokens (speculative decoding accepts a + # step's drafts at once; the frontend merges outputs when it + # lags the engine). Emitting them as one chunk hides the + # token granularity from streaming consumers, so split the + # delta into per-token slices. Text concatenation, token + # accounting and finish_reason placement are unchanged. + 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) + n_pieces = len(pieces) + for piece_idx, (piece_text, piece_ids) in enumerate(pieces): + is_last = piece_idx == n_pieces - 1 + 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, + 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] + completion_tokens = tokens_before + 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) @@ -689,3 +716,39 @@ class OpenAIServingCompletion(OpenAIServing): tokens=out_tokens, top_logprobs=out_top_logprobs, ) + + +def _split_delta_by_token( + tokenizer: TokenizerLike | None, + token_ids: list[int], + text: str, +) -> list[tuple[str, list[int]]]: + """Slice a multi-token delta into per-token (text, [token_id]) pieces. + + Boundaries come from incrementally decoding the delta's token prefix; they + are clamped and made monotone so the pieces always concatenate to exactly + ``text``. The final piece absorbs any remainder, so nothing is lost or + duplicated even when a token boundary does not fall on a character + boundary (multi-byte sequences, leading-space normalisation). + """ + n = len(token_ids) + if n <= 1 or not text: + return [(text, list(token_ids))] + bounds: list[int] = [] + if tokenizer is not None: + try: + for j in range(1, n): + prefix = tokenizer.decode(token_ids[:j], skip_special_tokens=True) + bounds.append(len(prefix)) + except Exception: # noqa: BLE001 - fall back to an even split + bounds = [] + if len(bounds) != n - 1: + bounds = [(len(text) * j) // n for j in range(1, n)] + pieces: list[tuple[str, list[int]]] = [] + start = 0 + for j, b in enumerate(bounds): + b = min(max(b, start), len(text)) + pieces.append((text[start:b], [token_ids[j]])) + start = b + pieces.append((text[start:], [token_ids[n - 1]])) + return pieces diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4ac8d49cd..bca89ea4d 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, + # One kernel for every batch size: no M-dependent dispatch, and no + # TensorRT-LLM nvcc JIT on the cold-start path. DeepGemmFp8BlockScaledMMKernel, + FlashInferFp8DeepGEMMDynamicBlockScaledKernel, CutlassFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 340a30403..cc6cd5e9c 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -177,7 +177,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] query_start_loc = m.query_start_loc query_start_loc_cpu = m.query_start_loc_cpu - context_lens_tensor = m.compute_num_computed_tokens() + # Dead-compute skip, gated off under speculative decoding. The + # context_lens_tensor build is only consumed by the prefill branch + # below, so decode-only steps can skip it when no speculative drafter + # is loaded. While spec decode (e.g. checkpoint-embedded MTP) is + # active we keep the baseline unconditional build: the drafter + # consumes late-layer state and we do not trust liveness derived + # without it. + context_lens_tensor = ( + m.compute_num_computed_tokens() if self.use_spec_decode else None + ) nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None block_table_tensor = mamba_get_block_table_tensor( m.block_table_tensor, @@ -387,6 +396,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) if num_prefills > 0: + if context_lens_tensor is None: + context_lens_tensor = m.compute_num_computed_tokens() has_initial_state = context_lens_tensor > 0 if spec_sequence_masks_cpu is not None: has_initial_state = has_initial_state[~spec_sequence_masks_cpu] diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index b3d6f5e4d..e4659104f 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -24,7 +24,8 @@ class AsyncOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.sampler_output = sampler_output self.num_sampled_tokens = num_sampled_tokens - self.copy_event = torch.cuda.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) @@ -81,7 +82,8 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.pooler_output = pooler_output self.is_valid = is_valid - self.copy_event = torch.cuda.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 4ab45b2ae..8e73ec8ab 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -12,7 +12,8 @@ class DraftTokensHandler: def __init__(self, device: torch.device | None = None): self.device = device self.copy_stream = torch.cuda.Stream(device) - self.copy_event = torch.cuda.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a823..8756c97a1 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -254,7 +254,8 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): self._invalid_req_indices = invalid_req_indices # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) # Keep a reference to the device tensor to avoid it being # deallocated until we finish copying it to the host. @@ -375,7 +376,8 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput): self._model_runner_output = model_runner_output # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) # Keep a reference to the device tensors to avoid them being # deallocated until we finish copying it to the host. @@ -697,7 +699,9 @@ class GPUModelRunner( self.prepare_inputs_event: torch.Event | None = None if self.use_async_scheduling: self.async_output_copy_stream = torch.cuda.Stream() - self.prepare_inputs_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock; + # under TP contention that spin can balloon and make the rank a straggler. + self.prepare_inputs_event = torch.cuda.Event(blocking=True) # self.cudagraph_batch_sizes sorts in ascending order. if (