270
vllm_ascend/spec_decode/dflash_proposer.py
Normal file
270
vllm_ascend/spec_decode/dflash_proposer.py
Normal file
@@ -0,0 +1,270 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.config import CUDAGraphMode, VllmConfig
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
|
||||
|
||||
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.ops.triton.spec_decode.utils import copy_and_expand_dflash_inputs_kernel_single_grid
|
||||
from vllm_ascend.spec_decode.eagle_proposer import AscendEagleProposer
|
||||
|
||||
|
||||
class AscendDflashProposer(AscendEagleProposer):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
device: torch.device,
|
||||
runner=None,
|
||||
):
|
||||
super().__init__(
|
||||
vllm_config,
|
||||
device,
|
||||
runner=runner,
|
||||
)
|
||||
|
||||
self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens)
|
||||
self.max_positions = self.max_num_tokens + self.max_query_tokens
|
||||
|
||||
self._context_slot_mapping_buffer = torch.zeros(
|
||||
self.max_num_tokens,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self._slot_mapping_buffer = torch.zeros(
|
||||
self.max_query_tokens,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self._context_positions_buffer = torch.zeros(
|
||||
self.max_num_tokens,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.positions = torch.zeros(
|
||||
self.max_query_tokens,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.arange_dflash = torch.arange(self.max_positions + 1, device=device, dtype=torch.int32)
|
||||
|
||||
self._dflash_hidden_states = torch.zeros(
|
||||
(self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=self.device
|
||||
)
|
||||
|
||||
self.parallel_drafting_hidden_state_tensor = None
|
||||
|
||||
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]:
|
||||
# DFlash cross-attention: context K/V from target hidden states,
|
||||
# Q from query embeddings (bonus + mask tokens).
|
||||
batch_size = cad.num_reqs
|
||||
num_context = target_token_ids.shape[0]
|
||||
num_query_per_req = 1 + self.num_speculative_tokens
|
||||
num_query_total = batch_size * num_query_per_req
|
||||
|
||||
self._dflash_num_context = num_context
|
||||
self._dflash_hidden_states[:num_context] = target_hidden_states
|
||||
|
||||
token_indices_to_sample = torch.empty(
|
||||
batch_size * self.num_speculative_tokens,
|
||||
dtype=torch.int32,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
has_num_rejected = num_rejected_tokens_gpu is not None
|
||||
|
||||
copy_and_expand_dflash_inputs_kernel_single_grid[1,](
|
||||
# Inputs
|
||||
next_token_ids_ptr=next_token_ids,
|
||||
target_positions_ptr=target_positions,
|
||||
context_slot_mapping_ptr=cad.slot_mapping,
|
||||
# Outputs
|
||||
out_input_ids_ptr=self.input_ids,
|
||||
out_context_positions_ptr=self._context_positions_buffer,
|
||||
out_query_positions_ptr=self.positions,
|
||||
out_context_slot_mapping_ptr=self._context_slot_mapping_buffer,
|
||||
out_query_slot_mapping_ptr=self._slot_mapping_buffer,
|
||||
out_token_indices_ptr=token_indices_to_sample,
|
||||
# Block table
|
||||
block_table_ptr=cad.block_table_tensor,
|
||||
block_table_stride=cad.block_table_tensor.stride(0),
|
||||
# Metadata
|
||||
query_start_loc_ptr=cad.query_start_loc,
|
||||
seq_lens_ptr=cad.seq_lens,
|
||||
num_rejected_tokens_ptr=(num_rejected_tokens_gpu if has_num_rejected else 0),
|
||||
# Scalars
|
||||
parallel_drafting_token_id=self.parallel_drafting_token_id,
|
||||
block_size=self.kernel_block_size,
|
||||
num_query_per_req=num_query_per_req,
|
||||
num_speculative_tokens=self.num_speculative_tokens,
|
||||
total_input_tokens=num_context,
|
||||
batch_size=batch_size,
|
||||
HAS_NUM_REJECTED=has_num_rejected,
|
||||
)
|
||||
|
||||
query_slot_mapping = self._slot_mapping_buffer[:num_query_total]
|
||||
new_query_start_loc = self.arange_dflash[: batch_size + 1] * num_query_per_req
|
||||
|
||||
effective_seq_lens = cad.seq_lens
|
||||
if has_num_rejected:
|
||||
effective_seq_lens = effective_seq_lens - num_rejected_tokens_gpu
|
||||
|
||||
cad.query_start_loc = new_query_start_loc
|
||||
cad.seq_lens = effective_seq_lens + num_query_per_req
|
||||
cad.query_start_loc_cpu = (
|
||||
torch.from_numpy(self.token_arange_np[: batch_size + 1]).clone() * num_query_per_req
|
||||
).to(torch.int32)
|
||||
|
||||
if hasattr(cad, "actual_seq_lengths_q"):
|
||||
cad.actual_seq_lengths_q = [num_query_per_req] * batch_size
|
||||
if hasattr(cad, "decode_token_per_req"):
|
||||
cad.decode_token_per_req = num_query_per_req
|
||||
|
||||
cad.num_actual_tokens = num_query_total
|
||||
cad.max_query_len = num_query_per_req
|
||||
cad.max_seq_len = cad.max_seq_len + num_query_per_req
|
||||
cad.slot_mapping = query_slot_mapping
|
||||
cad.causal = False
|
||||
cad.attn_mask = None
|
||||
cad.attn_state = AscendAttentionState.ChunkedPrefill
|
||||
|
||||
return num_query_total, token_indices_to_sample, cad, None
|
||||
|
||||
@torch.inference_mode()
|
||||
def dummy_run(
|
||||
self,
|
||||
num_tokens: int,
|
||||
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,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
num_query_tokens = min(num_tokens, self.max_query_tokens)
|
||||
|
||||
(
|
||||
num_input_tokens,
|
||||
num_tokens_across_dp,
|
||||
_,
|
||||
) = self.runner._sync_metadata_across_dp(num_query_tokens, is_draft_model=True)
|
||||
|
||||
if not self.use_cuda_graph:
|
||||
aclgraph_runtime_mode = CUDAGraphMode.NONE
|
||||
num_query_per_req = 1 + self.num_speculative_tokens
|
||||
num_query_total = num_reqs * num_query_per_req
|
||||
|
||||
context_positions = self._context_positions_buffer[:num_input_tokens]
|
||||
context_states = self.hidden_states[:num_input_tokens]
|
||||
|
||||
multi_steps_attn_metadata = []
|
||||
if aclgraph_runtime_mode == CUDAGraphMode.FULL and len(self.runner.attn_groups) > 0:
|
||||
builder = self.draft_attn_groups[0].get_metadata_builder()
|
||||
common_attn_metadata = AscendCommonAttentionMetadata(
|
||||
query_start_loc=self.arange_dflash[: num_reqs + 1] * num_query_per_req,
|
||||
query_start_loc_cpu=torch.from_numpy(self.token_arange_np[: num_reqs + 1]).clone() * num_query_per_req,
|
||||
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_query_tokens,
|
||||
max_query_len=num_query_per_req,
|
||||
max_seq_len=0,
|
||||
slot_mapping=self._slot_mapping_buffer[:num_query_total],
|
||||
attn_state=AscendAttentionState.ChunkedPrefill,
|
||||
causal=False,
|
||||
is_prefilling=torch.zeros(num_reqs, dtype=torch.bool),
|
||||
block_table_tensor=self.runner.input_batch.block_table[self.kv_cache_gid].get_device_tensor()[
|
||||
:num_reqs
|
||||
],
|
||||
)
|
||||
|
||||
attn_metadata_dflash = builder.build_for_graph_capture(
|
||||
common_attn_metadata,
|
||||
AscendAttentionState.ChunkedPrefill,
|
||||
)
|
||||
|
||||
attn_metadata_dflash.attn_mask = None
|
||||
attn_metadata_dflash.attn_state = AscendAttentionState.ChunkedPrefill
|
||||
|
||||
per_layer_attn_metadata = dict()
|
||||
for layer_name in self.attn_layer_names:
|
||||
per_layer_attn_metadata[layer_name] = attn_metadata_dflash
|
||||
multi_steps_attn_metadata.append(per_layer_attn_metadata)
|
||||
|
||||
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_input_tokens,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
num_actual_tokens=num_input_tokens,
|
||||
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,
|
||||
):
|
||||
if is_profile:
|
||||
self.model.precompute_and_store_context_kv(context_states, context_positions)
|
||||
self.model(
|
||||
input_ids=self.input_ids[:num_query_total],
|
||||
positions=self._get_positions(num_query_total),
|
||||
inputs_embeds=None,
|
||||
)
|
||||
|
||||
else:
|
||||
self._dflash_num_context = num_input_tokens
|
||||
self._runnable(
|
||||
num_input_tokens=num_input_tokens,
|
||||
batch_size=num_reqs,
|
||||
token_indices_to_sample=self.token_indices_to_sample[: num_reqs * self.num_speculative_tokens],
|
||||
target_positions=self._get_positions(num_input_tokens),
|
||||
inputs_embeds=None,
|
||||
multi_steps_attn_metadata=multi_steps_attn_metadata,
|
||||
num_tokens=num_input_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 build_model_inputs_first_pass(
|
||||
self,
|
||||
num_input_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
num_context = self._dflash_num_context
|
||||
|
||||
self.model.precompute_and_store_context_kv(
|
||||
self._dflash_hidden_states[:num_context],
|
||||
self._context_positions_buffer[:num_context],
|
||||
self._context_slot_mapping_buffer[:num_context],
|
||||
)
|
||||
|
||||
return dict(
|
||||
input_ids=self.input_ids[:num_input_tokens], positions=self.positions[:num_input_tokens], inputs_embeds=None
|
||||
)
|
||||
|
||||
def _raise_if_multimodal(self):
|
||||
pass
|
||||
Reference in New Issue
Block a user