diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index de505e122c..847ceac160 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy -from typing import TYPE_CHECKING, Any, Literal, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Literal, get_args from pydantic import Field, SkipValidation, field_validator, model_validator from typing_extensions import Self @@ -70,11 +70,55 @@ SpeculativeMethod = Literal[ RejectionSampleMethod = Literal["standard", "synthetic"] DraftSampleMethod = Literal["greedy", "probabilistic"] +BUNDLED_MTP_TOKENS_ENV = "SN10_MTP_K" +_DEFAULT_BUNDLED_MTP_TOKENS = 5 +_QWEN3_5_MODEL_TYPES = ("qwen3_5", "qwen3_5_moe") + + +def bundled_mtp_depth(model_config: ModelConfig) -> int: + """Number of MTP layers shipped inside a qwen3_5 checkpoint (0 if none).""" + hf_config = model_config.hf_config + if getattr(hf_config, "model_type", None) not in _QWEN3_5_MODEL_TYPES: + return 0 + for cfg in (model_config.hf_text_config, hf_config): + depth = getattr(cfg, "mtp_num_hidden_layers", None) + if isinstance(depth, int) and depth > 0: + return depth + return 0 + + +def default_bundled_mtp_args(model_config: ModelConfig) -> dict[str, Any] | None: + """Speculative-config kwargs enabling the checkpoint's own MTP head. + + Returns None when the model has no bundled MTP head or when + ``SN10_MTP_K`` is set to 0. + """ + import os + + if bundled_mtp_depth(model_config) == 0: + return None + raw = os.environ.get(BUNDLED_MTP_TOKENS_ENV, "") + try: + num_tokens = int(raw) if raw.strip() else _DEFAULT_BUNDLED_MTP_TOKENS + except ValueError: + logger.warning( + "Ignoring invalid %s=%r; using %d.", + BUNDLED_MTP_TOKENS_ENV, + raw, + _DEFAULT_BUNDLED_MTP_TOKENS, + ) + num_tokens = _DEFAULT_BUNDLED_MTP_TOKENS + if num_tokens <= 0: + return None + return {"method": "mtp", "num_speculative_tokens": num_tokens} + @config class SpeculativeConfig: """Configuration for speculative decoding.""" + MTP_TYPES_SHARING_TARGET_KV: ClassVar[tuple[str, ...]] = ("qwen3_5_mtp",) + enforce_eager: bool | None = None """Override the default enforce_eager from model_config""" # General speculative decoding control @@ -1109,6 +1153,24 @@ class SpeculativeConfig: def use_eagle(self) -> bool: return self.method in ("eagle", "eagle3", "mtp", "dflash") + def drafter_shares_target_kv_groups(self) -> bool: + """True when the drafter is an MTP layer bundled with the target model + whose attention layers live in the target's own KV cache groups.""" + if self.method != "mtp" or self.target_model_config is None: + return False + if ( + self.draft_model_config is None + or self.model != self.target_model_config.model + ): + return False + draft_type = getattr(self.draft_model_config.hf_config, "model_type", None) + return draft_type in self.MTP_TYPES_SHARING_TARGET_KV + + def requires_cache_tail_recompute(self) -> bool: + """Whether a prefix-cache hit must leave its last block uncached so the + drafter re-runs on it (the EAGLE last-block drop).""" + return self.use_eagle() and not self.drafter_shares_target_kv_groups() + def use_dflash(self) -> bool: return self.method == "dflash" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93b..07cf78036d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -285,6 +285,11 @@ OPTIMIZATION_LEVEL_TO_CONFIG = { OptimizationLevel.O3: OPTIMIZATION_LEVEL_03, } +# Chunked prefill with prefix caching produces steps of a short prompt tail +# plus the decode batch; graphs are captured up to this many tokens for them. +_MIXED_STEP_CAPTURE_LIMIT = 1024 +_MIXED_STEP_CAPTURE_STRIDE = 64 + @config(config=ConfigDict(arbitrary_types_allowed=True)) class VllmConfig: @@ -1685,12 +1690,22 @@ class VllmConfig: max_cudagraph_capture_size = ( self.compilation_config.max_cudagraph_capture_size ) + max_num_tokens = self.scheduler_config.max_num_batched_tokens + decode_query_len = 1 + self.num_speculative_tokens + decode_capture_limit = min( + self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 + ) + mixed_capture_limit = 0 + if ( + max_cudagraph_capture_size is None + and self.scheduler_config.enable_chunked_prefill + and self.cache_config.enable_prefix_caching + ): + mixed_capture_limit = min(max_num_tokens, _MIXED_STEP_CAPTURE_LIMIT) if max_cudagraph_capture_size is None: - decode_query_len = 1 + self.num_speculative_tokens - max_cudagraph_capture_size = min( - self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 + max_cudagraph_capture_size = max( + decode_capture_limit, mixed_capture_limit ) - max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) assert max_cudagraph_capture_size >= 1, ( @@ -1721,16 +1736,26 @@ class VllmConfig: cudagraph_capture_sizes = [ i for i in [1, 2, 4] if i <= max_cudagraph_capture_size ] - if max_cudagraph_capture_size >= 8: + dense_limit = max_cudagraph_capture_size + if mixed_capture_limit > decode_capture_limit: + dense_limit = min(dense_limit, decode_capture_limit) + if dense_limit >= 8: # Step size 8 for small batch sizes, up to 256(not included) cudagraph_capture_sizes += list( - range(8, min(max_cudagraph_capture_size + 1, 256), 8) + range(8, min(dense_limit + 1, 256), 8) ) - if max_cudagraph_capture_size >= 256: + if dense_limit >= 256: # Step size 16 for larger batch sizes + cudagraph_capture_sizes += list(range(256, dense_limit + 1, 16)) + if dense_limit < max_cudagraph_capture_size: + # Coarse ladder above the decode region for the mixed + # prefill-tail + decode steps of chunked prefill. + stride = _MIXED_STEP_CAPTURE_STRIDE + first = (dense_limit // stride + 1) * stride cudagraph_capture_sizes += list( - range(256, max_cudagraph_capture_size + 1, 16) + range(first, max_cudagraph_capture_size + 1, 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 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466b..f948bebd2a 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -95,6 +95,7 @@ from vllm.config.parallel import ( ExpertPlacementStrategy, ) from vllm.config.scheduler import SchedulerPolicy +from vllm.config.speculative import default_bundled_mtp_args from vllm.config.utils import get_field from vllm.config.vllm import OptimizationLevel, PerformanceMode from vllm.logger import init_logger, suppress_logging @@ -1719,7 +1720,13 @@ class EngineArgs: self.speculative_config[key] = value if self.speculative_config is None: - return None + self.speculative_config = default_bundled_mtp_args(target_model_config) + if self.speculative_config is None: + return None + logger.info( + "Enabling bundled MTP speculative decoding: %s", + self.speculative_config, + ) # Note(Shangming): These parameters are not obtained from the cli arg # '--speculative-config' and must be passed in when creating the engine diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index a16f522183..5d569ab6e5 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time import asyncio import importlib import inspect @@ -578,6 +579,8 @@ async def build_and_serve( app = build_app(args, supported_tasks, model_config) await init_app_state(engine_client, app.state, args, supported_tasks) + await _exercise_cached_prefix_paths(engine_client, model_config) + logger.info("Starting vLLM server on %s", listen_address) return await serve_http( @@ -663,6 +666,60 @@ async def run_server(args, **uvicorn_kwargs) -> None: await run_server_worker(listen_address, sock, args, **uvicorn_kwargs) + +_PREFIX_WARMUP_LENGTHS = (256, 640, 1024, 1300, 1700, 2100, 2600, 3200) +_PREFIX_WARMUP_TAIL = 37 +_PREFIX_WARMUP_MAX_TOKEN_ID = 50000 + + +async def _exercise_cached_prefix_paths(engine_client: EngineClient, model_config) -> None: + """Run the prefill-over-cached-prefix code paths once before serving. + + The engine's own warmup only sees cold prompts. The first prompt that hits + the prefix cache takes a different route (state restore + partial prefill of + the tail), whose kernels are compiled lazily and whose GEMM shapes are new, + so the first cache hits after startup pay several seconds of compile time. + Pay it here instead: run a few synthetic prompts twice, then once more with + a short extra tail, so both the exact-hit and the partial-hit shapes exist. + """ + if os.environ.get("SN10_PREFIX_WARMUP", "1").strip().lower() in ("0", "false", "no"): + return + try: + import random + + from vllm.inputs import TokensPrompt + from vllm.sampling_params import SamplingParams + + max_len = int(getattr(model_config, "max_model_len", 0) or 0) + vocab = int(model_config.get_vocab_size()) + top_id = max(2, min(vocab - 1, _PREFIX_WARMUP_MAX_TOKEN_ID)) + rng = random.Random(0x5310) + lengths = [n for n in _PREFIX_WARMUP_LENGTHS if n + _PREFIX_WARMUP_TAIL + 16 < max_len] + if not lengths: + return + bodies = [[rng.randrange(1, top_id) for _ in range(n)] for n in lengths] + params = SamplingParams(temperature=0.0, max_tokens=8, ignore_eos=True) + + async def one(ids: list[int], tag: str) -> None: + async for _ in engine_client.generate( + TokensPrompt(prompt_token_ids=ids), params, request_id=f"sn10-warm-{tag}" + ): + pass + + started = time.perf_counter() + for rnd in range(2): + await asyncio.gather(*(one(b, f"{rnd}-{i}") for i, b in enumerate(bodies))) + tails = [b + [rng.randrange(1, top_id) for _ in range(_PREFIX_WARMUP_TAIL)] for b in bodies] + await asyncio.gather(*(one(b, f"2-{i}") for i, b in enumerate(tails))) + logger.info( + "Cached-prefix warmup: %d prompts x 3 passes in %.1fs", + len(bodies), + time.perf_counter() - started, + ) + except Exception: + logger.warning("Cached-prefix warmup skipped", exc_info=True) + + async def run_server_worker( listen_address, sock, args, client_config=None, **uvicorn_kwargs ) -> None: diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef1741351..5d4f28ba4e 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -277,6 +277,36 @@ class OpenAIServingCompletion(OpenAIServing): return response + @staticmethod + def _split_delta_per_token( + delta_text: str, + token_ids: list[int], + tokenizer: TokenizerLike | None, + skip_special_tokens: bool, + ) -> list[str]: + """Split a multi-token delta into one text piece per token. + + The pieces always concatenate to ``delta_text``. When the delta cannot + be attributed to individual tokens (no tokenizer, or a token boundary + that does not fall on a character boundary), all text is carried by + the final piece. + """ + num_tokens = len(token_ids) + if tokenizer is not None: + pieces: list[str] = [] + seen = "" + for end in range(1, num_tokens + 1): + prefix = tokenizer.decode( + token_ids[:end], skip_special_tokens=skip_special_tokens + ) + if not prefix.startswith(seen): + break + pieces.append(prefix[len(seen) :]) + seen = prefix + if len(pieces) == num_tokens and seen == delta_text: + return pieces + return [""] * (num_tokens - 1) + [delta_text] + async def completion_stream_generator( self, request: CompletionRequest, @@ -332,7 +362,8 @@ class OpenAIServingCompletion(OpenAIServing): prompt_token_ids_to_return: list[int] | None = None assert request.max_tokens is not None - if request.echo and not has_echoed[i]: + echoes_prompt = request.echo and not has_echoed[i] + if echoes_prompt: assert prompt_token_ids is not None if request.return_token_ids: prompt_text = "" @@ -395,46 +426,72 @@ 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). + step_token_ids = as_list(output.token_ids) if ( - not include_usage - and self.system_fingerprint is not None - and finish_reason is not None + len(step_token_ids) > 1 + and logprobs is None + and not echoes_prompt ): - 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, + texts = self._split_delta_per_token( + delta_text, + step_token_ids, + tokenizer, + request.skip_special_tokens, + ) + id_groups = [[tid] for tid in step_token_ids] + else: + texts = [delta_text] + id_groups = [step_token_ids] + + tokens_so_far = previous_num_tokens[i] - len(step_token_ids) + last_idx = len(texts) - 1 + for piece_idx, piece_text in enumerate(texts): + at_end = piece_idx == last_idx + tokens_so_far += len(id_groups[piece_idx]) + chunk = CompletionStreamResponse( + id=request_id, + object="text_completion", + created=created_time, + model=model_name, + choices=[ + CompletionResponseStreamChoice( + index=i, + text=piece_text, + logprobs=logprobs if at_end else None, + finish_reason=finish_reason if at_end else None, + stop_reason=stop_reason if at_end else None, + prompt_token_ids=( + prompt_token_ids_to_return + if piece_idx == 0 + else None + ), + token_ids=( + id_groups[piece_idx] + 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 ( + at_end + 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_so_far, + total_tokens=prompt_tokens + tokens_so_far, + ) - 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/envs.py b/vllm/envs.py index 27a85bb3d0..9de2ff05b9 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -185,7 +185,7 @@ if TYPE_CHECKING: "relax", ] = "relax" VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True - VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True + VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = False VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto" @@ -1471,7 +1471,7 @@ environment_variables: dict[str, Callable[[], Any]] = { # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( - int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1")) + int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "0")) ), # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops. "VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool( diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 7e0c7e05ca..01303b971b 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -11,6 +11,8 @@ import torch from vllm.triton_utils import tl, triton +from .op import exp + @triton.heuristics( { @@ -26,8 +28,6 @@ def fused_sigmoid_gating_delta_rule_update_kernel( a, b, dt_bias, - beta, - threshold, q, k, v, @@ -51,6 +51,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( stride_final_state_token: tl.constexpr, stride_indices_seq: tl.constexpr, stride_indices_tok: tl.constexpr, + SOFTPLUS_BETA: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, # whether to use initial state INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace USE_QK_L2NORM_IN_KERNEL: tl.constexpr, @@ -125,25 +127,31 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) b_b = tl.load(p_b).to(tl.float32) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + # If the model is loaded in fp16, without the .float() here, A might be -inf x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32) - softplus_x = tl.where( - beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x - ) + if SOFTPLUS_BETA == 1.0: + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + else: + softplus_x = tl.where( + SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, + (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), + x, + ) b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x - # compute beta_output = sigmoid(b) - b_beta = tl.sigmoid(b_b.to(tl.float32)) + # sigmoid(b) rounded through b's storage dtype, as in the decode kernel + b_beta = tl.sigmoid(b_b).to(b.dtype.element_ty).to(tl.float32) - if USE_QK_L2NORM_IN_KERNEL: - b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)) - b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)) - b_q = b_q * scale # [BV, BK] if not IS_KDA: - b_h *= tl.exp(b_g) + b_h *= exp(b_g) else: - b_h *= tl.exp(b_g[None, :]) + b_h *= exp(b_g[None, :]) # [BV] b_v -= tl.sum(b_h * b_k[None, :], 1) b_v *= b_beta @@ -209,7 +217,7 @@ def fused_sigmoid_gating_delta_rule_update( NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" num_stages = 3 - num_warps = 4 + num_warps = 1 if cu_seqlens is not None and q.shape[0] != 1: raise ValueError( @@ -244,8 +252,6 @@ def fused_sigmoid_gating_delta_rule_update( a=a.contiguous(), b=b.contiguous(), dt_bias=dt_bias, - beta=beta, - threshold=threshold, q=q.contiguous(), k=k.contiguous(), v=v.contiguous(), @@ -269,6 +275,8 @@ def fused_sigmoid_gating_delta_rule_update( stride_final_state_token=stride_final_state_token, stride_indices_seq=stride_indices_seq, stride_indices_tok=stride_indices_tok, + SOFTPLUS_BETA=float(beta), + SOFTPLUS_THRESHOLD=float(threshold), INPLACE_FINAL_STATE=inplace_final_state, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_KDA=is_kda, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 90d93a110c..4c88603d78 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -227,6 +227,7 @@ class Scheduler(SchedulerInterface): speculative_config = vllm_config.speculative_config self.use_eagle = False + self.drop_cached_tail_block = False self.num_spec_tokens = vllm_config.num_speculative_tokens self.num_lookahead_tokens = 0 self.dynamic_sd_lookup: list[int] | None = None @@ -240,6 +241,9 @@ class Scheduler(SchedulerInterface): if speculative_config.use_eagle(): self.use_eagle = True self.num_lookahead_tokens = self.num_spec_tokens + self.drop_cached_tail_block = ( + speculative_config.requires_cache_tail_recompute() + ) if speculative_config.uses_draft_model(): self.num_lookahead_tokens = self.num_spec_tokens if speculative_config.use_dflash(): @@ -256,7 +260,7 @@ class Scheduler(SchedulerInterface): max_model_len=self.max_model_len, max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, enable_caching=self.cache_config.enable_prefix_caching, - use_eagle=self.use_eagle, + use_eagle=self.drop_cached_tail_block, log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, @@ -357,7 +361,7 @@ class Scheduler(SchedulerInterface): block_size = self.cache_config.block_size last_cache_position = request.num_tokens - request.num_tokens % block_size # eagle prune - if self.use_eagle: + if self.drop_cached_tail_block: last_cache_position = max(last_cache_position - block_size, 0) num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens if num_computed_tokens_after_sched < last_cache_position: