Checkpoint

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>

Author: divyanshsinghvi

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -23,8 +23,6 @@ stage_args:
       hf_config_name: thinker_config
       tensor_parallel_size: 2
       load_format: dummy
-    final_output: true
-    final_output_type: text
     is_comprehension: true
     default_sampling_params:
       temperature: 0.4
@@ -55,8 +53,8 @@ stage_args:
        load_format: dummy
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-    # final_output: true
-    # final_output_type: text
+    final_output: true
+    final_output_type: text
     default_sampling_params:
       temperature: 0.9
       top_k: 50

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -1,6 +1,7 @@
 import time
 from collections import defaultdict
 
+from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import (
@@ -16,6 +17,8 @@
 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.outputs import OmniModelRunnerOutput
 
+logger = init_logger(__name__)
+
 
 class OmniGenerationScheduler(VLLMScheduler):
     def schedule(self) -> SchedulerOutput:
@@ -185,7 +188,7 @@ def update_from_output(
                 # request is aborted while the model is executing it (e.g.,
                 # in pipeline parallelism).
                 continue
-
+            logger.info(f"Diffusion request completed: {req_id} {model_runner_output.req_id_to_index}")
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
 
@@ -248,9 +251,14 @@ def update_from_output(
             # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
             pooling_output_dict = None
             if pooler_output is not None:
-                # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
-                # which expects Optional[dict[str, torch.Tensor]]
-                pooling_output_dict = {"model_outputs": pooler_output}
+                # If pooler_output is already a dict (from stages that output multiple tensors),
+                # preserve it directly. Otherwise, wrap tensor in dict.
+                if isinstance(pooler_output, dict):
+                    logger.info(f"[DEBUG scheduler] pooler_output is dict with keys: {list(pooler_output.keys())}")
+                    pooling_output_dict = pooler_output
+                else:
+                    logger.info("[DEBUG scheduler] pooler_output is tensor, wrapping as model_outputs")
+                    pooling_output_dict = {"model_outputs": pooler_output}
             if new_token_ids or pooler_output is not None or kv_transfer_params:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(

vllm_omni/engine/output_processor.py

@@ -111,6 +111,10 @@ def _to_cpu(x):
                 return x
 
             if isinstance(payload, dict):
+                logger.info(
+                    f"[DEBUG add_multimodal_tensor] payload is dict with keys: {list(payload.keys())}, "
+                    f"mm_type={self.mm_type}"
+                )
                 incoming: Dict[str, Any] = {}
                 # Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type
                 # to keep a consistent key namespace per engine_core_output_type.
@@ -125,8 +129,10 @@ def _to_cpu(x):
                         incoming[k] = {str(sk): _to_cpu(sv) for sk, sv in v.items()}
                     else:
                         incoming[k] = _to_cpu(v)
+                logger.info(f"[DEBUG add_multimodal_tensor] incoming dict has keys: {list(incoming.keys())}")
             else:
                 key = self.mm_type or "hidden"
+                logger.info(f"[DEBUG add_multimodal_tensor] payload is NOT dict, wrapping as: {{'{key}': payload}}")
                 incoming = {key: _to_cpu(payload)}
 
             if self.mm_accumulated is None:
@@ -380,6 +386,14 @@ def process_outputs(
             # 2.5) Accumulate multimodal tensors in RequestState
             try:
                 mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
+                logger.info(
+                    f"[DEBUG process_outputs] req_id={req_id}, mm_type={mm_type},"
+                    f"pooling_output type: {type(pooling_output)}"
+                )
+                if isinstance(pooling_output, dict):
+                    logger.info(
+                        f"[DEBUG process_outputs] pooling_output is dictwith keys: {list(pooling_output.keys())}"
+                    )
                 if pooling_output is not None and isinstance(req_state, OmniRequestState):
                     req_state.add_multimodal_tensor(pooling_output, mm_type)
             except Exception:
@@ -497,9 +511,27 @@ def _process_text_image_output(self, eco: EngineCoreOutput) -> None:
     def _process_latents_output(self, eco: EngineCoreOutput) -> None:
         """Ensure latent tensors are surfaced via pooling_output."""
         if eco.pooling_output is None:
+            # DEBUG: Log what we're processing
+            mm = getattr(eco, "multimodal_outputs", None)
+            logger.info(f"[DEBUG _process_latents_output] multimodal_outputs type: {type(mm)}")
+            if isinstance(mm, dict):
+                logger.info(f"[DEBUG _process_latents_output] multimodal_outputs keys: {list(mm.keys())}")
+
             tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior"))
+            logger.info(
+                f"[DEBUG _process_latents_output] extracted tensor type: {type(tensor)}, "
+                f"is dict: {isinstance(tensor, dict)}"
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
+                logger.info("[DEBUG _process_latents_output] set eco.pooling_output to extracted tensor")
+        else:
+            # pooling_output already set (likely from scheduler with full dict)
+            logger.info(f"[DEBUG _process_latents_output] pooling_output already set, type: {type(eco.pooling_output)}")
+            if isinstance(eco.pooling_output, dict):
+                logger.info(
+                    f"[DEBUG _process_latents_output] pooling_output dict keys: {list(eco.pooling_output.keys())}"
+                )
 
     def _process_audio_output(self, eco: EngineCoreOutput) -> None:
         """Ensure audio tensors are surfaced via pooling_output."""
@@ -532,9 +564,15 @@ def _extract_from_multimodal_outputs(self, eco: EngineCoreOutput, keys: tuple[st
         for k in keys:
             v = mm.get(k)
             if isinstance(v, torch.Tensor):
+                logger.info(f"[DEBUG _extract_from_multimodal_outputs] Found key '{k}' in multimodal_outputs")
                 return v
         # Try the first tensor in the dict as a fallback
-        for v in mm.values():
+        logger.info(f"[DEBUG _extract_from_multimodal_outputs] No matching keys {keys}, using fallback (first tensor)")
+        for k, v in mm.items():
             if isinstance(v, torch.Tensor):
+                logger.info(
+                    f"[DEBUG _extract_from_multimodal_outputs] Fallback: extracted first"
+                    f" tensor with key '{k}', shape: {v.shape}"
+                )
                 return v
         return None

vllm_omni/entrypoints/omni_llm.py

@@ -351,10 +351,7 @@ def _run_generation(
                 req_id = result.get("request_id")
                 if "error" in result:
                     logger.error(
-                        "Stage %s error on request %s: %s",
-                        stage_id,
-                        req_id,
-                        result["error"],
+                        "Stage %s error on request %s: %s %s", stage_id, req_id, result["error"], result["error_tb"]
                     )
                     continue
 

vllm_omni/entrypoints/omni_stage.py

@@ -16,6 +16,7 @@
 import multiprocessing as mp
 import os
 import sys
+import traceback as _traceback
 from typing import Any
 
 from vllm.inputs import TextPrompt
@@ -720,12 +721,14 @@ def filter(self, record: _logging.LogRecord) -> bool:
                 )
         except Exception as e:
             _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
+            _tb = _traceback.format_exc()
             for rid in batch_request_ids:
                 out_q.put(
                     {
                         "request_id": rid,
                         "stage_id": stage_id,
                         "error": str(e),
+                        "error_tb": _tb,
                     }
                 )