diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93..2486b9d5f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -65,6 +65,13 @@ else: logger = init_logger(__name__) +# Default CUDA graph coverage for chunked-prefill steps (see +# `_set_cudagraph_sizes`). Mixed prompt-chunk + decode steps are captured up +# to this many tokens; the operator's `max_cudagraph_capture_size` still wins. +CHUNKED_PREFILL_GRAPH_CEILING = 2048 +# Capture sizes are spaced 16 apart up to here and 64 apart above it. +FINE_STRIDE_LIMIT = 512 + DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( { "Qwen3ForCausalLM", @@ -286,6 +293,18 @@ OPTIMIZATION_LEVEL_TO_CONFIG = { } +# Prompt-chunk defaults used by `_fit_prefill_chunk_to_cudagraphs`. Chunks are +# rounded down to a multiple of the granule (the FLA chunk size, and a multiple +# of every attention block size vLLM selects) and never made smaller than the +# minimum, below which prefill GEMMs lose too much efficiency to be worth it. +PREFILL_CHUNK_GRANULE = 64 +MIN_PREFILL_CHUNK = 256 +# Upper bound on the default prompt chunk. The graph-fit rule +# (largest captured size minus one step of every running sequence) +# only lowers it further; it never raises a chunk above this. +PREFILL_CHUNK_TARGET = 1536 + + @config(config=ConfigDict(arbitrary_types_allowed=True)) class VllmConfig: """Dataclass which contains all vllm-related configuration. This @@ -1298,6 +1317,8 @@ class VllmConfig: else: self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + self._fit_prefill_chunk_to_cudagraphs() + if self.cache_config.kv_sharing_fast_prefill: if ( self.speculative_config is not None @@ -1632,6 +1653,53 @@ class VllmConfig: f" got {max_num_batched_tokens=} and {max_num_scheduled_tokens=}." ) + def _fit_prefill_chunk_to_cudagraphs(self) -> None: + """Default `long_prefill_token_threshold` to a graph-friendly chunk. + + With chunked prefill the scheduler gives a new prompt whatever token + budget is left after the running decodes, so a single 2k-token prompt + turns that step into a >2k-token eager forward pass which every + running decode has to sit through. Capping the prompt chunk at + `PREFILL_CHUNK_TARGET` tokens -- or lower, so that the chunk plus one + step of every running sequence still fits the largest captured CUDA + graph -- keeps mixed prefill/decode steps on captured graphs and + spreads the prompt work over a few steps. + + Trade-off: prompts wider than the chunk need one extra step per chunk + before their first token (the prefill compute itself is unchanged), + while the decodes sharing those steps see several short stalls + instead of one long one. + + The operator's own threshold always wins. Nothing is done without + CUDA graphs (there is no size to fit), for pooling runners (no decodes + to interleave) or in mamba "align" cache mode, whose chunks must be + block aligned. + """ + sched = self.scheduler_config + if ( + sched.long_prefill_token_threshold > 0 + or not sched.enable_chunked_prefill + or sched.runner_type != "generate" + or self.cache_config.mamba_cache_mode == "align" + ): + return + largest_graph = self.compilation_config.max_cudagraph_capture_size + if not largest_graph: + return + decode_tokens = sched.max_num_seqs * (1 + self.num_speculative_tokens) + chunk = min(PREFILL_CHUNK_TARGET, largest_graph - decode_tokens) + chunk = chunk // PREFILL_CHUNK_GRANULE * PREFILL_CHUNK_GRANULE + if chunk < MIN_PREFILL_CHUNK or chunk >= sched.max_num_batched_tokens: + return + sched.long_prefill_token_threshold = chunk + logger.info( + "Prompt chunks default to %d tokens so that a chunk plus %d decode " + "tokens fits the largest captured CUDA graph (%d tokens).", + chunk, + decode_tokens, + largest_graph, + ) + def _set_cudagraph_sizes(self): """ vLLM defines the default candidate list of batch sizes for CUDA graph @@ -1690,6 +1758,21 @@ class VllmConfig: max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) + if self.scheduler_config.enable_chunked_prefill: + # With chunked prefill, a step normally carries one prompt + # chunk on top of every running decode, so it is hundreds + # of tokens wide even when max_num_seqs is small. A step + # wider than the largest captured size runs eagerly and + # pays per-layer dispatch overhead on exactly the steps + # that already stall the decodes, so stretch the default + # far enough to keep those mixed steps on captured graphs. + max_cudagraph_capture_size = max( + max_cudagraph_capture_size, + min( + self.scheduler_config.max_num_batched_tokens, + CHUNKED_PREFILL_GRAPH_CEILING, + ), + ) max_num_tokens = self.scheduler_config.max_num_batched_tokens max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) @@ -1727,10 +1810,25 @@ class VllmConfig: range(8, min(max_cudagraph_capture_size + 1, 256), 8) ) if max_cudagraph_capture_size >= 256: - # Step size 16 for larger batch sizes + # Step size 16 for larger batch sizes, up to 512 (included) + cudagraph_capture_sizes += list( + range( + 256, + min(max_cudagraph_capture_size, FINE_STRIDE_LIMIT) + 1, + 16, + ) + ) + if max_cudagraph_capture_size > FINE_STRIDE_LIMIT: + # Sizes above 512 only serve mixed prefill/decode steps, + # where a few dozen tokens of padding are cheap next to + # the prompt chunk itself. A 64-token stride keeps the + # number of graphs, capture time and graph memory bounded. + # The ceiling is captured exactly so the widest step that + # is meant to hit a graph never spills over to eager. cudagraph_capture_sizes += list( - range(256, max_cudagraph_capture_size + 1, 16) + range(FINE_STRIDE_LIMIT + 64, max_cudagraph_capture_size, 64) ) + cudagraph_capture_sizes.append(max_cudagraph_capture_size) # ensure max_num_tokens is captured if within max capture size if ( max_num_tokens <= max_cudagraph_capture_size diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef174135..a23be01ef 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -40,7 +40,7 @@ from vllm.exceptions import VLLMValidationError from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.logprobs import Logprob -from vllm.outputs import RequestOutput +from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.utils.async_utils import merge_async_iterators @@ -395,46 +395,61 @@ 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 - ), - ) - ], + parts = _split_delta_per_token( + output, delta_text, logprobs, previous_num_tokens[i] ) - # 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, + last_part = len(parts) - 1 + for part_idx, ( + part_text, + part_token_ids, + part_logprobs, + tokens_so_far, + ) in enumerate(parts): + is_last = part_idx == last_part + chunk = CompletionStreamResponse( + id=request_id, + object="text_completion", + created=created_time, + model=model_name, + choices=[ + CompletionResponseStreamChoice( + index=i, + text=part_text, + logprobs=part_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 part_idx == 0 + else None + ), + token_ids=( + as_list(part_token_ids) + if request.return_token_ids + else None + ), + ) + ], ) + # Stamp on terminal chunk only when no trailing usage + # chunk will follow (that one is the true final message). + if ( + 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] + chunk.usage = UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=tokens_so_far, + total_tokens=prompt_tokens + tokens_so_far, + ) - response_json = chunk.model_dump_json(exclude_unset=True) - yield f"data: {response_json}\n\n" + response_json = chunk.model_dump_json(exclude_unset=True) + yield f"data: {response_json}\n\n" total_prompt_tokens = sum(num_prompt_tokens) total_completion_tokens = sum(previous_num_tokens) @@ -689,3 +704,63 @@ class OpenAIServingCompletion(OpenAIServing): tokens=out_tokens, top_logprobs=out_top_logprobs, ) + + +def _split_delta_per_token( + output: CompletionOutput, + delta_text: str, + logprobs: CompletionLogProbs | None, + tokens_so_far: int, +) -> list[tuple[str, GenericSequence[int], CompletionLogProbs | None, int]]: + """Break a streamed delta into one part per generated token. + + An engine step can hand the frontend several tokens at once (speculative + decoding accepts a whole draft per step, and consecutive steps are merged + when the frontend falls behind). Sent as one chunk they would hide the + token granularity from the client, so the delta is cut along the + per-token text the incremental detokenizer recorded for it. That is the + very text a one-token-per-step engine would have streamed: nothing is + decoded a second time and the concatenation is unchanged. A delta with at + most one token, or one without a usable per-token record, is passed + through as a single part. + + Each part is ``(text, token_ids, logprobs, completion tokens streamed up + to and including this part)``. ``delta_text`` may start with echoed prompt + text; that prefix stays on the first part, as do the prompt entries of + echoed logprobs. + """ + token_ids = output.token_ids + num_tokens = len(token_ids) + pieces = output.per_token_text + if ( + num_tokens <= 1 + or pieces is None + or len(pieces) != num_tokens + or "".join(pieces) != output.text + ): + return [(delta_text, token_ids, logprobs, tokens_so_far)] + + lead = delta_text[: len(delta_text) - len(output.text)] + tokens_before = tokens_so_far - num_tokens + parts: list[tuple[str, GenericSequence[int], CompletionLogProbs | None, int]] = [] + for j, piece in enumerate(pieces): + part_logprobs = None + if logprobs is not None: + prompt_entries = len(logprobs.tokens) - num_tokens + lo = 0 if j == 0 else prompt_entries + j + hi = prompt_entries + j + 1 + part_logprobs = CompletionLogProbs( + text_offset=logprobs.text_offset[lo:hi], + token_logprobs=logprobs.token_logprobs[lo:hi], + tokens=logprobs.tokens[lo:hi], + top_logprobs=logprobs.top_logprobs[lo:hi], + ) + parts.append( + ( + lead + piece if j == 0 else piece, + [token_ids[j]], + part_logprobs, + tokens_before + j + 1, + ) + ) + return parts diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4ac8d49cd..e02737f96 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -321,8 +321,12 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] ] = { PlatformEnum.CUDA: [ - FlashInferFp8DeepGEMMDynamicBlockScaledKernel, + # DeepGEMM first: one kernel for every M, so no torch.cond dispatch on + # the batch size and no TRT-LLM/FlashInfer JIT build on the cold-start + # path. The FlashInfer/DeepGEMM hybrid stays available as the next + # choice; it is only skipped when DeepGEMM itself can serve the layer. DeepGemmFp8BlockScaledMMKernel, + FlashInferFp8DeepGEMMDynamicBlockScaledKernel, CutlassFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, diff --git a/vllm/model_executor/warmup/first_step_warmup.py b/vllm/model_executor/warmup/first_step_warmup.py new file mode 100644 index 000000000..4d2176ba1 --- /dev/null +++ b/vllm/model_executor/warmup/first_step_warmup.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compile the Triton kernels a first real request would otherwise JIT. + +Startup exercises the model through dummy runs that carry no attention +metadata (the profiling pass) and through CUDA graph capture of uniform +decode batches. On a hybrid GDN model that leaves a handful of Triton +kernels uncompiled until a real request reaches them, each stalling that +request for a few hundred milliseconds: + +* the prefill convolution (``_causal_conv1d_fwd_kernel``), which only runs + on prompt tokens; +* the fused post-conv prep kernel for prompt lengths that are not a + multiple of 16 -- Triton specialises its untyped ``L`` argument on + ``== 1`` and ``% 16 == 0``, and the profiling warmup only covers ``L == 64``; +* the sigmoid-gating recurrent kernel that serves the decodes of a mixed + prefill+decode step (pure decode steps use the packed kernel), plus the + single-token conv update for batch sizes outside the captured + specialisations; +* the slot-mapping and KV block zeroing kernels on the bookkeeping side of + every real step (the zeroing kernel is specialised on the block count); +* with speculative decoding, the all-greedy specialisation of the greedy + rejection kernel (``is_greedy=None``), which the dummy sampler run never + produces. + +Each group runs once, on freshly allocated tensors that mirror the dtypes, +shapes and strides of the real call (all part of Triton's compile key). +State caches are handed in only for their geometry: every cache index +passed alongside them is the null block, which the kernels skip, so no +engine buffer is read or written and every result is discarded. Groups are +isolated by try/except -- a failed warmup only costs the latency it was +meant to save. +""" + +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__) + +# Triton specialises an untyped integer argument into three classes: exactly +# 1, a multiple of 16, and everything else. One launch per class covers a +# kernel whose batch or block count is such an argument. +INT_SPECIALISATION_CLASSES = (1, 16, 3) +# Block 0 is the null block; the state kernels skip sequences that point at it. +NULL_STATE_INDEX = 0 + + +def _index_view( + rows: int, stride: int, value: int, device: torch.device +) -> torch.Tensor: + """1-D int32 index tensor whose element stride is ``stride``. + + The GDN metadata builder hands the kernels a ``block_table[:, 0]`` view, + so the element stride equals the block-table row length and is baked into + the kernels as a constexpr; the warmup has to reproduce it. + """ + full = torch.full((rows, stride), value, dtype=torch.int32, device=device) + return full[:, 0] + + +def _mamba_index_strides(runner) -> list[int]: + """Strides of every state-index view the GDN kernels can receive. + + Pure decode steps use the metadata builder's contiguous buffers (stride + 1); prefill and mixed steps use the block-table view of each mamba group. + """ + from vllm.v1.kv_cache_interface import MambaSpec + + strides = {1} + kv_cache_config = getattr(runner, "kv_cache_config", None) + block_tables = getattr(getattr(runner, "input_batch", None), "block_table", None) + if kv_cache_config is not None and block_tables is not None: + for gid, group in enumerate(kv_cache_config.kv_cache_groups): + if isinstance(group.kv_cache_spec, MambaSpec): + strides.add(int(block_tables[gid].block_table.gpu.stride(0))) + return sorted(strides) + + +def _find_gdn_layer(model: torch.nn.Module): + """First Qwen GDN layer whose state cache is already bound, if any.""" + from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, + ) + + for module in model.modules(): + if not isinstance(module, QwenGatedDeltaNetAttention): + continue + caches = getattr(module, "kv_cache", None) + if caches is not None and len(caches) >= 2 and caches[0].numel() > 0: + return module + return None + + +def _warm_gdn_kernels(worker: "Worker") -> None: + from vllm.model_executor.layers.fla.ops import ( + fused_post_conv_prep, + fused_recurrent_gated_delta_rule_packed_decode, + fused_sigmoid_gating_delta_rule_update, + ) + from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, + ) + + layer = _find_gdn_layer(worker.get_model()) + if layer is None: + return + runner = worker.model_runner + device = runner.device + act_dtype = worker.vllm_config.model_config.dtype + assert isinstance(act_dtype, torch.dtype) + + num_k_heads = layer.num_k_heads // layer.tp_size + num_v_heads = layer.num_v_heads // layer.tp_size + head_k_dim, head_v_dim = layer.head_k_dim, layer.head_v_dim + qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + conv_weight = layer.conv1d.weight.view( + layer.conv1d.weight.size(0), layer.conv1d.weight.size(2) + ) + assert conv_weight.size(0) == qkv_dim, (conv_weight.shape, qkv_dim) + # The conv kernels want (..., dim, width - 1); mirror the layer's own + # layout handling so the strides match the real launch. + conv_state = ( + layer.kv_cache[0] + if is_conv_state_dim_first() + else layer.kv_cache[0].transpose(-1, -2) + ) + ssm_state = layer.kv_cache[1] + index_strides = _mamba_index_strides(runner) + + def zeros(*shape: int, dtype: torch.dtype = act_dtype) -> torch.Tensor: + return torch.zeros(shape, dtype=dtype, device=device) + + # 1) Prefill: the varlen convolution over a small three-sequence batch, + # channel-last exactly like `mixed_qkv.transpose(0, 1)` in the layer. + # Every cache index is the null block, so the real conv cache only + # contributes its geometry (`num_cache_lines` and strides are + # constexprs) and is neither read nor written. + seq_len, num_seqs = 8, 3 + x = zeros(seq_len * num_seqs, qkv_dim).transpose(0, 1) + query_start_loc = torch.arange( + 0, seq_len * (num_seqs + 1), seq_len, dtype=torch.int32, device=device + ) + has_initial_state = zeros(num_seqs, dtype=torch.bool) + for stride in index_strides: + causal_conv1d_fn( + x, + conv_weight, + layer.conv1d.bias, + activation=layer.activation, + conv_states=conv_state, + has_initial_state=has_initial_state, + cache_indices=_index_view(num_seqs, stride, NULL_STATE_INDEX, device), + query_start_loc=query_start_loc, + ) + + # 2) Post-conv prep for every specialisation class of its token count. + for num_tokens in INT_SPECIALISATION_CLASSES: + fused_post_conv_prep( + conv_output=zeros(num_tokens, qkv_dim), + a=zeros(num_tokens, num_v_heads), + b=zeros(num_tokens, num_v_heads), + A_log=layer.A_log, + dt_bias=layer.dt_bias, + num_k_heads=num_k_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + apply_l2norm=True, + output_g_exp=False, + ) + + # 3) Recurrent kernels. A two-row stand-in with the real page stride + # replaces the SSM cache (only its row stride is a constexpr); row 1 + # is the one the dummies point at, row 0 stays the null block. + ssm_shadow = torch.empty_strided( + (2, *ssm_state.shape[1:]), + ssm_state.stride(), + dtype=ssm_state.dtype, + device=device, + ).zero_() + num_decodes = 3 + cu_seqlens = torch.arange(num_decodes + 1, dtype=torch.int32, device=device) + for stride in index_strides: + # Decodes sharing a step with a prefill chunk take this kernel. + fused_sigmoid_gating_delta_rule_update( + A_log=layer.A_log, + a=zeros(num_decodes, num_v_heads), + b=zeros(num_decodes, num_v_heads), + dt_bias=layer.dt_bias, + q=zeros(1, num_decodes, num_k_heads, head_k_dim), + k=zeros(1, num_decodes, num_k_heads, head_k_dim), + v=zeros(1, num_decodes, num_v_heads, head_v_dim), + initial_state=ssm_shadow, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=_index_view(num_decodes, stride, 1, device), + use_qk_l2norm_in_kernel=True, + ) + + # 4) Pure-decode path. Graph capture normally compiles these, but only + # for the captured batch sizes and the contiguous index buffer; this + # covers the remaining integer classes and the strided view, and keeps + # eager mode free of first-request JIT as well. + for batch in INT_SPECIALISATION_CLASSES: + for stride in index_strides: + causal_conv1d_update( + zeros(batch, qkv_dim), + conv_state, + conv_weight, + layer.conv1d.bias, + layer.activation, + conv_state_indices=_index_view(batch, stride, NULL_STATE_INDEX, device), + validate_data=False, + ) + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=zeros(num_decodes, qkv_dim), + a=zeros(num_decodes, num_v_heads), + b=zeros(num_decodes, num_v_heads), + A_log=layer.A_log, + dt_bias=layer.dt_bias, + scale=head_k_dim**-0.5, + initial_state=ssm_shadow, + out=zeros(num_decodes, 1, num_v_heads, head_v_dim), + ssm_state_indices=_index_view(num_decodes, 1, 1, device), + use_qk_l2norm_in_kernel=True, + ) + + +def _warm_slot_mapping(worker: "Worker") -> None: + from vllm.v1.worker.block_table import PAD_SLOT_ID, _compute_slot_mapping_kernel + + runner = worker.model_runner + block_tables = getattr(getattr(runner, "input_batch", None), "block_table", None) + if block_tables is None: + return + device = runner.device + num_reqs, num_tokens = 1, 8 + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + positions = torch.zeros(num_tokens, dtype=torch.int64, device=device) + slot_mapping = torch.zeros(num_tokens, dtype=torch.int64, device=device) + for table in block_tables.block_tables: + # The row stride and block size are specialised integers, so pass + # the real values; the table contents are a throwaway all-zero row. + row_stride = int(table.block_table.gpu.stride(0)) + dummy_table = torch.zeros( + (num_reqs, row_stride), dtype=torch.int32, device=device + ) + _compute_slot_mapping_kernel[(num_reqs + 1,)]( + num_tokens, + # max_num_tokens == num_tokens leaves the padding loop empty, so + # the dummy slot mapping is never written out of bounds. + num_tokens, + query_start_loc, + positions, + dummy_table, + row_stride, + table.block_size, + slot_mapping, + TOTAL_CP_WORLD_SIZE=table.pcp_world_size * table.dcp_world_size, + TOTAL_CP_RANK=table.pcp_rank * table.dcp_world_size + table.dcp_rank, + CP_KV_CACHE_INTERLEAVE_SIZE=table.cp_kv_cache_interleave_size, + PAD_ID=PAD_SLOT_ID, + BLOCK_SIZE=1024, + ) + + +def _warm_kv_block_zeroing(worker: "Worker") -> None: + from vllm.v1.worker.utils import _zero_kv_blocks_kernel + + zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) + meta = getattr(zeroer, "_meta", None) + if meta is None: + return + _, page_size_el, blk_size, n_segs = meta + device = worker.model_runner.device + # Point every segment at one scratch page; with block id 0 the kernel + # writes exactly that page, never the KV cache. + scratch = torch.zeros(page_size_el, dtype=torch.int32, device=device) + seg_addrs = torch.tensor( + [scratch.data_ptr()] * n_segs, dtype=torch.uint64, device=device + ) + programs_per_block = n_segs * (page_size_el // blk_size) + for n_blocks in INT_SPECIALISATION_CLASSES: + block_ids = torch.zeros(n_blocks, dtype=torch.int64, device=device) + _zero_kv_blocks_kernel[(n_blocks * programs_per_block,)]( + seg_addrs, + block_ids, + n_blocks, + N_SEGS=n_segs, + PAGE_SIZE_EL=page_size_el, + BLOCK_SIZE=blk_size, + ) + + +def _warm_greedy_rejection(worker: "Worker") -> None: + runner = worker.model_runner + sampler = getattr(runner, "rejection_sampler", None) + spec_config = worker.vllm_config.speculative_config + if sampler is None or spec_config is None: + return + from vllm.v1.sample.logits_processor import LogitsProcessors + from vllm.v1.sample.metadata import SamplingMetadata + from vllm.v1.sample.rejection_sampler import rejection_sample + + device = runner.device + num_draft = max(int(spec_config.num_speculative_tokens or 1), 1) + vocab_size = 8 + unused = torch.zeros(1, dtype=torch.float32, device=device) + all_greedy = 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=unused, + presence_penalties=unused, + repetition_penalties=unused, + output_token_ids=[[]], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + # Going through the wrapper reproduces the exact argument dtypes of the + # real call (int32 ids, float32 logits, int64 argmax) and the + # `is_greedy=None` specialisation an all-greedy batch triggers. + rejection_sample( + torch.zeros(num_draft, dtype=torch.int32, device=device), + [num_draft], + num_draft, + torch.tensor([num_draft], dtype=torch.int32, device=device), + None, + torch.zeros((num_draft, vocab_size), dtype=torch.float32, device=device), + torch.zeros((1, 1), dtype=torch.int32, device=device), + all_greedy, + synthetic_mode=sampler.synthetic_mode, + synthetic_conditional_rates=sampler.synthetic_conditional_rates, + use_fp64_gumbel=sampler.use_fp64_gumbel, + ) + + +_WARMUP_GROUPS = ( + ("GDN prefill and decode kernels", _warm_gdn_kernels), + ("slot-mapping kernel", _warm_slot_mapping), + ("KV block zeroing kernel", _warm_kv_block_zeroing), + ("all-greedy rejection kernel", _warm_greedy_rejection), +) + + +def first_step_kernel_warmup(worker: "Worker") -> None: + """Compile the kernels listed in the module docstring. + + Needs the KV cache bound (state geometry, block tables and the zeroing + table come from it) and should run before graph capture so no capture + pays for JIT either; ``kernel_warmup`` is called at exactly that point. + """ + for name, fn in _WARMUP_GROUPS: + try: + with torch.inference_mode(): + fn(worker) + torch.accelerator.synchronize() + except Exception: + logger.warning( + "Startup warmup of the %s failed; the first request will " + "compile it instead.", + name, + exc_info=True, + ) + else: + logger.debug("Startup warmup of the %s done.", name) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 754270e65..813ad39a2 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -109,6 +109,17 @@ def kernel_warmup(worker: "Worker"): create_mixed_batch=True, ) + # Compile the Triton kernels that only a real request would otherwise + # reach (prefill conv, mixed-step recurrent kernel, slot mapping, KV + # zeroing, all-greedy rejection). Imported lazily so this module keeps + # its import graph. + if current_platform.is_cuda(): + from vllm.model_executor.warmup.first_step_warmup import ( + first_step_kernel_warmup, + ) + + first_step_kernel_warmup(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/outputs.py b/vllm/outputs.py index 2c71d2afb..fc879d65c 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -35,6 +35,9 @@ class CompletionOutput: to stop, None if the completion finished for some other reason including encountering the EOS token. lora_request: The LoRA request that was used to generate the output. + per_token_text: For a streamed delta, the slice of `text` produced by + each entry of `token_ids` (same length, concatenates to `text`); + None when not tracked. """ index: int @@ -46,6 +49,7 @@ class CompletionOutput: finish_reason: str | None = None stop_reason: int | str | None = None lora_request: LoRARequest | None = None + per_token_text: list[str] | None = None def finished(self) -> bool: return self.finish_reason is not None @@ -153,6 +157,16 @@ class RequestOutput: if completion.index == next_completion.index: if aggregate: # Merge outputs with same index + if ( + completion.per_token_text is not None + and next_completion.per_token_text is not None + ): + completion.per_token_text = [ + *completion.per_token_text, + *next_completion.per_token_text, + ] + else: + completion.per_token_text = 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/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index 4700eecb5..cb5eccf7f 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -45,6 +45,12 @@ class IncrementalDetokenizer: def get_next_output_text(self, finished: bool, delta: bool) -> str: return "" + def get_next_output_delta( + self, finished: bool, num_tokens: int + ) -> tuple[str, list[str] | None]: + """Next delta text with one text slice per token of the delta.""" + return "", [""] * num_tokens + @classmethod def from_new_request( cls, @@ -88,6 +94,10 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): else: self.stop_buffer_length = 0 self._last_output_text_offset: int = 0 + # End offset in ``output_text`` of every token decoded since the last + # delta was released; ``get_next_output_delta`` turns these into the + # per-token text slices of the next delta. + self._pending_token_ends: list[int] = [] # Generation data self.output_text = "" @@ -117,13 +127,16 @@ 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._pending_token_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 stop token is still + # a generated token, so it gets a zero-width slice. self.token_ids.append(skipped_stop_token_id) + self._pending_token_ends.append(len(self.output_text)) # 2) Evaluate stop strings. stop_string = None @@ -138,6 +151,9 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): stop_string, truncate_to = stop if truncate_to != -1: self.output_text = self.output_text[:truncate_to] + self._pending_token_ends = [ + min(end, truncate_to) for end in self._pending_token_ends + ] return stop_string @@ -152,6 +168,8 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): # We return the full output text if the sequence is finished. buffer_length = 0 if finished else self.stop_buffer_length if not delta: + # Per-token slices only matter for deltas; drop the bookkeeping. + self._pending_token_ends.clear() if not buffer_length: return self.output_text return self.output_text[:-buffer_length] @@ -163,6 +181,38 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): return self.output_text[last_offset:length] return "" + def get_next_output_delta( + self, finished: bool, num_tokens: int + ) -> tuple[str, list[str] | None]: + """Return the next delta text together with its per-token slices. + + The slices concatenate to exactly the returned text and there is one + per token of the delta, so a streaming frontend can send one chunk per + generated token without decoding anything a second time. Text held + back for stop-string matching is released with a later delta and then + belongs to that delta's first slice; a stop token that is not + detokenized gets an empty slice. If the bookkeeping does not line up + with ``num_tokens`` the slices are ``None`` and the caller should treat + the delta as a single piece. + """ + start = self._last_output_text_offset + text = self.get_next_output_text(finished, delta=True) + ends = self._pending_token_ends + self._pending_token_ends = [] + if len(ends) != num_tokens: + return text, None + limit = start + len(text) + slices: list[str] = [] + for end in ends: + end = min(max(end, start), limit) + slices.append(self.output_text[start:end]) + start = end + if slices and start < limit: + # Only if an end offset was clamped below the released text; + # keep the concatenation exact by extending the last slice. + slices[-1] += self.output_text[start:limit] + return text, slices + 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..281fb20f5 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: + per_token_text: list[str] | None = None + if delta: + text, per_token_text = self.detokenizer.get_next_output_delta( + finished, 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, + per_token_text=per_token_text, ) def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: