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