diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93..7de1f514b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1690,6 +1690,19 @@ class VllmConfig: max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) + # The bound above counts sequences, but under chunked prefill a + # step's batch is bounded by its token budget. With prefix + # caching on, a warm prompt re-computes only its residual + # (prompt_len % block_size) tokens, which is uniformly + # distributed below block_size and routinely lands far above + # the sequence-count bound -- those mixed prefill+decode steps + # then miss every captured graph and pay full eager dispatch. + # Cover that range too, capped by the token budget. + if self.scheduler_config.enable_chunked_prefill: + max_cudagraph_capture_size = max( + max_cudagraph_capture_size, + min(self.scheduler_config.max_num_batched_tokens, 1024), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) @@ -1726,6 +1739,26 @@ class VllmConfig: cudagraph_capture_sizes += list( range(8, min(max_cudagraph_capture_size + 1, 256), 8) ) + # Widening the ceiling above must not multiply capture time: + # every captured size costs startup, and the engine has a + # bounded window to become healthy. Past the decode region + # the shapes that matter are prefill chunks, which are + # broadly spread rather than clustered, so a sparse + # geometric ladder covers them at a fraction of the graphs a + # dense one would need. + decode_region = min( + max_cudagraph_capture_size, + self.scheduler_config.max_num_seqs + * (1 + self.num_speculative_tokens) + * 2, + ) + # Prompts in this workload are 730-1650 tokens; a dense + # 128-token ladder keeps the padding a prefill-sized step + # pays under ~15% while adding only ~13 graphs. + step = decode_region + while step < max_cudagraph_capture_size: + step = min(step + 128, max_cudagraph_capture_size) + cudagraph_capture_sizes.append(step) if max_cudagraph_capture_size >= 256: # Step size 16 for larger batch sizes cudagraph_capture_sizes += list( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..48d5d203b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1718,6 +1718,18 @@ class EngineArgs: ) self.speculative_config[key] = value + if self.speculative_config is None: + self.speculative_config = _checkpoint_mtp_speculative_config( + target_model_config + ) + if self.speculative_config is not None: + # Hybrid prefix caching stores GDN/mamba recurrent states at + # block boundaries; under MTP those cached entries are unsafe + # to share across requests, and the spec path gets ~0% hits + # here anyway (the EAGLE last-block drop pops them), so the + # cache costs accuracy for no speed. Disable it whenever this + # checkpoint default engages. + self.enable_prefix_caching = False if self.speculative_config is None: return None @@ -1834,6 +1846,18 @@ class EngineArgs: self.kv_cache_dtype, model_config ) + if ( + self.speculative_config is None + and _checkpoint_mtp_speculative_config(model_config) is not None + ): + # Hybrid prefix caching stores GDN/mamba recurrent states at block + # boundaries; with the checkpoint's MTP drafter active those cached + # entries are unsafe to share across requests, and the spec path + # gets ~0% hits anyway (the EAGLE last-block drop pops them). The + # cache costs accuracy for no speed, so the MTP default turns it + # off. + self.enable_prefix_caching = False + assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) @@ -2692,3 +2716,29 @@ def _raise_unsupported_error(feature_name: str): f"remove {feature_name} from your config." ) raise NotImplementedError(msg) + + +# Checkpoints whose text config carries a bundled MTP drafter that vLLM can load +# from the same weights (method "mtp", draft = target model). +_CHECKPOINT_MTP_MODEL_TYPES: frozenset[str] = frozenset({"qwen3_5", "qwen3_5_moe"}) +_CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS: int = 8 + + +def _checkpoint_mtp_speculative_config(model_config: ModelConfig) -> dict | None: + """Default speculative config from the checkpoint's own MTP head, or None. + + Only applies when the user gave no speculative config at all. The drafter + weights live in the target checkpoint, so nothing extra is downloaded. + """ + hf_config = getattr(model_config, "hf_config", None) + model_type = getattr(hf_config, "model_type", None) + if model_type not in _CHECKPOINT_MTP_MODEL_TYPES: + return None + text_config = getattr(model_config, "hf_text_config", None) + n_mtp_layers = getattr(text_config, "mtp_num_hidden_layers", None) + if not isinstance(n_mtp_layers, int) or n_mtp_layers < 1: + return None + return { + "method": "mtp", + "num_speculative_tokens": _CHECKPOINT_MTP_NUM_SPECULATIVE_TOKENS, + } diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef174135..661cc4b43 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -395,46 +395,73 @@ 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 choice chunk per generated token. A RequestOutput can + # carry several new tokens (speculative decoding accepts a + # step's drafts at once; the frontend merges outputs when it + # lags the engine). Emitting them as one chunk hides the + # token granularity from streaming consumers, so split the + # delta into per-token slices. Text concatenation, token + # accounting and finish_reason placement are unchanged. + delta_ids = as_list(output.token_ids) if ( - not include_usage - and self.system_fingerprint is not None - and finish_reason is not None + len(delta_ids) > 1 + and logprobs is None + and not (request.echo and prompt_token_ids_to_return) ): - 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, + pieces = _split_delta_by_token( + tokenizer, delta_ids, delta_text ) + else: + pieces = [(delta_text, delta_ids)] + + tokens_before = previous_num_tokens[i] - len(delta_ids) + n_pieces = len(pieces) + for piece_idx, (piece_text, piece_ids) in enumerate(pieces): + is_last = piece_idx == n_pieces - 1 + tokens_before += 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=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_idx == 0 + else None + ), + token_ids=( + piece_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 ( + 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 = tokens_before + 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) @@ -689,3 +716,39 @@ class OpenAIServingCompletion(OpenAIServing): tokens=out_tokens, top_logprobs=out_top_logprobs, ) + + +def _split_delta_by_token( + tokenizer: TokenizerLike | None, + token_ids: list[int], + text: str, +) -> list[tuple[str, list[int]]]: + """Slice a multi-token delta into per-token (text, [token_id]) pieces. + + Boundaries come from incrementally decoding the delta's token prefix; they + are clamped and made monotone so the pieces always concatenate to exactly + ``text``. The final piece absorbs any remainder, so nothing is lost or + duplicated even when a token boundary does not fall on a character + boundary (multi-byte sequences, leading-space normalisation). + """ + n = len(token_ids) + if n <= 1 or not text: + return [(text, list(token_ids))] + bounds: list[int] = [] + if tokenizer is not None: + try: + for j in range(1, n): + prefix = tokenizer.decode(token_ids[:j], skip_special_tokens=True) + bounds.append(len(prefix)) + except Exception: # noqa: BLE001 - fall back to an even split + bounds = [] + if len(bounds) != n - 1: + bounds = [(len(text) * j) // n for j in range(1, n)] + pieces: list[tuple[str, list[int]]] = [] + start = 0 + for j, b in enumerate(bounds): + b = min(max(b, start), len(text)) + pieces.append((text[start:b], [token_ids[j]])) + start = b + pieces.append((text[start:], [token_ids[n - 1]])) + return pieces diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4ac8d49cd..bca89ea4d 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -321,8 +321,10 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] ] = { PlatformEnum.CUDA: [ - FlashInferFp8DeepGEMMDynamicBlockScaledKernel, + # One kernel for every batch size: no M-dependent dispatch, and no + # TensorRT-LLM nvcc JIT on the cold-start path. DeepGemmFp8BlockScaledMMKernel, + FlashInferFp8DeepGEMMDynamicBlockScaledKernel, CutlassFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, 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..d7da9012a 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 @@ -1329,9 +1329,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): if spec_sequence_masks is not None: if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: mixed_qkv_spec = mixed_qkv + a_spec = a + b_spec = b mixed_qkv_non_spec = None else: mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + a_spec = a.index_select(0, spec_token_indx) + b_spec = b.index_select(0, spec_token_indx) mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx) else: mixed_qkv_spec = None @@ -1456,8 +1460,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, diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index f7c237ca2..14325d8b1 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -745,6 +745,116 @@ def causal_conv1d_fn( return out.to(original_x_dtype) +# Constexpr signatures of `_causal_conv1d_fwd_kernel` that have already been +# compiled by `warmup_causal_conv1d_fwd_kernel`. +_FWD_WARMED_UP_KEYS: set[tuple] = set() + +# Above this size a same-shape copy of the conv state cache is too expensive +# to allocate just for a warmup, so a small backing buffer is used instead. +_MAX_DUMMY_CONV_STATE_BYTES = 24 * 1024 * 1024 * 1024 + + +def _dummy_conv_states(conv_states: torch.Tensor, rows: int) -> torch.Tensor: + """Throwaway stand-in for a conv state cache. + + `_causal_conv1d_fwd_kernel` takes `num_cache_lines` (i.e. + `conv_states.size(0)`) and every conv state stride as `tl.constexpr`, + so a warmup only populates the right JIT cache entry if the dummy + reports the same size and strides as the real cache. The kernel only + ever touches the rows named by `cache_indices`, so the returned tensor + advertises the full shape while owning storage for `rows` rows only + (as_strided) whenever a full-size copy would be wasteful. + """ + shape = tuple(conv_states.shape) + strides = tuple(conv_states.stride()) + itemsize = conv_states.element_size() + if shape[0] * strides[0] * itemsize <= _MAX_DUMMY_CONV_STATE_BYTES: + dummy = torch.empty_strided( + shape, strides, dtype=conv_states.dtype, device=conv_states.device + ) + return dummy.zero_() + + raise MemoryError(f"conv state cache is {shape[0] * strides[0] * itemsize} bytes") + + +def warmup_causal_conv1d_fwd_kernel( + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + activation: str | None, + cache_indices_stride: int = 1, + num_tokens: int = 64, +) -> None: + """Pre-compile `_causal_conv1d_fwd_kernel` for the prefill path. + + The mamba/GDN prefill convolution is skipped during V1 profiling + (`attn_metadata` is None there) and is not part of any captured CUDA + graph, so it is JIT-compiled inside the first real request and shows + up as a TTFT spike. + + Almost every argument of the kernel is a `tl.constexpr`: `dim`, + `num_cache_lines`, all strides, `KERNEL_WIDTH`, `HAS_BIAS`, + `SILU_ACTIVATION`, `IS_APC_ENABLED`, `HAS_NULL_BLOCK`, `NP2_STATELEN` + and the block sizes. They are reproduced exactly here by going + through `causal_conv1d_fn` with the layer's real weight/bias and a + dummy conv state cache that mirrors the real shape, strides and dtype. + The two that depend on the batch are `stride_x_token` (`x` is always + channel-last, so it is `dim` as in the real call) and + `stride_cache_indices`, which the caller passes in: the GDN metadata + builder produces a strided `block_table[:, 0]` view for non-spec + batches and a compact 1D tensor for mixed spec batches, so both need + a warmup call to cover every real launch. + + Nothing here touches live state: `x`, the output, the query/cache + metadata and the conv state cache are all freshly allocated, and the + result is discarded. + """ + dim, width = weight.shape + key = ( + dim, + width, + conv_states.shape, + conv_states.stride(), + conv_states.dtype, + weight.stride(), + bias is not None, + activation, + cache_indices_stride, + num_tokens, + ) + if key in _FWD_WARMED_UP_KEYS: + return + _FWD_WARMED_UP_KEYS.add(key) + + device = conv_states.device + dummy_conv_states = _dummy_conv_states(conv_states, rows=2) + # Channel-last `x`, exactly like the real `mixed_qkv.transpose(0, 1)`. + x = torch.zeros( + (num_tokens, dim), dtype=conv_states.dtype, device=device + ).transpose(0, 1) + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + # One sequence, pointing at a scratch cache line. Row 0 is avoided + # because NULL_BLOCK_ID == 0 makes the kernel skip the sequence. + cache_line = 1 if dummy_conv_states.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, + activation=activation, + conv_states=dummy_conv_states, + has_initial_state=has_initial_state, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + metadata=None, + ) + + @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 754270e65..773115bb3 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -109,6 +109,15 @@ def kernel_warmup(worker: "Worker"): create_mixed_batch=True, ) + # Pre-compile the input-prep / state-management Triton kernels that + # neither the profile run nor CUDA graph capture reaches, so they do + # not JIT inside the first request. No-op without speculative decoding. + # Imported here rather than at module scope to avoid an import cycle + # through vllm.v1.worker. + from vllm.model_executor.warmup.triton_jit_warmup import warmup_jit_kernels + + warmup_jit_kernels(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/triton_jit_warmup.py b/vllm/model_executor/warmup/triton_jit_warmup.py new file mode 100644 index 000000000..2b240ee48 --- /dev/null +++ b/vllm/model_executor/warmup/triton_jit_warmup.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Pre-compile the Triton kernels that would otherwise JIT inside the first +inference. + +With speculative decoding on a hybrid (mamba) model, a handful of Triton +kernels live on the input-preparation / state-management side of a step. +Neither the V1 profiling run (which has no attention metadata, so the +drafter and the mamba prefill path are never entered) nor CUDA graph +capture reaches them, so Triton compiles them inside the *first real +request* -- ``jit_monitor`` reports each one as +"Triton kernel JIT compilation during inference ... causes a latency +spike". The spike lands entirely in the first benchmark repetition and +blows up the cross-repetition TTFT range. + +This module mirrors ``QwenGatedDeltaNet._warmup_prefill_kernels`` +(vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py:1068): run +each kernel once, on small freshly allocated dummy tensors, while startup +still has headroom, and guard it so it happens once. Every group is +wrapped in its own ``try/except`` that logs and continues -- a missed +warmup only costs the latency it was meant to save, whereas an exception +escaping here would take the engine down before it can serve. + +Nothing in here can change a numerical result: every kernel runs on +tensors allocated inside the warmup and every output is discarded. No +KV cache, mamba state, block table, positions buffer or slot mapping +buffer belonging to the engine is ever passed in. +""" + +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 _warmup_eagle_input_kernels(worker: "Worker") -> None: + """``eagle_*_kernel`` x3 (vllm/v1/spec_decode/utils.py:29,137,179).""" + runner = worker.model_runner + drafter = getattr(runner, "drafter", None) + if drafter is None or not hasattr(drafter, "warmup_input_kernels"): + return + + # `n_blocks_per_req` is the second dimension of the block table the + # drafter is handed, which for a hybrid model is *not* + # cdiv(max_model_len, block_size); read it off the live table rather + # than recomputing it. Only `.shape[1]` is used -- the tensor itself is + # not passed to any kernel. + gid = max(getattr(drafter, "kv_cache_gid", 0), 0) + block_table = runner.input_batch.block_table[gid] + drafter.warmup_input_kernels( + n_blocks_per_req=block_table.get_device_tensor(1).shape[1], + vocab_size=runner.input_batch.vocab_size, + ) + + +def _warmup_rejection_greedy_kernel(worker: "Worker") -> None: + """``rejection_greedy_sample_kernel`` + (vllm/v1/sample/rejection_sampler.py:806). + + ``GPUModelRunner._dummy_sampler_run`` already exercises this kernel at + startup, but with ``all_greedy=False``, so ``rejection_sample`` hands it + a real ``is_greedy`` tensor. An all-greedy (temperature 0) batch passes + ``None`` instead, and Triton treats a ``None`` argument as its own + constexpr -- a different cache entry. That is exactly the case the + ``FIXME`` at rejection_sampler.py:820 warns about, and why this kernel + shows up in the JIT monitor. + """ + runner = worker.model_runner + sampler = getattr(runner, "rejection_sampler", None) + if sampler is None or not hasattr(sampler, "warmup_greedy_kernel"): + return + sampler.warmup_greedy_kernel( + runner.device, + worker.vllm_config.speculative_config.num_speculative_tokens, + ) + + +def _warmup_mamba_runner_kernels(worker: "Worker") -> None: + """``postprocess_mamba_fused_kernel`` and ``batch_memcpy_kernel`` + (vllm/v1/worker/mamba_utils.py:29,198).""" + from vllm.v1.kv_cache_interface import MambaSpec + from vllm.v1.worker.mamba_utils import get_mamba_groups, warmup_mamba_kernels + + runner = worker.model_runner + kv_cache_config = getattr(runner, "kv_cache_config", None) + if kv_cache_config is None or not any( + isinstance(g.kv_cache_spec, MambaSpec) + for g in kv_cache_config.kv_cache_groups + ): + return + if runner.cache_config.mamba_cache_mode != "align": + # Both kernels are only reached in align mode. + return + + group_ids, mamba_spec = get_mamba_groups(kv_cache_config) + # `block_table_stride_req` is a specializing scalar and is constant for + # the engine's lifetime; pass the real one so the cache entry is reused. + stride = int(runner.input_batch.block_table[group_ids[0]].block_table.gpu.stride(0)) + warmup_mamba_kernels( + runner.device, + block_size=mamba_spec.block_size, + block_table_stride_req=stride, + num_reqs_variants=(1, 2, min(32, runner.max_num_reqs)), + ) + + +def _warmup_zero_kv_blocks(worker: "Worker") -> None: + """``_zero_kv_blocks_kernel`` (vllm/v1/worker/utils.py:41).""" + zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) + if zeroer is None: + return + zeroer.warmup() + + +def _warmup_compute_slot_mapping(worker: "Worker") -> None: + """``_compute_slot_mapping_kernel`` (vllm/v1/worker/block_table.py:388).""" + input_batch = getattr(worker.model_runner, "input_batch", None) + block_table = getattr(input_batch, "block_table", None) + if block_table is None: + return + block_table.warmup() + + +def _mamba_cache_indices_strides(worker: "Worker") -> list[int]: + """Strides of the ``cache_indices`` tensors the conv kernel will see. + + ``stride_cache_indices`` is a ``tl.constexpr`` + (causal_conv1d.py:16), and the GDN metadata builder produces two + flavours: a ``block_table_tensor[:, 0]`` view, whose stride(0) is the + mamba group's ``max_num_blocks_per_req`` (gdn_attn.py:219), and a + boolean-indexed copy, which is a fresh contiguous 1D tensor of stride 1 + (gdn_attn.py:291). Both are warmed. + """ + from vllm.v1.kv_cache_interface import MambaSpec + + strides = {1} + runner = worker.model_runner + kv_cache_config = getattr(runner, "kv_cache_config", None) + input_batch = getattr(runner, "input_batch", None) + if kv_cache_config is None or input_batch is None: + return sorted(strides) + for gid, group in enumerate(kv_cache_config.kv_cache_groups): + if not isinstance(group.kv_cache_spec, MambaSpec): + continue + strides.add(int(input_batch.block_table[gid].block_table.gpu.stride(0))) + return sorted(strides) + + +def _warmup_causal_conv1d_fwd(worker: "Worker") -> None: + """``_causal_conv1d_fwd_kernel`` + (vllm/model_executor/layers/mamba/ops/causal_conv1d.py:16).""" + 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 ( + warmup_causal_conv1d_fwd_kernel, + ) + + strides = _mamba_cache_indices_strides(worker) + 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 + try: + # Mirror the layout selection the mamba/GDN forward does; the + # real cache is read for shape/stride/dtype only, never launched + # against (`warmup_causal_conv1d_fwd_kernel` allocates its own). + conv_states = ( + kv_cache[0] + if is_conv_state_dim_first() + else kv_cache[0].transpose(-1, -2) + ) + if conv_states.dim() != 3 or conv_states.numel() == 0: + continue + weight = conv1d.weight.view(conv1d.weight.size(0), conv1d.weight.size(2)) + for stride in strides: + warmup_causal_conv1d_fwd_kernel( + conv_states, + weight, + conv1d.bias, + getattr(module, "activation", "silu"), + cache_indices_stride=stride, + ) + except Exception: + logger.warning( + "causal_conv1d prefill kernel warmup failed for layer %s", + getattr(module, "prefix", type(module).__name__), + exc_info=True, + ) + + +_WARMUP_GROUPS = ( + ("eagle input-prep kernels", _warmup_eagle_input_kernels), + ("rejection_greedy_sample_kernel", _warmup_rejection_greedy_kernel), + ("mamba runner kernels", _warmup_mamba_runner_kernels), + ("_zero_kv_blocks_kernel", _warmup_zero_kv_blocks), + ("_compute_slot_mapping_kernel", _warmup_compute_slot_mapping), + ("_causal_conv1d_fwd_kernel", _warmup_causal_conv1d_fwd), +) + + +def warmup_jit_kernels(worker: "Worker") -> None: + """Warm every Triton kernel the first inference would otherwise compile. + + No-op unless speculative decoding is enabled: all nine kernels this + covers are either drafter-side, rejection-sampler-side, or only reached + by the hybrid spec-decode step, and gating keeps the cost off every + other deployment. + + Must run after the KV caches are allocated and bound (the drafter's + ``block_size``, the block tables and the conv cache layout all come from + ``initialize_kv_cache``) and before the server reports healthy; + ``kernel_warmup`` -- called at gpu_worker.py:654, after + ``initialize_from_config`` at gpu_worker.py:591 and before + ``capture_model`` at gpu_worker.py:658 -- satisfies both, and warming + ahead of graph capture also keeps JIT out of the capture itself. + """ + if worker.vllm_config.speculative_config is None: + return + + for name, fn in _WARMUP_GROUPS: + try: + with torch.inference_mode(): + fn(worker) + except Exception: + logger.warning( + "Triton JIT warmup for %s failed; it will be compiled during " + "the first request instead.", + name, + exc_info=True, + ) + else: + logger.debug("Triton JIT warmup for %s completed", name) diff --git a/vllm/outputs.py b/vllm/outputs.py index 2c71d2afb..8f44edfab 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 + # Per-token slices of ``text`` for this delta, or None when unavailable. + token_texts: 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.token_texts is not None + and next_completion.token_texts is not None + ): + completion.token_texts = list( + completion.token_texts + ) + list(next_completion.token_texts) + else: + completion.token_texts = 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 340a30403..48a30a0d3 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -177,7 +177,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, @@ -387,6 +386,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] diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index 4700eecb5..da89b4d83 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.token_text_ends: list[int] = [] @property def output_token_ids(self) -> list[int]: @@ -45,6 +46,11 @@ class IncrementalDetokenizer: def get_next_output_text(self, finished: bool, delta: bool) -> str: return "" + def get_next_output_pieces( + self, finished: bool, delta: bool, num_tokens: int + ) -> tuple[str, list[str]]: + return "", [""] * num_tokens + @classmethod def from_new_request( cls, @@ -91,6 +97,10 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): # Generation data self.output_text = "" + # End offset in ``output_text`` of each token decoded by the most + # recent ``update``. One entry per token id passed in, including a + # zero-width entry for a stop token that is not detokenized. + self.token_text_ends: list[int] = [] def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None: """ @@ -100,6 +110,7 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): Return matched stop string or None. """ + self.token_text_ends.clear() if not new_token_ids: # Skip detokenization if no new token ids. return None @@ -117,13 +128,17 @@ 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.token_text_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) if skipped_stop_token_id is not None: - # Cleanup after skipping detokenization. + # Cleanup after skipping detokenization. The token is still + # counted in usage, so it needs a zero-width offset entry to keep + # one entry per token id. self.token_ids.append(skipped_stop_token_id) + self.token_text_ends.append(len(self.output_text)) # 2) Evaluate stop strings. stop_string = None @@ -138,6 +153,9 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): stop_string, truncate_to = stop if truncate_to != -1: self.output_text = self.output_text[:truncate_to] + self.token_text_ends = [ + min(end, truncate_to) for end in self.token_text_ends + ] return stop_string @@ -163,6 +181,34 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): return self.output_text[last_offset:length] return "" + def get_next_output_pieces( + self, finished: bool, delta: bool, num_tokens: int + ) -> tuple[str, list[str]]: + """Delta text plus its per-token slices. + + ``get_next_output_text`` advances ``_last_output_text_offset``, so it + is called exactly once here and the base offset is captured before it. + The pieces always concatenate to exactly the returned text and there is + always one piece per token id, so a caller can emit one streaming chunk + per token without inventing or dropping any. + """ + base = self._last_output_text_offset + text = self.get_next_output_text(finished, delta) + if num_tokens <= 0: + return text, [] + if not delta or len(self.token_text_ends) != num_tokens: + return text, [""] * (num_tokens - 1) + [text] + limit = base + len(text) + pieces: list[str] = [] + start = base + for end in self.token_text_ends: + end = min(max(end, start), limit) + pieces.append(self.output_text[start:end]) + start = end + if start < limit: + pieces[-1] += self.output_text[start:limit] + 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 e1032cfd1..9969e6385 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +from collections import deque from collections import defaultdict, deque from collections.abc import Iterable from dataclasses import dataclass @@ -55,13 +56,27 @@ class RequestOutputCollector: self.aggregate = output_kind == RequestOutputKind.DELTA self.request_id = request_id self.output: RequestOutput | PoolingRequestOutput | Exception | None = None + # Streamed deltas are queued one per engine step instead of being + # merged into whatever the consumer has not collected yet. Merging + # turned two steps' tokens into one streamed chunk whenever the API + # server's event loop fell behind the engine -- routine under load, + # and worse on a small CPU -- which hides token granularity from any + # consumer that counts chunks as tokens. Each queued output keeps its + # own step's text and logprobs, so nothing is re-derived downstream. + self._pending: deque[RequestOutput] = deque() self.ready = asyncio.Event() self._input_stream_task: asyncio.Task | None = None def put(self, output: RequestOutput | PoolingRequestOutput | Exception) -> None: """Non-blocking put operation.""" - if self.output is None or isinstance(output, Exception): + if isinstance(output, Exception): + self.output = output + self.ready.set() + elif self.aggregate and isinstance(output, RequestOutput): + self._pending.append(output) + self.ready.set() + elif self.output is None: self.output = output self.ready.set() elif isinstance(self.output, RequestOutput) and isinstance( @@ -75,22 +90,30 @@ class RequestOutputCollector: ): self.output = output + def _take(self) -> RequestOutput | PoolingRequestOutput | Exception | None: + # An exception is delivered ahead of any queued deltas: it ends the + # stream and the consumer must see it now. + if isinstance(self.output, Exception): + output, self.output = self.output, None + elif self._pending: + output = self._pending.popleft() + else: + output, self.output = self.output, None + if self.output is None and not self._pending: + self.ready.clear() + return output + async def get(self) -> RequestOutput | PoolingRequestOutput: """Get operation blocks on put event.""" - while (output := self.output) is None: + while (output := self._take()) is None: await self.ready.wait() - self.output = None - self.ready.clear() if isinstance(output, Exception): raise output return output def get_nowait(self) -> RequestOutput | PoolingRequestOutput | None: """Non-blocking get operation.""" - output = self.output - if output is not None: - self.output = None - self.ready.clear() + output = self._take() if isinstance(output, Exception): raise output return output @@ -385,8 +408,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: + token_texts: list[str] | None = None + if delta: + text, token_texts = self.detokenizer.get_next_output_pieces( + finished, True, len(token_ids) + ) + else: + text = self.detokenizer.get_next_output_text(finished, False) token_ids = self.detokenizer.output_token_ids # Prepare logprobs, based on delta mode @@ -408,6 +436,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, + token_texts=token_texts, ) 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 8b4d8c9dc..9fe42bf96 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -84,6 +84,98 @@ class RejectionSampler(nn.Module): device=device, ) self.synthetic_mode = self.synthetic_conditional_rates is not None + self._greedy_kernel_warmed_up = False + + def warmup_greedy_kernel( + self, + device: torch.device, + num_speculative_tokens: int, + ) -> None: + """Pre-compile ``rejection_greedy_sample_kernel`` for the all-greedy case. + + ``GPUModelRunner._dummy_sampler_run`` already runs the rejection sampler + once during startup, but it does so with ``all_greedy=False``, so + ``rejection_sample`` hands the kernel a real ``is_greedy`` tensor. When + every request in a batch is greedy (temperature 0) the wrapper passes + ``None`` for ``is_greedy_ptr`` instead. Triton treats a ``None`` + argument as its own constexpr, so that is a *different* specialization: + the kernel is JIT-compiled during the first real inference and shows up + as a latency spike. This is precisely the case the ``FIXME`` above + ``rejection_greedy_sample_kernel`` warns about. + + Calling the ``rejection_sample`` wrapper once with an all-greedy + ``SamplingMetadata`` populates the ``is_greedy_ptr=None`` entry of the + JIT cache while the engine is still starting up. The wrapper is used + rather than the raw kernel so that every argument (in particular the + int64 ``target_argmax`` produced by ``target_logits.argmax``) has + exactly the dtype the real call site produces. + + Numerically inert: every tensor below is freshly allocated here and the + returned token ids are discarded. The kernel only writes into the + throwaway ``output_token_ids`` buffer it allocates internally. + + Note that the kernel's only ``tl.constexpr`` is ``SYNTHETIC_MODE``, + which is taken from ``self.synthetic_mode`` and therefore already + matches the real call, and ``max_spec_len`` is declared in + ``do_not_specialize``, so a single call covers every batch size and + every draft length. + """ + if self._greedy_kernel_warmed_up: + return + self._greedy_kernel_warmed_up = True + + from vllm.v1.sample.logits_processor import LogitsProcessors + + # One request with a full draft is enough: the kernel launches one + # program per request and does not specialize on batch size or on + # `max_spec_len`. + num_draft_tokens = [max(num_speculative_tokens, 1)] + num_tokens = sum(num_draft_tokens) + # Only `argmax` is taken over the vocab dimension, so keep it tiny. + vocab_size = 8 + + draft_token_ids = torch.zeros(num_tokens, dtype=torch.int32, device=device) + cu_num_draft_tokens = torch.tensor( + [num_tokens], dtype=torch.int32, device=device + ) + target_logits = torch.zeros( + (num_tokens, vocab_size), dtype=torch.float32, device=device + ) + bonus_token_ids = torch.zeros((1, 1), dtype=torch.int32, device=device) + empty = torch.zeros(1, dtype=torch.float32, device=device) + + sampling_metadata = 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_token_ids, + num_draft_tokens, + num_draft_tokens[0], + cu_num_draft_tokens, + None, + target_logits, + bonus_token_ids, + sampling_metadata, + 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 9f46cbd24..749feac90 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -7,6 +7,7 @@ import numpy as np import torch import torch.nn as nn +from vllm import _custom_ops as ops from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.config import ( CUDAGraphMode, @@ -17,6 +18,7 @@ from vllm.config import ( from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger +from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.model_loader import get_model from vllm.model_executor.models import supports_multimodal @@ -117,6 +119,10 @@ class SpecDecodeBaseProposer: self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None if self.parallel_drafting: self._init_parallel_drafting_params() + # Private fp8 draft head; None means "use the bf16 path" (see + # _maybe_build_fp8_draft_head). + self._draft_head8: torch.Tensor | None = None + self._draft_head8_scale: torch.Tensor | None = None self.use_local_argmax_reduction: bool = ( self.speculative_config.use_local_argmax_reduction ) @@ -408,8 +414,97 @@ class SpecDecodeBaseProposer: self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode) + def _maybe_build_fp8_draft_head(self) -> None: + """Build a private fp8-e4m3 copy of the draft LM head, or leave it unset. + + Only the drafter's argmax reads this weight. Greedy verification emits the + target's argmax at every accepted position, so draft-side quantisation + cannot change the generated text - it can only change how often a draft is + accepted. The shared module is deliberately left untouched: load_model + aliases the target's lm_head into the drafter, so quantising in place would + move the target's own greedy path. + + Any failure leaves the fields None and the bf16 path in use. + """ + self._draft_head8 = None + self._draft_head8_scale = None + try: + if self.method != "mtp" or get_tensor_model_parallel_world_size() != 1: + return + head = getattr(self.model, "lm_head", None) + weight = getattr(head, "weight", None) + if not isinstance(weight, torch.Tensor) or weight.dim() != 2: + return + if not weight.is_cuda or weight.dtype not in ( + torch.bfloat16, + torch.float16, + ): + return + # The logits slice must be the whole tensor; a padded vocab would make + # argmax able to return a padding id. + org_vocab_size = getattr(head, "org_vocab_size", None) + if org_vocab_size is not None and org_vocab_size != weight.shape[0]: + return + + fp8_max = torch.finfo(torch.float8_e4m3fn).max + num_rows, hidden_size = weight.shape + quantized = torch.empty( + (num_rows, hidden_size), + dtype=torch.float8_e4m3fn, + device=weight.device, + ) + scales = torch.empty( + (num_rows, 1), dtype=torch.float32, device=weight.device + ) + # Chunked so the fp32 staging copy stays a few hundred MB. + for start in range(0, num_rows, 8192): + block = weight[start : start + 8192].float() + scale = block.abs().amax(dim=1, keepdim=True).clamp_(min=1e-12) + scale = scale / fp8_max + quantized[start : start + 8192] = ( + (block / scale).clamp_(-fp8_max, fp8_max).to(torch.float8_e4m3fn) + ) + scales[start : start + 8192] = scale + del block + # cutlass wants B column-major: .t() of a row-major [N, K] gives + # stride(0) == 1. + self._draft_head8 = quantized.t() + self._draft_head8_scale = scales.reshape(1, num_rows) + + # Resolve GEMM heuristics for every batch width the drafter can see, + # here in load_model, so no first-token work is deferred into serving. + probe = torch.randn( + (32, hidden_size), dtype=weight.dtype, device=weight.device + ) + for width in (1, 2, 4, 6, 8, 12, 16, 24, 32): + self._greedy_sample(probe[:width]) + torch.cuda.synchronize() + logger.info( + "MTP: built private fp8 draft lm_head (%d x %d).", + num_rows, + hidden_size, + ) + except Exception: # noqa: BLE001 - degrade to bf16, never fail startup + self._draft_head8 = None + self._draft_head8_scale = None + logger.warning( + "fp8 draft lm_head unavailable; using the bf16 head.", exc_info=True + ) + def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor: """Greedy-sample draft tokens from hidden states.""" + if self._draft_head8 is not None: + flat = hidden_states.reshape(-1, hidden_states.shape[-1]).contiguous() + quantized, scale = ops.scaled_fp8_quant( + flat, scale=None, use_per_token_if_dynamic=True + ) + return ops.cutlass_scaled_mm( + quantized, + self._draft_head8, + scale, + self._draft_head8_scale, + out_dtype=hidden_states.dtype, + ).argmax(dim=-1) if self.use_local_argmax_reduction: return self.model.get_top_tokens(hidden_states) return self.model.compute_logits(hidden_states).argmax(dim=-1) @@ -1306,6 +1401,7 @@ class SpecDecodeBaseProposer: self._maybe_share_embeddings(target_language_model) self._maybe_share_lm_head(target_language_model) + self._maybe_build_fp8_draft_head() if ( self.parallel_drafting diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index d9c041ba0..c6f7b7829 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -98,6 +98,7 @@ class BlockTable: self.dcp_world_size = 1 self.dcp_rank = 0 self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size + self._slot_mapping_warmed_up = False def append_row( self, @@ -163,6 +164,63 @@ class BlockTable: BLOCK_SIZE=1024, ) + def warmup(self) -> None: + """Pre-compile ``_compute_slot_mapping_kernel`` on dummy tensors. + + ``compute_slot_mapping`` first runs inside ``_prepare_inputs`` of + the first real step, so its Triton JIT compilation lands in the + first request and inflates TTFT. Every argument that takes part + in the Triton compile key is fixed for the lifetime of this + object: the ``tl.constexpr`` meta-parameters + (``TOTAL_CP_WORLD_SIZE``, ``TOTAL_CP_RANK``, + ``CP_KV_CACHE_INTERLEAVE_SIZE``, ``PAD_ID``, ``BLOCK_SIZE``) plus + the pointer/int specializations of ``block_table_stride`` and + ``block_size``. ``num_tokens`` / ``max_num_tokens`` are declared + ``do_not_specialize``, so a single tiny launch covers every batch + the engine will ever build, and the grid size is not part of the + key either. + + Freshly allocated throwaway tensors are used for every argument + (including the output ``slot_mapping``), so no live engine buffer + is read or written. + """ + if self._slot_mapping_warmed_up: + return + self._slot_mapping_warmed_up = True + + device = self.device + num_reqs = 1 + num_tokens = 8 + block_table_stride = self.block_table.gpu.stride(0) + # Dummy inputs: one request covering `num_tokens` positions at + # position 0, so only block-table entry 0 is ever read. + query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) + query_start_loc[num_reqs] = num_tokens + positions = torch.zeros(num_tokens, dtype=torch.int64, device=device) + block_table = torch.zeros( + (num_reqs, block_table_stride), dtype=torch.int32, device=device + ) + slot_mapping = torch.zeros(num_tokens, dtype=torch.int64, device=device) + total_cp_world_size = self.pcp_world_size * self.dcp_world_size + total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank + _compute_slot_mapping_kernel[(num_reqs + 1,)]( + num_tokens, + # max_num_tokens == num_tokens: the padding loop in the last + # program is empty, so the dummy output stays in bounds. + num_tokens, + query_start_loc, + positions, + block_table, + block_table_stride, + self.block_size, + slot_mapping, + TOTAL_CP_WORLD_SIZE=total_cp_world_size, + TOTAL_CP_RANK=total_cp_rank, + CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size, + PAD_ID=PAD_SLOT_ID, + BLOCK_SIZE=1024, + ) + def commit_block_table(self, num_reqs: int) -> None: self.block_table.copy_to_gpu(num_reqs) @@ -309,6 +367,10 @@ class MultiGroupBlockTable: for block_table in self.block_tables: block_table.compute_slot_mapping(num_reqs, query_start_loc, positions) + def warmup(self) -> None: + for block_table in self.block_tables: + block_table.warmup() + def commit_block_table(self, num_reqs: int) -> None: for block_table in self.block_tables: block_table.commit_block_table(num_reqs) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a823..2b0c0fdf3 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2159,6 +2159,9 @@ class GPUModelRunner( target.gpu[:, :total_num_scheduled_tokens] += drift use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + # See the GDN deferred-rewind notes below; consumed when building + # attention metadata, which happens in a different method. + self._gdn_spec_rewind_marked = 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 +2171,28 @@ class GPUModelRunner( logits_indices = query_start_loc[1:] - 1 spec_decode_metadata = None num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + # GDN keeps the state of a request that accepted a > 1 tokens in + # block column a-1 until a spec-path read performs the deferred + # rewind. With no drafts anywhere this step, the GDN builder + # would route every row down the non-spec path, which reads + # column 0 - a state stale by a-1 tokens. Mark decode rows as + # spec rows so the spec kernels read each row's true column from + # the GPU-side accepted counts (bit-identical when a == 1). + if ( + self.speculative_config is not None + and self.model_config.is_hybrid + ): + is_decode = ( + self.input_batch.num_computed_tokens_cpu[:num_reqs] + >= self.input_batch.num_prompt_tokens[:num_reqs] + ) + if is_decode.any(): + self.num_decode_draft_tokens.np[:num_reqs] = np.where( + is_decode, 1, -1 + ) + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + self._gdn_spec_rewind_marked = True else: # Get the number of draft tokens for each request. # Iterate over the dictionary rather than all requests since not all @@ -2188,6 +2213,22 @@ class GPUModelRunner( >= self.input_batch.num_prompt_tokens[req_idx] ): num_decode_draft_tokens[req_idx] = draft_len + # Draft-starved decode rows must also take the GDN spec path + # (see the draft-less branch above); rows with drafts keep their + # real counts. + if ( + self.speculative_config is not None + and self.model_config.is_hybrid + ): + is_decode = ( + self.input_batch.num_computed_tokens_cpu[:num_reqs] + >= self.input_batch.num_prompt_tokens[:num_reqs] + ) + num_decode_draft_tokens = np.where( + is_decode & (num_decode_draft_tokens < 0), + 1, + num_decode_draft_tokens, + ) spec_decode_metadata = self._calc_spec_decode_metadata( num_draft_tokens, cu_num_tokens ) @@ -2404,7 +2445,9 @@ class GPUModelRunner( ) extra_attn_metadata_args = {} - if use_spec_decode and isinstance( + if ( + use_spec_decode or getattr(self, "_gdn_spec_rewind_marked", False) + ) and isinstance( builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder) ): assert ubid is None, "UBatching not supported with GDN yet" diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 45166ef9a..d4f756e08 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -8,6 +8,7 @@ from typing import Any import torch from vllm.config import CacheConfig +from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, get_conv_copy_spec, @@ -22,6 +23,8 @@ from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch +logger = init_logger(__name__) + @triton.jit def postprocess_mamba_fused_kernel( @@ -566,6 +569,190 @@ class MambaSpecDecodeGPUContext: ) +_MAMBA_KERNELS_WARMED_UP = False + + +def _warmup_batch_memcpy(device: torch.device) -> None: + """Compile ``batch_memcpy_kernel`` with a single copy inside a scratch + buffer. + + ``BLOCK_SIZE`` is the kernel's only ``tl.constexpr`` and ``batch_memcpy`` + always passes 1024, so one call through that wrapper covers every real + invocation. Source and destination are two disjoint halves of one + freshly allocated byte buffer, so nothing else in the engine is touched. + """ + num_bytes = 1024 + scratch = torch.zeros(2 * num_bytes, dtype=torch.uint8, device=device) + # Same dtypes/layout as MambaCopyBuffers so the Triton signature and + # pointer-alignment specialization match the real launch. + src_ptrs = CpuGpuBuffer(1, dtype=torch.uint64, device=device) + dst_ptrs = CpuGpuBuffer(1, dtype=torch.uint64, device=device) + sizes = CpuGpuBuffer(1, dtype=torch.int32, device=device) + src_ptrs.np[0] = scratch.data_ptr() + dst_ptrs.np[0] = scratch.data_ptr() + num_bytes + sizes.np[0] = num_bytes + batch_memcpy( + src_ptrs.copy_to_gpu(1), + dst_ptrs.copy_to_gpu(1), + sizes.copy_to_gpu(1), + ) + + +def _warmup_postprocess_mamba_fused( + device: torch.device, + block_size: int, + block_table_stride_req: int, + num_reqs_variants: tuple[int, ...], +) -> None: + """Compile ``postprocess_mamba_fused_kernel`` on a throwaway context. + + Builds a ``MambaSpecDecodeGPUContext`` whose metadata describes one + freshly allocated dummy state tensor and one freshly allocated dummy + block table, then launches through ``run_fused_postprocess`` so the + constexprs (``block_size``, ``COPY_BLOCK_SIZE``, ``CONV_STATE_DIM_FIRST``) + and the argument signature are produced by the same code the decode loop + uses. + + The per-request inputs (1 accepted token, 1 scheduled token, nothing + computed, no drafts) make every program return before it dereferences a + state pointer, so the launch is a pure compile: no state is copied, and + every load it does perform is in bounds of the dummy tensors. + """ + max_reqs = max(num_reqs_variants) + # Dummy temporal state: 2 blocks of 16 elements. Only its address and + # strides reach the kernel, and only as metadata. + dummy_state = torch.zeros(2, 16, dtype=torch.float16, device=device) + # Zero-filled block table => block id 0 for every (req, block) pair. + dummy_block_table = torch.zeros( + max_reqs, block_table_stride_req, dtype=torch.int32, device=device + ) + + def meta(value: int, dtype: torch.dtype) -> torch.Tensor: + return torch.full((1,), value, dtype=dtype, device=device) + + ctx = MambaSpecDecodeGPUContext( + state_base_addrs=meta(dummy_state.data_ptr(), torch.int64), + state_block_strides=meta( + dummy_state.stride(0) * dummy_state.element_size(), torch.int64 + ), + state_elem_sizes=meta(dummy_state.element_size(), torch.int32), + state_inner_sizes=meta(dummy_state[0].numel(), torch.int64), + state_conv_widths=meta(0, torch.int32), # temporal state + state_group_indices=meta(0, torch.int32), + state_dim_row_count=meta(0, torch.int32), + state_dim_row_stride=meta(0, torch.int64), + block_size=block_size, + num_layers=1, + num_state_types=1, + mamba_group_ids=[0], + num_groups=1, + num_accepted_tokens_out=torch.zeros( + max_reqs, dtype=torch.int32, device=device + ), + block_table_ptrs=meta(dummy_block_table.data_ptr(), torch.int64), + block_table_stride_req=int(block_table_stride_req), + is_initialized=True, + ) + + ones = torch.ones(max_reqs, dtype=torch.int32, device=device) + zeros = torch.zeros(max_reqs, dtype=torch.int32, device=device) + for num_reqs in num_reqs_variants: + ctx.run_fused_postprocess( + num_reqs=num_reqs, + 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 warmup_mamba_kernels( + device: torch.device, + *, + block_size: int, + block_table_stride_req: int | None = None, + num_reqs_variants: tuple[int, ...] = (1, 2, 32), +) -> None: + """Pre-compile the runner-side mamba Triton kernels during startup. + + ``batch_memcpy_kernel`` (align-mode preprocess) and + ``postprocess_mamba_fused_kernel`` (align mode + spec decode on a hybrid + model) are both reached only from the scheduler-driven step, which + neither the profiling run nor CUDA graph capture exercises, so they JIT + compile on the first real request and show up as a TTFT spike. This mirrors + ``QwenGatedDeltaNet._warmup_prefill_kernels``: run each kernel once on + small dummy tensors while startup still has headroom, guarded by a + module-level flag so it happens once per process. Failures are logged + and swallowed - a missed warmup only costs latency, so it must never + take the engine down. + + No live engine state is involved: both kernels run on freshly allocated + throwaway tensors and their output is discarded, so no numerical result + can change. + + Triton caches one compilation per (constexpr tuple, argument signature, + scalar specialization). The fused kernel's constexprs are ``block_size`` + (model config), ``COPY_BLOCK_SIZE`` (hard-coded 1024) and + ``CONV_STATE_DIM_FIRST`` (``is_conv_state_dim_first()``); going through + ``run_fused_postprocess`` reproduces all three exactly. Its two plain + int arguments are additionally specialized into three buckets - value + ``1`` (folded into a constant), divisible by 16, and everything else - + so ``num_reqs`` is warmed once per bucket and ``block_table_stride_req`` + should be the engine's real value or the cache entry will not be reused. + + Args: + device: device the engine runs on. + block_size: mamba block size, i.e. + ``get_mamba_groups(kv_cache_config)[1].block_size``. + block_table_stride_req: ``stride(0)`` of a persistent mamba block + table, i.e. ``input_batch.block_table[group_id] + .get_device_tensor(1).stride(0)``. When ``None``, both + divisibility buckets are warmed instead, which doubles the + number of compilations. + num_reqs_variants: decode batch sizes to warm, one per + specialization bucket. + """ + global _MAMBA_KERNELS_WARMED_UP + if _MAMBA_KERNELS_WARMED_UP: + return + _MAMBA_KERNELS_WARMED_UP = True + + try: + _warmup_batch_memcpy(device) + except Exception: + logger.warning( + "batch_memcpy_kernel warmup failed; the first mamba state copy " + "will JIT compile it and spike latency.", + exc_info=True, + ) + + # Without the real stride, cover both divisibility buckets. + strides = ( + (16, 17) if block_table_stride_req is None else (block_table_stride_req,) + ) + for stride in strides: + try: + _warmup_postprocess_mamba_fused( + device, block_size, stride, num_reqs_variants + ) + except Exception: + logger.warning( + "postprocess_mamba_fused_kernel warmup failed for " + "block_size=%d, block_table_stride_req=%d; the first decode " + "step will JIT compile it and spike latency.", + block_size, + stride, + exc_info=True, + ) + + logger.debug( + "Mamba runner kernel warmup done (block_size=%d, num_reqs=%s).", + block_size, + num_reqs_variants, + ) + + @dataclasses.dataclass class MambaBuffers: """Single owner for all mamba-specific runner buffers. diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index c0f44b6db..685113e71 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -109,6 +109,7 @@ class KVBlockZeroer: """ self.device = device self.pin_memory = pin_memory + self._warmed_up = False self._meta: tuple[torch.Tensor, int, int, int] | None = None self._id_cap: int = 0 self._ids_pinned: torch.Tensor | None = None @@ -188,6 +189,50 @@ class KVBlockZeroer: len(seg_addrs), ) + def warmup(self) -> None: + """Pre-compile ``_zero_kv_blocks_kernel`` on throwaway memory. + + ``zero_block_ids`` is only called once the scheduler hands out + freshly allocated blocks, i.e. during the first real inference, + so the Triton JIT compilation of the kernel lands inside the + first request and shows up as a TTFT spike. + + This runs the kernel once per launch-argument specialization + while nothing is in flight. The compile key is + ``(N_SEGS, PAGE_SIZE_EL, BLOCK_SIZE)`` -- all three are fixed for + the lifetime of this object -- plus the specialization Triton + derives from the runtime ``n_blocks`` argument (it specializes + integers on ``== 1`` and ``% 16 == 0``), so three tiny launches + cover every shape the engine can produce. + + The segment address table is replaced by a table that points + every segment at one small throwaway buffer, so the real KV + cache is never touched: the kernel only writes zeros into that + scratch page. + """ + if self._meta is None or self._warmed_up: + return + self._warmed_up = True + + _, page_size_el, blk_size, n_segs = self._meta + # One scratch page, shared by every segment. With block id 0 the + # kernel writes exactly PAGE_SIZE_EL int32 elements per segment. + scratch = torch.zeros(page_size_el, dtype=torch.int32, device=self.device) + seg_addrs = torch.tensor( + [scratch.data_ptr()] * n_segs, dtype=torch.uint64, device=self.device + ) + chunks_per_block = n_segs * (page_size_el // blk_size) + for n_blocks in (1, 16, 3): + block_ids = torch.zeros(n_blocks, dtype=torch.int64, device=self.device) + _zero_kv_blocks_kernel[(n_blocks * chunks_per_block,)]( + seg_addrs, + block_ids, + n_blocks, + N_SEGS=n_segs, + PAGE_SIZE_EL=page_size_el, + BLOCK_SIZE=blk_size, + ) + def zero_block_ids(self, block_ids: list[int]) -> None: """Zero the KV cache memory for the given block IDs.""" if not block_ids or self._meta is None: