Merge remote-tracking branch 'origin/main' into dev/vllm-align

README.md

@@ -17,7 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
-- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
+- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.

docker/Dockerfile.npu

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1
+ARG VLLM_ASCEND_TAG=v0.17.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
-
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
-
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docs/.nav.yml

@@ -100,6 +100,7 @@ nav:
   - api/vllm_omni
 - CLI Reference: cli
 - Community:
+  - Governance: community/governance.md
   - community/*
   - Slack: https://slack.vllm.ai
   - Blog: https://blog.vllm.ai

docs/community/governance.md

@@ -0,0 +1,52 @@
+# Governance
+
+vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
+
+## Values
+
+vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
+
+### Design Values
+
+1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
+2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
+3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
+4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
+5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
+
+### Collaboration Values
+
+1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
+2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
+
+## Project Maintainers
+
+### Lead Maintainers
+
+Lead maintainers are responsible for the overall direction and strategy of the project:
+
+- [@Gaohan123](https://github.com/Gaohan123)
+- [@hsliuustc0106](https://github.com/hsliuustc0106)
+- [@ywang96](https://github.com/ywang96)
+
+### Active Committers
+
+Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
+
+- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
+- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
+- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
+- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@princepride](https://github.com/princepride): Diffusion and Omni Support
+- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
+- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
+- [@wtomin](https://github.com/wtomin):
+- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
+
+## Meetings
+
+Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
+
+## Committer Nomination Process
+
+Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.

docs/configuration/README.md

@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
 
+## Multi-Stage Recipes
+
+- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
+
 ## Optimization Features
 
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss

docs/configuration/pd_disaggregation.md

@@ -0,0 +1,176 @@
+# Prefill-Decode (PD) Disaggregation
+
+PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
+stages so prompt processing and token generation can run on different workers.
+
+This is documented as a stage-config recipe instead of a bundled YAML because the
+deployment-specific values usually change per environment:
+
+- GPU placement
+- `tensor_parallel_size`
+- connector backend and connector ports
+- connector IPs or bootstrap addresses
+
+Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
+and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
+changes below.
+
+## Requirements
+
+- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
+- A KV connector supported by vLLM, such as `MooncakeConnector`
+- Matching `tensor_parallel_size` on the prefill and decode thinker stages
+
+## 1. Split the thinker into prefill and decode stages
+
+Replace the original thinker stage with two stages:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_prefill_only: true
+    runtime:
+      devices: "0"
+      max_batch_size: 16
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_producer"
+        kv_rank: 0
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25201
+    final_output: false
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm
+    is_decode_only: true
+    runtime:
+      devices: "1"
+      max_batch_size: 64
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_consumer"
+        kv_rank: 1
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25202
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+```
+
+Notes:
+
+- `is_prefill_only: true` marks the thinker stage that only saves KV.
+- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
+- `kv_transfer_config` is required on both stages.
+- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
+  prefill side only processes the prompt and exports KV.
+
+## 2. Shift the downstream stages by one index
+
+After inserting the extra thinker stage, renumber the remaining stages:
+
+```yaml
+  - stage_id: 2
+    runtime:
+      devices: "2"
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+
+  - stage_id: 3
+    runtime:
+      devices: "2"
+      max_batch_size: 1
+    engine_input_source: [2]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+```
+
+Compared with the default Qwen3-Omni config:
+
+- the talker becomes stage `2` instead of stage `1`
+- the code2wav stage becomes stage `3` instead of stage `2`
+- the talker now reads from decode stage `1`
+
+## 3. Add runtime edges for the four-stage pipeline
+
+```yaml
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1
+    - from: 2
+      to: 3
+      window_size: -1
+```
+
+## 4. Launch with your custom config
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --stage-configs-path /path/to/qwen3_omni_pd.yaml
+```
+
+## Operational Notes
+
+- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
+  pair. Keep prefill and decode at the same `tensor_parallel_size`.
+- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
+  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
+- Choose connector ports and addresses that match your deployment. The values
+  shown above are examples only.

docs/getting_started/installation/npu/npu.inc.md

@@ -33,18 +33,25 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+cd /vllm-workspace/vllm
+git pull origin main
+git fetch origin --tags
+git checkout v0.16.0
+
 # Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
 # we have to pin vllm-ascend at the current commit.
 cd /vllm-workspace/vllm-ascend
+git pull origin main
 git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
 pip install -v -e .
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
-
 cd vllm-omni
-pip install -v -e .
+pip install -v -e . --no-build-isolation
+# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
@@ -61,22 +68,22 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 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 #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.16.0
+# Pin vLLM version to 0.17.0
 cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
-git checkout v0.16.0
+git checkout v0.17.0
 VLLM_TARGET_DEVICE=empty pip install -v -e .
 
 # 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
+cd /vllm-workspace/vllm-ascend
 git pull origin main
 git fetch origin --tags
-git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+git checkout v0.17.0
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
-cd ../vllm-omni
+cd /vllm-workspace/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 .

vllm_omni/distributed/kv_transfer/__init__.py

@@ -0,0 +1,13 @@
+"""Patched KV transfer connectors for PD disaggregation.
+
+This package provides monkey-patched versions of vLLM's native KV transfer
+connectors (e.g. MooncakeConnector) that fix the request-ID mismatch problem
+in prefill-decode disaggregation.
+
+vLLM's ``InputProcessor.assign_request_id()`` appends a random 8-char suffix
+to each request ID internally.  The prefill engine stores KV under its own
+suffix, but the decode engine generates a *different* suffix — so it can never
+find the KV data.  The patched connector threads the prefill engine's internal
+``remote_request_id`` through ``kv_transfer_params`` so the decode side can
+reference the correct KV entry.
+"""

vllm_omni/distributed/kv_transfer/monkey_patch.py

@@ -0,0 +1,197 @@
+"""Monkey-patch vLLM's MooncakeConnector to fix request-ID mismatch in PD disaggregation.
+
+vLLM's InputProcessor appends a random suffix to each request ID. The prefill
+engine stores KV under its suffix, but the decode engine generates a different
+suffix. This patch threads ``remote_request_id`` through ``kv_transfer_params``
+so the decode side references the correct KV entry.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from dataclasses import dataclass
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_patched: bool = False
+
+
+@dataclass
+class PatchedRecvReqMeta:
+    """Receive-request metadata carrying the prefill engine's request ID."""
+
+    request_id: str
+    remote_request_id: str
+    local_block_ids: list[int]
+    kv_transfer_params: dict[str, Any]
+
+
+def _import_mooncake_module():
+    """Import MooncakeConnector module, supporting both vLLM >=0.16 and older."""
+    try:
+        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import mooncake_connector
+
+        return mooncake_connector
+    except ImportError:
+        pass
+    try:
+        from vllm.distributed.kv_transfer.kv_connector.v1 import mooncake_connector
+
+        return mooncake_connector
+    except ImportError:
+        return None
+
+
+def _create_patched_mooncake_connector():
+    """Return a subclass of MooncakeConnector with remote_request_id support."""
+    try:
+        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
+            MooncakeConnector as _OriginalMooncakeConnector,
+        )
+    except (ImportError, AttributeError):
+        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector import (
+            MooncakeConnector as _OriginalMooncakeConnector,
+        )
+
+    class PatchedMooncakeConnector(_OriginalMooncakeConnector):
+        """Fixes request-ID mismatch in PD disaggregation by injecting
+        remote_request_id on the prefill side and using it for KV lookup
+        on the decode side.
+        """
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            super().__init__(*args, **kwargs)
+            self.remote_to_local_req: dict[str, str] = {}
+            logger.info("[PatchedMooncakeConnector] Initialized")
+
+        def request_finished(
+            self,
+            request: Any,
+            block_ids: list[int],
+        ) -> tuple[bool, dict[str, Any] | None]:
+            result = super().request_finished(request, block_ids)
+
+            if isinstance(result, tuple) and len(result) == 2:
+                delay_free, kv_params = result
+            else:
+                delay_free, kv_params = False, result
+
+            # Normalise _reqs_need_send values
+            req_id = getattr(request, "request_id", None)
+            if req_id and hasattr(self, "_reqs_need_send"):
+                entry = self._reqs_need_send.get(req_id)
+                if isinstance(entry, tuple) and len(entry) == 2:
+                    self._reqs_need_send[req_id] = entry[1]
+
+            # Inject remote_request_id into kv_transfer_params
+            if kv_params is not None and isinstance(kv_params, dict):
+                kv_params["remote_request_id"] = req_id or "NOT_SET"
+                if hasattr(self, "side_channel_host"):
+                    kv_params.setdefault("remote_host", self.side_channel_host)
+                if hasattr(self, "side_channel_port"):
+                    kv_params.setdefault("remote_port", self.side_channel_port)
+
+            return delay_free, kv_params
+
+        def add_new_req(
+            self,
+            request_id: str,
+            local_block_ids: list[int],
+            kv_transfer_params: dict[str, Any] | None = None,
+            **kwargs: Any,
+        ) -> None:
+            super().add_new_req(request_id, local_block_ids, kv_transfer_params, **kwargs)
+
+            kv_transfer_params = kv_transfer_params or {}
+            load_remote_cache = kv_transfer_params.get(
+                "do_remote_prefill",
+                kv_transfer_params.get("load_remote_cache", False),
+            )
+
+            if load_remote_cache:
+                remote_request_id = kv_transfer_params.get("remote_request_id", request_id)
+                meta = PatchedRecvReqMeta(
+                    request_id=request_id,
+                    remote_request_id=remote_request_id,
+                    local_block_ids=local_block_ids,
+                    kv_transfer_params=kv_transfer_params,
+                )
+                if not hasattr(self, "_reqs_need_recv"):
+                    self._reqs_need_recv = {}
+                self._reqs_need_recv[request_id] = meta
+
+        def group_kv_pull(self, metadata: Any | None = None) -> None:
+            """Use remote_request_id as ZMQ lookup key via save-patch-restore."""
+            if not hasattr(self, "_reqs_need_recv") or not self._reqs_need_recv:
+                return
+
+            original_recv = self._reqs_need_recv.copy()
+            patched_recv: dict[str, Any] = {}
+
+            for local_id, meta in original_recv.items():
+                if isinstance(meta, PatchedRecvReqMeta):
+                    remote_id = meta.remote_request_id
+                    self.remote_to_local_req[remote_id] = local_id
+                    patched_meta = type(meta)(
+                        request_id=remote_id,
+                        remote_request_id=remote_id,
+                        local_block_ids=meta.local_block_ids,
+                        kv_transfer_params=meta.kv_transfer_params,
+                    )
+                    patched_recv[remote_id] = patched_meta
+                else:
+                    patched_recv[local_id] = meta
+
+            self._reqs_need_recv = patched_recv
+            super().group_kv_pull(metadata)
+
+            # Restore unconsumed entries to original local keys
+            for remote_id, local_id in list(self.remote_to_local_req.items()):
+                if remote_id in self._reqs_need_recv:
+                    entry = self._reqs_need_recv.pop(remote_id)
+                    self._reqs_need_recv[local_id] = original_recv.get(local_id, entry)
+
+        def receive_kv(self, path: Any = None, req_blocks: Any = None) -> Any:
+            result = super().receive_kv(path, req_blocks)
+
+            if self.remote_to_local_req:
+                completed = [
+                    rid
+                    for rid, lid in self.remote_to_local_req.items()
+                    if not hasattr(self, "_reqs_need_recv") or lid not in self._reqs_need_recv
+                ]
+                for remote_id in completed:
+                    self.remote_to_local_req.pop(remote_id, None)
+
+            return result
+
+    PatchedMooncakeConnector.__qualname__ = _OriginalMooncakeConnector.__qualname__
+
+    return PatchedMooncakeConnector
+
+
+def apply_mooncake_connector_patch() -> bool:
+    """Replace vLLM's MooncakeConnector with the patched version."""
+    global _patched
+    if _patched:
+        return True
+
+    _mc_module = _import_mooncake_module()
+    if _mc_module is None:
+        logger.warning("[monkey_patch] Cannot import MooncakeConnector — patch NOT applied.")
+        return False
+
+    _OriginalClass = _mc_module.MooncakeConnector
+
+    PatchedClass = _create_patched_mooncake_connector()
+
+    _mc_module.MooncakeConnector = PatchedClass
+    for _, module in sys.modules.items():
+        if hasattr(module, "MooncakeConnector") and module.MooncakeConnector is _OriginalClass:
+            module.MooncakeConnector = PatchedClass
+
+    _patched = True
+    logger.info("[monkey_patch] MooncakeConnector patch applied")
+    return True

vllm_omni/entrypoints/pd_utils.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PD (Prefill-Decode) disaggregation helpers.
+
+Mixin for OmniBase — keeps omni.py focused on orchestration.
+"""
+
+import logging
+import threading
+from dataclasses import asdict, is_dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from vllm import SamplingParams
+
+    from vllm_omni.entrypoints.omni_stage import OmniStage
+
+_DEFAULT_MOONCAKE_BOOTSTRAP_PORT = 25201
+
+logger = logging.getLogger(__name__)
+
+
+class PDDisaggregationMixin:
+    """Mixin supplying PD disaggregation helpers to OmniBase."""
+
+    def _init_pd_state(self) -> None:
+        """Initialise PD disaggregation state."""
+        self._pd_separation_pair: tuple[int, int] | None = self._detect_pd_separation()
+        self._pd_connector_info: dict[str, Any] | None = None
+        self._pd_kv_params_by_req: dict[str, dict[str, Any]] = {}
+        self._pd_kv_params_lock = threading.Lock()
+        if self._pd_separation_pair is not None:
+            self._validate_pd_separation_config()
+            self._pd_connector_info = self._get_pd_connector_info()
+            p_id, d_id = self._pd_separation_pair
+            logger.info(
+                "[%s] PD disaggregation detected: prefill=stage-%d, decode=stage-%d",
+                self._name,
+                p_id,
+                d_id,
+            )
+
+    def _detect_pd_separation(self) -> tuple[int, int] | None:
+        """Scan stage_list for a prefill/decode pair. Returns (p_id, d_id) or None."""
+        prefill_by_id: dict[int, int] = {}
+        decode_indices: list[int] = []
+        for i, stage in enumerate(self.stage_list):
+            if getattr(stage, "is_prefill_only", False):
+                prefill_by_id[i] = i
+                sid = getattr(stage, "stage_id", i)
+                if sid != i:
+                    prefill_by_id[sid] = i
+            if getattr(stage, "is_decode_only", False):
+                decode_indices.append(i)
+
+        pd_pairs: list[tuple[int, int]] = []
+        for j in decode_indices:
+            source_ids = getattr(self.stage_list[j], "engine_input_source", [])
+            for src in source_ids:
+                if src in prefill_by_id:
+                    pd_pairs.append((prefill_by_id[src], j))
+                    break
+
+        if len(pd_pairs) > 1:
+            raise ValueError(
+                f"Multiple PD pairs detected ({pd_pairs}); only a single PD pair per pipeline is supported"
+            )
+        return pd_pairs[0] if pd_pairs else None
+
+    @staticmethod
+    def _to_dict(obj: Any, default: Any = None) -> dict[str, Any] | None:
+        """Convert obj to a plain dict, trying several strategies."""
+        if obj is None:
+            return default
+        if isinstance(obj, dict):
+            return obj
+        if is_dataclass(obj):
+            try:
+                return asdict(obj)
+            except Exception:
+                return default
+        for attr in ("model_dump", "dict"):
+            if hasattr(obj, attr):
+                try:
+                    return getattr(obj, attr)()
+                except Exception:
+                    pass
+        if hasattr(obj, "items"):
+            try:
+                return dict(obj)
+            except Exception:
+                pass
+        try:
+            return dict(obj)
+        except Exception:
+            try:
+                return vars(obj)
+            except Exception:
+                logger.debug("Unable to convert %s to dict", type(obj))
+                return default
+
+    def _kv_cfg_to_dict(self, kv_cfg: Any) -> dict[str, Any]:
+        return self._to_dict(kv_cfg, default={}) or {}
+
+    def _normalize_kv_transfer_params(self, kv_params: Any) -> dict[str, Any] | None:
+        return self._to_dict(kv_params)
+
+    def _validate_pd_separation_config(self) -> None:
+        """Validate PD stage configurations are consistent."""
+        assert self._pd_separation_pair is not None
+        p_id, d_id = self._pd_separation_pair
+        p_stage = self.stage_list[p_id]
+        d_stage = self.stage_list[d_id]
+
+        def _get_kv_cfg(stage: "OmniStage") -> dict[str, Any]:
+            ea = stage.engine_args
+            cfg = getattr(ea, "kv_transfer_config", None)
+            if cfg is None:
+                cfg = ea.get("kv_transfer_config", None) if hasattr(ea, "get") else None
+            if cfg is None:
+                raise ValueError(f"Stage-{stage.stage_id} is marked for PD but has no 'kv_transfer_config'")
+            cfg_dict = self._kv_cfg_to_dict(cfg)
+            if not cfg_dict:
+                raise ValueError(f"Stage-{stage.stage_id} kv_transfer_config could not be parsed")
+            return cfg_dict
+
+        p_cfg = _get_kv_cfg(p_stage)
+        d_cfg = _get_kv_cfg(d_stage)
+
+        p_role = p_cfg.get("kv_role")
+        d_role = d_cfg.get("kv_role")
+        if p_role not in ("kv_producer", "kv_both"):
+            raise ValueError(f"Prefill stage-{p_id} kv_role must be 'kv_producer' or 'kv_both', got '{p_role}'")
+        if d_role not in ("kv_consumer", "kv_both"):
+            raise ValueError(f"Decode stage-{d_id} kv_role must be 'kv_consumer' or 'kv_both', got '{d_role}'")
+
+        d_sources = list(getattr(d_stage, "engine_input_source", []) or [])
+        if p_id not in d_sources and p_stage.stage_id not in d_sources:
+            raise ValueError(f"Decode stage-{d_id} must list prefill stage-{p_id} in engine_input_source")
+
+        p_conn = p_cfg.get("kv_connector")
+        d_conn = d_cfg.get("kv_connector")
+        if p_conn != d_conn:
+            raise ValueError(f"PD connector mismatch: prefill uses '{p_conn}', decode uses '{d_conn}'")
+        if not p_conn:
+            raise ValueError("PD requires kv_connector in kv_transfer_config")
+
+        for key in ("kv_buffer_device", "kv_buffer_size"):
+            p_val = p_cfg.get(key)
+            d_val = d_cfg.get(key)
+            if p_val is not None and d_val is not None and p_val != d_val:
+                raise ValueError(f"PD {key} mismatch: prefill='{p_val}', decode='{d_val}'")
+
+        p_tp = getattr(getattr(p_stage, "engine_args", None), "tensor_parallel_size", 1)
+        d_tp = getattr(getattr(d_stage, "engine_args", None), "tensor_parallel_size", 1)
+        if p_tp != d_tp:
+            raise ValueError(f"PD stages must have matching tensor_parallel_size: prefill={p_tp}, decode={d_tp}")
+
+    def _get_pd_connector_info(self) -> dict[str, Any] | None:
+        """Extract prefill engine KV connector info."""
+        if self._pd_separation_pair is None:
+            return None
+
+        p_id, _ = self._pd_separation_pair
+        p_stage = self.stage_list[p_id]
+
+        ea = p_stage.engine_args
+        kv_cfg = getattr(ea, "kv_transfer_config", None)
+        if kv_cfg is None and hasattr(ea, "get"):
+            kv_cfg = ea.get("kv_transfer_config")
+        if kv_cfg is None:
+            return None
+
+        kv_cfg_dict = self._kv_cfg_to_dict(kv_cfg)
+        if not kv_cfg_dict:
+            return None
+
+        kv_connector = str(kv_cfg_dict.get("kv_connector", "") or "")
+        extra_cfg = kv_cfg_dict.get("kv_connector_extra_config", {}) or {}
+        if not isinstance(extra_cfg, dict):
+            extra_cfg = self._kv_cfg_to_dict(extra_cfg)
+
+        info: dict[str, Any] = {}
+
+        if "mooncake" in kv_connector.lower():
+            bootstrap_port = extra_cfg.get("mooncake_bootstrap_port", _DEFAULT_MOONCAKE_BOOTSTRAP_PORT)
+            kv_ip = kv_cfg_dict.get("kv_ip") or "127.0.0.1"
+            info["prefill_bootstrap_addr"] = f"{kv_ip}:{bootstrap_port}"
+
+        return info
+
+    def _prepare_prefill_sampling_params(self, req_id: str, sp: "SamplingParams") -> "SamplingParams":
+        sp = sp.clone()
+        sp.max_tokens = 1
+        if hasattr(sp, "min_tokens"):
+            try:
+                sp.min_tokens = 1
+            except Exception:
+                pass
+        # MooncakeConnector cancels KV transfer unless finish_reason='length'
+        sp.stop = []
+        sp.stop_token_ids = []
+        sp.include_stop_str_in_output = False
+        if sp.extra_args is None:
+            sp.extra_args = {}
+        kv_params = self._normalize_kv_transfer_params(sp.extra_args.get("kv_transfer_params"))
+        merged: dict[str, Any] = {}
+        if kv_params:
+            merged.update(kv_params)
+        merged.update(
+            {
+                "do_remote_decode": True,
+                "do_remote_prefill": False,
+                "transfer_id": f"xfer-{req_id}",
+            }
+        )
+        sp.extra_args["kv_transfer_params"] = merged
+        logger.debug("[PD] prefill SP: req=%s kv_transfer_params=%s", req_id, merged)
+        return sp
+
+    def _pop_pd_kv_params(self, req_id: str, fallback: Any | None = None) -> dict[str, Any] | None:
+        kv_params = self._normalize_kv_transfer_params(fallback)
+        with self._pd_kv_params_lock:
+            stored = self._pd_kv_params_by_req.pop(req_id, None)
+        if kv_params is None:
+            kv_params = stored
+        return kv_params
+
+    def _drop_pd_kv_params(self, req_id: str) -> None:
+        with self._pd_kv_params_lock:
+            self._pd_kv_params_by_req.pop(req_id, None)
+
+    def _extract_kv_transfer_params(self, engine_outputs: Any) -> dict[str, Any] | None:
+        """Extract kv_transfer_params from engine outputs (if available)."""
+        outputs = engine_outputs if isinstance(engine_outputs, list) else [engine_outputs]
+        for output in outputs:
+            kv_params = getattr(output, "kv_transfer_params", None)
+            if kv_params is not None:
+                return self._normalize_kv_transfer_params(kv_params)
+        return None
+
+    def _is_pd_routing(self, stage_id: int, next_stage_id: int) -> bool:
+        """True when edge stage_id → next_stage_id is the prefill→decode boundary."""
+        return self._pd_separation_pair is not None and self._pd_separation_pair == (
+            stage_id,
+            next_stage_id,
+        )
+
+    def _maybe_expand_sampling_params(self, sampling_params_list: list) -> list:
+        """Auto-duplicate thinker SP for decode stage when user provides N-1 params."""
+        if self._pd_separation_pair is None:
+            return sampling_params_list
+        if len(sampling_params_list) != len(self.stage_list) - 1:
+            return sampling_params_list
+        p_id, d_id = self._pd_separation_pair
+        sp_list = list(sampling_params_list)
+        sp_list.insert(d_id, sp_list[p_id])
+        return sp_list
+
+    def _build_decode_kv_params(
+        self,
+        req_id: str,
+        sp_next: "SamplingParams",
+        prefill_kv_params: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Build decode-side kv_transfer_params dict."""
+        decode_kv_params: dict[str, Any] = {
+            "do_remote_decode": False,
+            "do_remote_prefill": True,
+            "transfer_id": f"xfer-{req_id}",
+        }
+
+        if sp_next.extra_args:
+            existing = self._normalize_kv_transfer_params(sp_next.extra_args.get("kv_transfer_params"))
+            if existing:
+                decode_kv_params.update(existing)
+
+        if self._pd_connector_info:
+            baddr = self._pd_connector_info.get("prefill_bootstrap_addr")
+            if baddr is not None and "remote_bootstrap_addr" not in decode_kv_params:
+                decode_kv_params["remote_bootstrap_addr"] = baddr
+
+        if prefill_kv_params:
+            decode_kv_params.update(prefill_kv_params)
+
+        decode_kv_params["do_remote_prefill"] = True
+        decode_kv_params["do_remote_decode"] = False
+        if not decode_kv_params.get("transfer_id"):
+            decode_kv_params["transfer_id"] = f"xfer-{req_id}"
+
+        return decode_kv_params
+
+    def _prepare_pd_decode_routing(
+        self,
+        req_id: str,
+        prompt: Any,
+        sp_template: "SamplingParams",
+        prefill_kv_fallback: Any | None = None,
+        engine_outputs: Any | None = None,
+    ) -> tuple[list, "SamplingParams"]:
+        """Prepare next_inputs and sampling params for prefill→decode routing.
+
+        Used by both sync (omni.py) and async (async_omni.py) paths.
+
+        Returns:
+            (next_inputs, sp_next) ready to submit to the decode stage.
+        """
+        next_inputs = [prompt] if not isinstance(prompt, list) else prompt
+
+        sp_next = sp_template.clone()
+        if sp_next.extra_args is None:
+            sp_next.extra_args = {}
+
+        if engine_outputs is not None:
+            prefill_kv = self._extract_kv_transfer_params(engine_outputs)
+        else:
+            prefill_kv = self._pop_pd_kv_params(req_id, prefill_kv_fallback)
+
+        decode_kv_params = self._build_decode_kv_params(req_id, sp_next, prefill_kv)
+        sp_next.extra_args["kv_transfer_params"] = decode_kv_params
+
+        logger.debug(
+            "[PD] routing req=%s, remote_request_id=%s",
+            req_id,
+            decode_kv_params.get("remote_request_id", "NOT SET"),
+        )
+        if "remote_request_id" not in decode_kv_params:
+            logger.warning(
+                "[PD] remote_request_id NOT SET for req %s. Apply mooncake_connector.py patch to fix.",
+                req_id,
+            )
+
+        return next_inputs, sp_next

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py

@@ -22,6 +22,8 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 
+from vllm_omni.platforms import current_omni_platform
+
 logger = init_logger(__name__)
 
 
@@ -343,6 +345,10 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
     def _ensure_model_fwd(self) -> None:
         if self._model_fwd is not None:
             return
+        if not current_omni_platform.supports_torch_inductor():
+            logger.warning_once("code_predictor: torch.compile disabled")
+            self._model_fwd = self.model.forward
+            return
         self._model_fwd = torch.compile(
             self.model.forward,
             mode="default",

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -21,6 +21,8 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 from vllm.model_executor.models.utils import is_pp_missing_parameter
 
+from vllm_omni.platforms import current_omni_platform
+
 from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
 
 logger = init_logger(__name__)
@@ -410,6 +412,10 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
         """
         if self._compiled_model_fwd is not None:
             return
+        if not current_omni_platform.supports_torch_inductor():
+            logger.warning_once("code_predictor: torch.compile disabled")
+            self._compiled_model_fwd = self.model.forward
+            return
         self._compiled_model_fwd = torch.compile(
             self.model.forward,
             mode="default",

vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe_async_chunk.yaml

@@ -0,0 +1,101 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
+# Stage 2: Code2Wav (16-layer RVQ codes → audio waveform)
+
+# The following config has been verified on 2x H100-80G GPUs.
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm  # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "0,1"
+      max_batch_size: 10
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: false
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm  # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "2"
+      max_batch_size: 10
+    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.6
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output codec codes for code2wav
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      distributed_executor_backend: "mp"
+      hf_config_name: talker_config
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
+    engine_input_source: [0]
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.0
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    stage_type: llm  # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "2"
+      max_batch_size: 10
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: audio  # Final output: audio waveform
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 51200 # [TODO] if max_num_batch_tokens < max_batch_size * 800, there will be precision problem.
+      hf_config_name: thinker_config
+    engine_input_source: [1]
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -14,6 +14,7 @@ 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.logger import logger
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import RoutedExpertsCapturer
 from vllm.sequence import IntermediateTensors
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import (
@@ -33,7 +34,7 @@ from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
 # yapf conflicts with isort for this block
 # yapf: disable
 from vllm_ascend.ops.rotary_embedding import update_cos_sin
-from vllm_ascend.utils import enable_sp
+from vllm_ascend.utils import enable_sp, global_stream
 
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -55,7 +56,7 @@ class ExecuteModelState(NamedTuple):
     positions: torch.Tensor
     ec_connector_output: ECConnectorOutput | None
     cudagraph_stats: CUDAGraphStat | None
-    multimodal_outputs: Any
+    multimodal_outputs: Any # Omni-Specific
 
 class NPUARModelRunner(OmniNPUModelRunner):
     """Autoregressive NPU model runner that returns hidden states per request."""
@@ -88,13 +89,36 @@ class NPUARModelRunner(OmniNPUModelRunner):
         scheduler_output: SchedulerOutput,
         intermediate_tensors: IntermediateTensors | None = None,
     ) -> OmniModelRunnerOutput | IntermediateTensors | None:
+        if self.vllm_config.model_config.enable_return_routed_experts:
+            capturer = RoutedExpertsCapturer.get_instance()
+            if capturer is not None:
+                capturer.clear_buffer()
+            else:
+                logger.warning("RoutedExpertsCapturer is not initialized.")
         if self.execute_model_state is not 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)
+        if not getattr(self, "_warmup_state_cleared", False):
+            self._warmup_state_cleared = True
+            if hasattr(self.model, "_clear_warmup_state"):
+                self.model._clear_warmup_state()
+
+        # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
+        finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
+        if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
+            for req_id, data in finished_reqs.items():
+                try:
+                    model_meta = self.model.get_kv_transfer_metadata(req_id)
+                    if model_meta:
+                        existing = data.get("custom_metadata") or {}
+                        existing.update(model_meta)
+                        data["custom_metadata"] = existing
+                except Exception as e:
+                    logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}")
         self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
-            finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
+            finished_reqs=finished_reqs,
             kv_caches=self.kv_caches,
             block_size=self.cache_config.block_size,
             cache_dtype=str(self.cache_config.cache_dtype),
@@ -115,10 +139,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
         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.
+                self._update_states(scheduler_output)
 
                 if has_ec_transfer() and get_ec_transfer().is_producer:
                     with self.maybe_get_ec_connector_output(
@@ -150,6 +172,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
                         "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]
@@ -159,6 +182,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 (
                     logits_indices,
                     spec_decode_metadata,
+                    total_num_scheduled_tokens,
                 ) = self._prepare_inputs(
                     scheduler_output,
                     num_scheduled_tokens_np,
@@ -224,10 +248,19 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     # 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)
+                    old_num_reqs_padded = num_reqs_padded
+                    num_reqs_padded = self._pad_query_start_loc_for_fia(
+                        num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs
+                    )
+                    if enable_sp() and num_tokens_padded == num_tokens_unpadded:
+                        if num_reqs_padded > old_num_reqs_padded:
+                            num_reqs_padded = old_num_reqs_padded
+                            self.query_start_loc.np[num_reqs_padded + 1] = 0
 
                 (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
-                    num_tokens=num_tokens_unpadded,
+                    num_tokens=num_tokens_unpadded
+                    if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn)
+                    else total_num_scheduled_tokens,
                     num_tokens_padded=num_tokens_padded,
                     num_reqs=num_reqs,
                     num_reqs_padded=num_reqs_padded,
@@ -247,9 +280,21 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 intermediate_tensors,
                 model_kwargs,
                 ec_connector_output,
-            ) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)
+            ) = self._preprocess(
+                scheduler_output,
+                num_tokens_padded
+                if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn)
+                else total_num_scheduled_tokens,
+                intermediate_tensors,
+            )
+
             # update global cos, sin
             update_cos_sin(positions)
+
+        if self.dynamic_eplb:
+            with record_function_or_nullcontext("EPLB weight D2D"):
+                self.eplb_updator.forward_before()
+
         # Set cudagraph mode to none if calc_kv_scales is true.
         # KV scales calculation involves dynamic operations that are incompatible
         # with CUDA graph capture.
@@ -278,6 +323,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
         has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
 
         # Run forward pass
+        clear_kv_metadata = self.speculative_config is None
         with (
             record_function_or_nullcontext("forward"),
             set_ascend_forward_context(
@@ -289,15 +335,22 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 batch_descriptor=batch_desc,
                 num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
                 model_instance=self.model,
+                max_tokens_across_pcp=0 if self.pcp_size == 1 else self.pcp_manager.max_num_tokens_across_pcp,
                 skip_compiled=has_encoder_input,
             ),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+            self.maybe_get_kv_connector_output(
+                scheduler_output, clear_metadata=clear_kv_metadata
+            ) 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 -------------------------------------------------
+            # [Omni] Map pending ropes metadata to req_ids.
+            if hasattr(self.model, "flush_pending_metadata"):
+                self.model.flush_pending_metadata(list(req_ids))
+
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
 
             if multimodal_outputs is not None:
@@ -310,13 +363,18 @@ class NPUARModelRunner(OmniNPUModelRunner):
             else:
                 logger.debug("[AR] execute_model: multimodal_outputs is None")
             #  -------------------------------------- Omni-new -------------------------------------------------
+            aux_hidden_states = None
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, aux_hidden_states = hidden_states
             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 aux_hidden_states is not None:
+                    aux_hidden_states = [
+                        self.pcp_manager.get_restore_hidden_states(aux_hidden_states_pcp)
+                        for aux_hidden_states_pcp in aux_hidden_states
+                    ]
 
             if not self.broadcast_pp_output:
                 # Common case.
@@ -392,7 +450,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 positions,
                 ec_connector_output,
                 cudagraph_stats,
-                multimodal_outputs,
+                multimodal_outputs, # Omni-specific
             )
             self.kv_connector_output = kv_connector_output
         return None
@@ -412,6 +470,11 @@ class NPUARModelRunner(OmniNPUModelRunner):
 
         if self.execute_model_state is None:
             # Nothing to do (PP non-final rank case), output isn't used.
+            # receive sampled token ids from the last PP rank when using
+            # async scheduling + pipeline parallelism so downstream code
+            # (e.g., PCP input preparation) can access them.
+            if self.use_async_scheduling and get_pp_group().world_size > 1:
+                self._pp_receive_prev_sampled_token_ids_to_input_batch()
             if not kv_connector_output:
                 return None  # noqa
             # In case of PP with kv transfer, we need to pass through the
@@ -436,7 +499,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
             positions,
             ec_connector_output,
             cudagraph_stats,
-            multimodal_outputs,
+            multimodal_outputs, # Omni-Specific
         ) = self.execute_model_state
         # Clear ephemeral state.
         self.execute_model_state = None
@@ -450,9 +513,27 @@ class NPUARModelRunner(OmniNPUModelRunner):
             apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits)
             logits = logits.to(self.device).to(logits_dtype)
 
+        #  -------------------------------------- Omni-new -------------------------------------------------
+        # Correct padding values of prompt_token_ids to match the logits vocabulary size.
+        if logits is not None and not self.input_batch.sampling_metadata.no_penalties:
+            smd = self.input_batch.sampling_metadata
+            if smd.prompt_token_ids is not None:
+                logits_vocab = logits.shape[-1]
+                if self.input_batch.vocab_size > logits_vocab:
+                    smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab)
+        #  -------------------------------------- Omni-new -------------------------------------------------
+
+
         with record_function_or_nullcontext("sample_token"):
             sampler_output = self._sample(logits, spec_decode_metadata)
 
+        if self.need_accepted_tokens:
+            if self.sampling_done_event is None:
+                self.sampling_done_event = torch.npu.Event()
+
+            assert self.sampling_done_event is not None
+            self.sampling_done_event.record()
+
         def propose_draft_token_ids(sampled_token_ids):
             assert spec_decode_common_attn_metadata is not None
             self._draft_token_ids = self.propose_draft_token_ids(
@@ -464,7 +545,6 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 positions,
                 scheduler_output.total_num_scheduled_tokens,
                 hidden_states,
-                attn_metadata,
                 aux_hidden_states,
                 sample_hidden_states,
             )
@@ -488,22 +568,30 @@ class NPUARModelRunner(OmniNPUModelRunner):
 
         with record_function_or_nullcontext("draft_token"):
             if self.speculative_config:
-                use_padded_batch_for_eagle = (
+                use_padded_batch = (
                     self.speculative_config
-                    and self.speculative_config.use_eagle()
+                    and (self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model())
                     and not self.speculative_config.disable_padded_drafter_batch
                 )
-                if use_padded_batch_for_eagle:
+                if use_padded_batch:
                     # EAGLE speculative decoding can use the GPU sampled tokens
                     # as inputs, and does not need to wait for bookkeeping to finish.
                     propose_draft_token_ids(sampler_output.sampled_token_ids)
-                if self.speculative_config and not use_padded_batch_for_eagle:
+                if self.speculative_config and not use_padded_batch:
                     # ngram and other speculative decoding methods use the sampled
                     # tokens on the CPU, so they are run after bookkeeping.
                     propose_draft_token_ids(valid_sampled_token_ids)
 
             if has_kv_transfer_group():
                 get_kv_transfer_group().clear_connector_metadata()
+
+        if self.model_config.enable_return_routed_experts:
+            capturer = RoutedExpertsCapturer.get_instance()
+            if capturer is not None:
+                capturer.save_captured_experts(indices=self.cpu_slot_mapping)
+            else:
+                logger.warning("RoutedExpertsCapturer is not initialized.")
+
         #  -------------------------------------- Omni-new -------------------------------------------------
         hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
         num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
@@ -562,13 +650,31 @@ class NPUARModelRunner(OmniNPUModelRunner):
         model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
         #  -------------------------------------- Omni-new -------------------------------------------------
 
-
         if self.dynamic_eplb:
-            self.eplb_updator.forward_end()
+            with record_function_or_nullcontext("EPLB update"):
+                self.eplb_updator.forward_end()
 
         if self.debugger is not None:
             self.debugger.stop()
             self.debugger.step()
+
+        if self.need_accepted_tokens:
+            assert self.sampling_done_event is not None
+            with (
+                record_function_or_nullcontext("async_state_update"),
+                torch.npu.stream(global_stream()),
+            ):
+                global_stream().wait_event(self.sampling_done_event)
+                self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output)
+
+        # In async scheduling + PP, broadcast sampled token ids from the
+        # last PP rank so other PP ranks can receive them without going
+        # through the scheduler/engine IPC path.
+        if self.use_async_scheduling:
+            pp = get_pp_group()
+            if pp.world_size > 1 and pp.is_last_rank:
+                self._pp_broadcast_prev_sampled_token_ids(sampler_output.sampled_token_ids)
+
         if not self.use_async_scheduling:
             return model_runner_output
         return AsyncGPUModelRunnerOutput(

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -14,6 +14,7 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
 from vllm.distributed.kv_transfer import has_kv_transfer_group
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.logger import logger
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import RoutedExpertsCapturer
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -69,7 +70,13 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         self,
         scheduler_output: SchedulerOutput,
         intermediate_tensors: IntermediateTensors | None = None,
-    ) -> OmniModelRunnerOutput | IntermediateTensors:
+    ) -> OmniModelRunnerOutput | IntermediateTensors | None:
+        if self.vllm_config.model_config.enable_return_routed_experts:
+            capturer = RoutedExpertsCapturer.get_instance()
+            if capturer is not None:
+                capturer.clear_buffer()
+            else:
+                logger.warning("RoutedExpertsCapturer is not initialized.")
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
         # self._draft_token_ids is None when `input_fits_in_drafter=False`
@@ -123,6 +130,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                         "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]
@@ -132,6 +140,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 (
                     logits_indices,
                     spec_decode_metadata,
+                    total_num_scheduled_tokens,
                 ) = self._prepare_inputs(
                     scheduler_output,
                     num_scheduled_tokens_np,
@@ -197,10 +206,19 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                     # 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)
+                    old_num_reqs_padded = num_reqs_padded
+                    num_reqs_padded = self._pad_query_start_loc_for_fia(
+                        num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs
+                    )
+                    if enable_sp() and num_tokens_padded == num_tokens_unpadded:
+                        if num_reqs_padded > old_num_reqs_padded:
+                            num_reqs_padded = old_num_reqs_padded
+                            self.query_start_loc.np[num_reqs_padded + 1] = 0
 
                 (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
-                    num_tokens=num_tokens_unpadded,
+                    num_tokens=num_tokens_unpadded
+                    if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn)
+                    else total_num_scheduled_tokens,
                     num_tokens_padded=num_tokens_padded,
                     num_reqs=num_reqs,
                     num_reqs_padded=num_reqs_padded,
@@ -220,7 +238,13 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 intermediate_tensors,
                 model_kwargs,
                 ec_connector_output,
-            ) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)
+            ) = self._preprocess(
+                scheduler_output,
+                num_tokens_padded
+                if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn)
+                else total_num_scheduled_tokens,
+                intermediate_tensors,
+            )
 
             # [Omni] Pass token counts per request for code2wav output slicing
             model_kwargs["seq_token_counts"] = tokens
@@ -228,6 +252,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             # update global cos, sin
             update_cos_sin(positions)
 
+        if self.dynamic_eplb:
+            with record_function_or_nullcontext("EPLB weight D2D"):
+                self.eplb_updator.forward_before()
+
         # Set cudagraph mode to none if calc_kv_scales is true.
         # KV scales calculation involves dynamic operations that are incompatible
         # with CUDA graph capture.
@@ -256,6 +284,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         has_encoder_input = self.model_config.is_encoder_decoder and num_encoder_reqs > 0
 
         # Run forward pass
+        clear_kv_metadata = self.speculative_config is None
         with (
             record_function_or_nullcontext("forward"),
             set_ascend_forward_context(
@@ -267,9 +296,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 batch_descriptor=batch_desc,
                 num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
                 model_instance=self.model,
+                max_tokens_across_pcp=0 if self.pcp_size == 1 else self.pcp_manager.max_num_tokens_across_pcp,
                 skip_compiled=has_encoder_input,
             ),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+            self.maybe_get_kv_connector_output(
+                scheduler_output, clear_metadata=clear_kv_metadata
+            ) as kv_connector_output,
         ):
             #  -------------------------------------- Omni-new -------------------------------------------------
             outputs = self._run_generation_model(
@@ -296,7 +328,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 positions,
                 ec_connector_output,
                 cudagraph_stats,
-                multimodal_outputs,  # Omni-new: pass multimodal_outputs to ExecuteModelState
+                multimodal_outputs,  # Omni-specific: pass multimodal_outputs to ExecuteModelState
             )
             self.kv_connector_output = kv_connector_output
         return None
@@ -310,6 +342,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
 
         if self.execute_model_state is None:
             # Nothing to do (PP non-final rank case), output isn't used.
+            # receive sampled token ids from the last PP rank when using
+            # async scheduling + pipeline parallelism so downstream code
+            # (e.g., PCP input preparation) can access them.
+            if self.use_async_scheduling and get_pp_group().world_size > 1:
+                self._pp_receive_prev_sampled_token_ids_to_input_batch()
             if not kv_connector_output:
                 return None  # noqa
             # In case of PP with kv transfer, we need to pass through the
@@ -334,7 +371,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             positions,
             ec_connector_output,
             cudagraph_stats,
-            multimodal_outputs,
+            multimodal_outputs,  # Omni-Specific
         ) = self.execute_model_state
         # Clear ephemeral state.
         self.execute_model_state = None
@@ -464,7 +501,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         allow_microbatching: bool = True,
         skip_eplb: bool = False,
         remove_lora: bool = True,
-        activate_lora: bool = False,
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
     ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -505,6 +541,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
 
+        if not is_profile and self.dynamic_eplb:
+            self.eplb_updator.forward_before()
+
         num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
         self.query_lens = torch.from_numpy(num_scheduled_tokens)
         num_tokens_unpadded = int(num_scheduled_tokens.sum())
@@ -525,7 +564,8 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             # `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,
+            force_has_lora=num_active_loras > 0,
+            force_num_active_loras=num_active_loras,
         )
         if self.use_cp:
             self.pcp_manager.init_batch_info(
@@ -551,9 +591,8 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         # 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:
+        # Build attention metadata for dummy_run
+        if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode):
             if create_mixed_batch:
                 raise NotImplementedError(
                     "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it"
@@ -582,8 +621,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             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()
-
-            num_reqs_padded = self._pad_query_start_loc_for_fia(num_tokens_padded, num_reqs_padded, num_reqs)
+            num_reqs_padded = self._pad_query_start_loc_for_fia(
+                num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs
+            )
 
             pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
             attn_metadata, _ = self._build_attention_metadata(
@@ -600,6 +640,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             self.lora_config,
             num_scheduled_tokens,
             num_sampled_tokens,
+            remove_lora,
+            # TODO: The next line is a temporary workaround
+            # to fix the accuracy issue of test_llama32_lora.py,
+            # which is introduced by vllm-project/vllm#32005
+            num_active_loras=(self.lora_config.max_loras if self.lora_config is not None else num_active_loras),
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
@@ -610,6 +655,19 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
 
+            # -------------------------------------- Omni-new -------------------------------------------------
+            model_kwargs = self._init_model_kwargs()
+            # Some generation-stage models (e.g. MammothModa2DiTPipeline) require
+            # model-specific runtime information (such as image size and conditioning
+            # embeddings) even during the dummy profiling run that vLLM uses to
+            # estimate KV-cache capacity.  get_dummy_runtime_additional_information
+            # provides placeholder values of the correct shape so that the profiling
+            # run does not raise an error due to missing inputs.
+            if hasattr(self.model, "get_dummy_runtime_additional_information"):
+                runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs)
+                model_kwargs["runtime_additional_information"] = runtime_addi
+            # -------------------------------------- Omni-new -------------------------------------------------
+
             if self.uses_mrope:
                 positions = self.mrope_positions.gpu[:, :num_tokens_padded]
             elif self.uses_xdrope_dim > 0:
@@ -656,8 +714,6 @@ 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,
@@ -674,9 +730,8 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                     positions=positions,
                     intermediate_tensors=intermediate_tensors,
                     inputs_embeds=inputs_embeds,
-                    **model_kwargs,
+                    **model_kwargs,  # Omni-specific
                 )
-
             if self.use_aux_hidden_state_outputs:
                 hidden_states, _ = outputs
             else:
@@ -698,7 +753,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             if is_profile and self.dynamic_eplb:
                 self.model.clear_all_moe_loads()
             if self.dynamic_eplb:
-                self.eplb_updator.take_update_info_from_eplb_process()
                 self.eplb_updator.forward_end()
 
             # -------------------------------------- Omni-new -------------------------------------------------

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -71,7 +71,6 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         allow_microbatching: bool = True,
         skip_eplb: bool = False,
         remove_lora: bool = True,
-        activate_lora: bool = False,
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
     ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -112,6 +111,9 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
 
+        if not is_profile and self.dynamic_eplb:
+            self.eplb_updator.forward_before()
+
         num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
         self.query_lens = torch.from_numpy(num_scheduled_tokens)
         num_tokens_unpadded = int(num_scheduled_tokens.sum())
@@ -132,7 +134,8 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             # `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,
+            force_has_lora=num_active_loras > 0,
+            force_num_active_loras=num_active_loras,
         )
         if self.use_cp:
             self.pcp_manager.init_batch_info(
@@ -158,9 +161,8 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
         # 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:
+        # Build attention metadata for dummy_run
+        if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode):
             if create_mixed_batch:
                 raise NotImplementedError(
                     "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it"
@@ -189,8 +191,9 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             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()
-
-            num_reqs_padded = self._pad_query_start_loc_for_fia(num_tokens_padded, num_reqs_padded, num_reqs)
+            num_reqs_padded = self._pad_query_start_loc_for_fia(
+                num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs
+            )
 
             pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
             attn_metadata, _ = self._build_attention_metadata(
@@ -207,6 +210,11 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             self.lora_config,
             num_scheduled_tokens,
             num_sampled_tokens,
+            remove_lora,
+            # TODO: The next line is a temporary workaround
+            # to fix the accuracy issue of test_llama32_lora.py,
+            # which is introduced by vllm-project/vllm#32005
+            num_active_loras=(self.lora_config.max_loras if self.lora_config is not None else num_active_loras),
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
@@ -287,9 +295,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
                     )
                     self.compilation_config.cache_dir = None
                 # ---------------------------------------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.
+                # Call self.model() directly (like GPU) to avoid make_omni_output during dummy_run
                 outputs = self.model(
                     input_ids=input_ids,
                     positions=positions,
@@ -320,7 +326,6 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             if is_profile and self.dynamic_eplb:
                 self.model.clear_all_moe_loads()
             if self.dynamic_eplb:
-                self.eplb_updator.take_update_info_from_eplb_process()
                 self.eplb_updator.forward_end()
             return hidden_states, hidden_states
 
@@ -382,7 +387,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
             )
 
         # NPU-specific: all-gather for sequence parallelism
-        if get_forward_context().sp_enabled and not isinstance(model_output, IntermediateTensors):
+        if get_forward_context().flash_comm_v1_enabled and not isinstance(model_output, IntermediateTensors):
             model_output = self._all_gather_hidden_states_and_aux(model_output)
 
         return model_output