2096 lines
103 KiB
Python
2096 lines
103 KiB
Python
# 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
|