# SPDX-License-Identifier: Apache-2.0 import copy from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import replace from functools import partial from typing import Any, cast import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import vllm.distributed.parallel_state as _ps # type: ignore[import-not-found] from vllm.config import CompilationMode, CUDAGraphMode, VllmConfig, get_layers_from_vllm_config from vllm.distributed.parallel_state import ( get_pp_group, get_tp_group, get_world_group, init_model_parallel_group, ) from vllm.forward_context import BatchDescriptor, ForwardContext, get_forward_context from vllm.logger import logger 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.deepseek_v2 import DeepseekV32IndexerCache from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.triton_utils import HAS_TRITON, triton from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.spec_decode.utils import ( PADDING_SLOT_ID, compute_new_slot_mapping, extend_all_queries_by_N, ) from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm_ascend.ascend_config import get_ascend_config from vllm_ascend.ascend_forward_context import _EXTRA_CTX, set_ascend_forward_context from vllm_ascend.attention.attention_mask import AttentionMaskBuilder from vllm_ascend.attention.attention_v1 import AscendAttentionState from vllm_ascend.attention.utils import AscendCommonAttentionMetadata from vllm_ascend.compilation.acl_graph import ACLGraphWrapper, update_full_graph_params from vllm_ascend.device.device_op import DeviceOperator from vllm_ascend.distributed.parallel_state import get_lmhead_tp_group from vllm_ascend.models.llama_eagle3_vwn import Eagle3VwnLlamaForCausalLM from vllm_ascend.ops.triton.spec_decode.utils import prepare_inputs_padded_kernel from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num from vllm_ascend.utils import check_gdn_layer, enable_sp, lmhead_tp_enable, shared_expert_dp_enabled from vllm_ascend.worker.utils import copy_snapshot_to_gpu @contextmanager def patch_tensor_parallel_group(tp_group): """Temporarily swap the global TP group for draft-model spec decode. vllm-ascend local implementation for swapping the global TP group so the draft model can run with a TP degree that differs from the target model. """ old_tp_group = _ps.get_tp_group() _ps._TP_STATE_PATCHED = True _ps._TP = tp_group try: yield finally: _ps._TP_STATE_PATCHED = False _ps._TP = old_tp_group # Currently we will fix block size to a small one since `num_reqs` can't be too large _PREPARE_INPUTS_BLOCK_SIZE = 4 # TODO: Remove it when the bug of fx-graph is solved # patch vllm_config to be in CompilationMode.NONE temporarily @contextmanager def _maybe_eager_context(vllm_config): target_compilation_config = vllm_config.compilation_config draft_compilation_config = replace( target_compilation_config, mode=CompilationMode.NONE, ) # Model layers use these registries even when compilation is disabled. draft_compilation_config.static_forward_context = target_compilation_config.static_forward_context draft_compilation_config.static_all_moe_layers = target_compilation_config.static_all_moe_layers vllm_config.compilation_config = draft_compilation_config try: yield finally: vllm_config.compilation_config = target_compilation_config # split hidden states along dimension of sequence def split_inputs_tp_to_sp(hidden_states, out): # tp and sp share the same group group = get_tp_group() world_size = group.world_size rank = group.rank num_tokens = hidden_states.shape[0] # the size per rank after padded padded_num_tokens_per_rank = (num_tokens + world_size - 1) // world_size # compute the start and end of slice start = padded_num_tokens_per_rank * rank end = padded_num_tokens_per_rank * (rank + 1) # copy only hidden_states in current rank hidden_states_curr_rank = hidden_states[start:end] out[: hidden_states_curr_rank.shape[0]] = hidden_states_curr_rank return out[:padded_num_tokens_per_rank] def greedy_sample(logits: torch.Tensor) -> torch.Tensor: tp_group = get_tp_group() B, V_local = logits.shape rank = tp_group.rank_in_group local_max_logits, local_max_indices = logits.max(dim=-1) local_global_idx = local_max_indices + rank * V_local # [B] # [B, world_size] gathered_logits = tp_group.all_gather(local_max_logits.unsqueeze(-1), dim=-1) gathered_global_idx = tp_group.all_gather(local_global_idx.unsqueeze(-1), dim=-1) # [B, world_size] global_max_rank = gathered_logits.argmax(dim=-1) # [B] target_argmax = gathered_global_idx.gather(dim=-1, index=global_max_rank.unsqueeze(-1)).squeeze(-1) # [B] return target_argmax # TODO(lilinsiman): Remove this code segment after future versions of the GLM # series models support graph input for speculative inference. def _is_glm_model(model_config) -> bool: """Return True if the target model belongs to the GLM series. Detection is based on the model_type string (covers glm, chatglm, glm4, glm4_moe, glm4_moe_lite, glm4_1v, glm_ocr, glm_moe_dsa, etc). """ hf_text_config = getattr(model_config, "hf_text_config", None) model_type = getattr(hf_text_config, "model_type", "") or "" return "glm" in str(model_type).lower() class AscendSpecDecodeBaseProposer(SpecDecodeBaseProposer): _runnable: ACLGraphWrapper | Callable def __init__(self, vllm_config: VllmConfig, device: torch.device, pass_hidden_states_to_model: bool, runner=None): super().__init__(vllm_config, device, pass_hidden_states_to_model, runner=runner) # Assign runner before it's used in the methods below self.runner = runner logger.debug( "[spec_decode/base] Initializing spec decode proposer: method=%s," " num_speculative_tokens=%s, hidden_size=%s, pass_hidden_states=%s," " parallel_drafting=%s, use_cuda_graph=%s, device=%s", self.method, self.num_speculative_tokens, self.hidden_size, pass_hidden_states_to_model, self.speculative_config.parallel_drafting if self.speculative_config else False, runner._use_aclgraph() if runner else False, device, ) self.use_async_scheduling = self.vllm_config.scheduler_config.async_scheduling self.use_compress = hasattr(self.vllm_config.model_config.hf_config, "compress_ratios") self.has_gdn = check_gdn_layer(self.vllm_config) self.pass_hidden_states_to_model = pass_hidden_states_to_model self.decode_threshold = 1 + self.num_speculative_tokens self.query_start_loc = self.runner._make_buffer(self.runner.max_num_reqs + 2, dtype=torch.int32) self.arange_cpu = torch.arange(self.arange.shape[0], device="cpu", dtype=torch.int32) self.attn_mask_builder = AttentionMaskBuilder(self.device) self.enable_shared_expert_dp = shared_expert_dp_enabled() self.pcp_size = self.runner.pcp_size self.dcp_size = self.runner.dcp_size self.use_sparse = hasattr(vllm_config.model_config.hf_text_config, "index_topk") self._share_mtp_indices = False spec_config = self.vllm_config.speculative_config draft_model_config = getattr(spec_config, "draft_model_config", None) draft_hf_config = draft_model_config.hf_config if draft_model_config is not None else None self._share_mtp_indices = getattr(draft_hf_config, "index_share_for_mtp_iteration", False) # NOTE: # `draft_tensor_parallel_size` does not take effect for Eagle: # the draft model uses the same TP size as the target model in practice. # so we applied this patch to set tp=1 of draft model separately. # Due to verification of `_verify_and_get_draft_tp` in vllm, # the value of `draft_tensor_parallel_size` here will either be 1 separately # or the same as target model. # TODO(zhaomingyu13): If we want to adapt to the case where draft model tp # is not 1 and differs from target model, this part should be rewritten. if vllm_config.parallel_config.tensor_parallel_size != self.speculative_config.draft_tensor_parallel_size: tp_group = init_model_parallel_group( [[get_world_group().rank]], get_world_group().rank, torch.distributed.get_backend(get_world_group().device_group), use_message_queue_broadcaster=True, group_name="tp", ) self.tp_group_context: AbstractContextManager[Any] = patch_tensor_parallel_group(tp_group) else: self.tp_group_context = nullcontext() self.use_cuda_graph = self.runner._use_aclgraph() and not self.speculative_config.enforce_eager self._raise_if_padded_drafter_batch_disabled_and_full_graph_enabled() # GLM series models: speculative decoding does not yet support running # the draft model in graph mode. Force the draft model to always use # eager mode. This is equivalent to the user adding # `"enforce_eager": true` to the `--speculative-config`, and keeps # the target model's graph-mode setting untouched. # TODO(lilinsiman): Remove this code segment after future versions of the GLM # series models support graph input for speculative inference. if _is_glm_model(self.vllm_config.model_config): if self.use_cuda_graph: logger.warning( "GLM series models with speculative decoding currently do " "not support graph mode. The draft model has been " "automatically switched to eager mode " "(enforce_eager=true). Graph mode support for GLM " "speculative decoding will be added in a future release. " ) self.use_cuda_graph = False # TODO: Remove it when the bug of fx-graph is solved self.maybe_eager_context: AbstractContextManager[Any] = nullcontext() if not self.use_cuda_graph and enable_sp(vllm_config): self.maybe_eager_context = _maybe_eager_context(vllm_config) self.token_indices_to_sample = torch.zeros( self.vllm_config.scheduler_config.max_num_batched_tokens, dtype=torch.int32, device=device ) metadata_lens = self.runner.max_num_tokens + 2 * self.pcp_size * self.runner.max_num_reqs num_mtp_draft_slots = max(self.num_speculative_tokens - 1, 0) * self.runner.max_num_reqs slot_mapping_lens = metadata_lens + num_mtp_draft_slots self.slot_mapping_group = [ torch.zeros(slot_mapping_lens, dtype=torch.int32, device=device, pin_memory=self.runner.pin_memory) for _ in range(self.num_speculative_tokens) ] # dsv32 needs seq_lens and query_start_loc persistent tensors for full graph mode self.seq_lens_group = [ torch.zeros(metadata_lens, dtype=torch.int32, device=device, pin_memory=self.runner.pin_memory) for _ in range(self.num_speculative_tokens) ] self.query_start_loc_group = [ torch.zeros(metadata_lens, dtype=torch.int32, device=device, pin_memory=self.runner.pin_memory) for _ in range(self.num_speculative_tokens) ] # pcp needs independent block table tensor in step=0 and step>0, and the following is for step>0 # since final block table tensor is not ready in __init__, it is delayed until dummy_run self.block_table_tensor_clone: torch.Tensor | None = None self._runnable = self._run_merged_draft self.is_multimodal_model = self.vllm_config.model_config.is_multimodal_model if self.uses_mrope: self.mrope_positions = torch.zeros((3, self.max_num_tokens + 1), dtype=torch.int32, 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), dtype=torch.int32, device=device, ) else: # RoPE need (max_num_tokens,) self.positions = torch.zeros(self.max_num_tokens, dtype=torch.int32, device=device) self.token_arange_np = np.arange(self.max_num_tokens + 1, dtype=np.int32) self.enable_enpu = self.runner.enable_enpu self.use_eagle = self.runner.use_eagle def _raise_if_padded_drafter_batch_disabled_and_full_graph_enabled(self): if ( self.speculative_config.disable_padded_drafter_batch and self.use_cuda_graph and self.compilation_config.cudagraph_mode.has_full_cudagraphs() ): raise NotImplementedError( "Speculative Decoding with cudagraph mode containing full cudagraphs only " "supports padded drafter batch. Please unset " "disable_padded_drafter_batch in the speculative_config." ) def _get_model(self) -> nn.Module: """ Default method to call get_model(). Can be overridden by subclasses which need to customize model loading. """ from vllm.compilation.backends import set_model_tag draft_vllm_config = self._create_draft_vllm_config() draft_load_config = self.speculative_config.draft_load_config logger.info( "[spec_decode/base] Loading draft model: method=%s, load_format=%s, model=%s", self.method, getattr(draft_load_config, "load_format", None), getattr(self.speculative_config.draft_model_config, "model", None), ) with set_model_tag("eagle_head"): model = get_model( vllm_config=draft_vllm_config, model_config=self.speculative_config.draft_model_config, load_config=self.speculative_config.draft_load_config, ) return model def load_model(self, model: nn.Module) -> None: assert get_pp_group().is_last_rank, f"{self.method} drafter must be loaded on the last pipeline stage." target_attn_layer_names = set(get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()) with self.maybe_eager_context: self.model = self._get_model() # Find draft layers (attention layers added by draft model) all_attn_layers = get_layers_from_vllm_config( self.vllm_config, AttentionLayerBase, # type: ignore[type-abstract] ) all_indexer_layer_names = set(get_layers_from_vllm_config(self.vllm_config, DeepseekV32IndexerCache).keys()) # Filter to only layers that have KV cache specs. self._draft_attn_layer_names = { name for name in (set(all_attn_layers.keys()) - target_attn_layer_names) if all_attn_layers[name].get_kv_cache_spec(self.vllm_config) is not None } - all_indexer_layer_names self.attn_layer_names = list(sorted(self._draft_attn_layer_names)) draft_attn_layers_dict = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase) # initialized for mamba models self.kernel_block_size = ( draft_attn_layers_dict[self.attn_layer_names[0]].get_attn_backend().get_supported_kernel_block_sizes()[0] ) self.piece_all_attn_layer_name = [] for _ in range(self.num_speculative_tokens): self.piece_all_attn_layer_name.append([name for name in self.attn_layer_names]) if supports_multimodal(model): # handle multimodality if self.get_model_name(model) in [ "Qwen2_5_VLForConditionalGeneration", "Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", "Step3p7ForConditionalGeneration", ]: self.model.config.image_token_index = model.config.image_token_id elif self.get_model_name(model) == "PixtralForConditionalGeneration": self.model.config.image_token_index = model.config.vision_config.image_token_id elif self.get_model_name(model) == "KimiK25ForConditionalGeneration": self.model.config.image_token_index = model.config.media_placeholder_token_id else: self.model.config.image_token_index = model.config.image_token_index target_language_model = model.get_language_model() else: target_language_model = model # share embed_tokens with the target model if needed self._maybe_share_embeddings(target_language_model) self._maybe_share_topk_indices(target_language_model) self._maybe_share_lm_head(model) if ( self.parallel_drafting and self.pass_hidden_states_to_model and 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.eagle3_use_aux_hidden_state else self.model.mask_hidden.view(self.hidden_size) ) def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: """ Some draft models may not have their own embedding layers, and some may have a duplicate copy of the target model's embedding layers. In these cases, we share the target model's embedding layers with the draft model to save memory. """ if get_pp_group().world_size == 1: if hasattr(target_language_model.model, "embed_tokens"): target_embed_tokens = target_language_model.model.embed_tokens elif hasattr(target_language_model.model, "embedding"): target_embed_tokens = target_language_model.model.embedding else: raise AttributeError("Target model does not have 'embed_tokens' or 'embedding' attribute") # If pp>1, the weights of mtp and the main model's embedding are not on the same device. # check if mtp model use main model's embedding and LMhead share_embeddings = False if hasattr(self.model, "has_own_embed_tokens"): # EAGLE model if not self.model.has_own_embed_tokens: share_embeddings = True logger.info( "[spec_decode/base] Detected EAGLE model without its own" " embed_tokens in the checkpoint. Sharing target model" " embedding weights with the draft model." ) elif ( isinstance(target_embed_tokens.weight, torch.Tensor) and isinstance(self.model.model.embed_tokens.weight, torch.Tensor) # TODO: Offload to CPU for comparison to avoid extra NPU memory # usage in CI testing environments with limited NPU memory and torch.equal( target_embed_tokens.weight.cpu(), self.model.model.embed_tokens.weight.cpu(), ) ): share_embeddings = True logger.info( "[spec_decode/base] Detected EAGLE model with embed_tokens" " identical to the target model. Sharing target model embedding" " weights with the draft model." ) else: logger.info( "[spec_decode/base] Detected EAGLE model with distinct" " embed_tokens weights. Keeping separate embedding weights" " from the target model." ) else: # MTP model share_embeddings = not self.use_compress if share_embeddings: logger.info( "[spec_decode/base] Detected MTP model. Sharing target model" " embedding weights with the draft model." ) if share_embeddings: if hasattr(self.model.model, "embed_tokens"): del self.model.model.embed_tokens self.model.model.embed_tokens = target_embed_tokens else: logger.info( "[spec_decode/base] PP>1: draft model loaded its own vocab embedding" " weights instead of sharing them with the target model." ) # share lm_head with the target model if needed def _maybe_share_lm_head(self, model: nn.Module) -> None: # some model definition do not define lm_head explicitly # and reuse embed_tokens for lm_head, e.g., CohereForCausalLM if self.method in ("eagle", "dflash"): # For DFlash drafters trained with a reduced draft vocabulary, the # draft model ships its own lm_head of shape [draft_vocab_size, # hidden] whose rows map to a trained subset of the target vocab via # the draft_id_to_target_id (d2t) buffer. Overwriting it with the # target lm_head ([target_vocab_size, hidden]) makes the draft emit # logits over the wrong vocabulary, so the verifier rejects almost # every speculative token. Keep the draft's own lm_head in that case. draft_has_own_lm_head = ( self.method == "dflash" and getattr(self.model, "draft_id_to_target_id", None) is not None ) if draft_has_own_lm_head: logger.info( "[spec_decode/base] DFlash draft uses d2t vocab remapping;" " keeping the draft's own lm_head instead of sharing the target" " lm_head." ) else: logger.info("[spec_decode/base] Loading EAGLE/DFLASH LM head weights from the target model.") if hasattr(model, "lm_head"): self.model.lm_head = model.lm_head elif hasattr(model, "get_language_model") and hasattr(model.get_language_model(), "lm_head"): self.model.lm_head = model.get_language_model().lm_head else: logger.warning( "[spec_decode/base] Target model has no accessible lm_head" " for sharing. Draft model will use its own lm_head." " This may cause incorrect logits if the draft lm_head" " is not trained." ) if self.method == "mtp" and self.vllm_config.model_config.is_deepseek_mla: for _, layer_module in self.model.model.layers.items(): if torch.equal(layer_module.shared_head.head.weight, model.lm_head.weight): layer_module.shared_head.head = model.lm_head if self.vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs() and self.use_cuda_graph: logger.info( "[spec_decode/base] Wrapping draft model with ACLGraphWrapper:" " runtime_mode=FULL, use_eagle=%s, enable_enpu=%s", self.use_eagle, self.enable_enpu, ) self.update_stream = torch.npu.Stream() self._runnable = ACLGraphWrapper( self._run_merged_draft, self.vllm_config, runtime_mode=CUDAGraphMode.FULL, use_eagle=self.use_eagle, enable_enpu=self.enable_enpu, ) def _maybe_share_topk_indices(self, target_language_model: nn.Module) -> None: if hasattr(target_language_model.model, "topk_indices_buffer"): if hasattr(self.model.model, "topk_indices_buffer"): del self.model.model.topk_indices_buffer self.model.model.topk_indices_buffer = target_language_model.model.topk_indices_buffer logger.info( "[spec_decode/base] Detected MTP model with topk_indices_buffer." " Sharing target model topk_indices_buffer with the draft model." ) target_buffer = target_language_model.model.topk_indices_buffer draft_model = getattr(self.model, "model", None) if target_buffer is not None and draft_model is not None: for _, module in draft_model.named_modules(): if hasattr(module, "topk_indices_buffer"): module.topk_indices_buffer = target_buffer def get_model(self) -> nn.Module: # get raw model out of the aclgraph wrapper. if isinstance(self.model, ACLGraphWrapper): return self.model.unwrap() return self.model def shallow_copy_metadata(self, attn_metadata): # Currently, new objects will be assigned to the lists in attn_metadata # when update. So we can use the shallow copy. return copy.copy(attn_metadata) def _freeze_draft_index_attn_metadata(self, attn_metadata): decode_metadata = getattr(attn_metadata, "decode", None) if decode_metadata is not None: if decode_metadata.sas_metadata is not None: decode_metadata.sas_metadata = decode_metadata.sas_metadata.clone() return attn_metadata @torch.inference_mode() def dummy_run( self, num_tokens: int, with_prefill: bool = False, in_graph_capturing: bool = False, num_reqs: int = 0, num_tokens_across_dp: torch.Tensor | None = None, aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, batch_descriptor=None, dummy_compute_logits=lambda hidden_states: None, is_profile=False, ): ( num_tokens, num_tokens_across_dp, _, ) = self.runner._sync_metadata_across_dp(num_tokens, is_draft_model=True) pcp_manager = getattr(self.runner, "pcp_manager", None) multi_steps_attn_metadata = [] if not self.use_cuda_graph: aclgraph_runtime_mode = CUDAGraphMode.NONE # init block table tensor clone is only available after profile run and is only used for graph mode if ( self.pcp_size * self.dcp_size > 1 and self.use_cuda_graph and not is_profile and self.block_table_tensor_clone is None ): self.block_table_tensor_clone = torch.zeros( ( self.runner.max_num_tokens + 2 * self.pcp_size * self.runner.max_num_reqs, self.runner.input_batch.block_table[0].get_device_tensor().shape[1], ), dtype=torch.int32, device=self.device, pin_memory=self.runner.pin_memory, ) if aclgraph_runtime_mode == CUDAGraphMode.FULL and len(self.runner.attn_groups) > 0: num_computed_tokens_cpu = self.runner.input_batch.num_computed_tokens_cpu_tensor[:num_reqs] # num_reqs is already the padded version self.query_start_loc.cpu[: num_reqs + 1].copy_(self.runner.query_start_loc.cpu[: num_reqs + 1]) copy_snapshot_to_gpu(self.query_start_loc) common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], seq_lens_cpu=self.runner.optimistic_seq_lens_cpu, _seq_lens_cpu=self.runner.optimistic_seq_lens_cpu, seq_lens_cpu_upper_bound=self.runner.optimistic_seq_lens_cpu, seq_lens=self.runner.seq_lens[:num_reqs], num_reqs=num_reqs, num_actual_tokens=num_tokens, num_input_tokens=num_tokens, max_query_len=self.num_speculative_tokens + 1, num_computed_tokens_cpu=num_computed_tokens_cpu, actual_seq_lengths_q=self.runner.actual_seq_lengths_q, block_table_tensor=self.runner.input_batch.block_table[self.kv_cache_gid].get_device_tensor()[ :num_reqs ], # This is used to hold a position. slot_mapping=self.runner.input_batch.block_table[self.kv_cache_gid].slot_mapping.gpu, positions=self.runner.positions, positions_cpu=self.runner._dsa_positions_cpu_buf if self.use_compress else None, attn_state=self.runner.attn_state, decode_token_per_req=self.runner.decode_token_per_req, is_prefilling=torch.zeros(num_reqs, dtype=torch.bool), max_seq_len=0, ) if pcp_manager is not None: # update long_seq related params and flatten block_table common_attn_metadata.prefill_context_parallel_metadata = pcp_manager.long_seq_metadata assert len(self.draft_attn_groups) > 0 builder = self.draft_attn_groups[0].get_metadata_builder() kv_cache_spec = self.draft_attn_groups[0].kv_cache_spec # update the tensor's address for each step. for draft_index in range(self.num_speculative_tokens): common_attn_metadata = self.shallow_copy_metadata(common_attn_metadata) extra_attn_metadata_args: dict = {} if self.use_compress: extra_attn_metadata_args.update( prefill_ratio_to_sas_metadata=dict(), decode_ratio_to_sas_metadata=dict(), common_ratio_to_sas_metadata=dict(), block_size=kv_cache_spec.block_size, ) # Set the real slot_mapping. slot_mapping_lens = common_attn_metadata.slot_mapping.shape[0] self.slot_mapping_group[draft_index][:slot_mapping_lens].copy_(common_attn_metadata.slot_mapping) self.slot_mapping_group[draft_index][slot_mapping_lens:].fill_(PADDING_SLOT_ID) common_attn_metadata.slot_mapping = self.slot_mapping_group[draft_index] self.seq_lens_group[draft_index][:num_reqs].copy_(common_attn_metadata.seq_lens) self.seq_lens_group[draft_index][num_reqs:].fill_(0) common_attn_metadata.seq_lens = self.seq_lens_group[draft_index][:num_reqs] self.query_start_loc_group[draft_index][: num_reqs + 1].copy_(common_attn_metadata.query_start_loc) self.query_start_loc_group[draft_index][num_reqs + 1 :].fill_(0) common_attn_metadata.query_start_loc = self.query_start_loc_group[draft_index][: num_reqs + 1] if self.pcp_size * self.dcp_size > 1 and draft_index > 0: assert self.block_table_tensor_clone is not None, "block_table_tensor_clone is not init" common_attn_metadata.block_table_tensor = self.block_table_tensor_clone[:num_reqs] if not self.use_compress or draft_index == 0: attn_metadata_eagle = builder.build_for_graph_capture( common_attn_metadata, AscendAttentionState.SpecDecoding if self.method == "mtp" else AscendAttentionState.ChunkedPrefill, **extra_attn_metadata_args, ) else: attn_metadata_eagle = builder.build_for_drafting( common_attn_metadata, draft_index, **extra_attn_metadata_args, ) per_layer_attn_metadata = dict() for layer_name in self.attn_layer_names: per_layer_attn_metadata[layer_name] = attn_metadata_eagle multi_steps_attn_metadata.append(per_layer_attn_metadata) model_positions = self._get_positions(num_tokens) batch_size = max(num_tokens // (self.num_speculative_tokens + 1), 1) # TODO: temporarily hack here, we should find out batch_size for profile_run if is_profile: batch_size = min(batch_size, self.runner.max_num_reqs) if self.supports_mm_inputs: mm_embeds, is_mm_embed = (None, None) inputs_embeds = self.model.embed_input_ids( self.input_ids[:num_tokens], multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed ) self.inputs_embeds[:num_tokens] = inputs_embeds inputs_embeds = self.inputs_embeds[:num_tokens] else: inputs_embeds = None self.token_indices_to_sample.fill_(0) with set_ascend_forward_context( multi_steps_attn_metadata[0] if multi_steps_attn_metadata else None, self.vllm_config, num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, num_actual_tokens=0, in_profile_run=is_profile, batch_descriptor=batch_descriptor, aclgraph_runtime_mode=aclgraph_runtime_mode, is_draft_model=True, draft_attn_metadatas=multi_steps_attn_metadata, ): # Reset MOE layer index before first model call forward_context = get_forward_context() if forward_context is not None: forward_context.moe_layer_index = 0 self._runnable( num_input_tokens=num_tokens, batch_size=batch_size, token_indices_to_sample=self.token_indices_to_sample[: batch_size * self.extra_slots_per_request], # The target_position's address is same as the model_positions's target_positions=model_positions, inputs_embeds=inputs_embeds, multi_steps_attn_metadata=multi_steps_attn_metadata, num_tokens=num_tokens, ) forward_context = get_forward_context() if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL and not _EXTRA_CTX.capturing: self._update_full_graph_params(forward_context, num_tokens, multi_steps_attn_metadata) def _update_full_graph_params_if_needed( self, forward_context: ForwardContext, num_input_tokens: int, multi_steps_attn_metadata: list[dict[str, Any]], ) -> None: if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL: self._update_full_graph_params(forward_context, num_input_tokens, multi_steps_attn_metadata) def _propose( self, # [num_tokens] target_token_ids: torch.Tensor, # [num_tokens] or [3, num_tokens] when M-RoPE is enabled target_positions: torch.Tensor, # [num_tokens, hidden_size] target_hidden_states: torch.Tensor, # [batch_size] next_token_ids: torch.Tensor, token_indices_to_sample: torch.Tensor | None, common_attn_metadata: CommonAttentionMetadata, target_model_batch_desc: BatchDescriptor, sampling_metadata: SamplingMetadata, mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, req_scheduled_tokens=None, long_seq_metadata=None, num_prefill_reqs=0, num_decode_reqs=0, scheduler_output: SchedulerOutput = None, num_scheduled_tokens: int = 0, num_rejected_tokens_gpu: torch.Tensor | None = None, ) -> torch.Tensor: batch_size = common_attn_metadata.batch_size() if token_indices_to_sample is None: token_indices_to_sample = common_attn_metadata.query_start_loc[1:] - 1 if self.method in ("eagle3", "dflash"): assert isinstance( self.get_model(), ( Eagle3LlamaForCausalLM, DFlashQwen3ForCausalLM, Eagle3VwnLlamaForCausalLM, Eagle3DeepseekV2ForCausalLM, ), ) target_hidden_states = self.model.combine_hidden_states(target_hidden_states) assert target_hidden_states.shape[-1] == self.hidden_size num_tokens, token_indices_to_sample, common_attn_metadata, long_seq_args = self.set_inputs_first_pass( target_token_ids=target_token_ids, next_token_ids=next_token_ids, target_positions=target_positions, target_hidden_states=target_hidden_states, token_indices_to_sample=token_indices_to_sample, cad=common_attn_metadata, num_rejected_tokens_gpu=num_rejected_tokens_gpu, req_scheduled_tokens=req_scheduled_tokens, long_seq_metadata=long_seq_metadata, num_prefill_reqs=num_prefill_reqs, num_decode_reqs=num_decode_reqs, ) assert self.runner is not None pcp_manager = getattr(self.runner, "pcp_manager", None) if pcp_manager is not None: assert long_seq_args is not None _, ori_token_indices_to_sample = long_seq_args has_lora = len(self.runner.input_batch.lora_id_to_lora_request) > 0 uniform_decode = target_model_batch_desc.uniform if self.use_cuda_graph: _, batch_descriptor = self.runner.cudagraph_dispatcher.dispatch( num_tokens=num_tokens, uniform_decode=uniform_decode, has_lora=has_lora ) num_input_tokens = batch_descriptor.num_tokens else: num_input_tokens = num_tokens ( num_input_tokens, num_tokens_across_dp, _, ) = self.runner._sync_metadata_across_dp(num_input_tokens, is_draft_model=True) if self.use_cuda_graph: aclgraph_runtime_mode, batch_descriptor = self.runner.cudagraph_dispatcher.dispatch( num_tokens=num_input_tokens, uniform_decode=uniform_decode, has_lora=has_lora ) num_input_tokens = batch_descriptor.num_tokens else: aclgraph_runtime_mode = CUDAGraphMode.NONE batch_descriptor = None if aclgraph_runtime_mode == CUDAGraphMode.FULL: # TODO: Due to the inconsistency between the proposer `dispatcher` and model runner, this padding # should have been done in model runner but not. For example, at prefill stage, target model # is run in eager mode currently, which means `_pad_query_start_loc_for_fia` is not called, # while draft model is run in graph model, which means we should pad the `query_start_loc`. # Need to be fixed in the future. num_reqs = common_attn_metadata.query_start_loc.shape[0] self.query_start_loc.gpu[:num_reqs].copy_(common_attn_metadata.query_start_loc) self.query_start_loc.cpu[:num_reqs].copy_(common_attn_metadata.query_start_loc_cpu) num_reqs_padded = self.runner._pad_query_start_loc_for_fia( self.query_start_loc, num_input_tokens, batch_descriptor.num_reqs if batch_descriptor.num_reqs is not None else common_attn_metadata.num_reqs, common_attn_metadata.num_reqs, aclgraph_runtime_mode, batch_descriptor.num_reqs, ) common_attn_metadata.num_reqs = num_reqs_padded common_attn_metadata.query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1] common_attn_metadata.query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs_padded + 1] slicing_length = ( num_reqs_padded * self.decode_threshold if self.pcp_size * self.dcp_size > 1 else num_reqs_padded ) common_attn_metadata.block_table_tensor = self._adjust_tensor( common_attn_metadata.block_table_tensor, slicing_length ) if self.method == "dflash": common_attn_metadata.seq_lens = self._adjust_tensor(common_attn_metadata.seq_lens, num_reqs_padded) else: common_attn_metadata.seq_lens = self._adjust_tensor(self.runner.seq_lens, num_reqs_padded) common_attn_metadata.seq_lens_cpu = self._adjust_tensor( self.runner.optimistic_seq_lens_cpu, num_reqs_padded ) # Keep the upstream-canonical mirror length-aligned with the # padded subclass field, but only if the caller already # populated it (production cm_base does; some unit-test mocks # leave it None and assert it stays None). ``.clone()`` keeps # the two fields independent so per-step in-place updates in # ``attn_update_stack_num_spec_norm`` don't double-count. if common_attn_metadata._seq_lens_cpu is not None: common_attn_metadata._seq_lens_cpu = common_attn_metadata.seq_lens_cpu.clone() if common_attn_metadata.num_computed_tokens_cpu is not None: common_attn_metadata.num_computed_tokens_cpu = self._adjust_tensor( common_attn_metadata.num_computed_tokens_cpu, num_reqs_padded ) if pcp_manager is not None and self.pcp_size > 1: pcp_manager.mask_spec_decode_restore_idx_for_graph( common_attn_metadata.prefill_context_parallel_metadata.pcp_allgather_restore_idx ) else: num_reqs_padded = common_attn_metadata.num_reqs # In the below scenario, padding has been applied by _pad_query_start_loc_for_fia in the model runner. # We need to unpad here for eager mode to maintain compatibility. if not self.vllm_config.model_config.use_mla and self.pcp_size * self.dcp_size == 1: common_attn_metadata.block_table_tensor = self._adjust_tensor( common_attn_metadata.block_table_tensor, num_reqs_padded ) if self.supports_mm_inputs: mm_embeds, is_mm_embed = mm_embed_inputs or (None, None) inputs_embeds = self.model.embed_input_ids( self.input_ids[:num_tokens], multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed ) self.inputs_embeds[:num_tokens] = inputs_embeds inputs_embeds = self.inputs_embeds[:num_input_tokens] else: inputs_embeds = None # Update slot_mapping for different speculative. # NOTE: Currently, we only remake the slot_mapping, because it's the # only tensor which will be used in current FIA. # Strictly speaking, `query_start_loc`, `seq_lens` should also have # their memory allocated separately for each step just like `slot_mapping`. slot_mapping_lens = common_attn_metadata.slot_mapping.shape[0] self.slot_mapping_group[0][:slot_mapping_lens].copy_(common_attn_metadata.slot_mapping) self.slot_mapping_group[0][slot_mapping_lens:].fill_(-1) common_attn_metadata.slot_mapping = self.slot_mapping_group[0] self.seq_lens_group[0][:num_reqs_padded].copy_(common_attn_metadata.seq_lens) self.seq_lens_group[0][num_reqs_padded:].fill_(0) common_attn_metadata.seq_lens = self.seq_lens_group[0][:num_reqs_padded] self.query_start_loc_group[0][: num_reqs_padded + 1].copy_(common_attn_metadata.query_start_loc) self.query_start_loc_group[0][num_reqs_padded + 1 :].fill_(0) common_attn_metadata.query_start_loc = self.query_start_loc_group[0][: num_reqs_padded + 1] common_attn_metadata.num_input_tokens = num_input_tokens # FIXME(woosuk): The below two ops cause synchronization. Optimize. assert len(self.draft_attn_groups) > 0 builder = self.draft_attn_groups[0].get_metadata_builder() extra_attn_metadata_args: dict = {} if self.use_compress: extra_attn_metadata_args = dict( prefill_ratio_to_sas_metadata=dict(), decode_ratio_to_sas_metadata=dict(), common_ratio_to_sas_metadata=dict(), block_size=self.draft_attn_groups[0].kv_cache_spec.block_size, ) attn_metadata = builder.build(0, common_attn_metadata, self.runner.get_model(), **extra_attn_metadata_args) if hasattr(attn_metadata, "causal") and not attn_metadata.causal: attn_metadata.attn_mask = None if self.uses_mrope: used_update_positions = self.mrope_positions[:, token_indices_to_sample] else: used_update_positions = self.positions[token_indices_to_sample] per_layer_attn_metadata = dict() # The first step of speculative. for layer_name in self.attn_layer_names: per_layer_attn_metadata[layer_name] = attn_metadata multi_steps_attn_metadata = [per_layer_attn_metadata] # Copy the old attn_metadata and update attn_metadata_i = per_layer_attn_metadata[self.attn_layer_names[0]] # Clone the data so that when calculating the data at position 2 and position 3 # in the merged graph, it does not affect position 1 # FIXME(lilinsiman) if self.pcp_size * self.dcp_size > 1 and self.use_cuda_graph: assert self.block_table_tensor_clone is not None, "block_table_tensor_clone is not init" self.block_table_tensor_clone[: common_attn_metadata.block_table_tensor.shape[0]] = ( common_attn_metadata.block_table_tensor ) common_attn_metadata.block_table_tensor = self.block_table_tensor_clone[ : common_attn_metadata.block_table_tensor.shape[0] ] else: common_attn_metadata.block_table_tensor = common_attn_metadata.block_table_tensor.clone() metadata_has_prefill = bool(getattr(attn_metadata_i, "num_prefills", 0)) is_prefill_batch = num_prefill_reqs > 0 or metadata_has_prefill pcp_mtp_inputs = None draft_cp_kwargs = { "ori_seq_len": None, "ori_seq_len_cpu": None, "slot_indices": None, "mtp_slot_mapping": None, } if pcp_manager is not None: pcp_mtp_inputs = pcp_manager.prepare_spec_decode_mtp_drafting_inputs( common_attn_metadata=common_attn_metadata, attn_metadata=attn_metadata_i, ori_token_indices_to_sample=ori_token_indices_to_sample, batch_size=batch_size, num_decode_reqs=num_decode_reqs, is_prefill_batch=is_prefill_batch, num_speculative_tokens=self.num_speculative_tokens, ) if pcp_mtp_inputs is not None: draft_cp_kwargs.update( ori_seq_len=pcp_mtp_inputs.seq_lens, ori_seq_len_cpu=pcp_mtp_inputs.seq_lens_cpu, slot_indices=pcp_mtp_inputs.slot_indices, mtp_slot_mapping=pcp_mtp_inputs.slot_mapping, ) should_update_next_steps = not self.parallel_drafting and ( self.pcp_size * self.dcp_size == 1 or pcp_mtp_inputs is not None ) if should_update_next_steps: # Copy the old attn_metadata and update for draft_index in range(1, self.num_speculative_tokens): per_layer_attn_metadata = dict() for attn_group in self.draft_attn_groups: common_attn_metadata, attn_metadata = self.attn_update_stack_num_spec_norm( draft_index, attn_metadata, common_attn_metadata, batch_size, num_input_tokens, used_update_positions, aclgraph_runtime_mode, **draft_cp_kwargs, attn_group=attn_group, ) for layer_name in self.attn_layer_names: per_layer_attn_metadata[layer_name] = attn_metadata multi_steps_attn_metadata.append(per_layer_attn_metadata) token_indices_to_sample_len = token_indices_to_sample.shape[0] self.token_indices_to_sample[:token_indices_to_sample_len].copy_(token_indices_to_sample) self.token_indices_to_sample[token_indices_to_sample_len:].fill_(0) with set_ascend_forward_context( multi_steps_attn_metadata[0], self.vllm_config, num_tokens=num_input_tokens, num_tokens_across_dp=num_tokens_across_dp, num_actual_tokens=num_tokens, batch_descriptor=batch_descriptor, aclgraph_runtime_mode=aclgraph_runtime_mode, is_draft_model=True, draft_attn_metadatas=multi_steps_attn_metadata, eplb_heat_collection_status=( self.runner.eplb_heat_collection_status if self.runner.dynamic_eplb else False ), ): # Reset MOE layer index for forward pass forward_context = get_forward_context() if forward_context is not None: forward_context.moe_layer_index = 0 model_inputs: dict[str, Any] = { "num_input_tokens": num_input_tokens, "batch_size": batch_size, "token_indices_to_sample": self.token_indices_to_sample[:token_indices_to_sample_len], "target_positions": target_positions, "inputs_embeds": inputs_embeds, "multi_steps_attn_metadata": multi_steps_attn_metadata, "num_tokens": num_tokens, "is_prefill": is_prefill_batch, } runnable = cast(Callable[..., Any], self._runnable) run_draft: Callable[[], Any] = partial(runnable, **model_inputs) if self.enable_enpu: self._update_full_graph_params_if_needed(forward_context, num_input_tokens, multi_steps_attn_metadata) draft_token_ids = run_draft() else: draft_token_ids = run_draft() self._update_full_graph_params_if_needed(forward_context, num_input_tokens, multi_steps_attn_metadata) return draft_token_ids def compute_draft_token_ids(self, hidden_states: torch.Tensor): if self.method in ("eagle3", "dflash"): logits = self.model.logits_processor(self.model.lm_head, hidden_states) if not hasattr(self.model, "draft_id_to_target_id") or self.model.draft_id_to_target_id is None: return greedy_sample(logits) logits = logits.contiguous() next_token = greedy_sample(logits) bias = torch.index_select(self.model.draft_id_to_target_id, dim=0, index=next_token.view(-1)).view( next_token.shape ) return next_token + bias else: logits = self.model.compute_logits(hidden_states) return greedy_sample(logits) def _run_merged_draft( self, num_input_tokens, batch_size, token_indices_to_sample, target_positions, inputs_embeds, multi_steps_attn_metadata, num_tokens, is_prefill=None, ) -> torch.Tensor: # The lifecycle of `input_ids`, `positions`, `hidden_states` runs through all # speculative tokens' proposings. `model_input_ids`, `model_positions` and # `model_hidden_states` represent the speculative model inputs. model_input_ids = self.input_ids[:num_input_tokens] model_positions = self._get_positions(num_input_tokens) if self.method == "dflash": model_kwargs = self.build_model_inputs_first_pass(num_input_tokens) else: model_kwargs = { "input_ids": model_input_ids, "positions": model_positions, "inputs_embeds": inputs_embeds, } if self.pass_hidden_states_to_model: model_hidden_states = self.hidden_states[:num_input_tokens] model_hidden_states, model_positions = self.maybe_pad_and_reduce(model_hidden_states, model_positions) model_kwargs["hidden_states"] = model_hidden_states if self.method == "mtp": model_kwargs["positions"] = model_positions # step 0 draft_model = getattr(self.model, "model", None) if self._share_mtp_indices and draft_model is not None and hasattr(draft_model, "set_skip_topk"): draft_model.set_skip_topk(False) ret_hidden_states = self.model(**model_kwargs) if not self.model_returns_tuple(): last_hidden_states = ret_hidden_states hidden_states = last_hidden_states else: last_hidden_states, hidden_states = ret_hidden_states # step 1+ skip indexer draft_model = getattr(self.model, "model", None) if self._share_mtp_indices and draft_model is not None and hasattr(draft_model, "set_skip_topk"): draft_model.set_skip_topk(True) if self.method != "dflash": last_hidden_states, model_positions, hidden_states = self.maybe_all_gather_and_unpad( last_hidden_states, model_positions, hidden_states ) num_indices = token_indices_to_sample.shape[0] if lmhead_tp_enable(): max_num_reqs_across_dp = ( self.vllm_config.scheduler_config.max_num_seqs * self.runner.uniform_decode_query_len ) # It is necessary to evaluate the case where num_indices becomes large # in the context of the dummy‑run accompaniment of p‑eagle. if num_indices > max_num_reqs_across_dp: ori_token_indices_to_sample = token_indices_to_sample else: ori_token_indices_to_sample = None pcp_manager = getattr(self.runner, "pcp_manager", None) if pcp_manager is not None and self.pcp_size > 1: # remove graph padding before all_gather hidden_states = pcp_manager.get_restore_hidden_states( hidden_states, num_input_tokens=num_input_tokens, ) if self.method == "mtp": last_hidden_states = hidden_states else: # eagle and eagle3 need allgather last_hidden_states. last_hidden_states = pcp_manager.get_restore_hidden_states( last_hidden_states, num_input_tokens=num_input_tokens, ) if lmhead_tp_enable(): token_indices_to_sample = nn.functional.pad( token_indices_to_sample, (0, max_num_reqs_across_dp - num_indices) ) sample_hidden_states = last_hidden_states[token_indices_to_sample] if get_ascend_config().enable_reduce_sample: if self.method in ("eagle3", "dflash", "mtp"): draft_token_ids = self.compute_draft_token_ids(sample_hidden_states) if lmhead_tp_enable(): draft_token_ids, token_indices_to_sample = self._align_tensor_and_indices( draft_token_ids, num_indices, token_indices_to_sample, ori_token_indices_to_sample, is_logits=False, ) else: logits = self.model.compute_logits(sample_hidden_states) if lmhead_tp_enable(): logits = get_lmhead_tp_group().all_to_all(logits) else: logits = self.model.model.logits_processor._gather_logits(logits) if lmhead_tp_enable(): logits, token_indices_to_sample = self._align_tensor_and_indices( logits, num_indices, token_indices_to_sample, ori_token_indices_to_sample, is_logits=True, ) draft_token_ids = logits.argmax(dim=-1) else: logits = self.model.compute_logits(sample_hidden_states) if lmhead_tp_enable(): logits, token_indices_to_sample = self._align_tensor_and_indices( logits, num_indices, token_indices_to_sample, ori_token_indices_to_sample, is_logits=True, ) draft_token_ids = logits.argmax(dim=-1) # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1 or self.parallel_drafting: # [batch_size, 1] return draft_token_ids.view(-1, self.num_speculative_tokens) if self.pcp_size > 1 and is_prefill: draft_token_ids_list = [] for _ in range(self.num_speculative_tokens): draft_token_ids_list.append(draft_token_ids) return torch.stack(draft_token_ids_list, dim=1) # The logits are split and then merged only when lmhead_tp_enable() is enabled. # As a result, the batch size length becomes the actual length 32. # However, when lmhead_tp_enable() is disabled, the batch size uses the length after padding. # To decouple the scenarios, a judgment is required. # That is, the batch size needs to be modified only when lmhead_tp_enable() is enabled. if lmhead_tp_enable() and self.method == "mtp": batch_size = draft_token_ids.shape[0] # Generate the remaining draft tokens. draft_token_ids_tensor = torch.zeros( (self.num_speculative_tokens, *draft_token_ids.shape), dtype=draft_token_ids.dtype, device=self.device ) draft_token_ids_tensor[0] = draft_token_ids if self.uses_mrope: positions = self.mrope_positions[:, token_indices_to_sample] else: positions = self.positions[token_indices_to_sample] hidden_states = hidden_states[token_indices_to_sample] token_indices_to_sample = self.arange[:batch_size] input_batch_size = num_input_tokens if (self.method == "mtp" or self.use_cuda_graph) else batch_size forward_context = get_forward_context() _EXTRA_CTX.num_tokens = input_batch_size _EXTRA_CTX.num_accept_tokens = batch_size for draft_index in range(self.num_speculative_tokens - 1): # Reset MOE layer index for each draft step iteration forward_context = get_forward_context() if forward_context is not None: forward_context.moe_layer_index = 0 # Update the inputs. # cast to int32 is crucial when eagle model is compiled. # tensor.argmax() returns int64 by default. input_ids = draft_token_ids_tensor[draft_index] 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. if self.uses_mrope: exceeds_max_model_len = positions[0] >= self.vllm_config.model_config.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 ) else: exceeds_max_model_len = positions >= self.vllm_config.model_config.max_model_len clamped_positions = torch.where(exceeds_max_model_len, 0, positions) # 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.view(batch_size, -1) if self.supports_mm_inputs: self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids) input_ids = self.input_ids[:input_batch_size] inputs_embeds = self.inputs_embeds[:input_batch_size] else: input_ids = self.input_ids[:input_batch_size] inputs_embeds = None # Run the model. # The lifecycle of `input_ids`, `positions`, `hidden_states` runs through all # speculative tokens' proposings. `model_input_ids`, `model_positions` and # `model_hidden_states` represent the speculative model inputs. model_input_ids = self.input_ids[:input_batch_size] model_positions = self._get_positions(input_batch_size) model_hidden_states = self.hidden_states[:input_batch_size] model_hidden_states, model_positions = self.maybe_pad_and_reduce(model_hidden_states, model_positions) forward_context.attn_metadata = ( multi_steps_attn_metadata[draft_index + 1] if multi_steps_attn_metadata else None ) model_kwargs = { "input_ids": model_input_ids, "positions": model_positions, "inputs_embeds": inputs_embeds, } if self.pass_hidden_states_to_model: model_kwargs["hidden_states"] = model_hidden_states ret_hidden_states = self.model(**model_kwargs) if not self.model_returns_tuple(): last_hidden_states = ret_hidden_states hidden_states = last_hidden_states else: last_hidden_states, hidden_states = ret_hidden_states last_hidden_states, model_positions, hidden_states = self.maybe_all_gather_and_unpad( last_hidden_states, model_positions, hidden_states ) num_indices = token_indices_to_sample.shape[0] if lmhead_tp_enable(): max_num_reqs_across_dp = ( self.vllm_config.scheduler_config.max_num_seqs * self.runner.uniform_decode_query_len ) token_indices_to_sample = nn.functional.pad( token_indices_to_sample, (0, max_num_reqs_across_dp - num_indices), ) sample_hidden_states = last_hidden_states[token_indices_to_sample] if get_ascend_config().enable_reduce_sample: if self.method in ("eagle3", "dflash", "mtp"): draft_token_ids = self.compute_draft_token_ids(sample_hidden_states) if lmhead_tp_enable() and num_indices < draft_token_ids.shape[0]: draft_token_ids = draft_token_ids[:num_indices] token_indices_to_sample = token_indices_to_sample[:num_indices] else: logits = self.model.compute_logits(sample_hidden_states) if lmhead_tp_enable(): logits = get_lmhead_tp_group().all_to_all(logits) else: logits = self.model.model.logits_processor._gather_logits(logits) if lmhead_tp_enable() and num_indices < logits.shape[0]: logits = logits[:num_indices] token_indices_to_sample = token_indices_to_sample[:num_indices] draft_token_ids = logits.argmax(dim=-1) else: logits = self.model.compute_logits(sample_hidden_states) if lmhead_tp_enable() and num_indices < logits.shape[0]: logits = logits[:num_indices] token_indices_to_sample = token_indices_to_sample[:num_indices] draft_token_ids = logits.argmax(dim=-1) # TODO(wenlong): get more than one token for tree attention hidden_states = hidden_states[:batch_size] draft_token_ids_tensor[draft_index + 1] = draft_token_ids # [batch_size, num_speculative_tokens] draft_token_ids = draft_token_ids_tensor.swapaxes(0, 1) return draft_token_ids def set_inputs_first_pass( self, target_token_ids: torch.Tensor, next_token_ids: torch.Tensor, target_positions: torch.Tensor, target_hidden_states: torch.Tensor, token_indices_to_sample: torch.Tensor | None, cad: CommonAttentionMetadata, num_rejected_tokens_gpu: torch.Tensor | None, req_scheduled_tokens=None, long_seq_metadata=None, num_prefill_reqs=0, num_decode_reqs=0, ) -> tuple[int, torch.Tensor, CommonAttentionMetadata, tuple[Any, Any] | None]: if not self.needs_extra_input_slots: # Default EAGLE pathway: no reshaping of input tensors needed. # Simply rotate the input ids and leave the positions unchanged, # Inserting the next token ids at the last slot in each request. if token_indices_to_sample is None: token_indices_to_sample = cad.query_start_loc[1:] - 1 num_tokens = target_token_ids.shape[0] # Shift the input ids by one token. # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3] self.input_ids[: num_tokens - 1] = target_token_ids[1:] # Replace the last token with the next token. # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4] self.input_ids[token_indices_to_sample] = next_token_ids assert self.runner is not None pcp_manager = getattr(self.runner, "pcp_manager", None) long_seq_args = None if pcp_manager is not None: first_pass_inputs = pcp_manager.prepare_spec_decode_first_pass_inputs( input_ids=self.input_ids[:num_tokens], target_positions=target_positions, target_hidden_states=target_hidden_states, token_indices_to_sample=token_indices_to_sample, common_attn_metadata=cad, long_seq_metadata=long_seq_metadata, req_scheduled_tokens=req_scheduled_tokens, req_ids=self.runner.input_batch.req_ids, logits_indices=self.runner.logits_indices, num_tokens=num_tokens, num_prefill_reqs=num_prefill_reqs, num_decode_reqs=num_decode_reqs, uses_mrope=self.uses_mrope, ) num_tokens = first_pass_inputs.num_tokens target_positions = first_pass_inputs.target_positions target_hidden_states = first_pass_inputs.target_hidden_states token_indices_to_sample = first_pass_inputs.token_indices_to_sample self.input_ids[:num_tokens].copy_(first_pass_inputs.input_ids) long_seq_args = first_pass_inputs.long_seq_args # copy inputs to buffer for cudagraph if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim == 0: target_positions = target_positions[0] self._set_positions(num_tokens, target_positions) self.hidden_states[:num_tokens] = target_hidden_states.view(num_tokens, -1) return num_tokens, token_indices_to_sample, cad, long_seq_args else: assert self.is_rejected_token_mask is not None assert self.is_masked_token_mask is not None # 1. # Call the CopyAndExpandEagleInputs AscendC operator to copy # input_ids and positions into the correct slots in the # preallocated buffers self.input_ids, self.positions. batch_size = cad.batch_size() total_num_input_tokens = target_token_ids.shape[0] total_num_output_tokens = total_num_input_tokens + (self.net_num_new_slots_per_request * batch_size) query_start_loc = cad.query_start_loc query_end_loc = cad.query_start_loc[1:] - 1 if num_rejected_tokens_gpu is not None: query_end_loc = query_end_loc - num_rejected_tokens_gpu ( out_input_ids, out_positions, out_is_rejected_token_mask, out_is_masked_token_mask, token_indices_to_sample, out_hidden_state_mapping, ) = torch.ops._C_ascend.npu_copy_and_expand_eagle_inputs( target_token_ids, target_positions.to(torch.int32), next_token_ids, query_start_loc, query_end_loc, 0, # padding_token_id self.parallel_drafting_token_id, self.extra_slots_per_request, self.pass_hidden_states_to_model, total_num_output_tokens, ) # Copy returned tensors into pre-allocated buffers self.input_ids[:total_num_output_tokens].copy_(out_input_ids) self.positions[:total_num_output_tokens].copy_(out_positions) self.is_rejected_token_mask[:total_num_output_tokens].copy_(out_is_rejected_token_mask) self.is_masked_token_mask[:total_num_output_tokens].copy_(out_is_masked_token_mask) if self.pass_hidden_states_to_model: assert self.parallel_drafting_hidden_state_tensor is not None self.hidden_states[out_hidden_state_mapping] = target_hidden_states # Use torch.where to avoid DtoH sync from boolean indexing mask = self.is_masked_token_mask[:total_num_output_tokens] torch.where( mask.unsqueeze(1), # type: ignore self.parallel_drafting_hidden_state_tensor, self.hidden_states[:total_num_output_tokens], out=self.hidden_states[:total_num_output_tokens], ) # 2. # Recompute the slot mapping based on the new positions and # rejection mask. # Use the first draft attention group's kv_cache_spec for block_size # (all draft layers share the same kv-cache group) assert len(self.draft_attn_groups) > 0 block_size = self.draft_attn_groups[0].kv_cache_spec.block_size new_slot_mapping = compute_new_slot_mapping( cad=cad, new_positions=self.positions[:total_num_output_tokens], is_rejected_token_mask=self.is_rejected_token_mask[:total_num_output_tokens], block_size=block_size, num_new_tokens=self.net_num_new_slots_per_request, max_model_len=self.max_model_len, ) # 3. Update the common attention metadata with the new (meta)data new_cad = extend_all_queries_by_N( cad, N=self.net_num_new_slots_per_request, arange=self.arange, new_slot_mapping=new_slot_mapping, ) # ``extend_all_queries_by_N`` adds N to every per-row GPU # ``seq_lens`` but cannot touch the host-side mirrors (it # only knows about upstream's deprecated ``_seq_lens_cpu`` # field; the Ascend subclass also has its own # ``seq_lens_cpu`` field, which would be silently stale # after the dataclass ``replace``). # # NPU attention backends (MLA, AscendAttention, SFA) read # those CPU mirrors as kernel input, so they MUST be in # sync with the GPU view. We compute the +N update on CPU # to avoid an extra GPU->CPU sync (which was the original # FIXME): every consumer that needs the post-extend value # already had a valid pre-extend mirror, so a CPU-only # ``+N`` keeps both in lock-step at zero device-side cost. N = self.net_num_new_slots_per_request if cad._seq_lens_cpu is not None: new_cad._seq_lens_cpu = cad._seq_lens_cpu + N elif cad.seq_lens_cpu is not None: # Parent field absent but Ascend subclass field set: # populate ``_seq_lens_cpu`` so upstream code paths # that prefer the parent field still get a fresh value. new_cad._seq_lens_cpu = cad.seq_lens_cpu + N if cad.seq_lens_cpu is not None: new_cad.seq_lens_cpu = cad.seq_lens_cpu + N return total_num_output_tokens, token_indices_to_sample, new_cad, None def model_returns_tuple(self) -> bool: if self.method == "mtp": # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final- # norm hidden, so its forward returns (logit_hidden, # recycle_hidden). Other MTP families return a single tensor. draft_model_config = getattr(self, "draft_model_config", None) hf_config = getattr(draft_model_config, "hf_config", None) architectures = getattr(hf_config, "architectures", []) or [] return "DeepSeekMTPModel" in architectures return self.method not in ("mtp", "draft_model", "dflash") def attn_update_stack_num_spec_norm( self, # `draft_index` must start from `1`, no `0` draft_index, old_attn_metadata, old_common_metadata, batch_size, input_batch_size, used_update_positions, aclgraph_runtime_mode, ori_seq_len=None, ori_seq_len_cpu=None, slot_indices=None, mtp_slot_mapping=None, attn_group=None, ): assert draft_index > 0 assert attn_group is not None, "vllm-ascend v0.17.0rc1 requires attn_group" common_attn_metadata = self.shallow_copy_metadata(old_common_metadata) if draft_index == 1: if aclgraph_runtime_mode == CUDAGraphMode.FULL: common_attn_metadata.num_reqs = input_batch_size common_attn_metadata.block_table_tensor = self._adjust_tensor( common_attn_metadata.block_table_tensor, input_batch_size ) common_attn_metadata.seq_lens = self._adjust_tensor(common_attn_metadata.seq_lens, input_batch_size) common_attn_metadata.seq_lens_cpu = self._adjust_tensor( common_attn_metadata.seq_lens_cpu, input_batch_size ) if common_attn_metadata._seq_lens_cpu is not None: common_attn_metadata._seq_lens_cpu = self._adjust_tensor( common_attn_metadata._seq_lens_cpu, input_batch_size ) if common_attn_metadata.num_computed_tokens_cpu is not None: common_attn_metadata.num_computed_tokens_cpu = self._adjust_tensor( common_attn_metadata.num_computed_tokens_cpu, input_batch_size ) common_attn_metadata.query_start_loc = self.arange[: input_batch_size + 1] common_attn_metadata.query_start_loc_cpu = torch.from_numpy( self.token_arange_np[: input_batch_size + 1] ).clone() else: common_attn_metadata.query_start_loc = self.arange[: batch_size + 1] common_attn_metadata.query_start_loc_cpu = torch.from_numpy( self.token_arange_np[: batch_size + 1] ).clone() common_attn_metadata.num_actual_tokens = batch_size common_attn_metadata.max_query_len = 1 common_attn_metadata.decode_token_per_req = 1 common_attn_metadata.attn_state = ( AscendAttentionState.SpecDecoding if self.method == "mtp" else AscendAttentionState.ChunkedPrefill ) common_attn_metadata.graph_pad_size = -1 common_attn_metadata.num_input_tokens = input_batch_size # The loop part used_update_positions += 1 # Clone the data so that when calculating the data at position 2 and position 3 # in the merged graph, it does not affect position 1 # FIXME(lilinsiman) common_attn_metadata.seq_lens = common_attn_metadata.seq_lens.clone() if common_attn_metadata.seq_lens_cpu is not None: common_attn_metadata.seq_lens_cpu = common_attn_metadata.seq_lens_cpu.clone() if common_attn_metadata._seq_lens_cpu is not None: common_attn_metadata._seq_lens_cpu = common_attn_metadata._seq_lens_cpu.clone() if common_attn_metadata.num_computed_tokens_cpu is not None: common_attn_metadata.num_computed_tokens_cpu = common_attn_metadata.num_computed_tokens_cpu.clone() common_attn_metadata.positions = common_attn_metadata.positions.clone() # 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. if self.uses_mrope: exceeds_max_model_len = used_update_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(used_update_positions), used_update_positions ) else: exceeds_max_model_len = used_update_positions >= self.max_model_len clamped_positions = torch.where(exceeds_max_model_len, 0, used_update_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[:batch_size] += 1 # For the requests that exceed the max model length, we set the # sequence length to 1 to minimize their overheads in attention. exceeds_mask = common_attn_metadata.seq_lens[:batch_size] > self.max_model_len common_attn_metadata.seq_lens[:batch_size].masked_fill_(exceeds_mask, 1) if common_attn_metadata.seq_lens_cpu is not None: common_attn_metadata.seq_lens_cpu[:batch_size] = common_attn_metadata.seq_lens_cpu[:batch_size] + 1 exceeds_mask_cpu = common_attn_metadata.seq_lens_cpu[:batch_size] > self.max_model_len common_attn_metadata.seq_lens_cpu[:batch_size].masked_fill_(exceeds_mask_cpu, 1) if common_attn_metadata._seq_lens_cpu is not None: common_attn_metadata._seq_lens_cpu[:batch_size] = common_attn_metadata._seq_lens_cpu[:batch_size] + 1 exceeds_mask_internal_cpu = common_attn_metadata._seq_lens_cpu[:batch_size] > self.max_model_len common_attn_metadata._seq_lens_cpu[:batch_size].masked_fill_(exceeds_mask_internal_cpu, 1) if common_attn_metadata.num_computed_tokens_cpu is not None: common_attn_metadata.num_computed_tokens_cpu[:batch_size] += 1 if self.uses_mrope: common_attn_metadata.positions[:batch_size].copy_(clamped_positions[0]) else: common_attn_metadata.positions[:batch_size].copy_(clamped_positions) pcp_manager = getattr(self.runner, "pcp_manager", None) if pcp_manager is not None: kv_cache_spec = getattr(attn_group, "kv_cache_spec", self.draft_attn_groups[0].kv_cache_spec) # update slot_mapping slot_indices += self.pcp_size slot_mapping = mtp_slot_mapping[slot_indices] self.slot_mapping_group[draft_index][: batch_size * self.pcp_size] = slot_mapping self.slot_mapping_group[draft_index][batch_size * self.pcp_size :].fill_(PADDING_SLOT_ID) common_attn_metadata.slot_mapping = self.slot_mapping_group[draft_index] else: # NOTE: In vllm, `block_size = attn_metadata_builder.kv_cache_spec.block_size`. # However, in vllm-ascend, the above value can be multiple of `kernel_block_size`, # which is not correct for computing `slot_mapping` below. if self.has_gdn: block_size = self.kernel_block_size else: block_size = self.block_size # Compute the slot mapping. if self.uses_mrope: block_numbers = clamped_positions[0] // block_size else: block_numbers = clamped_positions // block_size block_ids = old_common_metadata.block_table_tensor.gather(dim=1, index=block_numbers.view(-1, 1)) block_ids = block_ids.view(-1) if self.uses_mrope: slot_mapping = block_ids * block_size + clamped_positions[0] % block_size else: 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. slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID) self.slot_mapping_group[draft_index][: slot_mapping.shape[0]].copy_(slot_mapping.to(torch.int32)) self.slot_mapping_group[draft_index][slot_mapping.shape[0] :].fill_(PADDING_SLOT_ID) # Set the address of the attn_metadata.slot_mapping to the self.slot_mapping_group[idx] common_attn_metadata.slot_mapping = self.slot_mapping_group[draft_index] self.seq_lens_group[draft_index][: common_attn_metadata.seq_lens.shape[0]].copy_(common_attn_metadata.seq_lens) self.seq_lens_group[draft_index][common_attn_metadata.seq_lens.shape[0] :].fill_(0) common_attn_metadata.seq_lens = self.seq_lens_group[draft_index][: common_attn_metadata.seq_lens.shape[0]] self.query_start_loc_group[draft_index][: common_attn_metadata.query_start_loc.shape[0]].copy_( common_attn_metadata.query_start_loc ) self.query_start_loc_group[draft_index][common_attn_metadata.query_start_loc.shape[0] :].fill_(0) common_attn_metadata.query_start_loc = self.query_start_loc_group[draft_index][ : common_attn_metadata.query_start_loc.shape[0] ] attn_metadata_builder = attn_group.get_metadata_builder() extra_attn_metadata_args = {} if self.use_compress: extra_attn_metadata_args = dict( prefill_ratio_to_sas_metadata=dict(), decode_ratio_to_sas_metadata=dict(), common_ratio_to_sas_metadata=dict(), block_size=self.draft_attn_groups[0].kv_cache_spec.block_size, ) attn_metadata = attn_metadata_builder.build_for_drafting( common_attn_metadata, draft_index, **extra_attn_metadata_args, ) if pcp_manager is not None: pcp_manager.update_spec_decode_drafting_cp_metadata( attn_metadata=attn_metadata, kv_cache_spec=kv_cache_spec, seq_lens=ori_seq_len, draft_index=draft_index, seq_lens_cpu=ori_seq_len_cpu, attn_metadata_builder=attn_metadata_builder, ) return common_attn_metadata, attn_metadata def prepare_next_token_ids_padded( self, sampled_token_ids: torch.Tensor, requests: dict[str, CachedRequestState], gpu_input_batch: InputBatch, discard_request_indices: torch.Tensor, num_discarded_requests: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ This function is used to prepare the inputs for speculative decoding. It calculates the next token ids and the number of valid sampled tokens for each request, considering the "discarded" requests whose next token is not sampled and comes from `request.get_token_id()` instead. It also accounts for the rejected tokens in `sampled_token_ids`. This function must use device functions to operate on the inputs, and should not introduce any blocking CPU-GPU synchronization. """ # TODO(Ben): Combine this into a custom fused kernel # 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(seq_lens_list[i]) for i in range(num_reqs)] ) self.backup_next_token_ids.copy_to_gpu(num_reqs) # Mask out the sampled tokens indices that should not be sampled. discard_sampled_tokens_req_indices = discard_request_indices[:num_discarded_requests] valid_sampled_token_ids_gpu = sampled_token_ids.clone() valid_sampled_token_ids_gpu = DeviceOperator.index_fill( valid_sampled_token_ids_gpu, 0, discard_sampled_tokens_req_indices, -1, ) # Generate a mask for all valid tokens within those requests valid_mask = (valid_sampled_token_ids_gpu != -1) & (valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size) # Count the number of valid tokens in each request valid_sampled_tokens_count = valid_mask.sum(dim=1) # Get the rightmost valid index per row last_valid_indices = valid_sampled_tokens_count - 1 last_valid_indices_safe = torch.clamp(last_valid_indices, min=0) # Get last valid token from each row # (assume undefined state where there is no valid token) selected_tokens = torch.gather(valid_sampled_token_ids_gpu, 1, last_valid_indices_safe.unsqueeze(1)).squeeze(1) # Use last token if valid, pre-computed backup if not batch_size = valid_sampled_token_ids_gpu.shape[0] next_token_ids = torch.where( last_valid_indices != -1, selected_tokens, self.backup_next_token_ids.gpu[:batch_size], ) return next_token_ids, valid_sampled_tokens_count def prepare_inputs( self, common_attn_metadata: CommonAttentionMetadata, sampled_token_ids: list[list[int]], num_draft_tokens: list[int], ) -> tuple[CommonAttentionMetadata, torch.Tensor]: """ This function is used to prepare the inputs for speculative decoding. It updates to the common_attn_metadata to account for the rejected tokens (and newly sampled tokens). It also returns the token indices of the tokens that should be fed to the speculator. """ # E.g. # common_attn_metadata.query_start_loc{_cpu}: # [0, q1, q1 + q2, q1 + q2 + q3] # common_attn_metadata.seq_lens{_cpu}: [s1, s2, s3] # num_rejected_tokens: [n1, n2, n3] # This function computes the intermediate values: # num_tokens_per_req: [q1 - n1, q2 - n2, q3 - n3] # And returns: # common_attn_metadata.query_start_loc{_cpu}: # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3] # common_attn_metadata.seq_lens{_cpu}: # [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1] # token_indices: [0, 1, ..., q1 - n1 - 1, # q1, q1 + 1, ..., q1 + q2 - n2 - 1, # q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1] num_actual_reqs = len(num_draft_tokens) num_rejected_tokens = [ n + 1 - len(sampled_token_ids[i]) if n > 0 else 0 for i, n in enumerate(num_draft_tokens) ] num_rejected_tokens = torch.tensor(num_rejected_tokens, dtype=torch.int32) device = common_attn_metadata.query_start_loc.device query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[: num_actual_reqs + 1] # Prefer the upstream-canonical ``_seq_lens_cpu``; fall back to the # Ascend subclass field. In async-spec mode the model runner only # populates ``_seq_lens_cpu`` (optimistic_seq_lens_cpu) and leaves # ``seq_lens_cpu`` as None, so an unguarded read here would crash. if common_attn_metadata._seq_lens_cpu is not None: seq_lens_cpu = common_attn_metadata._seq_lens_cpu[:num_actual_reqs] else: seq_lens_cpu = common_attn_metadata.seq_lens_cpu[:num_actual_reqs] new_seq_lens_cpu = seq_lens_cpu - num_rejected_tokens # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3] new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] # [q1, q2, q3] -> [q1 - n1, q2 - n2, q3 - n3] new_num_tokens_per_req = new_query_len_per_req - num_rejected_tokens new_num_tokens_per_req_np = new_num_tokens_per_req.numpy() # [q1 - n1, q2 - n2, q3 - n3] -> # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3] new_query_start_loc_cpu = torch.zeros( query_start_loc_cpu.shape, dtype=torch.int32, pin_memory=is_pin_memory_available(), ) new_query_start_loc_np = new_query_start_loc_cpu.numpy() np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:]) total_num_tokens = new_query_start_loc_np[-1] # Example assuming num_tokens_per_req_np = [2, 4, 3] # this implies that `new_query_start_locs` is: # [0, 2, 6, 9] -> # [0, 0, 2, 2, 2, 2, 6, 6, 6] # _r1_ ____r2____ ___r3__ new_query_start_locs_expanded = np.repeat(new_query_start_loc_np[:-1], new_num_tokens_per_req_np) # [0, 1, 2, 3, 4, 5, 6, 7, 8] -> # [0, 1, 0, 1, 2, 3, 0, 1, 2] # _r1_ ____r2____ ___r3__ token_offsets = self.token_arange_np[:total_num_tokens] - new_query_start_locs_expanded # Expand starting positions to match token pattern # [0, q1, q1 + q2] -> # [0, 0, q1, q1, q1, q1, q1 + q2, q1 + q2, q1 + q2] # _r1_ _____r2_______ ___________r3____________ old_query_start_locs_expanded = np.repeat(query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np) # Final token indices are: # [0, 1, // req 1 # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2 # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3 token_indices_np = token_offsets + old_query_start_locs_expanded token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True) common_attn_metadata.slot_mapping[: token_indices.shape[0]].copy_( common_attn_metadata.slot_mapping[token_indices] ) common_attn_metadata.slot_mapping[token_indices.shape[0] :].fill_(-1) # NOTE: Currently positions and seq_lens are not used in attn forward # so we do not need to fixed them. But if they are used in the future, # we should fixed them. # Mirror ``new_seq_lens_cpu`` into the upstream-canonical # ``_seq_lens_cpu`` slot so consumers preferring the parent field # (e.g. attention_cp builder) see the rejection-adjusted value. spec_common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True), query_start_loc_cpu=new_query_start_loc_cpu, seq_lens=new_seq_lens_cpu.to(device, non_blocking=True), seq_lens_cpu=new_seq_lens_cpu, _seq_lens_cpu=new_seq_lens_cpu, num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu, _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, seq_lens_cpu_upper_bound=new_seq_lens_cpu, num_reqs=common_attn_metadata.num_reqs, num_actual_tokens=total_num_tokens, num_input_tokens=common_attn_metadata.num_input_tokens, max_query_len=new_query_len_per_req.max().item(), block_table_tensor=common_attn_metadata.block_table_tensor, slot_mapping=common_attn_metadata.slot_mapping, actual_seq_lengths_q=self.runner.actual_seq_lengths_q, positions=common_attn_metadata.positions[token_indices], positions_cpu=common_attn_metadata.positions_cpu[token_indices] if common_attn_metadata.positions_cpu is not None else None, attn_state=self.runner.attn_state, decode_token_per_req=self.runner.decode_token_per_req, is_prefilling=common_attn_metadata.is_prefilling, max_seq_len=0, ) return spec_common_attn_metadata, token_indices def prepare_inputs_padded( self, common_attn_metadata: CommonAttentionMetadata, spec_decode_metadata: SpecDecodeMetadata, valid_sampled_tokens_count: torch.Tensor, ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor, torch.Tensor]: """ This function is used to prepare the inputs for speculative decoding It updates the common_attn_metadata for speculative decoding, but does not consider the rejected tokens. Instead, all tokens are included as inputs to the speculator, with the rejected tokens used as padding and filtered out later by `token_indices_to_sample`. No blocking CPU operations should be introduced in this function. """ if HAS_TRITON: num_reqs = common_attn_metadata.num_reqs device = valid_sampled_tokens_count.device token_indices_to_sample = torch.empty((num_reqs,), dtype=torch.int32, device=device) num_rejected_tokens_gpu = torch.empty((num_reqs,), dtype=torch.int32, device=device) num_blocks_needed = triton.cdiv(num_reqs, _PREPARE_INPUTS_BLOCK_SIZE) num_vector_core = get_vectorcore_num() grid_size = min(num_blocks_needed, num_vector_core) grid = (grid_size,) prepare_inputs_padded_kernel[grid]( spec_decode_metadata.cu_num_draft_tokens, valid_sampled_tokens_count, common_attn_metadata.query_start_loc, token_indices_to_sample, num_rejected_tokens_gpu, num_reqs, BLOCK_SIZE=_PREPARE_INPUTS_BLOCK_SIZE, ) else: num_draft_tokens_gpu = torch.cat( [ spec_decode_metadata.cu_num_draft_tokens[0:1], spec_decode_metadata.cu_num_draft_tokens[1:] - spec_decode_metadata.cu_num_draft_tokens[:-1], ] ) num_rejected_tokens_gpu = torch.where( num_draft_tokens_gpu > 0, num_draft_tokens_gpu + 1 - valid_sampled_tokens_count, torch.zeros_like(num_draft_tokens_gpu), ) token_indices_to_sample = common_attn_metadata.query_start_loc[1:] - 1 - num_rejected_tokens_gpu query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] total_num_tokens = query_start_loc_cpu[-1].item() token_indices = self.arange[:total_num_tokens] # NOTE: Currently positions and seq_lens are not used in attn forward # so we do not need to fixed them. But if they are used in the future, # we should fixed them. # ``prepare_inputs_padded`` does not change ``seq_lens`` (rejected # tokens are kept as padding and filtered out later). Pass through # both the subclass ``seq_lens_cpu`` field and the upstream-canonical # ``_seq_lens_cpu`` field unchanged. In async-spec mode only the # latter is populated (subclass field is None to signal "GPU is # authoritative"); dropping ``_seq_lens_cpu`` here causes downstream # backends (e.g. attention_cp) to crash on a None subscript. spec_common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=common_attn_metadata.query_start_loc, query_start_loc_cpu=query_start_loc_cpu, seq_lens_cpu=common_attn_metadata.seq_lens_cpu, _seq_lens_cpu=common_attn_metadata._seq_lens_cpu, seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound, num_reqs=common_attn_metadata.num_reqs, num_actual_tokens=common_attn_metadata.num_actual_tokens if self.pcp_size > 1 else total_num_tokens, num_input_tokens=common_attn_metadata.num_input_tokens, max_query_len=new_query_len_per_req.max().item(), actual_seq_lengths_q=self.runner.actual_seq_lengths_q, block_table_tensor=common_attn_metadata.block_table_tensor, slot_mapping=common_attn_metadata.slot_mapping, positions=common_attn_metadata.positions, positions_cpu=common_attn_metadata.positions_cpu, attn_state=self.runner.attn_state, decode_token_per_req=self.runner.decode_token_per_req, num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu, _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, seq_lens=common_attn_metadata.seq_lens, is_prefilling=common_attn_metadata.is_prefilling, max_seq_len=0, ) return spec_common_attn_metadata, token_indices, token_indices_to_sample, num_rejected_tokens_gpu # update full-graph params for one spec token def _update_full_graph_params(self, forward_context, num_tokens, draft_attn_metadatas=None): assert len(self.draft_attn_groups) > 0 attn_backend = self.draft_attn_groups[0].backend update_full_graph_params( attn_backend, self.update_stream, forward_context, num_tokens, self.vllm_config, self.vllm_config.speculative_config, draft_attn_metadatas=draft_attn_metadatas, ) # adjusting tensor into desired size def _adjust_tensor(self, tensor, desired_size): pad_size = desired_size - tensor.shape[0] if pad_size > 0: pad = [0] * (2 * tensor.dim() - 1) + [pad_size] tensor = F.pad(tensor, pad, mode="constant", value=0) else: tensor = tensor[:desired_size] return tensor def maybe_pad_and_reduce( self, hidden_states: torch.Tensor, positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if self.method == "mtp": if _EXTRA_CTX.flash_comm_v1_enabled and not self.is_multimodal_model: hidden_states = torch.ops.vllm.maybe_pad_and_reduce(hidden_states) positions = positions.unsqueeze(-1) positions = torch.ops.vllm.maybe_pad_and_reduce(positions) positions = positions.squeeze(-1) else: if _EXTRA_CTX.flash_comm_v1_enabled: hidden_states = split_inputs_tp_to_sp(hidden_states, hidden_states) return hidden_states, positions def maybe_all_gather_and_unpad( self, last_hidden_states: torch.Tensor, positions: torch.Tensor, hidden_states: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: if self.method == "mtp": if self.enable_shared_expert_dp: last_hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad( last_hidden_states.contiguous(), True ) # in mm model, positions not need allgather, because it not reduced before(see maybe_pad_and_reduce()) if not self.is_multimodal_model: positions = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(positions.contiguous(), True) if hidden_states is not None: hidden_states = last_hidden_states else: if _EXTRA_CTX.flash_comm_v1_enabled: last_hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad( last_hidden_states.contiguous(), True ) if hidden_states is not None: hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(hidden_states.contiguous(), True) return last_hidden_states, positions, hidden_states # In the context of the dummy‑run accompaniment of p‑eagle, when num_indices becomes large, # enabling the LM head feature causes token_indices_to_sample to switch from padding to trimming. # The trimmed length may not be an integer multiple of the speculative length, # in which case padding is required to restore it to the original length. def _align_tensor_and_indices( self, tensor, num_indices, token_indices_to_sample, ori_token_indices_to_sample, is_logits=False, ): """ Align the tensor (either draft_token_ids or logits) and token_indices_to_sample to the length specified by num_indices. Args: tensor: The tensor to be aligned (draft_token_ids or logits) num_indices: The target length token_indices_to_sample: The current index tensor ori_token_indices_to_sample: The original index tensor (used for restoration) is_logits: Whether the tensor is logits (affects the padding dimension and padding value) Returns: The adjusted tensor and token_indices_to_sample """ if tensor.shape[0] == num_indices: return tensor, token_indices_to_sample if tensor.shape[0] > num_indices: # Trim to the target length. tensor = tensor[:num_indices] token_indices_to_sample = token_indices_to_sample[:num_indices] else: # Padding to the target length. pad_size = num_indices - tensor.shape[0] if is_logits: # logits: shape [seq_len, vocab_size], Padding at the end of the seq dimension. tensor = nn.functional.pad(tensor, (0, 0, 0, pad_size), value=-1e9) else: # draft_token_ids: shape [seq_len], Padding at the end tensor = nn.functional.pad(tensor, (0, pad_size)) token_indices_to_sample = ori_token_indices_to_sample return tensor, token_indices_to_sample