# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations import math from typing import TYPE_CHECKING, Any import numpy as np import torch import torch.nn as nn from vllm.config import CUDAGraphMode from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import BatchDescriptor from vllm.logger import logger from vllm.multimodal.inputs import MultiModalKwargs from vllm.sequence import IntermediateTensors from vllm.utils import ( cdiv, ) from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder from vllm.v1.core.sched.output import SchedulerOutput # yapf conflicts with isort for this block # yapf: disable from vllm.v1.kv_cache_interface import ( EncoderOnlyAttentionSpec, ) # yapf: enable from vllm.v1.outputs import ( EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, ModelRunnerOutput, ) from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.attention.attention_v1 import AscendAttentionState from vllm_ascend.attention.utils import AscendCommonAttentionMetadata from vllm_ascend.spec_decode.interface import SpecDcodeType from vllm_ascend.utils import ( ProfileExecuteDuration, enable_sp, lmhead_tp_enable, ) from vllm_ascend.worker.model_runner_v1 import AsyncNPUModelRunnerOutput from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput class NPUDiffusionModelRunner(OmniNPUModelRunner): """Diffusion model runner for vLLM-omni on NPU (non-autoregressive).""" def _prepare_inputs( self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> tuple[ dict[str, Any], torch.Tensor, np.ndarray, int, torch.Tensor, int, torch.Tensor, SpecDecodeMetadata, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, int, dict[str, Any], ]: total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens assert total_num_scheduled_tokens > 0 num_reqs = self.input_batch.num_reqs assert num_reqs > 0 # OPTIMIZATION: Start copying the block table first. # This way, we can overlap the copy with the following CPU operations. self.input_batch.block_table.commit_block_table(num_reqs) # Get the number of scheduled tokens for each request. req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] num_scheduled_tokens = np.array(tokens, dtype=np.int32) max_num_scheduled_tokens = num_scheduled_tokens.max() num_valid_tokens = np.array( [ num_tokens - len(scheduler_output.scheduled_spec_decode_tokens.get(i, [])) for num_tokens, i in zip(tokens, req_ids) ], dtype=np.int32, ) if self.use_aclgraph and total_num_scheduled_tokens <= self.aclgraph_batch_sizes[-1]: # Add padding to the batch size. num_input_tokens = self.vllm_config.pad_for_cudagraph(total_num_scheduled_tokens) elif self.use_aclgraph and enable_sp(self.vllm_config): # When using aclgraph, if total_num_scheduled_tokens exceeds the maximum graph size, # the model will fall back to running its FX graph in eager mode. # In this case, when sequence parallelism is enabled, we need to pad tokens to align # with tp_size because pad_size cannot be captured by the FX graph tp_size = self.vllm_config.parallel_config.tensor_parallel_size num_input_tokens = math.ceil(total_num_scheduled_tokens / tp_size) * tp_size else: # Eager mode. num_input_tokens = total_num_scheduled_tokens # Get the attention state. attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens) self.attn_state = attn_state # type: ignore # Determine if it's a splitfuse batch with_prefill = attn_state not in [AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding] self.query_lens = torch.from_numpy(num_scheduled_tokens) enable_dbo = self._check_dbo_is_valid(self.query_lens.tolist(), attn_state, total_num_scheduled_tokens) # Get info across DP ranks. # NOTE: maybe_padded_num_tokens is only used when using TorchAir with DP, # Otherwise, it's just max_tokens_across_dp_cpu (maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo) = self._sync_metadata_across_dp( num_input_tokens, with_prefill, enable_dbo ) # TODO: Now that num_input_tokens is basically identical with maybe_padded_num_tokens # We should consider removing maybe_padded_num_tokens later num_input_tokens = maybe_padded_num_tokens # Hot-Swap lora model if self.lora_config: self.set_active_loras(self.input_batch, num_scheduled_tokens) # Get request indices. # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2] req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens) # cu_num_tokens: [2, 5, 3] -> [2, 7, 10] # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens) positions_np = self.positions_np[:total_num_scheduled_tokens] np.add(self.input_batch.num_computed_tokens_cpu[req_indices], arange, out=positions_np) # Calculate M-RoPE positions. # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: self._calc_mrope_positions(scheduler_output) # Only relevant for models using M-RoPE (e.g, Qwen2-VL) self.mrope_positions[:, :total_num_scheduled_tokens].copy_( self.mrope_positions_cpu[:, :total_num_scheduled_tokens], non_blocking=True ) # Get token indices. # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2] # where M is the max_model_len. token_indices = positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1] # Prepare input_ids. # NOTE(woosuk): We use torch.index_select instead of np.take here # because torch.index_select is much faster than np.take for large # tensors. torch.index_select( self.input_batch.token_ids_cpu_tensor.flatten(), 0, torch.from_numpy(token_indices), out=self.input_ids_cpu[:total_num_scheduled_tokens], ) # Prepare some information for building Attention-Metadata # Compute and commit slot mapping self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np) self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens) self.query_start_loc_np[0] = 0 self.query_start_loc_np[1 : num_reqs + 1] = cu_num_tokens self.query_start_loc[: num_reqs + 1].copy_(self.query_start_loc_cpu[: num_reqs + 1], non_blocking=True) self.seq_lens_np[:num_reqs] = self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], non_blocking=True) # Fill unused with -1. Needed for reshape_and_cache self.query_start_loc[num_reqs + 1 :].fill_(-1) self.seq_lens[num_reqs:].fill_(0) self.query_lens = torch.from_numpy(num_scheduled_tokens) # Copy the tensors to the NPU. self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens) self.positions_cpu[total_num_scheduled_tokens:num_input_tokens].zero_() self.positions[:num_input_tokens].copy_(self.positions_cpu[:num_input_tokens], non_blocking=True) # Make Attention metadata positions_cpu = self.positions_cpu[:num_input_tokens] positions = self.positions[:num_input_tokens] seq_lens_cpu = self.seq_lens_cpu[:num_reqs] attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens) self.attn_mask = self._make_attention_mask(seq_lens=seq_lens_cpu, position=positions_cpu, attn_state=attn_state) self.attn_state = attn_state # type: ignore self.with_prefill = with_prefill self.num_tokens_across_dp = num_tokens_across_dp self._update_graph_pad_size(with_prefill, maybe_padded_num_tokens) attn_metadata: dict[str, Any] = {} # _prepare_inputs may reorder the batch, so we must gather # multi-modal outputs after that to ensure the correct order if self.is_multimodal_model: # Run the multimodal encoder if any. self._execute_mm_encoder(scheduler_output) mm_embeds = self._gather_mm_embeddings(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- # NOTE(woosuk): To unify token ids and soft tokens (vision # embeddings), we always use embeddings (rather than token ids) # as input to the multimodal model, even when the input is text. input_ids = self.input_ids[:num_input_tokens] if mm_embeds: inputs_embeds = self.model.get_input_embeddings(input_ids, mm_embeds) else: inputs_embeds = self.model.get_input_embeddings(input_ids) # TODO(woosuk): Avoid the copy. Optimize. self.inputs_embeds[:num_input_tokens].copy_(inputs_embeds) inputs_embeds = self.inputs_embeds[:num_input_tokens] model_kwargs = { **self._init_model_kwargs(num_input_tokens), **self._extract_mm_kwargs(scheduler_output), } # (NOTE) Omni-new: input_ids isn't set as None # ------------------------------------------------------------------------------------------------- else: # For text-only models, we use token ids as input. # While it is possible to use embeddings as input just like the # multimodal models, it is not desirable for performance since # then the embedding layer is not included in the ACL graph. input_ids = self.input_ids[:num_input_tokens] inputs_embeds = None model_kwargs = self._init_model_kwargs(num_input_tokens) positions = self.positions[:num_input_tokens] input_ids, positions = self._update_input_ids_and_positions( input_ids, positions, num_input_tokens, with_prefill, maybe_padded_num_tokens ) if get_pp_group().is_first_rank: intermediate_tensors = None else: assert intermediate_tensors is not None assert self.intermediate_tensors is not None for k, v in intermediate_tensors.items(): self.intermediate_tensors[k][:num_input_tokens].copy_(v[:num_input_tokens], non_blocking=True) intermediate_tensors = IntermediateTensors( {k: v[:num_input_tokens] for k, v in self.intermediate_tensors.items()} ) use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 if not use_spec_decode: # NOTE(woosuk): Due to chunked prefills, the batch may contain # partial requests. While we should not sample any token # from these partial requests, we do so for simplicity. # We will ignore the sampled tokens from the partial requests. # TODO: Support prompt logprobs. spec_decode_metadata = None logits_indices = torch.from_numpy(cu_num_tokens - 1).to(self.device, non_blocking=True) else: # Get the number of draft tokens for each request. # Iterate over the dictionary rather than all requests since not all # requests have draft tokens. num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) for req_id, draft_token_ids in scheduler_output.scheduled_spec_decode_tokens.items(): req_idx = self.input_batch.req_id_to_index[req_id] num_draft_tokens[req_idx] = len(draft_token_ids) spec_decode_metadata = self._calc_spec_decode_metadata(num_draft_tokens, cu_num_tokens) logits_indices = spec_decode_metadata.logits_indices self.num_draft_tokens.np[:num_reqs] = num_draft_tokens self.num_draft_tokens.np[num_reqs:].fill(0) self.num_draft_tokens.copy_to_gpu() # Used in the below loop. # query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1] num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs] spec_decode_common_attn_metadata = None if use_spec_decode and self.need_accepted_tokens: self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] self.num_accepted_tokens.np[num_reqs:].fill(1) self.num_accepted_tokens.copy_to_gpu() # Prepare the attention metadata for each KV cache group and make layers # in the same group share the same metadata. for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups): if isinstance(kv_cache_group_spec.kv_cache_spec, EncoderOnlyAttentionSpec): # Encoder-only layers do not have KV cache, so we need to # create a dummy block table and slot mapping for them. blk_table_tensor = torch.zeros( (num_reqs, 1), dtype=torch.int32, device=self.device, ) slot_mapping = torch.zeros( (total_num_scheduled_tokens,), dtype=torch.int64, device=self.device, ) else: blk_table = self.input_batch.block_table[kv_cache_group_id] blk_table_tensor = blk_table.get_device_tensor() slot_mapping = blk_table.slot_mapping_cpu[:total_num_scheduled_tokens] self.slot_mapping[:total_num_scheduled_tokens].copy_( slot_mapping[:total_num_scheduled_tokens], non_blocking=True, ) self.slot_mapping[total_num_scheduled_tokens:].fill_(0) # Make AscendCommonAttentionMetadata common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=self.query_start_loc[: num_reqs + 1], query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1], seq_lens_cpu=self.seq_lens_cpu, seq_lens=self.seq_lens_cpu[:num_reqs], num_reqs=num_reqs, num_actual_tokens=total_num_scheduled_tokens, num_input_tokens=num_input_tokens, actual_seq_lengths_q=self.actual_seq_lengths_q, # TODO: change this to the right block table for linear attn block_table_tensor=blk_table_tensor[:num_reqs], slot_mapping=self.slot_mapping, num_computed_tokens_cpu=num_computed_tokens_cpu, positions=self.positions, attn_mask=self.attn_mask, spec_attn_mask=self.spec_attn_mask, attn_state=self.attn_state, enable_dbo_across_dp=enable_dbo, is_only_prefill=bool(np.all(num_valid_tokens != 1)), max_query_len=max_num_scheduled_tokens, graph_pad_size=self.graph_pad_size, decode_token_per_req=self.decode_token_per_req, cos=self.cos, sin=self.sin, ) if self.speculative_config and spec_decode_common_attn_metadata is None: spec_decode_common_attn_metadata = common_attn_metadata for attn_group in self.attn_groups[kv_cache_group_id]: common_prefix_len = 0 extra_attn_metadata_args = {} builder = attn_group.get_metadata_builder() if isinstance(builder, GDNAttentionMetadataBuilder) or self.model_config.runner_type == "pooling": if use_spec_decode: extra_attn_metadata_args = dict( num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs], num_draft_tokens=self.num_draft_tokens.gpu[:num_reqs], ) attn_metadata_i = builder.build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, **extra_attn_metadata_args, ) else: attn_metadata_i = builder.build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, model=self.get_model(), **extra_attn_metadata_args, ) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = attn_metadata_i if lmhead_tp_enable(): max_num_reqs_across_dp = maybe_padded_num_tokens if not with_prefill else self.max_num_reqs logits_indices = nn.functional.pad(logits_indices, (0, max_num_reqs_across_dp - logits_indices.shape[0])) return ( attn_metadata, positions, num_input_tokens, num_input_tokens, num_tokens_across_dp, maybe_padded_num_tokens, logits_indices, spec_decode_metadata, input_ids, inputs_embeds, intermediate_tensors, max_num_scheduled_tokens, model_kwargs, ) @torch.inference_mode() def execute_model( self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: with ProfileExecuteDuration().capture_async("prepare input"): self._update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: return EMPTY_MODEL_RUNNER_OUTPUT if self.dynamic_eplb: self.eplb_updator.forward_before() ( attn_metadata, positions, num_scheduled_tokens_np, num_input_tokens, num_tokens_across_dp, maybe_padded_num_tokens, logits_indices, spec_decode_metadata, input_ids, inputs_embeds, intermediate_tensors, max_query_len, model_kwargs, ) = self._prepare_inputs(scheduler_output, intermediate_tensors) if self.dynamic_eplb: self.eplb_updator.take_update_info_from_eplb_process() moe_comm_type = self._select_moe_comm_method(num_input_tokens, self.with_prefill) # -------------------------------------- Omni-new ------------------------------------------------- # Omni-new: don't use cudagraph_dispatcher # and remove ubatch_slices aclgraph_runtime_mode = CUDAGraphMode.NONE # ------------------------------------------------------------------------------------------------- # Run forward pass with ProfileExecuteDuration().capture_async("forward"): with set_ascend_forward_context( attn_metadata, self.vllm_config, num_tokens=num_input_tokens, num_tokens_across_dp=num_tokens_across_dp, with_prefill=self.with_prefill, reserved_mc2_mask=self.reserved_mc2_mask, moe_comm_type=moe_comm_type, aclgraph_runtime_mode=aclgraph_runtime_mode, batch_descriptor=None, num_actual_tokens=scheduler_output.total_num_scheduled_tokens, prefetch_stream=self.prefetch_stream, model_instance=self.model, weight_prefetch_method=self.weight_prefetch_method, ): self.maybe_setup_kv_connector(scheduler_output) outputs = self._run_diffusion( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, multimodal_kwargs=model_kwargs, logits_indices=logits_indices, ) self.maybe_wait_for_kv_save() finished_sending, finished_recving = self.get_finished_kv_transfer(scheduler_output) aux_hidden_states = None if self.drafter and self.drafter.name == SpecDcodeType.EAGLE3: hidden_states, aux_hidden_states = outputs kv_connector_output = KVConnectorOutput(finished_sending=finished_sending, finished_recving=finished_recving) finished_sending = None finished_recving = None # -------------------------------------- Omni-new ------------------------------------------------- # Omni-new: extract_multimodal_outputs _, multimodal_outputs = self.extract_multimodal_outputs(outputs) # Ensure one tensor per request, map to CPU for output struct pooler_output: list[torch.Tensor | None] = [] if isinstance(multimodal_outputs, torch.Tensor): # If model returned a single stacked tensor, split by requests assert multimodal_outputs.shape[0] == self.input_batch.num_reqs for i in range(self.input_batch.num_reqs): pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous()) elif isinstance(multimodal_outputs, list): for out in multimodal_outputs: pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) elif isinstance(multimodal_outputs, dict): for out in multimodal_outputs.values(): pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) else: raise RuntimeError("Unsupported diffusion output type") output = OmniModelRunnerOutput( req_ids=self.input_batch.req_ids, req_id_to_index=self.input_batch.req_id_to_index, sampled_token_ids=[], logprobs=None, prompt_logprobs_dict={}, pooler_output=pooler_output, kv_connector_output=kv_connector_output, num_nans_in_logits={}, ) # ------------------------------------------------------------------------------------------------- durations = ProfileExecuteDuration().pop_captured_sync() if durations: dr_str = [f"[{tag}]:{duration:.2f}ms" for tag, duration in durations.items()] captured_name = "Decode" if self.attn_state == AscendAttentionState.DecodeOnly else "Prefill" logger.info("Profile execute duration [%s]:%s", captured_name, " ".join(dr_str)) if self.dynamic_eplb: self.eplb_updator.forward_end() if not self.use_async_scheduling: return output return AsyncNPUModelRunnerOutput( model_runner_output=output, sampled_token_ids=[], invalid_req_indices=[], async_output_copy_stream=self.async_output_copy_stream, ) def _run_diffusion( self, *, input_ids: torch.Tensor, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None, multimodal_kwargs: dict, logits_indices: torch.Tensor, ) -> torch.Tensor | list[torch.Tensor]: """Runs the diffusion process and returns per-request tensors. Tries model interfaces in the following order for maximal compatibility: 1) model.sample(condition=..., **kwargs) 2) model.forward(condition=..., **kwargs) 3) model.diffuse(condition=..., **kwargs) """ kwargs = dict( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device), sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, ) if hasattr(self.model, "forward"): return self.model.forward(**kwargs) # TODO: add the diffuse method for other models raise RuntimeError( "The loaded model does not expose diffusion interfaces 'sample', " "'forward', or 'diffuse'. Please implement one of them or adapt the runner." ) @torch.inference_mode() def _dummy_run( self, num_tokens: int, with_prefill: bool = False, is_torchair_compile: bool = False, aclgraph_runtime_mode: CUDAGraphMode | None = None, force_attention: bool = False, uniform_decode: bool = False, ) -> torch.Tensor: # only support eager mode and piecewise graph now assert aclgraph_runtime_mode is None or aclgraph_runtime_mode in { CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL, } # In multi-DP scenarios, there may be situations where all DP groups are executing dummy runs. # If sequence parallelism is enabled, it is essential to ensure that num_tokens is divisible by tp_size. if self.use_aclgraph and enable_sp(self.vllm_config): tp_size = self.vllm_config.parallel_config.tensor_parallel_size num_tokens = math.ceil(num_tokens / tp_size) * tp_size # Force dummy run on prefill stage when this node is deemed as kv producer. if self.is_kv_producer and not self.is_kv_consumer: with_prefill = True # Padding for DP (num_tokens, num_tokens_across_dp, with_prefill, _) = self._sync_metadata_across_dp( num_tokens, with_prefill, False ) moe_comm_type = self._select_moe_comm_method(num_tokens, with_prefill) # If cudagraph_mode.decode_mode() == FULL and # cudagraph_mode.separate_routine(). This means that we are using # different graphs and/or modes for mixed prefill-decode batches vs. # uniform decode batches. A uniform decode batch means that all # requests have identical query length, except a potential virtual # request (shorter) in the batch account for padding. # Uniform decode batch could either be common pure decode, where # max_query_len == 1, or speculative decode, where # max_query_len == 1 + num_spec_decode_tokens. # When setting max_query_len = 1, we switch to and capture the optimized # routine of FA2 for pure decode, i.e., Flashdecode + an optimization # for GQA/MQA. max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. assert num_tokens <= self.scheduler_config.max_num_batched_tokens max_num_reqs = self.max_num_reqs if uniform_decode: num_reqs = cdiv(num_tokens, max_query_len) num_scheduled_tokens_list = [max_query_len] * num_reqs if num_tokens % max_query_len != 0: num_scheduled_tokens_list[-1] = num_tokens % max_query_len else: if with_prefill: num_reqs = num_tokens else: num_reqs = (num_tokens + self.decode_token_per_req - 1) // self.decode_token_per_req num_reqs = min(num_reqs, max_num_reqs) min_tokens_per_req = num_tokens // num_reqs num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs num_scheduled_tokens_list[-1] += num_tokens % num_reqs assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) if not self.in_profile_run and self.dynamic_eplb: self.eplb_updator.forward_before() with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): if self.is_multimodal_model: input_ids = None inputs_embeds = self.inputs_embeds[:num_tokens] else: input_ids = self.input_ids[:num_tokens] inputs_embeds = None if self.uses_mrope: positions = self.mrope_positions[:, :num_tokens] else: positions = self.positions[:num_tokens] if get_pp_group().is_first_rank: intermediate_tensors = None else: if self.intermediate_tensors is None: self.intermediate_tensors = self.model.make_empty_intermediate_tensors( batch_size=num_tokens, dtype=self.dtype, device=self.device ) intermediate_tensors = IntermediateTensors( {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} ) # filter out the valid batch descriptor _ag_mode, batch_descriptor = self.aclgraph_dispatcher.dispatch( BatchDescriptor(num_tokens=num_tokens, uniform_decode=uniform_decode) ) if aclgraph_runtime_mode is not None: # we allow forcing NONE when the dispatcher disagrees to support # warm ups for aclgraph capture assert aclgraph_runtime_mode == CUDAGraphMode.NONE or aclgraph_runtime_mode == _ag_mode, ( f"Aclgraph runtime mode mismatch at dummy_run. " f"Expected {_ag_mode}, but got {aclgraph_runtime_mode}." ) else: aclgraph_runtime_mode = _ag_mode # TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup # and not supported in ASCEND now. We could remove it in the future. attn_metadata = self._build_dummy_attn_metadata( False, num_reqs=num_reqs, num_tokens=num_tokens, max_query_len=max_query_len, aclgraph_runtime_mode=aclgraph_runtime_mode, force_attention=force_attention, ) need_dummy_logits = not self.in_profile_run and lmhead_tp_enable() if need_dummy_logits: max_num_reqs_across_dp = num_tokens if not with_prefill else max_num_reqs dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32) def dummy_compute_logits(hidden_states): return self.model.compute_logits(hidden_states[dummy_indices]) with set_ascend_forward_context( attn_metadata, self.vllm_config, num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, with_prefill=with_prefill, in_profile_run=self.in_profile_run, reserved_mc2_mask=self.reserved_mc2_mask, moe_comm_type=moe_comm_type, num_actual_tokens=0, aclgraph_runtime_mode=aclgraph_runtime_mode, batch_descriptor=batch_descriptor, prefetch_stream=self.prefetch_stream, model_instance=self.model, weight_prefetch_method=self.weight_prefetch_method, ): hidden_states = self._generate_dummy_run_hidden_states( with_prefill, is_torchair_compile, input_ids, positions, attn_metadata, num_tokens, intermediate_tensors, inputs_embeds, ) if need_dummy_logits: dummy_compute_logits(hidden_states) if self.drafter: self.drafter.dummy_run( num_tokens=num_tokens, with_prefill=with_prefill, skip_attn=True, num_reqs=num_reqs, num_tokens_across_dp=num_tokens_across_dp, aclgraph_runtime_mode=aclgraph_runtime_mode, batch_descriptor=batch_descriptor, ) if need_dummy_logits: self.drafter.model.compute_logits(hidden_states[dummy_indices]) if self.in_profile_run and self.dynamic_eplb: self.model.clear_all_moe_loads() if not self.in_profile_run and self.dynamic_eplb: self.eplb_updator.take_update_info_from_eplb_process() self.eplb_updator.forward_end() # -------------------------------------- Omni-new ------------------------------------------------- hidden_states, _ = self.extract_multimodal_outputs(hidden_states) # ------------------------------------------------------------------------------------------------- return hidden_states