diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml index f2205b8be..2812e99a0 100644 --- a/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml @@ -9,7 +9,7 @@ stage_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker - scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true # Now we only support eager mode trust_remote_code: true @@ -35,7 +35,7 @@ stage_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker - scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true trust_remote_code: true @@ -60,8 +60,8 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_diffusion_worker.NPUDiffusionWorker - scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler + worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/worker/npu/npu_diffusion_worker.py b/vllm_omni/worker/npu/npu_diffusion_worker.py deleted file mode 100644 index 9974354ad..000000000 --- a/vllm_omni/worker/npu/npu_diffusion_worker.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm_ascend.worker.worker_v1 import NPUWorker - -from vllm_omni.worker.npu.npu_diffusion_model_runner import NPUDiffusionModelRunner - - -class NPUDiffusionWorker(NPUWorker): - """NPU diffusion worker for code2wav stage in Omni model.""" - - def init_device(self): - device = self._init_device() - - self.model_runner: NPUDiffusionModelRunner = NPUDiffusionModelRunner(self.vllm_config, device) diff --git a/vllm_omni/worker/npu/npu_diffusion_model_runner.py b/vllm_omni/worker/npu/npu_generation_model_runner.py similarity index 98% rename from vllm_omni/worker/npu/npu_diffusion_model_runner.py rename to vllm_omni/worker/npu/npu_generation_model_runner.py index ccda654e2..33fc8e63a 100644 --- a/vllm_omni/worker/npu/npu_diffusion_model_runner.py +++ b/vllm_omni/worker/npu/npu_generation_model_runner.py @@ -53,8 +53,8 @@ if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput -class NPUDiffusionModelRunner(OmniNPUModelRunner): - """Diffusion model runner for vLLM-omni on NPU (non-autoregressive).""" +class NPUGenerationModelRunner(OmniNPUModelRunner): + """Generation model runner for vLLM-omni on NPU (non-autoregressive).""" def _prepare_inputs( self, @@ -455,7 +455,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner): ): self.maybe_setup_kv_connector(scheduler_output) - outputs = self._run_diffusion( + outputs = self._run_generation( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, @@ -525,7 +525,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner): async_output_copy_stream=self.async_output_copy_stream, ) - def _run_diffusion( + def _run_generation( self, *, input_ids: torch.Tensor, @@ -535,7 +535,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner): multimodal_kwargs: dict, logits_indices: torch.Tensor, ) -> torch.Tensor | list[torch.Tensor]: - """Runs the diffusion process and returns per-request tensors. + """Runs the generation process and returns per-request tensors. Tries model interfaces in the following order for maximal compatibility: 1) model.sample(condition=..., **kwargs) @@ -558,7 +558,7 @@ class NPUDiffusionModelRunner(OmniNPUModelRunner): # TODO: add the diffuse method for other models raise RuntimeError( - "The loaded model does not expose diffusion interfaces 'sample', " + "The loaded model does not expose generation interfaces 'sample', " "'forward', or 'diffuse'. Please implement one of them or adapt the runner." ) diff --git a/vllm_omni/worker/npu/npu_generation_worker.py b/vllm_omni/worker/npu/npu_generation_worker.py new file mode 100644 index 000000000..23cdf41c7 --- /dev/null +++ b/vllm_omni/worker/npu/npu_generation_worker.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_ascend.worker.worker_v1 import NPUWorker + +from vllm_omni.worker.npu.npu_generation_model_runner import NPUGenerationModelRunner + + +class NPUGenerationWorker(NPUWorker): + """NPU generation worker for code2wav stage in Omni model.""" + + def init_device(self): + device = self._init_device() + + self.model_runner: NPUGenerationModelRunner = NPUGenerationModelRunner(self.vllm_config, device) diff --git a/vllm_omni/worker/npu/npu_model_runner.py b/vllm_omni/worker/npu/npu_model_runner.py index 3b1fdb486..ee6ef83f4 100644 --- a/vllm_omni/worker/npu/npu_model_runner.py +++ b/vllm_omni/worker/npu/npu_model_runner.py @@ -66,17 +66,26 @@ class OmniNPUModelRunner(NPUModelRunner): if use_audio_in_video_value is not None: use_audio_in_video = bool(use_audio_in_video_value.item()) - assert supports_mrope(self.get_model()), "M-RoPE support is not implemented." - - req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor( - req_state.prompt_token_ids, - hf_config=self.model_config.hf_config, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - audio_feature_lengths=audio_feature_lengths, - use_audio_in_video=use_audio_in_video, - ) + if supports_mrope(self.model): + req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions( + req_state.prompt_token_ids, + hf_config=self.model_config.hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + else: + req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor( + req_state.prompt_token_ids, + hf_config=self.model_config.hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) def _update_states(self, scheduler_output: "SchedulerOutput") -> None: # Remove finished requests from the cached states.