diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 38808648e..b0c120ca7 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -48,6 +48,39 @@ from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element +def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]: + """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``. + + Non-tensor scalars (int, float, bool) are wrapped with ``torch.tensor()``. + Values that cannot be safely converted are dropped with a warning. + This enforces the tensor-only invariant required by the + ``OmniEngineCoreOutput.multimodal_output`` wire field and msgspec + serialization. Mirrors ``gpu_ar_model_runner._ensure_tensor_values``. + """ + result: dict[str, torch.Tensor] = {} + for key, val in payload.items(): + if isinstance(val, torch.Tensor): + result[key] = val + elif isinstance(val, (int, float, bool)): + result[key] = torch.tensor(val) + elif isinstance(val, (list, tuple)): + try: + result[key] = torch.tensor(val) + except (ValueError, TypeError, RuntimeError): + logger.warning( + "Dropping non-tensorizable multimodal output key '%s' (type=%s) from wire payload.", + key, + type(val).__name__, + ) + else: + logger.warning( + "Dropping non-tensor multimodal output key '%s' (type=%s) from wire payload.", + key, + type(val).__name__, + ) + return result + + class ExecuteModelState(NamedTuple): """Ephemeral cached state transferred between execute_model() and sample_tokens(), after execute_model() returns None.""" @@ -216,6 +249,18 @@ class NPUARModelRunner(OmniNPUModelRunner): return combined_hidden_states[rid] return hidden_states_cpu[start:end] + + def _build_multimodal_outputs( + self, + per_req_payloads: list[dict[str, object]] | None, + ) -> list[dict[str, torch.Tensor]] | None: + if self.vllm_config.model_config.engine_output_type == "text": + return None + if per_req_payloads is None: + return None + return [_ensure_tensor_values(payload) if payload else {} for payload in per_req_payloads] + + def _request_final_stage_id(self, req_id: str) -> int | None: info = self.model_intermediate_buffer.get(req_id) if not isinstance(info, dict): @@ -1087,13 +1132,15 @@ class NPUARModelRunner(OmniNPUModelRunner): payload.update(mm_payload) pooler_output.append(flatten_payload(payload)) + multimodal_outputs = self._build_multimodal_outputs(pooler_output) model_runner_output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, sampled_token_ids=valid_sampled_token_ids, logprobs=logprobs_lists, prompt_logprobs_dict=prompt_logprobs_dict, - pooler_output=(pooler_output if engine_output_type != "text" and needs_pooler_payload else None), + pooler_output=None, + multimodal_outputs=multimodal_outputs, kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, cudagraph_stats=cudagraph_stats, diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 341b79d36..d2eeeb578 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -33,7 +33,7 @@ from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState +from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner @@ -427,33 +427,35 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): ec_connector_output, cudagraph_stats, _batch_desc, - multimodal_outputs, # Omni-Specific + multimodal_outputs_raw, # Omni-Specific ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None # -------------------------------------- Omni-new ------------------------------------------------- - pooler_output: list[object] = [] - if isinstance(multimodal_outputs, torch.Tensor): - assert multimodal_outputs.shape[0] == 1, ( + # Build per-request multimodal_outputs list (dedicated channel). + # pooler_output is no longer used for multimodal data. + multimodal_outputs: list[dict[str, object]] = [] + if isinstance(multimodal_outputs_raw, torch.Tensor): + assert multimodal_outputs_raw.shape[0] == 1, ( "model should return a single tensor, to return multiple tensors, use a dict" ) - assert multimodal_outputs.shape[0] == self.input_batch.num_reqs + assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs for i in range(self.input_batch.num_reqs): - pooler_output.append({"model_outputs": multimodal_outputs[i].detach().to("cpu").contiguous()}) - elif isinstance(multimodal_outputs, list): - assert len(multimodal_outputs) == 1, ( + multimodal_outputs.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) + elif isinstance(multimodal_outputs_raw, list): + assert len(multimodal_outputs_raw) == 1, ( "model should return a single list, to return multiple lists, use a dict" ) - for out in multimodal_outputs: - pooler_output.append( + for out in multimodal_outputs_raw: + multimodal_outputs.append( {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} ) - elif isinstance(multimodal_outputs, Mapping): + elif isinstance(multimodal_outputs_raw, Mapping): num_reqs = self.input_batch.num_reqs for i in range(num_reqs): mm_payload = {} - for key, out in multimodal_outputs.items(): + for key, out in multimodal_outputs_raw.items(): if isinstance(out, list): if len(out) != num_reqs: raise ValueError( @@ -465,7 +467,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): mm_payload[key] = out.detach().to("cpu").contiguous() else: logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") - pooler_output.append(mm_payload) + multimodal_outputs.append(_ensure_tensor_values(mm_payload)) else: raise RuntimeError("Unsupported diffusion output type") # [Omni] Copy req_id mappings to avoid async scheduling mutation. @@ -482,10 +484,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): sampled_token_ids=[], logprobs=None, prompt_logprobs_dict={}, - pooler_output=pooler_output, + pooler_output=None, + multimodal_outputs=multimodal_outputs, kv_connector_output=kv_connector_output, - ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, + num_nans_in_logits={}, cudagraph_stats=cudagraph_stats, + ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) output.routed_experts = routed_experts_lists # -------------------------------------- Omni-new -------------------------------------------------