diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93..cbaa9bb00 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1690,6 +1690,20 @@ class VllmConfig: max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) + # The bound above sizes a step that is purely decode: one + # query position per running sequence, times the speculative + # width. Under chunked prefill a step also carries prompt + # tokens, so its size is bounded by the token budget instead. + # Raise the default toward that budget while retaining a + # ceiling on capture time and device memory. An explicit + # max_cudagraph_capture_size still wins outright. + max_cudagraph_capture_size = max( + max_cudagraph_capture_size, + min( + self.scheduler_config.max_num_batched_tokens, + _cudagraph_capture_ceiling(), + ), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) @@ -1727,10 +1741,26 @@ class VllmConfig: range(8, min(max_cudagraph_capture_size + 1, 256), 8) ) if max_cudagraph_capture_size >= 256: - # Step size 16 for larger batch sizes + # Step size 16 for larger batch sizes, up to the dense + # limit (included) + dense_max = min( + max_cudagraph_capture_size, _CUDAGRAPH_CAPTURE_DENSE_LIMIT + ) + cudagraph_capture_sizes += list(range(256, dense_max + 1, 16)) + if max_cudagraph_capture_size > _CUDAGRAPH_CAPTURE_DENSE_LIMIT: + # Above the dense limit the sizes are prefill-shaped mixed + # steps whose cost is dominated by the token count itself, + # not by launch overhead, so capture a coarse ladder there + # (bounded padding, few graphs) and always include the top. + coarse_stride = _cudagraph_capture_coarse_stride() cudagraph_capture_sizes += list( - range(256, max_cudagraph_capture_size + 1, 16) + range( + _CUDAGRAPH_CAPTURE_DENSE_LIMIT + coarse_stride, + max_cudagraph_capture_size + 1, + coarse_stride, + ) ) + cudagraph_capture_sizes.append(max_cudagraph_capture_size) # ensure max_num_tokens is captured if within max capture size if ( max_num_tokens <= max_cudagraph_capture_size @@ -2292,3 +2322,70 @@ def get_layers_from_vllm_config( for layer_name in layer_names if isinstance(layer := forward_context.get(layer_name), layer_type) } + + +# Ceiling on the token-keyed CUDA graph capture range chosen in +# _set_cudagraph_sizes when the config pins no explicit size. Capture cost +# grows with the number and size of captured graphs, so the token budget is +# only honoured this far. Overridable so a sweep needs one build. +_CUDAGRAPH_CAPTURE_CEILING_ENV = "VLLM_CUDAGRAPH_CAPTURE_CEILING" +_CUDAGRAPH_CAPTURE_CEILING_DEFAULT = 2048 + +# Largest size captured on the dense 16-token stride. Above it the default +# list continues on a coarse stride (see _cudagraph_capture_coarse_stride) so +# that raising the ceiling costs a handful of graphs rather than dozens. +_CUDAGRAPH_CAPTURE_DENSE_LIMIT = 1024 +_CUDAGRAPH_CAPTURE_COARSE_STRIDE_ENV = "VLLM_CUDAGRAPH_CAPTURE_COARSE_STRIDE" +_CUDAGRAPH_CAPTURE_COARSE_STRIDE_DEFAULT = 256 + + +def _cudagraph_capture_ceiling() -> int: + """Return the configured positive capture ceiling.""" + raw = os.environ.get(_CUDAGRAPH_CAPTURE_CEILING_ENV) + if raw is None: + return _CUDAGRAPH_CAPTURE_CEILING_DEFAULT + try: + ceiling = int(raw) + except ValueError: + logger.warning( + "%s=%r is not an integer; using %d", + _CUDAGRAPH_CAPTURE_CEILING_ENV, + raw, + _CUDAGRAPH_CAPTURE_CEILING_DEFAULT, + ) + return _CUDAGRAPH_CAPTURE_CEILING_DEFAULT + if ceiling < 1: + logger.warning( + "%s=%d is not positive; using %d", + _CUDAGRAPH_CAPTURE_CEILING_ENV, + ceiling, + _CUDAGRAPH_CAPTURE_CEILING_DEFAULT, + ) + return _CUDAGRAPH_CAPTURE_CEILING_DEFAULT + return ceiling + + +def _cudagraph_capture_coarse_stride() -> int: + """Return the configured positive stride used above the dense limit.""" + raw = os.environ.get(_CUDAGRAPH_CAPTURE_COARSE_STRIDE_ENV) + if raw is None: + return _CUDAGRAPH_CAPTURE_COARSE_STRIDE_DEFAULT + try: + stride = int(raw) + except ValueError: + logger.warning( + "%s=%r is not an integer; using %d", + _CUDAGRAPH_CAPTURE_COARSE_STRIDE_ENV, + raw, + _CUDAGRAPH_CAPTURE_COARSE_STRIDE_DEFAULT, + ) + return _CUDAGRAPH_CAPTURE_COARSE_STRIDE_DEFAULT + if stride < 1: + logger.warning( + "%s=%d is not positive; using %d", + _CUDAGRAPH_CAPTURE_COARSE_STRIDE_ENV, + stride, + _CUDAGRAPH_CAPTURE_COARSE_STRIDE_DEFAULT, + ) + return _CUDAGRAPH_CAPTURE_COARSE_STRIDE_DEFAULT + return stride diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..a47b60cbd 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1838,6 +1838,8 @@ class EngineArgs: "enable_prefix_caching must be set by this point" ) + self._set_gdn_hybrid_prefill_defaults(model_config) + cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] gpu_memory_utilization=self.gpu_memory_utilization, @@ -2525,6 +2527,68 @@ class EngineArgs: ) self.enable_prefix_caching = False + def _set_gdn_hybrid_prefill_defaults(self, model_config: ModelConfig) -> None: + """Prefill-side defaults for hybrid gated-delta-net (GDN) models. + + Applies only to hybrid models that interleave ``linear_attention`` + (GDN) layers with full attention and whose implementation does not + support Mamba prefix caching (``SupportsMambaPrefixCaching``). For + those models, prefix caching can only run in the experimental + ``"align"`` Mamba cache mode, which forces every prompt spanning a + full block to be prefilled in two scheduler steps (the last aligned + chunk is scheduled separately so its recurrent state can be cached), + runs a per-request Python state pass on every step and copies the + recurrent state whenever a request's state block moves. Those costs + land on the TTFT path of every request, while the cache itself can + only pay off through cross-request prefix reuse. Unless the operator + asked for ``"align"`` mode explicitly, prefer the plain single-step + prefill: the recurrent state stays private to its request and there + is nothing to align. + + The same models default to the FLA/Triton GDN prefill kernel when the + operator pins ``--gdn-prefill-backend triton``. When speculative + decoding is off, hand the choice back to the kernel resolver + (``"auto"``), which picks the fused FlashInfer kernel where it is + supported and verified (Hopper; Blackwell with an intact CuteDSL + install) and stays on Triton everywhere else. No hardware is + hard-coded here: the resolver owns the support matrix. + + Both adjustments are independently opt-out for sweeps: + ``VLLM_GDN_HYBRID_KEEP_PREFIX_CACHING=1`` and + ``VLLM_GDN_HYBRID_KEEP_PREFILL_BACKEND=1``. + """ + if not current_platform.is_cuda(): + return + if not _is_gdn_hybrid_without_mamba_prefix_caching(model_config): + return + + if ( + self.enable_prefix_caching + and self.mamba_cache_mode == "none" + and not _env_flag("VLLM_GDN_HYBRID_KEEP_PREFIX_CACHING") + ): + logger.info( + "Disabling prefix caching for %s: the model does not support " + "Mamba prefix caching, so caching would run in the experimental " + "'align' mode and split every prompt into two prefill steps. " + "Set --mamba-cache-mode align to keep it.", + model_config.architecture, + ) + self.enable_prefix_caching = False + + if ( + self.gdn_prefill_backend == "triton" + and self.speculative_config is None + and not _env_flag("VLLM_GDN_HYBRID_KEEP_PREFILL_BACKEND") + ): + logger.info( + "Letting the GDN prefill kernel resolver choose the backend for " + "%s instead of pinning Triton/FLA. Set " + "VLLM_GDN_HYBRID_KEEP_PREFILL_BACKEND=1 to keep the pin.", + model_config.architecture, + ) + self.gdn_prefill_backend = None + def _set_default_reasoning_config_args(self): if not self.reasoning_parser: return @@ -2692,3 +2756,29 @@ def _raise_unsupported_error(feature_name: str): f"remove {feature_name} from your config." ) raise NotImplementedError(msg) + + +def _env_flag(name: str) -> bool: + """True when the environment variable is set to a truthy value.""" + return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on") + + +def _is_gdn_hybrid_without_mamba_prefix_caching(model_config: ModelConfig) -> bool: + """Hybrid model with gated-delta-net layers and no Mamba prefix caching. + + Detection is structural, not by model name: the model must be registered + as hybrid, must not implement ``SupportsMambaPrefixCaching`` (so prefix + caching could only run in ``"align"`` mode) and its text config must + declare ``linear_attention`` layers with a GDN head layout + (``linear_key_head_dim``), which is what the GDN prefill kernels key on. + """ + try: + if not model_config.is_hybrid or model_config.supports_mamba_prefix_caching: + return False + except (AttributeError, ValueError): + return False + text_config = getattr(model_config, "hf_text_config", None) + layer_types = getattr(text_config, "layer_types", None) + if not layer_types or "linear_attention" not in layer_types: + return False + return getattr(text_config, "linear_key_head_dim", None) is not None diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef174135..e86647a9d 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,43 @@ 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: + return [(text, list(token_ids))] + if not text: + # All-empty-text multi-token delta (e.g. UTF-8 byte fragments): emit one + # chunk per token so inter-token gap count always >= completion_tokens - 1. + return [("", [tid]) for tid in 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