diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 921f31466..9dc82a0c3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -530,7 +530,7 @@ class EngineArgs: max_logprobs: int = ModelConfig.max_logprobs logprobs_mode: LogprobsMode = ModelConfig.logprobs_mode use_fp64_gumbel: bool = ModelConfig.use_fp64_gumbel - disable_log_stats: bool = False + disable_log_stats: bool = True aggregate_engine_logging: bool = False revision: str | None = ModelConfig.revision code_revision: str | None = ModelConfig.code_revision @@ -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,27 @@ 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 + # Mixed chunked-prefill steps carry prompt tokens, so their size + # is bounded by the token budget, not by max_num_seqs x the + # speculative width the capture default is derived from. Steps + # larger than the largest captured graph run with per-layer + # Python dispatch. Widen capture toward the token budget so those + # steps stay on graphs; an explicit setting from the operator + # still wins because this only fills the unset default. + if self.compilation_config.max_cudagraph_capture_size is None: + self.compilation_config.max_cudagraph_capture_size = 1024 + assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) @@ -2310,6 +2343,18 @@ class EngineArgs: ), ) + if ( + self.gdn_prefill_backend == "triton" + and current_platform.is_cuda() + and current_platform.is_device_capability(90) + ): + # On Hopper the FlashInfer GDN prefill kernel is supported with no + # further conditions (see _resolve_gdn_prefill_backend), and it is + # what "auto" resolves to. The GDN chunk prefill path is the largest + # single cost in a prefill step for this model, so keeping the + # slower Triton kernel here is a pure loss on this hardware. + self.gdn_prefill_backend = "flashinfer" + if self.gdn_prefill_backend is not None: self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend @@ -2692,3 +2737,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 = 5 + + +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..d632d2953 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,45 @@ 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: + return [(text, list(token_ids))] + if not text: + # A multi-token delta can still decode to nothing: held-back bytes of an + # incomplete UTF-8 sequence, or a token counted in usage but never + # detokenized. Collapsing it into one chunk would put the stream below + # one chunk per token, so emit one empty piece per token instead. + return [("", [tid]) for tid in 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/fla/ops/chunk.py b/vllm/model_executor/layers/fla/ops/chunk.py index caf8b0c97..2bf86a8b4 100644 --- a/vllm/model_executor/layers/fla/ops/chunk.py +++ b/vllm/model_executor/layers/fla/ops/chunk.py @@ -86,6 +86,66 @@ def chunk_gated_delta_rule_fwd( return g, o, A, final_state, w, h, v_new +def _chunk_gated_delta_rule_inference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + core_attn_out: torch.Tensor | None = None, +): + """ChunkGatedDeltaRuleFunction.forward without the autograd machinery. + + Same operations in the same order, so results are bit-for-bit identical. + The contiguity coercion that `input_guard` performs is kept explicitly; + it is the identity on an already-contiguous tensor. + """ + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + g = g.contiguous() + beta = beta.contiguous() + if initial_state is not None: + initial_state = initial_state.contiguous() + if cu_seqlens is not None: + cu_seqlens = cu_seqlens.contiguous() + if chunk_indices is not None: + chunk_indices = chunk_indices.contiguous() + if chunk_offsets is not None: + chunk_offsets = chunk_offsets.contiguous() + if core_attn_out is not None: + core_attn_out = core_attn_out.contiguous() + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + + g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + core_attn_out=core_attn_out, + ) + if core_attn_out is not None: + assert q.dtype == o.dtype, "Incompatible dtype for inplace computation" + return o.to(q.dtype), final_state + + class ChunkGatedDeltaRuleFunction(torch.autograd.Function): @staticmethod @input_guard @@ -227,19 +287,36 @@ def chunk_gated_delta_rule( ) if scale is None: scale = k.shape[-1] ** -0.5 - o, final_state = ChunkGatedDeltaRuleFunction.apply( - q, - k, - v, - g, - beta, - scale, - initial_state, - output_final_state, - cu_seqlens, - chunk_indices, - chunk_offsets, - use_qk_l2norm_in_kernel, - core_attn_out, - ) + if torch.is_grad_enabled(): + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + chunk_indices, + chunk_offsets, + use_qk_l2norm_in_kernel, + core_attn_out, + ) + else: + o, final_state = _chunk_gated_delta_rule_inference( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + chunk_indices, + chunk_offsets, + use_qk_l2norm_in_kernel, + core_attn_out, + ) return o, final_state diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c5d..18d4acdeb 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, @@ -1507,11 +1511,11 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # when decodes are peeled off, else the full non-spec batch), so they # don't need to be re-derived per layer. prefill_state_indices = attn_metadata.prefill_state_indices - prefill_has_initial_state = attn_metadata.prefill_has_initial_state + prefill_no_initial_state = attn_metadata.prefill_no_initial_state assert prefill_state_indices is not None - assert prefill_has_initial_state is not None + assert prefill_no_initial_state is not None initial_state = ssm_state[prefill_state_indices] - initial_state[~prefill_has_initial_state, ...] = 0 + initial_state.masked_fill_(prefill_no_initial_state, 0) ( core_attn_out_non_spec, last_recurrent_state, @@ -1562,14 +1566,18 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # 3. Merge core attention output if spec_sequence_masks is not None and core_attn_out_non_spec is not None: - merged_out = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_non_spec.dtype, - device=core_attn_out_non_spec.device, - ) - merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec) - merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) + # spec_token_indx and non_spec_token_indx are the two halves of + # one argsort permutation, so together they cover every scheduled + # token exactly once. Scattering straight into the caller's buffer + # drops a full-size allocation and a device-to-device copy per + # layer-call. core_attn_out is zero-allocated on every path that + # reaches here, so on a token-padded step the rows the permutation + # does not reach stay zero rather than carrying the uninitialised + # values the temporary used to propagate. index_copy_ does not cast. + assert core_attn_out.dtype == core_attn_out_non_spec.dtype + out_view = core_attn_out[:num_actual_tokens].unsqueeze(0) + out_view.index_copy_(1, spec_token_indx, core_attn_out_spec) + out_view.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) elif spec_sequence_masks is not None: core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) else: diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 340a30403..6f279ca4a 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -72,6 +72,10 @@ class GDNAttentionMetadata: prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + # ~prefill_has_initial_state, pre-shaped [P, 1, 1, 1] so it broadcasts + # against the gathered [P, HV, V, K] prefill state. Computed once per + # step here instead of once per GDN layer inside _forward_core. + prefill_no_initial_state: torch.Tensor | None = None # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -177,7 +181,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 +390,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) if num_prefills > 0: + context_lens_tensor = m.compute_num_computed_tokens() has_initial_state = context_lens_tensor > 0 if spec_sequence_masks_cpu is not None: has_initial_state = has_initial_state[~spec_sequence_masks_cpu] @@ -482,6 +486,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] non_spec_query_start_loc = self.non_spec_query_start_loc[: batch_size + 1] non_spec_query_start_loc[num_decodes + 1 :].fill_(non_spec_num_query_tokens) + prefill_no_initial_state = ( + + None + + if prefill_has_initial_state is None + + else (~prefill_has_initial_state).view(-1, 1, 1, 1) + + ) + attn_metadata = GDNAttentionMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -495,6 +509,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] chunk_offsets=chunk_offsets, prefill_query_start_loc=prefill_query_start_loc, prefill_state_indices=prefill_state_indices, + prefill_no_initial_state=prefill_no_initial_state, prefill_has_initial_state=prefill_has_initial_state, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9f46cbd24..26530504e 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 @@ -58,6 +60,12 @@ from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) +# Draft-only vocab shortlist (see apply_draft_head_shortlist.py): the merge-ordered +# head of the vocabulary plus the tail rows that hold the chat-template specials. +_DRAFT_HEAD_ROWS = 98304 +_DRAFT_HEAD_TAIL = 2048 + + class SpecDecodeBaseProposer: def __init__( self, @@ -117,6 +125,11 @@ 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._draft_head8_ids: torch.Tensor | None = None self.use_local_argmax_reduction: bool = ( self.speculative_config.use_local_argmax_reduction ) @@ -408,8 +421,122 @@ 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 + self._draft_head8_ids = 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 + # Draft-only vocab shortlist: the head of the merge-ordered vocab + # plus the tail that holds the chat-template specials. Both bounds + # are multiples of 16 (cutlass output-stride requirement). + keep = None + if num_rows > _DRAFT_HEAD_ROWS + _DRAFT_HEAD_TAIL: + keep = torch.cat( + ( + torch.arange(0, _DRAFT_HEAD_ROWS, dtype=torch.int64), + torch.arange( + num_rows - _DRAFT_HEAD_TAIL, num_rows, dtype=torch.int64 + ), + ) + ).to(weight.device) + n_keep = num_rows if keep is None else int(keep.numel()) + quantized = torch.empty( + (n_keep, hidden_size), + dtype=torch.float8_e4m3fn, + device=weight.device, + ) + scales = torch.empty( + (n_keep, 1), dtype=torch.float32, device=weight.device + ) + # Chunked so the fp32 staging copy stays a few hundred MB. + for start in range(0, n_keep, 8192): + if keep is None: + src = weight[start : start + 8192] + else: + src = weight.index_select(0, keep[start : start + 8192]) + block = src.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, src + # 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, n_keep) + self._draft_head8_ids = keep + + # 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 of %d rows).", + n_keep, + hidden_size, + num_rows, + ) + except Exception: # noqa: BLE001 - degrade to bf16, never fail startup + self._draft_head8 = None + self._draft_head8_scale = None + self._draft_head8_ids = 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 + ) + local = ops.cutlass_scaled_mm( + quantized, + self._draft_head8, + scale, + self._draft_head8_scale, + out_dtype=hidden_states.dtype, + ).argmax(dim=-1) + if self._draft_head8_ids is None: + return local + return torch.index_select(self._draft_head8_ids, 0, local) 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 +1433,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/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"