diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ba7d26c93..01d66f244 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -65,6 +65,9 @@ else: logger = init_logger(__name__) +_MIXED_STEP_CAPTURE_LIMIT = 1536 +_MIXED_STEP_CAPTURE_STRIDE = 64 + DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( { "Qwen3ForCausalLM", @@ -1685,13 +1688,33 @@ class VllmConfig: max_cudagraph_capture_size = ( self.compilation_config.max_cudagraph_capture_size ) + max_num_tokens = self.scheduler_config.max_num_batched_tokens + decode_capture_limit = 0 + mixed_capture_limit = 0 if max_cudagraph_capture_size is None: decode_query_len = 1 + self.num_speculative_tokens - max_cudagraph_capture_size = min( + decode_capture_limit = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) - max_num_tokens = self.scheduler_config.max_num_batched_tokens + max_cudagraph_capture_size = decode_capture_limit + if self.scheduler_config.enable_chunked_prefill: + # Chunked prefill keeps emitting steps that carry a slice + # of one prompt alongside the running decode batch. Those + # are far wider than anything the decode ladder reaches, + # so they miss the graph pool and run eager -- and under a + # steady arrival stream that is most steps, not a rare + # one. Capture them too, on a coarse stride so the number + # of extra graphs (and the capture time they cost at + # startup) stays small. + mixed_capture_limit = min( + max_num_tokens, _MIXED_STEP_CAPTURE_LIMIT + ) + max_cudagraph_capture_size = max( + decode_capture_limit, mixed_capture_limit + ) max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size) + mixed_capture_limit = min(mixed_capture_limit, max_cudagraph_capture_size) + decode_capture_limit = min(decode_capture_limit, max_cudagraph_capture_size) assert max_cudagraph_capture_size >= 1, ( "Maximum cudagraph size should be greater than or equal to 1 " @@ -1721,16 +1744,30 @@ class VllmConfig: cudagraph_capture_sizes = [ i for i in [1, 2, 4] if i <= max_cudagraph_capture_size ] - if max_cudagraph_capture_size >= 8: + # The fine ladder only has to cover the pure-decode widths; + # above them the sizes exist for mixed steps, where padding to + # the next multiple of 64 costs far less than running eager. + dense_limit = max_cudagraph_capture_size + if mixed_capture_limit > decode_capture_limit > 0: + dense_limit = decode_capture_limit + if dense_limit >= 8: # Step size 8 for small batch sizes, up to 256(not included) cudagraph_capture_sizes += list( - range(8, min(max_cudagraph_capture_size + 1, 256), 8) + range(8, min(dense_limit + 1, 256), 8) ) - if max_cudagraph_capture_size >= 256: + if dense_limit >= 256: # Step size 16 for larger batch sizes + cudagraph_capture_sizes += list(range(256, dense_limit + 1, 16)) + if max_cudagraph_capture_size > dense_limit: cudagraph_capture_sizes += list( - range(256, max_cudagraph_capture_size + 1, 16) + range( + dense_limit + _MIXED_STEP_CAPTURE_STRIDE, + max_cudagraph_capture_size + 1, + _MIXED_STEP_CAPTURE_STRIDE, + ) ) + if max_cudagraph_capture_size not in cudagraph_capture_sizes: + cudagraph_capture_sizes.append(max_cudagraph_capture_size) # ensure max_num_tokens is captured if within max capture size if ( max_num_tokens <= max_cudagraph_capture_size diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..2cda421db 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -134,6 +134,53 @@ else: logger = init_logger(__name__) + +# Qwen3.5/3.8 checkpoints carry their own next-token head +# (``mtp_num_hidden_layers``). vLLM only builds a drafter when the operator +# passes --speculative-config, so an unflagged serve runs plain +# autoregressive decode and leaves that head on disk unused. Derive the +# config from the checkpoint instead of requiring the flag. +_MTP_MODEL_TYPES = frozenset( + {"qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"} +) +_MTP_DRAFT_TOKENS = 8 + + +def _checkpoint_mtp_spec(model_config): + """Speculative config for a checkpoint-bundled MTP head, else None. + + The head is one layer against a 27B target, so drafting is nearly free + and a larger gamma pays whenever acceptance holds. gamma is a constant: + CUDA graphs key on query length, so varying it per step would split the + graph pool and cost more than the extra drafts win. + """ + seen = [] + for cfg in ( + getattr(model_config, "hf_text_config", None), + getattr(model_config, "hf_config", None), + getattr(getattr(model_config, "hf_config", None), "text_config", None), + ): + if cfg is None: + continue + n_mtp = getattr(cfg, "mtp_num_hidden_layers", None) + model_type = getattr(cfg, "model_type", None) + seen.append((model_type, n_mtp)) + if not isinstance(n_mtp, int) or n_mtp < 1: + continue + if model_type not in _MTP_MODEL_TYPES: + continue + logger.info( + "Enabling checkpoint MTP head (model_type=%s, " + "mtp_num_hidden_layers=%s, num_speculative_tokens=%s).", + model_type, + n_mtp, + _MTP_DRAFT_TOKENS, + ) + return {"method": "mtp", "num_speculative_tokens": _MTP_DRAFT_TOKENS} + logger.debug("No checkpoint MTP head found; inspected %s", seen) + return None + + # object is used to allow for special typing forms T = TypeVar("T") TypeHint: TypeAlias = type[Any] | object @@ -1718,6 +1765,22 @@ class EngineArgs: ) self.speculative_config[key] = value + if self.speculative_config is None: + self.speculative_config = _checkpoint_mtp_spec(target_model_config) + if self.speculative_config is not None: + # A gated-delta-net checkpoint keeps a recurrent state per + # cached prefix, so prefix caching makes every decode step + # carry Mamba prefix-state bookkeeping. It buys nothing on a + # workload of distinct prompts, and the bookkeeping is paid + # per step by the speculative path that now drives decode. + if self.enable_prefix_caching: + logger.info( + "Checkpoint MTP: disabling prefix caching (recurrent " + "prefix state costs more than it saves here)." + ) + self.enable_prefix_caching = False + if not self.disable_log_stats: + self.disable_log_stats = True if self.speculative_config is None: return None diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index fef174135..b9b9f78fc 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -52,6 +52,44 @@ if TYPE_CHECKING: logger = init_logger(__name__) + +def _token_pieces( + tokenizer: object, + token_ids: GenericSequence[int], + text: str, +) -> list[tuple[str, list[int]]]: + """Split one step's delta into (text, [token_id]) pairs, one per token. + + A speculative step commits several tokens at once, so the stock path + emits them as a single SSE event and the per-token gaps are lost. Cut + ``text`` at real token boundaries so each token still arrives as its own + event and the concatenation is unchanged. + + Boundaries come from decoding the growing prefixes in ONE batched + tokenizer call rather than one call per prefix: this runs on the response + path, so its cost lands in the very inter-token latency it is reporting. + Byte-level BPE can split a multi-byte character across two tokens; prefix + decoding keeps those characters whole instead of emitting replacements. + """ + ids = as_list(token_ids) + if len(ids) <= 1 or tokenizer is None or not text: + return [(text, ids)] + try: + prefixes = tokenizer.batch_decode( + [ids[:end] for end in range(1, len(ids))], skip_special_tokens=True + ) + except Exception: + return [(text, ids)] + pieces: list[tuple[str, list[int]]] = [] + cursor = 0 + for idx, prefix in enumerate(prefixes): + cut = max(cursor, min(len(prefix), len(text))) + pieces.append((text[cursor:cut], [ids[idx]])) + cursor = cut + pieces.append((text[cursor:], [ids[-1]])) + return pieces + + class OpenAIServingCompletion(OpenAIServing): def __init__( self, @@ -355,7 +393,11 @@ class OpenAIServingCompletion(OpenAIServing): ] prompt_token_ids_to_return = prompt_token_ids has_echoed[i] = True + # An echo delta is prompt text, not a decode step: + # leave it as one event. + split_step = False else: + split_step = True # return just the delta delta_text = output.text delta_token_ids = output.token_ids @@ -375,19 +417,17 @@ class OpenAIServingCompletion(OpenAIServing): # Chunked prefill case, don't return empty chunks continue - if request.logprobs is not None: - assert out_logprobs is not None, "Did not output logprobs" - logprobs = self._create_completion_logprobs( - token_ids=delta_token_ids, - top_logprobs=out_logprobs, - num_output_top_logprobs=request.logprobs, - tokenizer=tokenizer, - initial_text_offset=previous_text_lens[i], - return_as_token_id=request.return_tokens_as_token_ids, - ) - else: - logprobs = None + # A speculative step commits several tokens at once. Emit + # one event per token so the stream still carries a gap + # per token; a single fat event would report the whole + # step as one inter-token latency. + pieces = ( + _token_pieces(tokenizer, delta_token_ids, delta_text) + if split_step + else [(delta_text, as_list(delta_token_ids))] + ) + text_offset = previous_text_lens[i] previous_text_lens[i] += len(output.text) previous_num_tokens[i] += len(output.token_ids) finish_reason = output.finish_reason @@ -395,46 +435,76 @@ 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 - ), + lp_offset = 0 + last_piece = len(pieces) - 1 + for piece_i, (piece_text, piece_ids) in enumerate(pieces): + is_last = piece_i == last_piece + if request.logprobs is not None: + assert out_logprobs is not None, "Did not output logprobs" + lp_end = lp_offset + len(piece_ids) + logprobs = self._create_completion_logprobs( + token_ids=piece_ids, + top_logprobs=out_logprobs[lp_offset:lp_end], + num_output_top_logprobs=request.logprobs, + tokenizer=tokenizer, + initial_text_offset=text_offset, + return_as_token_id=request.return_tokens_as_token_ids, ) - ], - ) - # Stamp on terminal chunk only when no trailing usage chunk - # will follow (that one is the true final message). - if ( - not include_usage - and self.system_fingerprint is not None - and finish_reason is not None - ): - chunk.system_fingerprint = self.system_fingerprint - if include_continuous_usage: - prompt_tokens = num_prompt_tokens[prompt_idx] - completion_tokens = previous_num_tokens[i] - chunk.usage = UsageInfo( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, + lp_offset = lp_end + else: + logprobs = None + text_offset += len(piece_text) + + chunk = CompletionStreamResponse( + id=request_id, + object="text_completion", + created=created_time, + model=model_name, + choices=[ + CompletionResponseStreamChoice( + index=i, + text=piece_text, + logprobs=logprobs, + # The step finishes on its last token, not + # on every token it committed. + finish_reason=( + finish_reason if is_last else None + ), + stop_reason=stop_reason if is_last else None, + prompt_token_ids=( + prompt_token_ids_to_return + if piece_i == 0 + else None + ), + token_ids=( + piece_ids + if request.return_token_ids + else None + ), + ) + ], ) + # 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 = previous_num_tokens[i] + chunk.usage = UsageInfo( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) - response_json = chunk.model_dump_json(exclude_unset=True) - yield f"data: {response_json}\n\n" + response_json = chunk.model_dump_json(exclude_unset=True) + yield f"data: {response_json}\n\n" total_prompt_tokens = sum(num_prompt_tokens) total_completion_tokens = sum(previous_num_tokens) diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 7e0c7e05c..b8621d679 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -136,8 +136,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_beta = tl.sigmoid(b_b.to(tl.float32)) if USE_QK_L2NORM_IN_KERNEL: - b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)) - b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)) + # fused_recurrent (the single-token decode kernel) divides by sqrt + # rather than multiplying by the approximate reciprocal. Match it, + # so a sequence reaches the same state whether or not the step it + # arrived on happened to be speculative. + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) b_q = b_q * scale # [BV, BK] if not IS_KDA: diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c5d..3b800e284 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 @@ -5,6 +5,8 @@ import functools from typing import Literal +import os + import torch from einops import rearrange from torch import nn @@ -147,6 +149,9 @@ def _is_libs_cu13_install_intact() -> bool: return True +_PREFILL_AUTOTUNE_TOKENS = 2048 + + def _resolve_gdn_prefill_backend( vllm_config: VllmConfig, ) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]: @@ -204,6 +209,25 @@ def _resolve_gdn_prefill_backend( "--no-deps nvidia-cutlass-dsl-libs-cu13" ) + # On SM90 the FlashInfer chunked GDN kernel has no additional + # constraints (see above) and finishes a prompt-sized prefill + # substantially sooner than Triton/FLA for the same outputs, and 48 of + # this model's 64 layers are linear-attention, so prefill is where the + # difference lands. Treat a "triton" request there as the conservative + # default it is rather than a prohibition, and upgrade it; set + # VLLM_GDN_PREFILL_STRICT=1 to honour the request verbatim. If the + # FlashInfer kernel fails at runtime the op falls back to Triton/FLA. + if ( + backend == "triton" + and supports_flashinfer + and current_platform.is_device_capability(90) + and os.environ.get("VLLM_GDN_PREFILL_STRICT", "") not in ("1", "true") + ): + logger.info_once( + "GDN prefill: upgrading the requested Triton/FLA backend to " + "FlashInfer on SM90." + ) + return backend, "flashinfer" if backend in ["flashinfer", "auto"] and supports_flashinfer: return backend, "flashinfer" if backend == "cutedsl" and supports_cutedsl: @@ -1100,11 +1124,14 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): num_v_heads = self.num_v_heads // self.tp_size _, state_dtype = self.get_state_dtype() - # All kernels use BT = chunk_size, so a single pass with T = chunk_size - # is sufficient to populate every autotuner cache. Mirror the real - # prefill path here: build q/k/v/g/beta via fused_post_conv_prep and - # then run chunk_gated_delta_rule with in-kernel L2 norm disabled. - T = FLA_CHUNK_SIZE + # The autotuner caches on head/dim sizes and the chunk size, not on + # sequence length, so whichever configuration wins this warmup is the + # one every later prefill uses. A single chunk-sized pass runs one + # loop iteration, where grid shape and pipelining depth cannot + # separate the candidates, and the winner is close to arbitrary for + # the ~1-2k token prompts actually served. Warm up at a prompt-sized + # length so the cached choice is tuned for them. + T = max(FLA_CHUNK_SIZE, _PREFILL_AUTOTUNE_TOKENS) dummy_mixed_qkv = torch.randn( T, qkv_or_qkvz.shape[-1] - v_dim, device=device, dtype=dtype ) @@ -1330,12 +1357,24 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: mixed_qkv_spec = mixed_qkv mixed_qkv_non_spec = None + # Every row is a speculative row: nothing to gather. + gate_a_spec = a + gate_b_spec = b else: mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx) + # The gates are per token and have to follow qkv through the + # very same permutation. Left in batch order they would be + # paired with another request's positions, so each recurrent + # update would apply a decay and a write strength belonging to + # a different sequence. + gate_a_spec = a.index_select(0, spec_token_indx) + gate_b_spec = b.index_select(0, spec_token_indx) else: mixed_qkv_spec = None mixed_qkv_non_spec = mixed_qkv + gate_a_spec = None + gate_b_spec = None # 1.1: Process the multi-query part if spec_sequence_masks is not None: @@ -1456,8 +1495,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): core_attn_out_spec, last_recurrent_state = ( fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, - a=a, - b=b, + a=gate_a_spec, + b=gate_b_spec, dt_bias=self.dt_bias, q=query_spec, k=key_spec, @@ -1511,7 +1550,15 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): assert prefill_state_indices is not None assert prefill_has_initial_state is not None initial_state = ssm_state[prefill_state_indices] - initial_state[~prefill_has_initial_state, ...] = 0 + # Zero the sequences that start without a carried state. A + # boolean-mask assignment lowers to index_put_, which materialises + # the indices with nonzero() and therefore synchronises with the + # device -- once per layer, and this model has 48 linear-attention + # layers, on every step that carries a prefill chunk. masked_fill_ + # performs the same zeroing in one launch with nothing to wait on. + initial_state.masked_fill_( + (~prefill_has_initial_state).view(-1, 1, 1, 1), 0 + ) ( core_attn_out_non_spec, last_recurrent_state, diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 021462f3e..c0f607ca6 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -437,6 +437,97 @@ class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): ) return hidden_states + # Proposals are re-scored by the target model before anything is + # emitted, so narrowing the vocabulary the drafter may argmax over can + # only cost acceptance -- never correctness. The full output embedding is + # vocab x hidden and is re-read on every one of the gamma draft steps, + # which makes it the single largest term in the drafter's cost; a + # byte-level BPE vocabulary orders its ids roughly by merge frequency, so + # a low-id prefix covers almost every token a draft would propose. The + # added/special tokens (end-of-turn, tool-call markers) sit at the very + # top of the vocabulary and have to come along, or a draft can never + # propose the token that ends a turn. + DRAFT_VOCAB_PREFIX = 96256 + DRAFT_VOCAB_TAIL = 4096 + + def _draft_head(self): + """(weight, ids) for the narrowed draft head, or (None, None).""" + cached = getattr(self, "_draft_head_cache", None) + if cached is not None: + return cached + weight = getattr(self.lm_head, "weight", None) + keep, tail = self.DRAFT_VOCAB_PREFIX, self.DRAFT_VOCAB_TAIL + if ( + weight is None + or weight.dim() != 2 + or keep + tail >= weight.shape[0] + or weight.dtype not in (torch.bfloat16, torch.float16, torch.float32) + ): + self._draft_head_cache = (None, None) + return self._draft_head_cache + vocab = weight.shape[0] + ids = torch.cat( + [ + torch.arange(keep, device=weight.device), + torch.arange(vocab - tail, vocab, device=weight.device), + ] + ) + narrowed = weight.index_select(0, ids).contiguous() + # e4m3 halves the bytes read per draft step again. A draft is only a + # proposal -- the target model re-scores it -- so the worst an + # imprecise argmax can do is lose a little acceptance near a tie. + scale = None + try: + amax = narrowed.abs().amax().clamp(min=1e-6).float() + scale = (amax / 448.0).to(torch.float32) + narrowed_fp8 = (narrowed.float() / scale).clamp(-448.0, 448.0).to( + torch.float8_e4m3fn + ) + except Exception: + narrowed_fp8, scale = None, None + logger.info( + "MTP draft head narrowed to %d of %d rows (%.0f%% of the bytes " + "read per draft step).", + ids.numel(), + vocab, + 100.0 * ids.numel() / vocab, + ) + self._draft_head_fp8 = narrowed_fp8 + self._draft_head_scale = scale + self._draft_head_cache = (narrowed, ids) + return self._draft_head_cache + + def draft_argmax(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + """Greedy draft token ids from the narrowed head, or None.""" + weight, ids = self._draft_head() + if weight is None: + return None + w8 = getattr(self, "_draft_head_fp8", None) + w8_scale = getattr(self, "_draft_head_scale", None) + if w8 is not None and w8_scale is not None: + try: + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + x_amax = x.abs().amax().clamp(min=1e-6).float() + x_scale = (x_amax / 448.0).to(torch.float32) + x8 = (x.float() / x_scale).clamp(-448.0, 448.0).to( + torch.float8_e4m3fn + ) + logits = torch._scaled_mm( + x8, + w8.t(), + scale_a=x_scale, + scale_b=w8_scale, + out_dtype=torch.bfloat16, + ) + return ids[logits.argmax(dim=-1)] + except Exception: + logger.warning_once( + "MTP draft head: fp8 path unavailable, using bf16." + ) + self._draft_head_fp8 = None + logits = torch.matmul(hidden_states.to(weight.dtype), weight.t()) + return ids[logits.argmax(dim=-1)] + def compute_logits( self, hidden_states: torch.Tensor, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f46cbd24..4c4c41002 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -412,6 +412,14 @@ class SpecDecodeBaseProposer: """Greedy-sample draft tokens from hidden states.""" if self.use_local_argmax_reduction: return self.model.get_top_tokens(hidden_states) + # A draft proposal is re-scored by the target before anything is + # emitted, so a model may offer a cheaper head to argmax over. Full + # logits are still what compute_logits returns for every other caller. + draft_argmax = getattr(self.model, "draft_argmax", None) + if draft_argmax is not None: + tokens = draft_argmax(hidden_states) + if tokens is not None: + return tokens return self.model.compute_logits(hidden_states).argmax(dim=-1) def _sample_from_logits(