[Core] Allow Broadcasting Scalars

tests/utils/test_mm_outputs.py

@@ -1,8 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 import torch
 
 from vllm_omni.utils import mm_outputs as mm_outputs_mod
-from vllm_omni.utils.mm_outputs import build_mm_cpu
+from vllm_omni.utils.mm_outputs import build_mm_cpu, split_mapping_outputs, to_payload_element
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -145,3 +148,36 @@ def test_build_mm_cpu_preserves_nested_and_non_tensor_payload_semantics():
     assert torch.equal(result["grad_key"], grad_tensor.detach())
 
     assert result["scalar"] == 42
+
+
+@pytest.mark.parametrize("sr", [24000, torch.tensor(24000)])
+def test_split_mapping_outputs_broadcasts_invariants_and_indexes_lists(sr):
+    waveforms = [torch.zeros(4), torch.ones(4)]
+
+    payloads = split_mapping_outputs(
+        {"model_outputs": waveforms, "sr": sr},
+        num_reqs=2,
+        copy_tensor=torch.Tensor.clone,
+    )
+
+    # Both the int and the int wrapped as a 0d tensor should be canonicalized to a wrapped tensor
+    for payload, waveform in zip(payloads, waveforms, strict=True):
+        model_outputs = payload["model_outputs"]
+        assert isinstance(model_outputs, torch.Tensor) and torch.equal(model_outputs, waveform)
+        assert torch.equal(torch.as_tensor(payload["sr"]), torch.tensor(24000))
+
+
+def test_split_mapping_outputs_rejects_misaligned_lists():
+    with pytest.raises(ValueError, match="one entry per request"):
+        split_mapping_outputs(
+            {"model_outputs": [torch.zeros(4)]},
+            num_reqs=2,
+            copy_tensor=torch.Tensor.clone,
+        )
+
+
+def test_to_payload_element_shares_0d_tensor_with_every_request():
+    sr = torch.tensor(24000)
+    for idx in range(2):
+        payload = to_payload_element(sr, idx=idx, start=idx, end=idx + 1, seq_len=2, scheduled_seq_len=2)
+        assert torch.equal(payload, sr)

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -38,7 +38,7 @@
 from vllm_omni.outputs import OmniModelRunnerOutput
 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
-from vllm_omni.utils.mm_outputs import partition_payload_list
+from vllm_omni.utils.mm_outputs import partition_payload_list, split_mapping_outputs
 from vllm_omni.worker.mixins import maybe_unpad_input_ids
 from vllm_omni.worker.omni_connector_model_runner_mixin import (
     OmniConnectorModelRunnerMixin,
@@ -591,22 +591,14 @@ def sample_tokens(
                 for out in multimodal_outputs_raw:
                     per_req_payloads.append({"model_outputs": to_host.copy(out) if out is not None else None})
             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_raw.items():
-                        if isinstance(out, list):
-                            if len(out) != num_reqs:
-                                raise ValueError(
-                                    f"Multimodal output list for key '{key}' has length {len(out)} "
-                                    f"but expected {num_reqs} (one entry per request)."
-                                )
-                            mm_payload[key] = to_host.copy(out[i])
-                        elif isinstance(out, torch.Tensor):
-                            mm_payload[key] = to_host.copy(out)
-                        else:
-                            logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
-                    per_req_payloads.append(_ensure_tensor_values(mm_payload))
+                per_req_payloads = [
+                    _ensure_tensor_values(payload)
+                    for payload in split_mapping_outputs(
+                        multimodal_outputs_raw,
+                        self.input_batch.num_reqs,
+                        lambda tensor: tensor.detach().to("cpu").contiguous(),
+                    )
+                ]
             else:
                 raise RuntimeError("Unsupported diffusion output type")
         if not use_async_output:

vllm_omni/utils/mm_outputs.py

@@ -1,8 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Utilities for handling multimodal outputs / building multimodal output
 payloads, most of which are shared by the prefix cache / no prefix cache path.
 """
 
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping, Sequence
 
 import torch
 from vllm.logger import init_logger
@@ -66,7 +69,7 @@ def partition_flat_payload(
 
 
 def partition_payload_list(
-    payloads: list[dict[str, object]],
+    payloads: Sequence[Mapping[str, object]],
 ) -> tuple[list[dict[str, object] | None] | None, list[dict[str, object] | None] | None]:
     inter_stage_list: list[dict[str, object] | None] = []
     client_mm_list: list[dict[str, object] | None] = []
@@ -180,6 +183,38 @@ def _to_cpu(value):
     return value
 
 
+def split_mapping_outputs(
+    outputs: Mapping[str, object],
+    num_reqs: int,
+    copy_tensor: Callable[[torch.Tensor], torch.Tensor],
+) -> list[dict[str, object]]:
+    """Split a generation stage's batched output mapping into one payload per request."""
+    batch_payloads: list[dict[str, object]] = [{} for _ in range(num_reqs)]
+    for key, out in outputs.items():
+        if isinstance(out, list):
+            if len(out) != num_reqs:
+                raise ValueError(
+                    f"Multimodal output list for key '{key}' has length {len(out)} "
+                    f"but expected {num_reqs} (one entry per request)."
+                )
+            for req_payload, req_item in zip(batch_payloads, out, strict=True):
+                req_payload[key] = copy_tensor(req_item)
+        elif isinstance(out, torch.Tensor):
+            # TODO (Alex): For now, we copy all tensors to all requests, because previously
+            # we used this to pass primitives by wrapping them as 0-d tensors. This is probably
+            # not optimal and should be revisited, but we need to ensure it doesn't cause issues
+            # for cases like max_num_seqs > 1 on audio producing models
+            for req_payload in batch_payloads:
+                req_payload[key] = copy_tensor(out)
+        elif isinstance(out, (int, float, bool)):
+            # primitives are constant across all requests in the batch, e.g., sampling rate
+            for req_payload in batch_payloads:
+                req_payload[key] = out
+        else:
+            logger.warning("Unsupported multimodal output type for key '%s': %s", key, type(out))
+    return batch_payloads
+
+
 def to_payload_element(
     element: object,
     idx: int,
@@ -216,9 +251,13 @@ def to_payload_element(
     # Cached per-token tensors are merged elsewhere; here a first dim equal to
     # either the hidden-aligned or scheduler-aligned batch length means a
     # per-request slice is required.
-    if isinstance(element, torch.Tensor) and (
-        (seq_len is not None and element.shape[0] == seq_len)
-        or (scheduled_seq_len is not None and element.shape[0] == scheduled_seq_len)
+    if (
+        isinstance(element, torch.Tensor)
+        and element.ndim > 0
+        and (
+            (seq_len is not None and element.shape[0] == seq_len)
+            or (scheduled_seq_len is not None and element.shape[0] == scheduled_seq_len)
+        )
     ):
         return element[start:end].contiguous()
     # Every other case is shared between prefix cache (passthrough data)

vllm_omni/worker/gpu_generation_model_runner.py

@@ -41,7 +41,7 @@
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.utils.mm_outputs import partition_payload_list
+from vllm_omni.utils.mm_outputs import partition_payload_list, split_mapping_outputs
 from vllm_omni.worker.gpu_ar_model_runner import (
     ExecuteModelState,
     _AsyncCPUPayloadSnapshot,
@@ -527,7 +527,7 @@ def sample_tokens(
 
         # Build per-request multimodal_outputs list (dedicated channel).
         # pooler_output is no longer used for multimodal data.
-        per_req_payloads: list[dict[str, object]] = []
+        per_req_payloads: list[Mapping[str, object]] = []
         pin_memory = is_pin_memory_available()
         accumulate_full_payload = self._should_accumulate_full_payload_output()
         use_async_output = (
@@ -564,22 +564,10 @@ def sample_tokens(
             for out in multimodal_outputs_raw:
                 per_req_payloads.append({"model_outputs": copy_tensor(out) if out is not None else None})
         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_raw.items():
-                    if isinstance(out, list):
-                        if len(out) != num_reqs:
-                            raise ValueError(
-                                f"Multimodal output list for key '{key}' has length {len(out)} "
-                                f"but expected {num_reqs} (one entry per request)."
-                            )
-                        mm_payload[key] = copy_tensor(out[i])
-                    elif isinstance(out, torch.Tensor):
-                        mm_payload[key] = copy_tensor(out)
-                    else:
-                        logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
-                per_req_payloads.append(_ensure_tensor_values(mm_payload))
+            per_req_payloads = [
+                _ensure_tensor_values(payload)
+                for payload in split_mapping_outputs(multimodal_outputs_raw, self.input_batch.num_reqs, copy_tensor)
+            ]
         else:
             raise RuntimeError("Unsupported diffusion output type")
         snapshot = None