831 lines
38 KiB
Python
831 lines
38 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from collections.abc import Callable
|
|
from functools import partial
|
|
from typing import Any
|
|
|
|
import torch
|
|
from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config, replace
|
|
from vllm.forward_context import BatchDescriptor, get_forward_context
|
|
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
|
from vllm.model_executor.models.utils import get_draft_quant_config
|
|
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec, UniformTypeKVCacheSpecs
|
|
from vllm.v1.sample.metadata import SamplingMetadata
|
|
from vllm.v1.spec_decode.llm_base_proposer import compute_probs_and_sample_next_token
|
|
from vllm.v1.spec_decode.utils import PADDING_SLOT_ID
|
|
from vllm.v1.worker.utils import AttentionGroup
|
|
|
|
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_v1 import AscendAttentionState
|
|
from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
|
|
from vllm_ascend.distributed.parallel_state import get_lmhead_tp_group
|
|
from vllm_ascend.spec_decode.eagle_proposer import AscendEagleProposer
|
|
from vllm_ascend.utils import lmhead_tp_enable
|
|
from vllm_ascend.worker.utils import copy_snapshot_to_gpu
|
|
|
|
|
|
class AscendStep3p5MTPProposer(AscendEagleProposer):
|
|
"""Step3.5 MTP proposer with per-MTP-layer independent KV cache groups.
|
|
|
|
Each Step3.5 MTP layer owns a different attention/KV-cache group. The
|
|
generic Ascend MTP proposer assumes one draft KV-cache group and builds one
|
|
attention metadata object for all draft layers. Step3.5 therefore keeps the
|
|
whole propose flow in this subclass: it builds per-group metadata for all MTP
|
|
layers, reuses that step0 metadata across layer calls, and refreshes the
|
|
metadata tensors after each in-graph input update.
|
|
"""
|
|
|
|
_runnable: Callable
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
runner=None,
|
|
) -> None:
|
|
super().__init__(vllm_config, device, runner=runner)
|
|
# Per KV cache group block tables / slot mappings captured from the
|
|
# model runner each step (see set_per_group_attn_metadata).
|
|
self._per_group_block_tables: dict[int, torch.Tensor] = {}
|
|
self._per_group_slot_mappings: dict[int, torch.Tensor] = {}
|
|
# Slot-mapping buffers for additional KV cache groups. FULL graph
|
|
# capture records tensor addresses for reshape/cache, so every group
|
|
# that can appear in the shared step0 metadata needs a persistent
|
|
# buffer.
|
|
self._group_slot_buffers: dict[int, torch.Tensor] = {}
|
|
|
|
def set_per_group_attn_metadata(
|
|
self,
|
|
gid: int,
|
|
block_table: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
) -> None:
|
|
self._per_group_block_tables[gid] = block_table
|
|
self._per_group_slot_mappings[gid] = slot_mapping
|
|
|
|
def _seed_graph_capture_per_group_metadata(
|
|
self,
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
) -> None:
|
|
if self.runner is None or not self.draft_attn_groups:
|
|
return
|
|
|
|
block_tables = self.runner.input_batch.block_table
|
|
for attn_group in self.draft_attn_groups:
|
|
gid = attn_group.kv_cache_group_id
|
|
try:
|
|
block_table = block_tables[gid]
|
|
except (IndexError, KeyError):
|
|
continue
|
|
block_table_tensor = block_table.get_device_tensor()
|
|
if num_reqs > 0:
|
|
block_table_tensor = block_table_tensor[:num_reqs]
|
|
slot_mapping = block_table.slot_mapping.gpu[:num_tokens]
|
|
self.set_per_group_attn_metadata(gid, block_table_tensor, slot_mapping)
|
|
|
|
def _slot_mapping_buffer_for_group(
|
|
self,
|
|
attn_group: AttentionGroup,
|
|
) -> torch.Tensor:
|
|
gid = attn_group.kv_cache_group_id
|
|
if gid == self.kv_cache_gid:
|
|
return self.slot_mapping_group[0]
|
|
buf = self._group_slot_buffers.get(gid)
|
|
if buf is None:
|
|
buf = torch.zeros_like(self.slot_mapping_group[0])
|
|
self._group_slot_buffers[gid] = buf
|
|
return buf
|
|
|
|
def _common_attn_metadata_for_group(
|
|
self,
|
|
common_attn_metadata: CommonAttentionMetadata,
|
|
attn_group: AttentionGroup,
|
|
) -> CommonAttentionMetadata:
|
|
group_common_attn_metadata = self.shallow_copy_metadata(common_attn_metadata)
|
|
gid = attn_group.kv_cache_group_id
|
|
|
|
block_table = self._per_group_block_tables.get(gid)
|
|
if block_table is not None:
|
|
group_common_attn_metadata.block_table_tensor = block_table[: group_common_attn_metadata.num_reqs]
|
|
|
|
slot_mapping = self._per_group_slot_mappings.get(gid)
|
|
if slot_mapping is None:
|
|
return group_common_attn_metadata
|
|
slot_mapping_buffer = self._slot_mapping_buffer_for_group(attn_group)
|
|
slot_mapping_len = slot_mapping.shape[0]
|
|
if slot_mapping.data_ptr() != slot_mapping_buffer.data_ptr():
|
|
slot_mapping_buffer[:slot_mapping_len].copy_(slot_mapping.to(torch.int32))
|
|
slot_mapping_buffer[slot_mapping_len:].fill_(PADDING_SLOT_ID)
|
|
slot_mapping_view = slot_mapping_buffer[: group_common_attn_metadata.num_actual_tokens]
|
|
self._per_group_slot_mappings[gid] = slot_mapping_view
|
|
group_common_attn_metadata.slot_mapping = slot_mapping_view
|
|
|
|
return group_common_attn_metadata
|
|
|
|
def _build_step_attn_metadatas(
|
|
self,
|
|
common_attn_metadata: CommonAttentionMetadata,
|
|
*,
|
|
graph_capture: bool = False,
|
|
) -> tuple[list[Any], list[dict[str, Any]]]:
|
|
"""Build base-proposer-style per-step metadata for Step3.5 MTP.
|
|
|
|
The full-window path still reuses the same logical window for every MTP
|
|
layer, but graph-param update/replay expects draft metadata to be
|
|
indexed by speculative step. Each Step3.5 MTP attention group maps to
|
|
one draft step/KV-cache group, so return:
|
|
|
|
``[{layer0: meta0}, {layer1: meta1}, {layer2: meta2}]``
|
|
|
|
instead of one dict containing all MTP layers.
|
|
"""
|
|
per_group_attn_metadata: list[Any] = []
|
|
multi_steps_attn_metadata: list[dict[str, Any]] = []
|
|
extra_attn_metadata_args: dict[str, Any] = {}
|
|
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,
|
|
)
|
|
|
|
for attn_group in self.draft_attn_groups:
|
|
group_common_attn_metadata = self._common_attn_metadata_for_group(common_attn_metadata, attn_group)
|
|
builder = attn_group.get_metadata_builder()
|
|
if graph_capture:
|
|
attn_metadata = builder.build_for_graph_capture(
|
|
group_common_attn_metadata,
|
|
AscendAttentionState.SpecDecoding,
|
|
)
|
|
else:
|
|
attn_metadata = builder.build(
|
|
0,
|
|
group_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
|
|
per_group_attn_metadata.append(attn_metadata)
|
|
per_step_attn_metadata: dict[str, Any] = {}
|
|
for layer_name in attn_group.layer_names:
|
|
per_step_attn_metadata[layer_name] = attn_metadata
|
|
multi_steps_attn_metadata.append(per_step_attn_metadata)
|
|
return per_group_attn_metadata, multi_steps_attn_metadata
|
|
|
|
def _sample_draft_tokens_for_step(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
sampling_metadata: SamplingMetadata,
|
|
spec_step_idx: int,
|
|
num_indices: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""GPU Step3.5 sampling semantics with Ascend TP/reduce-sample paths."""
|
|
logits: torch.Tensor | None = None
|
|
if get_ascend_config().enable_reduce_sample and self.method == "mtp":
|
|
if not hasattr(self.model.model, "compute_logits"):
|
|
draft_token_ids = self.compute_draft_token_ids(hidden_states)
|
|
if lmhead_tp_enable() and num_indices < draft_token_ids.shape[0]:
|
|
draft_token_ids = draft_token_ids[:num_indices]
|
|
return draft_token_ids, None
|
|
logits = self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx)
|
|
if lmhead_tp_enable():
|
|
logits = get_lmhead_tp_group().all_to_all(logits)
|
|
else:
|
|
logits = self.model.model.logits_processor._gather_logits(logits)
|
|
else:
|
|
logits = self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx)
|
|
|
|
if lmhead_tp_enable() and num_indices < logits.shape[0]:
|
|
logits = logits[:num_indices]
|
|
if not self._enable_probabilistic_draft_probs or sampling_metadata.all_greedy:
|
|
return logits.argmax(dim=-1), None
|
|
return compute_probs_and_sample_next_token(logits, sampling_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)
|
|
|
|
multi_steps_attn_metadata: list[dict[str, Any]] = []
|
|
if not self.use_cuda_graph:
|
|
aclgraph_runtime_mode = CUDAGraphMode.NONE
|
|
|
|
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,
|
|
)
|
|
|
|
batch_size = max(num_tokens // (self.num_speculative_tokens + 1), 1)
|
|
if is_profile:
|
|
batch_size = min(batch_size, self.runner.max_num_reqs)
|
|
|
|
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]
|
|
|
|
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_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[0].get_device_tensor()[:num_reqs],
|
|
slot_mapping=self.runner.input_batch.block_table[0].slot_mapping.gpu,
|
|
positions=self.runner.positions,
|
|
attn_state=self.runner.attn_state,
|
|
decode_token_per_req=self.runner.decode_token_per_req,
|
|
max_seq_len=0,
|
|
)
|
|
if self.pcp_size * self.dcp_size > 1:
|
|
common_attn_metadata.prefill_context_parallel_metadata = self.runner.pcp_manager.long_seq_metadata
|
|
|
|
common_attn_metadata = self.shallow_copy_metadata(common_attn_metadata)
|
|
common_attn_metadata.slot_mapping = self.slot_mapping_group[0]
|
|
common_attn_metadata.seq_lens = self.seq_lens_group[0][:num_reqs]
|
|
common_attn_metadata.query_start_loc = self.query_start_loc_group[0][: num_reqs + 1]
|
|
self._seed_graph_capture_per_group_metadata(num_reqs, num_tokens)
|
|
_, multi_steps_attn_metadata = self._build_step_attn_metadatas(common_attn_metadata, graph_capture=True)
|
|
|
|
model_positions = self._get_positions(num_tokens)
|
|
|
|
if self.supports_mm_inputs:
|
|
inputs_embeds = self.model.embed_input_ids(self.input_ids[:num_tokens])
|
|
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,
|
|
):
|
|
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],
|
|
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 _propose(
|
|
self,
|
|
target_token_ids: torch.Tensor,
|
|
target_positions: torch.Tensor,
|
|
target_hidden_states: torch.Tensor,
|
|
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:
|
|
self._last_draft_probs = None
|
|
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
|
|
|
|
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,
|
|
)
|
|
if self.pcp_size * self.dcp_size > 1:
|
|
assert long_seq_args is not None
|
|
assert self.runner is not None
|
|
|
|
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:
|
|
num_reqs_padded = self.runner._pad_query_start_loc_for_fia(
|
|
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.runner.query_start_loc.gpu[: num_reqs_padded + 1]
|
|
common_attn_metadata.query_start_loc_cpu = self.runner.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
|
|
)
|
|
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
|
|
)
|
|
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 self.pcp_size > 1:
|
|
pcp_allgather_restore_idx = (
|
|
common_attn_metadata.prefill_context_parallel_metadata.pcp_allgather_restore_idx
|
|
)
|
|
index = torch.arange(
|
|
pcp_allgather_restore_idx.shape[0],
|
|
device=pcp_allgather_restore_idx.device,
|
|
)
|
|
mask = (index % (self.pcp_size * self.decode_threshold)) >= self.decode_threshold
|
|
pcp_allgather_restore_idx[mask] = 0
|
|
self.runner.pcp_manager.pcp_allgather_restore_idx.gpu[: pcp_allgather_restore_idx.shape[0]] = (
|
|
pcp_allgather_restore_idx
|
|
)
|
|
self.runner.pcp_manager.pcp_allgather_restore_idx.gpu[pcp_allgather_restore_idx.shape[0] :] = 0
|
|
else:
|
|
num_reqs_padded = common_attn_metadata.num_reqs
|
|
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:
|
|
inputs_embeds = self.model.embed_input_ids(self.input_ids[:num_tokens])
|
|
self.inputs_embeds[:num_tokens] = inputs_embeds
|
|
inputs_embeds = self.inputs_embeds[:num_input_tokens]
|
|
else:
|
|
inputs_embeds = None
|
|
|
|
slot_mapping_len = common_attn_metadata.slot_mapping.shape[0]
|
|
self.slot_mapping_group[0][:slot_mapping_len].copy_(common_attn_metadata.slot_mapping)
|
|
self.slot_mapping_group[0][slot_mapping_len:].fill_(PADDING_SLOT_ID)
|
|
common_attn_metadata.slot_mapping = self.slot_mapping_group[0]
|
|
self._per_group_slot_mappings[self.kv_cache_gid] = self.slot_mapping_group[0][
|
|
: common_attn_metadata.num_actual_tokens
|
|
]
|
|
|
|
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
|
|
_, multi_steps_attn_metadata = self._build_step_attn_metadatas(common_attn_metadata)
|
|
attn_metadata_i = next(iter(multi_steps_attn_metadata[0].values()))
|
|
|
|
if not self.use_cuda_graph:
|
|
common_attn_metadata.block_table_tensor = common_attn_metadata.block_table_tensor.clone()
|
|
|
|
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,
|
|
):
|
|
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": attn_metadata_i.num_prefills,
|
|
}
|
|
run_draft = partial(self._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 _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:
|
|
"""Base MTP execution flow with Step3.5 step-aware layer/head selection."""
|
|
self._last_draft_probs = None
|
|
sampling_metadata = self.runner.input_batch.sampling_metadata
|
|
model_input_ids = self.input_ids[:num_input_tokens]
|
|
model_positions = self._get_positions(num_input_tokens)
|
|
model_kwargs = {
|
|
"input_ids": model_input_ids,
|
|
"positions": model_positions,
|
|
"inputs_embeds": inputs_embeds,
|
|
"spec_step_idx": 0,
|
|
}
|
|
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
|
|
model_kwargs["positions"] = model_positions
|
|
|
|
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 = torch.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]
|
|
draft_token_ids, draft_probs = self._sample_draft_tokens_for_step(
|
|
sample_hidden_states,
|
|
sampling_metadata,
|
|
spec_step_idx=0,
|
|
num_indices=num_indices,
|
|
)
|
|
if self.num_speculative_tokens == 1 or self.parallel_drafting:
|
|
if draft_probs is not None:
|
|
self._last_draft_probs = draft_probs.view(
|
|
-1, self.num_speculative_tokens, draft_probs.shape[-1]
|
|
).contiguous()
|
|
return draft_token_ids.view(-1, self.num_speculative_tokens)
|
|
|
|
if self.pcp_size * self.dcp_size > 1 and is_prefill:
|
|
draft_token_ids_list = [draft_token_ids for _ in range(self.num_speculative_tokens)]
|
|
return torch.stack(draft_token_ids_list, dim=1)
|
|
|
|
return self._run_window_draft_steps(
|
|
first_draft_token_ids=draft_token_ids,
|
|
first_draft_probs=draft_probs,
|
|
first_hidden_states=hidden_states,
|
|
num_input_tokens=num_input_tokens,
|
|
batch_size=batch_size,
|
|
token_indices_to_sample=token_indices_to_sample,
|
|
target_positions=target_positions,
|
|
num_tokens=num_tokens,
|
|
multi_steps_attn_metadata=multi_steps_attn_metadata,
|
|
inputs_embeds=inputs_embeds,
|
|
sampling_metadata=sampling_metadata,
|
|
)
|
|
|
|
def _roll_window_inputs_only(
|
|
self,
|
|
*,
|
|
prev_token_ids: torch.Tensor,
|
|
next_token_ids: torch.Tensor,
|
|
previous_hidden_states: torch.Tensor,
|
|
token_indices_to_sample: torch.Tensor,
|
|
num_tokens: int,
|
|
input_batch_size: int,
|
|
) -> torch.Tensor | None:
|
|
self.input_ids[: num_tokens - 1] = prev_token_ids[1:]
|
|
self.input_ids[token_indices_to_sample] = next_token_ids
|
|
self.hidden_states[:num_tokens] = previous_hidden_states.view(num_tokens, -1)
|
|
|
|
if self.supports_mm_inputs:
|
|
self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(self.input_ids[:num_tokens])
|
|
return self.inputs_embeds[:input_batch_size]
|
|
return None
|
|
|
|
def _run_window_draft_steps(
|
|
self,
|
|
*,
|
|
first_draft_token_ids: torch.Tensor,
|
|
first_draft_probs: torch.Tensor | None,
|
|
first_hidden_states: torch.Tensor,
|
|
num_input_tokens: int,
|
|
batch_size: int,
|
|
token_indices_to_sample: torch.Tensor,
|
|
target_positions: torch.Tensor,
|
|
num_tokens: int,
|
|
multi_steps_attn_metadata,
|
|
inputs_embeds,
|
|
sampling_metadata: SamplingMetadata,
|
|
) -> torch.Tensor:
|
|
"""Run the Step3.5 MTP full-window layer chain.
|
|
|
|
This is the canonical Step3.5 MTP draft path, matching the full-window
|
|
multi-layer proposer loop: every MTP layer reprocesses the whole draft
|
|
window so each independent KV-cache group is seeded over the full
|
|
window. Attention metadata is built once before this loop; layer-to-layer
|
|
state transition only rolls the model input buffers by shifting the
|
|
previous window and placing the newly drafted token at
|
|
token_indices_to_sample.
|
|
"""
|
|
draft_probs_list = None if first_draft_probs is None else [first_draft_probs]
|
|
draft_token_ids_list = [first_draft_token_ids]
|
|
full_hidden_states = first_hidden_states[:num_tokens]
|
|
input_batch_size = num_input_tokens
|
|
|
|
forward_context = get_forward_context()
|
|
_EXTRA_CTX.num_tokens = input_batch_size
|
|
_EXTRA_CTX.num_accept_tokens = batch_size
|
|
|
|
for draft_step in range(self.num_speculative_tokens - 1):
|
|
forward_context = get_forward_context()
|
|
if forward_context is not None:
|
|
forward_context.moe_layer_index = 0
|
|
|
|
spec_step_idx = draft_step + 1
|
|
prev_token_ids = self.input_ids[:num_tokens].clone()
|
|
next_token_ids = draft_token_ids_list[-1].int()
|
|
inputs_embeds = self._roll_window_inputs_only(
|
|
prev_token_ids=prev_token_ids,
|
|
next_token_ids=next_token_ids,
|
|
previous_hidden_states=full_hidden_states,
|
|
token_indices_to_sample=token_indices_to_sample,
|
|
num_tokens=num_tokens,
|
|
input_batch_size=input_batch_size,
|
|
)
|
|
|
|
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,
|
|
)
|
|
if forward_context is not None and multi_steps_attn_metadata:
|
|
if spec_step_idx >= len(multi_steps_attn_metadata):
|
|
raise AssertionError("Step3.5 MTP metadata must contain one entry per draft step")
|
|
forward_context.attn_metadata = multi_steps_attn_metadata[spec_step_idx]
|
|
|
|
model_kwargs = {
|
|
"input_ids": model_input_ids,
|
|
"positions": model_positions,
|
|
"inputs_embeds": inputs_embeds,
|
|
"spec_step_idx": spec_step_idx,
|
|
}
|
|
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 = ret_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]
|
|
sample_hidden_states = last_hidden_states[token_indices_to_sample]
|
|
draft_token_ids, draft_probs = self._sample_draft_tokens_for_step(
|
|
sample_hidden_states,
|
|
sampling_metadata,
|
|
spec_step_idx=spec_step_idx,
|
|
num_indices=num_indices,
|
|
)
|
|
if draft_probs is not None:
|
|
assert draft_probs_list is not None
|
|
draft_probs_list.append(draft_probs)
|
|
full_hidden_states = hidden_states[:num_tokens]
|
|
draft_token_ids_list.append(draft_token_ids)
|
|
|
|
draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
|
|
if draft_probs_list is not None:
|
|
self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous()
|
|
return draft_token_ids
|
|
|
|
# -- overrides matching the GPU Step3p5MTPProposer ----------------------
|
|
|
|
def _ensure_draft_layer_types_cover_mtp_layers(self) -> None:
|
|
hf_config = self.draft_model_config.hf_config
|
|
layer_types = getattr(hf_config, "layer_types", None)
|
|
num_hidden_layers = getattr(hf_config, "num_hidden_layers", None)
|
|
num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0)
|
|
|
|
if layer_types is None or num_hidden_layers is None or num_mtp_layers is None:
|
|
return
|
|
|
|
needed_num_layer_types = num_hidden_layers + num_mtp_layers
|
|
if len(layer_types) >= needed_num_layer_types:
|
|
return
|
|
|
|
hf_config.layer_types = list(layer_types) + ["sliding_attention"] * (needed_num_layer_types - len(layer_types))
|
|
|
|
def _create_draft_vllm_config(self) -> VllmConfig:
|
|
base = super()._create_draft_vllm_config()
|
|
# Transformers validates Step3.5 layer_types against num_hidden_layers
|
|
# and may leave only the base-layer entries. The MTP draft model builds
|
|
# layers at num_hidden_layers + k, so make the draft config cover those
|
|
# layer indices before the model is constructed.
|
|
self._ensure_draft_layer_types_cover_mtp_layers()
|
|
# Ascend ModelSlim keeps quant metadata in quant_model_description.json,
|
|
# not in config.json's quantization_config, so get_quant_config falls
|
|
# through to the hf_overrides branch. The verifier model_config has
|
|
# hf_overrides={}, but the draft_model_config leaves it None, which that
|
|
# branch rejects. Normalize to {} so the draft takes the verifier's path.
|
|
if not isinstance(getattr(self.draft_model_config, "hf_overrides", None), dict):
|
|
self.draft_model_config.hf_overrides = {}
|
|
return replace(
|
|
base,
|
|
model_config=self.draft_model_config,
|
|
quant_config=get_draft_quant_config(base),
|
|
)
|
|
|
|
def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
|
|
"""Step3.5 MTP draft layers may span multiple KV cache groups."""
|
|
return
|
|
|
|
def initialize_attn_backend(
|
|
self,
|
|
kv_cache_config: KVCacheConfig,
|
|
kernel_block_sizes: list[int] | None = None,
|
|
) -> None:
|
|
all_attn_layers = get_layers_from_vllm_config(
|
|
self.vllm_config,
|
|
AttentionLayerBase, # type: ignore[type-abstract]
|
|
)
|
|
|
|
layer_to_gid: dict[str, int] = {}
|
|
layer_to_spec: dict[str, KVCacheSpec] = {}
|
|
for gid, group in enumerate(kv_cache_config.kv_cache_groups):
|
|
group_spec = group.kv_cache_spec
|
|
for layer_name in group.layer_names:
|
|
layer_to_gid[layer_name] = gid
|
|
if isinstance(group_spec, UniformTypeKVCacheSpecs):
|
|
if layer_name in group_spec.kv_cache_specs:
|
|
layer_to_spec[layer_name] = group_spec.kv_cache_specs[layer_name]
|
|
else:
|
|
target_layer_name = getattr(
|
|
all_attn_layers.get(layer_name),
|
|
"kv_sharing_target_layer_name",
|
|
None,
|
|
)
|
|
if target_layer_name and target_layer_name in group_spec.kv_cache_specs:
|
|
layer_to_spec[layer_name] = group_spec.kv_cache_specs[target_layer_name]
|
|
else:
|
|
layer_to_spec[layer_name] = group_spec
|
|
else:
|
|
layer_to_spec[layer_name] = group_spec
|
|
|
|
attention_groups: dict[tuple[str, int], AttentionGroup] = {}
|
|
for layer_name in sorted(self._draft_attn_layer_names):
|
|
if layer_name not in layer_to_spec:
|
|
continue
|
|
attn_layer = all_attn_layers[layer_name]
|
|
attn_backend = attn_layer.get_attn_backend()
|
|
spec = layer_to_spec[layer_name]
|
|
gid = layer_to_gid[layer_name]
|
|
group_key = (attn_backend.full_cls_name(), gid)
|
|
|
|
if group_key not in attention_groups:
|
|
kernel_block_size = (
|
|
kernel_block_sizes[gid]
|
|
if kernel_block_sizes is not None and gid < len(kernel_block_sizes)
|
|
else None
|
|
)
|
|
attn_group = AttentionGroup(
|
|
backend=attn_backend,
|
|
layer_names=[layer_name],
|
|
kv_cache_spec=spec,
|
|
kv_cache_group_id=gid,
|
|
)
|
|
attn_group.create_metadata_builders(
|
|
self.vllm_config,
|
|
self.device,
|
|
kernel_block_size=kernel_block_size,
|
|
)
|
|
attention_groups[group_key] = attn_group
|
|
else:
|
|
attention_groups[group_key].layer_names.append(layer_name)
|
|
|
|
self.draft_attn_groups = list(attention_groups.values())
|
|
if self.draft_attn_groups:
|
|
self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id
|
|
self.block_size = self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size
|
|
else:
|
|
self.kv_cache_gid = 0
|
|
self.block_size = kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
|