diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py index c58ce31bd2..fa8b5f7349 100644 --- a/vllm/compilation/passes/fusion/act_quant_fusion.py +++ b/vllm/compilation/passes/fusion/act_quant_fusion.py @@ -299,6 +299,13 @@ class ActivationQuantFusionPass(VllmFusionPatternMatcherPass): self.register(SiluMulNvfp4QuantPattern()) if current_platform.is_cuda(): + # SN10: the fused silu_and_mul_per_block_quant kernel stores + # transposed scales as scales[group * num_tokens + token], i.e. + # it ignores the padded leading dimension of a TMA-aligned scale + # tensor (stride round_up(num_tokens, 4)). Fusing such a quant + # therefore scrambles the scales whenever num_tokens % 4 != 0, so + # the TMA-aligned variants are not registered; that quant stays a + # standalone per_token_group_fp8_quant, which honours the stride. for ( quant_key, is_scale_transposed, @@ -308,7 +315,7 @@ class ActivationQuantFusionPass(VllmFusionPatternMatcherPass): [kFp8Dynamic128Sym, kFp8Dynamic64Sym], [False, True], [True, False], - [False, True], + [False], ): self.register( SiluMulBlockQuantPattern( diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index de505e122c..585869ac74 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 = 9 +_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 ba7d26c93b..f3b28f551c 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 921f31466b..bffa2939d2 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 a16f522183..d67d3720e3 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 fef1741351..e41fbaff6a 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/envs.py b/vllm/envs.py index 27a85bb3d0..ccdc638afa 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -46,7 +46,7 @@ if TYPE_CHECKING: NO_COLOR: bool = False VLLM_LOG_STATS_INTERVAL: float = 10.0 VLLM_TRACE_FUNCTION: int = 0 - VLLM_USE_FLASHINFER_SAMPLER: bool = True + VLLM_USE_FLASHINFER_SAMPLER: bool = False VLLM_PP_LAYER_PARTITION: str | None = None VLLM_CPU_KVCACHE_SPACE: int | None = 0 VLLM_CPU_OMP_THREADS_BIND: str = "auto" @@ -793,10 +793,13 @@ environment_variables: dict[str, Callable[[], Any]] = { # Whether to use the FlashInfer top-k / top-p sampler on CUDA. Enabled # by default when the hardware supports it — set to 0 to opt out # explicitly, which forces the PyTorch-native (Triton for bs>=8) path. + # Off unless asked for: the sampler kernels are JIT-compiled during the + # start-up profiling run, which costs about a minute of a cold start, + # and greedy serving never calls them. Set to 1 to opt back in. "VLLM_USE_FLASHINFER_SAMPLER": lambda: ( bool(int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"])) if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ - else True + else False ), # Pipeline stage partition strategy "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None), diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4ac8d49cd5..5aad0cd5c2 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -154,6 +154,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import ( FlashInferFp8DeepGEMMDynamicBlockScaledKernel, FlashInferFP8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.sn10_w8a8 import ( + SN10FlashInferW8A8BlockScaledKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.marlin import ( MarlinFP8ScaledMMLinearKernel, ) @@ -321,6 +324,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] ] = { PlatformEnum.CUDA: [ + SN10FlashInferW8A8BlockScaledKernel, FlashInferFp8DeepGEMMDynamicBlockScaledKernel, DeepGemmFp8BlockScaledMMKernel, CutlassFp8BlockScaledMMKernel, diff --git a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py index 72a3b84984..fc0236dc15 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py @@ -301,6 +301,20 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl( if envs.VLLM_BATCH_INVARIANT: return run_deepgemm(input, weight, weight_scale) + # Small verify/decode batches on narrow outputs: the swapAB kernel below + # runs one CTA per 128 output columns over the full K and is latency + # bound there; a split-K kernel with fused activation quant keeps the + # whole GPU busy. Declines (returns None) for shapes it has no tuned + # configuration for, so everything else takes the stock path. + if input.shape[0] < 32: + from vllm.model_executor.kernels.linear.scaled_mm.sn10_splitk import ( + small_m_gemm, + ) + + out = small_m_gemm(input, weight, weight_scale) + if out is not None: + return out + condition = input.shape[0] < 32 # PyTorch's torch.compile cannot handle input-dependent control flow in standard diff --git a/vllm/model_executor/kernels/linear/scaled_mm/sn10_splitk.py b/vllm/model_executor/kernels/linear/scaled_mm/sn10_splitk.py new file mode 100644 index 0000000000..0e17a9f111 --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/sn10_splitk.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Split-K FP8 block-scaled GEMM for small M (decode / speculative verify). + +Target: A[M, K] (bf16 activations, M <= 32) x W[N, K]^T (fp8 e4m3, 128x128 +block scales) -> C[M, N] bf16. The stock small-M path (FlashInfer swapAB) +tiles N in 128-wide blocks and runs one CTA per block over the whole K, so a +narrow output (N = 5120) gets 40 CTAs on a 132-SM part and the kernel is +latency-bound at ~20 us regardless of K. This kernel splits K across +SPLIT_K programs, tiles N with BLOCK_N, and folds the per-token-group +(1 x 128) activation quantisation into the main loop (the group is exactly +one K block, so every program can derive its own scales from the bf16 rows +it loads anyway). Partials are reduced in a fixed order, so the result is +deterministic run to run. + +Numerics: activation quant reproduces vLLM's per_token_group_quant_fp8 +(absmax clamped at eps, scale = absmax / 448, clamp to +-448, RNE to e4m3); +products accumulate in fp32 per 128-block and are scaled by a_s * b_s, +exactly like the stock Triton block-scaled kernel. The split partials are +summed in fp32 in ascending split order. +""" +from __future__ import annotations + +import functools +import os + +import torch +import triton +import triton.language as tl + +FP8_MAX = 448.0 +FP8_MIN = -448.0 +EPS = 1e-10 +GROUP_K = 128 + + +@triton.jit +def _sk_gemm_kernel( + A, B, Bs, W, C, + M, N, K, nkb, nkb_per_split, num_tiles, num_pid_n, + stride_am, stride_bn, stride_bs_n, stride_bs_k, stride_ws, stride_wm, stride_cm, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, GK: tl.constexpr, + SPLIT_K: tl.constexpr, FP8_MAX_C: tl.constexpr, EPS_C: tl.constexpr, + INTERP: tl.constexpr, +): + # Persistent loop over (n-block, k-split) tiles. + for tile in range(tl.program_id(0), num_tiles, tl.num_programs(0)): + pid_n = tile % num_pid_n + pid_k = tile // num_pid_n + offs_m = tl.arange(0, BLOCK_M) + m_mask = offs_m < M + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = offs_n < N + offs_k = tl.arange(0, GK) + kb0 = pid_k * nkb_per_split + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for i in range(0, nkb_per_split): + kb = kb0 + i + kvalid = kb < nkb + k0 = kb * GK + a = tl.load( + A + offs_m[:, None] * stride_am + (k0 + offs_k)[None, :], + mask=m_mask[:, None] & kvalid, other=0.0, + ).to(tl.float32) + amax = tl.maximum(tl.max(tl.abs(a), axis=1), EPS_C) + a_s = amax * (1.0 / FP8_MAX_C) + a_q = tl.clamp(a / a_s[:, None], -FP8_MAX_C, FP8_MAX_C) + b = tl.load( + B + (k0 + offs_k)[:, None] + offs_n[None, :] * stride_bn, + mask=n_mask[None, :] & kvalid, other=0.0, + ) + b_s = tl.load( + Bs + (offs_n // 128) * stride_bs_n + kb * stride_bs_k, + mask=n_mask & kvalid, other=0.0, + ) + if INTERP: + prod = tl.dot(a_q, b.to(tl.float32)) + else: + prod = tl.dot(a_q.to(tl.float8e4nv), b) + acc += prod * a_s[:, None] * b_s[None, :] + if SPLIT_K == 1: + tl.store( + C + offs_m[:, None] * stride_cm + offs_n[None, :], + acc.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :], + ) + else: + tl.store( + W + pid_k * stride_ws + offs_m[:, None] * stride_wm + offs_n[None, :], + acc, mask=m_mask[:, None] & n_mask[None, :], + ) + + +@triton.jit +def _sk_reduce_kernel(W, C, MN, N, stride_ws, SPLIT_K: tl.constexpr, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < MN + acc = tl.zeros((BLOCK,), dtype=tl.float32) + for s in range(0, SPLIT_K): + acc += tl.load(W + s * stride_ws + offs, mask=mask, other=0.0) + tl.store(C + offs, acc.to(tl.bfloat16), mask=mask) + + +@functools.lru_cache(maxsize=None) +def _num_sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +# (BLOCK_N, SPLIT_K, num_warps, num_stages) per (N, K), from the H200 +# micro-benchmark. Shapes that are not listed keep the stock kernel unless +# SN10_SPLITK=all is set, in which case _default_config is used for them. +_CONFIGS: dict[tuple[int, int], tuple[int, int, int, int]] = {} + +_MODE = os.environ.get("SN10_SPLITK", "table").strip().lower() # table | all | 0 + + +def small_m_gemm(a: torch.Tensor, w: torch.Tensor, ws: torch.Tensor): + """Entry point for the M < 32 dispatch. Returns None to decline.""" + if _MODE in ("0", "off", "false"): + return None + M, K = a.shape + N = w.shape[0] + if M > 32 or K % GROUP_K != 0 or a.dtype != torch.bfloat16: + return None + if w.dtype != torch.float8_e4m3fn or ws.dtype != torch.float32: + return None + cfg = _CONFIGS.get((N, K)) + if cfg is None and _MODE != "all": + return None + return splitk_blockscaled_gemm(a, w, ws, config=cfg) + + +def _default_config(M: int, N: int, K: int, num_sms: int): + nkb = K // GROUP_K + block_n = 64 if N <= 8192 else 128 + tiles_n = triton.cdiv(N, block_n) + # aim for ~2.5 waves of tiles so the persistent loop balances well + split_k = 1 + while tiles_n * split_k < 2.5 * num_sms and split_k * 2 <= nkb and split_k < 8: + split_k *= 2 + return block_n, split_k, 4, 3 + + +def splitk_blockscaled_gemm( + a: torch.Tensor, w: torch.Tensor, ws: torch.Tensor, + config: tuple[int, int, int, int] | None = None, interp: bool = False, +) -> torch.Tensor: + """a: [M, K] bf16 (M <= 32); w: [N, K] fp8 e4m3; ws: [N/128, K/128] fp32.""" + M, K = a.shape + N = w.shape[0] + assert a.stride(1) == 1 and w.stride(1) == 1 + assert K % GROUP_K == 0 + num_sms = 132 if interp else _num_sms(a.device.index or 0) + if config is None: + config = _CONFIGS.get((N, K)) or _default_config(M, N, K, num_sms) + block_n, split_k, num_warps, num_stages = config + block_m = 16 if M <= 16 else 32 + nkb = K // GROUP_K + nkb_per_split = triton.cdiv(nkb, split_k) + num_pid_n = triton.cdiv(N, block_n) + num_tiles = num_pid_n * split_k + grid = (min(num_tiles, num_sms),) + c = torch.empty((M, N), dtype=torch.bfloat16, device=a.device) + if split_k > 1: + wsp = torch.empty((split_k, M, N), dtype=torch.float32, device=a.device) + else: + wsp = c # unused + _sk_gemm_kernel[grid]( + a, w, ws, wsp, c, + M, N, K, nkb, nkb_per_split, num_tiles, num_pid_n, + a.stride(0), w.stride(0), ws.stride(0), ws.stride(1), + wsp.stride(0) if split_k > 1 else 0, wsp.stride(1) if split_k > 1 else 0, c.stride(0), + BLOCK_M=block_m, BLOCK_N=block_n, GK=GROUP_K, SPLIT_K=split_k, + FP8_MAX_C=FP8_MAX, EPS_C=EPS, INTERP=interp, + num_warps=num_warps, num_stages=num_stages, + ) + if split_k > 1: + MN = M * N + BLOCK = 1024 + _sk_reduce_kernel[(triton.cdiv(MN, BLOCK),)]( + wsp, c, MN, N, wsp.stride(0), SPLIT_K=split_k, BLOCK=BLOCK, num_warps=4, + ) + return c diff --git a/vllm/model_executor/kernels/linear/scaled_mm/sn10_w8a8.py b/vllm/model_executor/kernels/linear/scaled_mm/sn10_w8a8.py new file mode 100644 index 0000000000..8e60e62013 --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/sn10_w8a8.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP8 block-scaled linear kernel that hands FlashInfer / DeepGEMM an +activation that is already quantised. + +The stock dynamic kernel feeds FlashInfer's swapAB GEMM with BF16 activations +and lets it quantise them internally: one extra pass over the activation and +a separate scale kernel for every one of the ~280 GEMMs of a decode step, and +no chance for torch.compile to fold the quantisation into the preceding +RMSNorm / SiLU. With ``apply_input_quant`` on, the base class quantises in +the usual per-token-group (1x128) way -- which the norm/activation fusion +passes then absorb -- and both branches below consume the FP8 tensor +directly. FlashInfer's small-M kernel expects the scales in the same +TMA-aligned column-major layout DeepGEMM uses (M padded to a multiple of +4), so a single quantisation serves both. + +Numerics are unchanged in kind: the same e4m3 rounding with per-token-group +fp32 scales that FlashInfer's internal quantisation uses, just performed once +instead of inside every GEMM. (The stock DeepGEMM branch rounds activation +scales to powers of two on this platform; full-precision fp32 scales are used +here for both branches -- DeepGEMM on Hopper reads fp32 scales natively. +Weights keep the stock post-processing, so the GEMM operands are identical +to the stock kernel's.) +""" +from __future__ import annotations + +import os +from typing import ClassVar + +import torch + +import vllm.envs as envs +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import fp8_gemm_nt, is_deep_gemm_supported +from vllm.utils.flashinfer import ( + flashinfer_fp8_blockscale_gemm, + is_flashinfer_fp8_blockscale_gemm_supported, +) +from vllm.utils.torch_utils import direct_register_custom_op + +from .BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, +) +from .deep_gemm import DeepGemmFp8BlockScaledMMKernel + +# Below this many rows FlashInfer's swapAB kernel wins; above it DeepGEMM. +_SMALL_M = 32 + + +def w8a8_enabled() -> bool: + return os.environ.get("SN10_W8A8", "0").strip().lower() in ("1", "on", "true") + + +class SN10FlashInferW8A8BlockScaledKernel(Fp8BlockScaledMMLinearKernel): + apply_input_quant: ClassVar[bool] = True + + def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None: + super().__init__(config) + # Full-precision fp32 activation scales for both branches (see above). + self.act_ue8m0 = False + act_scale_descriptor = config.activation_quant_key.scale + # TMA-aligned column-major scales: the layout both GEMM paths read + # (FlashInfer's W8A8 kernel needs the M dimension padded to 4). + self.quant_fp8 = QuantFP8( + static=False, + group_shape=act_scale_descriptor.group_shape, + use_ue8m0=self.act_ue8m0, + tma_aligned_scales=True, + column_major_scales=True, + ) + self._deepgemm = DeepGemmFp8BlockScaledMMKernel(config) + + @classmethod + def is_supported(cls, compute_capability=None): + if not w8a8_enabled(): + return False, "disabled via SN10_W8A8=0" + if not current_platform.is_cuda() or not current_platform.is_device_capability(90): + return False, "only Hopper is supported" + if not is_flashinfer_fp8_blockscale_gemm_supported(): + return False, "FlashInfer block-scale FP8 GEMM is not available." + if not is_deep_gemm_supported(): + return False, "DeepGEMM is not available." + return True, None + + @classmethod + def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): + ok, reason = super().can_implement(config) + if not ok: + return ok, reason + if config.out_dtype != torch.bfloat16: + return False, "Supports only output dtype of bfloat16" + if config.activation_quant_key.scale.group_shape != GroupShape(1, 128): + return False, "Supports only per-token-group (1,128) activation scales." + return DeepGemmFp8BlockScaledMMKernel.can_implement(config) + + def process_weights_after_loading(self, layer: torch.nn.Module): + # Same tensor layout for both branches; DeepGEMM's post-processing + # is what the stock dynamic kernel applies too. + self._deepgemm.process_weights_after_loading(layer) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.vllm.sn10_w8a8_blockscale_gemm(A, As, B, Bs, self.act_ue8m0) + + +def _sn10_w8a8_blockscale_gemm_impl( + q_input: torch.Tensor, + input_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + use_deep_gemm_e8m0: bool, +) -> torch.Tensor: + if q_input.shape[0] < _SMALL_M and not envs.VLLM_BATCH_INVARIANT: + return flashinfer_fp8_blockscale_gemm( + input=q_input, + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + out_dtype=torch.bfloat16, + ) + output = torch.empty( + (q_input.shape[0], weight.shape[0]), + dtype=torch.bfloat16, + device=q_input.device, + ) + fp8_gemm_nt( + (q_input, input_scale), + (weight, weight_scale), + output, + is_deep_gemm_e8m0_used=use_deep_gemm_e8m0, + ) + return output + + +def _sn10_w8a8_blockscale_gemm_fake( + q_input: torch.Tensor, + input_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + use_deep_gemm_e8m0: bool, +) -> torch.Tensor: + return torch.empty( + q_input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=q_input.device + ) + + +direct_register_custom_op( + "sn10_w8a8_blockscale_gemm", + _sn10_w8a8_blockscale_gemm_impl, + fake_impl=_sn10_w8a8_blockscale_gemm_fake, +) diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 7e0c7e05ca..9ba86ff2b2 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -51,6 +51,15 @@ def fused_sigmoid_gating_delta_rule_update_kernel( stride_final_state_token: tl.constexpr, stride_indices_seq: tl.constexpr, stride_indices_tok: tl.constexpr, + # Distance (in elements) between two consecutive tokens of q/k/v/a/b. + # Packed inputs use H*K / HV*V / HV; views cut out of a wider row-major + # buffer (e.g. the [T, q|k|v] projection) pass that buffer's row width. + # Within one token the heads are always packed (head stride K / V / 1). + stride_q_tok: tl.constexpr, + stride_k_tok: tl.constexpr, + stride_v_tok: tl.constexpr, + stride_a_tok: tl.constexpr, + stride_b_tok: 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, @@ -80,19 +89,19 @@ def fused_sigmoid_gating_delta_rule_update_kernel( o_k = i_k * BK + tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) - p_q = q + (bos * H + i_h) * K + o_k - p_k = k + (bos * H + i_h) * K + o_k - p_v = v + (bos * HV + i_hv) * V + o_v + p_q = q + bos * stride_q_tok + i_h * K + o_k + p_k = k + bos * stride_k_tok + i_h * K + o_k + p_v = v + bos * stride_v_tok + i_hv * V + o_v p_A_log = A_log + i_hv if not IS_KDA: - p_a = a + bos * HV + i_hv + p_a = a + bos * stride_a_tok + i_hv p_dt_bias = dt_bias + i_hv else: - p_a = a + (bos * HV + i_hv) * K + o_k + p_a = a + bos * stride_a_tok + i_hv * K + o_k p_dt_bias = dt_bias + i_hv * K + o_k - p_b = b + bos * HV + i_hv + p_b = b + bos * stride_b_tok + i_hv p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v mask_k = o_k < K @@ -170,12 +179,62 @@ def fused_sigmoid_gating_delta_rule_update_kernel( tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) # Update pointers for next timestep - p_q += H * K - p_k += H * K + p_q += stride_q_tok + p_k += stride_k_tok p_o += HV * V - p_v += HV * V - p_b += HV - p_a += HV + p_v += stride_v_tok + p_b += stride_b_tok + p_a += stride_a_tok + + +def _inplace_token_stride(x: torch.Tensor, T: int, inner: tuple[int, ...]) -> int | None: + """Inter-token stride at which the kernel can read ``x`` without a copy. + + ``x`` is ``[T, *inner]`` or ``[B, T, *inner]``; the kernel addresses it as + one flat run of ``B*T`` tokens whose elements within a token are packed + (row-major, unit stride), consecutive tokens being ``stride`` elements + apart. Returns ``None`` when ``x`` does not have that form, in which case + the caller has to fall back to a packed copy. + """ + n_inner = len(inner) + if x.dim() == n_inner + 1: + B, T_x = 1, x.shape[0] + elif x.dim() == n_inner + 2: + B, T_x = x.shape[0], x.shape[1] + else: + return None + if T_x != T or tuple(x.shape[-n_inner:]) != tuple(inner): + return None + # Elements of one token must sit at their packed offsets. + expected = 1 + for size, stride in zip(reversed(inner), reversed(x.stride()[-n_inner:])): + if size > 1 and stride != expected: + return None + expected *= size + # Strides of size-1 dims carry no information, so read the token stride + # from the batch dim when there is a single token per batch entry. + if T > 1: + stride_tok = x.stride(-n_inner - 1) + if B > 1 and x.stride(0) != T * stride_tok: + return None + elif B > 1: + stride_tok = x.stride(0) + else: + stride_tok = expected + return stride_tok + + +def _packed_or_copy( + x: torch.Tensor, T: int, inner: tuple[int, ...] +) -> tuple[torch.Tensor, int]: + """Return ``x`` plus its token stride, copying to a packed layout if needed.""" + stride_tok = _inplace_token_stride(x, T, inner) + if stride_tok is None: + x = x.contiguous() + stride_tok = 1 + for size in inner: + stride_tok *= size + return x, stride_tok def fused_sigmoid_gating_delta_rule_update( @@ -238,17 +297,28 @@ def fused_sigmoid_gating_delta_rule_update( else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() + # q/k/v and the gating columns are read in place whenever they are packed + # per token (heads x dim at unit stride) and only the inter-token stride + # differs from the dense layout, e.g. views into the [T, q|k|v] projection + # buffer or the b/a halves of the [T, 2*HV] projection. Anything else is + # copied to a packed layout exactly as before. + q, stride_q_tok = _packed_or_copy(q, T, (H, K)) + k, stride_k_tok = _packed_or_copy(k, T, (H, K)) + v, stride_v_tok = _packed_or_copy(v, T, (HV, V)) + a, stride_a_tok = _packed_or_copy(a, T, (HV, K) if is_kda else (HV,)) + b, stride_b_tok = _packed_or_copy(b, T, (HV,)) + grid = (NK, NV, N * HV) fused_sigmoid_gating_delta_rule_update_kernel[grid]( A_log=A_log, - a=a.contiguous(), - b=b.contiguous(), + a=a, + b=b, dt_bias=dt_bias, beta=beta, threshold=threshold, - q=q.contiguous(), - k=k.contiguous(), - v=v.contiguous(), + q=q, + k=k, + v=v, o=o, h0=initial_state, ht=final_state, @@ -269,6 +339,11 @@ 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, + stride_q_tok=stride_q_tok, + stride_k_tok=stride_k_tok, + stride_v_tok=stride_v_tok, + stride_a_tok=stride_a_tok, + stride_b_tok=stride_b_tok, 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 06bfe5c5de..49d7592b3b 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) @@ -841,6 +895,43 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): return query, key, value + def rearrange_mixed_qkv_views(self, mixed_qkv): + """Split packed qkv into (1, seq, heads, dim) views, without copying. + + The recurrent decode kernel reads its inputs through per-tensor token + strides, so the q/k/v column blocks of the ``[seq, q|k|v]`` buffer can + be handed over as they are: each head's ``dim`` elements are already + adjacent, and only the distance between two tokens (the row width of + ``mixed_qkv``) differs from a dense layout. This saves the three + gather copies plus the concatenation ``rearrange_mixed_qkv`` performs + per layer and step. Buffers that are not row-major with a unit stride + along the feature dim take the copying route instead. + """ + if mixed_qkv is None: + return None, None, None + + q_dim = self.key_dim // self.tp_size + k_dim = self.key_dim // self.tp_size + v_dim = self.value_dim // self.tp_size + if ( + mixed_qkv.dim() != 2 + or mixed_qkv.shape[-1] != q_dim + k_dim + v_dim + or mixed_qkv.stride(-1) != 1 + ): + return self.rearrange_mixed_qkv(mixed_qkv) + + seq_len = mixed_qkv.shape[0] + num_k_heads = q_dim // self.head_k_dim + num_v_heads = v_dim // self.head_v_dim + query = mixed_qkv[:, :q_dim].view(1, seq_len, num_k_heads, self.head_k_dim) + key = mixed_qkv[:, q_dim : q_dim + k_dim].view( + 1, seq_len, num_k_heads, self.head_k_dim + ) + value = mixed_qkv[:, q_dim + k_dim :].view( + 1, seq_len, num_v_heads, self.head_v_dim + ) + return query, key, value + def forward( self, hidden_states: torch.Tensor, @@ -938,9 +1029,12 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): z_size = self.value_dim // self.tp_size mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1) z = z.reshape(z.size(0), -1, self.head_v_dim) + # b and a stay column views of the [tokens, 2*HV] projection: + # every consumer inside the core op (fused_post_conv_prep, the + # packed non-spec decode kernel and the spec/decode recurrent + # kernel) reads them through their token stride, so the two + # per-layer copies a .contiguous() would launch are not needed. b, a = self.split_ba(ba) - b = b.contiguous() - a = a.contiguous() # ============================================================ # Part 2: Core Attention (Custom Op) @@ -1104,7 +1198,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 +1427,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: @@ -1389,7 +1503,10 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): else: mixed_qkv_non_spec = None - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + # The recurrent kernel reads q/k/v straight out of the conv output. + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv_views( + mixed_qkv_spec + ) # Split mixed non-spec-decode+prefill to process independently split_non_spec = ( @@ -1443,8 +1560,10 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): g_non_spec = g_non_spec.unsqueeze(0) beta_non_spec = beta_non_spec.unsqueeze(0) else: - query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv( - mixed_qkv_non_spec + # No prefill in this step: these only feed the recurrent kernel, + # which takes the packed conv output through views. + query_non_spec, key_non_spec, value_non_spec = ( + self.rearrange_mixed_qkv_views(mixed_qkv_non_spec) ) g_non_spec = None beta_non_spec = None @@ -1456,8 +1575,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, @@ -1478,7 +1597,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # 2.2: Process non-spec-decode part if split_non_spec: - query_decode, key_decode, value_decode = self.rearrange_mixed_qkv( + query_decode, key_decode, value_decode = self.rearrange_mixed_qkv_views( mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] ) core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( @@ -1511,7 +1630,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 021462f3ee..868e10dda4 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,48 @@ class Qwen3_5MultiTokenPredictor(nn.Module): return loaded_params +_DRAFT_ARGMAX_STATE = {"on": None} + + +def _draft_argmax_enabled() -> bool: + if _DRAFT_ARGMAX_STATE["on"] is None: + _DRAFT_ARGMAX_STATE["on"] = os.environ.get( + "SN10_DRAFT_ARGMAX", "1" + ).strip().lower() not in ("0", "off", "false") + return _DRAFT_ARGMAX_STATE["on"] + + +def _disable_draft_argmax() -> None: + _DRAFT_ARGMAX_STATE["on"] = False + + +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 +487,145 @@ 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_top_tokens(logits) + 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_top_tokens(logits) + + def _draft_top_tokens(self, logits: torch.Tensor) -> torch.Tensor: + """ids[argmax(logits)] -- two-stage Triton reduction fused with the + id gather (first index on ties, exactly like torch.argmax); falls back + to the torch path on any error or when SN10_DRAFT_ARGMAX=0.""" + if _draft_argmax_enabled(): + try: + from vllm.model_executor.models.sn10_draft_argmax import ( + draft_argmax_gather, + ) + + return draft_argmax_gather(logits, self._draft_ids) + except Exception: + logger.warning( + "MTP draft head: fused argmax failed, using torch.argmax.", + exc_info=True, + ) + _disable_draft_argmax() + 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 +639,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): diff --git a/vllm/model_executor/models/sn10_draft_argmax.py b/vllm/model_executor/models/sn10_draft_argmax.py new file mode 100644 index 0000000000..e1720c9c57 --- /dev/null +++ b/vllm/model_executor/models/sn10_draft_argmax.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Row-wise argmax (first index on ties, like torch.argmax) fused with the +draft-vocabulary id gather, for the MTP draft head. + +torch's reduce kernel for a (rows<=32, ~100k) argmax is latency bound (~30us); +a two-stage Triton reduction over many blocks takes ~5us. Ties resolve to +the smallest index, matching torch.argmax, so the drafts are unchanged. +""" +import torch +import triton +import triton.language as tl + +_BLOCK = 4096 + + +@triton.jit +def _partial_argmax_kernel( + x_ptr, + stride_row, + n_cols, + pmax_ptr, + pidx_ptr, + n_blocks, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + blk = tl.program_id(1) + offs = blk * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_cols + x = tl.load(x_ptr + row * stride_row + offs, mask=mask, other=float("-inf")) + x = x.to(tl.float32) + m = tl.max(x, axis=0) + # first index attaining the block max (masked lanes can never match: -inf) + cand = tl.where(x == m, offs, n_cols) + idx = tl.min(cand, axis=0) + tl.store(pmax_ptr + row * n_blocks + blk, m) + tl.store(pidx_ptr + row * n_blocks + blk, idx) + + +@triton.jit +def _final_argmax_kernel( + pmax_ptr, + pidx_ptr, + n_blocks, + ids_ptr, + out_ptr, + n_cols, + NB: tl.constexpr, +): + row = tl.program_id(0) + offs = tl.arange(0, NB) + mask = offs < n_blocks + m = tl.load(pmax_ptr + row * n_blocks + offs, mask=mask, other=float("-inf")) + idx = tl.load(pidx_ptr + row * n_blocks + offs, mask=mask, other=n_cols) + gmax = tl.max(m, axis=0) + cand = tl.where(m == gmax, idx, n_cols) + best = tl.min(cand, axis=0) + tok = tl.load(ids_ptr + best) + tl.store(out_ptr + row, tok) + + +def draft_argmax_gather(logits: torch.Tensor, ids: torch.Tensor) -> torch.Tensor: + """Equivalent of ``ids[logits.argmax(dim=-1)]`` (int64 output).""" + assert logits.dim() == 2 and logits.stride(1) == 1 + rows, n_cols = logits.shape + n_blocks = triton.cdiv(n_cols, _BLOCK) + pmax = torch.empty((rows, n_blocks), dtype=torch.float32, device=logits.device) + pidx = torch.empty((rows, n_blocks), dtype=torch.int32, device=logits.device) + out = torch.empty((rows,), dtype=torch.int64, device=logits.device) + _partial_argmax_kernel[(rows, n_blocks)]( + logits, logits.stride(0), n_cols, pmax, pidx, n_blocks, BLOCK=_BLOCK + ) + _final_argmax_kernel[(rows,)]( + pmax, pidx, n_blocks, ids, out, n_cols, NB=triton.next_power_of_2(n_blocks) + ) + return out diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f46cbd242..97d331147e 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from importlib.util import find_spec from typing import Any, cast @@ -7,7 +8,10 @@ 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.config import ( CUDAGraphMode, VllmConfig, @@ -243,6 +247,36 @@ class SpecDecodeBaseProposer: self.max_positions, dtype=torch.int64, device=device ) + # SN10 prototype: capture the whole (k-1)-iteration draft loop as ONE + # CUDA graph per (padded batch size, k). Off unless SN10_DRAFT_GRAPH=1. + # Every per-iteration input of the loop already lives in a persistent + # buffer (positions / slot_mapping / seq_lens are advanced in place by + # eagle_step_slot_mapping_metadata_kernel; input_ids / hidden_states + # are the drafter's own static buffers), so the Python loop is pure + # launch overhead. Only the three loop *entry* values (first + # positions, first hidden states, first draft token) are temporaries; + # they are staged into the buffers below before replay. + self._loop_graph_enabled = ( + os.environ.get("SN10_DRAFT_GRAPH", "1").strip() not in ("0", "off", "false") + ) + self._loop_graphs: dict[tuple[int, int, int], torch.cuda.CUDAGraph] = {} + self._loop_graph_warm: set[tuple[int, int, int]] = set() + if self._loop_graph_enabled: + pos_shape = ( + (3, self.max_batch_size) if self.uses_mrope else (self.max_batch_size,) + ) + self._loop_pos_in = torch.zeros(pos_shape, dtype=torch.int64, device=device) + self._loop_hid_in = torch.zeros( + (self.max_batch_size, self.hidden_size), + dtype=self.dtype, + device=device, + ) + self._loop_tokens = torch.zeros( + (self.max_batch_size, max(self.num_speculative_tokens, 1)), + dtype=torch.int64, + device=device, + ) + # Determine allowed attention backends once during initialization. self.allowed_attn_types: tuple | None = None if current_platform.is_rocm(): @@ -610,6 +644,19 @@ class SpecDecodeBaseProposer: block_size = self.block_size assert block_size > 0, "block_size has not been initialized." + + if self._loop_graph_applicable( + cudagraph_runtime_mode, batch_size_across_dp, sampling_metadata + ): + return self._propose_with_loop_graph( + positions, + hidden_states, + draft_token_ids, + common_attn_metadata, + batch_size, + input_batch_size, + ) + for token_index in range(self.num_speculative_tokens - 1): # Update the inputs. # cast to int32 is crucial when eagle model is compiled. @@ -686,6 +733,215 @@ class SpecDecodeBaseProposer: self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous() return draft_token_ids + # ------------------------------------------------------------------ + # SN10 prototype: whole-draft-loop CUDA graph (SN10_DRAFT_GRAPH=1). + # ------------------------------------------------------------------ + def _loop_graph_applicable( + self, + cudagraph_runtime_mode: CUDAGraphMode, + batch_size_across_dp: torch.Tensor | None, + sampling_metadata: SamplingMetadata, + ) -> bool: + """Conservative allow-list; anything else takes the eager loop.""" + if not self._loop_graph_enabled: + return False + if cudagraph_runtime_mode == CUDAGraphMode.NONE: + # Batch larger than the capture ladder / cudagraphs disabled. + return False + if self.num_speculative_tokens > self._loop_tokens.shape[1]: + return False + checks = { + "pass_hidden_states": self.pass_hidden_states_to_model, + "no_xdrope": not ( + self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0 + ), + "no_constant_positions": not self.constant_draft_positions, + "no_shared_mtp_indices": not self._share_mtp_indices, + "no_dp_padding": batch_size_across_dp is None, + "dp1": self.vllm_config.parallel_config.data_parallel_size == 1, + # TP>1 collectives need vLLM's graph_capture() context; not + # wired up in the prototype. + "tp1": self.vllm_config.parallel_config.tensor_parallel_size == 1, + "no_breakable_graph": not is_breakable_cudagraph_enabled(), + "greedy": ( + not self._enable_probabilistic_draft_probs + or sampling_metadata.all_greedy + ), + } + ok = all(checks.values()) + if not ok: + logger.info_once( + "SN10_DRAFT_GRAPH: eager draft loop kept, failed checks: %s", + [k for k, v in checks.items() if not v], + ) + return ok + + def _draft_loop_body( + self, + batch_size: int, + input_batch_size: int, + common_attn_metadata: CommonAttentionMetadata, + num_iters: int, + runtime_mode: CUDAGraphMode, + ) -> None: + """The (k-1) draft iterations, operating only on persistent buffers. + + Same kernels/shapes as the eager loop in `propose` (so the draft tokens + are bit-identical), but: + * loop entry values come from `_loop_pos_in`/`_loop_hid_in`/ + `_loop_tokens[:, 0]` instead of temporaries, + * each iteration's draft token is written into `_loop_tokens[:, i+1]` + instead of being collected in a Python list, + * the forward context carries `runtime_mode` (FULL while capturing: + the piecewise CUDAGraphWrappers then pass through and the whole + forward -- attention custom op included -- is recorded by the outer + torch.cuda.graph; NONE for the eager warm-up run). + """ + block_size = self.block_size + positions = ( + self._loop_pos_in[:, :batch_size] + if self.uses_mrope + else self._loop_pos_in[:batch_size] + ) + hidden_states = self._loop_hid_in[:batch_size] + for token_index in range(num_iters): + # int64 -> int32 cast folded into the copy (eager: .int() + copy). + self.input_ids[:batch_size].copy_( + self._loop_tokens[:batch_size, token_index] + ) + positions = self._update_positions_dependent_metadata( + positions, + common_attn_metadata, + batch_size, + input_batch_size, + block_size, + ) + # Only host-side dataclass construction (fast_build => no AOT + # scheduler kernel); consumed at capture time only. + _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata( + common_attn_metadata, draft_index=token_index + 1 + ) + self.hidden_states[:batch_size].copy_(hidden_states) + if self.supports_mm_inputs: + # Same as the eager loop: text-only drafts through the + # embedding lookup (a gather; captured like any kernel). + self.inputs_embeds[:batch_size].copy_( + self.model.embed_input_ids(self.input_ids[:batch_size]) + ) + model_kwargs = { + "input_ids": None, + "positions": self._get_positions(input_batch_size), + "inputs_embeds": self.inputs_embeds[:input_batch_size], + "hidden_states": self.hidden_states[:input_batch_size], + } + else: + model_kwargs = { + "input_ids": self.input_ids[:input_batch_size], + "positions": self._get_positions(input_batch_size), + "inputs_embeds": None, + "hidden_states": self.hidden_states[:input_batch_size], + } + with set_forward_context( + per_layer_attn_metadata, + self.vllm_config, + num_tokens=input_batch_size, + num_tokens_across_dp=None, + cudagraph_runtime_mode=runtime_mode, + slot_mapping=self._get_slot_mapping(input_batch_size), + ): + ret_hidden_states = self.model(**model_kwargs) + if not self.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 = self._greedy_sample(last_hidden_states[:batch_size]) + self._loop_tokens[:batch_size, token_index + 1].copy_(draft_token_ids) + + def _propose_with_loop_graph( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + first_draft_token_ids: torch.Tensor, + common_attn_metadata: CommonAttentionMetadata, + batch_size: int, + input_batch_size: int, + ) -> torch.Tensor: + k = self.num_speculative_tokens + num_iters = k - 1 + # Stage the loop-entry temporaries into persistent buffers. + if self.uses_mrope: + self._loop_pos_in[:, :batch_size].copy_(positions) + else: + self._loop_pos_in[:batch_size].copy_(positions) + self._loop_hid_in[:batch_size].copy_(hidden_states) + self._loop_tokens[:batch_size, 0].copy_(first_draft_token_ids) + + key = (batch_size, input_batch_size, k) + graph = self._loop_graphs.get(key) + if graph is not None: + graph.replay() + elif key not in self._loop_graph_warm: + # First time this shape is seen: run the very same body eagerly. + # This is the warm-up (Triton JIT / inductor specialisations for + # this padded batch) and produces the step's drafts. + self._loop_graph_warm.add(key) + self._draft_loop_body( + batch_size, + input_batch_size, + common_attn_metadata, + num_iters, + CUDAGraphMode.NONE, + ) + else: + # Second time: capture (kernels are recorded, NOT executed) and + # replay right away so this step still gets its drafts. + # UNTESTED: lazy capture happens after capture_model() finished + # (bypasses validate_cudagraph_capturing_enabled on purpose); + # torch.cuda.graph() does a one-time device synchronize + + # empty_cache before recording. + graph = torch.cuda.CUDAGraph() + pool = current_platform.get_global_graph_pool() + logger.info( + "SN10_DRAFT_GRAPH: capturing draft loop graph for " + "batch %d (padded %d), k=%d", + batch_size, + input_batch_size, + k, + ) + try: + with torch.cuda.graph(graph, pool=pool): + self._draft_loop_body( + batch_size, + input_batch_size, + common_attn_metadata, + num_iters, + CUDAGraphMode.FULL, + ) + except Exception: + # Fail open: this step and all later ones take the eager loop. + logger.warning( + "SN10_DRAFT_GRAPH: capture failed; eager draft loop from now on.", + exc_info=True, + ) + self._loop_graph_enabled = False + self._loop_graphs.clear() + torch.cuda.synchronize() + self._draft_loop_body( + batch_size, + input_batch_size, + common_attn_metadata, + num_iters, + CUDAGraphMode.NONE, + ) + return self._loop_tokens[:batch_size, :k].clone() + self._loop_graphs[key] = graph + graph.replay() + # Clone: the runner keeps the returned tensor across the step + # boundary, and the buffer is rewritten by the next replay. + return self._loop_tokens[:batch_size, :k].clone() + def _update_positions_dependent_metadata( self, positions: torch.Tensor,