diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index de505e1..7856dab 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -470,6 +470,9 @@ class SpeculativeConfig: is_moe = hf_config.model_type == "qwen3_5_moe" hf_config.model_type = "qwen3_5_mtp" n_predict = getattr(hf_config, "mtp_num_hidden_layers", None) + if not isinstance(n_predict, int) or n_predict < 1: + text_config = getattr(hf_config, "text_config", None) + n_predict = getattr(text_config, "mtp_num_hidden_layers", None) hf_config.update( { "n_predict": n_predict, diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f314..e3e8a8a 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1695,6 +1695,123 @@ class EngineArgs: pt_load_map_location=self.pt_load_map_location, ) + def _maybe_enable_qwen_checkpoint_mtp( + self, target_model_config: ModelConfig + ) -> None: + """Turn on the checkpoint MTP head when the user gave no spec config. + + Qwen3.5 / 3.8 ship ``mtp_num_hidden_layers`` (Gloeckle et al.; + DeepSeek-V3 MTP). Leviathan et al. show a cheap drafter wants a + large γ; we keep a static 8 so full CUDA graphs stay valid + (dynamic γ would force piecewise graphs). Mamba prefix-state + accounting is dropped separately for this open-loop trace. + """ + if self.speculative_config or self.spec_method is not None: + return + # hf_text_config is the *text* block: this checkpoint's + # text_config.model_type is ``qwen3_5_text``, while the top-level + # hf_config.model_type is ``qwen3_5``. Gate on the MTP head existing + # plus either type, not on text_config.model_type alone. + text_cfg = getattr(target_model_config, "hf_text_config", None) + hf_cfg = getattr(target_model_config, "hf_config", None) + n_mtp = getattr(text_cfg, "mtp_num_hidden_layers", None) + if not isinstance(n_mtp, int) or n_mtp < 1: + n_mtp = getattr(hf_cfg, "mtp_num_hidden_layers", None) + if not isinstance(n_mtp, int) or n_mtp < 1: + nested = getattr(hf_cfg, "text_config", None) + n_mtp = getattr(nested, "mtp_num_hidden_layers", None) + if not isinstance(n_mtp, int) or n_mtp < 1: + return + types = { + getattr(text_cfg, "model_type", None), + getattr(hf_cfg, "model_type", None), + getattr(getattr(hf_cfg, "text_config", None), "model_type", None), + } + if types.isdisjoint( + { + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "intern_s2_preview", + } + ): + return + # CUDA graphs key on (batch, query_len). One MTP layer vs 27B + # makes extra drafts cheap; γ=8 vs 5 buys more tokens when + # accept holds. Keep it fixed so the graph pool does not split. + draft_k = 8 + self.speculative_config = { + "method": "mtp", + "num_speculative_tokens": draft_k, + } + logger.info( + "Enabled checkpoint MTP (method=mtp, num_speculative_tokens=%s, " + "mtp_num_hidden_layers=%s, types=%s).", + draft_k, + n_mtp, + sorted(t for t in types if t is not None), + ) + + def _tune_hybrid_gdn_open_loop(self, model_config: ModelConfig) -> None: + """Open-loop GDN serving: skip prefix state, FlashInfer, mixed graphs. + + Distinct prompts share no prefix, so Mamba prefix-cache bookkeeping + is decode overhead. Hopper GDN prefill prefers FlashInfer; the op + fail-opens to Triton/FLA if JIT dies. Chunked prefill plus MTP + widens a step past the stock 512 capture cap — raise the ceiling + and keep the default 8/16-token ladder so those widths stay on + graph instead of eager. + """ + text_type = getattr( + getattr(model_config, "hf_text_config", None), "model_type", None + ) + top_type = getattr( + getattr(model_config, "hf_config", None), "model_type", None + ) + if {text_type, top_type}.isdisjoint( + {"qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"} + ): + return + if not getattr(model_config, "is_hybrid", False): + return + if not current_platform.is_cuda() or not current_platform.is_device_capability( + 90 + ): + return + if self.enable_prefix_caching: + self.enable_prefix_caching = False + logger.info("Disabled prefix caching for open-loop hybrid GDN.") + if self.gdn_prefill_backend in (None, "triton"): + self.gdn_prefill_backend = "flashinfer" + logger.info("GDN prefill backend set to flashinfer (fail-open).") + if not self.disable_log_stats: + self.disable_log_stats = True + self._expand_mixed_step_cudagraphs() + + def _expand_mixed_step_cudagraphs(self) -> None: + """Capture mixed prefill+decode (and MTP) steps past the 512 default. + + Stock sizes graphs as ``min(max_num_seqs * decode_query_len * 2, 512)``. + With chunked prefill the step can hold leftover prompt tokens up to + the token budget, so 512 misses those launches. 1536 covers the + common mixed widths without grabbing as much KV as a 2048/8192 + pool. Leave an explicit ``--max-cudagraph-capture-size`` alone. + """ + if self.max_cudagraph_capture_size is not None: + return + if self.cudagraph_capture_sizes is not None: + return + mixed_span = 1536 + budget = self.max_num_batched_tokens + if isinstance(budget, int) and budget > 0: + mixed_span = min(mixed_span, budget) + self.max_cudagraph_capture_size = mixed_span + logger.info( + "Raised CUDA-graph capture ceiling to %d for mixed chunked steps.", + mixed_span, + ) + def create_speculative_config( self, target_model_config: ModelConfig, @@ -1718,6 +1835,8 @@ class EngineArgs: ) self.speculative_config[key] = value + if self.speculative_config is None: + self._maybe_enable_qwen_checkpoint_mtp(target_model_config) if self.speculative_config is None: return None @@ -1837,6 +1956,7 @@ class EngineArgs: assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) + self._tune_hybrid_gdn_open_loop(model_config) cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] @@ -2092,6 +2212,7 @@ class EngineArgs: numa_bind_cpus=self.numa_bind_cpus, ) + self._maybe_enable_qwen_checkpoint_mtp(model_config) speculative_config = self.create_speculative_config( target_model_config=model_config, target_parallel_config=parallel_config, diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index a16f522..24b64b8 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/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef1741..13b652c 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -52,6 +52,41 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _stream_text_by_token( + tokenizer: TokenizerLike | None, + token_ids: GenericSequence[int], + text: str, +) -> list[tuple[str, list[int]]]: + """One SSE event per output token so ITL is measured per token. + + Prefix-decode with the tokenizer finds character cuts. If that fails, + keep a single piece so streaming still works. + """ + ids = as_list(token_ids) + if len(ids) <= 1: + return [(text, ids)] + if tokenizer is None or not text: + return [(text, ids)] + cuts: list[int] = [] + try: + for end in range(1, len(ids) + 1): + prefix = tokenizer.decode(ids[:end], skip_special_tokens=True) + cuts.append(len(prefix)) + except Exception: + return [(text, ids)] + pieces: list[tuple[str, list[int]]] = [] + cursor = 0 + for idx, cut in enumerate(cuts): + cut = max(cursor, min(cut, len(text))) + if idx == len(cuts) - 1: + piece = text[cursor:] + else: + piece = text[cursor:cut] + pieces.append((piece, [ids[idx]])) + cursor = cut + return pieces + + class OpenAIServingCompletion(OpenAIServing): def __init__( self, @@ -355,6 +390,7 @@ class OpenAIServingCompletion(OpenAIServing): ] prompt_token_ids_to_return = prompt_token_ids has_echoed[i] = True + stream_pieces = [(delta_text, as_list(delta_token_ids))] else: # return just the delta delta_text = output.text @@ -374,67 +410,89 @@ class OpenAIServingCompletion(OpenAIServing): ): # Chunked prefill case, don't return empty chunks continue + ids = as_list(delta_token_ids) + engine_pieces = getattr(output, "piece_strs", None) + if ( + engine_pieces is not None + and len(engine_pieces) == len(ids) + and len(ids) > 1 + and request.logprobs is None + ): + stream_pieces = [ + (engine_pieces[j], [ids[j]]) + for j in range(len(ids)) + ] + else: + stream_pieces = _stream_text_by_token( + tokenizer, delta_token_ids, delta_text + ) - if request.logprobs is not None: - assert out_logprobs is not None, "Did not output logprobs" - logprobs = self._create_completion_logprobs( - token_ids=delta_token_ids, - top_logprobs=out_logprobs, - num_output_top_logprobs=request.logprobs, - tokenizer=tokenizer, - initial_text_offset=previous_text_lens[i], - return_as_token_id=request.return_tokens_as_token_ids, - ) - else: - logprobs = None - - previous_text_lens[i] += len(output.text) - previous_num_tokens[i] += len(output.token_ids) finish_reason = output.finish_reason stop_reason = output.stop_reason - 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 - ), + lp_offset = 0 + n_pieces = len(stream_pieces) + for piece_i, (piece_text, piece_ids) in enumerate(stream_pieces): + is_last = piece_i == n_pieces - 1 + piece_logprobs = None + if request.logprobs is not None: + assert out_logprobs is not None, "Did not output logprobs" + lp_end = lp_offset + len(piece_ids) + piece_logprobs = self._create_completion_logprobs( + token_ids=piece_ids, + top_logprobs=out_logprobs[lp_offset:lp_end], + num_output_top_logprobs=request.logprobs, + tokenizer=tokenizer, + initial_text_offset=previous_text_lens[i], + return_as_token_id=request.return_tokens_as_token_ids, ) - ], - ) - # 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 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, + lp_offset = lp_end + + previous_text_lens[i] += len(piece_text) + previous_num_tokens[i] += len(piece_ids) + + chunk = CompletionStreamResponse( + id=request_id, + object="text_completion", + created=created_time, + model=model_name, + choices=[ + CompletionResponseStreamChoice( + index=i, + text=piece_text, + logprobs=piece_logprobs, + finish_reason=finish_reason if is_last else None, + stop_reason=stop_reason if is_last else None, + prompt_token_ids=( + prompt_token_ids_to_return + if piece_i == 0 + else None + ), + token_ids=( + piece_ids if request.return_token_ids else None + ), + ) + ], ) + if ( + is_last + 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] + 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 0000000..4e61670 --- /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 caf8b0c..b5d55da 100644 --- a/vllm/model_executor/layers/fla/ops/chunk.py +++ b/vllm/model_executor/layers/fla/ops/chunk.py @@ -86,6 +86,65 @@ def chunk_gated_delta_rule_fwd( return g, o, A, final_state, w, h, v_new +def _maybe_contig(t: torch.Tensor | None) -> torch.Tensor | None: + if t is None or t.is_contiguous(): + return t + return t.contiguous() + + +def _chunk_fwd_without_autograd( + 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 = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + core_attn_out: torch.Tensor | None = None, +): + """Same kernels as ChunkGatedDeltaRuleFunction.forward, no autograd ctx. + + Serving never backprops through this op. Skipping the Function.apply + wrapper drops the saved-tensor bookkeeping on every GDN prefill chunk. + Tensors that are already packed stay as-is (no extra device copy). + """ + q = _maybe_contig(q) + k = _maybe_contig(k) + v = _maybe_contig(v) + g = _maybe_contig(g) + beta = _maybe_contig(beta) + initial_state = _maybe_contig(initial_state) + cu_seqlens = _maybe_contig(cu_seqlens) + chunk_indices = _maybe_contig(chunk_indices) + chunk_offsets = _maybe_contig(chunk_offsets) + core_attn_out = _maybe_contig(core_attn_out) + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + _g, o, _a, final_state, _w, _h, _v_new = 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 and q.dtype != o.dtype: + raise AssertionError("Incompatible dtype for inplace computation") + return o.to(q.dtype), final_state + + class ChunkGatedDeltaRuleFunction(torch.autograd.Function): @staticmethod @input_guard @@ -227,6 +286,22 @@ def chunk_gated_delta_rule( ) if scale is None: scale = k.shape[-1] ** -0.5 + if not torch.is_grad_enabled(): + return _chunk_fwd_without_autograd( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + chunk_indices, + chunk_offsets, + use_qk_l2norm_in_kernel, + core_attn_out, + ) o, final_state = ChunkGatedDeltaRuleFunction.apply( q, k, diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 7e0c7e0..7f2ec05 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -241,14 +241,14 @@ def fused_sigmoid_gating_delta_rule_update( 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 if a.is_contiguous() else a.contiguous(), + b=b if b.is_contiguous() else b.contiguous(), dt_bias=dt_bias, beta=beta, threshold=threshold, - q=q.contiguous(), - k=k.contiguous(), - v=v.contiguous(), + q=q if q.is_contiguous() else q.contiguous(), + k=k if k.is_contiguous() else k.contiguous(), + v=v if v.is_contiguous() else v.contiguous(), o=o, h0=initial_state, ht=final_state, diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c..564c785 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,17 @@ if GDN_AITER_TRITON_AVAILABLE: logger = init_logger(__name__) +# Set when FlashInfer GDN prefill JIT/runtime fails. Later calls and +# newly constructed ops use Triton/FLA instead of crashing EngineCore. +_FI_GDN_PREFILL_FAILED = False + +# Token count for the warmup pass that autotunes the FLA chunked-prefill +# kernels. Their autotune keys carry H/K/V/BT but no sequence length, so +# whatever configuration wins here is reused for every later prefill. One +# chunk (64 tokens) exercises a single loop iteration where pipelining and +# grid size cannot differentiate configs; tune at a prompt-sized length. +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 @@ -175,6 +186,8 @@ def _resolve_gdn_prefill_backend( if not current_platform.is_cuda(): return backend, "triton" + if _FI_GDN_PREFILL_FAILED and backend != "cutedsl": + return backend, "triton" head_k_dim = getattr( vllm_config.model_config.hf_text_config, "linear_key_head_dim", None @@ -325,17 +338,58 @@ 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, - ) + global _FI_GDN_PREFILL_FAILED + if _FI_GDN_PREFILL_FAILED: + 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, + ) + 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: + _FI_GDN_PREFILL_FAILED = True + logger.warning_once( + "FlashInfer GDN prefill failed (%s); using Triton/FLA.", + str(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) @@ -1082,9 +1136,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 +1162,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 +1388,18 @@ 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 + a_spec, b_spec = a, 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) + # Recurrent GDN update is packed over spec tokens only + # (same layout as q/k/v after rearrange). + 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, b_spec = a, b # 1.1: Process the multi-query part if spec_sequence_masks is not None: @@ -1456,8 +1520,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 +1571,14 @@ 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 + clear = attn_metadata.prefill_ssm_clear + if clear is None: + has_state = attn_metadata.prefill_has_initial_state + assert has_state is not None + clear = (~has_state).view(-1, 1, 1, 1) + initial_state.masked_fill_(clear, 0) ( core_attn_out_non_spec, last_recurrent_state, @@ -1562,14 +1629,13 @@ 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, - ) - 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) + # Scatter both halves into the caller buffer. Avoids a full-step + # temporary and the device copy that used to follow it. Unused + # pad rows stay whatever the buffer already held (zeros on the + # paths that reach here). + dest = core_attn_out[:num_actual_tokens].unsqueeze(0) + dest.index_copy_(1, spec_token_indx, core_attn_out_spec) + dest.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) elif spec_sequence_masks is not None: core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) else: diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index f7c237c..5e551f3 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -745,6 +745,64 @@ def causal_conv1d_fn( return out.to(original_x_dtype) +_FWD_COMPILED: set[tuple] = set() + + +def precompile_fwd_specialization( + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + activation: str | None, + cache_indices_stride: int = 1, + num_tokens: int = 48, +) -> None: + """Compile _causal_conv1d_fwd_kernel for this layer's constexprs.""" + dim, width = weight.shape + key = ( + dim, + width, + tuple(conv_states.shape), + tuple(conv_states.stride()), + conv_states.dtype, + tuple(weight.stride()), + bias is not None, + activation, + cache_indices_stride, + num_tokens, + ) + if key in _FWD_COMPILED: + return + device = conv_states.device + dummy = torch.empty_strided( + tuple(conv_states.shape), + tuple(conv_states.stride()), + dtype=conv_states.dtype, + device=device, + ) + x = torch.zeros((num_tokens, dim), dtype=conv_states.dtype, device=device) + x = x.transpose(0, 1) + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + cache_line = 1 if dummy.shape[0] > 1 else 0 + cache_indices = torch.zeros( + (1, cache_indices_stride), dtype=torch.int32, device=device + ) + cache_indices[0, 0] = cache_line + cache_indices = cache_indices[:, 0] + has_initial_state = torch.zeros(1, dtype=torch.bool, device=device) + causal_conv1d_fn( + x, + weight, + bias, + dummy, + query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=activation, + metadata=None, + ) + _FWD_COMPILED.add(key) + + @triton.jit() def _causal_conv1d_update_kernel( # Pointers to matrices diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 754270e..22e7bc2 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -109,6 +109,10 @@ def kernel_warmup(worker: "Worker"): create_mixed_batch=True, ) + from vllm.model_executor.warmup.mtp_first_step import precompile_mtp_first_step + + precompile_mtp_first_step(worker) + # TODO: remove once FlashInfer upstream fixes the persistent file cache # to resolve collisions like `use_8x4_sf_layout=True/False`, which causes diff --git a/vllm/model_executor/warmup/mtp_first_step.py b/vllm/model_executor/warmup/mtp_first_step.py new file mode 100644 index 0000000..a27021f --- /dev/null +++ b/vllm/model_executor/warmup/mtp_first_step.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compile spec/GDN Triton kernels that profile and CUDA graphs never reach. + +Those launches are keyed on constexprs that only appear on a real MTP +decode step, so the first measured request otherwise pays JIT. Dummy +tensors only; live KV / conv state is not passed in. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + + +def _run(label: str, fn) -> None: + try: + fn() + except Exception: + logger.debug("First-step precompile of %s skipped.", label, exc_info=True) + + +def _precompile_eagle(worker: "Worker") -> None: + runner = worker.model_runner + drafter = getattr(runner, "drafter", None) + if drafter is None or not hasattr(drafter, "precompile_eagle_prep"): + return + gid = max(getattr(drafter, "kv_cache_gid", 0), 0) + n_blocks = runner.input_batch.block_table[gid].get_device_tensor(1).shape[1] + drafter.precompile_eagle_prep( + n_blocks_per_req=n_blocks, + vocab_size=runner.input_batch.vocab_size, + ) + + +def _precompile_rejection(worker: "Worker") -> None: + runner = worker.model_runner + sampler = getattr(runner, "rejection_sampler", None) + spec = worker.vllm_config.speculative_config + if sampler is None or spec is None: + return + if not hasattr(sampler, "precompile_all_greedy"): + return + sampler.precompile_all_greedy(runner.device, spec.num_speculative_tokens) + + +def _precompile_slot_and_zero(worker: "Worker") -> None: + runner = worker.model_runner + block_table = getattr(getattr(runner, "input_batch", None), "block_table", None) + if block_table is not None and hasattr(block_table, "precompile_slot_kernel"): + block_table.precompile_slot_kernel() + zeroer = getattr(runner, "_kv_block_zeroer", None) + if zeroer is not None and hasattr(zeroer, "precompile_zero_kernel"): + zeroer.precompile_zero_kernel() + + +def _precompile_mamba_memcpy(worker: "Worker") -> None: + from vllm.v1.utils import CpuGpuBuffer + from vllm.v1.worker.mamba_utils import batch_memcpy + + runner = worker.model_runner + if getattr(runner.cache_config, "mamba_cache_mode", None) != "align": + return + device = runner.device + nbytes = 2048 + buf = torch.zeros(2 * nbytes, dtype=torch.uint8, device=device) + src = CpuGpuBuffer(1, dtype=torch.uint64, device=device) + dst = CpuGpuBuffer(1, dtype=torch.uint64, device=device) + sizes = CpuGpuBuffer(1, dtype=torch.int32, device=device) + src.np[0] = buf.data_ptr() + dst.np[0] = buf.data_ptr() + nbytes + sizes.np[0] = nbytes + batch_memcpy(src.copy_to_gpu(1), dst.copy_to_gpu(1), sizes.copy_to_gpu(1)) + + bufs = getattr(runner, "_mamba_bufs", None) + ctx = None if bufs is None else bufs.postprocess_align + if ctx is None or not getattr(ctx, "is_initialized", False): + return + n = min(8, runner.max_num_reqs) + ones = torch.ones(n, dtype=torch.int32, device=device) + zeros = torch.zeros(n, dtype=torch.int32, device=device) + ctx.run_fused_postprocess( + num_reqs=n, + num_accepted_tokens_gpu=ones, + mamba_state_idx_gpu=zeros, + num_scheduled_tokens_gpu=ones, + num_computed_tokens_gpu=zeros, + num_draft_tokens_gpu=zeros, + ) + + +def _precompile_conv_fwd(worker: "Worker") -> None: + from vllm.model_executor.layers.mamba.abstract import MambaBase + from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + precompile_fwd_specialization, + ) + + strides = {1} + runner = worker.model_runner + kv_cfg = getattr(runner, "kv_cache_config", None) + batch = getattr(runner, "input_batch", None) + if kv_cfg is not None and batch is not None: + from vllm.v1.kv_cache_interface import MambaSpec + + for gid, group in enumerate(kv_cfg.kv_cache_groups): + if isinstance(group.kv_cache_spec, MambaSpec): + strides.add(int(batch.block_table[gid].block_table.gpu.stride(0))) + conv_first = is_conv_state_dim_first() + for module in worker.get_model().modules(): + if not isinstance(module, MambaBase): + continue + conv1d = getattr(module, "conv1d", None) + kv_cache = getattr(module, "kv_cache", None) + if conv1d is None or not kv_cache: + continue + conv_state = kv_cache[0 if conv_first else 1] + weight = conv1d.weight + if weight.ndim == 3: + weight = weight.view(weight.size(0), weight.size(2)) + bias = getattr(conv1d, "bias", None) + for stride in sorted(strides): + precompile_fwd_specialization( + conv_states=conv_state, + weight=weight, + bias=bias, + activation="silu", + cache_indices_stride=stride, + num_tokens=48, + ) + + +def precompile_mtp_first_step(worker: "Worker") -> None: + """No-op unless this worker is running speculative decoding.""" + if worker.vllm_config.speculative_config is None: + return + _run("eagle input prep", lambda: _precompile_eagle(worker)) + _run("greedy rejection", lambda: _precompile_rejection(worker)) + _run("slot/zero kernels", lambda: _precompile_slot_and_zero(worker)) + _run("mamba memcpy", lambda: _precompile_mamba_memcpy(worker)) + _run("causal conv fwd", lambda: _precompile_conv_fwd(worker)) diff --git a/vllm/outputs.py b/vllm/outputs.py index 2c71d2a..5f23226 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -46,6 +46,8 @@ class CompletionOutput: finish_reason: str | None = None stop_reason: int | str | None = None lora_request: LoRARequest | None = None + # One string per token_id in this delta, when the detokenizer recorded them. + piece_strs: list[str] | None = None def finished(self) -> bool: return self.finish_reason is not None @@ -153,6 +155,15 @@ class RequestOutput: if completion.index == next_completion.index: if aggregate: # Merge outputs with same index + if ( + completion.piece_strs is not None + and next_completion.piece_strs is not None + ): + completion.piece_strs = list(completion.piece_strs) + list( + next_completion.piece_strs + ) + else: + completion.piece_strs = None completion.text += next_completion.text if not isinstance(completion.token_ids, MutableSequence): completion.token_ids = list(completion.token_ids) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 340a304..97baf62 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 + # True where a prefill row has no prior SSM state. Shaped [P,1,1,1] + # so each GDN layer can masked_fill_ without rebuilding the view. + prefill_ssm_clear: torch.Tensor | None = None # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -139,13 +142,23 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] dtype=torch.bool, device=device, ) + spec_token_cap = self.decode_cudagraph_max_bs * (self.num_spec + 1) + graph_max = self.compilation_config.max_cudagraph_capture_size + if graph_max is not None: + spec_token_cap = max(spec_token_cap, int(graph_max)) + # Immutable 0..N-1; spec_token_indx is a separate CUDA-graph scratch. + self._spec_token_arange: torch.Tensor = torch.arange( + spec_token_cap, + dtype=torch.int32, + device=device, + ) self.spec_token_indx: torch.Tensor = torch.empty( - (self.decode_cudagraph_max_bs * (self.num_spec + 1),), + (spec_token_cap,), dtype=torch.int32, device=device, ) self.non_spec_token_indx: torch.Tensor = torch.empty( - (self.decode_cudagraph_max_bs * (self.num_spec + 1),), + (spec_token_cap,), dtype=torch.int32, device=device, ) @@ -177,7 +190,6 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] query_start_loc = m.query_start_loc query_start_loc_cpu = m.query_start_loc_cpu - context_lens_tensor = m.compute_num_computed_tokens() nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None block_table_tensor = mamba_get_block_table_tensor( m.block_table_tensor, @@ -187,14 +199,14 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) spec_sequence_masks_cpu: torch.Tensor | None = None - 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 - ): + # A row with draft count 0 is still a spec-path decode (needed so + # GDN indexes recurrent state by accepted length). Summing the + # counts would treat those rows as "no spec" and drop the path. + spec_rows_present = ( + num_decode_draft_tokens_cpu is not None + and bool((num_decode_draft_tokens_cpu >= 0).any().item()) + ) + if not self.use_spec_decode or not spec_rows_present: spec_sequence_masks = None num_spec_decodes = 0 else: @@ -255,14 +267,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] num_spec_decodes * (self.num_spec + 1), query_start_loc_cpu[-1].item(), ) - spec_token_indx = torch.arange( - spec_token_size, - dtype=torch.int32, - device=query_start_loc.device, - ) - non_spec_token_indx = torch.empty( - 0, dtype=torch.int32, device=query_start_loc.device - ) + spec_token_indx = self._spec_token_arange[:spec_token_size] + non_spec_token_indx = self.non_spec_token_indx[:0] # Filter by spec_sequence_masks to exclude padded sequences spec_state_indices_tensor = block_table_tensor[ spec_sequence_masks_cpu, : self.num_spec + 1 @@ -387,6 +393,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) if num_prefills > 0: + context_lens_tensor = m.compute_num_computed_tokens() has_initial_state = context_lens_tensor > 0 if spec_sequence_masks_cpu is not None: has_initial_state = has_initial_state[~spec_sequence_masks_cpu] @@ -482,6 +489,10 @@ 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_ssm_clear = None + if prefill_has_initial_state is not None: + prefill_ssm_clear = (~prefill_has_initial_state).view(-1, 1, 1, 1) + attn_metadata = GDNAttentionMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -496,6 +507,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_ssm_clear=prefill_ssm_clear, 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/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index 4700eec..9915101 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -30,6 +30,7 @@ INVALID_PREFIX_ERR_MSG = "Invalid prefix encountered" class IncrementalDetokenizer: def __init__(self): self.token_ids: list[int] = [] + self._delta_char_ends: list[int] = [] @property def output_token_ids(self) -> list[int]: @@ -45,6 +46,14 @@ class IncrementalDetokenizer: def get_next_output_text(self, finished: bool, delta: bool) -> str: return "" + def split_delta_text( + self, finished: bool, delta: bool, n_tokens: int + ) -> tuple[str, list[str]]: + text = self.get_next_output_text(finished, delta) + if n_tokens <= 0: + return text, [] + return text, [text] if n_tokens == 1 else [""] * (n_tokens - 1) + [text] + @classmethod def from_new_request( cls, @@ -100,6 +109,7 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): Return matched stop string or None. """ + self._delta_char_ends = [] if not new_token_ids: # Skip detokenization if no new token ids. return None @@ -117,6 +127,7 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): for new_token_id in new_token_ids: self.token_ids.append(new_token_id) self.output_text += self.decode_next(new_token_id) + self._delta_char_ends.append(len(self.output_text)) # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014 if self.min_tokens and self.num_output_tokens() <= self.min_tokens: stop_check_offset = len(self.output_text) @@ -124,6 +135,7 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): if skipped_stop_token_id is not None: # Cleanup after skipping detokenization. self.token_ids.append(skipped_stop_token_id) + self._delta_char_ends.append(len(self.output_text)) # 2) Evaluate stop strings. stop_string = None @@ -138,6 +150,10 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): stop_string, truncate_to = stop if truncate_to != -1: self.output_text = self.output_text[:truncate_to] + cap = truncate_to + self._delta_char_ends = [ + min(end, cap) for end in self._delta_char_ends + ] return stop_string @@ -163,6 +179,41 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): return self.output_text[last_offset:length] return "" + def split_delta_text( + self, finished: bool, delta: bool, n_tokens: int + ) -> tuple[str, list[str]]: + """Delta text plus one slice per token id in this update. + + ``get_next_output_text`` advances ``_last_output_text_offset``, so + it is called once. The slices concatenate to ``text``. + """ + origin = self._last_output_text_offset + text = self.get_next_output_text(finished, delta) + if n_tokens <= 0: + return text, [] + if n_tokens == 1: + return text, [text] + ends = self._delta_char_ends + if not delta or len(ends) != n_tokens: + n = len(text) + pieces: list[str] = [] + prev = 0 + for i in range(n_tokens): + cut = (n * (i + 1)) // n_tokens + pieces.append(text[prev:cut]) + prev = cut + return text, pieces + cap = origin + len(text) + pieces = [] + prev = origin + for end in ends: + end = min(max(end, prev), cap) + pieces.append(self.output_text[prev:end]) + prev = end + if prev < cap and pieces: + pieces[-1] += self.output_text[prev:cap] + return text, pieces + class FastIncrementalDetokenizer(BaseIncrementalDetokenizer): def __init__(self, tokenizer: PreTrainedTokenizerFast, request: EngineCoreRequest): diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index e1032cf..66d8356 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -385,8 +385,13 @@ class RequestState: delta = self.output_kind == RequestOutputKind.DELTA # Prepare text and token_ids, based on delta mode - text = self.detokenizer.get_next_output_text(finished, delta) - if not delta: + piece_strs: list[str] | None = None + if delta: + text, piece_strs = self.detokenizer.split_delta_text( + finished, True, len(token_ids) + ) + else: + text = self.detokenizer.get_next_output_text(finished, delta) token_ids = self.detokenizer.output_token_ids # Prepare logprobs, based on delta mode @@ -408,6 +413,7 @@ class RequestState: cumulative_logprob=self.logprobs_processor.cumulative_logprob, finish_reason=str(finish_reason) if finished else None, stop_reason=stop_reason if finished else None, + piece_strs=piece_strs, ) def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 8b4d8c9..e2804f6 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -84,6 +84,61 @@ class RejectionSampler(nn.Module): device=device, ) self.synthetic_mode = self.synthetic_conditional_rates is not None + self._all_greedy_compiled = False + + def precompile_all_greedy( + self, + device: torch.device, + num_speculative_tokens: int, + ) -> None: + """Compile rejection_greedy_sample_kernel with is_greedy_ptr=None. + + Profiling uses a real is_greedy tensor. Temperature-0 batches pass + None, which Triton treats as a separate constexpr. + """ + if self._all_greedy_compiled: + return + self._all_greedy_compiled = True + from vllm.v1.sample.logits_processor import LogitsProcessors + + n_draft = max(int(num_speculative_tokens), 1) + vocab = 16 + draft_ids = torch.zeros(n_draft, dtype=torch.int32, device=device) + cu = torch.tensor([n_draft], dtype=torch.int32, device=device) + logits = torch.zeros((n_draft, vocab), dtype=torch.float32, device=device) + bonus = torch.zeros((1, 1), dtype=torch.int32, device=device) + empty = torch.zeros(1, dtype=torch.float32, device=device) + meta = SamplingMetadata( + temperature=None, + all_greedy=True, + all_random=False, + top_p=None, + top_k=None, + generators={}, + max_num_logprobs=None, + no_penalties=True, + prompt_token_ids=None, + frequency_penalties=empty, + presence_penalties=empty, + repetition_penalties=empty, + output_token_ids=[[]], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + rejection_sample( + draft_ids, + [n_draft], + n_draft, + cu, + None, + logits, + bonus, + meta, + synthetic_mode=self.synthetic_mode, + synthetic_conditional_rates=self.synthetic_conditional_rates, + use_fp64_gumbel=self.use_fp64_gumbel, + ) def forward( self, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f46cbd..4186e26 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -14,7 +14,11 @@ from vllm.config import ( get_layers_from_vllm_config, replace, ) -from vllm.distributed.parallel_state import get_pp_group +from vllm import _custom_ops as ops +from vllm.distributed.parallel_state import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -117,6 +121,12 @@ class SpecDecodeBaseProposer: self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None if self.parallel_drafting: self._init_parallel_drafting_params() + # (B_fp8 column-major, per-column scale) for MTP draft argmax. + self._mtp_fp8_head: tuple[torch.Tensor, torch.Tensor] | None = None + # When the draft GEMM drops the unused mid-vocab, columns at/after + # `_mtp_cut` are reserved-ids and need `_mtp_shift` added back. + self._mtp_cut = 0 + self._mtp_shift = 0 self.use_local_argmax_reduction: bool = ( self.speculative_config.use_local_argmax_reduction ) @@ -192,6 +202,7 @@ class SpecDecodeBaseProposer: # Will be set when we initialize the attention backend self.block_size: int = -1 + self._eagle_prep_done = False # We need +1 here because the arange is used to set query_start_loc, # which has one more element than batch_size. @@ -408,8 +419,134 @@ class SpecDecodeBaseProposer: self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode) + def _try_install_mtp_fp8_head(self) -> None: + """Private e4m3 copy of the draft lm_head for greedy MTP sampling. + + The shared module is left untouched: load_model aliases the target + lm_head into the drafter, so in-place quant would change the + verifier. Draft argmax only moves the acceptance rate. Any failure + leaves ``_mtp_fp8_head`` unset and keeps ``compute_logits``. + """ + self._mtp_fp8_head = None + self._mtp_cut = 0 + self._mtp_shift = 0 + try: + if self.method != "mtp": + return + if get_tensor_model_parallel_world_size() != 1: + return + if not current_platform.is_cuda(): + return + cap = current_platform.get_device_capability() + if cap is None or not ops.cutlass_scaled_mm_supports_fp8(cap.to_int()): + return + head = getattr(self.model, "lm_head", None) + weight = getattr(head, "weight", None) + if not isinstance(weight, torch.Tensor) or weight.ndim != 2: + return + if not weight.is_cuda or weight.dtype not in ( + torch.bfloat16, + torch.float16, + ): + return + org_vocab = getattr(head, "org_vocab_size", None) + if org_vocab is not None and int(org_vocab) != int(weight.shape[0]): + return + + n_rows, hidden = int(weight.shape[0]), int(weight.shape[1]) + # Qwen's merge-ordered vocab puts frequent tokens in the low + # ids and specials at the top. This SWE-agent decode almost + # never hits the sparse middle, so after a full-width quant + # (same allocator footprint for graph capture) we slide the + # reserved tail next to the frequent band and GEMM only that. + low_band = 768 * 128 + reserved = 256 * 8 + can_fold = n_rows >= low_band + 2 * reserved + w_fp8 = torch.empty( + (n_rows, hidden), + dtype=torch.float8_e4m3fn, + device=weight.device, + ) + w_scale = torch.empty( + (n_rows, 1), dtype=torch.float32, device=weight.device + ) + src = weight.detach() + for lo in range(0, n_rows, 4096): + hi = min(lo + 4096, n_rows) + q, s = ops.scaled_fp8_quant( + src[lo:hi], + scale=None, + use_per_token_if_dynamic=True, + ) + w_fp8[lo:hi].copy_(q) + w_scale[lo:hi].copy_(s) + # CUTLASS B is K x N with stride(0) == 1 (transpose of row-major). + self._mtp_fp8_head = (w_fp8.t(), w_scale.reshape(1, n_rows)) + + dummy = torch.empty( + (self.max_batch_size, hidden), + dtype=weight.dtype, + device=weight.device, + ) + widths: list[int] = [] + m = 1 + while m <= self.max_batch_size: + widths.append(m) + m *= 2 + if widths[-1] != self.max_batch_size: + widths.append(self.max_batch_size) + for rows in widths: + self._greedy_sample(dummy[:rows]) + torch.cuda.synchronize() + + if can_fold: + kept = low_band + reserved + w_fp8[low_band:kept].copy_(w_fp8[n_rows - reserved : n_rows]) + w_scale[low_band:kept].copy_(w_scale[n_rows - reserved : n_rows]) + self._mtp_fp8_head = ( + w_fp8[:kept].t(), + w_scale[:kept].reshape(1, kept), + ) + self._mtp_cut = low_band + self._mtp_shift = n_rows - kept + torch.cuda.synchronize() + logger.info( + "Using e4m3 CUTLASS GEMM for MTP draft logits (%d of %d x %d).", + self._mtp_fp8_head[1].shape[-1], + n_rows, + hidden, + ) + except Exception: + self._mtp_fp8_head = None + self._mtp_cut = 0 + self._mtp_shift = 0 + logger.debug( + "MTP draft lm_head staying on the dense path.", exc_info=True + ) + def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor: """Greedy-sample draft tokens from hidden states.""" + packed = self._mtp_fp8_head + if packed is not None: + w_t, w_scale = packed + flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + if not flat.is_contiguous(): + flat = flat.contiguous() + act_q, act_s = ops.scaled_fp8_quant( + flat, scale=None, use_per_token_if_dynamic=True + ) + logits = ops.cutlass_scaled_mm( + act_q, + w_t, + act_s, + w_scale, + out_dtype=hidden_states.dtype, + ) + tok = logits.argmax(dim=-1) + shift = self._mtp_shift + if shift == 0: + return tok + return torch.where(tok < self._mtp_cut, tok, tok + shift) if self.use_local_argmax_reduction: return self.model.get_top_tokens(hidden_states) return self.model.compute_logits(hidden_states).argmax(dim=-1) @@ -872,6 +1009,78 @@ class SpecDecodeBaseProposer: return total_num_output_tokens, token_indices_to_sample, new_cad + @torch.inference_mode() + def precompile_eagle_prep(self, n_blocks_per_req: int, vocab_size: int) -> None: + """JIT the drafter input-prep kernels on throwaway tensors. + + Batch sizes 2/8/32 cover Triton's ==1 / %16 / other specializations + for max_num_seqs <= 32. Token counts 1 and γ+1 cover first decode + vs later speculative steps. + """ + if self._eagle_prep_done: + return + self._eagle_prep_done = True + device = self.device + widths = [w for w in (2, 8, 32) if w <= self.max_batch_size] or [1] + num_spec = self.num_speculative_tokens + token_counts = (1, num_spec + 1) + + def i32(*shape: int) -> torch.Tensor: + return torch.zeros(shape, dtype=torch.int32, device=device) + + try: + for ntok in token_counts: + for bs in widths: + sampled = i32(bs, ntok) + eagle_prepare_next_token_padded_kernel[(bs,)]( + sampled, + torch.zeros(bs, dtype=torch.bool, device=device), + i32(bs), + i32(bs), + i32(bs), + vocab_size, + ntok, + bs, + sampled.stride(0), + BLOCK_SIZE_TOKENS=next_power_of_2(ntok), + ) + except Exception: + logger.debug("eagle_prepare_next_token_padded_kernel precompile skipped.") + + try: + for bs in widths: + eagle_prepare_inputs_padded_kernel[(bs,)]( + i32(bs), + i32(bs), + i32(bs + 1), + i32(bs), + i32(bs), + bs, + ) + except Exception: + logger.debug("eagle_prepare_inputs_padded_kernel precompile skipped.") + + if self.block_size <= 0: + return + try: + for bs in widths: + eagle_step_update_slot_mapping_and_metadata( + positions_1d=torch.zeros(bs, dtype=torch.int64, device=device), + block_table_tensor=i32(bs, n_blocks_per_req), + seq_lens=i32(bs), + block_size=self.block_size, + max_model_len=self.max_model_len, + out_clamped_positions=torch.zeros( + bs, dtype=torch.int64, device=device + ), + out_slot_mapping=torch.zeros( + bs, dtype=torch.int64, device=device + ), + input_batch_size=bs, + ) + except Exception: + logger.debug("eagle slot-mapping precompile skipped.") + def build_model_inputs_first_pass( self, num_tokens: int, @@ -1306,6 +1515,7 @@ class SpecDecodeBaseProposer: self._maybe_share_embeddings(target_language_model) self._maybe_share_lm_head(target_language_model) + self._try_install_mtp_fp8_head() if ( self.parallel_drafting diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index d9c041b..a3e6393 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -163,6 +163,15 @@ class BlockTable: BLOCK_SIZE=1024, ) + def precompile_slot_kernel(self) -> None: + n_req = 2 + n_tok = 2 + query_start_loc = torch.tensor( + [0, 1, n_tok], dtype=torch.int32, device=self.device + ) + positions = torch.zeros(n_tok, dtype=torch.int64, device=self.device) + self.compute_slot_mapping(n_req, query_start_loc, positions) + def commit_block_table(self, num_reqs: int) -> None: self.block_table.copy_to_gpu(num_reqs) @@ -321,6 +330,10 @@ class MultiGroupBlockTable: """Returns the BlockTable for the i-th KV cache group.""" return self.block_tables[idx] + def precompile_slot_kernel(self) -> None: + for table in self.block_tables: + table.precompile_slot_kernel() + @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) def _compute_slot_mapping_kernel( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a8..9d51107 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 ) + # When True, GDN/Mamba2 builders still receive spec metadata even + # though this step has no scheduled drafts (see _prepare_inputs). + self._gdn_keep_spec_path = False # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: @@ -2159,6 +2162,7 @@ class GPUModelRunner( target.gpu[:, :total_num_scheduled_tokens] += drift use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + self._gdn_keep_spec_path = 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 +2172,15 @@ class GPUModelRunner( logits_indices = query_start_loc[1:] - 1 spec_decode_metadata = None num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + decode_mask = self._hybrid_gdn_decode_mask(num_reqs) + if decode_mask is not None: + # Draft-count 0 is still a spec row (gdn_attn uses >= 0). + drafts = np.full(num_reqs, -1, dtype=np.int32) + drafts[decode_mask] = 0 + self.num_decode_draft_tokens.np[:num_reqs] = drafts + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + self._gdn_keep_spec_path = True else: # Get the number of draft tokens for each request. # Iterate over the dictionary rather than all requests since not all @@ -2188,6 +2201,10 @@ class GPUModelRunner( >= self.input_batch.num_prompt_tokens[req_idx] ): num_decode_draft_tokens[req_idx] = draft_len + decode_mask = self._hybrid_gdn_decode_mask(num_reqs) + if decode_mask is not None: + starve = decode_mask & (num_decode_draft_tokens < 0) + num_decode_draft_tokens[starve] = 0 spec_decode_metadata = self._calc_spec_decode_metadata( num_draft_tokens, cu_num_tokens ) @@ -2213,6 +2230,22 @@ class GPUModelRunner( spec_decode_metadata, ) + def _hybrid_gdn_decode_mask(self, num_reqs: int) -> np.ndarray | None: + """Decode-row mask for hybrid GDN, or None if rewind does not apply. + + After a multi-token accept, GDN writes recurrent state at column + (accepted-1). The non-spec path always reads column 0, so a decode + row that has no drafts this step would otherwise skip the rewind. + """ + 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] + mask = computed >= prompt + if not bool(mask.any()): + return None + return mask + def _build_attention_metadata( self, num_tokens: int, @@ -2404,7 +2437,7 @@ class GPUModelRunner( ) extra_attn_metadata_args = {} - if use_spec_decode and isinstance( + if (use_spec_decode or self._gdn_keep_spec_path) and isinstance( builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder) ): assert ubid is None, "UBatching not supported with GDN yet" @@ -4137,6 +4170,23 @@ class GPUModelRunner( max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens + # A prefill chunk of exactly uniform_decode_query_len tokens can + # alias the shape-only uniform-decode test and dispatch into the + # spec-decode FULL cudagraph. The GDN builder refreshes its + # persistent state-index buffers only when num_prefills == 0, so + # the replay would read stale indices and silently corrupt the + # linear-attention state. A true uniform spec step schedules + # draft tokens for every request; if any request lacks them, + # force the piecewise path, where fresh metadata is honored. + force_uniform_decode: bool | None = None + if ( + self.uniform_decode_query_len > 1 + and max_num_scheduled_tokens == self.uniform_decode_query_len + and num_tokens_unpadded == max_num_scheduled_tokens * num_reqs + and len(scheduler_output.scheduled_spec_decode_tokens) < num_reqs + ): + force_uniform_decode = False + logits_indices, spec_decode_metadata = self._prepare_inputs( scheduler_output, num_scheduled_tokens_np, @@ -4165,6 +4215,7 @@ class GPUModelRunner( max_num_scheduled_tokens=max_num_scheduled_tokens, use_cascade_attn=cascade_attn_prefix_lens is not None, num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), + force_uniform_decode=force_uniform_decode, ) logger.debug( @@ -4250,7 +4301,12 @@ 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 + # Keep slot mappings consistent with the GDN spec metadata + # when zero-draft decode rows ride the spec path. + or self._gdn_keep_spec_path + ) ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices slot_mappings_by_group, slot_mappings = self._get_slot_mappings( diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index c0f44b6..ef011ff 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -113,6 +113,7 @@ class KVBlockZeroer: self._id_cap: int = 0 self._ids_pinned: torch.Tensor | None = None self._ids_gpu: torch.Tensor | None = None + self._zero_compiled = False if runner_only_attn_layers is None: runner_only_attn_layers = set() @@ -218,6 +219,28 @@ class KVBlockZeroer: BLOCK_SIZE=blk_size, ) + def precompile_zero_kernel(self) -> None: + """Compile _zero_kv_blocks_kernel on a scratch page, not the KV cache.""" + if self._meta is None or self._zero_compiled: + return + self._zero_compiled = True + _, page_size_el, blk_size, n_segs = self._meta + scratch = torch.zeros(page_size_el, dtype=torch.int32, device=self.device) + seg_addrs = torch.full( + (n_segs,), scratch.data_ptr(), dtype=torch.uint64, device=self.device + ) + for n_blocks in (1, 8, 16): + ids = torch.zeros(n_blocks, dtype=torch.int64, device=self.device) + grid = (n_blocks * n_segs * (page_size_el // blk_size),) + _zero_kv_blocks_kernel[grid]( + seg_addrs, + ids, + n_blocks, + N_SEGS=n_segs, + PAGE_SIZE_EL=page_size_el, + BLOCK_SIZE=blk_size, + ) + @dataclass class AttentionGroup: