diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..7b0bcd7c0 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -134,6 +134,77 @@ 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, so depth keeps paying well past +# the point where per-position acceptance starts to fall. +BUNDLED_MTP_NUM_SPECULATIVE_TOKENS: int = 7 + +# 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 = 64 + + +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_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 +1766,72 @@ 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 for the linear-attention state stores recurrent + states at block boundaries and reconciles them against speculative + state slots on every step. That bookkeeping is pure overhead for + requests that share no prefix, so it stays off unless requested + through an explicit cache mode. + * 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.enable_prefix_caching and self.mamba_cache_mode == "none": + self.enable_prefix_caching = False + logger.info( + "Prefix caching disabled for the bundled MTP head: the hybrid " + "state cache adds per-step bookkeeping without reuse." + ) + 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) + self.cudagraph_capture_sizes = sizes + logger.info( + "CUDA-graph capture sizes: %d sizes up to %d tokens.", + len(sizes), + sizes[-1], + ) + if not self.disable_log_stats: + self.disable_log_stats = True def create_speculative_config( self, target_model_config: ModelConfig, @@ -1837,6 +1974,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..4e616704e --- /dev/null +++ b/vllm/entrypoints/openai/serving_warmup.py @@ -0,0 +1,131 @@ +# 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], ...] = ( + (1, 0.0, 640), + (2, 0.05, 900), + (3, 0.03, 700), + (5, 0.02, 1100), + (16, 0.0, 256), + (32, 0.0, 128), + (1, 0.0, 1200), + (2, 0.15, 800), +) +_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 + for phase_index, (width, stagger, length) 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 = [] + for i in range(width): + seq += 1 + token_ids = _prompt_token_ids(length + (i % 3) * 7, seq, 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 + 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/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/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/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a823..0aac492ed 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, @@ -4250,7 +4297,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(