From 69799cf2b65a06767d2cfcc224aaa1f3a51aa74d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 11 Apr 2026 01:19:13 +0000 Subject: [PATCH 1/2] [GLM-5] Fix NextN MTP model-side parity --- vllm/model_executor/models/deepseek_mtp.py | 126 +++++++++++++++------ vllm/model_executor/models/deepseek_v2.py | 29 ++++- 2 files changed, 113 insertions(+), 42 deletions(-) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index c75ee1a1bbfe..ec6b6aa2a932 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import typing from collections.abc import Callable, Iterable +from copy import copy import torch import torch.nn as nn @@ -10,7 +11,7 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.logger import init_logger +from vllm.config.utils import replace from vllm.model_executor.layers.fused_moe import SharedFusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -32,9 +33,25 @@ DeepseekV2MoE, get_spec_layer_idx_from_weight_name, ) -from .utils import maybe_prefix +from .utils import get_draft_quant_config, maybe_prefix -logger = init_logger(__name__) + +def _get_nextn_vllm_config( + vllm_config: VllmConfig, +) -> tuple[VllmConfig, QuantizationConfig | None]: + quant_config = get_draft_quant_config(vllm_config) + # GLM-5 stores NextN/MTP weights in a separate `mtp.safetensors` file and + # those tensors are BF16 rather than NVFP4-packed. Keep the draft path + # unquantized so MoE/attention params match the checkpoint layout. + if quant_config is not None and quant_config.get_name() == "modelopt_fp4": + nextn_vllm_config = copy(vllm_config) + nextn_vllm_config.quant_config = None + return nextn_vllm_config, None + + if quant_config is vllm_config.quant_config: + return vllm_config, quant_config + + return replace(vllm_config, quant_config=quant_config), quant_config class SharedHead(nn.Module): @@ -63,7 +80,8 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: config = vllm_config.speculative_config.draft_model_config.hf_config self.config = config - quant_config = vllm_config.quant_config + mtp_vllm_config, quant_config = _get_nextn_vllm_config(vllm_config) + layer_idx = int(prefix.rsplit(".", 1)[-1]) self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -87,10 +105,12 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: config=config, prefix=prefix, quant_config=quant_config ) self.mtp_block = DeepseekV2DecoderLayer( - vllm_config, - prefix, - config=self.config, + mtp_vllm_config, + f"{prefix}.mtp_block", + config=config, topk_indices_buffer=topk_indices_buffer, + layer_idx_override=layer_idx, + is_nextn=True, ) def forward( @@ -99,23 +119,24 @@ def forward( positions: torch.Tensor, previous_hidden_states: torch.Tensor, inputs_embeds: torch.Tensor | None = None, - spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - # masking inputs at position 0, as not needed by MTP - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) - inputs_embeds = self.enorm(inputs_embeds) - previous_hidden_states = self.hnorm(previous_hidden_states) + normed_inputs_embeds = self.enorm(inputs_embeds) + normed_previous_hidden_states = self.hnorm(previous_hidden_states) - hidden_states = self.eh_proj( - torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + eh_proj_hidden_states = self.eh_proj( + torch.cat([normed_inputs_embeds, normed_previous_hidden_states], dim=-1) ) - hidden_states, residual = self.mtp_block( - positions=positions, hidden_states=hidden_states, residual=None + decoder_hidden_states, residual = self.mtp_block( + positions=positions, hidden_states=eh_proj_hidden_states, residual=None ) - hidden_states = residual + hidden_states - return hidden_states + hidden_states = ( + decoder_hidden_states + if residual is None + else decoder_hidden_states + residual + ) + return self.shared_head(hidden_states) class DeepSeekMultiTokenPredictor(nn.Module): @@ -124,17 +145,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.mtp_start_layer_idx = config.num_hidden_layers self.num_mtp_layers = config.num_nextn_predict_layers - # to map the exact layer index from weights + # Keep internal layer ids zero-based so the NextN decoder behaves like + # SGLang's single-layer draft path while weight loading still maps from + # checkpoint layers appended after the target stack. + self.mtp_internal_start_idx = 0 self.layers = torch.nn.ModuleDict( { str(idx): DeepSeekMultiTokenPredictorLayer( vllm_config, f"{prefix}.layers.{idx}" ) - for idx in range( - self.mtp_start_layer_idx, - self.mtp_start_layer_idx + self.num_mtp_layers, - ) + for idx in range(self.mtp_internal_start_idx, self.num_mtp_layers) } ) self.embed_tokens = VocabParallelEmbedding( @@ -158,12 +179,11 @@ def forward( if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) current_step_idx = spec_step_idx % self.num_mtp_layers - return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + return self.layers[str(current_step_idx)]( input_ids, positions, previous_hidden_states, inputs_embeds, - current_step_idx, ) def compute_logits( @@ -172,11 +192,23 @@ def compute_logits( spec_step_idx: int = 0, ) -> torch.Tensor: current_step_idx = spec_step_idx % self.num_mtp_layers - mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] - logits = self.logits_processor( - mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states) + mtp_layer = self.layers[str(current_step_idx)] + return self.logits_processor( + mtp_layer.shared_head.head, + hidden_states, + ) + + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(current_step_idx)] + return self.logits_processor.get_top_tokens( + mtp_layer.shared_head.head, + hidden_states, ) - return logits @support_torch_compile @@ -232,6 +264,13 @@ def compute_logits( ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model.get_top_tokens(hidden_states, spec_step_idx) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: rocm_aiter_moe_shared_expert_enabled = ( rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() @@ -258,6 +297,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + loaded_spec_layers: set[int] = set() for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue @@ -300,6 +340,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) + loaded_spec_layers.add(spec_layer) break else: # Special handling: when AITER fusion_shared_experts is enabled, @@ -380,6 +421,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return_success=True, ) if success: + loaded_spec_layers.add(spec_layer) if not is_fusion_moe_shared_experts_layer: name = name_mapped else: @@ -413,20 +455,16 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) + loaded_spec_layers.add(spec_layer) if not is_fusion_moe_shared_experts_layer: loaded_params.add(name) # Validate that weights were loaded for each expected MTP layer. - loaded_layers: set[int] = set() - for param_name in loaded_params: - spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) - if spec_layer is not None: - loaded_layers.add(spec_layer) for layer_idx in range( self.model.mtp_start_layer_idx, self.model.mtp_start_layer_idx + self.model.num_mtp_layers, ): - if layer_idx not in loaded_layers: + if layer_idx not in loaded_spec_layers: raise ValueError( f"MTP speculative decoding layer {layer_idx} weights " f"missing from checkpoint. The checkpoint may have " @@ -443,6 +481,13 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: Add .mtp_block for modules in transformer layer block for spec layer and rename shared layer weights to be top level. """ + internal_layer_idx = spec_layer - self.model.mtp_start_layer_idx + if internal_layer_idx < 0 or internal_layer_idx >= self.model.num_mtp_layers: + raise ValueError( + f"Unexpected MTP checkpoint layer {spec_layer}; expected " + f"{self.model.mtp_start_layer_idx}.." + f"{self.model.mtp_start_layer_idx + self.model.num_mtp_layers - 1}." + ) spec_layer_weight_names = [ "embed_tokens", "enorm", @@ -462,9 +507,18 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: if not spec_layer_weight: # treat rest weights as weights for transformer layer block name = name.replace( - f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." + f"model.layers.{spec_layer}.", + f"model.layers.{internal_layer_idx}.mtp_block.", ) elif shared_weight: # treat shared weights as top level weights name = name.replace(f"model.layers.{spec_layer}.", "model.") + else: + name = name.replace( + f"model.layers.{spec_layer}.", f"model.layers.{internal_layer_idx}." + ) return name + + +class DeepSeekNextNMTP(DeepSeekMTP): + pass diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 8277e99fdc37..f5407427c47d 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -993,6 +993,8 @@ def __init__( prefix: str, config: DeepseekV2Config | None = None, topk_indices_buffer: torch.Tensor | None = None, + layer_idx_override: int | None = None, + is_nextn: bool = False, ) -> None: super().__init__() @@ -1008,7 +1010,11 @@ def __init__( moe_layer_freq = getattr(config, "moe_layer_freq", 1) # DecoderLayers are created with `make_layers` which passes the prefix # with the layer's index. - layer_idx = int(prefix.split(sep=".")[-1]) + layer_idx = ( + layer_idx_override + if layer_idx_override is not None + else int(prefix.split(sep=".")[-1]) + ) self.layer_idx = layer_idx # verify MLA attention specific fields @@ -1045,11 +1051,7 @@ def __init__( topk_indices_buffer=topk_indices_buffer, ) - if ( - config.n_routed_experts is not None - and layer_idx >= config.first_k_dense_replace - and layer_idx % moe_layer_freq == 0 - ): + if _should_use_nextn_moe_layer(config, layer_idx, moe_layer_freq, is_nextn): self.mlp = DeepseekV2MoE( config=config, parallel_config=parallel_config, @@ -1607,3 +1609,18 @@ def get_spec_layer_idx_from_weight_name( if weight_name.startswith(f"model.layers.{layer_idx + i}."): return layer_idx + i return None + + +def _should_use_nextn_moe_layer( + config: DeepseekV2Config | DeepseekV3Config, + layer_idx: int, + moe_layer_freq: int, + is_nextn: bool, +) -> bool: + return config.n_routed_experts is not None and ( + is_nextn + or ( + layer_idx >= config.first_k_dense_replace + and layer_idx % moe_layer_freq == 0 + ) + ) From 546080ee7a750bfbb6d55c9ea5e4f812e9423cac Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 12 Apr 2026 02:18:13 +0000 Subject: [PATCH 2/2] [MTP][Runtime] Reuse draft attention metadata across MTP steps --- .../layers/attention/mla_attention.py | 231 ++++++++-- vllm/v1/attention/backends/mla/indexer.py | 241 ++++++++--- vllm/v1/spec_decode/eagle.py | 396 ++++++++++++------ 3 files changed, 639 insertions(+), 229 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b1dc1a860501..8f6a9eda9c0d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -227,7 +227,6 @@ maybe_transfer_kv_layer, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant from vllm.model_executor.layers.linear import ( ColumnParallelLinear, ) @@ -242,6 +241,7 @@ from vllm.utils.math_utils import cdiv, round_down from vllm.utils.torch_utils import ( direct_register_custom_op, + is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, ) from vllm.v1.attention.backend import ( @@ -343,7 +343,7 @@ def __init__( # Automatically convert fp8 kv-cache format to "fp8_ds_mla" if ( self.attn_backend.get_name() == "FLASHMLA_SPARSE" - and kv_cache_dtype.startswith("fp8") + and is_quantized_kv_cache(kv_cache_dtype) and kv_cache_dtype != "fp8_ds_mla" ): assert cache_config is not None @@ -357,7 +357,7 @@ def __init__( if ( self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE" - and kv_cache_dtype.startswith("fp8") + and is_quantized_kv_cache(kv_cache_dtype) ): logger.info_once( "Using standard fp8 KV cache format. To use DeepSeek's fp8_ds_mla " @@ -365,6 +365,10 @@ def __init__( ) # Initialize KV cache quantization attributes + # The C++ cache ops only accept "auto" or fp8 variants. + # Map non-quantized dtypes to "auto". + if kv_cache_dtype in ("bfloat16", "float16"): + kv_cache_dtype = "auto" self.kv_cache_dtype = kv_cache_dtype self.calculate_kv_scales = calculate_kv_scales _init_kv_cache_quant(self, quant_config, prefix) @@ -372,7 +376,7 @@ def __init__( if ( cache_config is not None and cache_config.enable_prefix_caching - and vllm_is_batch_invariant() + and envs.VLLM_BATCH_INVARIANT and ( self.attn_backend.get_name() == "TRITON_MLA" or self.attn_backend.get_name() == "FLASHINFER" @@ -416,12 +420,7 @@ def __init__( raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - self.kv_cache = [ - torch.tensor([]) - for _ in range( - get_current_vllm_config().parallel_config.pipeline_parallel_size - ) - ] + self.kv_cache = torch.tensor([]) self.use_sparse = use_sparse @@ -442,6 +441,7 @@ def __init__( # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported self.is_aiter_triton_fp4_bmm_enabled = ( rocm_aiter_ops.is_fp4bmm_enabled() + and hasattr(self.kv_b_proj, "weight") and self.kv_b_proj.weight.dtype == torch.bfloat16 ) @@ -453,6 +453,11 @@ def __init__( group_shape=GroupShape.PER_TENSOR, compile_native=True, ) + self._quant_fp8_op = QuantFP8( + static=True, + group_shape=GroupShape.PER_TENSOR, + compile_native=True, + ) @property def chunked_prefill_workspace_size(self) -> int: @@ -479,7 +484,7 @@ def forward( attn_metadata = forward_context.attn_metadata if isinstance(attn_metadata, dict): attn_metadata = attn_metadata[self.layer_name] - self_kv_cache = self.kv_cache[forward_context.virtual_engine] + self_kv_cache = self.kv_cache slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict), ( @@ -549,9 +554,19 @@ def forward_impl( ) -> torch.Tensor: assert output is not None, "Output tensor must be provided." - if output_scale is not None or output_block_scale is not None: - raise NotImplementedError( - "fused output quantization is not yet supported for MLA" + use_quant = output_scale is not None or output_block_scale is not None + if use_quant: + # The fusion pass has allocated output with quantized dtype + # (FP8 or uint8 for FP4). We can't write into it directly, + # so we swap in a temp buffer for computation, then quantize + # into the real output at the end. + # NOTE(carlyou): this is temporary until kernels support fp8 output + quant_output = output + output = torch.empty( + output.shape[0], + self.num_heads * self.v_head_dim, + dtype=q.dtype, + device=output.device, ) if attn_metadata is None: @@ -571,12 +586,14 @@ def forward_impl( # The zero fill is required when used with DP + EP # to ensure all ranks within a DP group compute the # same expert outputs. + if use_quant: + return quant_output.fill_(0) return output.fill_(0) if self.impl.dcp_world_size == -1: self.impl.dcp_world_size = get_dcp_group().world_size - fp8_attention = self.kv_cache_dtype.startswith("fp8") + fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype) num_actual_toks = attn_metadata.num_actual_tokens @@ -710,6 +727,21 @@ def forward_impl( # v_up projection self._v_up_proj(attn_out, out=mqa_output_slice) + + if use_quant: + # Quantize the BF16 computation result into the quantized output + actual = output[:num_actual_toks] + if output_block_scale is not None: + # NVFP4: two FP4 values packed into one uint8 + fp4_data, fp4_scales = ops.scaled_fp4_quant(actual, output_scale) + quant_output[:num_actual_toks].copy_(fp4_data) + output_block_scale.copy_(fp4_scales) + else: + # Static FP8 quantization + fp8_data, _ = self._quant_fp8_op(actual, output_scale) + quant_output[:num_actual_toks].copy_(fp8_data) + return quant_output + return output_padded def process_weights_after_loading(self, act_dtype: torch.dtype): @@ -934,12 +966,14 @@ def unified_mla_kv_cache_update( the data dependency between them to ensure torch.compile preserves ordering. """ forward_context = get_forward_context() - if forward_context.attn_metadata is None: - # Dummy/profile forwards should not update live KV cache pages. - return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype) - attn_layer = forward_context.no_compile_layers[layer_name] - kv_cache = attn_layer.kv_cache[forward_context.virtual_engine] + kv_cache = attn_layer.kv_cache + + # This needs to run even when we don't have metadata yet, so that the op + # is correctly captured. + if kv_cache.numel() == 0: + # Can't update an empty KV cache. + return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype) slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict), ( @@ -1064,6 +1098,10 @@ class QueryLenSupport(Enum): "MLA models using TRITON_MLA will require flash_attn. " "AITER_MLA backends use aiter kernels instead." ) + elif current_platform.is_xpu(): + from vllm._xpu_ops import xpu_ops as ops + + flash_attn_varlen_func = ops.flash_attn_varlen_func # type: ignore[no-redef] def dynamic_per_batched_tensor_quant( @@ -1141,14 +1179,16 @@ def get_kv_cache_shape( def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets - # us from `get_kv_cache_shape` to the actual memory layout we want. - # (num_blocks, num_layers, block_size, head_size) - return (1, 0, 2, 3) if include_num_layers_dimension else (0, 1, 2) + if include_num_layers_dimension: + # MLA kernels require contiguous per-layer KV cache views. + # Identity permutation keeps num_layers first in physical + # layout, signaling cross-layer allocation is unsupported. + return (0, 1, 2, 3) + return (0, 1, 2) @classmethod def get_supported_head_sizes(cls) -> list[int]: - return [576] + return [320, 576] @classmethod def is_mla(cls) -> bool: @@ -1178,6 +1218,7 @@ class ChunkedContextMetadata: padded_local_cu_seq_lens: torch.Tensor | None = None cu_seq_lens_lst: list[list[int]] | None = None chunk_size: int | None = None + prefill_tokens_with_context: int | None = None block_table: torch.Tensor query_start_loc: torch.Tensor @@ -1282,8 +1323,6 @@ def is_deepseek_r1_mla_compatible(vllm_config: VllmConfig) -> bool: @functools.cache def use_flashinfer_prefill() -> bool: - # For blackwell default to flashinfer prefill if it's available since - # it is faster than FA2. from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() @@ -1381,6 +1420,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # - VARLEN: Supports variable-length queries (spec decode with mixed lengths) # If set to UNIFORM or VARLEN, this will increase `reorder_batch_threshold` when # speculative decoding is enabled. + supports_update_for_drafting: ClassVar[bool] = True query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY # The threshold for reordering the batch into decode and prefill requests. @@ -1432,7 +1472,7 @@ def determine_prefill_query_data_type( is enabled, else model dtype. """ use_fp8 = ( - vllm_config.cache_config.cache_dtype.startswith("fp8") + is_quantized_kv_cache(vllm_config.cache_config.cache_dtype) and vllm_config.attention_config.use_prefill_query_quantization and backend_supports_prefill_query_quantization() ) @@ -1694,6 +1734,83 @@ def build_for_cudagraph_capture( return self.build(0, m) + def build_for_drafting( + self, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int, + ) -> M: + """Build attention metadata for draft model without DCP fields. + + Draft models (MTP layer) have their own local KV cache and don't + participate in decode context parallelism. Temporarily disable DCP + to prevent incompatible tensor shapes in draft attention kernels. + """ + if self.dcp_world_size > 1: + saved = self.dcp_world_size + self.dcp_world_size = 1 + try: + return self.build(0, common_attn_metadata, fast_build=True) + finally: + self.dcp_world_size = saved + return self.build(0, common_attn_metadata, fast_build=True) + + def update_for_drafting( + self, + attn_metadata: M, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int, + ) -> M: + num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens + max_seq_len = common_attn_metadata.max_seq_len + query_start_loc = common_attn_metadata.query_start_loc + seq_lens = common_attn_metadata.seq_lens + dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens + + num_decodes, num_prefills, num_decode_tokens, _ = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold, + require_uniform=(self.query_len_support != QueryLenSupport.VARLEN), + ) + ) + + if num_prefills > 0: + return self.build_for_drafting(common_attn_metadata, draft_index) + + dcp_tot_seq_lens_device = None + if self.dcp_world_size > 1: + dcp_tot_seq_lens_device = seq_lens[:num_decodes] + seq_lens = dcp_local_seq_lens + + num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size + max_seq_len = ( + (max_seq_len + num_partitions - 1) // num_partitions + ) * self.cp_kv_cache_interleave_size + + attn_metadata.num_reqs = num_reqs + attn_metadata.max_query_len = common_attn_metadata.max_query_len + attn_metadata.max_seq_len = max_seq_len + attn_metadata.num_actual_tokens = num_tokens + attn_metadata.query_start_loc = query_start_loc + attn_metadata.seq_lens = seq_lens + attn_metadata.slot_mapping = common_attn_metadata.slot_mapping + attn_metadata.num_decodes = num_decodes + attn_metadata.num_decode_tokens = num_decode_tokens + attn_metadata.num_prefills = num_prefills + attn_metadata.num_prefill_tokens = 0 + attn_metadata.prefill = None + + decode = attn_metadata.decode + if decode is None: + return self.build_for_drafting(common_attn_metadata, draft_index) + + decode.block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] + decode.seq_lens = seq_lens + decode.dcp_tot_seq_lens = dcp_tot_seq_lens_device + + return attn_metadata + def build( self, common_prefix_len: int, @@ -1742,6 +1859,9 @@ def build( prefill_query_start_loc = ( query_start_loc[reqs_start:] - query_start_loc[reqs_start] ) + prefill_query_start_loc_cpu = ( + query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start] + ) chunked_context_metadata = None if max_context_len_cpu > 0: @@ -1863,6 +1983,11 @@ def build( if self._use_cudnn_prefill else MLACommonPrefillMetadata.ChunkedContextMetadata ) + prefill_tokens_with_context = None + if num_prefills_with_context_cpu > 0: + prefill_tokens_with_context = prefill_query_start_loc_cpu[ + num_prefills_with_context_cpu + ].item() if self.dcp_world_size > 1: chunked_context_metadata = chunked_context_metadata_cls( cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True), @@ -1882,6 +2007,7 @@ def build( ), cu_seq_lens_lst=cu_seq_lens_cpu.tolist(), chunk_size=padded_local_max_context_chunk_across_ranks, + prefill_tokens_with_context=prefill_tokens_with_context, ) else: chunked_context_metadata = chunked_context_metadata_cls( @@ -1895,6 +2021,7 @@ def build( ), chunk_total_token=chunk_total_token, workspace=self.chunked_prefill_workspace, + prefill_tokens_with_context=prefill_tokens_with_context, ) if self._use_cudnn_prefill: @@ -2056,6 +2183,14 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): understand this class """ + def fused_output_quant_supported(self, quant_key): + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, + kNvfp4Dynamic, + ) + + return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic) + def __init__( self, num_heads: int, @@ -2154,13 +2289,16 @@ def __init__( # For MLA the v head dim is smaller than qk head dim so we pad out # v with 0s to match the qk head dim for attention backends that do - # not support different headdims - # We don't need to pad V if we are on a hopper system with FA3 + # not support different headdims. + # FA3 on Hopper (SM90) and FA4 natively handle diff headdims. device_capability = current_platform.get_device_capability() self._pad_v = self.vllm_flash_attn_version is None or not ( - self.vllm_flash_attn_version == 3 - and device_capability is not None - and device_capability[0] == 9 + ( + self.vllm_flash_attn_version == 3 + and device_capability is not None + and device_capability[0] == 9 + ) + or self.vllm_flash_attn_version == 4 ) self.dcp_world_size: int = -1 @@ -2184,7 +2322,7 @@ def _flash_attn_varlen_diff_headdims( # ROCm leverages the upstream flash_attn, which takes a parameter # called "return_attn_probs" instead of return_softmax_lse kwargs["return_attn_probs"] = return_softmax_lse - if vllm_is_batch_invariant(): + if envs.VLLM_BATCH_INVARIANT: kwargs["num_splits"] = 1 attn_out = self.flash_attn_varlen_func( @@ -2378,14 +2516,13 @@ def _run_prefill_context_chunk_trtllm_ragged( assert prefill.chunked_context.seq_lens[chunk_idx] is not None assert prefill.workspace_buffer is not None - out = torch.zeros( + out = torch.empty( q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=prefill.output_dtype, ) - prefill.workspace_buffer.fill_(0) attn_out, lse = trtllm_ragged_attention_deepseek( query=q, @@ -2491,10 +2628,18 @@ def _compute_prefill_context( kv_c_normed = workspace[:toks][..., : self.kv_lora_rank] # When FP8 weights are used without FP8 prefill, kv_b_proj expects # model dtype input and will quantize internally. + # For quantized layers (AWQ/GPTQ) that lack a .weight attribute, + # use params_dtype which is the expected input dtype. + _kv_b_proj_w_dtype = ( + self.kv_b_proj.weight.dtype + if hasattr(self.kv_b_proj, "weight") + else self.kv_b_proj.params_dtype + ) + # For NVFP4, weights are packed uint8 — keep input in model dtype + # since the NVFP4 linear layer quantizes internally. if ( - use_fp8_prefill - or self.kv_b_proj.weight.dtype != current_platform.fp8_dtype() - ): + use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype() + ) and _kv_b_proj_w_dtype != torch.uint8: kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype) k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1) @@ -2683,8 +2828,13 @@ def forward_mha( ) if has_context: + assert prefill_metadata.chunked_context is not None suffix_output, suffix_lse = output_prefill - if self.dcp_world_size > 1: + # Check metadata for DCP fields instead of impl.dcp_world_size + # because draft models (MTP) build metadata without DCP + use_dcp_prefill = (self.dcp_world_size > 1 + and prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None) + if use_dcp_prefill: context_output, context_lse = ( self._context_parallel_compute_prefill_context( q, @@ -2711,6 +2861,7 @@ def forward_mha( prefix_lse=context_lse, suffix_output=suffix_output, suffix_lse=suffix_lse, + prefill_tokens_with_context=prefill_metadata.chunked_context.prefill_tokens_with_context, ) else: output_prefill = output_prefill[..., : v.shape[-1]].flatten(start_dim=-2) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index e84312970989..831a60a5b69a 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1,16 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass -from typing import ClassVar import torch +import vllm.envs as envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( get_paged_mqa_logits_metadata, - is_deep_gemm_supported, + has_deep_gemm, ) from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import num_compute_units @@ -23,13 +23,62 @@ ) from vllm.v1.attention.backends.utils import ( split_decodes_and_prefills, - split_prefill_chunks, ) +from vllm.v1.kv_cache_interface import AttentionSpec from vllm.v1.worker.cp_utils import get_total_cp_world_size logger = init_logger(__name__) +def split_indexer_prefill_chunks( + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + workspace_size: int, + max_logits_bytes: int, + request_offset: int = 0, +) -> list[tuple[slice, slice]]: + """ + Split prefill requests into chunks for the sparse indexer, respecting: + - N constraint: total_seq_lens <= workspace_size (existing O(N) workspace) + - Logits constraint: M * N * 4 <= max_logits_bytes + + When a single request-level chunk still exceeds the logits budget, + sub-chunks on the query dimension (M) to bound peak memory. + + Returns list of (req_slice, query_slice) tuples. + """ + chunks: list[tuple[slice, slice]] = [] + n = len(seq_lens_cpu) + max_logits_elems = max_logits_bytes // 4 + end = 0 + + while end < n: + start, chunk_m, chunk_n = end, 0, 0 + + while end < n: + q, s = query_lens_cpu[end].item(), seq_lens_cpu[end].item() + new_m, new_n = chunk_m + q, chunk_n + s + if new_n <= workspace_size and new_m * new_n <= max_logits_elems: + chunk_m, chunk_n = new_m, new_n + end += 1 + else: + break + + # A single request can exceed the budget, requiring sub-chunking + # on the query dimension. + if end == start: + chunk_m, chunk_n = query_lens_cpu[end].item(), seq_lens_cpu[end].item() + end += 1 + + req_slice = slice(start + request_offset, end + request_offset) + max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m + for q_off in range(0, chunk_m, max_q): + sub_m = min(max_q, chunk_m - q_off) + chunks.append((req_slice, slice(q_off, q_off + sub_m))) + + return chunks + + class DeepseekV32IndexerBackend(AttentionBackend): @staticmethod def get_name() -> str: @@ -63,6 +112,9 @@ def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: if include_num_layers_dimension: + # DeepseekV32Indexer kernels do not support cross-layer + # KV cache layout. Identity permutation keeps num_layers + # first, signaling incompatibility. return (0, 1, 2, 3) return (0, 1, 2) @@ -78,6 +130,7 @@ class DeepseekV32IndexerPrefillChunkMetadata: token_start: int token_end: int num_reqs: int + skip_kv_gather: bool = False @dataclass @@ -202,9 +255,18 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig): class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH - + supports_update_for_drafting = True reorder_batch_threshold: int = 1 + natively_supported_next_n: list[int] = [1, 2] + # TODO (matt): integrate kernel with next_n = 4 support + + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: AttentionSpec, + ) -> AttentionCGSupport: + return AttentionCGSupport.UNIFORM_BATCH def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -216,7 +278,9 @@ def __init__(self, *args, **kwargs): if self.vllm_config.speculative_config else 0 ) + next_n = self.num_speculative_tokens + 1 self.reorder_batch_threshold += self.num_speculative_tokens + self.use_flattening = next_n not in self.natively_supported_next_n sm_count = num_compute_units(self.device.index) self.num_sms = sm_count @@ -226,10 +290,11 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) - - # Pre-allocated buffers for flattening (spec decode). + self.offsets_buffer = torch.arange( + next_n, device=self.device, dtype=torch.int32 + ) self.arange_buffer = torch.arange( - scheduler_config.max_num_seqs * (1 + self.num_speculative_tokens), + scheduler_config.max_num_seqs * next_n, dtype=torch.int32, device=self.device, ) @@ -257,43 +322,51 @@ def __init__(self, *args, **kwargs): ) def build_one_prefill_chunk( - self, reqs_start, reqs_end, query_start_loc_cpu, seq_lens_cpu, block_table - ): + self, + req_slice: slice, + query_slice: slice, + query_start_loc_cpu, + seq_lens_cpu, + block_table, + skip_kv_gather: bool = False, + ) -> DeepseekV32IndexerPrefillChunkMetadata: prefill_query_start_loc = ( - query_start_loc_cpu[reqs_start : reqs_end + 1] - - query_start_loc_cpu[reqs_start] + query_start_loc_cpu[req_slice.start : req_slice.stop + 1] + - query_start_loc_cpu[req_slice.start] ) cu_seqlen_ks, cu_seqlen_ke = kv_spans_from_batches( - prefill_query_start_loc, seq_lens_cpu[reqs_start:reqs_end], self.device + prefill_query_start_loc, seq_lens_cpu[req_slice], self.device + ) + token_start = query_start_loc_cpu[req_slice.start].item() + total_seq_lens = seq_lens_cpu[req_slice].sum() + num_reqs = req_slice.stop - req_slice.start + seq_idx = torch.arange(0, num_reqs, dtype=torch.int32) + token_to_seq = torch.repeat_interleave(seq_idx, seq_lens_cpu[req_slice]).to( + self.device ) - token_start = query_start_loc_cpu[reqs_start].item() - token_end = query_start_loc_cpu[reqs_end].item() - total_seq_lens = seq_lens_cpu[reqs_start:reqs_end].sum() - seq_idx = torch.arange(0, reqs_end - reqs_start, dtype=torch.int32) - token_to_seq = torch.repeat_interleave( - seq_idx, seq_lens_cpu[reqs_start:reqs_end] - ).to(self.device) assert total_seq_lens <= self.max_prefill_buffer_size cu_seq_lens = ( torch.cat( [ torch.zeros(1, dtype=torch.int32), - seq_lens_cpu[reqs_start:reqs_end].cumsum(dim=0), + seq_lens_cpu[req_slice].cumsum(dim=0), ] ) .to(torch.int32) .to(self.device) ) + return DeepseekV32IndexerPrefillChunkMetadata( - cu_seqlen_ks=cu_seqlen_ks, - cu_seqlen_ke=cu_seqlen_ke, + cu_seqlen_ks=cu_seqlen_ks[query_slice], + cu_seqlen_ke=cu_seqlen_ke[query_slice], cu_seq_lens=cu_seq_lens, token_to_seq=token_to_seq, total_seq_lens=total_seq_lens, - block_table=block_table[reqs_start:reqs_end], - token_start=token_start, - token_end=token_end, - num_reqs=reqs_end - reqs_start, + block_table=block_table[req_slice], + token_start=token_start + query_slice.start, + token_end=token_start + query_slice.stop, + num_reqs=num_reqs, + skip_kv_gather=skip_kv_gather, ) def build( @@ -308,7 +381,9 @@ def build( query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( split_decodes_and_prefills( - common_attn_metadata, decode_threshold=self.reorder_batch_threshold + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold, + require_uniform=not self.use_flattening, ) ) @@ -317,20 +392,27 @@ def build( prefill_metadata = None if num_prefills > 0: - chunk_seq_ids = split_prefill_chunks( + prefill_query_lens_cpu = torch.diff( + query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1] + ) + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + chunk_specs = split_indexer_prefill_chunks( common_attn_metadata.seq_lens_cpu[num_decodes:], + prefill_query_lens_cpu, self.max_prefill_buffer_size, + max_logits_bytes, request_offset=num_decodes, ) chunks = [ self.build_one_prefill_chunk( - reqs_start, - reqs_end, + req_slice, + query_slice, query_start_loc_cpu, common_attn_metadata.seq_lens_cpu, common_attn_metadata.block_table_tensor, + skip_kv_gather=query_slice.start > 0, ) - for reqs_start, reqs_end in chunk_seq_ids + for req_slice, query_slice in chunk_specs ] prefill_metadata = DeepseekV32IndexerPrefillMetadata( chunks=chunks, @@ -350,18 +432,22 @@ def build( seq_lens = common_attn_metadata.seq_lens[:num_decodes] block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] - # Padded CUDA graph requests have block_table entries of -1. - # Clamp to 0 to prevent OOB access in the DeepGEMM kernel. - # This is safe because padded requests have seq_lens=0, so the - # kernel produces no meaningful output for those rows. - block_table.clamp_(min=0) - max_decode_len = int(decode_lens_cpu.max().item()) - if max_decode_len > 1: + next_n = 1 + self.num_speculative_tokens + use_native = not self.use_flattening and max_decode_len == next_n + + if use_native and next_n > 1: + offsets = self.offsets_buffer + batch_size = num_decodes + elif max_decode_len > 1: # Flatten multi-token decode requests into single-token # batch entries, expanding seq_lens and block tables so # the kernel always sees next_n=1. + # Also handles the edge case where use_flattening=False + # but max_decode_len != next_n (e.g. a batch containing some + # short prefills (q_len < next_n) and no true decodes). + # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is # padding) and decode_lens [3, 1, 4, 0] in the below example comments. # The context lengths are therefore @@ -372,12 +458,14 @@ def build( # [7, 6, 8, 0] -> [7, 7, 7, 6, 8, 8, 8, 8] expanded_base = torch.repeat_interleave( - seq_lens - decode_lens, decode_lens + seq_lens - decode_lens, decode_lens, output_size=actual_expanded ) # [0, 3, 4, 8] -> [0, 0, 0, 3, 4, 4, 4, 4] expanded_starts = torch.repeat_interleave( - common_attn_metadata.query_start_loc[:num_decodes], decode_lens + common_attn_metadata.query_start_loc[:num_decodes], + decode_lens, + output_size=actual_expanded, ) # [0, 1, 2, 0, 0, 1, 2, 3] @@ -395,7 +483,9 @@ def build( # Give each of the flattened entries the same block table row as the # original request. self.expanded_block_table_buffer[:actual_expanded] = ( - torch.repeat_interleave(block_table, decode_lens, dim=0) + torch.repeat_interleave( + block_table, decode_lens, dim=0, output_size=actual_expanded + ) ) if actual_expanded < num_decode_tokens: self.expanded_block_table_buffer[ @@ -409,17 +499,11 @@ def build( offsets = None batch_size = num_decode_tokens else: - next_n = 1 + self.num_speculative_tokens - if next_n > 1: - offsets = torch.arange( - next_n, device=self.device, dtype=torch.int32 - ) - else: - offsets = None + offsets = None batch_size = num_decodes # DeepGEMM is required for the paged MQA logits on CUDA devices - if current_platform.is_cuda() and is_deep_gemm_supported(): + if current_platform.is_cuda() and has_deep_gemm(): self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata( seq_lens, self.kv_cache_spec.block_size, @@ -463,3 +547,60 @@ def build( # if get_tensor_model_parallel_rank() == 0: # logger.info(f"attn_metadata: {attn_metadata}") return attn_metadata + + def update_for_drafting( + self, + attn_metadata: DeepseekV32IndexerMetadata, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int, + ) -> DeepseekV32IndexerMetadata: + num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold, + ) + ) + + if num_prefills > 0 or num_decode_tokens != num_decodes: + return self.build_for_drafting(common_attn_metadata, draft_index) + + max_seq_len = common_attn_metadata.max_seq_len + seq_lens = common_attn_metadata.seq_lens[:num_decodes] + + attn_metadata.seq_lens = common_attn_metadata.seq_lens + attn_metadata.num_reqs = num_reqs + attn_metadata.max_query_len = common_attn_metadata.max_query_len + attn_metadata.max_seq_len = max_seq_len + attn_metadata.num_actual_tokens = num_tokens + attn_metadata.query_start_loc = common_attn_metadata.query_start_loc + attn_metadata.slot_mapping = common_attn_metadata.slot_mapping + attn_metadata.num_decodes = num_decodes + attn_metadata.num_decode_tokens = num_decode_tokens + attn_metadata.num_prefills = num_prefills + attn_metadata.num_prefill_tokens = num_prefill_tokens + attn_metadata.prefill = None + + decode = attn_metadata.decode + if decode is None: + return self.build_for_drafting(common_attn_metadata, draft_index) + + decode.block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] + decode.seq_lens = seq_lens + self.decode_lens_buffer[:num_decode_tokens] = 1 + decode.decode_lens = self.decode_lens_buffer[:num_decode_tokens] + decode.requires_padding = False + decode.use_large_context_topk = ( + num_decodes <= 128 and max_seq_len > 8192 + ) + decode.offsets = None + + if current_platform.is_cuda() and has_deep_gemm(): + decode.schedule_metadata = get_paged_mqa_logits_metadata( + seq_lens, + self.kv_cache_spec.block_size, + self.num_sms, + ) + + return attn_metadata diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 89c9c80ce0aa..68ab97e4503b 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import ast -from dataclasses import replace from importlib.util import find_spec -from typing import cast +from typing import Any, cast import numpy as np import torch @@ -13,6 +12,7 @@ CUDAGraphMode, VllmConfig, get_layers_from_vllm_config, + replace, ) from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context @@ -20,8 +20,10 @@ 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 +from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausalLM from vllm.model_executor.models.interfaces import SupportsMultiModal from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.triton_utils import triton @@ -44,6 +46,7 @@ copy_and_expand_eagle_inputs_kernel, eagle_prepare_inputs_padded_kernel, eagle_prepare_next_token_padded_kernel, + eagle_step_update_slot_mapping_and_metadata, extend_all_queries_by_N, ) from vllm.v1.utils import CpuGpuBuffer @@ -54,6 +57,10 @@ logger = init_logger(__name__) +def _reuse_mtp_draft_attn_metadata() -> bool: + return True + + class SpecDecodeBaseProposer: def __init__( self, @@ -69,7 +76,6 @@ def __init__( self.method = self.speculative_config.method self.pass_hidden_states_to_model = pass_hidden_states_to_model - self.runner = runner self.device = device self.dtype = vllm_config.model_config.dtype self.max_model_len = vllm_config.model_config.max_model_len @@ -82,13 +88,15 @@ def __init__( self.hidden_size = self.draft_model_config.get_hidden_size() self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size() - # Unifying eagle, draft model, and parallel drafting support + # Unifying eagle, draft model, and parallel drafting support. + # DFlash always uses parallel drafting (all tokens in one pass), + # but has an additional slot for the next_token_id (does not shift like EAGLE) self.parallel_drafting: bool = self.speculative_config.parallel_drafting self.extra_slots_per_request = ( 1 if not self.parallel_drafting else self.num_speculative_tokens ) self.net_num_new_slots_per_request = self.extra_slots_per_request - ( - 1 if self.pass_hidden_states_to_model else 0 + 1 if (self.pass_hidden_states_to_model and self.method != "dflash") else 0 ) self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0 @@ -100,10 +108,14 @@ def __init__( self.speculative_config.use_local_argmax_reduction ) - max_batch_size = vllm_config.scheduler_config.max_num_seqs + self.max_batch_size = vllm_config.scheduler_config.max_num_seqs self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens self.token_arange_np = np.arange(self.max_num_tokens) + # Can be specialized by methods like DFlash to reduce the limit + self.max_query_tokens = self.max_num_tokens + self.max_positions = self.max_num_tokens + # Multi-modal data support self.mm_registry = MULTIMODAL_REGISTRY self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( @@ -145,18 +157,20 @@ def __init__( # 1D-RoPE. # See page 5 of https://arxiv.org/abs/2409.12191 self.mrope_positions = torch.zeros( - (3, self.max_num_tokens + 1), dtype=torch.int64, device=device + (3, self.max_positions + 1), dtype=torch.int64, device=device ) elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0: self.xdrope_positions = torch.zeros( - (self.uses_xdrope_dim, self.max_num_tokens + 1), + (self.uses_xdrope_dim, self.max_positions + 1), dtype=torch.int64, device=device, ) else: # RoPE need (max_num_tokens,) self.positions = torch.zeros( - self.max_num_tokens, dtype=torch.int64, device=device + self.max_positions, + dtype=torch.int64, + device=device, ) self.hidden_states = torch.zeros( (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device @@ -167,7 +181,7 @@ def __init__( # We need +1 here because the arange is used to set query_start_loc, # which has one more element than batch_size. - max_num_slots_for_arange = max(max_batch_size + 1, self.max_num_tokens) + max_num_slots_for_arange = max(self.max_batch_size + 1, self.max_num_tokens) self.arange = torch.arange( max_num_slots_for_arange, device=device, dtype=torch.int32 ) @@ -199,7 +213,7 @@ def __init__( ) self.backup_next_token_ids = CpuGpuBuffer( - max_batch_size, + self.max_batch_size, dtype=torch.int32, pin_memory=is_pin_memory_available(), device=device, @@ -207,17 +221,27 @@ def __init__( ) self._slot_mapping_buffer = torch.zeros( - self.max_num_tokens, dtype=torch.int64, device=device + self.max_positions, + dtype=torch.int64, + device=device, ) # Determine allowed attention backends once during initialization. self.allowed_attn_types: tuple | None = None if current_platform.is_rocm(): + from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + ) + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseMetadata, + ) from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata rocm_types = [ TritonAttentionMetadata, RocmAttentionMetadata, + ROCMAiterMLASparseMetadata, + DeepseekV32IndexerMetadata, ] # ROCM_AITER_FA is an optional backend # We check is_enabled() here to avoid importing the backend module during @@ -270,7 +294,7 @@ def __init__( # Precompute draft position offsets in flattened tree. self.tree_draft_pos_offsets = torch.arange( 1, len(self.tree_choices) + 1, device=device, dtype=torch.int32 - ).repeat(max_batch_size, 1) + ).repeat(self.max_batch_size, 1) def _raise_if_padded_drafter_batch_disabled(self): if self.speculative_config.disable_padded_drafter_batch: @@ -300,14 +324,19 @@ def _init_parallel_drafting_params(self): # for those masked slots. model_hf_config = self.draft_model_config.hf_config - if hasattr(model_hf_config, "pard_token"): + # DFlash stores mask_token_id in dflash_config + dflash_config = getattr(model_hf_config, "dflash_config", None) + if dflash_config and "mask_token_id" in dflash_config: + self.parallel_drafting_token_id = dflash_config["mask_token_id"] + elif hasattr(model_hf_config, "pard_token"): self.parallel_drafting_token_id = model_hf_config.pard_token elif hasattr(model_hf_config, "ptd_token_id"): self.parallel_drafting_token_id = model_hf_config.ptd_token_id else: raise ValueError( "For parallel drafting, the draft model config must have " - "`pard_token` or `ptd_token_id` specified in its config.json." + "`pard_token`, `ptd_token_id`, or " + "`dflash_config.mask_token_id` specified in its config.json." ) if self.pass_hidden_states_to_model: @@ -376,6 +405,48 @@ def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.model.get_top_tokens(hidden_states) return self.model.compute_logits(hidden_states).argmax(dim=-1) + def _should_reuse_draft_attn_metadata( + self, per_layer_attn_metadata: dict[str, object] + ) -> bool: + if not _reuse_mtp_draft_attn_metadata(): + return False + if self.method != "mtp": + return False + if self.num_speculative_tokens <= 1 or self.parallel_drafting: + return False + if not per_layer_attn_metadata: + return False + return all( + getattr(attn_group.get_metadata_builder(), + "supports_update_for_drafting", False) + for attn_group in self.draft_attn_groups + ) + + def _update_reused_draft_attn_metadata( + self, + per_layer_attn_metadata: dict[str, object], + common_attn_metadata: CommonAttentionMetadata, + draft_index: int, + ) -> None: + updated_ids: set[int] = set() + for attn_group in self.draft_attn_groups: + if not attn_group.layer_names: + continue + layer_name = attn_group.layer_names[0] + attn_metadata = per_layer_attn_metadata[layer_name] + metadata_id = id(attn_metadata) + if metadata_id in updated_ids: + continue + builder = attn_group.get_metadata_builder() + attn_metadata = builder.update_for_drafting( + attn_metadata, + common_attn_metadata, + draft_index, + ) + for group_layer_name in attn_group.layer_names: + per_layer_attn_metadata[group_layer_name] = attn_metadata + updated_ids.add(metadata_id) + def propose( self, # [num_tokens] @@ -397,8 +468,15 @@ def propose( ) -> torch.Tensor: batch_size = common_attn_metadata.batch_size() - if self.method == "eagle3": - assert isinstance(self.model, Eagle3LlamaForCausalLM) + if self.method in ("eagle3", "dflash"): + assert isinstance( + self.model, + ( + Eagle3LlamaForCausalLM, + Eagle3DeepseekV2ForCausalLM, + DFlashQwen3ForCausalLM, + ), + ) target_hidden_states = self.model.combine_hidden_states( target_hidden_states ) @@ -416,42 +494,21 @@ def propose( ) ) - assert self.runner is not None - - per_layer_attn_metadata: dict[str, object] = {} - for attn_group in self.draft_attn_groups: - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( - common_attn_metadata=common_attn_metadata, draft_index=0 - ) - for layer_name in attn_group.layer_names: - per_layer_attn_metadata[layer_name] = attn_metadata + per_group_attn_metadata, per_layer_attn_metadata = ( + self.build_per_group_and_layer_attn_metadata(common_attn_metadata) + ) + can_reuse_draft_attn_metadata = self._should_reuse_draft_attn_metadata( + per_layer_attn_metadata + ) + reuse_draft_attn_metadata_ready = False cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = ( self._determine_batch_execution_and_padding(num_tokens) ) - if self.supports_mm_inputs: - mm_embeds, is_mm_embed = mm_embed_inputs or (None, None) - - self.inputs_embeds[:num_tokens] = self.model.embed_input_ids( - self.input_ids[:num_tokens], - multimodal_embeddings=mm_embeds, - is_multimodal=is_mm_embed, - ) - - input_ids = None - inputs_embeds = self.inputs_embeds[:num_input_tokens] - else: - input_ids = self.input_ids[:num_input_tokens] - inputs_embeds = None - - model_kwargs = { - "input_ids": input_ids, - "positions": self._get_positions(num_input_tokens), - "inputs_embeds": inputs_embeds, - } - if self.pass_hidden_states_to_model: - model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens] + model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass( + num_tokens, num_input_tokens, mm_embed_inputs + ) with set_forward_context( per_layer_attn_metadata, @@ -460,7 +517,7 @@ def propose( num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, slot_mapping=self._get_slot_mapping( - num_input_tokens, common_attn_metadata.slot_mapping + slot_mapping_size, common_attn_metadata.slot_mapping ), ): ret_hidden_states = self.model(**model_kwargs) @@ -483,7 +540,7 @@ def propose( positions = self.positions[token_indices_to_sample] hidden_states = hidden_states[token_indices_to_sample] - if isinstance(attn_metadata, TreeAttentionMetadata): + if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata): # Draft using tree attention - requires full logits for top-k logits = self.model.compute_logits(sample_hidden_states) draft_token_ids_list = self.propose_tree( @@ -499,15 +556,15 @@ def propose( draft_token_ids = self._greedy_sample(sample_hidden_states) - if self.allowed_attn_types is not None and not isinstance( - attn_metadata, self.allowed_attn_types - ): - raise ValueError( - f"Unsupported attention metadata type for speculative " - "decoding with num_speculative_tokens > 1: " - f"{type(attn_metadata)}. Supported types are: " - f"{self.allowed_attn_types}" - ) + if self.allowed_attn_types is not None: + for group_md in per_group_attn_metadata: + if not isinstance(group_md, self.allowed_attn_types): + raise ValueError( + f"Unsupported attention metadata type for speculative " + "decoding with num_speculative_tokens > 1: " + f"{type(group_md)}. Supported types are: " + f"{self.allowed_attn_types}" + ) # Generate the remaining draft tokens. draft_token_ids_list = [draft_token_ids] @@ -533,41 +590,46 @@ def propose( common_attn_metadata._seq_lens_cpu = None common_attn_metadata._num_computed_tokens_cpu = None + block_size = self.block_size + assert block_size > 0, "block_size has not been initialized." for token_index in range(self.num_speculative_tokens - 1): # Update the inputs. # cast to int32 is crucial when eagle model is compiled. # tensor.argmax() returns int64 by default. input_ids = draft_token_ids_list[-1].int() + # Use fused kernel for slot mapping and metadata updates. + # Write clamped positions directly into the positions buffer to + # avoid an extra D2D copy for the common (non-mrope) case. + positions_1d = positions[0] if self.uses_mrope else positions if self.uses_mrope: - positions += 1 - # NOTE(woosuk): We should handle the case where the draft model - # generates tokens beyond the max model length. - # Since it is complex to remove such requests from the batch, - # we keep them in the batch but adjust the position ids - # and slot mappings to avoid the - # out-of-range access during the model execution. - # The draft tokens generated with this adjustment - # should be ignored. - exceeds_max_model_len = positions[0] >= self.max_model_len - # Mask out the position ids that exceed the max model length. - # Otherwise, we may get out-of-range error in RoPE. - clamped_positions = torch.where( - exceeds_max_model_len.unsqueeze(0), - torch.zeros_like(positions), - positions, - ) + out_pos = self.mrope_positions[0, :batch_size] + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0: + out_pos = self.xdrope_positions[0, :batch_size] else: - positions += 1 - exceeds_max_model_len = positions >= self.max_model_len - clamped_positions = torch.where(exceeds_max_model_len, 0, positions) - # For data integrity when async scheduling, we shouldn't use in place - # operations in case they are modified in next step's `prepare_input` - # of main model. - # Increment the sequence lengths. - common_attn_metadata.seq_lens += 1 - # For the requests that exceed the max model length, we set the - # sequence length to 1 to minimize their overheads in attention. - common_attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1) + out_pos = self.positions[:batch_size] + eagle_step_update_slot_mapping_and_metadata( + positions_1d=positions_1d, + block_table_tensor=common_attn_metadata.block_table_tensor, + seq_lens=common_attn_metadata.seq_lens, + block_size=block_size, + max_model_len=self.max_model_len, + out_clamped_positions=out_pos, + out_slot_mapping=self._slot_mapping_buffer[:input_batch_size], + input_batch_size=input_batch_size, + ) + common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size] + if self.uses_mrope: + self.mrope_positions[1:, :batch_size] = self.mrope_positions[ + 0, :batch_size + ] + positions = self.mrope_positions[:, :batch_size] + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0: + self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[ + 0, :batch_size + ] + positions = self.xdrope_positions[0, :batch_size] + else: + positions = self.positions[:batch_size] # Increment the maximum sequence length. We increment max_seq_len # unconditionally even though some seq_lens may have been capped above, # as max_seq_len serves as an upper bound for sequence lengths. @@ -582,45 +644,24 @@ def propose( if common_attn_metadata._num_computed_tokens_cpu is not None: common_attn_metadata._num_computed_tokens_cpu += 1 - # Compute the slot mapping. - block_size = self.block_size - assert block_size > 0, "block_size has not been initialized." - if self.uses_mrope: - # all dimensions of positions are the same - block_numbers = clamped_positions[0] // block_size - else: - block_numbers = clamped_positions // block_size - block_ids = common_attn_metadata.block_table_tensor.gather( - dim=1, index=block_numbers.view(-1, 1) - ) - block_ids = block_ids.view(-1) - if self.uses_mrope: - common_attn_metadata.slot_mapping = ( - block_ids * block_size + clamped_positions[0] % block_size + if can_reuse_draft_attn_metadata and reuse_draft_attn_metadata_ready: + self._update_reused_draft_attn_metadata( + per_layer_attn_metadata, + common_attn_metadata, + token_index + 1, ) else: - common_attn_metadata.slot_mapping = ( - block_ids * block_size + clamped_positions % block_size - ) - # Mask out the slot mappings that exceed the max model length. - # Otherwise, the KV cache will be inadvertently updated with the - # padding tokens. - common_attn_metadata.slot_mapping.masked_fill_( - exceeds_max_model_len, PADDING_SLOT_ID - ) - - # Rebuild attention metadata - for attn_group in self.draft_attn_groups: - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( - common_attn_metadata=common_attn_metadata, - draft_index=token_index + 1, + _, per_layer_attn_metadata = ( + self.build_per_group_and_layer_attn_metadata( + common_attn_metadata, + draft_index=token_index + 1, + ) ) - for layer_name in attn_group.layer_names: - per_layer_attn_metadata[layer_name] = attn_metadata + if can_reuse_draft_attn_metadata: + reuse_draft_attn_metadata_ready = True # copy inputs to buffer for cudagraph self.input_ids[:batch_size] = input_ids - self._set_positions(batch_size, clamped_positions) self.hidden_states[:batch_size] = hidden_states if self.supports_mm_inputs: self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids) @@ -646,9 +687,7 @@ def propose( num_tokens=input_batch_size, num_tokens_across_dp=batch_size_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, - slot_mapping=self._get_slot_mapping( - input_batch_size, common_attn_metadata.slot_mapping - ), + slot_mapping=self._get_slot_mapping(input_batch_size), ): ret_hidden_states = self.model(**model_kwargs) if not self.model_returns_tuple(): @@ -800,8 +839,53 @@ def set_inputs_first_pass( return total_num_output_tokens, token_indices_to_sample, new_cad + def build_model_inputs_first_pass( + self, + num_tokens: int, + num_input_tokens: int, + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None, + ) -> tuple[dict[str, Any], int]: + if self.supports_mm_inputs: + mm_embeds, is_mm_embed = mm_embed_inputs or (None, None) + + self.inputs_embeds[:num_tokens] = self.model.embed_input_ids( + self.input_ids[:num_tokens], + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) + + input_ids = None + inputs_embeds = self.inputs_embeds[:num_input_tokens] + else: + input_ids = self.input_ids[:num_input_tokens] + inputs_embeds = None + + model_kwargs = { + "input_ids": input_ids, + "positions": self._get_positions(num_input_tokens), + "inputs_embeds": inputs_embeds, + } + if self.pass_hidden_states_to_model: + model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens] + + return model_kwargs, num_input_tokens + + def build_per_group_and_layer_attn_metadata( + self, common_attn_metadata: CommonAttentionMetadata, draft_index: int = 0 + ) -> tuple[list[object], dict[str, object]]: + per_group_attn_metadata: list[object] = [] + per_layer_attn_metadata: dict[str, object] = {} + for attn_group in self.draft_attn_groups: + attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + common_attn_metadata=common_attn_metadata, draft_index=draft_index + ) + per_group_attn_metadata.append(attn_metadata) + for layer_name in attn_group.layer_names: + per_layer_attn_metadata[layer_name] = attn_metadata + return per_group_attn_metadata, per_layer_attn_metadata + def model_returns_tuple(self) -> bool: - return self.method not in ("mtp", "draft_model") + return self.method not in ("mtp", "draft_model", "dflash") def prepare_next_token_ids_cpu( self, @@ -838,7 +922,6 @@ def prepare_next_token_ids_cpu( def prepare_next_token_ids_padded( self, - common_attn_metadata: CommonAttentionMetadata, sampled_token_ids: torch.Tensor, requests: dict[str, CachedRequestState], gpu_input_batch: InputBatch, @@ -853,11 +936,10 @@ def prepare_next_token_ids_padded( """ # Precompute get_token_id for when there is no valid next token num_reqs = gpu_input_batch.num_reqs + seq_lens_list = (gpu_input_batch.num_tokens_no_spec[:num_reqs] - 1).tolist() self.backup_next_token_ids.np[:num_reqs] = np.array( [ - requests[gpu_input_batch.req_ids[i]].get_token_id( - common_attn_metadata.seq_lens_cpu[i].item() - ) + requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i]) for i in range(num_reqs) ], dtype=np.int32, @@ -942,7 +1024,7 @@ def prepare_inputs_padded( num_reqs=common_attn_metadata.num_reqs, num_actual_tokens=total_num_tokens, max_query_len=new_query_len_per_req.max().item(), - max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(), + max_seq_len=common_attn_metadata.max_seq_len, block_table_tensor=common_attn_metadata.block_table_tensor, slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens], causal=True, @@ -1234,6 +1316,21 @@ def get_model_name(self, model: nn.Module) -> str: model = model.module return model.__class__.__name__ + def _create_draft_vllm_config(self) -> VllmConfig: + """Return a VllmConfig with kernel-level overrides for the proposer. + Subclasses may override to apply additional config changes. + """ + spec_cfg = self.speculative_config + if spec_cfg.moe_backend is not None: + return replace( + self.vllm_config, + kernel_config=replace( + self.vllm_config.kernel_config, + moe_backend=spec_cfg.moe_backend, + ), + ) + return self.vllm_config + def _get_model(self) -> nn.Module: """ Default method to call get_model(). Can be overridden by subclasses which @@ -1241,9 +1338,10 @@ def _get_model(self) -> nn.Module: """ from vllm.compilation.backends import set_model_tag + draft_vllm_config = self._create_draft_vllm_config() with set_model_tag("eagle_head"): model = get_model( - vllm_config=self.vllm_config, + vllm_config=draft_vllm_config, model_config=self.speculative_config.draft_model_config, load_config=self.speculative_config.draft_load_config, ) @@ -1268,6 +1366,14 @@ def load_model(self, target_model: nn.Module) -> None: set(all_attn_layers.keys()) - target_attn_layer_names ) + # Disable DCP on draft model attention layers. + # Draft models (MTP) have their own local KV cache and don't + # participate in decode context parallelism. + for layer_name in self._draft_attn_layer_names: + layer = all_attn_layers[layer_name] + if hasattr(layer, 'impl') and hasattr(layer.impl, 'dcp_world_size'): + layer.impl.dcp_world_size = 1 + if self.supports_mm_inputs: # Even if the target model is multimodal, we can also use # text-only draft models @@ -1298,6 +1404,10 @@ def load_model(self, target_model: nn.Module) -> None: self.model.config.image_token_index = ( target_model.config.vision_config.image_token_id ) + elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration": + self.model.config.image_token_index = ( + target_model.config.media_placeholder_token_id + ) else: self.model.config.image_token_index = ( target_model.config.image_token_index @@ -1311,15 +1421,20 @@ def load_model(self, target_model: nn.Module) -> None: self._maybe_share_embeddings(target_language_model) self._maybe_share_lm_head(target_language_model) - if self.parallel_drafting and self.pass_hidden_states_to_model: - assert self.parallel_drafting_hidden_state_tensor is not None - self.parallel_drafting_hidden_state_tensor.copy_( - self.model.combine_hidden_states( - self.model.mask_hidden.view(3 * self.hidden_size) + if ( + self.parallel_drafting + and self.pass_hidden_states_to_model + and self.parallel_drafting_hidden_state_tensor is not None + ): + flat_mask = self.model.mask_hidden.view(-1) + if self.eagle3_use_aux_hidden_state: + # EAGLE3: mask_hidden stores all aux hidden states, + # project through combine_hidden_states + self.parallel_drafting_hidden_state_tensor.copy_( + self.model.combine_hidden_states(flat_mask) ) - if self.eagle3_use_aux_hidden_state - else self.model.mask_hidden.view(self.hidden_size) - ) + else: + self.parallel_drafting_hidden_state_tensor.copy_(flat_mask) def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: """ @@ -1407,6 +1522,8 @@ def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None: ) elif ( hasattr(target_language_model, "lm_head") + and hasattr(target_language_model.lm_head, "weight") + and hasattr(self.model.lm_head, "weight") and isinstance(target_language_model.lm_head.weight, torch.Tensor) and isinstance(self.model.lm_head.weight, torch.Tensor) # TODO: Offload to CPU for comparison to avoid extra GPU memory @@ -1492,8 +1609,9 @@ def dummy_run( ) -> None: # FIXME: when using tree-based specdec, adjust number of forward-passes # according to the depth of the tree. + only_one_forward_pass = is_graph_capturing or self.parallel_drafting for fwd_idx in range( - self.num_speculative_tokens if not is_graph_capturing else 1 + 1 if only_one_forward_pass else self.num_speculative_tokens ): if fwd_idx <= 1: cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (