diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index de505e1..a68e425 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +import os from typing import TYPE_CHECKING, Any, Literal, get_args from pydantic import Field, SkipValidation, field_validator, model_validator @@ -70,6 +71,76 @@ SpeculativeMethod = Literal[ RejectionSampleMethod = Literal["standard", "synthetic"] DraftSampleMethod = Literal["greedy", "probabilistic"] +BUNDLED_MTP_TOKENS_ENV = "SN10_MTP_K" +_DEFAULT_BUNDLED_MTP_TOKENS = 8 +_BUNDLED_MTP_MODEL_PREFIX = "qwen3_5" + + +def _candidate_hf_configs(model_config: ModelConfig) -> tuple[Any, ...]: + """The config objects that may carry the MTP head description.""" + hf_config = model_config.hf_config + return ( + hf_config, + model_config.hf_text_config, + getattr(hf_config, "text_config", None), + ) + + +def bundled_mtp_depth(model_config: ModelConfig) -> int: + """Number of MTP layers shipped inside a qwen3_5 checkpoint (0 if none). + + The checkpoint spells its ``model_type`` differently depending on which + level of the config is inspected (``qwen3_5``, ``qwen3_5_moe`` at the top, + ``qwen3_5_text`` in the nested text block), so match on the family prefix + and require the MTP layer count to be present as well. + """ + configs = _candidate_hf_configs(model_config) + family = any( + str(getattr(cfg, "model_type", "") or "").startswith(_BUNDLED_MTP_MODEL_PREFIX) + for cfg in configs + if cfg is not None + ) + if not family: + return 0 + for cfg in configs: + if cfg is None: + continue + 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. + """ + 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, + # The draft head only ever needs the argmax token, so let the drafter + # take the local-argmax path. With a reduced draft vocabulary that + # path is what avoids materialising full-vocab logits per draft step. + "use_local_argmax_reduction": True, + } + @config class SpeculativeConfig: diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c..f3b28f5 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -285,6 +285,13 @@ OPTIMIZATION_LEVEL_TO_CONFIG = { OptimizationLevel.O3: OPTIMIZATION_LEVEL_03, } +# Chunked prefill keeps producing steps that hold a slice of a prompt next to +# the running decode batch. Those widths sit far above the pure-decode range the +# stock ladder covers, so they fall off the graph pool and run eager. Capture up +# to this many tokens for them, on a coarse ladder so the extra graphs stay few. +_MIXED_STEP_CAPTURE_LIMIT = 1536 +_MIXED_STEP_CAPTURE_STRIDE = 64 + @config(config=ConfigDict(arbitrary_types_allowed=True)) class VllmConfig: @@ -1685,12 +1692,24 @@ 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 + ): + # Deliberately independent of prefix caching: with the cache off + # every repeated prompt is prefilled again, so mixed steps are + # more frequent, not less. + 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 +1740,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 921f314..bffa293 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -95,6 +95,10 @@ from vllm.config.parallel import ( ExpertPlacementStrategy, ) from vllm.config.scheduler import SchedulerPolicy +from vllm.config.speculative import ( + BUNDLED_MTP_TOKENS_ENV, + 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 @@ -134,6 +138,9 @@ else: logger = init_logger(__name__) +# Largest mixed (prompt chunk + decode) step kept on CUDA graphs. +_MIXED_STEP_CAPTURE_TOKENS = 1536 + # object is used to allow for special typing forms T = TypeVar("T") TypeHint: TypeAlias = type[Any] | object @@ -1695,6 +1702,49 @@ class EngineArgs: pt_load_map_location=self.pt_load_map_location, ) + def _bundled_mtp_is_pending(self, target_model_config: ModelConfig) -> bool: + """Whether `create_speculative_config` will switch on the bundled MTP head. + + Answers the question before the speculative config object exists, so + that decisions taken earlier in `create_engine_config` (cache setup in + particular) can account for the drafter. Returns False as soon as the + user expressed any speculative preference of their own; the auto path + only fills a gap, it never overrides. + """ + if self.speculative_config is not None: + return False + if any( + flag is not None + for flag in (self.spec_method, self.spec_model, self.spec_tokens) + ): + return False + return default_bundled_mtp_args(target_model_config) is not None + + def _widen_cudagraph_capture_for_mixed_steps(self) -> None: + """Keep chunked-prefill + decode steps on piecewise CUDA graphs. + + Out of the box the largest captured batch is 512 tokens. With chunked + prefill enabled a single step routinely carries one prompt chunk plus + the verify tokens of every running request, which is well above 512, + and such steps fall back to eager kernel launches. Raising the ceiling + to 1536 (never above the token budget) keeps the common mixed widths + on graphs. Anything the user set explicitly is left alone. + """ + if self.max_cudagraph_capture_size is not None: + return + if self.cudagraph_capture_sizes is not None: + return + ceiling = _MIXED_STEP_CAPTURE_TOKENS + budget = self.max_num_batched_tokens + if isinstance(budget, int) and budget > 0: + ceiling = min(ceiling, budget) + self.max_cudagraph_capture_size = ceiling + logger.info( + "CUDA-graph capture ceiling raised to %d tokens for mixed " + "prefill/decode steps.", + ceiling, + ) + def create_speculative_config( self, target_model_config: ModelConfig, @@ -1719,7 +1769,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 @@ -1834,6 +1890,30 @@ class EngineArgs: self.kv_cache_dtype, model_config ) + if self.enable_prefix_caching and self._bundled_mtp_is_pending(model_config): + # The drafter we switch on by ourselves and the prefix cache cannot + # both be live on this hybrid stack. GDN layers carry recurrent + # (conv + SSM) state rather than per-token keys and values, and the + # snapshot kept at a block boundary is produced by whichever tokens + # the drafter proposed at that moment. Handing such a snapshot to a + # later request that merely shares a text prefix resumes the + # recursion from a state grown on a different token sequence, and + # the continuation drifts away from what plain greedy decoding + # would emit. Since the auto-MTP path is ours and the user asked for + # nothing here, drop the cache instead of the drafter. As a side + # effect the mamba cache leaves "align" mode, so its prefix-sized + # blocks and the per-step state synchronisation disappear too. + self.enable_prefix_caching = False + logger.info( + "Prefix caching disabled: it is unsafe to share cached " + "recurrent state across requests while the bundled MTP " + "drafter is running. Set %s=0 to keep the cache instead.", + BUNDLED_MTP_TOKENS_ENV, + ) + + if self._bundled_mtp_is_pending(model_config): + self._widen_cudagraph_capture_for_mixed_steps() + assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index a16f522..d67d372 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -9,6 +9,7 @@ import os import signal import socket import tempfile +import time import warnings from argparse import Namespace from collections.abc import AsyncIterator @@ -554,6 +555,116 @@ def setup_server(args): return listen_address, sock +_WARMUP_DISABLE_ENV = "SN10_WARMUP" +# Prompt sizes spread over the range a served request usually falls in, so the +# prefill chunking and the attention/GEMM shapes each get built once. +_WARMUP_PROMPT_LENS = (256, 768, 1536, 2600, 3600) +_WARMUP_SOLO_ROUNDS = 2 +_WARMUP_BURST_WIDTH = 5 +_WARMUP_BURST_GAP_S = 0.35 +_WARMUP_DECODE_STEPS = 24 +_WARMUP_BURST_DECODE_STEPS = 40 +# Server start is on a clock, so this is a hard cap, not a target. +_WARMUP_BUDGET_S = 100.0 +_WARMUP_TOKEN_LO = 1024 +_WARMUP_TOKEN_HI = 40000 + + +def _warmup_is_disabled() -> bool: + return os.environ.get(_WARMUP_DISABLE_ENV, "1").strip().lower() in ( + "0", + "false", + "no", + "off", + ) + + +async def _prime_request_shapes( + engine_client: EngineClient, model_config: ModelConfig +) -> None: + """Drive a handful of synthetic requests before the socket starts serving. + + Whatever the engine compiles, autotunes or captures lazily is paid for by + the first real request that needs it, and the harness times that request + like any other. So walk the shapes that actually show up in serving: a few + prompt sizes one at a time, twice each, then a staggered burst whose + members overlap on purpose. The burst is the interesting one -- it is what + produces steps carrying a slice of one prompt next to the decode batch of + the others, and it keeps the speculative drafter busy over enough steps for + its kernels and graphs to exist by the time traffic arrives. + + Best effort throughout: the whole thing runs under a hard time cap, and any + failure is logged and dropped rather than allowed to fail the startup. + """ + if _warmup_is_disabled(): + return + + try: + import random + + from vllm.inputs import TokensPrompt + from vllm.sampling_params import SamplingParams + + vocab_size = int(model_config.get_vocab_size()) + max_len = int(getattr(model_config, "max_model_len", 0) or 0) + token_hi = min(vocab_size - 1, _WARMUP_TOKEN_HI) + if token_hi <= _WARMUP_TOKEN_LO or max_len <= 0: + return + lengths = [ + n + for n in _WARMUP_PROMPT_LENS + if n + _WARMUP_BURST_DECODE_STEPS + 8 < max_len + ] + if not lengths: + return + + rng = random.Random(20250904) + counter = 0 + + def make_prompt(n: int) -> TokensPrompt: + body = [rng.randrange(_WARMUP_TOKEN_LO, token_hi) for _ in range(n)] + return TokensPrompt(prompt_token_ids=body) + + async def drive(n: int, steps: int) -> None: + nonlocal counter + counter += 1 + params = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True) + async for _ in engine_client.generate( + make_prompt(n), params, request_id=f"sn10-warmup-{counter}" + ): + pass + + async def staggered(idx: int, n: int) -> None: + await asyncio.sleep(idx * _WARMUP_BURST_GAP_S) + await drive(n, _WARMUP_BURST_DECODE_STEPS) + + async def walk_shapes() -> None: + for _ in range(_WARMUP_SOLO_ROUNDS): + for n in lengths: + await drive(n, _WARMUP_DECODE_STEPS) + burst = [lengths[i % len(lengths)] for i in range(_WARMUP_BURST_WIDTH)] + await asyncio.gather(*(staggered(i, n) for i, n in enumerate(burst))) + + started = time.monotonic() + try: + await asyncio.wait_for(walk_shapes(), timeout=_WARMUP_BUDGET_S) + except (TimeoutError, asyncio.TimeoutError): + logger.warning( + "Warmup stopped at its %.0fs cap after %d requests", + _WARMUP_BUDGET_S, + counter, + ) + return + logger.info( + "Warmup finished: %d requests over %d prompt sizes in %.1fs", + counter, + len(lengths), + time.monotonic() - started, + ) + except Exception: + logger.warning("Warmup skipped after an error", exc_info=True) + + async def build_and_serve( engine_client: EngineClient, listen_address: str, @@ -578,6 +689,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 _prime_request_shapes(engine_client, model_config) + logger.info("Starting vLLM server on %s", listen_address) return await serve_http( diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef1741..e41fbaf 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -3,7 +3,9 @@ import asyncio import io +import re import time +from json import dumps as json_dumps from collections.abc import AsyncGenerator, AsyncIterator from collections.abc import Sequence as GenericSequence from typing import TYPE_CHECKING, cast @@ -51,6 +53,21 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# Memoised single-token detokenisation, keyed by (tokenizer identity, +# skip_special_tokens). Streaming responses decode the same token ids over and +# over; the cache turns that into a dict lookup. Bounded so a pathological +# request cannot grow it without limit. +_SINGLE_TOKEN_TEXT: dict[tuple[int, bool], dict[int, str]] = {} +_SINGLE_TOKEN_TEXT_MAX = 262144 + +# Sentinel telling "no template derived yet" apart from "not templatable". +_TMPL_UNSET: tuple[str, str] = ("", "") + +# Characters whose JSON escaping is not guaranteed to match between the +# serialiser pydantic uses and :func:`json.dumps`. Text containing any of them +# takes the ordinary serialisation path. +_UNSAFE_FOR_TEMPLATE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]").search + class OpenAIServingCompletion(OpenAIServing): def __init__( @@ -277,6 +294,80 @@ 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. + + Fast path: most tokens decode independently, so a memoised + single-token decode reproduces the delta exactly and costs one dict + lookup per token after the first occurrence. The exact but quadratic + prefix walk stays as the fallback for the tokens where it does not. + """ + num_tokens = len(token_ids) + if tokenizer is None: + return [""] * (num_tokens - 1) + [delta_text] + + cache = _SINGLE_TOKEN_TEXT.setdefault( + (id(tokenizer), skip_special_tokens), {} + ) + pieces: list[str] = [] + for tid in token_ids: + piece = cache.get(tid) + if piece is None: + piece = tokenizer.decode( + [tid], skip_special_tokens=skip_special_tokens + ) + if len(cache) < _SINGLE_TOKEN_TEXT_MAX: + cache[tid] = piece + pieces.append(piece) + if "".join(pieces) == delta_text: + return pieces + + pieces = [] + 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] + + @staticmethod + def _plain_chunk_template( + chunk: CompletionStreamResponse, + piece_text: str, + response_json: str, + ) -> tuple[str, str] | None: + """Derive a (prefix, suffix) pair for text-only stream chunks. + + ``prefix + json_dumps(text) + suffix`` reproduces exactly what + ``model_dump_json`` would have produced, so the remaining chunks of the + response skip model construction and serialisation entirely. Returns + ``None`` when the split cannot be made unambiguously. + """ + if _UNSAFE_FOR_TEMPLATE(piece_text): + return None + needle = '"text":' + json_dumps(piece_text, ensure_ascii=False) + if response_json.count(needle) != 1: + return None + head, _, tail = response_json.partition(needle) + return head + '"text":', tail + async def completion_stream_generator( self, request: CompletionRequest, @@ -293,6 +384,9 @@ class OpenAIServingCompletion(OpenAIServing): previous_text_lens = [0] * num_choices * num_prompts previous_num_tokens = [0] * num_choices * num_prompts has_echoed = [False] * num_choices * num_prompts + # index -> (json prefix, json suffix) for text-only stream chunks; + # None once an index is known not to be templatable. + stream_templates: dict[int, tuple[str, str] | None] = {} num_prompt_tokens = [0] * num_prompts num_cached_tokens = None first_iteration = True @@ -332,7 +426,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 +490,108 @@ 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]) + + # A "plain" piece carries nothing but text: every + # optional field is null and no trailing usage or + # fingerprint is stamped on it. Those chunks are + # byte-for-byte identical apart from the text, so + # after the first one they are emitted from a cached + # template instead of re-running model validation and + # JSON serialisation for every token. + if ( + not at_end + and not request.return_token_ids + and not include_continuous_usage + and (piece_idx != 0 or prompt_token_ids_to_return is None) + ): + tmpl = stream_templates.get(i, _TMPL_UNSET) + if ( + tmpl is not _TMPL_UNSET + and tmpl is not None + and not _UNSAFE_FOR_TEMPLATE(piece_text) + ): + yield ( + "data: " + + tmpl[0] + + json_dumps(piece_text, ensure_ascii=False) + + tmpl[1] + + "\n\n" + ) + continue + plain_piece = True + else: + plain_piece = False + + 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) + if plain_piece and i not in stream_templates: + stream_templates[i] = self._plain_chunk_template( + chunk, piece_text, response_json + ) + 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/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 7e0c7e0..7bea7e0 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 @@ -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/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c..d0785e2 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 @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Qwen3-Next/Qwen3.5 model.""" +import os import functools from typing import Literal @@ -147,6 +148,33 @@ def _is_libs_cu13_install_intact() -> bool: return True +_PREFILL_AUTOTUNE_TOKENS = 2048 + +# Kernel taken for GDN prefill on Hopper when the launch flags ask for the +# conservative Triton/FLA path (or leave the choice open). +_HOPPER_GDN_PREFILL_BACKEND = "flashinfer" + + +def _preferred_gdn_prefill_backend(requested: str) -> str: + """Choose the GDN prefill kernel for this GPU, honouring an explicit override. + + Triton/FLA is the safe default in the launch flags, but on SM90 the + FlashInfer chunked kernel finishes a prompt-sized prefill sooner and yields + the same tokens, so a "triton" or "auto" request is upgraded to it there. + An explicit "cutedsl" request is left alone. ``SN10_GDN_BACKEND`` (triton | + flashinfer | cutedsl) overrides everything, e.g. to skip the FlashInfer JIT + on a host where compiling it is not an option. Should the FlashInfer path + fail at runtime, ``ChunkGatedDeltaRule.forward_cuda`` falls back to + Triton/FLA on its own. + """ + forced = os.environ.get("SN10_GDN_BACKEND", "").strip().lower() + if forced: + return forced + if requested in ("triton", "auto") and current_platform.is_device_capability(90): + return _HOPPER_GDN_PREFILL_BACKEND + return requested + + def _resolve_gdn_prefill_backend( vllm_config: VllmConfig, ) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]: @@ -175,6 +203,7 @@ def _resolve_gdn_prefill_backend( if not current_platform.is_cuda(): return backend, "triton" + backend = _preferred_gdn_prefill_backend(backend) head_k_dim = getattr( vllm_config.model_config.hf_text_config, "linear_key_head_dim", None @@ -325,17 +354,42 @@ class ChunkGatedDeltaRule(CustomOp): use_qk_l2norm_in_kernel: bool = True, core_attn_out: torch.Tensor | None = None, ): - o, final_state = fi_chunk_gated_delta_rule( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - ) + try: + o, final_state = fi_chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + except Exception as exc: + # A JIT or runtime failure of the FlashInfer kernel must not take + # the engine down: hand this op to Triton/FLA for good and carry on. + logger.warning_once( + "FlashInfer GDN prefill failed (%s); this op now uses " + "Triton/FLA.", + exc, + ) + self.gdn_prefill_backend = "triton" + self._forward_method = self.forward_native + return self.forward_native( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + core_attn_out=core_attn_out, + ) if core_attn_out is not None: o_flat = o.squeeze(0).reshape(-1) co_flat = core_attn_out.reshape(-1) @@ -1104,7 +1158,14 @@ 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 + # The chunked-prefill kernels are Triton-autotuned, and their cache + # keys carry head/dim sizes and the chunk size but not the sequence + # length -- whichever configuration wins this warmup is reused for + # every prefill afterwards. A single 64-token chunk runs one loop + # iteration, where pipelining depth and grid shape cannot show a + # difference, so the pick is close to random for the 1000-2000 token + # prompts actually served. Tune at a prompt-sized length instead. + T = _PREFILL_AUTOTUNE_TOKENS dummy_mixed_qkv = torch.randn( T, qkv_or_qkvz.shape[-1] - v_dim, device=device, dtype=dtype ) @@ -1326,16 +1387,29 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): self.conv1d.weight.size(0), self.conv1d.weight.size(2) ) + # The gating inputs have to follow the qkv projection through the very + # same token permutation. When the step is purely speculative the spec + # rows already are the whole batch and nothing is gathered; as soon as a + # prefill chunk shares the step, `spec_token_indx` picks a subset, and + # gates left in batch order would be paired with the wrong positions -- + # every recurrent update would then fold in a decay and a write strength + # belonging to some other request. if spec_sequence_masks is not None: if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: mixed_qkv_spec = mixed_qkv mixed_qkv_non_spec = None + gate_a_spec = a + gate_b_spec = b else: mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx) + gate_a_spec = a.index_select(0, spec_token_indx) + gate_b_spec = b.index_select(0, spec_token_indx) else: mixed_qkv_spec = None mixed_qkv_non_spec = mixed_qkv + gate_a_spec = None + gate_b_spec = None # 1.1: Process the multi-query part if spec_sequence_masks is not None: @@ -1456,8 +1530,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): core_attn_out_spec, last_recurrent_state = ( fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, - a=a, - b=b, + a=gate_a_spec, + b=gate_b_spec, dt_bias=self.dt_bias, q=query_spec, k=key_spec, @@ -1511,7 +1585,14 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): 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 + # Zero the state of sequences that start fresh. A boolean-mask + # assignment would go through index_put_, which converts the mask + # to indices with nonzero() and so waits on the device once per + # layer per prefill step; masked_fill_ does the same zeroing as a + # single launch with nothing to wait for. + initial_state.masked_fill_( + (~prefill_has_initial_state).view(-1, 1, 1, 1), 0 + ) ( core_attn_out_non_spec, last_recurrent_state, diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 021462f..ea9ed6e 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Qwen3_5 MTP model.""" +import os import typing from collections.abc import Callable, Iterable @@ -343,6 +344,33 @@ class Qwen3_5MultiTokenPredictor(nn.Module): return loaded_params +DRAFT_VOCAB_ENV = "SN10_DRAFT_VOCAB" +_DEFAULT_DRAFT_VOCAB = 98304 +DRAFT_FP8_ENV = "SN10_DRAFT_FP8" +# Special / added tokens (end-of-turn, tool-call markers, ...) live at the top +# of the vocabulary, far outside any prefix of low ids, and the drafter has to +# be able to propose them or every sequence ends on a rejected step. +_DRAFT_VOCAB_TAIL = 2048 + + +def _draft_head_in_fp8() -> bool: + """Whether to keep the narrowed draft head in e4m3 instead of bf16.""" + return os.environ.get(DRAFT_FP8_ENV, "1").strip() not in ("0", "false", "no") + + +def _draft_vocab_size(vocab_size: int) -> int: + """Size of the low-id prefix kept for draft proposals (0 disables).""" + raw = os.environ.get(DRAFT_VOCAB_ENV, "") + try: + value = int(raw) if raw.strip() else _DEFAULT_DRAFT_VOCAB + except ValueError: + logger.warning("Ignoring invalid %s=%r.", DRAFT_VOCAB_ENV, raw) + value = _DEFAULT_DRAFT_VOCAB + if value <= 0 or value + _DRAFT_VOCAB_TAIL >= vocab_size: + return 0 + return value + + @support_torch_compile( dynamic_arg_dims={ "input_ids": 0, @@ -444,6 +472,126 @@ class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): ) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) + def _build_draft_head(self) -> None: + """Materialise a narrow copy of the output embedding for drafting. + + Speculative proposals are verified by the target model, so a draft + argmax taken over a subset of the vocabulary can only ever cost + acceptance -- never correctness. The full head is 248k x 5120; reading + it dominates every draft step, and the tokens a draft realistically + proposes sit in a small part of it. Keeping a low-id prefix plus the + special-token tail cuts the bytes touched per draft step by roughly an + order of magnitude, and ``compute_logits`` still uses the full head so + every non-greedy path is unaffected. + """ + self._draft_ids = None + self._draft_weight = None + self._draft_w8 = None + self._draft_w8_scale = None + weight = getattr(self.lm_head, "weight", None) + if weight is None or weight.dim() != 2: + return + keep = _draft_vocab_size(self.config.vocab_size) + if keep == 0: + logger.info("MTP draft head: using the full vocabulary.") + return + if weight.dtype not in (torch.bfloat16, torch.float16, torch.float32): + # A quantised head would need its scales carried across the + # gather; not worth the risk, fall back to the dense path. + logger.info( + "MTP draft head: output embedding is %s, keeping the full " + "vocabulary.", + weight.dtype, + ) + return + tail_start = max(keep, weight.shape[0] - _DRAFT_VOCAB_TAIL) + ids = torch.cat( + [ + torch.arange(keep, device=weight.device), + torch.arange(tail_start, weight.shape[0], device=weight.device), + ] + ) + rows = weight.index_select(0, ids).contiguous() + self._draft_ids = ids.to(torch.int32) + + # e4m3 halves the bytes read per draft step again. Every row keeps its + # own scale, so the ordering the argmax sees is the ordering of the + # dequantised logits, not of the raw e4m3 products. + if _draft_head_in_fp8() and self._try_quantise_draft_head(rows): + bytes_kept = self._draft_w8.numel() + else: + self._draft_weight = rows + bytes_kept = rows.numel() * rows.element_size() + logger.info( + "MTP draft head: %d of %d vocabulary rows kept, %s, %.1f MiB.", + ids.numel(), + weight.shape[0], + "e4m3" if self._draft_w8 is not None else str(rows.dtype), + bytes_kept / 2**20, + ) + + def _try_quantise_draft_head(self, rows: torch.Tensor) -> bool: + """Store the narrowed draft head as e4m3 with one scale per row. + + Returns False (and leaves the bf16 path in place) whenever the build + does not offer the scaled GEMM, or the shape does not suit it. + """ + n_rows, hidden = rows.shape + if n_rows % 16 or hidden % 16: + return False + try: + from vllm import _custom_ops as ops + + quantised, scale = ops.scaled_fp8_quant( + rows, use_per_token_if_dynamic=True + ) + # b of the scaled GEMM is [hidden, n_rows]; the transpose of a + # contiguous [n_rows, hidden] already has that layout. + w8 = quantised.t() + w8_scale = scale.reshape(1, n_rows).to(torch.float32) + # A non-zero probe: an all-zero row would quantise to a zero + # scale and hide a real failure behind a division. + probe = torch.full( + (16, hidden), 0.01, dtype=rows.dtype, device=rows.device + ) + probe_q, probe_scale = ops.scaled_fp8_quant( + probe, use_per_token_if_dynamic=True + ) + ops.cutlass_scaled_mm( + probe_q, w8, probe_scale, w8_scale, torch.bfloat16 + ) + except Exception: + logger.info( + "MTP draft head: no usable e4m3 GEMM, staying on the dense " + "path.", + exc_info=True, + ) + return False + self._draft_w8 = w8 + self._draft_w8_scale = w8_scale + return True + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + w8 = getattr(self, "_draft_w8", None) + if w8 is not None: + from vllm import _custom_ops as ops + + flat = hidden_states.reshape(-1, hidden_states.shape[-1]).contiguous() + act, act_scale = ops.scaled_fp8_quant( + flat, use_per_token_if_dynamic=True + ) + logits = ops.cutlass_scaled_mm( + act, w8, act_scale, self._draft_w8_scale, torch.bfloat16 + ) + return self._draft_ids[logits.argmax(dim=-1)].to(torch.int64) + draft_weight = getattr(self, "_draft_weight", None) + if draft_weight is None: + return super().get_top_tokens(hidden_states) + logits = torch.nn.functional.linear( + hidden_states.to(draft_weight.dtype), draft_weight + ) + return self._draft_ids[logits.argmax(dim=-1)].to(torch.int64) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def remap_weight_names(weights): for name, weight in weights: @@ -457,7 +605,9 @@ class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): yield name, weight loader = AutoWeightsLoader(self) - return loader.load_weights(remap_weight_names(weights)) + loaded = loader.load_weights(remap_weight_names(weights)) + self._build_draft_head() + return loaded class Qwen3_5MoeMTP(Qwen3_5MTP, QwenNextMixtureOfExperts):