[BugFix] Restore pre-#9572 NPU graph behavior (cap cudagraph_mode to PIECEWISE) (#4674)

Signed-off-by: Weiming Liao <liaowm5@gmail.com>

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -77,13 +77,19 @@ platforms:
         gpu_memory_utilization: 0.6
         tensor_parallel_size: 2
         max_num_batched_tokens: 8192
+        compilation_config:
+          cudagraph_mode: PIECEWISE   # Qwen3-Omni not FULL-cudagraph-safe on NPU
         devices: "0,1"
       - stage_id: 1
         gpu_memory_utilization: 0.6
         max_num_batched_tokens: 8192
+        compilation_config:
+          cudagraph_mode: PIECEWISE   # code_predictor MaskedSelect / inner NPU graph replay
         devices: "2"
       - stage_id: 2
         gpu_memory_utilization: 0.3
+        compilation_config:
+          cudagraph_mode: PIECEWISE
         devices: "2"
 
   # Code2Wav decode uses UpSample1d, which calls F.conv_transpose1d.

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -742,7 +742,15 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
-            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds:
+            model_kwargs = self._init_model_kwargs()
+            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
+                input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)
+
+                model_kwargs = {
+                    **model_kwargs,
+                    **self._dummy_mm_kwargs(num_reqs),
+                }
+            elif self.enable_prompt_embeds:
                 input_ids = None
                 inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
             else:
@@ -750,7 +758,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 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