From e0fc22996127f2e23dcb3f7d1ebfce1dcf97b322 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 16 Oct 2026 08:40:01 +0000 Subject: [PATCH] broadcast scalars for mrv1 Signed-off-by: Alex Brooks --- tests/utils/test_mm_outputs.py | 38 +++++++++++++- .../npu/worker/npu_generation_model_runner.py | 26 ++++------ vllm_omni/utils/mm_outputs.py | 49 +++++++++++++++++-- .../worker/gpu_generation_model_runner.py | 24 +++------ 4 files changed, 96 insertions(+), 41 deletions(-) diff --git a/tests/utils/test_mm_outputs.py b/tests/utils/test_mm_outputs.py index a07effb98d8..2d96f342356 100644 --- a/tests/utils/test_mm_outputs.py +++ b/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) 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 63e53be2c2d..44a8f202dad 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/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: diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index ccdb040dff0..dfd99d5965e 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/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) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 0106fc19085..b6ca261c0c0 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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