# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations import gc import time from copy import copy, deepcopy import numpy as np import torch from vllm.config import CUDAGraphMode from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group from vllm.distributed.parallel_state import get_pp_group from vllm.logger import logger from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs 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 using_paged_attention from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner class NPUGenerationModelRunner(OmniNPUModelRunner): """Generation model runner for vLLM-omni on NPU (non-autoregressive).""" def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests for req_id in scheduler_output.finished_req_ids: self.input_batch.remove_request(req_id) scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() cached_req_ids = self.input_batch.req_id_to_index.keys() resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) for req_id in unscheduled_req_ids: self.input_batch.remove_request(req_id) cached_reqs = scheduler_output.scheduled_cached_reqs req_states = [] for req_id in cached_reqs.req_ids: req_state = self.requests.get(req_id) assert req_state is not None req_state.prompt_token_ids = cached_reqs.prompt_token_ids.get(req_id) req_states.append(req_state) # Remove the request from the current input batch only if it is still present. if req_id in self.input_batch.req_id_to_index: self.input_batch.remove_request(req_id) for req_state in req_states: # update the request state in self.input_batch self.input_batch.add_request(req_state) # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: self._init_mrope_positions(req_state) @torch.inference_mode() def execute_model( self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: if self.vllm_config.model_config.enable_return_routed_experts: capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() if self.ascend_config.profiling_chunk_config.enabled: self._sync_device() self._execution_start_time = time.perf_counter() if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") # self._draft_token_ids is None when `input_fits_in_drafter=False` # and there is no draft tokens scheduled. so it need to update the # spec_decoding info in scheduler_output with async_scheduling. # use deepcopy to avoid the modification has influence on the # scheduler_output in engine core process. # TODO(Ronald1995): deepcopy is expensive when there is a large # number of requests, optimize it later. if ( self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type] ) or ( # NOTE: This branch specifically triggers a deepcopy during the prefill phase # only for PCP (Parallel Context Processing) + Multi-Modal (MM) scenarios. # It does not affect other use cases. This is a temporary workaround and # will be removed once upstream vLLM provides native support for PCP + MM. self.pcp_size > 1 and self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder ): scheduler_output = deepcopy(scheduler_output) if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata if kv_connector_metadata is not None: get_kv_transfer_group().handle_preemptions(kv_connector_metadata) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with record_function_or_nullcontext("prepare input"): # -------------------------------------- Omni-new ------------------------------------------------- if self.model_config.async_chunk and num_scheduled_tokens: self._update_request_states(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- with self.synchronize_input_prep(): # Update persistent batch states. deferred_state_corrections_fn = self._update_states(scheduler_output) if has_ec_transfer() and get_ec_transfer().is_producer: with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) return make_empty_encoder_model_runner_output(scheduler_output) if not num_scheduled_tokens: if ( self.parallel_config.distributed_executor_backend == "external_launcher" and self.parallel_config.data_parallel_size > 1 ): # this is a corner case when both external launcher # and DP are enabled, num_scheduled_tokens could be # 0, and has_unfinished_requests in the outer loop # returns True. before returning early here we call # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. self._dummy_run(1) if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. return EMPTY_MODEL_RUNNER_OUTPUT return self.kv_connector_no_forward(scheduler_output, self.vllm_config) if self.cache_config.kv_sharing_fast_prefill: assert not self.num_prompt_logprobs, ( "--kv-sharing-fast-prefill produces incorrect " "logprobs for prompt tokens, tokens, please disable " "it when the requests need prompt logprobs" ) num_reqs = self.input_batch.num_reqs req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) ( logits_indices, spec_decode_metadata, total_num_scheduled_tokens, num_scheduled_tokens_compressed_list, ) = self._prepare_inputs( scheduler_output, num_scheduled_tokens_np, ) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens if self.pcp_size > 1: num_tokens_unpadded = self.pcp_manager.total_num_sampled_tokens_pcp cascade_attn_prefix_lens = None # Disable cascade attention when using microbatching (DBO) if self.cascade_attn_enabled and not self.parallel_config.enable_dbo: # Pre-compute cascade attention prefix lengths cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens( num_scheduled_tokens_np, self.input_batch.num_computed_tokens_cpu[:num_reqs], scheduler_output.num_common_prefix_blocks, ) ( cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, cudagraph_stats, ) = self._determine_batch_execution_and_padding( num_tokens=num_tokens_unpadded, num_reqs=num_reqs, num_scheduled_tokens_np=num_scheduled_tokens_np, max_num_scheduled_tokens=max_num_scheduled_tokens, use_cascade_attn=cascade_attn_prefix_lens is not None, force_eager=self.model_config.enforce_eager, num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), ) logger.debug( "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " "should_ubatch: %s, num_tokens_across_dp: %s", cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( should_ubatch, num_scheduled_tokens_np, num_tokens_padded, num_reqs_padded, self.parallel_config.num_ubatches, ) pad_attn = cudagraph_mode == CUDAGraphMode.FULL # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877, # there should be a corresponding 'postprocess_mamba'. However, it is called inside # '_update_states_after_model_execute', which is not overridden in vLLM-Ascend. # We simply utilize the implementation in vLLM. if self.cache_config.mamba_cache_mode == "align": # preprocess_mamba reads req_state.num_computed_tokens (CPU) # to decide copy operations, so we must apply deferred # corrections before it runs. if deferred_state_corrections_fn: deferred_state_corrections_fn() deferred_state_corrections_fn = None preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, self.mamba_state_idx, self.input_batch, self.requests, self.compilation_config.static_forward_context, self.model.get_mamba_state_copy_func(), self._get_mamba_copy_bufs(), ) # preprocess_mamba resets num_accepted_tokens_cpu to 1 # for requests whose state was copied to a new block. # Re-sync to GPU so the mamba kernel reads from the # correct initial state slot (init_token_idx = 0). self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] self.num_accepted_tokens.copy_to_gpu(num_reqs) use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices if ( cudagraph_mode == CUDAGraphMode.FULL or (enable_sp() and not self.model_config.use_mla) and self.pcp_size * self.dcp_size == 1 ): # Currently, Graph Mode and SP will both pad num_tokens, # Another possible condition is num_tokens_padded != num_tokens_unpadded # but this scope is way too big and the consequences are unpredictable num_reqs_padded = self._pad_query_start_loc_for_fia( num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs ) (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( num_tokens=num_tokens_unpadded if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) else total_num_scheduled_tokens, num_tokens_padded=num_tokens_padded, num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, max_query_len=max_num_scheduled_tokens, ubatch_slices=ubatch_slices_attn, logits_indices=logits_indices, use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, num_scheduled_tokens_np=num_scheduled_tokens_np, cascade_attn_prefix_lens=cascade_attn_prefix_lens, ) ( input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ec_connector_output, ) = self._preprocess( scheduler_output, num_tokens_padded if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) else total_num_scheduled_tokens, intermediate_tensors, ) # [Omni] Pass token counts per request for code2wav output slicing model_kwargs["seq_token_counts"] = tokens # update global cos, sin update_cos_sin(positions) if self.dynamic_eplb: with record_function_or_nullcontext("EPLB weight D2D"): self.eplb_updator.forward_before() # Set cudagraph mode to none if calc_kv_scales is true. # KV scales calculation involves dynamic operations that are incompatible # with CUDA graph capture. if self.calculate_kv_scales: # type: ignore[has-type] cudagraph_mode = CUDAGraphMode.NONE # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False # type: ignore[has-type] # prevent debugger is None if self.debugger is not None: dbg_cfg = getattr(self.debugger, "config", None) dump_level = str(getattr(dbg_cfg, "level", "L1")).upper() if dbg_cfg is not None else "L1" if dump_level in ("L0", "MIX"): self.debugger.start(model=self.model) else: self.debugger.start() if self.ascend_config.enable_async_exponential: self.sampler.do_async_exponential( b_s=logits_indices.shape[0], head_dim=self.model_config.get_vocab_size(), generators=self.input_batch.sampling_metadata.generators, ) # Encoder-decoder models can only compile the pure decode steps where no # encoder inputs are present. Use eager for the first pass. num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0 # Run forward pass clear_kv_metadata = self.speculative_config is None with ( record_function_or_nullcontext("forward"), set_ascend_forward_context( attn_metadata, self.vllm_config, num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, aclgraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, num_actual_tokens=scheduler_output.total_num_scheduled_tokens, model_instance=self.model, max_tokens_across_pcp=0 if self.pcp_size == 1 else self.pcp_manager.max_num_tokens_across_pcp, skip_compiled=has_encoder_input, ), self.maybe_get_kv_connector_output( scheduler_output, **({"defer_finalize": not clear_kv_metadata}), ) as kv_connector_output, ): # -------------------------------------- Omni-new ------------------------------------------------- outputs = self._run_generation_model( num_tokens_padded=num_tokens_padded, input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, model_kwargs=model_kwargs, logits_indices=logits_indices, ) _, multimodal_outputs = self.extract_multimodal_outputs(outputs) # -------------------------------------- Omni-new ------------------------------------------------- # Apply structured output bitmasks if present self.execute_model_state = ExecuteModelState( scheduler_output, None, spec_decode_metadata, spec_decode_common_attn_metadata, None, None, None, attn_metadata, positions, ec_connector_output, cudagraph_stats, batch_desc, multimodal_outputs, # Omni-specific: pass multimodal_outputs to ExecuteModelState ) self.kv_connector_output = kv_connector_output # Now the batch has been launched we can wait for corrections from the # previous model forward without breaking async scheduling. if deferred_state_corrections_fn: deferred_state_corrections_fn() if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): self._omni_routed_experts_d2h(scheduler_output) return None @torch.inference_mode() def sample_tokens( self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: kv_connector_output = self.kv_connector_output self.kv_connector_output = None if self.execute_model_state is None: # Nothing to do (PP non-final rank case), output isn't used. # receive sampled token ids from the last PP rank when using # async scheduling + pipeline parallelism so downstream code # (e.g., PCP input preparation) can access them. if self.use_async_scheduling and get_pp_group().world_size > 1: self._pp_receive_prev_sampled_token_ids_to_input_batch() if not kv_connector_output: return None # noqa # In case of PP with kv transfer, we need to pass through the # kv_connector_output if kv_connector_output.is_empty(): return EMPTY_MODEL_RUNNER_OUTPUT output = copy(EMPTY_MODEL_RUNNER_OUTPUT) output.kv_connector_output = kv_connector_output return output # Unpack ephemeral state. ( scheduler_output, logits, spec_decode_metadata, spec_decode_common_attn_metadata, hidden_states, sample_hidden_states, aux_hidden_states, attn_metadata, positions, ec_connector_output, cudagraph_stats, _batch_desc, multimodal_outputs, # Omni-Specific ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None # -------------------------------------- Omni-new ------------------------------------------------- pooler_output: list[object] = [] if isinstance(multimodal_outputs, torch.Tensor): assert multimodal_outputs.shape[0] == 1, ( "model should return a single tensor, to return multiple tensors, use a dict" ) assert multimodal_outputs.shape[0] == self.input_batch.num_reqs for i in range(self.input_batch.num_reqs): pooler_output.append({"model_outputs": multimodal_outputs[i].detach().to("cpu").contiguous()}) elif isinstance(multimodal_outputs, list): assert len(multimodal_outputs) == 1, ( "model should return a single list, to return multiple lists, use a dict" ) for out in multimodal_outputs: pooler_output.append( {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} ) elif isinstance(multimodal_outputs, dict): num_reqs = self.input_batch.num_reqs for i in range(num_reqs): mm_payload = {} for key, out in multimodal_outputs.items(): if isinstance(out, list): if len(out) != num_reqs: raise ValueError( f"Multimodal output list for key '{key}' has length {len(out)} " f"but expected {num_reqs} (one entry per request)." ) mm_payload[key] = out[i].detach().to("cpu").contiguous() elif isinstance(out, torch.Tensor): mm_payload[key] = out.detach().to("cpu").contiguous() else: logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") pooler_output.append(mm_payload) else: raise RuntimeError("Unsupported diffusion output type") # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() routed_experts_lists = None if self.vllm_config.model_config.enable_return_routed_experts and hasattr( self.input_batch, "num_tokens_no_spec" ): routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, sampled_token_ids=[], logprobs=None, prompt_logprobs_dict={}, pooler_output=pooler_output, kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, cudagraph_stats=cudagraph_stats, ) output.routed_experts = routed_experts_lists # -------------------------------------- Omni-new ------------------------------------------------- if self.speculative_config is not None: self.finalize_kv_connector() if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"): self._sync_device() output.execution_time_ms = (time.perf_counter() - self._execution_start_time) * 1000.0 if self.dynamic_eplb: with record_function_or_nullcontext("EPLB update"): self.eplb_updator.forward_end() if self.debugger is not None: self.debugger.stop() self.debugger.step() if not self.use_async_scheduling: return output async_output = AsyncGPUModelRunnerOutput( model_runner_output=output, sampled_token_ids=torch.tensor([], device=self.device), invalid_req_indices=[], async_output_copy_stream=self.async_output_copy_stream, vocab_size=self.input_batch.vocab_size, logprobs_tensors=None, ) self.input_batch.set_async_sampled_token_ids( async_output.sampled_token_ids_cpu, async_output.async_copy_ready_event, ) return async_output def _run_generation_model( self, *, num_tokens_padded: int, input_ids: torch.Tensor | None, positions: torch.Tensor | None, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None, model_kwargs: dict, logits_indices: torch.Tensor, ) -> torch.Tensor | list[torch.Tensor]: """Run generation from codec codes to waveforms. Args: scheduler_output: Contains codec codes in input_ids or additional info intermediate_tensors: PP intermediate tensors if applicable Returns: Audio waveforms: [batch, 1, waveform_len] or list of tensors """ # Keep inputs identical to AR runner kwargs = dict( num_tokens_padded=num_tokens_padded, input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, **model_kwargs, sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, ) if hasattr(self.model, "forward"): return self._model_forward(**kwargs) 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_sampler_run(self, hidden_states: torch.Tensor) -> None: logger.warning("Dummy sampler run is not implemented for generation model") return None @torch.inference_mode() def _dummy_run( self, num_tokens: int, with_prefill: bool = False, cudagraph_runtime_mode: CUDAGraphMode | None = None, force_attention: bool = False, uniform_decode: bool = False, is_profile: bool = False, create_mixed_batch: bool = False, allow_microbatching: bool = True, skip_eplb: bool = False, remove_lora: bool = True, is_graph_capturing: bool = False, num_active_loras: int = 0, profile_seq_lens: int | None = None, profile_cpp: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() # 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.scheduler_config.max_num_seqs if create_mixed_batch: raise NotImplementedError("create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it") elif uniform_decode: num_reqs = min(max_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 elif profile_cpp: num_reqs = 1 num_scheduled_tokens_list = [num_tokens] * num_reqs else: num_reqs = min(num_tokens, 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 if not is_profile and self.dynamic_eplb: self.eplb_updator.forward_before() num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) self.query_lens = torch.from_numpy(num_scheduled_tokens) num_tokens_unpadded = int(num_scheduled_tokens.sum()) num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) _cudagraph_mode, batch_desc, _, num_tokens_across_dp, _ = self._determine_batch_execution_and_padding( num_tokens=num_tokens_unpadded, num_reqs=num_reqs, num_scheduled_tokens_np=num_scheduled_tokens, max_num_scheduled_tokens=max_query_len, use_cascade_attn=False, allow_microbatching=allow_microbatching, force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE) or profile_cpp, # `force_uniform_decode` is used for cudagraph capture; because for # capturing mixed prefill-decode batches, we sometimes use # num_tokens == num_reqs which looks like a uniform decode batch to the # dispatcher; but we actually want to capture a piecewise cudagraph force_uniform_decode=uniform_decode, # `force_has_lora` is used for cudagraph capture; because LoRA is # activated later in the context manager, but we need to know the # LoRA state when determining the batch descriptor for capture force_has_lora=num_active_loras > 0, force_num_active_loras=num_active_loras, ) if self.use_cp: self.pcp_manager.init_batch_info( num_scheduled_tokens, num_reqs, ) if self.speculative_config: self.pcp_manager.query_lens_pcp_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) self.pcp_manager.query_lens_pcp_full.copy_to_gpu() if cudagraph_runtime_mode is None: cudagraph_runtime_mode = _cudagraph_mode else: assert cudagraph_runtime_mode == _cudagraph_mode, ( f"Cudagraph runtime mode mismatch in dummy_run. " f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}." ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs if num_tokens_across_dp is not None and num_tokens_padded != num_tokens: # pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher num_tokens_across_dp[:] = num_tokens_padded num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded) # vllm-ascend does not support ubatch now ubatch_slices, ubatch_slices_padded = None, None attn_metadata: PerLayerAttnMetadata | None = None # Build attention metadata for dummy_run if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): if create_mixed_batch: raise NotImplementedError( "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" ) self.attn_state = AscendAttentionState.DecodeOnly if self.speculative_config and self.speculative_config.method == "mtp": # `AscendAttentionState.SpecDecoding` is only designed for mla if self.vllm_config.model_config.use_mla: self.attn_state = AscendAttentionState.SpecDecoding else: self.attn_state = AscendAttentionState.ChunkedPrefill # The reason why we use a fixed seq_len rather than max_query_len is that # _npu_paged_attention_get_workspace only returns max workspace with specific # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len # in inference. This will be removed once npu_fused_infer_attention_score # outperforms _npu_paged_attention on all cases. if profile_seq_lens is not None: seq_lens = profile_seq_lens else: seq_lens = ( SEQ_LEN_WITH_MAX_PA_WORKSPACE if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) else max_query_len ) # type: ignore[assignment] self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() if self._has_gdn: self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.gdn_query_start_loc.copy_to_gpu() if not profile_cpp: num_reqs_padded = self._pad_query_start_loc_for_fia( num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs ) pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL attn_metadata, _ = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded, num_reqs=num_reqs_padded, max_query_len=max_query_len, ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, num_scheduled_tokens_np=num_scheduled_tokens, ) with self.maybe_dummy_run_with_lora( self.lora_config, num_scheduled_tokens, num_sampled_tokens, remove_lora, # TODO: The next line is a temporary workaround # to fix the accuracy issue of test_llama32_lora.py, # which is introduced by vllm-project/vllm#32005 num_active_loras=(self.lora_config.max_loras if self.lora_config is not None else num_active_loras), ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens if self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None # -------------------------------------- Omni-new ------------------------------------------------- model_kwargs = self._init_model_kwargs() # Some generation-stage models (e.g. MammothModa2DiTPipeline) require # model-specific runtime information (such as image size and conditioning # embeddings) even during the dummy profiling run that vLLM uses to # estimate KV-cache capacity. get_dummy_runtime_additional_information # provides placeholder values of the correct shape so that the profiling # run does not raise an error due to missing inputs. if hasattr(self.model, "get_dummy_runtime_additional_information"): runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs) model_kwargs["runtime_additional_information"] = runtime_addi # -------------------------------------- Omni-new ------------------------------------------------- if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] elif self.uses_xdrope_dim > 0: positions = self.xdrope_positions.gpu[:, :num_tokens_padded] else: positions = self.positions[:num_tokens_padded] # update global cos, sin update_cos_sin(positions) if get_pp_group().is_first_rank: intermediate_tensors = None else: # When PP and flashcomm1 are enabled, during dummy_run the estimated space should divide num_tokens by # tp_size; otherwise, on non-first PP ranks it would effectively perform an extra all-gather, leading # to incorrect memory estimation and potentially causing OOM. intermediate_tokens = num_tokens_padded if enable_sp(): tp_size = get_tensor_model_parallel_world_size() intermediate_tokens = (num_tokens_padded + tp_size - 1) // tp_size if self.intermediate_tensors is None: max_actual_tokens = self.max_num_tokens if enable_sp(): max_actual_tokens = (self.max_num_tokens + tp_size - 1) // tp_size self.intermediate_tensors = self.model.make_empty_intermediate_tensors( batch_size=max_actual_tokens, dtype=self.dtype, device=self.device ) intermediate_tensors = IntermediateTensors( {k: v[:intermediate_tokens] for k, v in self.intermediate_tensors.items()} ) need_dummy_logits = not is_profile and lmhead_tp_enable() max_num_reqs_across_dp = max_num_reqs * self.uniform_decode_query_len dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32) def dummy_compute_logits(hidden_states): if not need_dummy_logits: return None return self.model.compute_logits(hidden_states[dummy_indices]) def dummy_drafter_compute_logits(hidden_states): if not need_dummy_logits or self.drafter is None: return if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"): return self.drafter.model.compute_logits(hidden_states[dummy_indices]) with set_ascend_forward_context( attn_metadata, self.vllm_config, num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, in_profile_run=is_profile, num_actual_tokens=num_tokens_padded, aclgraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, model_instance=self.model, ): outputs = self.model( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, **model_kwargs, # Omni-specific ) if self.use_aux_hidden_state_outputs: hidden_states, _ = outputs else: hidden_states = outputs # -------------------------------------- Omni-new ------------------------------------------------- hidden_states, _ = self.extract_multimodal_outputs(hidden_states) # ------------------------------------------------------------------------------------------------- dummy_compute_logits(hidden_states) if self.drafter: self.drafter.dummy_run( num_tokens=num_tokens_padded, with_prefill=with_prefill, num_reqs=num_reqs_padded, num_tokens_across_dp=num_tokens_across_dp, aclgraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, dummy_compute_logits=dummy_drafter_compute_logits, in_graph_capturing=not force_attention, is_profile=is_profile, ) if is_profile and self.dynamic_eplb: target = self.model.language_model if hasattr(self.model, "language_model") else self.model target.clear_all_moe_loads() if self.dynamic_eplb: self.eplb_updator.forward_end() return hidden_states, None def profile_run(self) -> None: # Profile with multimodal encoder & encoder cache. if self.supports_mm_inputs: mm_config = self.model_config.multimodal_config if mm_config is not None and mm_config.skip_mm_profiling: logger.info("Skipping memory profiling for multimodal encoder and encoder cache.") else: mm_budget = self.mm_budget assert mm_budget is not None if (encoder_budget := mm_budget.get_encoder_budget()) > 0: if not mm_budget.mm_max_toks_per_item: # All modality limits are 0 — embedding-only mode. # Budget is non-zero for embedding storage, but # there's no encoder to profile. logger.info( "Skipping encoder profiling for embedding-only " "mode (all modality limits=0 with " "enable_mm_embeds=True).", ) else: # NOTE: Currently model is profiled with a single non-text # modality with the max possible input tokens even when # it supports multiple. dummy_modality = mm_budget.get_modality_with_max_tokens() max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[dummy_modality] logger.info( "Encoder cache will be initialized with a budget of " "%s tokens, and profiled with %s %s items of the " "maximum feature size.", encoder_budget, max_mm_items_per_batch, dummy_modality, ) # Create dummy batch of multimodal inputs. batched_dummy_mm_inputs = self._get_mm_dummy_batch( dummy_modality, max_mm_items_per_batch, ) # Run multimodal encoder. dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs) sanity_check_mm_encoder_outputs( dummy_encoder_outputs, expected_num_items=max_mm_items_per_batch, ) for i, output in enumerate(dummy_encoder_outputs): self.encoder_cache[f"tmp_{i}"] = output # Add `is_profile` here to pre-allocate communication buffers hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) output = None self._sync_device() del hidden_states self.encoder_cache.clear() gc.collect()