diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml index 5ebdddd18..db51e6d7d 100644 --- a/vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml +++ b/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 diff --git a/vllm_omni/worker/npu/npu_ar_model_runner.py b/vllm_omni/worker/npu/npu_ar_model_runner.py index ac677e3cd..fcaf36795 100644 --- a/vllm_omni/worker/npu/npu_ar_model_runner.py +++ b/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 \ diff --git a/vllm_omni/worker/npu/npu_model_runner.py b/vllm_omni/worker/npu/npu_model_runner.py index b6f12d84f..788458a1e 100644 --- a/vllm_omni/worker/npu/npu_model_runner.py +++ b/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