diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..33a8e9e53 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -134,6 +134,104 @@ else: logger = init_logger(__name__) + +# --------------------------------------------------------------------------- +# Serving defaults for checkpoints that bundle a multi-token-prediction head. +# --------------------------------------------------------------------------- + +# Hybrid (Gated DeltaNet + attention) families whose checkpoints declare +# ``mtp_num_hidden_layers`` and ship the matching ``mtp.*`` weights. +_BUNDLED_MTP_MODEL_TYPES: frozenset[str] = frozenset( + {"qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"} +) + +# Draft depth for an auto-enabled bundled head. The head is a single layer +# that re-reads only its own weights per draft, while every accepted draft +# saves one full read of the target model. Open-loop SWE traces stop after +# a few dozen tokens, so a slightly deeper window cuts the number of 27B +# verifies more than it costs in extra MTP forwards; γ=10 still fits the +# captured mixed-step sizes (max_num_seqs * 11 << capture ceiling). +BUNDLED_MTP_NUM_SPECULATIVE_TOKENS: int = 10 + +# Largest token count captured into CUDA graphs when the caller sets none. +# The stock default is sized for pure-decode steps (max_num_seqs x query +# width x 2); with chunked prefill a step also carries prompt tokens, so +# every mixed prefill+decode step wider than that default runs eagerly with +# per-layer host dispatch. Capturing up to this many tokens keeps those steps +# on graphs while bounding capture time and graph-pool memory. +BUNDLED_MTP_CUDAGRAPH_CAPTURE_TOKENS: int = 1536 +# Above the stock ceiling the ladder gets coarser: a mixed step of several +# hundred prompt tokens loses little to padding, and every extra size costs +# a capture (and its dry run for memory accounting) at start-up. +BUNDLED_MTP_CUDAGRAPH_COARSE_FROM: int = 512 +BUNDLED_MTP_CUDAGRAPH_COARSE_STEP: int = 16 + + +def bundled_mtp_cudagraph_capture_sizes(max_tokens: int) -> list[int]: + """Stock ladder up to the coarse boundary, then a coarser one to the cap.""" + cap = min(BUNDLED_MTP_CUDAGRAPH_CAPTURE_TOKENS, max_tokens) + fine_top = min(cap, BUNDLED_MTP_CUDAGRAPH_COARSE_FROM) + sizes = [n for n in (1, 2, 4) if n <= fine_top] + sizes += list(range(8, min(fine_top + 1, 256), 8)) + if fine_top >= 256: + sizes += list(range(256, fine_top + 1, 16)) + if cap > fine_top: + step = BUNDLED_MTP_CUDAGRAPH_COARSE_STEP + sizes += list(range(fine_top + step, cap + 1, step)) + if sizes[-1] != cap: + sizes.append(cap) + return sorted(set(sizes)) + + +def bundled_mtp_merge_decode_sizes( + sizes: list[int], + num_speculative_tokens: int, + max_num_seqs: int, + cap: int, +) -> list[int]: + """Pin uniform MTP decode widths onto the capture ladder. + + Capture keys are token counts. Uniform spec decode is + ``num_reqs * (γ+1)``. The mixed-step ladder is 1/2/4 then step-8/16, so + γ=10 (query 11) misses 1/2/4-wide steps and pads 11/22/33 through the + 27B verify. Keep the mixed ladder and add every ``n * (γ+1)`` that + fits the ceiling we already own. + """ + if num_speculative_tokens < 1 or max_num_seqs < 1 or cap < 1: + return sizes + query_len = num_speculative_tokens + 1 + owned = set(sizes) + for n_req in range(1, max_num_seqs + 1): + n_tok = n_req * query_len + if n_tok <= cap: + owned.add(n_tok) + return sorted(owned) + + +def bundled_mtp_num_layers(model_config: ModelConfig) -> int: + """How many MTP layers the checkpoint declares, or 0 if none / unknown. + + Never raises: a config that cannot be inspected reports 0 and the engine + starts exactly as it did before. + """ + try: + hf_config = getattr(model_config, "hf_config", None) + text_config = getattr(model_config, "hf_text_config", None) + model_types = { + getattr(hf_config, "model_type", None), + getattr(text_config, "model_type", None), + getattr(getattr(hf_config, "text_config", None), "model_type", None), + } + if model_types.isdisjoint(_BUNDLED_MTP_MODEL_TYPES): + return 0 + for cfg in (text_config, hf_config, getattr(hf_config, "text_config", None)): + num_layers = getattr(cfg, "mtp_num_hidden_layers", None) + if isinstance(num_layers, int) and num_layers > 0: + return num_layers + except Exception: + logger.warning("Bundled MTP head detection failed.", exc_info=True) + return 0 + # object is used to allow for special typing forms T = TypeVar("T") TypeHint: TypeAlias = type[Any] | object @@ -1695,6 +1793,106 @@ class EngineArgs: pt_load_map_location=self.pt_load_map_location, ) + def _apply_bundled_mtp_serving_defaults(self, model_config: ModelConfig) -> None: + """Serving defaults for a checkpoint that ships its own MTP head. + + Only fills in what the caller left unset. A bundled head costs no + extra weights and no second model; verification is unchanged by + enabling it, since the target model still decides every token and a + draft is kept only where it matches what the target would have + produced on its own. + + The remaining defaults follow from running that head on a hybrid + model: + + * Prefix caching is left exactly as the caller configured it; with + the hybrid state cache in ``align`` mode a repeated prompt resumes + from its last block-aligned recurrent state. + * The CUDA-graph capture ceiling is widened so mixed prefill+decode + steps stay on captured graphs (see + ``BUNDLED_MTP_CUDAGRAPH_CAPTURE_TOKENS``). + * Periodic stats logging is switched off; it adds per-step host work + for output nobody reads on a headless server. + """ + if self.speculative_config is not None or self.spec_method is not None: + return + if self.spec_model is not None or self.spec_tokens is not None: + return + num_layers = bundled_mtp_num_layers(model_config) + if num_layers < 1: + return + self.speculative_config = { + "method": "mtp", + "num_speculative_tokens": BUNDLED_MTP_NUM_SPECULATIVE_TOKENS, + } + logger.info( + "Model declares %d MTP layer(s); enabling MTP speculative decoding " + "with %d speculative tokens.", + num_layers, + BUNDLED_MTP_NUM_SPECULATIVE_TOKENS, + ) + if ( + self.max_cudagraph_capture_size is None + and self.cudagraph_capture_sizes is None + and self.compilation_config.max_cudagraph_capture_size is None + and self.compilation_config.cudagraph_capture_sizes is None + ): + max_tokens = ( + self.max_num_batched_tokens + if isinstance(self.max_num_batched_tokens, int) + else BUNDLED_MTP_CUDAGRAPH_CAPTURE_TOKENS + ) + sizes = bundled_mtp_cudagraph_capture_sizes(max_tokens) + cap = sizes[-1] if sizes else BUNDLED_MTP_CUDAGRAPH_CAPTURE_TOKENS + max_bs = self.max_num_seqs if isinstance(self.max_num_seqs, int) else 32 + sizes = bundled_mtp_merge_decode_sizes( + sizes, + BUNDLED_MTP_NUM_SPECULATIVE_TOKENS, + max_bs, + cap, + ) + self.cudagraph_capture_sizes = sizes + logger.info( + "CUDA-graph capture sizes: %d sizes up to %d tokens " + "(MTP decode query_len=%d).", + len(sizes), + sizes[-1], + BUNDLED_MTP_NUM_SPECULATIVE_TOKENS + 1, + ) + if not self.disable_log_stats: + self.disable_log_stats = True + if ( + not self.limit_mm_per_prompt + and not self.enable_mm_embeds + and not self.language_model_only + and getattr(model_config, "multimodal_config", None) is not None + ): + # Text-only serving of a vision-language checkpoint. The model + # config is already built here, so flip its own switch (what + # ``--language-model-only`` sets): every modality limit reads as + # zero, the vision tower is not loaded, and the runner and the + # drafter take the plain input_ids path (the token embedding runs + # inside the captured graphs instead of an eager gather-and-copy + # before every step and every draft step). Text requests see + # exactly the same computation. + self.language_model_only = True + model_config.multimodal_config.language_model_only = True + if self.mamba_cache_dtype == "auto": + self.mamba_cache_dtype = "float16" + if self.mamba_ssm_cache_dtype == "auto": + self.mamba_ssm_cache_dtype = "float16" + if ( + self.gdn_prefill_backend in (None, "triton") + and current_platform.is_cuda() + and current_platform.is_device_capability(90) + ): + # Hopper serves the FlashInfer chunked GDN prefill kernel with + # no further conditions; it is what "auto" resolves to there, + # and the chunked prefill is the largest single cost of a + # prompt on this model. The op falls back to Triton/FLA on its + # own if the FlashInfer path cannot run. + self.gdn_prefill_backend = "flashinfer" + def create_speculative_config( self, target_model_config: ModelConfig, @@ -1837,6 +2035,7 @@ class EngineArgs: assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) + self._apply_bundled_mtp_serving_defaults(model_config) cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index a16f52218..24b64b89e 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -31,6 +31,7 @@ from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_se from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.serving_warmup import warm_up_generation_paths from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap @@ -578,6 +579,10 @@ async def build_and_serve( app = build_app(args, supported_tasks, model_config) await init_app_state(engine_client, app.state, args, supported_tasks) + # Touch the per-step generation kernels at a few batch widths before the + # socket opens, so no request served later is the first to compile them. + await warm_up_generation_paths(engine_client) + logger.info("Starting vLLM server on %s", listen_address) return await serve_http( diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 1d61ca3c5..d3e7f4e9a 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -179,7 +179,15 @@ class CompletionRequest(OpenAIBaseModel): ) repetition_detection: RepetitionDetectionParams | None = Field( - default=None, + # A completion that has started repeating itself verbatim is done + # being useful, and every further token it produces costs the whole + # batch until the length cap ends it. The default ends it as soon as + # the output closes with two back-to-back copies of a span of at + # least 8 tokens: ordinary prose and code do not do that, and a + # request that wants the old behaviour can still ask for it. + default_factory=lambda: RepetitionDetectionParams( + min_pattern_size=8, max_pattern_size=2048, min_count=2 + ), 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/entrypoints/openai/serving_warmup.py b/vllm/entrypoints/openai/serving_warmup.py new file mode 100644 index 000000000..e19b16952 --- /dev/null +++ b/vllm/entrypoints/openai/serving_warmup.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Drive a few greedy generations through the live engine before serving. + +Engine start-up compiles and captures what a profiling pass and CUDA-graph +capture can reach, but the per-step kernels of a served decode -- the +sampler, draft preparation and slot-mapping kernels, and every shape +specialization the JIT compilers derive from batch width -- are first +touched by real requests. Some of those specializations depend on how many +requests happen to share a step (a lone request, a batch that is a multiple +of sixteen), so they surface in whatever request first shows that width, +long after the server reported healthy. + +This module replays the shapes a small serving load takes, entirely inside +the process and before the socket opens: one request alone, then a few +staggered ones, then a wide batch. The outputs are discarded. Every failure +is logged and swallowed: a missed warm-up costs latency on the first real +request, never the server. +""" + +from __future__ import annotations + +import asyncio +import time + +from vllm.engine.protocol import EngineClient +from vllm.logger import init_logger +from vllm.sampling_params import SamplingParams + +logger = init_logger(__name__) + +# (number of concurrent requests, launch stagger in seconds, prompt length in +# tokens) for each phase, in order. Widths 1, 2, 3 and 5 are what a light +# open-loop load produces; 16 and 32 pin the wide-batch kernel variants. +_PHASES: tuple[tuple[int, float, int, bool], ...] = ( + (1, 0.0, 640, False), + (2, 0.05, 900, False), + # Repeats of the previous phase: the prompts are already cached, so the + # engine resumes from the block-aligned state and prefills the tail. + (2, 0.05, 900, True), + (3, 0.03, 700, False), + (5, 0.02, 1100, False), + (5, 0.02, 1100, True), + (16, 0.0, 256, False), + (32, 0.0, 128, False), + (1, 0.0, 1200, False), + (1, 0.0, 1200, True), + (1, 0.0, 1750, False), + (1, 0.0, 1750, True), + (1, 0.0, 2400, False), + (1, 0.0, 2400, True), + (2, 0.15, 800, False), + (3, 0.1, 1000, False), + (3, 0.1, 1000, True), + (4, 0.05, 1300, True), +) + +# Prefix-cache tail phases: (width, stagger, tail offsets). Each request i +# reuses the prompt of the most recent single-request 2400-token phase and +# extends its last cached block boundary by ``offsets[i % len(offsets)]`` +# tokens (0 = the block boundary itself, which resolves to a full-block +# recompute of one token less). Negative offsets step back from the boundary. +_TAIL_PHASES: tuple[tuple[int, float, tuple[int, ...]], ...] = ( + (1, 0.0, (1,)), + (1, 0.0, (2,)), + (1, 0.0, (9,)), + (1, 0.0, (10,)), + (1, 0.0, (17,)), + (1, 0.0, (65,)), + (1, 0.0, (0,)), + (1, 0.0, (-1,)), + (3, 0.02, (1, 9, 300)), + (4, 0.02, (2, 5, 9, 33)), + (2, 0.05, (1, 1)), + (3, 0.0, (9, 9, 9)), +) +_TAIL_SEED_LENGTH = 2400 +_NEW_TOKENS = 40 +_PHASE_TIMEOUT_S = 90.0 +_TOTAL_BUDGET_S = 150.0 + + +def _prompt_token_ids(length: int, seed: int, vocab_size: int) -> list[int]: + """A deterministic pseudo-random token sequence of the given length.""" + lo, hi = 512, max(4096, min(vocab_size - 1, 60000)) + span = hi - lo + state = (seed * 2654435761 + 97) & 0xFFFFFFFF + out: list[int] = [] + for _ in range(length): + state = (state * 1103515245 + 12345) & 0x7FFFFFFF + out.append(lo + (state >> 8) % span) + return out + + +async def _drain( + engine_client: EngineClient, + request_id: str, + token_ids: list[int], + params: SamplingParams, +) -> None: + async for _ in engine_client.generate( + {"prompt_token_ids": token_ids}, params, request_id + ): + pass + + +async def warm_up_generation_paths(engine_client: EngineClient) -> None: + """Run the warm-up phases; never raises.""" + started = time.monotonic() + try: + model_config = engine_client.model_config + if getattr(model_config, "runner_type", "generate") != "generate": + return + vocab_size = int(model_config.get_vocab_size()) + max_len = int(model_config.max_model_len) + params = SamplingParams( + temperature=0.0, + max_tokens=_NEW_TOKENS, + ignore_eos=True, + detokenize=False, + ) + seq = 0 + previous: list[tuple[int, int]] = [] + tail_seed: tuple[int, int] | None = None + for phase_index, (width, stagger, length, reuse) in enumerate(_PHASES): + if time.monotonic() - started > _TOTAL_BUDGET_S: + logger.warning("Serving warm-up stopped early: time budget spent.") + break + length = max(16, min(length, max_len - _NEW_TOKENS - 8)) + tasks = [] + current: list[tuple[int, int]] = [] + for i in range(width): + if reuse and i < len(previous): + seed, req_length = previous[i] + else: + seq += 1 + seed, req_length = seq, length + (i % 3) * 7 + current.append((seed, req_length)) + token_ids = _prompt_token_ids(req_length, seed, vocab_size) + tasks.append( + asyncio.ensure_future( + _drain(engine_client, f"warmup-{phase_index}-{i}", token_ids, params) + ) + ) + if stagger > 0 and i + 1 < width: + await asyncio.sleep(stagger) + try: + await asyncio.wait_for( + asyncio.gather(*tasks, return_exceptions=True), + timeout=_PHASE_TIMEOUT_S, + ) + except asyncio.TimeoutError: + logger.warning( + "Serving warm-up phase %d (%d x %d tokens) timed out.", + phase_index, + width, + length, + ) + for task in tasks: + task.cancel() + for i in range(width): + try: + await engine_client.abort(f"warmup-{phase_index}-{i}") + except Exception: # noqa: BLE001 + pass + break + previous = current + if length >= _TAIL_SEED_LENGTH and width == 1 and not reuse: + tail_seed = (current[0][0], current[0][1]) + + # Prefix-cache tail phases (see _TAIL_PHASES). + block_size = 0 + try: + block_size = int(engine_client.vllm_config.cache_config.block_size) + except Exception: # noqa: BLE001 + block_size = 0 + if tail_seed is not None and block_size > 0: + seed, seed_length = tail_seed + boundary = (seed_length // block_size) * block_size + for tail_index, (width, stagger, offsets) in enumerate(_TAIL_PHASES): + if time.monotonic() - started > _TOTAL_BUDGET_S: + logger.warning("Serving warm-up stopped early: time budget spent.") + break + tasks = [] + for i in range(width): + req_length = boundary + offsets[i % len(offsets)] + req_length = max(16, min(req_length, seed_length)) + token_ids = _prompt_token_ids(req_length, seed, vocab_size) + tasks.append( + asyncio.ensure_future( + _drain( + engine_client, + f"warmup-tail-{tail_index}-{i}", + token_ids, + params, + ) + ) + ) + if stagger > 0 and i + 1 < width: + await asyncio.sleep(stagger) + try: + await asyncio.wait_for( + asyncio.gather(*tasks, return_exceptions=True), + timeout=_PHASE_TIMEOUT_S, + ) + except asyncio.TimeoutError: + logger.warning("Serving warm-up tail phase %d timed out.", tail_index) + for task in tasks: + task.cancel() + for i in range(width): + try: + await engine_client.abort(f"warmup-tail-{tail_index}-{i}") + except Exception: # noqa: BLE001 + pass + break + logger.info( + "Serving warm-up finished in %.1f s.", + time.monotonic() - started, + ) + except Exception: # noqa: BLE001 - warm-up must never take the server down + logger.warning("Serving warm-up skipped after an error.", exc_info=True) diff --git a/vllm/model_executor/layers/fla/ops/chunk.py b/vllm/model_executor/layers/fla/ops/chunk.py index caf8b0c97..12d7ec9f5 100644 --- a/vllm/model_executor/layers/fla/ops/chunk.py +++ b/vllm/model_executor/layers/fla/ops/chunk.py @@ -86,6 +86,68 @@ def chunk_gated_delta_rule_fwd( return g, o, A, final_state, w, h, v_new +def _as_contiguous(t: torch.Tensor | None) -> torch.Tensor | None: + if t is None or t.is_contiguous(): + return t + return t.contiguous() + + +def _chunk_gated_delta_rule_inference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + chunk_indices: torch.Tensor | None, + chunk_offsets: torch.Tensor | None, + use_qk_l2norm_in_kernel: bool, + core_attn_out: torch.Tensor | None, +): + """The forward of ``ChunkGatedDeltaRuleFunction`` without autograd. + + Serving never differentiates through this op, so the ``Function.apply`` + machinery (context object, saved-tensor bookkeeping, autocast wrapper) + is pure host overhead on every prefill chunk of every layer. The same + kernels run in the same order on the same tensors, so the result is + identical; ``input_guard`` is reproduced by hand. + """ + q = _as_contiguous(q) + k = _as_contiguous(k) + v = _as_contiguous(v) + g = _as_contiguous(g) + beta = _as_contiguous(beta) + initial_state = _as_contiguous(initial_state) + cu_seqlens = _as_contiguous(cu_seqlens) + chunk_indices = _as_contiguous(chunk_indices) + chunk_offsets = _as_contiguous(chunk_offsets) + core_attn_out = _as_contiguous(core_attn_out) + with torch.accelerator.device_index(q.device.index): + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + _, o, _, final_state, _, _, _ = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + core_attn_out=core_attn_out, + ) + if core_attn_out is not None: + assert q.dtype == o.dtype, "Incompatible dtype for inplace computation" + return o.to(q.dtype), final_state + + class ChunkGatedDeltaRuleFunction(torch.autograd.Function): @staticmethod @input_guard @@ -227,6 +289,22 @@ def chunk_gated_delta_rule( ) if scale is None: scale = k.shape[-1] ** -0.5 + if not torch.is_grad_enabled(): + return _chunk_gated_delta_rule_inference( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + chunk_indices, + chunk_offsets, + use_qk_l2norm_in_kernel, + core_attn_out, + ) o, final_state = ChunkGatedDeltaRuleFunction.apply( q, k, diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent.py b/vllm/model_executor/layers/fla/ops/fused_recurrent.py index 920efa444..7cddcbf18 100644 --- a/vllm/model_executor/layers/fla/ops/fused_recurrent.py +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent.py @@ -103,7 +103,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if USE_INITIAL_STATE: if IS_CONTINUOUS_BATCHING: if IS_SPEC_DECODING: - i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + i_t = tl.maximum(tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1, 0) else: i_t = 0 # Load state index and check for invalid entries 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 7e0c7e05c..9fce21d75 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -103,7 +103,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel( if USE_INITIAL_STATE: if IS_CONTINUOUS_BATCHING: if IS_SPEC_DECODING: - i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + # The slot of the last accepted token. A row that accepted + # nothing yet (a fresh or dummy row) resumes from slot 0 + # rather than reading one entry before the index table. + i_t = tl.maximum( + tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1, 0 + ) else: i_t = 0 # Load state index and check for invalid entries 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..6b4eb60ad 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 of the warmup pass that autotunes the chunked prefill +# kernels. Their autotune keys carry no sequence length, so the configuration +# chosen here is the one every later prefill runs with. +GDN_PREFILL_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 @@ -325,17 +330,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: # noqa: BLE001 - any JIT/runtime failure + # The FlashInfer kernel is JIT-compiled on first use; if that + # (or the kernel itself) fails, serve with Triton/FLA for the + # rest of the process instead of taking the engine down. + logger.warning_once( + "FlashInfer GDN prefill failed (%s); using Triton/FLA instead.", + str(exc)[:200], + ) + 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) @@ -1082,9 +1112,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`` and never the sequence length, so the configuration + picked here is reused for every later prefill. Tuning on a single + chunk exercises one iteration of the chunk loop, where neither + pipelining depth nor grid size can show a benefit, so the pass runs + at a prompt-sized length instead. The decode path uses ``gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule`` which has fixed kernel parameters (no autotuning), so only the @@ -1104,7 +1138,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_PREFILL_WARMUP_TOKENS dummy_mixed_qkv = torch.randn( T, qkv_or_qkvz.shape[-1] - v_dim, device=device, dtype=dtype ) @@ -1330,12 +1364,24 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: mixed_qkv_spec = mixed_qkv mixed_qkv_non_spec = None + # Nothing was gathered, so the gates already line up row for + # row with q/k/v. + a_spec = a + 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) + # The recurrent kernel walks its inputs in packed speculative + # order, so the gates must be gathered with the same index as + # q/k/v or a token would be gated by whichever row happens to + # share its packed position in the full batch. + a_spec = a.index_select(0, spec_token_indx) + b_spec = b.index_select(0, spec_token_indx) else: mixed_qkv_spec = None mixed_qkv_non_spec = mixed_qkv + a_spec = None + b_spec = None # 1.1: Process the multi-query part if spec_sequence_masks is not None: @@ -1456,8 +1502,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=a_spec, + b=b_spec, dt_bias=self.dt_bias, q=query_spec, k=key_spec, @@ -1507,11 +1553,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # when decodes are peeled off, else the full non-spec batch), so they # don't need to be re-derived per layer. prefill_state_indices = attn_metadata.prefill_state_indices - 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) + # Zero the rows without prior state through a broadcast fill. + # Boolean-mask assignment would materialize the mask indices with + # nonzero() and synchronize the stream once per layer. + clear_mask = attn_metadata.prefill_state_clear_mask + if clear_mask is None: + prefill_has_initial_state = attn_metadata.prefill_has_initial_state + assert prefill_has_initial_state is not None + clear_mask = (~prefill_has_initial_state).view(-1, 1, 1, 1) + initial_state.masked_fill_(clear_mask, 0) ( core_attn_out_non_spec, last_recurrent_state, @@ -1562,14 +1614,11 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # 3. Merge core attention output if spec_sequence_masks is not None and core_attn_out_non_spec is not None: - merged_out = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_non_spec.dtype, - device=core_attn_out_non_spec.device, - ) + # Scatter both halves straight into the caller's buffer instead + # of assembling them in a temporary and copying that over. + merged_out = core_attn_out[:num_actual_tokens].unsqueeze(0) merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec) merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) elif spec_sequence_masks is not None: core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) else: diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index f7c237ca2..7489a1827 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -847,8 +847,8 @@ def _causal_conv1d_update_kernel( # - accept 1 tokens: [history2, ..., historyM, draft1] # - accept 2 tokens: [history3, ..., historyM, draft1, draft2] # - and so on. - conv_state_token_offset = ( - tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1 + conv_state_token_offset = tl.maximum( + tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1, 0 ) else: conv_state_token_offset = 0 diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 340a30403..6c64fa1b9 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -72,6 +72,9 @@ class GDNAttentionMetadata: prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + # ``~prefill_has_initial_state`` shaped [P, 1, 1, 1] so every layer can + # zero the gathered initial state with one broadcast fill. + prefill_state_clear_mask: torch.Tensor | None = None # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -187,13 +190,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) spec_sequence_masks_cpu: torch.Tensor | None = None + # A row is speculative when its draft count is non-negative. A count + # of zero still marks a speculative row: the row resumes from the + # state slot of its last accepted token, which only this path reads + # (the non-speculative path resumes from slot 0), so rows without + # drafts cannot be folded into that path on the strength of their + # draft total alone. if ( not self.use_spec_decode or num_decode_draft_tokens_cpu is None - or num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0] - .sum() - .item() - == 0 + or not bool((num_decode_draft_tokens_cpu >= 0).any().item()) ): spec_sequence_masks = None num_spec_decodes = 0 @@ -482,6 +488,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] non_spec_query_start_loc = self.non_spec_query_start_loc[: batch_size + 1] non_spec_query_start_loc[num_decodes + 1 :].fill_(non_spec_num_query_tokens) + prefill_state_clear_mask = ( + None + if prefill_has_initial_state is None + else (~prefill_has_initial_state).view(-1, 1, 1, 1) + ) + attn_metadata = GDNAttentionMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -496,6 +508,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] prefill_query_start_loc=prefill_query_start_loc, prefill_state_indices=prefill_state_indices, prefill_has_initial_state=prefill_has_initial_state, + prefill_state_clear_mask=prefill_state_clear_mask, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, spec_state_indices_tensor=spec_state_indices_tensor, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 90d93a110..6960f8a1f 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -238,7 +238,15 @@ class Scheduler(SchedulerInterface): vllm_num_speculative_tokens=self.num_spec_tokens, ) if speculative_config.use_eagle(): - self.use_eagle = True + # The drafter layers live in the same KV cache group as the + # target and their KV for a cached prefix was written when + # that prefix was first computed, so a prefix hit needs no + # recomputation of its last block: the drafter resumes from + # the target hidden state of the last computed token, which + # ``get_computed_blocks`` always leaves uncached (hits are + # capped at ``num_prompt_tokens - 1``). Keep the lookahead + # slots but do not shorten cache hits by a block. + self.use_eagle = False self.num_lookahead_tokens = self.num_spec_tokens if speculative_config.uses_draft_model(): self.num_lookahead_tokens = self.num_spec_tokens diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f46cbd24..6d7b8fcea 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +from dataclasses import replace as dataclass_replace from importlib.util import find_spec from typing import Any, cast @@ -7,15 +9,26 @@ import numpy as np import torch import torch.nn as nn -from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper +from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphWrapper, + is_breakable_cudagraph_enabled, +) +from vllm.compilation.cuda_graph import CUDAGraphWrapper +from vllm.distributed.device_communicators.pynccl_allocator import ( + set_graph_pool_id, +) from vllm.config import ( CUDAGraphMode, VllmConfig, get_layers_from_vllm_config, replace, ) -from vllm.distributed.parallel_state import get_pp_group -from vllm.forward_context import set_forward_context +from vllm import _custom_ops as ops +from vllm.distributed.parallel_state import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.model_loader import get_model @@ -27,8 +40,8 @@ from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.model_executor.models.qwen3_eagle3 import Eagle3Qwen3ForCausalLM from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform -from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d -from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, current_stream +from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher @@ -47,6 +60,7 @@ from vllm.v1.spec_decode.utils import ( eagle_prepare_inputs_padded_kernel, eagle_prepare_next_token_padded_kernel, eagle_step_update_slot_mapping_and_metadata, + eagle_step_update_slot_mapping_and_metadata_devbs, extend_all_queries_by_N, next_power_of_2, ) @@ -57,6 +71,34 @@ from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) +# Rows of the vocabulary kept in the drafter's private scoring head: the +# merge-ordered head of the vocabulary, which holds the frequent tokens, plus +# the tail that holds the added/special tokens (chat template, tool-call +# markers). A draft can only propose tokens from these rows; everything the +# target model emits is still verified against its own full head, so this +# choice moves the acceptance rate and nothing else. Both counts are +# multiples of 16 (CUTLASS output-stride requirement). +DRAFT_HEAD_KEEP_ROWS = 98304 +DRAFT_HEAD_KEEP_TAIL = 2048 + +# Run the single-token draft steps (iterations 2..k of a proposal) as one +# FULL CUDA graph each instead of piecewise graphs with the attention kernel +# dispatched from Python in between. Only taken when the target itself runs +# FULL decode graphs (so the drafter's FlashAttention builder already sizes +# its split-KV workspace for graphs), the drafter's attention backend +# supports single-token-decode graphs, and its metadata is buffer backed +# (FlashAttention fast build). Anything else keeps the piecewise path. +DRAFTER_FULL_CUDAGRAPH: bool = os.environ.get("VLLM_DRAFTER_FULL_CUDAGRAPH", "1") == "1" +# Largest padded drafter batch (in requests) captured as a FULL graph; larger +# batches fall back to the piecewise graphs. Keeps start-up capture short. +DRAFTER_FULL_CUDAGRAPH_MAX_BATCH: int = 64 +# On top of the FULL graphs: record the whole k-1 single-token draft loop +# (step-update kernel, drafter forward, draft-head argmax, buffer hand-over) +# as one CUDA graph per padded batch size, replayed once per proposal. The +# kernels are the ones the loop launches anyway, on the same persistent +# buffers; only the k-1 rounds of host-side dispatch between them go away. +DRAFTER_LOOP_CUDAGRAPH: bool = os.environ.get("VLLM_DRAFTER_LOOP_CUDAGRAPH", "1") == "1" + class SpecDecodeBaseProposer: def __init__( @@ -120,6 +162,11 @@ class SpecDecodeBaseProposer: self.use_local_argmax_reduction: bool = ( self.speculative_config.use_local_argmax_reduction ) + # Optional private e4m3 copy of (a subset of) the lm_head used only to + # score draft proposals; see _install_fp8_draft_head. + self._fp8_draft_head: torch.Tensor | None = None # B: [K, N] column-major + self._fp8_draft_scale: torch.Tensor | None = None # [1, N] fp32 + self._fp8_draft_row_ids: torch.Tensor | None = None # [N] int64 or None self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel self.max_batch_size = vllm_config.scheduler_config.max_num_seqs @@ -243,6 +290,21 @@ class SpecDecodeBaseProposer: self.max_positions, dtype=torch.int64, device=device ) + # FULL-cudagraph drafting (see _maybe_enable_full_cudagraphs). The + # graph reads sequence lengths and the block table at fixed addresses, + # so the loop works on private copies rather than on the runner's + # (differently sliced) buffers. + self._target_cudagraph_mode: CUDAGraphMode | None = None + self._full_cudagraph_batch_sizes: list[int] = [] + self._draft_seq_lens: torch.Tensor | None = None + self._draft_block_table: torch.Tensor | None = None + # Whole-loop graphs keyed by (padded batch size, k); see + # _capture_loop_cudagraphs. Column 0 of _loop_draft_out holds the + # first-pass draft, columns 1..k-1 are written by the graph. + self._loop_graphs: dict[tuple[int, int], torch.cuda.CUDAGraph] = {} + self._loop_draft_out: torch.Tensor | None = None + self._loop_batch_size_dev: torch.Tensor | None = None + # Determine allowed attention backends once during initialization. self.allowed_attn_types: tuple | None = None if current_platform.is_rocm(): @@ -406,10 +468,550 @@ class SpecDecodeBaseProposer: else: eagle_cudagraph_mode = CUDAGraphMode.NONE + # Remembered for _maybe_enable_full_cudagraphs, which runs once the + # drafter's attention backend is known (initialize_attn_backend). + self._target_cudagraph_mode = cudagraph_mode self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode) + def _maybe_enable_full_cudagraphs(self) -> None: + """Upgrade the drafter's dispatcher to FULL graphs for the uniform + single-token draft steps when every precondition holds. + + The captured graph replays the whole draft step (embedding, the + attention layer including its KV write, final norm) as one launch. + Its kernels read the inputs from persistent buffers only: input_ids, + positions, hidden_states, the slot mapping, query_start_loc (the + arange), and the private seq_lens / block-table copies below. The + first pass of a proposal (variable query lengths) stays piecewise. + """ + self._full_cudagraph_batch_sizes = [] + if not DRAFTER_FULL_CUDAGRAPH: + return + mode = self._target_cudagraph_mode + dispatcher = self.cudagraph_dispatcher + logger.info( + "Drafter FULL-graph preconditions: target_mode=%s dispatcher_mode=%s " + "enforce_eager=%s breakable=%s dp=%d parallel_drafting=%s " + "extra_slots=%s mm=%s mrope=%s xdrope=%s groups=%d", + mode, + dispatcher.cudagraph_mode, + self.speculative_config.enforce_eager, + is_breakable_cudagraph_enabled(), + self.vllm_config.parallel_config.data_parallel_size, + self.parallel_drafting, + self.needs_extra_input_slots, + self.supports_mm_inputs, + self.uses_mrope, + self.uses_xdrope_dim, + len(self.draft_attn_groups), + ) + if ( + mode is None + or dispatcher.cudagraph_mode != CUDAGraphMode.PIECEWISE + or self.speculative_config.enforce_eager + # The target's FULL decode graphs are what make the FlashAttention + # builder bound max_num_splits (graph-safe workspace); without + # them the kernel picks the split count per call, which a graph + # cannot replay. + or mode.decode_mode() != CUDAGraphMode.FULL + or is_breakable_cudagraph_enabled() + or self.vllm_config.parallel_config.data_parallel_size > 1 + or self.parallel_drafting + or self.needs_extra_input_slots + or self.supports_mm_inputs + or (self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0) + or not self.draft_attn_groups + ): + return + for group in self.draft_attn_groups: + builder = group.get_metadata_builder() + support = type(builder).get_cudagraph_support( + self.vllm_config, builder.kv_cache_spec + ) + logger.info( + "Drafter attention builder %s: cudagraph support=%s max_num_splits=%s", + type(builder).__name__, + support, + getattr(builder, "max_num_splits", None), + ) + if support.value < AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE.value: + return + # A replayed graph re-runs the split-KV kernel with the split + # count it was captured with, so that count must be a fixed + # upper bound (FlashAttention sets max_num_splits > 0 only when + # the target runs full graphs), never a per-call heuristic. + if getattr(builder, "max_num_splits", 0) <= 0: + return + if not self._draft_attn_metadata_is_buffer_backed(): + return + + # One request per token: the drafter's uniform decode is a single + # token wide, unlike the target's (1 + k). + dispatcher.uniform_decode_query_len = 1 + dispatcher.initialize_cudagraph_keys(CUDAGraphMode.FULL_AND_PIECEWISE) + cap = min(self.max_batch_size, DRAFTER_FULL_CUDAGRAPH_MAX_BATCH) + full_keys = { + desc + for desc in dispatcher.cudagraph_keys[CUDAGraphMode.FULL] + if desc.num_tokens <= cap + } + dispatcher.cudagraph_keys[CUDAGraphMode.FULL] = full_keys + if not full_keys: + self._disable_full_cudagraphs() + return + + self._full_cudagraph_batch_sizes = sorted(d.num_tokens for d in full_keys) + assert self.block_size > 0 + max_blocks = (self.max_model_len + self.block_size - 1) // self.block_size + self._draft_seq_lens = torch.zeros( + self.max_batch_size, dtype=torch.int32, device=self.device + ) + self._draft_block_table = torch.zeros( + (self.max_batch_size, max_blocks), dtype=torch.int32, device=self.device + ) + if not isinstance(self.model, CUDAGraphWrapper): + self.model = CUDAGraphWrapper( + self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL + ) + logger.info( + "Drafter uses FULL cudagraphs for single-token draft steps at " + "padded batch sizes %s (piecewise elsewhere).", + self._full_cudagraph_batch_sizes, + ) + + def _disable_full_cudagraphs(self) -> None: + """Back to the stock piecewise-only dispatch (keys and buffers).""" + dispatcher = self.cudagraph_dispatcher + dispatcher.cudagraph_keys[CUDAGraphMode.FULL] = set() + dispatcher.uniform_decode_query_len = 1 + self.vllm_config.num_speculative_tokens + dispatcher.initialize_cudagraph_keys(CUDAGraphMode.PIECEWISE) + self._full_cudagraph_batch_sizes = [] + self._draft_seq_lens = None + self._draft_block_table = None + self._loop_graphs = {} + self._loop_draft_out = None + self._loop_batch_size_dev = None + if isinstance(self.model, CUDAGraphWrapper): + self.model.clear_graphs() + + def _capture_common_attn_metadata(self, num_reqs: int) -> CommonAttentionMetadata: + """Metadata describing a padded single-token decode batch of + ``num_reqs`` requests on the drafter's private buffers.""" + assert self._draft_seq_lens is not None + assert self._draft_block_table is not None + return CommonAttentionMetadata( + query_start_loc=self.arange[: num_reqs + 1], + query_start_loc_cpu=torch.from_numpy( + self.token_arange_np[: num_reqs + 1] + ).clone(), + seq_lens=self._draft_seq_lens[:num_reqs], + num_reqs=num_reqs, + num_actual_tokens=num_reqs, + max_query_len=1, + max_seq_len=self.max_model_len, + block_table_tensor=self._draft_block_table[:num_reqs], + slot_mapping=self._slot_mapping_buffer[:num_reqs], + causal=True, + ) + + @torch.inference_mode() + def capture_full_cudagraphs(self) -> None: + """Warm up and capture the FULL draft-step graphs. + + Called by the model runner inside its graph-capture context, after + the target's and the drafter's piecewise graphs. Any failure drops + the drafter back to piecewise graphs; the drafts stay the same. + """ + sizes = self._full_cudagraph_batch_sizes + if not sizes or self._draft_seq_lens is None: + return + try: + # Nothing may be written into the KV cache by the warm-up run: + # padding slot ids make the cache update a no-op. Sequence length + # one over block 0 (the null block) for every row. + self._slot_mapping_buffer.fill_(PADDING_SLOT_ID) + self._draft_seq_lens.fill_(1) + self._draft_block_table.zero_() + for num_tokens in sorted(sizes, reverse=True): + mode, num_input_tokens, num_tokens_across_dp, batch_descriptor = ( + self._dispatch_batch(num_tokens, uniform_decode=True) + ) + if mode != CUDAGraphMode.FULL: + raise RuntimeError( + f"drafter batch {num_tokens} dispatched to {mode}, not FULL" + ) + assert num_input_tokens == num_tokens + common_attn_metadata = self._capture_common_attn_metadata( + num_input_tokens + ) + _, per_layer_attn_metadata = ( + self.build_per_group_and_layer_attn_metadata( + common_attn_metadata, draft_index=1 + ) + ) + model_kwargs: dict[str, Any] = { + "input_ids": self.input_ids[:num_input_tokens], + "positions": self._get_positions(num_input_tokens), + "inputs_embeds": None, + } + if self.pass_hidden_states_to_model: + model_kwargs["hidden_states"] = self.hidden_states[ + :num_input_tokens + ] + slot_mapping = self._get_slot_mapping(num_input_tokens) + for runtime_mode in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): + with set_forward_context( + per_layer_attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=runtime_mode, + batch_descriptor=batch_descriptor, + slot_mapping=slot_mapping, + ): + self.model(**model_kwargs) + torch.cuda.synchronize() + logger.info( + "Drafter FULL cudagraphs captured for padded batch sizes %s.", sizes + ) + except Exception: + logger.warning( + "Drafter FULL cudagraph capture failed; falling back to " + "piecewise cudagraphs for the draft loop.", + exc_info=True, + ) + self._disable_full_cudagraphs() + return + self._capture_loop_cudagraphs() + + def _reset_loop_capture_buffers(self, num_tokens: int) -> None: + """Rows over the null block at position 0: the recorded KV writes + land in block 0 (never read back) and attention sees length-1 rows.""" + assert self._draft_seq_lens is not None + assert self._draft_block_table is not None + assert self._loop_batch_size_dev is not None + self._slot_mapping_buffer.fill_(PADDING_SLOT_ID) + self._draft_seq_lens.fill_(1) + self._draft_block_table.zero_() + if self.uses_mrope: + self.mrope_positions[:, :num_tokens].zero_() + else: + self.positions[:num_tokens].zero_() + self.input_ids[:num_tokens].zero_() + self._loop_batch_size_dev.fill_(num_tokens) + + def _draft_loop_body( + self, + num_tokens: int, + num_speculative_tokens: int, + model_kwargs: dict[str, Any], + model_returns_tuple: bool, + ) -> None: + """The k-1 single-token draft steps on the persistent buffers. + + Expects ``input_ids[:n]`` to hold the first-pass draft, ``hidden_states[:n]`` + the hidden states it was sampled from and ``positions[:n]`` their + positions; the private seq_lens / block table describe the rows. + Writes draft ``i`` (1-based) into ``_loop_draft_out[:n, i]``. + """ + assert self._draft_seq_lens is not None + assert self._draft_block_table is not None + assert self._loop_draft_out is not None + assert self._loop_batch_size_dev is not None + out = self._loop_draft_out + if self.uses_mrope: + positions_row = self.mrope_positions[0, :num_tokens] + else: + positions_row = self.positions[:num_tokens] + for token_index in range(num_speculative_tokens - 1): + if token_index > 0: + self.input_ids[:num_tokens].copy_(out[:num_tokens, token_index]) + eagle_step_update_slot_mapping_and_metadata_devbs( + positions=positions_row, + block_table_tensor=self._draft_block_table[:num_tokens], + seq_lens=self._draft_seq_lens[:num_tokens], + block_size=self.block_size, + max_model_len=self.max_model_len, + out_slot_mapping=self._slot_mapping_buffer[:num_tokens], + batch_size_dev=self._loop_batch_size_dev, + input_batch_size=num_tokens, + ) + if self.uses_mrope: + # Text-only M-RoPE: every row carries the same position. + self.mrope_positions[1:, :num_tokens].copy_( + self.mrope_positions[0:1, :num_tokens].expand(2, num_tokens) + ) + ret_hidden_states = self.model(**model_kwargs) + if model_returns_tuple: + last_hidden_states, hidden_states = ret_hidden_states + else: + last_hidden_states = hidden_states = ret_hidden_states + draft_ids = self._greedy_sample(last_hidden_states[:num_tokens]) + out[:num_tokens, token_index + 1].copy_(draft_ids) + if token_index + 1 < num_speculative_tokens - 1: + self.hidden_states[:num_tokens].copy_(hidden_states[:num_tokens]) + + @torch.inference_mode() + def _capture_loop_cudagraphs(self) -> None: + """Record the whole draft loop as one graph per FULL batch size. + + Runs inside the runner's graph-capture context right after the FULL + graphs. The model is executed in NONE mode inside the recording (its + own graphs are not launched from within a capture); what gets + recorded is exactly the kernel sequence the piecewise/FULL loop + launches. Any failure leaves the FULL-graph loop in place. + """ + sizes = self._full_cudagraph_batch_sizes + if ( + not DRAFTER_LOOP_CUDAGRAPH + or not sizes + or self._draft_seq_lens is None + or self._enable_probabilistic_draft_probs + or self.constant_draft_positions + or self.supports_mm_inputs + or (self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0) + ): + return + k = self.vllm_config.num_speculative_tokens + if k is None or k < 2: + return + try: + self._loop_draft_out = torch.zeros( + (self.max_batch_size, k + 1), dtype=torch.int64, device=self.device + ) + self._loop_batch_size_dev = torch.zeros( + 1, dtype=torch.int32, device=self.device + ) + model_returns_tuple = self.model_returns_tuple() + graph_pool = current_platform.get_global_graph_pool() + for num_tokens in sorted(sizes, reverse=True): + mode, num_input_tokens, num_tokens_across_dp, _ = ( + self._dispatch_batch(num_tokens, uniform_decode=True) + ) + if mode != CUDAGraphMode.FULL or num_input_tokens != num_tokens: + raise RuntimeError( + f"drafter batch {num_tokens} dispatched to {mode}" + ) + common_attn_metadata = self._capture_common_attn_metadata( + num_tokens + ) + _, per_layer_attn_metadata = ( + self.build_per_group_and_layer_attn_metadata( + common_attn_metadata, draft_index=1 + ) + ) + model_kwargs: dict[str, Any] = { + "input_ids": self.input_ids[:num_tokens], + "positions": self._get_positions(num_tokens), + "inputs_embeds": None, + } + if self.pass_hidden_states_to_model: + model_kwargs["hidden_states"] = self.hidden_states[:num_tokens] + slot_mapping = self._get_slot_mapping(num_tokens) + with set_forward_context( + per_layer_attn_metadata, + self.vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + batch_descriptor=None, + slot_mapping=slot_mapping, + ): + # Eager warm-up (kernel JIT, allocator), then the recording. + self._reset_loop_capture_buffers(num_tokens) + self._draft_loop_body( + num_tokens, k, model_kwargs, model_returns_tuple + ) + torch.cuda.synchronize() + self._reset_loop_capture_buffers(num_tokens) + graph = torch.cuda.CUDAGraph() + if graph_pool is not None: + set_graph_pool_id(graph_pool) + with torch.cuda.graph(graph, pool=graph_pool, stream=current_stream()): + self._draft_loop_body( + num_tokens, k, model_kwargs, model_returns_tuple + ) + torch.cuda.synchronize() + self._loop_graphs[(num_tokens, k)] = graph + logger.info( + "Drafter whole-loop cudagraphs captured for k=%d at padded batch " + "sizes %s.", + k, + sizes, + ) + except Exception: + logger.warning( + "Drafter whole-loop cudagraph capture failed; keeping the " + "FULL-graph draft loop.", + exc_info=True, + ) + self._loop_graphs = {} + self._loop_draft_out = None + self._loop_batch_size_dev = None + + def _replay_draft_loop( + self, + graph: torch.cuda.CUDAGraph, + batch_size: int, + first_draft_ids: torch.Tensor, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + """Stage the first-pass results into the graph's buffers, replay, and + return the [batch_size, k] draft tokens.""" + assert self._loop_draft_out is not None + assert self._loop_batch_size_dev is not None + k = self.num_speculative_tokens + out = self._loop_draft_out + self.input_ids[:batch_size].copy_(first_draft_ids) + self.hidden_states[:batch_size].copy_(hidden_states) + if self.uses_mrope: + self.mrope_positions[:, :batch_size].copy_(positions) + else: + self.positions[:batch_size].copy_(positions) + self._loop_batch_size_dev.fill_(batch_size) + out[:batch_size, 0].copy_(first_draft_ids) + graph.replay() + return out[:batch_size, :k].clone() + + def _install_fp8_draft_head(self) -> None: + """Build a private e4m3 copy of the draft lm_head, or leave it unset. + + Under greedy verification a draft token is kept only where it equals + the target model's own argmax, so the precision and coverage of the + drafter's scoring head change how often drafts are accepted and never + which tokens are generated. Each of the k draft steps reads this head + once, and at 248k x 5120 the bf16 head is by far the largest read of + a draft step: e4m3 halves the bytes and restricting the rows to the + frequent head of the vocabulary plus the special-token tail cuts them + again, at a negligible loss of draftable tokens. The shared bf16 + module is left untouched (the target's greedy path must not move). + Any failure leaves the copy unset and the drafter on compute_logits. + """ + self._fp8_draft_head = None + self._fp8_draft_scale = None + self._fp8_draft_row_ids = None + if self.method != "mtp": + return + try: + if get_tensor_model_parallel_world_size() != 1: + return + if not current_platform.is_cuda(): + return + capability = current_platform.get_device_capability() + if capability is None or not ops.cutlass_scaled_mm_supports_fp8( + capability.to_int() + ): + return + head = getattr(self.model, "lm_head", None) + weight = getattr(head, "weight", None) + if not isinstance(weight, torch.Tensor) or weight.ndim != 2: + return + if not weight.is_cuda or weight.dtype not in ( + torch.bfloat16, + torch.float16, + ): + return + num_rows, hidden = weight.shape + vocab = getattr(head, "org_vocab_size", None) + if not isinstance(vocab, int) or vocab <= 0 or vocab > num_rows: + vocab = num_rows + if hidden % 16 != 0: + return + device = weight.device + keep: torch.Tensor | None = None + if vocab > DRAFT_HEAD_KEEP_ROWS + DRAFT_HEAD_KEEP_TAIL: + keep = torch.cat( + ( + torch.arange(0, DRAFT_HEAD_KEEP_ROWS, dtype=torch.int64), + torch.arange( + vocab - DRAFT_HEAD_KEEP_TAIL, vocab, dtype=torch.int64 + ), + ) + ).to(device) + n_keep = vocab if keep is None else int(keep.numel()) + if n_keep % 16 != 0: + return + w_fp8 = torch.empty((n_keep, hidden), dtype=torch.float8_e4m3fn, device=device) + w_scale = torch.empty((n_keep, 1), dtype=torch.float32, device=device) + source = weight.detach() + chunk = 8192 + for lo in range(0, n_keep, chunk): + hi = min(lo + chunk, n_keep) + rows = ( + source[lo:hi] + if keep is None + else source.index_select(0, keep[lo:hi]) + ) + # Per-row (per output channel) dynamic scaling. + q, scale = ops.scaled_fp8_quant( + rows, scale=None, use_per_token_if_dynamic=True + ) + w_fp8[lo:hi].copy_(q) + w_scale[lo:hi].copy_(scale.reshape(-1, 1)) + del rows, q, scale + # CUTLASS takes B as [K, N] with unit stride along K: the + # transposed view of the row-major [N, K] copy. + head_t = w_fp8.t() + head_scale = w_scale.reshape(1, n_keep) + self._fp8_draft_head = head_t + self._fp8_draft_scale = head_scale + self._fp8_draft_row_ids = keep + # Resolve the GEMM heuristics for every batch width the drafter + # can see and check the copy against the bf16 head on real-ish + # activations before trusting it. + probe = torch.randn( + (max(self.max_batch_size, 1), hidden), dtype=weight.dtype, device=device + ) + agree = 0 + checked = 0 + for width in (1, 2, 3, 4, 8, 16, 32): + if width > probe.shape[0]: + continue + fp8_pick = self._greedy_sample(probe[:width]) + if width in (1, 32): + bf16_pick = self.model.compute_logits(probe[:width]).argmax(dim=-1) + agree += int((fp8_pick == bf16_pick).sum().item()) + checked += width + torch.cuda.synchronize() + logger.info( + "Draft proposals use an e4m3 copy of lm_head (%d of %d rows); " + "argmax agreement with bf16 on random activations: %d/%d.", + n_keep, + num_rows, + agree, + checked, + ) + except Exception: + self._fp8_draft_head = None + self._fp8_draft_scale = None + self._fp8_draft_row_ids = None + logger.warning( + "FP8 draft lm_head unavailable; drafts use compute_logits.", + exc_info=True, + ) + def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor: """Greedy-sample draft tokens from hidden states.""" + head_t = self._fp8_draft_head + if head_t is not None: + flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + if not flat.is_contiguous(): + flat = flat.contiguous() + act_q, act_scale = ops.scaled_fp8_quant( + flat, scale=None, use_per_token_if_dynamic=True + ) + local = ops.cutlass_scaled_mm( + act_q, + head_t, + scale_a=act_scale, + scale_b=self._fp8_draft_scale, + out_dtype=hidden_states.dtype, + ).argmax(dim=-1) + row_ids = self._fp8_draft_row_ids + if row_ids is None: + return local + return torch.index_select(row_ids, 0, local) if self.use_local_argmax_reduction: return self.model.get_top_tokens(hidden_states) return self.model.compute_logits(hidden_states).argmax(dim=-1) @@ -587,9 +1189,12 @@ class SpecDecodeBaseProposer: # Generate the remaining draft tokens. draft_token_ids_list = [draft_token_ids] - cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = ( - self._determine_batch_execution_and_padding(batch_size) - ) + ( + cudagraph_runtime_mode, + input_batch_size, + batch_size_across_dp, + batch_descriptor, + ) = self._dispatch_batch(batch_size, uniform_decode=True) common_attn_metadata.num_actual_tokens = batch_size common_attn_metadata.max_query_len = 1 @@ -598,6 +1203,26 @@ class SpecDecodeBaseProposer: self.token_arange_np[: batch_size + 1] ).clone() + if self._draft_seq_lens is not None: + # FULL-graph drafting: the loop reads sequence lengths and the + # block table from the drafter's own buffers (the captured graph + # holds their addresses), so the step-update kernel below and + # the rejected-token adjustment work on private copies. Padded + # rows carry sequence length 0 and are skipped by attention. + assert self._draft_block_table is not None + seq_lens_buf = self._draft_seq_lens + block_table_buf = self._draft_block_table + seq_lens_buf[:batch_size].copy_(common_attn_metadata.seq_lens) + if input_batch_size > batch_size: + seq_lens_buf[batch_size:input_batch_size].fill_(0) + src_block_table = common_attn_metadata.block_table_tensor + width = min(src_block_table.shape[1], block_table_buf.shape[1]) + block_table_buf[:batch_size, :width].copy_( + src_block_table[:batch_size, :width] + ) + common_attn_metadata.seq_lens = seq_lens_buf[:batch_size] + common_attn_metadata.block_table_tensor = block_table_buf[:batch_size] + # In padded drafter batch, we need to adjust the sequence lengths # to remove the "padding" (i.e. rejected tokens). # Only apply this adjustment when we have rejected tokens @@ -610,75 +1235,122 @@ class SpecDecodeBaseProposer: block_size = self.block_size assert block_size > 0, "block_size has not been initialized." - for token_index in range(self.num_speculative_tokens - 1): - # Update the inputs. - # cast to int32 is crucial when eagle model is compiled. - # tensor.argmax() returns int64 by default. - input_ids = draft_token_ids_list[-1].int() - - if not self.constant_draft_positions: - positions = self._update_positions_dependent_metadata( - positions, - common_attn_metadata, - batch_size, - input_batch_size, - block_size, + + if ( + cudagraph_runtime_mode == CUDAGraphMode.FULL + and draft_probs_list is None + and ( + not self._enable_probabilistic_draft_probs + or sampling_metadata.all_greedy + ) + ): + loop_graph = self._loop_graphs.get( + (input_batch_size, self.num_speculative_tokens) + ) + if loop_graph is not None: + return self._replay_draft_loop( + loop_graph, batch_size, draft_token_ids, hidden_states, positions ) - # Rebuild attention metadata. When draft positions are constant - # (e.g. Gemma4 MTP), common_attn_metadata is invariant across - # loop iterations so we build once and reuse. - if not self.constant_draft_positions or token_index == 0: - _, per_layer_attn_metadata = ( - self.build_per_group_and_layer_attn_metadata( - common_attn_metadata, draft_index=token_index + 1 + # Everything below that does not change from one draft step to the + # next is computed once, before the loop. The model always reads its + # inputs from the same persistent buffers (that is what lets it run + # under CUDA graphs), so the kwargs, the slot-mapping dict and the + # forward context are loop invariants; only the buffer contents move. + model_returns_tuple = self.model_returns_tuple() + input_ids_buf = self.input_ids[:batch_size] + hidden_states_buf = self.hidden_states[:batch_size] + if self.supports_mm_inputs: + model_input_ids = None + model_inputs_embeds = self.inputs_embeds[:input_batch_size] + else: + model_input_ids = self.input_ids[:input_batch_size] + model_inputs_embeds = None + model_kwargs = { + "input_ids": model_input_ids, + "positions": self._get_positions(input_batch_size), + "inputs_embeds": model_inputs_embeds, + } + if self.pass_hidden_states_to_model: + model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size] + # The attention layers look their metadata up in this dict at call + # time, so refreshing its entries in place is enough to hand every + # iteration its own metadata under a single forward context. + per_layer_attn_metadata: dict[str, object] = {} + reuse_attn_metadata = self._draft_attn_metadata_is_buffer_backed() + # A FULL graph replays the captured attention kernel on the buffers + # the step-update kernel advances; no Python-side metadata is read, + # so none is built for it. + need_attn_metadata = cudagraph_runtime_mode != CUDAGraphMode.FULL + + with set_forward_context( + per_layer_attn_metadata, + self.vllm_config, + num_tokens=input_batch_size, + num_tokens_across_dp=batch_size_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + slot_mapping=self._get_slot_mapping(input_batch_size), + ): + for token_index in range(self.num_speculative_tokens - 1): + if not self.constant_draft_positions: + positions = self._update_positions_dependent_metadata( + positions, + common_attn_metadata, + batch_size, + input_batch_size, + block_size, ) - ) - # copy inputs to buffer for cudagraph - self.input_ids[:batch_size] = input_ids - self.hidden_states[:batch_size] = hidden_states - if self.supports_mm_inputs: - self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids) + # Refresh the attention metadata. When draft positions are + # constant (e.g. Gemma4 MTP), common_attn_metadata is invariant + # across loop iterations so we build once and reuse. When the + # backend's metadata is nothing but views of the persistent + # buffers that the step-update kernel already advanced in + # place (FlashAttention, see + # _draft_attn_metadata_is_buffer_backed), later iterations only + # need the scalar max_seq_len moved forward. + if need_attn_metadata and ( + token_index == 0 or not self.constant_draft_positions + ): + if token_index > 0 and reuse_attn_metadata: + for name, md in per_layer_attn_metadata.items(): + per_layer_attn_metadata[name] = dataclass_replace( + md, # type: ignore[type-var] + max_seq_len=common_attn_metadata.max_seq_len, + ) + else: + _, fresh = self.build_per_group_and_layer_attn_metadata( + common_attn_metadata, draft_index=token_index + 1 + ) + per_layer_attn_metadata.update(fresh) - input_ids = None - inputs_embeds = self.inputs_embeds[:input_batch_size] - else: - input_ids = self.input_ids[:input_batch_size] - inputs_embeds = None - - # Run the model. - model_kwargs = { - "input_ids": input_ids, - "positions": self._get_positions(input_batch_size), - "inputs_embeds": inputs_embeds, - } - if self.pass_hidden_states_to_model: - model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size] + # Copy inputs to the persistent buffers for cudagraph. The + # int64 -> int32 cast of the sampled ids (crucial when the eagle + # model is compiled) happens inside the copy. + input_ids_buf.copy_(draft_token_ids_list[-1]) + hidden_states_buf.copy_(hidden_states) + if self.supports_mm_inputs: + self.inputs_embeds[:batch_size] = self.model.embed_input_ids( + input_ids_buf + ) - with set_forward_context( - per_layer_attn_metadata, - self.vllm_config, - num_tokens=input_batch_size, - num_tokens_across_dp=batch_size_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - slot_mapping=self._get_slot_mapping(input_batch_size), - ): + # Run the model. ret_hidden_states = self.model(**model_kwargs) - if not self.model_returns_tuple(): + if not model_returns_tuple: last_hidden_states = ret_hidden_states hidden_states = ret_hidden_states else: last_hidden_states, hidden_states = ret_hidden_states - hidden_states = hidden_states[:batch_size] - draft_token_ids, draft_probs = self._sample_draft_tokens( - last_hidden_states[:batch_size], sampling_metadata - ) - if draft_probs is not None: - assert draft_probs_list is not None - draft_probs_list.append(draft_probs) - draft_token_ids_list.append(draft_token_ids) + hidden_states = hidden_states[:batch_size] + draft_token_ids, draft_probs = self._sample_draft_tokens( + last_hidden_states[:batch_size], sampling_metadata + ) + if draft_probs is not None: + assert draft_probs_list is not None + draft_probs_list.append(draft_probs) + draft_token_ids_list.append(draft_token_ids) # [batch_size, num_speculative_tokens] draft_token_ids = torch.stack(draft_token_ids_list, dim=1) @@ -686,6 +1358,44 @@ class SpecDecodeBaseProposer: self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous() return draft_token_ids + def _draft_attn_metadata_is_buffer_backed(self) -> bool: + """True when the drafter's attention metadata can be carried across + draft steps by only advancing its scalar ``max_seq_len``. + + The step-update kernel advances positions, seq_lens and the slot + mapping in place, and FlashAttention's fast (non-AOT) drafting + metadata is nothing but views of those buffers plus Python scalars + that stay constant within the loop (query_start_loc, max_query_len=1, + num_actual_tokens, max_num_splits) and the one scalar that moves, + max_seq_len. Rebuilding it every step therefore produces an object + equal to ``replace(previous, max_seq_len=...)``. Backends that plan on + the host from sequence lengths (FlashInfer, MLA, cascade, DCP) are + not eligible and keep the per-step rebuild. + """ + cached = getattr(self, "_attn_metadata_buffer_backed", None) + if cached is not None: + return cached + eligible = False + try: + from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionMetadataBuilder, + ) + + builders = [g.get_metadata_builder() for g in self.draft_attn_groups] + eligible = bool(builders) and all( + type(b) is FlashAttentionMetadataBuilder + and getattr(b, "dcp_world_size", 1) == 1 + for b in builders + ) + except Exception: + eligible = False + self._attn_metadata_buffer_backed = eligible + if eligible: + logger.info( + "Drafter reuses FlashAttention metadata across draft steps " + "(only max_seq_len is advanced per step)." + ) + return eligible def _update_positions_dependent_metadata( self, positions: torch.Tensor, @@ -1306,6 +2016,7 @@ class SpecDecodeBaseProposer: self._maybe_share_embeddings(target_language_model) self._maybe_share_lm_head(target_language_model) + self._install_fp8_draft_head() if ( self.parallel_drafting @@ -1656,14 +2367,34 @@ class SpecDecodeBaseProposer: self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size ) logger.debug("Using block size %d for drafting layers", self.block_size) + self._maybe_enable_full_cudagraphs() def _determine_batch_execution_and_padding( self, num_tokens: int, use_cudagraphs: bool = True, ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]: + cudagraph_mode, num_tokens_padded, num_tokens_across_dp, _ = ( + self._dispatch_batch(num_tokens, use_cudagraphs=use_cudagraphs) + ) + return cudagraph_mode, num_tokens_padded, num_tokens_across_dp + + def _dispatch_batch( + self, + num_tokens: int, + use_cudagraphs: bool = True, + uniform_decode: bool = False, + ) -> tuple[CUDAGraphMode, int, torch.Tensor | None, BatchDescriptor]: + """Dispatch like _determine_batch_execution_and_padding and also + return the batch descriptor the graph wrappers must be given. + + ``uniform_decode`` marks a batch of one token per request (the draft + loop); it only changes the outcome when FULL keys exist, otherwise + the dispatcher relaxes it to the piecewise key as before. + """ cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch( num_tokens, + uniform_decode=uniform_decode, valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None), ) num_tokens_padded = batch_desc.num_tokens @@ -1692,6 +2423,7 @@ class SpecDecodeBaseProposer: # batch_descriptor cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch( num_tokens_padded, + uniform_decode=uniform_decode, valid_modes={CUDAGraphMode(synced_cudagraph_mode)}, ) # Assert to make sure the agreed upon token count is correct @@ -1699,7 +2431,7 @@ class SpecDecodeBaseProposer: assert batch_desc.num_tokens == num_tokens_padded num_tokens_across_dp[dp_rank] = num_tokens_padded - return cudagraph_mode, num_tokens_padded, num_tokens_across_dp + return cudagraph_mode, num_tokens_padded, num_tokens_across_dp, batch_desc # NOTE(woosuk): Currently, the below code is not used and we always use argmax diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py index 65b9408a8..5a8d458c4 100644 --- a/vllm/v1/spec_decode/utils.py +++ b/vllm/v1/spec_decode/utils.py @@ -599,3 +599,75 @@ def unconditional_to_conditional_rates(rates: list[float]) -> list[float]: """Convert per-position unconditional rates to per-position conditional rates for the early-terminating rejection loop (c_i = p_i / p_{i-1}).""" return [p / q if q > 0.0 else 0.0 for p, q in zip(rates, [1.0, *rates[:-1]])] + + +@triton.jit +def eagle_step_slot_mapping_metadata_devbs_kernel( + positions_ptr, # [input_batch_size] - read and write (in place) + block_table_ptr, # [input_batch_size, n_blocks_per_req] + block_table_stride, + seq_lens_ptr, # [input_batch_size] - read and write + out_slot_mapping_ptr, # [input_batch_size] (output) + batch_size_ptr, # [1] int32 - live batch size, read on device + block_size: tl.constexpr, + max_model_len: tl.constexpr, + n_blocks_per_req: tl.constexpr, + PAD_ID: tl.constexpr, +): + """``eagle_step_slot_mapping_metadata_kernel`` with the batch size read + from device memory and positions updated in place, so the launch can be + recorded into a CUDA graph once per padded batch size and replayed for + any smaller real batch: rows at or beyond the live batch size only write + PADDING_SLOT_ID and leave their position and sequence length alone.""" + req_idx = tl.program_id(0) + batch_size = tl.load(batch_size_ptr) + if req_idx >= batch_size: + tl.store(out_slot_mapping_ptr + req_idx, PAD_ID) + return + + position = tl.load(positions_ptr + req_idx) + new_position = position + 1 + exceeds_max = new_position >= max_model_len + clamped_position = tl.where(exceeds_max, 0, new_position) + + block_number = clamped_position // block_size + block_number = tl.minimum(block_number, n_blocks_per_req - 1) + block_id = tl.load(block_table_ptr + req_idx * block_table_stride + block_number) + slot_id = block_id * block_size + (clamped_position % block_size) + slot_id = tl.where(exceeds_max, PAD_ID, slot_id) + + seq_len = tl.load(seq_lens_ptr + req_idx) + new_seq_len = tl.where(exceeds_max, 1, seq_len + 1) + new_seq_len = tl.minimum(new_seq_len, max_model_len) + + tl.store(positions_ptr + req_idx, clamped_position) + tl.store(out_slot_mapping_ptr + req_idx, slot_id) + tl.store(seq_lens_ptr + req_idx, new_seq_len) + + +def eagle_step_update_slot_mapping_and_metadata_devbs( + positions: torch.Tensor, + block_table_tensor: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_model_len: int, + out_slot_mapping: torch.Tensor, + batch_size_dev: torch.Tensor, + input_batch_size: int, +) -> None: + """Graph-recordable variant of ``eagle_step_update_slot_mapping_and_metadata``: + ``positions`` is advanced in place and the live batch size comes from the + one-element device tensor ``batch_size_dev``.""" + n_blocks_per_req = block_table_tensor.shape[1] + eagle_step_slot_mapping_metadata_devbs_kernel[(input_batch_size,)]( + positions, + block_table_tensor, + block_table_tensor.stride(0), + seq_lens, + out_slot_mapping, + batch_size_dev, + block_size=block_size, + max_model_len=max_model_len, + n_blocks_per_req=n_blocks_per_req, + PAD_ID=PADDING_SLOT_ID, + ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a823..2728744cd 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -765,6 +765,9 @@ class GPUModelRunner( self.num_accepted_tokens = self._make_buffer( self.max_num_reqs, dtype=torch.int32 ) + # Set per step by _prepare_inputs when decode rows of a hybrid model + # need the speculative attention-metadata path without any drafts. + self._hybrid_spec_state_rows_pending = False # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: @@ -2159,6 +2162,11 @@ class GPUModelRunner( target.gpu[:, :total_num_scheduled_tokens] += drift use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + # Decode rows of a hybrid model whose recurrent state must be read + # back through the speculative metadata path this step, drafts or not + # (see _hybrid_decode_rows_with_spec_state). + spec_state_rows = self._hybrid_decode_rows_with_spec_state(num_reqs) + self._hybrid_spec_state_rows_pending = False if not use_spec_decode: # NOTE(woosuk): Due to chunked prefills, the batch may contain # partial requests. While we should not sample any token @@ -2168,6 +2176,17 @@ class GPUModelRunner( logits_indices = query_start_loc[1:] - 1 spec_decode_metadata = None num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + if spec_state_rows is not None: + # No drafts anywhere this step. Present the decode rows as + # speculative rows carrying zero drafts, so the linear + # attention backends index each row's state by its accepted + # count instead of assuming slot 0. + num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32) + num_decode_draft_tokens[spec_state_rows] = 0 + self.num_decode_draft_tokens.np[:num_reqs] = num_decode_draft_tokens + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + self._hybrid_spec_state_rows_pending = True else: # Get the number of draft tokens for each request. # Iterate over the dictionary rather than all requests since not all @@ -2188,6 +2207,13 @@ class GPUModelRunner( >= self.input_batch.num_prompt_tokens[req_idx] ): num_decode_draft_tokens[req_idx] = draft_len + if spec_state_rows is not None: + # A decode row that was scheduled without drafts this step + # still has to read its state by accepted count; zero drafts + # keeps it on the speculative path with a one-token query. + num_decode_draft_tokens[ + spec_state_rows & (num_decode_draft_tokens < 0) + ] = 0 spec_decode_metadata = self._calc_spec_decode_metadata( num_draft_tokens, cu_num_tokens ) @@ -2213,6 +2239,27 @@ class GPUModelRunner( spec_decode_metadata, ) + def _hybrid_decode_rows_with_spec_state(self, num_reqs: int) -> np.ndarray | None: + """Mask of decode rows whose recurrent state lives at an accepted offset. + + With speculative decoding, the recurrent (linear-attention) kernels + write one state per processed token into consecutive state slots and, + on the next step, resume from the slot of the last accepted token. + That resume happens only on the speculative metadata path: the + non-speculative path always resumes from slot 0. A decode row that + reaches a step without drafts (none scheduled for it, or none for the + whole step) would therefore continue from a state that predates the + tokens it already accepted. Returns None when this cannot apply. + """ + if self.speculative_config is None or not self.model_config.is_hybrid: + return None + computed = self.input_batch.num_computed_tokens_cpu[:num_reqs] + prompt = self.input_batch.num_prompt_tokens[:num_reqs] + rows = computed >= prompt + if not bool(rows.any()): + return None + return rows + def _build_attention_metadata( self, num_tokens: int, @@ -3848,6 +3895,23 @@ class GPUModelRunner( num_reqs=num_reqs, force_uniform_decode=force_uniform_decode, ) + if ( + uniform_decode + and force_uniform_decode is None + and self.uniform_decode_query_len > 1 + and num_reqs > 0 + ): + # A prompt chunk of exactly 1 + num_speculative_tokens tokens (a + # block-aligned remainder or the tail after a prefix-cache hit) + # has the geometry of a uniform speculative-decode batch, but its + # tokens are prompt tokens without drafts: it must take the + # prefill path, never the FULL verification graph. + ib = self.input_batch + still_prefilling = ( + ib.num_computed_tokens_cpu[:num_reqs] < ib.num_prompt_tokens[:num_reqs] + ) + if bool(still_prefilling.any()): + uniform_decode = False # Encoder-decoder models only support CG for decoder_step > 0 (no enc_output # is present). Also, chunked-prefill is disabled, so batch are uniform. has_encoder_output = ( @@ -4250,7 +4314,10 @@ class GPUModelRunner( self.mamba_state_idx, ) - use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + use_spec_decode = ( + len(scheduler_output.scheduled_spec_decode_tokens) > 0 + or self._hybrid_spec_state_rows_pending + ) ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices slot_mappings_by_group, slot_mappings = self._get_slot_mappings( @@ -5875,6 +5942,27 @@ class GPUModelRunner( # requests can corrupt Mamba state. self.input_batch.block_table.commit_block_table(num_reqs_padded) + if self.speculative_config is not None: + # Dummy batches bypass _prepare_inputs, so the per-step + # speculative scratch would otherwise hold whatever the + # buffers were created with. The recurrent-state kernels + # resume from state slot (num_accepted_tokens - 1) for + # every row on the speculative path, so an accepted count + # of zero indexes one slot before the table. Give every + # row the minimum valid count, and describe the rows the + # way a real step would: a uniform decode batch carries + # max_query_len - 1 drafts per row, anything else none. + self.num_accepted_tokens.np.fill(1) + self.num_accepted_tokens.copy_to_gpu() + if uniform_decode and max_query_len > 1: + self.num_decode_draft_tokens.np[:num_reqs].fill( + max_query_len - 1 + ) + else: + self.num_decode_draft_tokens.np[:num_reqs].fill(-1) + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL attn_metadata, _ = self._build_attention_metadata( num_tokens=num_tokens_unpadded, @@ -6627,6 +6715,16 @@ class GPUModelRunner( ) torch.accelerator.synchronize() + # The drafter's piecewise graphs were captured alongside the + # target's above (via dummy_run); its FULL graphs for the + # single-token draft steps, if enabled, are captured here. + capture_drafter_full = getattr( + getattr(self, "drafter", None), "capture_full_cudagraphs", None + ) + if capture_drafter_full is not None: + capture_drafter_full() + torch.accelerator.synchronize() + # Capture encoder CUDA graphs if enabled if self.encoder_cudagraph_manager is not None: encoder_graph_pool = current_platform.graph_pool_handle() diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 45166ef9a..1df9db087 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -702,6 +702,18 @@ def preprocess_mamba( for i, req_id in enumerate(input_batch.req_ids): req_state = requests[req_id] prev_state_idx = mamba_state_idx.get(req_id) + # A row only carries a speculative accepted count when the request's + # previous step was a speculative decode step. New and resumed + # requests, prefix-cache hits and chunked-prefill continuations + # resume from exactly the state stored for their last computed token + # (slot 0), whatever the reused batch row still holds from its + # previous occupant or from the asynchronous write-back of the last + # step's counts. + if ( + prev_state_idx is None + or req_state.num_computed_tokens <= req_state.num_prompt_tokens + ): + input_batch.num_accepted_tokens_cpu[i] = 1 if prev_state_idx is None: # new / resumed request, no previous state # if num_computed_tokens is 0, prev_state_idx will be -1