diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..a37196936 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -134,6 +134,64 @@ 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. Depth still pays well +# past the point acceptance starts falling, because the step amortizes one +# read of the model's weights over every token it does accept, and that read +# dominates a decode step. +AUTO_MTP_NUM_SPECULATIVE_TOKENS = 7 + + +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 +1776,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/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 1d61ca3c5..0ee0162c1 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -179,7 +179,16 @@ class CompletionRequest(OpenAIBaseModel): ) repetition_detection: RepetitionDetectionParams | None = Field( - default=None, + # A completion that has started repeating itself is finished being + # useful, and every further token costs the whole batch throughput + # until the length cap ends it. The default stops it instead. The + # pattern has to be at least 8 tokens long and repeat 4 times back to + # back, which is 32 tokens of exact repetition: prose and code do not + # do that, and a request that wants the old behaviour can still ask + # for it explicitly. + default_factory=lambda: RepetitionDetectionParams( + min_pattern_size=8, max_pattern_size=64, min_count=4 + ), description="Parameters for detecting repetitive N-gram patterns " "in output tokens. If such repetition is detected, generation will " "be ended early. LLMs can sometimes generate repetitive, unhelpful " 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) diff --git a/vllm/model_executor/layers/fla/ops/solve_tril.py b/vllm/model_executor/layers/fla/ops/solve_tril.py index 8d3811ca4..f9987428c 100644 --- a/vllm/model_executor/layers/fla/ops/solve_tril.py +++ b/vllm/model_executor/layers/fla/ops/solve_tril.py @@ -18,7 +18,7 @@ from .index import prepare_chunk_indices from .op import make_tensor_descriptor from .utils import input_guard, is_amd, is_tma_supported -FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "ieee") +FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "tf32") ALLOWED_TRIL_PRECISIONS = ["ieee", "tf32"] if is_amd else ["ieee", "tf32", "tf32x3"] assert FLA_TRIL_PRECISION in ALLOWED_TRIL_PRECISIONS, ( f"FLA_TRIL_PRECISION must be one of {ALLOWED_TRIL_PRECISIONS}, but got {FLA_TRIL_PRECISION}" diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c5d..34cad972f 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -82,6 +82,11 @@ if GDN_AITER_TRITON_AVAILABLE: logger = init_logger(__name__) +# Sequence length used to warm up (and therefore autotune) the chunked +# prefill kernels. Their autotune keys do not include the sequence length, +# so this is the shape every later prefill inherits its config from. +GDN_AUTOTUNE_WARMUP_TOKENS = 2048 + # TODO(arpera): remove ``_is_libs_cu13_install_intact`` and its caller in # ``_resolve_gdn_prefill_backend`` once the upstream packaging bug is @@ -1082,9 +1087,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): results are cached globally, so only the first layer incurs actual benchmarking cost. - All kernels including ``chunk_fwd_kernel_o`` now use a fixed - ``BT = chunk_size`` (64). A single warmup pass with T = 64 - is sufficient to populate the autotuner cache. + All kernels including ``chunk_fwd_kernel_o`` use a fixed + ``BT = chunk_size`` (64), but their autotune keys carry only + ``H/K/V/BT`` -- never the sequence length -- so the config chosen + during warmup is reused for every later shape. Warming up at one + chunk would tune against a single iteration of the chunk loop, where + pipelining depth and grid size cannot show any benefit. Use a + prompt-sized pass so the tuner sees the shape it is choosing for. The decode path uses ``gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule`` which has fixed kernel parameters (no autotuning), so only the @@ -1104,7 +1113,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # is sufficient to populate every autotuner cache. Mirror the real # prefill path here: build q/k/v/g/beta via fused_post_conv_prep and # then run chunk_gated_delta_rule with in-kernel L2 norm disabled. - T = FLA_CHUNK_SIZE + T = GDN_AUTOTUNE_WARMUP_TOKENS dummy_mixed_qkv = torch.randn( T, qkv_or_qkvz.shape[-1] - v_dim, device=device, dtype=dtype ) @@ -1510,8 +1519,12 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): prefill_has_initial_state = attn_metadata.prefill_has_initial_state assert prefill_state_indices is not None assert prefill_has_initial_state is not None - initial_state = ssm_state[prefill_state_indices] - initial_state[~prefill_has_initial_state, ...] = 0 + initial_state = ssm_state.index_select(0, prefill_state_indices) + initial_state = torch.where( + prefill_has_initial_state.view(-1, 1, 1, 1), + initial_state, + initial_state.new_zeros(()), + ) ( core_attn_out_non_spec, last_recurrent_state,