[Bugfix][NPU] Add _model_forward for ModelRunner (#505)

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

vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml

@@ -79,7 +79,7 @@ stage_args:
       engine_output_type: audio  # Final output: audio waveform
       gpu_memory_utilization: 0.1
       distributed_executor_backend: "mp"
-      max_num_batched_tokens: 4096
+      max_num_batched_tokens: 1000000
       hf_config_name: thinker_config
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav

vllm_omni/worker/npu/npu_ar_model_runner.py

@@ -1199,21 +1199,11 @@ class NPUARModelRunner(OmniNPUModelRunner):
                                              input_ids, positions,
                                              intermediate_tensors,
                                              inputs_embeds):
-        model_kwargs_extra = self._build_model_kwargs_extra()
-
-        runtime_info = model_kwargs_extra.get("runtime_additional_information", [])
-        if runtime_info:
-            for i, info in enumerate(runtime_info):
-                if info:
-                    logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}")
-
-        assert self.model is not None
-        hidden_states = self.model(input_ids=input_ids,
+        hidden_states = self._model_forward(input_ids=input_ids,
                                    positions=positions,
                                    intermediate_tensors=intermediate_tensors,
                                    inputs_embeds=inputs_embeds,
-                                   **self._init_model_kwargs(),
-                                   **model_kwargs_extra)
+                                   **self._init_model_kwargs())
 
         forward_context = get_forward_context()
         if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL \

vllm_omni/worker/npu/npu_model_runner.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import math
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import torch
@@ -636,3 +636,32 @@ class OmniNPUModelRunner(NPUModelRunner):
                 )
                 start_offset = int(self.query_start_loc.cpu[req_index])
                 self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
+
+    def _model_forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **model_kwargs: dict[str, Any],
+    ):
+        """Inject omni-specific kwargs into forward and cache model output"""
+        model_kwargs_extra = self._build_model_kwargs_extra()
+
+        runtime_info = model_kwargs_extra.get("runtime_additional_information", [])
+        if runtime_info:
+            for i, info in enumerate(runtime_info):
+                if info:
+                    logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}")
+
+        model_output = super()._model_forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **model_kwargs,
+            **model_kwargs_extra,
+        )
+        # Cache model output so later sample_tokens can consume multimodal results.
+        self._omni_last_model_output = model_output
+        return model_output