From 22c8cde3284c254f9fbebe532e5780dda260d521 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 26 Feb 2026 14:26:50 +0800 Subject: [PATCH] [NPU] Upgrade to v0.16.0 (#1375) Signed-off-by: gcanlin --- docker/Dockerfile.npu | 5 +- docker/Dockerfile.npu.a3 | 5 +- docs/getting_started/installation/npu.md | 8 +- .../installation/npu/npu.inc.md | 53 +- .../models/qwen3_omni/qwen3_omni.py | 35 +- .../qwen3_omni/qwen3_omni_moe_talker.py | 2 +- .../npu/stage_configs/qwen3_omni_moe.yaml | 12 +- .../npu/worker/npu_ar_model_runner.py | 434 ++++--- .../npu/worker/npu_generation_model_runner.py | 638 ++++++---- .../platforms/npu/worker/npu_model_runner.py | 1030 ++++------------- 10 files changed, 951 insertions(+), 1271 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index c555372d9..1c65d43b0 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -7,8 +7,9 @@ WORKDIR ${APP_DIR} COPY . . -# Install vllm-omni with dev dependencies -RUN pip install --no-cache-dir -e . --no-build-isolation +RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \ + && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml +RUN pip install -v -e . ENV VLLM_WORKER_MULTIPROC_METHOD=spawn diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 413ed88c3..33b5d3c4e 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -7,8 +7,9 @@ WORKDIR ${APP_DIR} COPY . . -# Install vllm-omni with dev dependencies -RUN pip install --no-cache-dir -e . --no-build-isolation +RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \ + && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml +RUN pip install -v -e . ENV VLLM_WORKER_MULTIPROC_METHOD=spawn diff --git a/docs/getting_started/installation/npu.md b/docs/getting_started/installation/npu.md index 5ee2c0ee4..0d538f25d 100644 --- a/docs/getting_started/installation/npu.md +++ b/docs/getting_started/installation/npu.md @@ -24,6 +24,10 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c ### Build wheel from source -=== "NPU" +=== "NPU release" - --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation" + --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-release" + +=== "NPU from main" + + --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-main" diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index c473820fe..50e18697b 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -3,14 +3,11 @@ For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html). # --8<-- [end:requirements] -# --8<-- [start:installation] +# --8<-- [start:installation-release] The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images: ```bash -# Update DEVICE according to your NPUs (/dev/davinci[0-7]) -export DEVICE0=/dev/davinci0 -export DEVICE1=/dev/davinci1 # Update the vllm-ascend image # Atlas A2: # export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1 @@ -20,8 +17,10 @@ export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1 docker run --rm \ --name vllm-omni-npu \ --shm-size=1g \ - --device $DEVICE0 \ - --device $DEVICE1 \ + --device /dev/davinci0 \ + --device /dev/davinci1 \ + --device /dev/davinci2 \ + --device /dev/davinci3 \ --device /dev/davinci_manager \ --device /dev/devmm_svm \ --device /dev/hisi_hdc \ @@ -39,8 +38,9 @@ cd /vllm-workspace git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git cd vllm-omni -VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . -# OR pip install -v -e . --no-build-isolation +sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \ + && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml +pip install -v -e . export VLLM_WORKER_MULTIPROC_METHOD=spawn ``` @@ -50,11 +50,38 @@ For other installation methods (pip installation, building from source, custom D We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there. -# --8<-- [end:installation] +# --8<-- [end:installation-release] + +# --8<-- [start:installation-main] + +You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) for the status of the latest commit of vLLM-Omni main branch on NPU.) + +```bash +# Pin vLLM version to 0.16.0 +cd /vllm-workspace/vllm +git pull origin main +git fetch origin --tags +git checkout v0.16.0 + +# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time. +cd ../vllm-ascend +git pull origin main +git fetch origin --tags +git checkout e2175d9c7e62b437391dfee996b1375674ba7c18 + +# Install vLLM-Omni from the latest main branch +cd ../vllm-omni +git clone https://github.com/vllm-project/vllm-omni.git +pip install -v -e . --no-build-isolation +# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . +export VLLM_WORKER_MULTIPROC_METHOD=spawn +``` + +# --8<-- [end:installation-main] # --8<-- [start:pre-built-images] -`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash. +vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-omni](https://quay.io/repository/ascend/vllm-omni?tab=tags) and run it with bash. Supported images as following. @@ -63,9 +90,7 @@ Supported images as following. | image-tag | Atlas A2 | Ubuntu | | image-tag-a3 | Atlas A3 | Ubuntu | -vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash. - -Here's an example deployment command that has been verified on 6 x NPUs: +Here's an example deployment command that has been verified on 4 x NPUs: ```bash # Atlas A2: @@ -80,8 +105,6 @@ docker run --rm \ --device /dev/davinci1 \ --device /dev/davinci2 \ --device /dev/davinci3 \ - --device /dev/davinci4 \ - --device /dev/davinci5 \ --device /dev/davinci_manager \ --device /dev/devmm_svm \ --device /dev/hisi_hdc \ diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 32137e9ae..4bd9bd3ba 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -37,7 +37,6 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape -from vllm_omni.platforms import current_omni_platform # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -271,27 +270,6 @@ class Qwen3OmniMoeForConditionalGeneration( # ========== Stage 1: Thinker ========== if self.model_stage == "thinker": thinker_dev = self._module_device(self.thinker) - if current_omni_platform.is_npu(): - # Normalize to batched inputs if needed - _added_batch_dim = False - if input_ids is not None and input_ids.ndim == 1: - input_ids = input_ids.unsqueeze(0) - _added_batch_dim = True - if positions is not None and positions.ndim == 1: - positions = positions.unsqueeze(0) - _added_batch_dim = True - if inputs_embeds is not None and inputs_embeds.ndim == 2: - inputs_embeds = inputs_embeds.unsqueeze(0) - _added_batch_dim = True - - # Handle None input_ids - if input_ids is None: - input_ids = torch.zeros( - inputs_embeds.shape[1], - dtype=torch.long, - device=thinker_dev, - ).unsqueeze(0) - _added_batch_dim = True # Move to thinker device if input_ids is not None and input_ids.device != thinker_dev: @@ -310,22 +288,13 @@ class Qwen3OmniMoeForConditionalGeneration( "capture_layer_indices": [0, int(accept_layer)], "return_hidden_states": True, } - if current_omni_platform.is_npu(): - # TODO: remove this hack when NPU supports batched inputs properly - thinker_input_ids = input_ids[0] if input_ids is not None and _added_batch_dim else input_ids - thinker_inputs_embeds = ( - inputs_embeds[0] if inputs_embeds is not None and _added_batch_dim else inputs_embeds - ) - else: - thinker_input_ids = input_ids - thinker_inputs_embeds = inputs_embeds # Run thinker text_hidden_states, captured_layer_dict = self.thinker( - input_ids=thinker_input_ids, + input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, - inputs_embeds=thinker_inputs_embeds, + inputs_embeds=inputs_embeds, **capture_kwargs, **kwargs, ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index ac7d8120d..400696f18 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -99,8 +99,8 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config + # rope_scaling already contains rope_theta, just assign it to rope_parameters talker_config.text_config.rope_parameters = talker_config.text_config.rope_scaling - talker_config.text_config.rope_parameters["rope_theta"] = talker_config.text_config.rope_theta self.quant_config = vllm_config.quant_config self.prefix = prefix self.vllm_config = vllm_config diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml index f99ed22e8..706bee5dd 100644 --- a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml @@ -7,7 +7,7 @@ stage_args: - stage_id: 0 runtime: - devices: "0,1,2,3" + devices: "0,1" max_batch_size: 1 engine_args: model_stage: thinker @@ -21,7 +21,7 @@ stage_args: distributed_executor_backend: "mp" enable_prefix_caching: false hf_config_name: thinker_config - tensor_parallel_size: 4 + tensor_parallel_size: 2 final_output: true final_output_type: text is_comprehension: true @@ -36,14 +36,14 @@ stage_args: - stage_id: 1 runtime: - devices: "4" + devices: "2" max_batch_size: 1 engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.2 + gpu_memory_utilization: 0.6 enforce_eager: true # haven't supported talker ACL graph on NPU trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav @@ -66,7 +66,7 @@ stage_args: - stage_id: 2 runtime: - devices: "0" + devices: "2" max_batch_size: 1 engine_args: model_stage: code2wav @@ -78,7 +78,7 @@ stage_args: async_scheduling: false enable_prefix_caching: false engine_output_type: audio # Final output: audio waveform - gpu_memory_utilization: 0.1 + gpu_memory_utilization: 0.3 distributed_executor_backend: "mp" max_num_batched_tokens: 1000000 hf_config_name: thinker_config diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index fbf9e632f..b0b4ea09d 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -3,16 +3,16 @@ from __future__ import annotations -from copy import copy +from copy import copy, deepcopy from typing import Any, NamedTuple import numpy as np import torch +from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.config import CUDAGraphMode 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, get_tp_group -from vllm.forward_context import get_forward_context from vllm.logger import logger from vllm.sequence import IntermediateTensors from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -24,16 +24,16 @@ from vllm.v1.outputs import ( ) from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.structured_output.utils import apply_grammar_bitmask -from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput -from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput +from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata +from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices 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 # yapf conflicts with isort for this block # yapf: disable -from vllm_ascend.compilation.acl_graph import update_full_graph_params from vllm_ascend.ops.rotary_embedding import update_cos_sin -from vllm_ascend.utils import ProfileExecuteDuration +from vllm_ascend.utils import enable_sp from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.outputs import OmniModelRunnerOutput @@ -47,13 +47,14 @@ class ExecuteModelState(NamedTuple): scheduler_output: SchedulerOutput logits: torch.Tensor spec_decode_metadata: SpecDecodeMetadata | None + spec_decode_common_attn_metadata: AscendCommonAttentionMetadata | None hidden_states: torch.Tensor sample_hidden_states: torch.Tensor aux_hidden_states: list[torch.Tensor] | None - kv_connector_output: KVConnectorOutput | None - attn_metadata: dict[str, Any] + attn_metadata: PerLayerAttnMetadata positions: torch.Tensor ec_connector_output: ECConnectorOutput | None + cudagraph_stats: CUDAGraphStat | None multimodal_outputs: Any class NPUARModelRunner(OmniNPUModelRunner): @@ -86,10 +87,9 @@ class NPUARModelRunner(OmniNPUModelRunner): self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, - ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None: + ) -> OmniModelRunnerOutput | IntermediateTensors | None: if self.execute_model_state is not None: - raise RuntimeError("State error: sample_tokens() must be called " - "after execute_model() returns None.") + raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") # -------------------------------------- Omni-new ------------------------------------------------- # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests) @@ -101,110 +101,202 @@ class NPUARModelRunner(OmniNPUModelRunner): request_id_resolver=self._resolve_global_request_id, ) # -------------------------------------- Omni-new ------------------------------------------------- + # 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] + ): + scheduler_output = deepcopy(scheduler_output) + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + with record_function_or_nullcontext("prepare input"): + with self.synchronize_input_prep(): + # -------------------------------------- Omni-new ------------------------------------------------- + self._update_states(scheduler_output) + # ------------------------------------------------------------------------------------------------ + # Update persistent batch states. - with ProfileExecuteDuration().capture_async("prepare input"): - # -------------------------------------- Omni-new ------------------------------------------------- - self._update_states(scheduler_output) - self._decode_and_store_request_payloads(scheduler_output) - # ------------------------------------------------------------------------------------------------ - - if has_ec_transfer() and get_ec_transfer().is_producer: - with self.maybe_get_ec_connector_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) + ) 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" + 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" ) - # 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) + 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()) - if self.dynamic_eplb: - self.eplb_updator.forward_before() + ( + logits_indices, + spec_decode_metadata, + ) = self._prepare_inputs( + scheduler_output, + num_scheduled_tokens_np, + ) - (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) + 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, + ) - (input_ids, inputs_embeds, positions, intermediate_tensors, - model_kwargs, ec_connector_output) = self._preprocess(scheduler_output, - num_input_tokens, - intermediate_tensors) + ( + 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, + 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 + + 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 == 1 # TODO(lxs): fix this + ): + # 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) + + (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + 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, intermediate_tensors) # update global cos, sin update_cos_sin(positions) - - if self.dynamic_eplb: - self.eplb_updator.take_update_info_from_eplb_process() - + # 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" + 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) + 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 - 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) - - hidden_states = self._generate_process_reqs_hidden_states( - num_input_tokens, input_ids, positions, - intermediate_tensors, inputs_embeds, model_kwargs) - - 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: - hidden_states, aux_hidden_states = hidden_states - - kv_connector_output = KVConnectorOutput( - finished_sending=finished_sending, - finished_recving=finished_recving) - finished_sending = None - finished_recving = None - with ProfileExecuteDuration().capture_async("post process"): + 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, + skip_compiled=has_encoder_input, + ), + self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, + ): + hidden_states = self._model_forward( + num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs + ) + with record_function_or_nullcontext("post process"): # -------------------------------------- Omni-new ------------------------------------------------- hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) @@ -218,36 +310,36 @@ class NPUARModelRunner(OmniNPUModelRunner): else: logger.debug("[AR] execute_model: multimodal_outputs is None") # -------------------------------------- Omni-new ------------------------------------------------- - # Broadcast PP output for external_launcher (torchrun) - # to make sure we are synced across pp ranks - # TODO: Support overlapping mirco-batches - # https://github.com/vllm-project/vllm/issues/18019 - broadcast_pp_output = \ - self.parallel_config.distributed_executor_backend \ - == "external_launcher" and len(get_pp_group().ranks) > 0 - if not get_pp_group().is_last_rank: - # For mid-pipeline stages, return the hidden states. - if not broadcast_pp_output: + if self.pcp_size > 1: + # NOTE we must `slice` hidden_states because pcp_allgather_restore_idx + # ignores the padding from CUDA Graph. + hidden_states = self.pcp_manager.get_restore_hidden_states(hidden_states) + aux_hidden_states = None + if self.use_aux_hidden_state_outputs: + hidden_states, aux_hidden_states = hidden_states + + if not self.broadcast_pp_output: + # Common case. + if not get_pp_group().is_last_rank: + # Return the intermediate tensors. + assert isinstance(hidden_states, IntermediateTensors) hidden_states.kv_connector_output = kv_connector_output self.kv_connector_output = kv_connector_output if self.debugger is not None: self.debugger.stop() self.debugger.step() return hidden_states - assert isinstance(hidden_states, IntermediateTensors) - get_pp_group().send_tensor_dict( - hidden_states.tensors, all_gather_group=get_tp_group()) - logits = None - else: - if self.input_batch.pooling_params: - pool_output = self._pool( - hidden_states, - scheduler_output.total_num_scheduled_tokens, - num_scheduled_tokens_np, kv_connector_output) + if self.is_pooling_model: + # Return the pooling output. + output = self._pool( + hidden_states, num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output + ) + output.kv_connector_output = kv_connector_output if self.debugger is not None: self.debugger.stop() self.debugger.step() - return pool_output + return output + sample_hidden_states = hidden_states[logits_indices] # -------------------------------------- Omni-new ------------------------------------------------- # Try with sampling_metadata first; fall back to without for models that don't support it @@ -258,37 +350,56 @@ class NPUARModelRunner(OmniNPUModelRunner): except TypeError: logits = self.model.compute_logits(sample_hidden_states) # -------------------------------------- Omni-new ------------------------------------------------- - if broadcast_pp_output: - model_output_broadcast_data = { - "logits": logits.contiguous(), - } if logits is not None else {} - model_output_broadcast_data = get_pp_group( - ).broadcast_tensor_dict(model_output_broadcast_data, - src=len(get_pp_group().ranks) - 1) - assert model_output_broadcast_data is not None - logits = model_output_broadcast_data["logits"] + else: + # Rare case. + assert not self.is_pooling_model + + if not get_pp_group().is_last_rank: + sample_hidden_states = hidden_states[logits_indices] + get_pp_group().send_tensor_dict(hidden_states.tensors, all_gather_group=get_tp_group()) + logits = None + else: + sample_hidden_states = hidden_states[logits_indices] + # -------------------------------------- Omni-new ------------------------------------------------- + # Try with sampling_metadata first; fall back to without for models that don't support it + try: + logits = self.model.compute_logits( + sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata + ) + except TypeError: + logits = self.model.compute_logits(sample_hidden_states) + # -------------------------------------- Omni-new ------------------------------------------------- + + model_output_broadcast_data: dict[str, Any] = {} + if logits is not None: + model_output_broadcast_data["logits"] = logits.contiguous() + broadcasted = get_pp_group().broadcast_tensor_dict( + model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 + ) + assert broadcasted is not None + logits = broadcasted["logits"] # Apply structured output bitmasks if present self.execute_model_state = ExecuteModelState( scheduler_output, logits, spec_decode_metadata, + spec_decode_common_attn_metadata, hidden_states, sample_hidden_states, aux_hidden_states, - kv_connector_output, attn_metadata, positions, ec_connector_output, - multimodal_outputs, # Omni-new + cudagraph_stats, + multimodal_outputs, ) self.kv_connector_output = kv_connector_output return None @torch.inference_mode() def sample_tokens( - self, - grammar_output: GrammarOutput | None = None, + self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -298,6 +409,7 @@ class NPUARModelRunner(OmniNPUModelRunner): self.kv_extracted_req_ids = None # -------------------------------------- Omni-new ------------------------------------------------- + if self.execute_model_state is None: # Nothing to do (PP non-final rank case), output isn't used. if not kv_connector_output: @@ -316,13 +428,14 @@ class NPUARModelRunner(OmniNPUModelRunner): scheduler_output, logits, spec_decode_metadata, + spec_decode_common_attn_metadata, hidden_states, sample_hidden_states, aux_hidden_states, - kv_connector_output, attn_metadata, positions, ec_connector_output, + cudagraph_stats, multimodal_outputs, ) = self.execute_model_state # Clear ephemeral state. @@ -334,26 +447,26 @@ class NPUARModelRunner(OmniNPUModelRunner): # the apply_grammar_bitmask uses torch.compile to optimize this,ascend does not support it now logits_dtype = logits.dtype logits = logits.to("cpu").float() - apply_grammar_bitmask(scheduler_output, grammar_output, - self.input_batch, logits) + apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits) logits = logits.to(self.device).to(logits_dtype) - with ProfileExecuteDuration().capture_async("Sample"): + with record_function_or_nullcontext("sample_token"): sampler_output = self._sample(logits, spec_decode_metadata) def propose_draft_token_ids(sampled_token_ids): - assert self.spec_decode_common_attn_metadata is not None + assert spec_decode_common_attn_metadata is not None self._draft_token_ids = self.propose_draft_token_ids( sampled_token_ids, self.input_batch.sampling_metadata, scheduler_output, spec_decode_metadata, + spec_decode_common_attn_metadata, positions, scheduler_output.total_num_scheduled_tokens, hidden_states, attn_metadata, aux_hidden_states, - sample_hidden_states + sample_hidden_states, ) self._copy_draft_token_ids_to_cpu(scheduler_output) @@ -373,11 +486,13 @@ class NPUARModelRunner(OmniNPUModelRunner): spec_decode_metadata, ) - with ProfileExecuteDuration().capture_async("Draft"): + with record_function_or_nullcontext("draft_token"): if self.speculative_config: - use_padded_batch_for_eagle = self.speculative_config and \ - self.speculative_config.use_eagle() and \ - not self.speculative_config.disable_padded_drafter_batch + use_padded_batch_for_eagle = ( + self.speculative_config + and self.speculative_config.use_eagle() + and not self.speculative_config.disable_padded_drafter_batch + ) if use_padded_batch_for_eagle: # EAGLE speculative decoding can use the GPU sampled tokens # as inputs, and does not need to wait for bookkeeping to finish. @@ -389,7 +504,6 @@ class NPUARModelRunner(OmniNPUModelRunner): if has_kv_transfer_group(): get_kv_transfer_group().clear_connector_metadata() - # -------------------------------------- Omni-new ------------------------------------------------- hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None) @@ -448,28 +562,15 @@ class NPUARModelRunner(OmniNPUModelRunner): model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids # -------------------------------------- Omni-new ------------------------------------------------- - 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: - if self.debugger is not None: - assert self.debugger is not None - self.debugger.stop() - self.debugger.step() - return model_runner_output if self.debugger is not None: - assert self.debugger is not None self.debugger.stop() self.debugger.step() + if not self.use_async_scheduling: + return model_runner_output return AsyncGPUModelRunnerOutput( model_runner_output=model_runner_output, sampled_token_ids=sampler_output.sampled_token_ids, @@ -479,33 +580,6 @@ class NPUARModelRunner(OmniNPUModelRunner): vocab_size=self.input_batch.vocab_size, ) - def _generate_process_reqs_hidden_states(self, num_input_tokens, - input_ids, positions, - intermediate_tensors, - inputs_embeds, model_kwargs): - assert self.model is not None - hidden_states = self._model_forward(input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs) - - forward_context = get_forward_context() - if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL \ - and not self.use_sparse: - if self.vllm_config.model_config.use_mla: - if self.pcp_size * self.dcp_size > 1: - update_full_graph_params(self.attn_backend, self.update_stream, forward_context, - num_input_tokens, self.vllm_config, - self.vllm_config.speculative_config) - - if get_forward_context().sp_enabled and not isinstance( - hidden_states, IntermediateTensors): - hidden_states = self._all_gather_hidden_states_and_aux( - hidden_states) - return hidden_states if self.pcp_size == 1 else self.pcp_manager.get_restore_hidden_states( - hidden_states) - def _resolve_global_request_id(self, req_id: str) -> str: """Resolve global request ID from request state.""" req_state = self.requests.get(req_id) diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index d263fb0d3..1fb1ac07a 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -4,8 +4,7 @@ from __future__ import annotations import gc -import math -from copy import copy +from copy import copy, deepcopy import numpy as np import torch @@ -19,12 +18,16 @@ 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.v1.utils import record_function_or_nullcontext +from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata +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.platform import NPUPlatform -from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable +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 @@ -35,15 +38,31 @@ 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 - for _, req_id in enumerate(cached_reqs.req_ids): + 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) - self.input_batch.remove_request(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) - self._init_mrope_positions(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( @@ -53,51 +72,169 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): ) -> 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"): + # 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] + ): + scheduler_output = deepcopy(scheduler_output) + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + with record_function_or_nullcontext("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( + with self.synchronize_input_prep(): + # Update persistent batch states. + 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, + ) = self._prepare_inputs( 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) + num_scheduled_tokens_np, + ) - 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) + 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, + ) - if self.dynamic_eplb: - self.eplb_updator.forward_before() + ( + 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, + 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 + + 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 == 1 # TODO(lxs): fix this + ): + # 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) + + (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + 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, + ) ( - 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_tokens_padded, intermediate_tensors) - (input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ec_connector_output) = ( - self._preprocess(scheduler_output, num_input_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: - self.eplb_updator.take_update_info_from_eplb_process() - + # 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) @@ -106,15 +243,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): 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], @@ -122,63 +250,60 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): 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 - with ProfileExecuteDuration().capture_async("forward"): - with set_ascend_forward_context( + with ( + record_function_or_nullcontext("forward"), + set_ascend_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_input_tokens, + num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, - aclgraph_runtime_mode=aclgraph_runtime_mode, - batch_descriptor=batch_descriptor, + aclgraph_runtime_mode=cudagraph_mode, + batch_descriptor=batch_desc, 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 + skip_compiled=has_encoder_input, + ), + self.maybe_get_kv_connector_output(scheduler_output) 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, + multimodal_outputs, # Omni-new: pass multimodal_outputs to ExecuteModelState + ) + self.kv_connector_output = kv_connector_output return None @torch.inference_mode() def sample_tokens( - self, - grammar_output: GrammarOutput | None = None, + self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -201,13 +326,14 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): scheduler_output, logits, spec_decode_metadata, + spec_decode_common_attn_metadata, hidden_states, sample_hidden_states, aux_hidden_states, - kv_connector_output, attn_metadata, positions, ec_connector_output, + cudagraph_stats, multimodal_outputs, ) = self.execute_model_state # Clear ephemeral state. @@ -231,11 +357,22 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): {"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) + 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. @@ -253,20 +390,25 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) # -------------------------------------- Omni-new ------------------------------------------------- + + if self.debugger is not None: + self.debugger.stop() + self.debugger.step() 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, + logprobs_tensors=None, ) def _run_generation_model( self, *, + num_tokens_padded: int, input_ids: torch.Tensor | None, positions: torch.Tensor | None, intermediate_tensors: IntermediateTensors | None, @@ -285,6 +427,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): """ # Keep inputs identical to AR runner kwargs = dict( + num_tokens_padded=num_tokens_padded, input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, @@ -317,38 +460,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): 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, activate_lora: bool = False, is_graph_capturing: bool = False, - ) -> torch.Tensor: + num_active_loras: int = 0, + ) -> tuple[torch.Tensor, 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 - ) - + 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. @@ -363,78 +484,126 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): # 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) + 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: + assert not create_mixed_batch + 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 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) + 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 + 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) - - 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 + _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), + # `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=activate_lora, + ) + if self.use_cp: + self.pcp_manager.init_batch_info( + num_scheduled_tokens, + num_reqs, ) - - 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 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) - - # 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}." + # vllm-ascend does not support ubatch now + ubatch_slices, ubatch_slices_padded = None, None + attn_metadata: PerLayerAttnMetadata | None = None + # If force_attention is True, we always capture attention. Otherwise, + # it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + if create_mixed_batch: + raise NotImplementedError( + "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" ) - else: - cudagraph_runtime_mode = _ag_mode + 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. + 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.seq_lens.np[:num_reqs_padded] = seq_lens + self.seq_lens.np[num_reqs_padded:] = 0 + self.seq_lens.copy_to_gpu() - # 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, - ) + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, num_sampled_tokens): + num_reqs_padded = self._pad_query_start_loc_for_fia(num_tokens_padded, num_reqs_padded, 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, + ): # 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: + if self.is_multimodal_model 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: @@ -454,20 +623,22 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): 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 + # 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() - actual_tokens = num_tokens // tp_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=actual_tokens, dtype=self.dtype, device=self.device + batch_size=max_actual_tokens, dtype=self.dtype, device=self.device ) intermediate_tensors = IntermediateTensors( - {k: v[:num_tokens_padded] for k, v in self.intermediate_tensors.items()} + {k: v[:intermediate_tokens] for k, v in self.intermediate_tensors.items()} ) need_dummy_logits = not is_profile and lmhead_tp_enable() @@ -485,21 +656,32 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"): return self.drafter.model.compute_logits(hidden_states[dummy_indices]) + model_kwargs = self._init_model_kwargs() + 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, + num_actual_tokens=num_tokens_padded, aclgraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + batch_descriptor=batch_desc, model_instance=self.model, ): - hidden_states = self._generate_dummy_run_hidden_states( - input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, ) - dummy_compute_logits(hidden_states) + + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + dummy_compute_logits(hidden_states) if self.drafter: self.drafter.dummy_run( @@ -508,7 +690,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): num_reqs=num_reqs_padded, num_tokens_across_dp=num_tokens_across_dp, aclgraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + batch_descriptor=batch_desc, dummy_compute_logits=dummy_drafter_compute_logits, in_graph_capturing=not force_attention, is_profile=is_profile, @@ -518,30 +700,68 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): 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 + return hidden_states, None 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) + # 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 - - NPUPlatform.synchronize() - del hidden_states, output + self._sync_device() + del hidden_states self.encoder_cache.clear() gc.collect() diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 83f476b3c..14ec559dd 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import math -from typing import TYPE_CHECKING, Any, cast +from typing import Any import numpy as np import torch @@ -10,40 +9,30 @@ from vllm.config import CUDAGraphMode from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger -from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding -from vllm.model_executor.models.interfaces import supports_mrope -from vllm.model_executor.models.interfaces_base import VllmModelForPooling -from vllm.sampling_params import SamplingType from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv -from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm_ascend.ascend_forward_context import set_ascend_forward_context -from vllm_ascend.compilation.acl_graph import ACLGraphWrapper +from vllm.v1.worker.gpu_model_runner import PerLayerAttnMetadata +from vllm_ascend.ascend_forward_context import get_forward_context, set_ascend_forward_context +from vllm_ascend.attention.attention_v1 import AscendAttentionState +from vllm_ascend.attention.utils import using_paged_attention +from vllm_ascend.compilation.acl_graph import ACLGraphWrapper, update_full_graph_params 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 NPUModelRunner +from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE, NPUModelRunner from vllm_omni.model_executor.models.output_templates import OmniOutput - -if TYPE_CHECKING: - from vllm.v1.core.sched.output import SchedulerOutput +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner logger = init_logger(__name__) -class OmniNPUModelRunner(NPUModelRunner): - """ - Base class for NPU model runners with multimodality support. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._omni_per_req_additional_information: dict[str, dict] | None = None - self._omni_num_scheduled_tokens_np: np.ndarray | None = None - self._omni_last_model_output: object | None = None - +class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): def load_model(self, *args, **kwargs) -> None: - super().load_model(*args, **kwargs) + NPUModelRunner.load_model(self, *args, **kwargs) + # Initialize enable_sp cache to avoid get_current_vllm_config() error + # in _pad_for_sequence_parallelism during execute_model. + # This is a workaround for vllm-ascend not passing vllm_config to enable_sp(). + enable_sp(self.vllm_config) # TODO move this model specific logic to a separate class if hasattr(self.model, "talker_mtp") and self.model.talker is not None: self.talker_mtp = self.model.talker_mtp @@ -62,346 +51,6 @@ class OmniNPUModelRunner(NPUModelRunner): self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False) - def _init_mrope_positions(self, req_state: CachedRequestState): - image_grid_thw = [] - video_grid_thw = [] - second_per_grid_ts = [] - audio_feature_lengths = [] - use_audio_in_video = False - for mm_feature in req_state.mm_features: - mm_item = mm_feature.data - if mm_item is None: - continue - mm_input = mm_item.get_data() - if (t := mm_input.get("image_grid_thw")) is not None: - image_grid_thw.append(t.tolist()) - if (t := mm_input.get("video_grid_thw")) is not None: - video_grid_thw.append(t.tolist()) - if (t := mm_input.get("second_per_grid_ts")) is not None: - second_per_grid_ts.append(t) - if (t := mm_input.get("audio_feature_lengths")) is not None: - audio_feature_lengths.append(t) - # Check for use_audio_in_video - use_audio_in_video_value = mm_input.get("use_audio_in_video") - if use_audio_in_video_value is not None: - use_audio_in_video = bool(use_audio_in_video_value.item()) - - if supports_mrope(self.get_model()): - req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions( - req_state.prompt_token_ids, - mm_features=req_state.mm_features, - hf_config=self.model_config.hf_config, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - audio_feature_lengths=audio_feature_lengths, - use_audio_in_video=use_audio_in_video, - ) - else: - req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor( - req_state.prompt_token_ids, - hf_config=self.model_config.hf_config, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - audio_feature_lengths=audio_feature_lengths, - use_audio_in_video=use_audio_in_video, - ) - - def _update_states(self, scheduler_output: "SchedulerOutput") -> None: - """Update the cached states and the persistent batch with the scheduler - output. - - The updated states are used by the `_prepare_inputs` function to create - the input GPU tensors for the model. - - The SamplingMetadata is updated and copied to the GPU if there is a - new/resumed/paused/finished request in the batch. - """ - # Remove finished requests from the cached states. - for req_id in scheduler_output.finished_req_ids: - self.requests.pop(req_id, None) - self.num_prompt_logprobs.pop(req_id, None) - # Remove the finished requests from the persistent batch. - # NOTE(woosuk): There could be an edge case where finished_req_ids and - # scheduled_req_ids overlap. This happens when a request is aborted and - # then resubmitted with the same ID. In this case, we treat them as two - # distinct requests - clearing the cached states for the first request - # and handling the second as a new request. - for req_id in scheduler_output.finished_req_ids: - self.input_batch.remove_request(req_id) - - # Free the cached encoder outputs. - for mm_hash in scheduler_output.free_encoder_mm_hashes: - self.encoder_cache.pop(mm_hash, None) - - # Remove the unscheduled requests from the persistent batch. - # NOTE(woosuk): The unscheduled requests are either preempted requests - # or running requests that are not scheduled in this step. We remove - # them from the persistent batch but keep their cached states since - # they will be scheduled again sometime in the future. - 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 - # NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint, - # so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds - # apart from the forced-preemption case in reset_prefix_cache. And in - # that case we include the resumed_req_ids in the unscheduled set so - # that they get cleared from the persistent batch before being re-scheduled - # in the normal resumed request path. - unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) - # NOTE(woosuk): The persistent batch optimization assumes that - # consecutive batches contain mostly the same requests. If batches - # have low request overlap (e.g., alternating between two distinct - # sets of requests), this optimization becomes very inefficient. - for req_id in unscheduled_req_ids: - self.input_batch.remove_request(req_id) - - reqs_to_add: list[CachedRequestState] = [] - # Add new requests to the cached states. - for new_req_data in scheduler_output.scheduled_new_reqs: - req_id = new_req_data.req_id - sampling_params = new_req_data.sampling_params - pooling_params = new_req_data.pooling_params - - if sampling_params and sampling_params.sampling_type == SamplingType.RANDOM_SEED: - generator = torch.Generator(device=self.device) - generator.manual_seed(sampling_params.seed) - else: - generator = None - - if self.is_pooling_model: - assert pooling_params is not None - task = pooling_params.task - assert task is not None, "You did not set `task` in the API" - - model = cast(VllmModelForPooling, self.get_model()) - to_update = model.pooler.get_pooling_updates(task) - to_update.apply(pooling_params) - - req_state = CachedRequestState( - req_id=req_id, - prompt_token_ids=new_req_data.prompt_token_ids, - prompt_embeds=new_req_data.prompt_embeds, - mm_features=new_req_data.mm_features, - sampling_params=sampling_params, - pooling_params=pooling_params, - generator=generator, - block_ids=new_req_data.block_ids, - num_computed_tokens=new_req_data.num_computed_tokens, - output_token_ids=[], - lora_request=new_req_data.lora_request, - ) - self.requests[req_id] = req_state - - # -------------------------------------- Omni-new ------------------------------------------------- - # If prompt embeddings are provided, decode and attach to inter_data - try: - if getattr(new_req_data, "prompt_embeds", None) is not None: - payload = new_req_data.prompt_embeds - dtype = getattr(np, payload.dtype) - arr = np.frombuffer(payload.data, dtype=dtype) - arr = arr.reshape(payload.shape) - pe_cpu = torch.from_numpy(arr) - # Store temporarily on CPU; later moved to device in builder - setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - # Also replace payload with Tensor for user visibility in - # scheduler_output - try: - new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment] - except Exception: - pass - except Exception as e: - logger.error(f"Error decoding prompt embeds: {e}") - # Decode additional_information payloads (dictionary) - try: - if getattr(new_req_data, "additional_information", None) is not None: - payload_info = new_req_data.additional_information - info_dict = {} - if isinstance(payload_info, dict): - info_dict = payload_info - else: - from vllm_omni.engine import AdditionalInformationPayload - - if isinstance(payload_info, AdditionalInformationPayload): - for k, entry in payload_info.entries.items(): - if entry.tensor_data is not None: - dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) - arr = np.frombuffer(entry.tensor_data, dtype=dt) - arr = arr.reshape(entry.tensor_shape) - info_dict[k] = torch.from_numpy(arr.copy()) - else: - info_dict[k] = entry.list_data - if info_dict: - setattr( - self.requests[req_id], - "additional_information_cpu", - info_dict, - ) - except Exception as e: - logger.error(f"Error decoding additional information: {e}") - pass - # -------------------------------------- Omni-new ------------------------------------------------- - - if sampling_params and sampling_params.prompt_logprobs is not None: - self.num_prompt_logprobs[req_id] = ( - self.input_batch.vocab_size - if sampling_params.prompt_logprobs == -1 - else sampling_params.prompt_logprobs - ) - # Only relevant for models using M-RoPE (e.g, Qwen2-VL) - if self.uses_mrope: - self._init_mrope_positions(req_state) - - # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) - if self.uses_xdrope_dim > 0: - self._init_xdrope_positions(req_state) - - reqs_to_add.append(self.requests[req_id]) - - # Update the states of the running/resumed requests. - is_last_rank = get_pp_group().is_last_rank - req_data = scheduler_output.scheduled_cached_reqs - scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens - - # Wait until valid_sampled_tokens_count is copied to cpu, - # then use it to update actual num_computed_tokens of each request. - valid_sampled_token_count = self._get_valid_sampled_token_count() - - for i, req_id in enumerate(req_data.req_ids): - req_state = self.requests[req_id] - num_computed_tokens = req_data.num_computed_tokens[i] - new_block_ids = req_data.new_block_ids[i] - resumed_from_preemption = req_id in req_data.resumed_req_ids - num_output_tokens = req_data.num_output_tokens[i] - req_index = self.input_batch.req_id_to_index.get(req_id) - - if req_state.prev_num_draft_len and self.use_async_scheduling: - # prev_num_draft_len is used in async scheduling mode with - # spec decode. it indicates if need to update num_computed_tokens - # of the request. for example: - # fist step: num_computed_tokens = 0, spec_tokens = [], - # prev_num_draft_len = 0. - # second step: num_computed_tokens = 100(prompt length), - # spec_tokens = [a,b], prev_num_draft_len = 0. - # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d], - # prev_num_draft_len = 2. - # num_computed_tokens in first step and second step does't contain - # the spec tokens length, but in third step it contains the - # spec tokens length. we only need to update num_computed_tokens - # when prev_num_draft_len > 0. - if req_index is None: - req_state.prev_num_draft_len = 0 - else: - assert self.input_batch.prev_req_id_to_index is not None - prev_req_index = self.input_batch.prev_req_id_to_index[req_id] - num_accepted = valid_sampled_token_count[prev_req_index] - 1 - num_rejected = req_state.prev_num_draft_len - num_accepted - num_computed_tokens -= num_rejected - req_state.output_token_ids.extend([-1] * num_accepted) - - # Update the cached states. - req_state.num_computed_tokens = num_computed_tokens - - if not is_last_rank: - # When using PP, the scheduler sends the sampled tokens back, - # because there's no direct communication between the first- - # stage worker and the last-stage worker. - new_token_ids = req_data.new_token_ids[i] - # Add the sampled token(s) from the previous step (if any). - # This doesn't include "unverified" tokens like spec tokens. - num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens - if num_new_tokens == 1: - # Avoid slicing list in most common case. - req_state.output_token_ids.append(new_token_ids[-1]) - elif num_new_tokens > 0: - req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:]) - elif num_output_tokens < len(req_state.output_token_ids): - # Some output tokens were discarded due to a sync-KV-load - # failure. Align the cached state. - del req_state.output_token_ids[num_output_tokens:] - if req_index is not None: - end_idx = self.input_batch.num_prompt_tokens[req_index] + num_output_tokens - self.input_batch.num_tokens_no_spec[req_index] = end_idx - - # Update the block IDs. - if not resumed_from_preemption: - if new_block_ids is not None: - # Append the new blocks to the existing block IDs. - for block_ids, new_ids in zip(req_state.block_ids, new_block_ids): - block_ids.extend(new_ids) - else: - assert req_index is None - assert new_block_ids is not None - # The request is resumed from preemption. - # Replace the existing block IDs with the new ones. - req_state.block_ids = new_block_ids - - req_index = self.input_batch.req_id_to_index.get(req_id) - if req_index is None: - # The request is not in the persistent batch. - # The request was either preempted and resumed later, or was not - # scheduled in the previous step and needs to be added again. - - if self.use_async_scheduling and num_output_tokens > 0: - # We must recover the output token ids for resumed requests in the - # async scheduling case, so that correct input_ids are obtained. - resumed_token_ids = req_data.all_token_ids[req_id] - req_state.output_token_ids = resumed_token_ids[-num_output_tokens:] - - reqs_to_add.append(req_state) - continue - - # Update the persistent batch. - self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens - if new_block_ids is not None: - self.input_batch.block_table.append_row(new_block_ids, req_index) - - # For the last rank, we don't need to update the token_ids_cpu - # because the sampled tokens are already cached. - if not is_last_rank: - # Add new_token_ids to token_ids_cpu. - start_token_index = num_computed_tokens - end_token_index = num_computed_tokens + len(new_token_ids) - self.input_batch.token_ids_cpu[req_index, start_token_index:end_token_index] = new_token_ids - self.input_batch.num_tokens_no_spec[req_index] = end_token_index - - # Add spec_token_ids to token_ids_cpu. - self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens) - - # Add the new or resumed requests to the persistent batch. - # The smaller empty indices are filled first. - for request in reqs_to_add: - self.input_batch.add_request(request) - self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens) - - # Condense the batched states if there are gaps left by removed requests - self.input_batch.condense() - # Allow attention backend to reorder the batch, potentially - self._may_reorder_batch(scheduler_output) - # Refresh batch metadata with any pending updates. - self.input_batch.refresh_metadata() - - @torch.inference_mode() - def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict: - if ( - hasattr(self.model, "have_multimodal_outputs") - and self.model.have_multimodal_outputs - and isinstance(hidden_states, OmniOutput) - ): - text_hidden_states = hidden_states.text_hidden_states - multimodal_outputs = hidden_states.multimodal_outputs - - elif isinstance(hidden_states, torch.Tensor): - text_hidden_states = hidden_states - multimodal_outputs = {} - elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple): - text_hidden_states = hidden_states[0] - multimodal_outputs = {} - else: - raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") - return text_hidden_states, multimodal_outputs - @torch.inference_mode() def _dummy_run( self, @@ -411,38 +60,16 @@ class OmniNPUModelRunner(NPUModelRunner): 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, activate_lora: bool = False, is_graph_capturing: bool = False, - ) -> torch.Tensor: + num_active_loras: int = 0, + ) -> tuple[torch.Tensor, 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 - ) - + 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. @@ -457,78 +84,126 @@ class OmniNPUModelRunner(NPUModelRunner): # 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) + 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: + assert not create_mixed_batch + 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 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) + 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 + 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) - - 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 + _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), + # `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=activate_lora, + ) + if self.use_cp: + self.pcp_manager.init_batch_info( + num_scheduled_tokens, + num_reqs, ) - - 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 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) - - # 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}." + # vllm-ascend does not support ubatch now + ubatch_slices, ubatch_slices_padded = None, None + attn_metadata: PerLayerAttnMetadata | None = None + # If force_attention is True, we always capture attention. Otherwise, + # it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + if create_mixed_batch: + raise NotImplementedError( + "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" ) - else: - cudagraph_runtime_mode = _ag_mode + 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. + 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.seq_lens.np[:num_reqs_padded] = seq_lens + self.seq_lens.np[num_reqs_padded:] = 0 + self.seq_lens.copy_to_gpu() - # 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, - ) + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, num_sampled_tokens): + num_reqs_padded = self._pad_query_start_loc_for_fia(num_tokens_padded, num_reqs_padded, 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, + ): # 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: + if self.is_multimodal_model 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: @@ -548,20 +223,22 @@ class OmniNPUModelRunner(NPUModelRunner): 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 + # 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() - actual_tokens = num_tokens // tp_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=actual_tokens, dtype=self.dtype, device=self.device + batch_size=max_actual_tokens, dtype=self.dtype, device=self.device ) intermediate_tensors = IntermediateTensors( - {k: v[:num_tokens_padded] for k, v in self.intermediate_tensors.items()} + {k: v[:intermediate_tokens] for k, v in self.intermediate_tensors.items()} ) need_dummy_logits = not is_profile and lmhead_tp_enable() @@ -585,28 +262,41 @@ class OmniNPUModelRunner(NPUModelRunner): num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, in_profile_run=is_profile, - num_actual_tokens=0, + num_actual_tokens=num_tokens_padded, aclgraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + batch_descriptor=batch_desc, model_instance=self.model, ): + # ---------------------------------------Omni-new---------------------------------------------- if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): num_tokens_padded_talker_mtp = num_tokens_padded if num_tokens_padded_talker_mtp == self.max_num_tokens: num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0] - hidden_states = self.talker_mtp( + outputs = self.talker_mtp( self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp], self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp], self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp], self.text_step.gpu[:num_tokens_padded_talker_mtp], ) self.compilation_config.cache_dir = None - hidden_states = self._generate_dummy_run_hidden_states( - input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds + # ---------------------------------------Omni-new---------------------------------------------- + # NOTE: Directly call self.model() instead of self._model_forward() to match + # GPU behavior. _model_forward contains Omni-specific logic (make_omni_output) + # that requires valid runtime_additional_information, which is empty during dummy run. + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, ) - dummy_compute_logits(hidden_states) - - hidden_states, _ = self.extract_multimodal_outputs(hidden_states) + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + # ---------------------------------------Omni-new---------------------------------------------- + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + # ---------------------------------------Omni-new---------------------------------------------- + dummy_compute_logits(hidden_states) if self.drafter: self.drafter.dummy_run( @@ -615,7 +305,7 @@ class OmniNPUModelRunner(NPUModelRunner): num_reqs=num_reqs_padded, num_tokens_across_dp=num_tokens_across_dp, aclgraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + batch_descriptor=batch_desc, dummy_compute_logits=dummy_drafter_compute_logits, in_graph_capturing=not force_attention, is_profile=is_profile, @@ -627,327 +317,74 @@ class OmniNPUModelRunner(NPUModelRunner): self.eplb_updator.forward_end() return hidden_states, hidden_states - def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None: - """Decode per-request prompt_embeds and additional_information for newly - scheduled requests and store them to CPU in the request state. - This version avoids hard dependency on payload classes by duck-typing.""" - try: - new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) - if not new_reqs: - return - for nr in new_reqs: - req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) - if req_id is None: - continue - # prompt_embeds - payload_pe = getattr(nr, "prompt_embeds", None) - pe_cpu = None - if payload_pe is not None: - if isinstance(payload_pe, torch.Tensor): - pe_cpu = payload_pe.detach().to("cpu").contiguous() - else: - # Try duck-typing a payload with data/shape/dtype - data = getattr(payload_pe, "data", None) - shape = getattr(payload_pe, "shape", None) - if data is not None and shape is not None: - dt = np.dtype(getattr(payload_pe, "dtype", "float32")) - arr = np.frombuffer(data, dtype=dt) - arr = arr.reshape(shape) - pe_cpu = torch.from_numpy(arr.copy()) - if pe_cpu is not None and req_id in self.requests: - setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - # additional_information - payload_info = getattr(nr, "additional_information", None) - if payload_info is not None: - info_dict = {} - if isinstance(payload_info, dict): - info_dict = payload_info - else: - # Try duck-typing a payload with entries, each entry may have - # tensor_data/tensor_dtype/tensor_shape or list_data - entries = getattr(payload_info, "entries", None) - if isinstance(entries, dict): - for k, entry in entries.items(): - tensor_data = getattr(entry, "tensor_data", None) - if tensor_data is not None: - dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) - arr = np.frombuffer(tensor_data, dtype=dt) - arr = arr.reshape(getattr(entry, "tensor_shape", ())) - info_dict[k] = torch.from_numpy(arr.copy()) - else: - info_dict[k] = getattr(entry, "list_data", None) - if info_dict and req_id in self.requests: - setattr(self.requests[req_id], "additional_information_cpu", info_dict) - except Exception as e: - logger.error(f"Error decoding prompt_embeds / additional_information: {e}") - - def _gather_runtime_additional_information(self) -> list[dict]: - """Gather per-request additional_information stored in request state in batch order.""" - per_req_runtime_info = [] - for req_id in self.input_batch.req_ids: - req_state = self.requests.get(req_id) - info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - if info and isinstance(info, dict): - per_req_runtime_info.append(info) - if "thinker_reply_part_per_request" in info: - q = info["thinker_reply_part_per_request"] - if hasattr(q, "shape"): - logger.debug(f"[OMNI] req={req_id} has thinker_reply_part_per_request queue shape: {q.shape}") - else: - per_req_runtime_info.append({}) - return per_req_runtime_info - - def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]: - """Compute (start, end) token spans for each request within the flattened step sequence.""" - req_token_spans: list[tuple[int, int]] = [] - for req_index in range(len(self.input_batch.req_ids)): - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - req_token_spans.append((start_offset, start_offset + sched_tokens)) - return req_token_spans - - def _build_model_kwargs_extra(self) -> dict: - """Build extra keyword arguments passed to the model for this step, including: - - runtime_additional_information: per-request additional information stored in request state - """ - model_kwargs_extra: dict[str, object] = {} - try: - model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information() - except Exception as e: - logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}") - import traceback - - traceback.print_exc() - return model_kwargs_extra - - def _process_additional_information_updates( + def _model_forward( self, - hidden_states: torch.Tensor, - multimodal_outputs: object, - num_scheduled_tokens_np: np.ndarray, - scheduler_output: "SchedulerOutput", - ) -> None: - """Process model-provided per-request additional_information updates and merge into request state.""" - try: - # execute the custom postprocess function - # TODO(Peiqi): do we have a more elegant way to do this? - if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: - for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests.get(req_id) - req_infos = ( - getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - ) - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - s, e = start_offset, start_offset + sched_tokens - # only consider to store data into update dict. - hidden_states_slice = hidden_states[s:e] - update_dict = self.model.postprocess(hidden_states_slice, **req_infos) - self._merge_additional_information_update(req_id, update_dict) - except Exception as e: - logger.error( - f"Error merging for requests:{self.input_batch.req_ids} " - f"additional information update: {e}, with the multimodal_outputs " - f"as {multimodal_outputs}" - ) - import traceback - - traceback.print_exc() - - def _collect_additional_information_for_prefill( - self, - num_scheduled_tokens_np: np.ndarray, - ) -> dict[str, dict]: - """Overlay per-request prompt_embeds for the prefill portion and collect - additional_information slices for this step. Returns a map req_id -> dict.""" - for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) - num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) - prompt_len = len(req_state.prompt_token_ids) - prompt_remaining = max(0, prompt_len - num_computed_tokens) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - overlay_len = min(sched_tokens, prompt_remaining) - if overlay_len <= 0: - continue - if overlay_len > 0 and pe_cpu is not None: - src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to( - dtype=self.dtype, device=self.device, non_blocking=True - ) - start_offset = int(self.query_start_loc.cpu[req_index]) - self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) - - def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None: - for new_req in scheduler_output.scheduled_new_reqs: - payload_info = getattr(new_req, "additional_information", None) - self._merge_additional_information_update(new_req.req_id, payload_info) - - if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"): - cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {}) - if isinstance(cached_infos, dict): - for req_id, req_infos in cached_infos.items(): - self._merge_additional_information_update(req_id, req_infos) - - def _preprocess( - self, - scheduler_output: "SchedulerOutput", - num_input_tokens: int, + num_tokens_padded: int, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **model_kwargs: dict[str, Any], ): - """Align with v0.14.0 preprocess and omni's additional information handling.""" - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - is_first_rank = get_pp_group().is_first_rank - is_encoder_decoder = self.model_config.is_encoder_decoder + """Override to combine NPUModelRunner's signature with OmniGPUModelRunner's logic. - # _prepare_inputs may reorder the batch, so we must gather multi - # modal outputs after that to ensure the correct order - ec_connector_output = None + This method: + 1. Accepts num_tokens_padded as required by NPUModelRunner + 2. Injects omni-specific kwargs (runtime_additional_information) + 3. Caches model output for multimodal results + 4. Handles NPU-specific post-forward logic (graph params update, SP all-gather) + """ + # Omni-specific: build and inject extra model kwargs + model_kwargs_extra = self._build_model_kwargs_extra() - if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder: - # Run the multimodal encoder if any. - with self.maybe_get_ec_connector_output( - scheduler_output, - encoder_cache=self.encoder_cache, - ) as ec_connector_output: - self._execute_mm_encoder(scheduler_output) - mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output) + runtime_info = model_kwargs_extra.get("runtime_additional_information", []) + if runtime_info: + for i, info in enumerate(runtime_info): + if info: + logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}") - # 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. - inputs_embeds_scheduled = self.model.embed_input_ids( - self.input_ids.gpu[:num_scheduled_tokens], - multimodal_embeddings=mm_embeds, - is_multimodal=is_mm_embed, + # Call the model forward (same as NPUModelRunner) + assert self.model is not None + model_output = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + **model_kwargs_extra, + ) + + # Omni-specific: wrap output if needed + if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + model_output = self.model.make_omni_output(model_output, **model_kwargs_extra) + + # Omni-specific: cache model output for later sample_tokens + self._omni_last_model_output = model_output + + # NPU-specific: update full graph params + forward_context = get_forward_context() + assert forward_context is not None + if ( + forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL + and not forward_context.capturing + and not self.use_sparse + ): + assert positions is not None + update_full_graph_params( + self.attn_backend, + self.update_stream, + forward_context, + num_tokens_padded, + self.vllm_config, + self.speculative_config, + positions.shape[0], ) - # TODO(woosuk): Avoid the copy. Optimize. - self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) + # NPU-specific: all-gather for sequence parallelism + if get_forward_context().sp_enabled and not isinstance(model_output, IntermediateTensors): + model_output = self._all_gather_hidden_states_and_aux(model_output) - input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens) - model_kwargs = { - **self._init_model_kwargs(), - **self._extract_mm_kwargs(scheduler_output), - } - elif self.enable_prompt_embeds and is_first_rank: - # Get the input embeddings for the tokens that are not input embeds, - # then put them into the appropriate positions. - # TODO(qthequartermasterman): Since even when prompt embeds are - # enabled, (a) not all requests will use prompt embeds, and (b) - # after the initial prompt is processed, the rest of the generated - # tokens will be token ids, it is not desirable to have the - # embedding layer outside of the CUDA graph all the time. The v0 - # engine avoids this by "double compiling" the CUDA graph, once - # with input_ids and again with inputs_embeds, for all num_tokens. - # If a batch only has token ids, then including the embedding layer - # in the CUDA graph will be more performant (like in the else case - # below). - token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1) - # Some tokens ids may need to become embeds - if token_ids_idx.numel() > 0: - token_ids = self.input_ids.gpu[token_ids_idx] - tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) - self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds - - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = self._init_model_kwargs() - input_ids = self.input_ids.gpu[:num_input_tokens] - 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 CUDA graph. - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = None - model_kwargs = self._init_model_kwargs() - - if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_input_tokens] - elif self.uses_xdrope_dim > 0: - positions = self.xdrope_positions.gpu[:, :num_input_tokens] - else: - positions = self.positions.gpu[:num_input_tokens] - - if is_first_rank: - intermediate_tensors = None - else: - assert intermediate_tensors is not None - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_input_tokens, intermediate_tensors, True - ) - - if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: - # Run the encoder, just like we do with other multimodal inputs. - # For an encoder-decoder model, our processing here is a bit - # simpler, because the outputs are just passed to the decoder. - # We are not doing any prompt replacement. We also will only - # ever have a single encoder input. - encoder_outputs = self._execute_mm_encoder(scheduler_output) - model_kwargs.update({"encoder_outputs": encoder_outputs}) - - req_ids = self.input_batch.req_ids - num_scheduled_tokens_np = np.array( - [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids], - dtype=np.int32, - ) - self._omni_num_scheduled_tokens_np = num_scheduled_tokens_np - - # Note: only prefill need collect additional_information for now. - # Decode don't need per_req_additional_information anymore. - if inputs_embeds is not None: - # Prefill: overlay prompt_embeds and collect additional_information - self._collect_additional_information_for_prefill(num_scheduled_tokens_np) - - if hasattr(self.model, "has_preprocess") and self.model.has_preprocess: - # Overlay custom prompt_embeds per request for the prompt portion; - # collect additional_information (tensor/list) for prefill portion only - decode_req_ids = [] - if self.vllm_config.model_config.async_chunk: - self._update_additional_information(scheduler_output) - for req_index, req_id in enumerate(self.input_batch.req_ids): - # Try to get additional_information from multiple sources - req_state = self.requests.get(req_id) - req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - s, e = start_offset, start_offset + sched_tokens - span_len = int(e) - int(s) - - # call the custom process function - req_input_ids, req_embeds, update_dict = self.model.preprocess( - input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos - ) - - if hasattr(self.model, "talker_mtp") and span_len == 1: - last_talker_hidden, text_step = update_dict.pop("mtp_inputs") - decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1) - self.talker_mtp_input_ids.gpu[decode_slice].copy_(req_input_ids) - self.talker_mtp_inputs_embeds.gpu[decode_slice].copy_(req_embeds) - self.last_talker_hidden.gpu[decode_slice].copy_(last_talker_hidden) - self.text_step.gpu[decode_slice].copy_(text_step) - decode_req_ids.append(req_id) - - # TODO(Peiqi): the merge stage could move out from the critical path - self._merge_additional_information_update(req_id, update_dict) - - # update the inputs_embeds and input_ids - seg_len = min(span_len, req_embeds.shape[0]) - inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] - if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: - input_ids[s : s + seg_len] = req_input_ids - - # run talker mtp decode - if hasattr(self.model, "talker_mtp"): - self._talker_mtp_forward(decode_req_ids, inputs_embeds) - - return ( - input_ids, - inputs_embeds, - positions, - intermediate_tensors, - model_kwargs, - ec_connector_output, - ) + return model_output def _talker_mtp_forward(self, decode_req_ids: list[str], inputs_embeds: torch.Tensor) -> None: decode_batch_size = len(decode_req_ids) @@ -960,10 +397,11 @@ class OmniNPUModelRunner(NPUModelRunner): max_num_scheduled_tokens=1, use_cascade_attn=False, ) - req_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size] - req_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size] - last_talker_hidden = self.last_talker_hidden.gpu[:decode_batch_size] - text_step = self.text_step.gpu[:decode_batch_size] + num_tokens_padded = batch_desc.num_tokens + req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded] + req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded] + last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded] + text_step = self.text_step.gpu[:num_tokens_padded] with set_ascend_forward_context( None, self.vllm_config, aclgraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): @@ -976,53 +414,3 @@ class OmniNPUModelRunner(NPUModelRunner): inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1] update_dict = {"code_predictor_codes": code_predictor_codes_cpu[idx : idx + 1]} self._merge_additional_information_update(req_id, update_dict) - - def _model_forward( - self, - input_ids: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **model_kwargs: dict[str, Any], - ): - """Inject omni-specific kwargs into forward and cache model output""" - model_kwargs_extra = self._build_model_kwargs_extra() - - runtime_info = model_kwargs_extra.get("runtime_additional_information", []) - if runtime_info: - for i, info in enumerate(runtime_info): - if info: - logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}") - - model_output = super()._model_forward( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - **model_kwargs_extra, - ) - if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): - model_output = self.model.make_omni_output(model_output, **model_kwargs_extra) - # Cache model output so later sample_tokens can consume multimodal results. - self._omni_last_model_output = model_output - return model_output - - def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: - req_state = self.requests.get(req_id) - if req_state is None: - return - existing = getattr(req_state, "additional_information_cpu", {}) - if not isinstance(existing, dict): - existing = {} - merged = dict(existing) - for k, v in upd.items(): - if isinstance(v, torch.Tensor): - merged[k] = v.detach().to("cpu").contiguous() - elif isinstance(v, list): - merged[k] = [ - (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v - ] - else: - merged[k] = v - setattr(req_state, "additional_information_cpu", merged)