EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Lint]Style: Convert vllm-ascend/ to ruff format(Batch #8) (#6129)

Browse Source 
### What this PR does / why we need it?
**Scope of Changes**:
| File Path |
| :--- |
| vllm_ascend/ops/\_\_init\_\_.py |
| vllm_ascend/ops/activation.py |
| vllm_ascend/ops/flashcomm2_oshard_manager.py |
| vllm_ascend/ops/layernorm.py |
| vllm_ascend/ops/mla.py |
| vllm_ascend/ops/mm_encoder_attention.py |
| vllm_ascend/ops/register_custom_ops.py |
| vllm_ascend/ops/vocab_parallel_embedding.py |
| vllm_ascend/ops/weight_prefetch.py |
| vllm_ascend/spec_decode/\_\_init\_\_.py |
| vllm_ascend/spec_decode/eagle_proposer.py |
| vllm_ascend/spec_decode/interface.py |
| vllm_ascend/spec_decode/mtp_proposer.py |
| vllm_ascend/spec_decode/ngram_proposer.py |
| vllm_ascend/spec_decode/suffix_proposer.py |

### Does this PR introduce _any_ user-facing change?

### How was this patch tested?

- vLLM version: v0.13.0
- vLLM main:
d68209402d

Signed-off-by: MrZ20 <2609716663@qq.com>
Signed-off-by: SILONG ZENG <2609716663@qq.com>
This commit is contained in:
SILONG ZENG
2026-02-06 15:25:08 +08:00
committed by GitHub
parent 99aedaff63
commit 4fb3d5e1b2
17 changed files with 948 additions and 1147 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
vllm_ascend/spec_decode/__init__.py
View File

@@ -30,10 +30,9 @@ def get_spec_decode_method(method, vllm_config, device, runner):
return EagleProposer(vllm_config, device, runner)
elif method == "mtp":
return MtpProposer(vllm_config, device, runner)
elif method == 'suffix':
elif method == "suffix":
return SuffixDecodingProposer(vllm_config, device, runner)
elif method == "medusa":
return MedusaProposer(vllm_config, device, runner)
else:
raise ValueError("Unknown speculative decoding method: "
f"{method}")
raise ValueError(f"Unknown speculative decoding method: {method}")

677
vllm_ascend/spec_decode/eagle_proposer.py
View File

File diff suppressed because it is too large Load Diff

47
vllm_ascend/spec_decode/interface.py
View File

@@ -1,5 +1,4 @@
import enum
from typing import Optional
import torch
from vllm.config import CUDAGraphMode, VllmConfig
@@ -18,11 +17,7 @@ class SpecDcodeType(enum.Enum):
class Proposer:
def __init__(self,
vllm_config: VllmConfig,
device: torch.device = None,
runner=None):
def __init__(self, vllm_config: VllmConfig, device: torch.device = None, runner=None):
pass
def load_model(self, model):
@@ -30,25 +25,29 @@ class Proposer:
raise NotImplementedError
@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: Optional[torch.Tensor] = None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None):
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,
):
"""Called by dummy_run in modle_runner"""
raise NotImplementedError
def generate_token_ids(self,
valid_sampled_token_ids: list[list[int]],
sampling_metadata: SamplingMetadata = None,
scheduler_output: SchedulerOutput = None,
spec_decode_metadata: SpecDecodeMetadata = None,
positions: torch.Tensor = None,
num_scheduled_tokens: int = 0,
hidden_states: torch.Tensor = None,
aux_hidden_states: torch.Tensor = None):
def generate_token_ids(
self,
valid_sampled_token_ids: list[list[int]],
sampling_metadata: SamplingMetadata = None,
scheduler_output: SchedulerOutput = None,
spec_decode_metadata: SpecDecodeMetadata = None,
positions: torch.Tensor = None,
num_scheduled_tokens: int = 0,
hidden_states: torch.Tensor = None,
aux_hidden_states: torch.Tensor = None,
):
"""Called by execute_model in model_runner"""
raise NotImplementedError
raise NotImplementedError

87
vllm_ascend/spec_decode/medusa_proposer.py
View File

@@ -1,14 +1,9 @@
from typing import Optional
import torch
import torch.nn as nn
from vllm.config import CUDAGraphMode, VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.model_loader import get_model
from vllm.model_executor.models.interfaces import is_mixture_of_experts
from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm.v1.spec_decode.medusa import MedusaProposer as VllmMedusaProposer
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
from vllm_ascend.spec_decode.interface import SpecDcodeType
@@ -22,72 +17,70 @@ class MedusaProposer(VllmMedusaProposer):
"""
def __init__(
self,
vllm_config: VllmConfig,
device: torch.device,
runner,
self,
vllm_config: VllmConfig,
device: torch.device,
runner,
):
# Save config parameters
self.name = SpecDcodeType.MEDUSA
self.vllm_config = vllm_config
self.device = device
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.hidden_size = (vllm_config.speculative_config.draft_model_config.
get_hidden_size())
self.hidden_size = vllm_config.speculative_config.draft_model_config.get_hidden_size()
self.dtype = vllm_config.model_config.dtype
self.runner = runner
@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: Optional[torch.Tensor] = None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None,
dummy_compute_logits=lambda hidden_states: None,
is_profile=False):
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,
):
hidden_states = torch.zeros(
(self.max_num_tokens, self.hidden_size),
dtype=self.dtype,
device=self.device,
)
with set_ascend_forward_context(
None,
self.vllm_config,
num_tokens=num_tokens,
num_actual_tokens=0,
in_profile_run=is_profile,
batch_descriptor=batch_descriptor,
aclgraph_runtime_mode=aclgraph_runtime_mode,
is_draft_model=True):
None,
self.vllm_config,
num_tokens=num_tokens,
num_actual_tokens=0,
in_profile_run=is_profile,
batch_descriptor=batch_descriptor,
aclgraph_runtime_mode=aclgraph_runtime_mode,
is_draft_model=True,
):
self.model(hidden_states)
dummy_compute_logits(hidden_states)
def generate_token_ids(self, valid_sampled_token_ids: list[list[int]],
sampling_metadata: SamplingMetadata,
spec_decode_metadata: SpecDecodeMetadata,
sample_hidden_states: torch.Tensor,
*args,
**kwargs
):
def generate_token_ids(
self,
valid_sampled_token_ids: list[list[int]],
sampling_metadata: SamplingMetadata,
spec_decode_metadata: SpecDecodeMetadata,
sample_hidden_states: torch.Tensor,
*args,
**kwargs,
):
if sample_hidden_states.shape[0] == len(valid_sampled_token_ids):
# The input to the target model does not include draft tokens.
hidden_states = sample_hidden_states
else:
num_accepted_tokens = torch.tensor(
[len(t) for t in valid_sampled_token_ids],
device=self.device,
dtype=torch.long)
num_draft_tokens = torch.tensor(
spec_decode_metadata.num_draft_tokens,
device=self.device,
dtype=torch.long)
[len(t) for t in valid_sampled_token_ids], device=self.device, dtype=torch.long
)
num_draft_tokens = torch.tensor(spec_decode_metadata.num_draft_tokens, device=self.device, dtype=torch.long)
offsets = torch.cumsum(num_draft_tokens + 1,
dim=0) - (num_draft_tokens + 1)
offsets = torch.cumsum(num_draft_tokens + 1, dim=0) - (num_draft_tokens + 1)
indices = offsets + num_accepted_tokens - 1
hidden_states = sample_hidden_states[indices]

413
vllm_ascend/spec_decode/mtp_proposer.py
View File

@@ -1,5 +1,3 @@
from typing import Optional, Union
import torch
import torch.nn as nn
from vllm.config import CUDAGraphMode
@@ -22,29 +20,33 @@ from vllm_ascend.utils import lmhead_tp_enable, vllm_version_is
class MtpProposer(EagleProposer):
# TODO: Find out why ModelRunner does not this explicit typing?
model: Union[nn.Module, ACLGraphWrapper]
model: nn.Module | ACLGraphWrapper
@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=None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None,
dummy_compute_logits=lambda hidden_states: None,
is_profile=False) -> None:
if (
self.pcp_size * self.dcp_size == 1
and not self.speculative_config.disable_padded_drafter_batch
):
def dummy_run(
self,
num_tokens: int,
with_prefill: bool = False,
in_graph_capturing: bool = False,
num_reqs: int = 0,
num_tokens_across_dp=None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None,
dummy_compute_logits=lambda hidden_states: None,
is_profile=False,
) -> None:
if self.pcp_size * self.dcp_size == 1 and not self.speculative_config.disable_padded_drafter_batch:
super().dummy_run(
num_tokens, with_prefill, in_graph_capturing, num_reqs,
num_tokens_across_dp, aclgraph_runtime_mode, batch_descriptor,
dummy_compute_logits, is_profile
num_tokens,
with_prefill,
in_graph_capturing,
num_reqs,
num_tokens_across_dp,
aclgraph_runtime_mode,
batch_descriptor,
dummy_compute_logits,
is_profile,
)
return
(
@@ -61,14 +63,10 @@ class MtpProposer(EagleProposer):
aclgraph_runtime_mode = CUDAGraphMode.NONE
if aclgraph_runtime_mode == CUDAGraphMode.FULL:
if len(self.runner.attn_groups) > 0:
num_computed_tokens_cpu = (
self.runner.input_batch.
num_computed_tokens_cpu_tensor[:num_reqs])
num_computed_tokens_cpu = self.runner.input_batch.num_computed_tokens_cpu_tensor[:num_reqs]
common_attn_metadata = AscendCommonAttentionMetadata(
query_start_loc=self.runner.query_start_loc.gpu[:num_reqs +
1],
query_start_loc_cpu=self.runner.query_start_loc.
cpu[:num_reqs + 1],
query_start_loc=self.runner.query_start_loc.gpu[: num_reqs + 1],
query_start_loc_cpu=self.runner.query_start_loc.cpu[: num_reqs + 1],
seq_lens_cpu=self.runner.seq_lens.cpu,
seq_lens=self.runner.seq_lens.gpu[:num_reqs],
num_reqs=num_reqs,
@@ -77,27 +75,29 @@ class MtpProposer(EagleProposer):
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(),
slot_mapping=self.runner.input_batch.block_table[0].
slot_mapping.gpu,
block_table_tensor=self.runner.input_batch.block_table[0].get_device_tensor(),
slot_mapping=self.runner.input_batch.block_table[0].slot_mapping.gpu,
positions=self.runner.positions.gpu,
attn_state=self.runner.attn_state,
decode_token_per_req=self.runner.decode_token_per_req,
max_seq_len=0)
max_seq_len=0,
)
if self.pcp_size * self.dcp_size > 1:
# update long_seq related params and flatten block_table
common_attn_metadata.prefill_context_parallel_metadata = \
self.runner.pcp_manager.long_seq_metadata
common_attn_metadata.block_table_tensor = \
self.runner.input_batch.block_table[0].get_device_tensor()[
:num_reqs * self.decode_threshold]
common_attn_metadata.prefill_context_parallel_metadata = self.runner.pcp_manager.long_seq_metadata
common_attn_metadata.block_table_tensor = self.runner.input_batch.block_table[
0
].get_device_tensor()[: num_reqs * self.decode_threshold]
builder = self.runner.attn_groups[0][0].get_metadata_builder()
# `AscendAttentionState.SpecDecoding` is only designed for mla, `AscendAttentionState.ChunkedPrefill` is used in self-attention.
attn_state = AscendAttentionState.SpecDecoding if self.vllm_config.model_config.use_mla else AscendAttentionState.ChunkedPrefill
attn_metadata_mtp = builder.build_for_graph_capture(
common_attn_metadata, attn_state)
# `AscendAttentionState.SpecDecoding` is only designed for mla,
# `AscendAttentionState.ChunkedPrefill` is used in self-attention.
attn_state = (
AscendAttentionState.SpecDecoding
if self.vllm_config.model_config.use_mla
else AscendAttentionState.ChunkedPrefill
)
attn_metadata_mtp = builder.build_for_graph_capture(common_attn_metadata, attn_state)
attn_metadata = {}
for layer_name in self.attn_layer_names:
attn_metadata[layer_name] = attn_metadata_mtp
@@ -113,32 +113,34 @@ class MtpProposer(EagleProposer):
if i > 0 and not in_graph_capturing and aclgraph_runtime_mode == CUDAGraphMode.FULL:
aclgraph_runtime_mode = CUDAGraphMode.NONE
with set_ascend_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp,
num_actual_tokens=0,
aclgraph_runtime_mode=aclgraph_runtime_mode,
batch_descriptor=batch_descriptor,
is_draft_model=True,
in_profile_run=is_profile):
attn_metadata,
self.vllm_config,
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp,
num_actual_tokens=0,
aclgraph_runtime_mode=aclgraph_runtime_mode,
batch_descriptor=batch_descriptor,
is_draft_model=True,
in_profile_run=is_profile,
):
if not vllm_version_is("v0.15.0"):
# Reset MOE layer index for each MTP step iteration
forward_context = get_forward_context()
if forward_context is not None:
forward_context.moe_layer_index = 0
previous_hidden_states, positions = self.maybe_pad_and_reduce(
previous_hidden_states, positions)
self.model(input_ids=input_ids,
positions=positions,
hidden_states=previous_hidden_states)
previous_hidden_states, positions = self.maybe_pad_and_reduce(previous_hidden_states, positions)
self.model(input_ids=input_ids, positions=positions, hidden_states=previous_hidden_states)
forward_context = get_forward_context()
if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL and \
not forward_context.capturing and not self.use_sparse:
if (
forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL
and not forward_context.capturing
and not self.use_sparse
):
self._update_full_graph_params(forward_context, num_tokens)
previous_hidden_states, positions, _ = self.maybe_all_gather_and_unpad(
previous_hidden_states, positions)
previous_hidden_states, positions
)
dummy_compute_logits(previous_hidden_states)
if with_prefill:
break
@@ -153,11 +155,10 @@ class MtpProposer(EagleProposer):
target_hidden_states: torch.Tensor,
# [batch_size]
next_token_ids: torch.Tensor,
last_token_indices: Optional[torch.Tensor],
last_token_indices: torch.Tensor | None,
common_attn_metadata: CommonAttentionMetadata,
sampling_metadata: SamplingMetadata,
mm_embed_inputs: Optional[tuple[list[torch.Tensor],
torch.Tensor]] = None,
mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
req_scheduled_tokens=None,
long_seq_metadata=None,
num_prefill_reqs=0,
@@ -165,16 +166,22 @@ class MtpProposer(EagleProposer):
scheduler_output: SchedulerOutput = None,
num_scheduled_tokens: int = 0,
) -> torch.Tensor:
if (
self.pcp_size * self.dcp_size == 1
and not self.speculative_config.disable_padded_drafter_batch
):
if self.pcp_size * self.dcp_size == 1 and not self.speculative_config.disable_padded_drafter_batch:
draft_token_ids = super()._propose(
target_token_ids, target_positions, target_hidden_states,
next_token_ids, last_token_indices, common_attn_metadata,
sampling_metadata, mm_embed_inputs, req_scheduled_tokens,
long_seq_metadata, num_prefill_reqs, num_decode_reqs,
scheduler_output, num_scheduled_tokens
target_token_ids,
target_positions,
target_hidden_states,
next_token_ids,
last_token_indices,
common_attn_metadata,
sampling_metadata,
mm_embed_inputs,
req_scheduled_tokens,
long_seq_metadata,
num_prefill_reqs,
num_decode_reqs,
scheduler_output,
num_scheduled_tokens,
)
return draft_token_ids
@@ -186,13 +193,12 @@ class MtpProposer(EagleProposer):
if self.method == "eagle3":
assert isinstance(self.model, Eagle3LlamaForCausalLM)
target_hidden_states = self.model.combine_hidden_states(
target_hidden_states)
target_hidden_states = self.model.combine_hidden_states(target_hidden_states)
assert target_hidden_states.shape[-1] == self.hidden_size
# 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:]
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[last_token_indices] = next_token_ids
@@ -213,20 +219,16 @@ class MtpProposer(EagleProposer):
num_tokens_d_padded = num_tokens_d * self.pcp_size
input_ids_d = self.input_ids[:num_tokens_d]
input_ids_p = self.input_ids[num_tokens_d:num_tokens]
target_hidden_states_d_padded = \
target_hidden_states[:num_tokens_d_padded]
target_hidden_states_d_padded = target_hidden_states[:num_tokens_d_padded]
if num_tokens_d:
# remove padding (from pcp all-gather) in decode part
mask_start_loc = torch.cat([
torch.tensor([0], dtype=torch.int32),
torch.cumsum(query_lens_d * self.pcp_size, dim=0)[:-1]
])
mask_start_loc = torch.cat(
[torch.tensor([0], dtype=torch.int32), torch.cumsum(query_lens_d * self.pcp_size, dim=0)[:-1]]
)
mask_len = query_lens_d
mask = []
for req_id in range(num_decode_reqs):
mask += list(
range(mask_start_loc[req_id],
mask_start_loc[req_id] + mask_len[req_id]))
mask += list(range(mask_start_loc[req_id], mask_start_loc[req_id] + mask_len[req_id]))
target_hidden_states_d = target_hidden_states_d_padded[mask]
else:
target_hidden_states_d = target_hidden_states_d_padded
@@ -234,46 +236,33 @@ class MtpProposer(EagleProposer):
req_scheduled_tokens_p = {}
for i, req_id in enumerate(self.runner.input_batch.req_ids):
if i >= num_decode_reqs:
req_scheduled_tokens_p[req_id] = \
req_scheduled_tokens[req_id]
(num_tokens_p, input_ids_p, target_hidden_states_p,
max_query_len_p, seq_lens_p, cu_num_tokens_p) = \
self._split_pcp_input(
req_scheduled_tokens_p, input_ids_p, target_hidden_states_p)
req_scheduled_tokens_p[req_id] = req_scheduled_tokens[req_id]
(num_tokens_p, input_ids_p, target_hidden_states_p, max_query_len_p, seq_lens_p, cu_num_tokens_p) = (
self._split_pcp_input(req_scheduled_tokens_p, input_ids_p, target_hidden_states_p)
)
num_tokens = num_tokens_d + num_tokens_p
target_positions = target_positions[:num_tokens]
self.input_ids[:num_tokens].copy_(
torch.cat([input_ids_d, input_ids_p], dim=0))
target_hidden_states = torch.cat(
[target_hidden_states_d, target_hidden_states_p], dim=0)
self.input_ids[:num_tokens].copy_(torch.cat([input_ids_d, input_ids_p], dim=0))
target_hidden_states = torch.cat([target_hidden_states_d, target_hidden_states_p], dim=0)
# 2. update sample_indices according to main model
if num_decode_reqs:
last_token_indices[:num_decode_reqs] = \
self.runner.logits_indices[last_token_indices[:num_decode_reqs]]
last_token_indices[:num_decode_reqs] = self.runner.logits_indices[last_token_indices[:num_decode_reqs]]
if num_prefill_reqs:
last_token_indices[-num_prefill_reqs:] = \
self.runner.logits_indices[-num_prefill_reqs:]
last_token_indices[-num_prefill_reqs:] = self.runner.logits_indices[-num_prefill_reqs:]
# 3. update attn_metadata params that may be influenced by pcp
common_attn_metadata.num_actual_tokens = num_tokens
common_attn_metadata.max_query_len = max(
self.decode_threshold, max_query_len_p)
common_attn_metadata.max_query_len = max(self.decode_threshold, max_query_len_p)
common_attn_metadata.seq_lens[-num_prefill_reqs:] = seq_lens_p
common_attn_metadata.seq_lens_cpu[
-num_prefill_reqs:] = seq_lens_p
query_start_loc_p = cu_num_tokens_p[1:] + \
common_attn_metadata.query_start_loc[num_decode_reqs].item()
common_attn_metadata.query_start_loc[-num_prefill_reqs:] = \
query_start_loc_p
common_attn_metadata.query_start_loc_cpu[-num_prefill_reqs:] = \
query_start_loc_p
common_attn_metadata.seq_lens_cpu[-num_prefill_reqs:] = seq_lens_p
query_start_loc_p = cu_num_tokens_p[1:] + common_attn_metadata.query_start_loc[num_decode_reqs].item()
common_attn_metadata.query_start_loc[-num_prefill_reqs:] = query_start_loc_p
common_attn_metadata.query_start_loc_cpu[-num_prefill_reqs:] = query_start_loc_p
assert self.runner is not None
# Note(qcs): We may need to refactor these check logics.
if self.use_cuda_graph and num_scheduled_tokens <= self.runner.cudagraph_batch_sizes[
-1]:
num_input_tokens = self.runner.cudagraph_dispatcher._bs_to_padded_graph_size[
num_scheduled_tokens]
if self.use_cuda_graph and num_scheduled_tokens <= self.runner.cudagraph_batch_sizes[-1]:
num_input_tokens = self.runner.cudagraph_dispatcher._bs_to_padded_graph_size[num_scheduled_tokens]
else:
# Eager mode, no padding needed
num_input_tokens = num_tokens
@@ -282,23 +271,23 @@ class MtpProposer(EagleProposer):
self._set_positions(num_tokens, target_positions)
self.hidden_states[:num_tokens] = target_hidden_states
# eager/acl piecewise mode need to update num_tokens_across_dp
(num_input_tokens, num_tokens_across_dp,
with_prefill) = self.runner._sync_metadata_across_dp(
num_input_tokens, self.runner.with_prefill)
(num_input_tokens, num_tokens_across_dp, with_prefill) = self.runner._sync_metadata_across_dp(
num_input_tokens, self.runner.with_prefill
)
# Enable shared_expert_dp and MTP FULL graph may cause accuracy issues.
if scheduler_output and not self.enable_shared_expert_dp:
max_query_len = common_attn_metadata.max_query_len
uniform_decode = (max_query_len in list(
range(1, self.num_speculative_tokens +
2))) and (scheduler_output.total_num_scheduled_tokens
== self.runner.input_batch.num_reqs *
(self.num_speculative_tokens + 1))
uniform_decode = (max_query_len in list(range(1, self.num_speculative_tokens + 2))) and (
scheduler_output.total_num_scheduled_tokens
== self.runner.input_batch.num_reqs * (self.num_speculative_tokens + 1)
)
else:
uniform_decode = False
has_lora = len(self.runner.input_batch.lora_id_to_lora_request) > 0
aclgraph_runtime_mode, batch_descriptor = \
self.runner.cudagraph_dispatcher.dispatch(num_tokens=num_input_tokens, uniform_decode=uniform_decode, has_lora=has_lora)
aclgraph_runtime_mode, batch_descriptor = self.runner.cudagraph_dispatcher.dispatch(
num_tokens=num_input_tokens, uniform_decode=uniform_decode, has_lora=has_lora
)
if not self.use_cuda_graph:
# there is synchronization between mtp steps when enabling aclgraph,
# disable aclgraph when use async scheduling to avoid the
@@ -307,8 +296,10 @@ class MtpProposer(EagleProposer):
# and _propose.
aclgraph_runtime_mode = CUDAGraphMode.NONE
if self.vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(
) and aclgraph_runtime_mode == CUDAGraphMode.FULL:
if (
self.vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
and aclgraph_runtime_mode == CUDAGraphMode.FULL
):
graph_pad_size = num_input_tokens
else:
graph_pad_size = -1
@@ -319,64 +310,58 @@ class MtpProposer(EagleProposer):
common_attn_metadata.graph_pad_size = graph_pad_size
common_attn_metadata.num_input_tokens = num_input_tokens
builder = self.runner.attn_groups[0][0].get_metadata_builder()
attn_metadata_mtp = builder.build(0, common_attn_metadata,
self.runner.get_model())
attn_metadata_mtp = builder.build(0, common_attn_metadata, self.runner.get_model())
attn_metadata = {}
for layer_name in self.attn_layer_names:
attn_metadata[layer_name] = attn_metadata_mtp
for step in range(self.num_speculative_tokens):
with set_ascend_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
aclgraph_runtime_mode=aclgraph_runtime_mode,
batch_descriptor=batch_descriptor,
num_actual_tokens=num_tokens,
is_draft_model=True):
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
aclgraph_runtime_mode=aclgraph_runtime_mode,
batch_descriptor=batch_descriptor,
num_actual_tokens=num_tokens,
is_draft_model=True,
):
if not vllm_version_is("v0.15.0"):
# Reset MOE layer index for each MTP step to match all_moe_layers registration
forward_context = get_forward_context()
if forward_context is not None:
forward_context.moe_layer_index = 0
with record_function_or_nullcontext('mtp_forward'):
with record_function_or_nullcontext("mtp_forward"):
model_kwargs = {}
model_kwargs["attn_metadata"] = attn_metadata
input_ids = self.input_ids[:num_input_tokens]
positions = self._get_positions(num_input_tokens)
hidden_states = self.hidden_states[:num_input_tokens]
hidden_states, positions = self.maybe_pad_and_reduce(
hidden_states, positions)
hidden_states, positions = self.maybe_pad_and_reduce(hidden_states, positions)
hidden_states = self.model(input_ids=input_ids,
positions=positions,
hidden_states=hidden_states)
forward_context = get_forward_context()
if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL and not self.use_sparse:
self._update_full_graph_params(forward_context,
num_input_tokens)
hidden_states = self.model(input_ids=input_ids, positions=positions, hidden_states=hidden_states)
forward_context = get_forward_context()
if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL and not self.use_sparse:
self._update_full_graph_params(forward_context, num_input_tokens)
hidden_states, positions, _ = self.maybe_all_gather_and_unpad(
hidden_states, positions)
hidden_states, positions, _ = self.maybe_all_gather_and_unpad(hidden_states, positions)
num_indices = last_token_indices.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
last_token_indices = nn.functional.pad(
last_token_indices,
(0, max_num_reqs_across_dp - num_indices))
max_num_reqs_across_dp = (
self.vllm_config.scheduler_config.max_num_seqs * self.runner.uniform_decode_query_len
)
last_token_indices = nn.functional.pad(last_token_indices, (0, max_num_reqs_across_dp - num_indices))
if self.pcp_size > 1 and step == 0:
# remove graph padding before all_gather
hidden_states = hidden_states[:num_tokens]
hidden_states = get_pcp_group().all_gather(hidden_states, 0)
hidden_states = torch.index_select(
hidden_states, 0, self.runner.pcp_manager.
pcp_allgather_restore_idx.gpu[:hidden_states.shape[0]])
hidden_states, 0, self.runner.pcp_manager.pcp_allgather_restore_idx.gpu[: hidden_states.shape[0]]
)
sample_hidden_states = hidden_states[last_token_indices]
logits = self.model.compute_logits(sample_hidden_states)
@@ -409,7 +394,7 @@ class MtpProposer(EagleProposer):
hidden_states = hidden_states[last_token_indices]
slot_mapping = attn_metadata_i.slot_mapping[last_token_indices]
attn_metadata_i.slot_mapping.fill_(-1)
attn_metadata_i.query_start_loc = self.arange[:batch_size + 1]
attn_metadata_i.query_start_loc = self.arange[: batch_size + 1]
last_token_indices = self.arange[:batch_size]
if getattr(attn_metadata_i, "num_decode_tokens", 0):
attn_metadata_i.num_decode_tokens = batch_size
@@ -420,44 +405,44 @@ class MtpProposer(EagleProposer):
# Instead, we pre-allocate mtp slot_mapping in model_runner
# (_generate_pcp_mtp_input), and use updated slot_indices
# to get corresponding slot_mapping in each step.
num_reject_tokens = torch.tensor(
self.runner.pcp_manager.cu_num_tokens_pcp_full,
dtype=torch.int32).to(
self.device) - ori_last_token_indices - 1
num_accept_tokens = \
query_lens_d.to(self.device) - num_reject_tokens
num_reject_tokens = (
torch.tensor(self.runner.pcp_manager.cu_num_tokens_pcp_full, dtype=torch.int32).to(self.device)
- ori_last_token_indices
- 1
)
num_accept_tokens = query_lens_d.to(self.device) - num_reject_tokens
ori_seq_len = attn_metadata_i.seq_lens
mtp_slot_mapping = self.runner.pcp_manager.mtp_slot_pad
# slot_mapping index base offset:
# scheduled tokens + pre-allocated mtp tokens + accepted tokens
slot_idx_base = (
torch.cat([
torch.tensor(
[0], dtype=torch.int32, device=self.device),
(torch.cumsum(query_lens_d, dim=0)[:-1] *
self.pcp_size).to(self.device)
]) +
torch.arange(num_decode_reqs, device=self.device) *
(self.num_speculative_tokens - 1) * self.pcp_size +
(num_accept_tokens - 1) * self.pcp_size)
torch.cat(
[
torch.tensor([0], dtype=torch.int32, device=self.device),
(torch.cumsum(query_lens_d, dim=0)[:-1] * self.pcp_size).to(self.device),
]
)
+ torch.arange(num_decode_reqs, device=self.device)
* (self.num_speculative_tokens - 1)
* self.pcp_size
+ (num_accept_tokens - 1) * self.pcp_size
)
slot_indices_list = []
for req_id in range(num_decode_reqs):
slot_indices_list.append(
torch.arange(slot_idx_base[req_id],
slot_idx_base[req_id] + self.pcp_size,
device=self.device))
torch.arange(
slot_idx_base[req_id], slot_idx_base[req_id] + self.pcp_size, device=self.device
)
)
slot_indices = torch.cat(slot_indices_list, dim=0)
# fold block_table (restore it to original size before flattened)
block_indices = torch.cat([
torch.tensor([0], dtype=torch.int32),
torch.cumsum(query_lens_d, dim=0)[:-1]
])
attn_metadata_i.decode.block_table[:batch_size] = \
attn_metadata_i.decode.block_table[block_indices]
attn_metadata_i.decode.block_table = \
attn_metadata_i.decode.block_table[:batch_size]
block_indices = torch.cat(
[torch.tensor([0], dtype=torch.int32), torch.cumsum(query_lens_d, dim=0)[:-1]]
)
attn_metadata_i.decode.block_table[:batch_size] = attn_metadata_i.decode.block_table[block_indices]
attn_metadata_i.decode.block_table = attn_metadata_i.decode.block_table[:batch_size]
input_ids = draft_token_ids_list[-1].int()
positions += 1
@@ -465,38 +450,32 @@ class MtpProposer(EagleProposer):
decode_metadata = getattr(attn_metadata_i, "decode", None)
prefill_metadata = getattr(attn_metadata_i, "prefill", None)
# When disable_padded_drafter_batch=False, it should not to be updating these params, maybe.
if decode_metadata is not None and (self.speculative_config.disable_padded_drafter_batch or \
aclgraph_runtime_mode != CUDAGraphMode.FULL):
decode_metadata.actual_seq_lengths_q = self.arange_cpu[
1:batch_size + 1].tolist()
if decode_metadata is not None and (
self.speculative_config.disable_padded_drafter_batch or aclgraph_runtime_mode != CUDAGraphMode.FULL
):
decode_metadata.actual_seq_lengths_q = self.arange_cpu[1 : batch_size + 1].tolist()
if aclgraph_runtime_mode == CUDAGraphMode.FULL:
decode_metadata.actual_seq_lengths_q = \
builder.pad_actual_seq_len_q_mtp_disable_pad(
graph_pad_size - batch_size,
batch_size,
decode_metadata.actual_seq_lengths_q)
decode_metadata.cos, decode_metadata.sin = get_cos_and_sin_mla(
positions[:batch_size])
decode_metadata.actual_seq_lengths_q = builder.pad_actual_seq_len_q_mtp_disable_pad(
graph_pad_size - batch_size, batch_size, decode_metadata.actual_seq_lengths_q
)
decode_metadata.cos, decode_metadata.sin = get_cos_and_sin_mla(positions[:batch_size])
# NOTE(woosuk): We should handle the case where the draft model
# generates tokens beyond the max model length. Since it is complex
# to remove such requests from the batch, we keep them in the batch
# but adjust the position ids and slot mappings to avoid the
# out-of-range access during the model execution. The draft tokens
# generated with this adjustment should be ignored.
exceeds_max_model_len = positions[:
batch_size] >= self.runner.model_config.max_model_len
exceeds_max_model_len = positions[:batch_size] >= self.runner.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, 0,
positions[:batch_size])
clamped_positions = torch.where(exceeds_max_model_len, 0, positions[:batch_size])
# Increment the sequence lengths.
# This is an out-of-place operation to avoid modifying the original tensor
# when enable async_scheduling.
attn_metadata_i.seq_lens = attn_metadata_i.seq_lens + 1
# For the requests that exceed the max model length, we set the
# sequence length to 1 to minimize their overheads in attention.
exceeds_mask = attn_metadata_i.seq_lens[:batch_size] > \
self.runner.model_config.max_model_len
exceeds_mask = attn_metadata_i.seq_lens[:batch_size] > self.runner.model_config.max_model_len
attn_metadata_i.seq_lens[:batch_size].masked_fill_(exceeds_mask, 1)
# Mask out the slot mappings that exceed the max model length.
# Otherwise, the KV cache will be inadvertently updated with the
@@ -504,13 +483,14 @@ class MtpProposer(EagleProposer):
slot_mapping += 1
if self.pcp_size > 1:
exceeds_max_model_len = exceeds_max_model_len.repeat_interleave(
slot_mapping.size(0) // exceeds_max_model_len.size(0))
slot_mapping.size(0) // exceeds_max_model_len.size(0)
)
slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID)
# copy inputs to buffer for cudagraph
self.input_ids[:batch_size] = input_ids
self._set_positions(batch_size, clamped_positions)
self.hidden_states[:hidden_states.shape[0]] = hidden_states
self.hidden_states[: hidden_states.shape[0]] = hidden_states
if self.pcp_size * self.dcp_size > 1:
# update local seq_len
num_computed_tokens_of_pcp_dcp = self.runner.pcp_manager._get_cp_local_seq_lens(
@@ -519,19 +499,17 @@ class MtpProposer(EagleProposer):
self.dcp_size,
self.runner.parallel_config.cp_kv_cache_interleave_size,
)
cp_seq_len = \
num_computed_tokens_of_pcp_dcp[:, self.pcp_rank, self.dcp_rank]
cp_seq_len = num_computed_tokens_of_pcp_dcp[:, self.pcp_rank, self.dcp_rank]
attn_metadata_i.decode.cp_seq_len = cp_seq_len
# update slot_mapping
slot_indices += self.pcp_size
slot_mapping = mtp_slot_mapping[slot_indices]
attn_metadata_i.slot_mapping[:batch_size *
self.pcp_size] = slot_mapping
attn_metadata_i.slot_mapping[: batch_size * self.pcp_size] = slot_mapping
else:
attn_metadata_i.slot_mapping[:batch_size] = slot_mapping
if self.speculative_config.disable_padded_drafter_batch:
if self.uses_mrope:
self.mrope_positions[:, batch_size:num_input_tokens] = 0
self.mrope_positions[:, batch_size:num_input_tokens] = 0
else:
self.positions[batch_size:num_input_tokens] = 0
self.input_ids[batch_size:num_input_tokens] = 0
@@ -539,31 +517,24 @@ class MtpProposer(EagleProposer):
if prefill_metadata is not None:
prefill_metadata.seq_lens = attn_metadata_i.seq_lens
prefill_metadata.seq_lens_list = prefill_metadata.seq_lens.tolist(
)
prefill_metadata.seq_lens_list = prefill_metadata.seq_lens.tolist()
prefill_metadata.context_lens = attn_metadata_i.seq_lens
prefill_metadata.input_positions = self._get_positions(
num_input_tokens)
prefill_metadata.input_positions = self._get_positions(num_input_tokens)
prefill_metadata.max_seq_lens += 1
prefill_metadata.max_seq_lens = min(
prefill_metadata.max_seq_lens,
self.runner.model_config.max_model_len)
prefill_metadata.max_seq_lens, self.runner.model_config.max_model_len
)
if decode_metadata is not None:
decode_metadata.seq_lens = attn_metadata_i.seq_lens
decode_metadata.seq_lens_list = decode_metadata.seq_lens.tolist(
)
decode_metadata.seq_lens_list = decode_metadata.seq_lens.tolist()
decode_seq_lens_list = decode_metadata.seq_lens_list
if aclgraph_runtime_mode == CUDAGraphMode.FULL and \
self.speculative_config.disable_padded_drafter_batch:
decode_metadata.seq_lens_list = decode_seq_lens_list + [
0
] * (graph_pad_size - len(decode_seq_lens_list))
decode_metadata.input_positions = self._get_positions(
num_input_tokens)
if aclgraph_runtime_mode == CUDAGraphMode.FULL and self.speculative_config.disable_padded_drafter_batch:
decode_metadata.seq_lens_list = decode_seq_lens_list + [0] * (
graph_pad_size - len(decode_seq_lens_list)
)
decode_metadata.input_positions = self._get_positions(num_input_tokens)
decode_metadata.max_seq_lens += 1
decode_metadata.max_seq_lens = min(
decode_metadata.max_seq_lens,
self.runner.model_config.max_model_len)
decode_metadata.max_seq_lens = min(decode_metadata.max_seq_lens, self.runner.model_config.max_model_len)
# mtp>1: [batch_size, k]
draft_token_ids = torch.stack(draft_token_ids_list, dim=1)

49
vllm_ascend/spec_decode/ngram_proposer.py
View File

@@ -1,13 +1,11 @@
import torch
from vllm.config import CUDAGraphMode
from vllm.v1.spec_decode.ngram_proposer import \
NgramProposer as VllmNgramProposer
from vllm.v1.spec_decode.ngram_proposer import NgramProposer as VllmNgramProposer
from vllm_ascend.spec_decode.interface import Proposer, SpecDcodeType
class NgramProposer(VllmNgramProposer, Proposer):
def __init__(self, vllm_config, device, runner):
super().__init__(vllm_config)
self.name = SpecDcodeType.NGRAM
@@ -19,27 +17,31 @@ class NgramProposer(VllmNgramProposer, Proposer):
pass
@torch.inference_mode()
def dummy_run(self,
num_tokens,
with_prefill=None,
in_graph_capturing=None,
num_reqs=None,
num_tokens_across_dp=None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None,
dummy_compute_logits=lambda hidden_states: None,
is_profile=False):
def dummy_run(
self,
num_tokens,
with_prefill=None,
in_graph_capturing=None,
num_reqs=None,
num_tokens_across_dp=None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None,
dummy_compute_logits=lambda hidden_states: None,
is_profile=False,
):
pass
def generate_token_ids(self,
valid_sampled_token_ids,
sampling_metadata=None,
scheduler_output=None,
spec_decode_metadata=None,
positions=None,
num_scheduled_tokens=None,
hidden_states=None,
aux_hidden_states=None) -> list[list[int]]:
def generate_token_ids(
self,
valid_sampled_token_ids,
sampling_metadata=None,
scheduler_output=None,
spec_decode_metadata=None,
positions=None,
num_scheduled_tokens=None,
hidden_states=None,
aux_hidden_states=None,
) -> list[list[int]]:
valid_ngram_requests = []
for i, sampled_ids in enumerate(valid_sampled_token_ids):
num_sampled_ids = len(sampled_ids)
@@ -57,8 +59,7 @@ class NgramProposer(VllmNgramProposer, Proposer):
start_idx = self.runner.input_batch.num_tokens_no_spec[i]
end_idx = start_idx + num_sampled_ids
self.runner.input_batch.token_ids_cpu[
i, start_idx:end_idx] = sampled_ids
self.runner.input_batch.token_ids_cpu[i, start_idx:end_idx] = sampled_ids
valid_ngram_requests.append(i)

49
vllm_ascend/spec_decode/suffix_proposer.py
View File

@@ -1,13 +1,11 @@
import torch
from vllm.config import CUDAGraphMode
from vllm.v1.spec_decode.suffix_decoding import \
SuffixDecodingProposer as VllmSuffixDecodingProposer
from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer as VllmSuffixDecodingProposer
from vllm_ascend.spec_decode.interface import Proposer, SpecDcodeType
class SuffixDecodingProposer(VllmSuffixDecodingProposer, Proposer):
def __init__(self, vllm_config, device, runner):
super().__init__(vllm_config)
self.name = SpecDcodeType.SUFFIX
@@ -19,27 +17,30 @@ class SuffixDecodingProposer(VllmSuffixDecodingProposer, Proposer):
pass
@torch.inference_mode()
def dummy_run(self,
num_tokens,
with_prefill=None,
in_graph_capturing=None,
num_reqs=None,
num_tokens_across_dp=None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None,
dummy_compute_logits=lambda hidden_states: None,
is_profile=False):
def dummy_run(
self,
num_tokens,
with_prefill=None,
in_graph_capturing=None,
num_reqs=None,
num_tokens_across_dp=None,
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
batch_descriptor=None,
dummy_compute_logits=lambda hidden_states: None,
is_profile=False,
):
pass
def generate_token_ids(self,
valid_sampled_token_ids,
sampling_metadata=None,
scheduler_output=None,
spec_decode_metadata=None,
positions=None,
num_scheduled_tokens=None,
hidden_states=None,
aux_hidden_states=None) -> list[list[int]]:
draft_token_ids = self.propose(self.runner.input_batch,
valid_sampled_token_ids)
def generate_token_ids(
self,
valid_sampled_token_ids,
sampling_metadata=None,
scheduler_output=None,
spec_decode_metadata=None,
positions=None,
num_scheduled_tokens=None,
hidden_states=None,
aux_hidden_states=None,
) -> list[list[int]]:
draft_token_ids = self.propose(self.runner.input_batch, valid_sampled_token_ids)
return draft_token_ids