[NPU][Refactor] Rename Diffusion* to Generation* (#211)

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

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

vllm_omni/worker/npu/npu_diffusion_worker.py

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

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."
         )
 

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)

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.