545 lines
25 KiB
Python
545 lines
25 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import math
|
|
from copy import copy
|
|
|
|
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 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.worker.gpu_model_runner import AsyncGPUModelRunnerOutput
|
|
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput
|
|
from vllm_ascend.ascend_forward_context import MoECommType, get_mc2_tokens_capacity, set_ascend_forward_context
|
|
from vllm_ascend.ops.rotary_embedding import update_cos_sin
|
|
from vllm_ascend.platform import NPUPlatform
|
|
from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable
|
|
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.worker.npu.npu_ar_model_runner import ExecuteModelState
|
|
from vllm_omni.worker.npu.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):
|
|
cached_reqs = scheduler_output.scheduled_cached_reqs
|
|
for _, req_id in enumerate(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)
|
|
self.input_batch.remove_request(req_id)
|
|
# update the request state in self.input_batch
|
|
self.input_batch.add_request(req_state)
|
|
self._init_mrope_positions(req_state)
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
) -> OmniModelRunnerOutput | IntermediateTensors:
|
|
if self.execute_model_state is not None:
|
|
raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
|
|
|
|
with ProfileExecuteDuration().capture_async("prepare input"):
|
|
# -------------------------------------- Omni-new -------------------------------------------------
|
|
if self.model_config.async_chunk:
|
|
self._update_request_states(scheduler_output)
|
|
# -------------------------------------- Omni-new -------------------------------------------------
|
|
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 scheduler_output.total_num_scheduled_tokens:
|
|
if not has_kv_transfer_group():
|
|
logger.debug("skip this step for we receive the data from remote disaggregate prefill node")
|
|
# Return empty ModelRunnerOutput if there's no work to do.
|
|
return EMPTY_MODEL_RUNNER_OUTPUT
|
|
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
|
|
|
|
if self.dynamic_eplb:
|
|
self.eplb_updator.forward_before()
|
|
|
|
(
|
|
attn_metadata,
|
|
num_scheduled_tokens_np,
|
|
num_input_tokens,
|
|
num_tokens_across_dp,
|
|
logits_indices,
|
|
spec_decode_metadata,
|
|
max_query_len,
|
|
) = self._prepare_inputs(scheduler_output)
|
|
|
|
(input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ec_connector_output) = (
|
|
self._preprocess(scheduler_output, num_input_tokens, intermediate_tensors)
|
|
)
|
|
|
|
# update global cos, sin
|
|
update_cos_sin(positions)
|
|
|
|
if self.dynamic_eplb:
|
|
self.eplb_updator.take_update_info_from_eplb_process()
|
|
|
|
# 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()
|
|
|
|
uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
|
|
scheduler_output.total_num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
|
|
)
|
|
has_lora = len(self.input_batch.lora_id_to_lora_request) > 0
|
|
aclgraph_runtime_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
|
|
num_tokens=num_input_tokens, uniform_decode=uniform_decode, has_lora=has_lora
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
# 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,
|
|
aclgraph_runtime_mode=aclgraph_runtime_mode,
|
|
batch_descriptor=batch_descriptor,
|
|
num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
|
|
model_instance=self.model,
|
|
):
|
|
self.maybe_setup_kv_connector(scheduler_output)
|
|
# -------------------------------------- Omni-new -------------------------------------------------
|
|
outputs = self._run_generation_model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
model_kwargs=model_kwargs,
|
|
logits_indices=logits_indices,
|
|
)
|
|
# -------------------------------------- Omni-new -------------------------------------------------
|
|
|
|
self.maybe_wait_for_kv_save()
|
|
finished_sending, finished_recving = self.get_finished_kv_transfer(scheduler_output)
|
|
|
|
aux_hidden_states = None
|
|
if self.use_aux_hidden_state_outputs:
|
|
outputs, aux_hidden_states = outputs
|
|
|
|
kv_connector_output = KVConnectorOutput(finished_sending=finished_sending, finished_recving=finished_recving)
|
|
finished_sending = None
|
|
finished_recving = None
|
|
|
|
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
|
|
# Apply structured output bitmasks if present
|
|
self.execute_model_state = ExecuteModelState(
|
|
scheduler_output,
|
|
None,
|
|
spec_decode_metadata,
|
|
outputs,
|
|
None,
|
|
aux_hidden_states,
|
|
kv_connector_output,
|
|
attn_metadata,
|
|
positions,
|
|
ec_connector_output,
|
|
multimodal_outputs,
|
|
)
|
|
self.kv_connector_output = kv_connector_output
|
|
return None
|
|
|
|
@torch.inference_mode()
|
|
def sample_tokens(
|
|
self,
|
|
grammar_output: GrammarOutput | None = 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.
|
|
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,
|
|
hidden_states,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
kv_connector_output,
|
|
attn_metadata,
|
|
positions,
|
|
ec_connector_output,
|
|
multimodal_outputs,
|
|
) = 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):
|
|
mm_payload = {}
|
|
for key, out in multimodal_outputs.items():
|
|
if out is not None and isinstance(out, torch.Tensor):
|
|
mm_payload[key] = out.detach().to("cpu").contiguous()
|
|
pooler_output.append(mm_payload)
|
|
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={},
|
|
ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
|
|
)
|
|
# -------------------------------------- Omni-new -------------------------------------------------
|
|
if not self.use_async_scheduling:
|
|
return output
|
|
return AsyncGPUModelRunnerOutput(
|
|
model_runner_output=output,
|
|
sampled_token_ids=torch.tensor([], device=self.device),
|
|
logprobs_tensors=None,
|
|
invalid_req_indices=[],
|
|
async_output_copy_stream=self.async_output_copy_stream,
|
|
vocab_size=self.input_batch.vocab_size,
|
|
)
|
|
|
|
def _run_generation_model(
|
|
self,
|
|
*,
|
|
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(
|
|
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,
|
|
allow_microbatching: bool = True,
|
|
skip_eplb: bool = False,
|
|
remove_lora: bool = True,
|
|
activate_lora: bool = False,
|
|
is_graph_capturing: bool = False,
|
|
) -> torch.Tensor:
|
|
# only support eager mode and piecewise graph now
|
|
assert cudagraph_runtime_mode is None or cudagraph_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
|
|
|
|
has_lora = True if self.lora_config and self.compilation_config.cudagraph_specialize_lora else False
|
|
_ag_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
|
|
num_tokens=num_tokens, uniform_decode=uniform_decode, has_lora=has_lora
|
|
)
|
|
|
|
# Padding for DP
|
|
(num_tokens, num_tokens_across_dp, with_prefill) = self._sync_metadata_across_dp(
|
|
batch_descriptor.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)
|
|
num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
|
|
|
|
if not is_profile and self.dynamic_eplb:
|
|
self.eplb_updator.forward_before()
|
|
|
|
if num_tokens != batch_descriptor.num_tokens:
|
|
_ag_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
|
|
num_tokens=num_tokens, uniform_decode=uniform_decode, has_lora=has_lora
|
|
)
|
|
|
|
num_tokens_padded = batch_descriptor.num_tokens
|
|
num_reqs_padded = batch_descriptor.num_reqs if batch_descriptor.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)
|
|
|
|
# filter out the valid batch descriptor
|
|
if cudagraph_runtime_mode is not None:
|
|
# we allow forcing NONE when the dispatcher disagrees to support
|
|
# warm ups for aclgraph capture
|
|
if cudagraph_runtime_mode != CUDAGraphMode.NONE and cudagraph_runtime_mode != _ag_mode:
|
|
raise ValueError(
|
|
f"Aclgraph runtime mode mismatch at dummy_run. "
|
|
f"Expected {_ag_mode}, but got {cudagraph_runtime_mode}."
|
|
)
|
|
else:
|
|
cudagraph_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_padded,
|
|
num_tokens=num_tokens_padded,
|
|
max_query_len=max_query_len,
|
|
aclgraph_runtime_mode=cudagraph_runtime_mode,
|
|
force_attention=force_attention,
|
|
is_graph_capturing=is_graph_capturing,
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
)
|
|
|
|
with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, num_sampled_tokens):
|
|
# Make sure padding doesn't exceed max_num_tokens
|
|
assert num_tokens_padded <= self.max_num_tokens
|
|
if self.is_multimodal_model and not self.model_config.is_encoder_decoder:
|
|
input_ids = None
|
|
inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
|
|
elif 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
|
|
|
|
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.gpu[: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.
|
|
actual_tokens = num_tokens
|
|
if enable_sp():
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
actual_tokens = num_tokens // tp_size
|
|
if self.intermediate_tensors is None:
|
|
self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
|
|
batch_size=actual_tokens, dtype=self.dtype, device=self.device
|
|
)
|
|
intermediate_tensors = IntermediateTensors(
|
|
{k: v[:num_tokens_padded] 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=0,
|
|
aclgraph_runtime_mode=cudagraph_runtime_mode,
|
|
batch_descriptor=batch_descriptor,
|
|
model_instance=self.model,
|
|
):
|
|
hidden_states = self._generate_dummy_run_hidden_states(
|
|
input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds
|
|
)
|
|
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_descriptor,
|
|
dummy_compute_logits=dummy_drafter_compute_logits,
|
|
in_graph_capturing=not force_attention,
|
|
is_profile=is_profile,
|
|
)
|
|
if is_profile and self.dynamic_eplb:
|
|
self.model.clear_all_moe_loads()
|
|
if 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
|
|
|
|
def profile_run(self) -> None:
|
|
# Trigger compilation for general shape.
|
|
with self.set_in_profile_run():
|
|
hidden_states = self._dummy_run(
|
|
self.max_num_tokens // self.pcp_size if self.pcp_size > 1 else self.max_num_tokens, with_prefill=True
|
|
)
|
|
# MC2 will consume additional NPU memory.
|
|
# Therefore, we need to run the MC2 path once here to complete its initialization,
|
|
# allowing vLLM to correctly estimate the maximum memory required.
|
|
mc2_tokens_capacity = get_mc2_tokens_capacity()
|
|
if (
|
|
self.max_num_tokens > mc2_tokens_capacity
|
|
and self._select_moe_comm_method(mc2_tokens_capacity) == MoECommType.MC2
|
|
):
|
|
self._dummy_run(mc2_tokens_capacity, with_prefill=True)
|
|
|
|
output = None
|
|
|
|
NPUPlatform.synchronize()
|
|
del hidden_states, output
|
|
self.encoder_cache.clear()
|
|
gc.collect()
|