[NPU] Upgrade to v0.16.0 (#1375)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

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
 

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
 

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"

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 \

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,
             )

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

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

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)

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()