583 lines
25 KiB
Python
583 lines
25 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""OmniGenerationModelRunner — non-autoregressive stage runner on MR V2.
|
|
|
|
Used for stages like Code2Wav that convert codec codes to audio waveforms.
|
|
No token sampling or logits computation — model output goes directly into
|
|
``pooler_output``. Inherits from ``OmniGPUModelRunner`` for intermediate
|
|
buffer and lifecycle hooks.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import replace
|
|
from typing import Any
|
|
|
|
import torch
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.forward_context import set_forward_context
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
|
get_ep_all2all_manager,
|
|
)
|
|
from vllm.utils.torch_utils import PIN_MEMORY
|
|
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
|
|
from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput
|
|
from vllm.v1.worker.gpu.model_runner import (
|
|
BatchDescriptor,
|
|
ExecuteModelState,
|
|
IntermediateTensors,
|
|
)
|
|
|
|
from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
|
|
from vllm_omni.data_entry_keys import flatten_payload
|
|
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.worker_v2.omni_ar_model_runner import (
|
|
_async_copy_mm,
|
|
_ensure_tensor_values,
|
|
)
|
|
from vllm_omni.worker_v2.omni_model_runner import (
|
|
OmniGPUModelRunner,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _contains_cuda_tensor(value: Any) -> bool:
|
|
"""Return whether a nested output payload still owns CUDA storage."""
|
|
if isinstance(value, torch.Tensor):
|
|
return value.is_cuda
|
|
if isinstance(value, dict):
|
|
return any(_contains_cuda_tensor(v) for v in value.values())
|
|
if isinstance(value, (list, tuple)):
|
|
return any(_contains_cuda_tensor(v) for v in value)
|
|
return False
|
|
|
|
|
|
def _materialize_generation_value(
|
|
value: Any,
|
|
req_index: int,
|
|
num_reqs: int,
|
|
*,
|
|
top_level: bool = True,
|
|
copy_cpu: bool = True,
|
|
) -> Any:
|
|
"""Recursively select and own one request's generation payload."""
|
|
if isinstance(value, torch.Tensor):
|
|
if value.ndim > 0 and value.shape[0] == num_reqs:
|
|
value = value[req_index]
|
|
return value.detach().to(device="cpu", copy=copy_cpu and value.device.type == "cpu").contiguous()
|
|
if isinstance(value, dict):
|
|
return {
|
|
key: _materialize_generation_value(item, req_index, num_reqs, top_level=False, copy_cpu=copy_cpu)
|
|
for key, item in value.items()
|
|
}
|
|
if isinstance(value, list):
|
|
# Tensor lists (e.g. variable-length waveforms) carry one item per
|
|
# request even inside a dict. Scalar lists remain payload metadata.
|
|
per_request = top_level or all(item is None or isinstance(item, torch.Tensor) for item in value)
|
|
if per_request and len(value) == num_reqs:
|
|
# The selected item no longer has a batch axis.
|
|
return _materialize_generation_value(value[req_index], 0, -1, top_level=False, copy_cpu=copy_cpu)
|
|
return [
|
|
_materialize_generation_value(item, req_index, num_reqs, top_level=False, copy_cpu=copy_cpu)
|
|
for item in value
|
|
]
|
|
if isinstance(value, tuple):
|
|
return tuple(
|
|
_materialize_generation_value(item, req_index, num_reqs, top_level=False, copy_cpu=copy_cpu)
|
|
for item in value
|
|
)
|
|
return value
|
|
|
|
|
|
class OmniGenerationAsyncOutput(AsyncModelRunnerOutput):
|
|
"""Overlap generation-stage D2H copies with the next model step."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model_runner_output: OmniModelRunnerOutput,
|
|
multimodal_outputs: dict[str, Any] | None,
|
|
num_reqs: int,
|
|
main_stream: torch.cuda.Stream,
|
|
copy_stream: torch.cuda.Stream,
|
|
finalize_output: Any | None = None,
|
|
check_ep_fault: bool = False,
|
|
) -> None:
|
|
self.model_runner_output = model_runner_output
|
|
self.num_reqs = num_reqs
|
|
self.finalize_output = finalize_output
|
|
self.copy_event = torch.cuda.Event(blocking=True)
|
|
self._has_fault: torch.Tensor | None = None
|
|
|
|
with torch.cuda.stream(copy_stream):
|
|
copy_stream.wait_stream(main_stream)
|
|
self.multimodal_outputs_cpu = _async_copy_mm(
|
|
multimodal_outputs,
|
|
total_tokens=0,
|
|
copy_stream=copy_stream,
|
|
pin_memory=PIN_MEMORY,
|
|
)
|
|
if check_ep_fault:
|
|
has_fault = get_ep_all2all_manager().query_fault()
|
|
self._has_fault = has_fault.to("cpu", non_blocking=True)
|
|
self.copy_event.record(copy_stream)
|
|
|
|
def get_output(self) -> OmniModelRunnerOutput:
|
|
self.copy_event.synchronize()
|
|
if self._has_fault is not None and self._has_fault.item():
|
|
mask = get_ep_all2all_manager().query_active_mask()
|
|
raise RuntimeError(
|
|
"Fault detected in EP all2all communication: one or more ranks "
|
|
f"timed out during dispatch/combine. Mask: {mask.cpu().tolist()}"
|
|
)
|
|
payloads = OmniGenerationModelRunner._build_pooler_output_from_cpu(
|
|
self.multimodal_outputs_cpu,
|
|
self.num_reqs,
|
|
)
|
|
self.model_runner_output.multimodal_outputs = [
|
|
_ensure_tensor_values(payload) if payload else {} for payload in payloads
|
|
]
|
|
if self.finalize_output is not None:
|
|
return self.finalize_output(self.model_runner_output)
|
|
return self.model_runner_output
|
|
|
|
|
|
class OmniGenerationModelRunner(OmniGPUModelRunner):
|
|
"""Non-autoregressive generation runner (e.g. Code2Wav).
|
|
|
|
Overrides ``execute_model`` to skip the tensor-only assertion and
|
|
``sample_tokens`` to construct ``pooler_output`` from multimodal
|
|
model outputs without performing token sampling.
|
|
"""
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self._gen_model_output: Any = None
|
|
self._gen_input_batch: Any = None
|
|
# Placeholder for ExecuteModelState.hidden_states — allocated
|
|
# once and reused every step to avoid per-forward allocation.
|
|
self._dummy_hidden = torch.zeros(1, dtype=self.dtype, device=self.device)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Async-chunk support: replace prompt_token_ids for cached requests
|
|
# ------------------------------------------------------------------
|
|
|
|
def _handle_async_chunk_updates(self, scheduler_output: SchedulerOutput) -> None:
|
|
"""In-place update cached requests whose prompt_token_ids changed.
|
|
|
|
In async_chunk mode, the ``ChunkTransferAdapter`` replaces
|
|
``Request.prompt_token_ids`` with new codec frames for each
|
|
chunk and resets ``num_computed_tokens`` to 0. The scheduler
|
|
propagates the new ``prompt_token_ids`` via
|
|
``OmniCachedRequestData``.
|
|
|
|
Instead of remove + re-add (which involves free_indices churn
|
|
and redundant model_state init), we update the existing slot
|
|
in-place. This is safe for Code2Wav because:
|
|
- No KV cache / rope state to reinitialize
|
|
- staged writes are applied once at the end
|
|
|
|
The old intermediate buffer for this slot is cleared here; the
|
|
inherited ``OmniGPUModelRunner.update_requests`` (called right after
|
|
this method in ``execute_model``) writes the current chunk state.
|
|
"""
|
|
cached = scheduler_output.scheduled_cached_reqs
|
|
if not cached.req_ids:
|
|
return
|
|
|
|
if not isinstance(cached, OmniCachedRequestData):
|
|
return
|
|
|
|
new_prompt_ids = cached.prompt_token_ids
|
|
if not new_prompt_ids:
|
|
return
|
|
|
|
updated = False
|
|
released_chunks: list[OmniNewRequestData] = []
|
|
|
|
for i, req_id in enumerate(cached.req_ids):
|
|
new_ids = new_prompt_ids.get(req_id)
|
|
if new_ids is None:
|
|
continue
|
|
|
|
req_idx = self.req_states.req_id_to_index.get(req_id)
|
|
if req_idx is None:
|
|
block_ids = cached.new_block_ids[i]
|
|
released_chunks.append(
|
|
OmniNewRequestData(
|
|
req_id=req_id,
|
|
prompt_token_ids=new_ids,
|
|
mm_features=[],
|
|
sampling_params=None,
|
|
pooling_params=None,
|
|
block_ids=block_ids if block_ids is not None else tuple(),
|
|
num_computed_tokens=0,
|
|
lora_request=None,
|
|
prompt_embeds=None,
|
|
prefill_token_ids=new_ids,
|
|
additional_information=cached.additional_information.get(req_id),
|
|
)
|
|
)
|
|
continue
|
|
|
|
self.model_state.remove_request(req_idx)
|
|
|
|
# In-place update token state — same slot, no remove/re-add.
|
|
# .np[] = direct write (no GPU buffer); stage_write = GPU-synced.
|
|
n = len(new_ids)
|
|
self.req_states.prompt_len.np[req_idx] = n
|
|
self.req_states.prefill_len.np[req_idx] = n
|
|
self.req_states.total_len.stage_write_elem(req_idx, n)
|
|
self.req_states.all_token_ids.stage_write(req_idx, 0, new_ids)
|
|
self.req_states.num_computed_tokens.stage_write_elem(req_idx, 0)
|
|
self.req_states.num_computed_prefill_tokens[req_idx] = 0
|
|
|
|
updated = True
|
|
|
|
if released_chunks:
|
|
self.add_requests(
|
|
replace(
|
|
scheduler_output,
|
|
scheduled_new_reqs=released_chunks,
|
|
)
|
|
)
|
|
if updated:
|
|
self.req_states.apply_staged_writes()
|
|
|
|
def _release_generation_slots(self, input_batch: Any) -> None:
|
|
if not getattr(self.model_config, "async_chunk", False):
|
|
return
|
|
model_state = getattr(self, "model_state", None)
|
|
remove_request = getattr(self, "_remove_request", None)
|
|
if model_state is None or remove_request is None:
|
|
return
|
|
for i in range(input_batch.num_reqs):
|
|
req_id = input_batch.req_ids[i]
|
|
req_idx = int(input_batch.idx_mapping_np[i])
|
|
model_state.remove_request(req_idx)
|
|
remove_request(req_id)
|
|
|
|
# ------------------------------------------------------------------
|
|
# profile / warmup — skip sampler since there are no logits
|
|
# ------------------------------------------------------------------
|
|
|
|
@torch.inference_mode()
|
|
def profile_run(self) -> None:
|
|
"""Generation models have no KV cache — skip profiling.
|
|
|
|
Code2Wav shares GPU memory with the Talker stage (same device);
|
|
its memory footprint is managed via ``gpu_memory_utilization``
|
|
config, not profiled dynamically. Running the real model with
|
|
random input_ids causes out-of-bounds indexing in codec lookup
|
|
tables.
|
|
"""
|
|
torch.accelerator.synchronize()
|
|
|
|
# ------------------------------------------------------------------
|
|
# execute_model — run the generation model, store raw output
|
|
# ------------------------------------------------------------------
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
dummy_run: bool = False,
|
|
skip_attn_for_dummy_run: bool = False,
|
|
is_profile: bool = False,
|
|
context_len: int = 0,
|
|
) -> ModelRunnerOutput | IntermediateTensors | None:
|
|
if not dummy_run:
|
|
self._prepare_native_data_plane(scheduler_output)
|
|
self.finish_requests(scheduler_output)
|
|
self.free_states(scheduler_output)
|
|
if (
|
|
scheduler_output.total_num_scheduled_tokens == 0
|
|
and not scheduler_output.scheduled_new_reqs
|
|
and not scheduler_output.scheduled_cached_reqs.req_ids
|
|
):
|
|
# Registration/terminal/abort work has been applied above.
|
|
# There is no new or cached request state to add or update.
|
|
self._apply_block_table_staged_writes_if_available()
|
|
empty_output = self.kv_connector.no_forward(scheduler_output)
|
|
return self._attach_native_data_plane_signals(
|
|
self._merge_ec_connector_no_forward(scheduler_output, empty_output)
|
|
)
|
|
# Handle async_chunk prompt_token_ids replacement for cached
|
|
# requests BEFORE add/update — update the existing slot
|
|
# in-place with the new chunk's tokens.
|
|
self._handle_async_chunk_updates(scheduler_output)
|
|
self.add_requests(scheduler_output)
|
|
self.update_requests(scheduler_output)
|
|
self._sync_native_data_plane_payloads(scheduler_output)
|
|
self._apply_block_table_staged_writes_if_available()
|
|
if scheduler_output.total_num_scheduled_tokens == 0:
|
|
empty_output = self.kv_connector.no_forward(scheduler_output)
|
|
return self._attach_native_data_plane_signals(
|
|
self._merge_ec_connector_no_forward(scheduler_output, empty_output)
|
|
)
|
|
|
|
num_reqs = len(scheduler_output.num_scheduled_tokens)
|
|
num_toks = scheduler_output.total_num_scheduled_tokens
|
|
max_query_len = max(scheduler_output.num_scheduled_tokens.values())
|
|
batch_req_state, uniform_tok_count = self.gather_batch_req_state(scheduler_output, dummy_run)
|
|
if batch_req_state is not None:
|
|
num_toks = batch_req_state.num_tokens
|
|
batch_desc, dp_sync = self._dispatch_batch_descriptor(
|
|
num_reqs=num_reqs,
|
|
num_toks=num_toks,
|
|
uniform_tok_count=uniform_tok_count,
|
|
num_active_loras=0,
|
|
# Generation stages return structured OmniOutput payloads and do
|
|
# not expose the tensor-only outer graph contract. Their model may
|
|
# still use an internal decoder graph, but the vLLM outer runner
|
|
# always uses eager dispatch here.
|
|
use_eager=True,
|
|
max_query_len=max_query_len,
|
|
)
|
|
|
|
if batch_desc.num_tokens == 0:
|
|
empty_output = self.kv_connector.no_forward(scheduler_output)
|
|
return self._attach_native_data_plane_signals(
|
|
self._merge_ec_connector_no_forward(scheduler_output, empty_output)
|
|
)
|
|
|
|
if not dummy_run:
|
|
assert batch_req_state is not None
|
|
input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
|
|
else:
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch
|
|
|
|
input_batch = InputBatch.make_dummy(
|
|
batch_desc.num_reqs or num_reqs,
|
|
batch_desc.num_tokens,
|
|
self.input_buffers,
|
|
max_query_len=batch_desc.max_query_len,
|
|
)
|
|
|
|
attn_metadata = None
|
|
slot_mappings_by_layer = None
|
|
|
|
input_ids, inputs_embeds, ec_connector_output = self._prepare_mm_inputs(
|
|
scheduler_output,
|
|
input_batch,
|
|
dummy_run=dummy_run,
|
|
)
|
|
|
|
model_inputs = {
|
|
"input_ids": input_ids,
|
|
"positions": input_batch.positions,
|
|
"inputs_embeds": inputs_embeds,
|
|
"intermediate_tensors": intermediate_tensors,
|
|
**self.model_state.prepare_inputs(input_batch, self.req_states),
|
|
}
|
|
self._add_legacy_forward_inputs(model_inputs, input_batch)
|
|
|
|
eplb = getattr(self, "eplb", None)
|
|
if eplb is not None:
|
|
eplb.prepare_forward(self.model_config, input_batch.num_tokens)
|
|
|
|
batch_descriptor = BatchDescriptor(
|
|
num_tokens=input_batch.num_tokens_after_padding,
|
|
has_lora=self.lora_config is not None,
|
|
num_active_loras=batch_desc.num_active_loras,
|
|
)
|
|
with set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=input_batch.num_tokens_after_padding,
|
|
cudagraph_runtime_mode=batch_desc.cg_mode,
|
|
batch_descriptor=batch_descriptor,
|
|
slot_mapping=slot_mappings_by_layer,
|
|
is_padding=input_batch.is_padding,
|
|
):
|
|
self.kv_connector.pre_forward(scheduler_output)
|
|
if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
|
|
assert self.cudagraph_manager is not None
|
|
model_output = self.cudagraph_manager.run_pw_graph(
|
|
self.model,
|
|
model_inputs,
|
|
)
|
|
else:
|
|
model_output = self.model(**model_inputs)
|
|
|
|
# Convert raw model output to OmniOutput.
|
|
if not isinstance(model_output, OmniOutput):
|
|
buffer_list = self.model_state.intermediate_buffer.gather(input_batch)
|
|
make_output_kwargs = {"model_intermediate_buffer": buffer_list}
|
|
make_output_kwargs["runtime_additional_information"] = buffer_list
|
|
model_output = self.model.make_omni_output(model_output, **make_output_kwargs)
|
|
|
|
self._gen_model_output = model_output
|
|
self._gen_input_batch = input_batch
|
|
|
|
# ExecuteModelState is required by the upstream engine loop
|
|
# (EngineCore checks execute_model_state is not None before
|
|
# calling sample_tokens).
|
|
self.execute_model_state = ExecuteModelState(
|
|
input_batch=input_batch,
|
|
attn_metadata=None,
|
|
slot_mappings_by_layer=None,
|
|
hidden_states=self._dummy_hidden,
|
|
aux_hidden_states=None,
|
|
finished_req_ids=scheduler_output.finished_req_ids,
|
|
dp_sync=dp_sync,
|
|
ec_connector_output=ec_connector_output,
|
|
routed_experts=None,
|
|
)
|
|
return None
|
|
|
|
# ------------------------------------------------------------------
|
|
# sample_tokens — build pooler_output, no actual sampling
|
|
# ------------------------------------------------------------------
|
|
|
|
@torch.inference_mode()
|
|
def sample_tokens(
|
|
self,
|
|
grammar_output: GrammarOutput | None = None,
|
|
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | None:
|
|
model_output = self._gen_model_output
|
|
input_batch = self._gen_input_batch
|
|
execute_model_state = self.execute_model_state
|
|
self._gen_model_output = None
|
|
self._gen_input_batch = None
|
|
self.execute_model_state = None
|
|
|
|
if model_output is None or input_batch is None or execute_model_state is None:
|
|
return None
|
|
|
|
num_reqs = input_batch.num_reqs
|
|
|
|
# Mark all scheduled tokens as computed so the scheduler does
|
|
# not re-schedule them. Unlike AR stages we do NOT call
|
|
# self.postprocess() — that kernel advances num_computed_tokens
|
|
# by 1 and emits sampled tokens, which would cause check_stop
|
|
# to fire. Instead, set num_computed_tokens = prompt_len
|
|
# directly, matching V1's behavior.
|
|
for i in range(num_reqs):
|
|
req_idx = int(input_batch.idx_mapping_np[i])
|
|
prompt_len = int(self.req_states.prompt_len.np[req_idx])
|
|
self.req_states.num_computed_tokens.stage_write_elem(req_idx, prompt_len)
|
|
self.req_states.num_computed_tokens.apply_write()
|
|
kv_connector_output = self.kv_connector.post_forward(execute_model_state.finished_req_ids)
|
|
|
|
# Async finalization can run after the input batch is reused for the
|
|
# next scheduler step, so output metadata must not alias the batch.
|
|
req_ids = list(input_batch.req_ids)
|
|
|
|
# Generation models don't do token sampling. Return one empty
|
|
# list per request so the scheduler does NOT trigger check_stop
|
|
# (which would prematurely finish the request). The request
|
|
# stays RUNNING until the orchestrator marks it done via
|
|
# chunk_transfer_adapter.finished_requests.
|
|
sampled_token_ids: list[list[int]] = [[] for _ in range(len(req_ids))]
|
|
output = OmniModelRunnerOutput(
|
|
req_ids=req_ids,
|
|
req_id_to_index={rid: i for i, rid in enumerate(req_ids)},
|
|
sampled_token_ids=sampled_token_ids,
|
|
# Match the V1 generation runner contract: final generation stages
|
|
# publish audio only through multimodal_outputs. If the same
|
|
# payload is also exposed as pooling_output, OutputProcessor routes
|
|
# it through the pooling path and can duplicate accumulated audio.
|
|
pooler_output=None,
|
|
multimodal_outputs=None,
|
|
kv_connector_output=kv_connector_output,
|
|
ec_connector_output=execute_model_state.ec_connector_output,
|
|
sampled_token_ids_materialized=True,
|
|
)
|
|
|
|
raw_multimodal_outputs = model_output.multimodal_outputs
|
|
can_copy_async = getattr(getattr(self, "device", None), "type", "cpu") == "cuda" and (
|
|
self.check_ep_fault or _contains_cuda_tensor(raw_multimodal_outputs)
|
|
)
|
|
if can_copy_async:
|
|
async_output = OmniGenerationAsyncOutput(
|
|
model_runner_output=output,
|
|
multimodal_outputs=raw_multimodal_outputs,
|
|
num_reqs=num_reqs,
|
|
main_stream=self.main_stream,
|
|
copy_stream=self.output_copy_stream,
|
|
finalize_output=self._finalize_native_data_plane_output,
|
|
check_ep_fault=self.check_ep_fault,
|
|
)
|
|
self._reserve_native_data_plane_outputs(list(req_ids))
|
|
self._release_generation_slots(input_batch)
|
|
return async_output
|
|
|
|
# CPU execution keeps the synchronous materialization path. Final
|
|
# generation stages publish audio only through multimodal_outputs,
|
|
# matching the V1 generation-runner contract.
|
|
self._reserve_native_data_plane_outputs(list(req_ids))
|
|
multimodal_outputs = self._build_pooler_output(model_output, num_reqs)
|
|
output.multimodal_outputs = [
|
|
_ensure_tensor_values(payload) if payload else {} for payload in multimodal_outputs
|
|
]
|
|
self._release_generation_slots(input_batch)
|
|
return self._finalize_native_data_plane_output(output)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _apply_block_table_staged_writes_if_available(self) -> None:
|
|
"""Flush block-table writes only when this stage has a writer.
|
|
|
|
Generation stages such as Code2Wav do not use KV cache or attention
|
|
metadata. Newer vLLM BlockTable implementations assert that a fused
|
|
writer exists before applying staged writes; no-KV generation stages
|
|
legitimately do not have one.
|
|
"""
|
|
block_tables = getattr(self, "block_tables", None)
|
|
if block_tables is None:
|
|
return
|
|
if getattr(block_tables, "fused_writer", None) is None:
|
|
return
|
|
block_tables.apply_staged_writes()
|
|
|
|
@staticmethod
|
|
def _build_pooler_output(
|
|
model_output: OmniOutput,
|
|
num_reqs: int,
|
|
) -> list[dict[str, Any] | None]:
|
|
"""Extract per-request pooler payloads from model output.
|
|
|
|
Code2Wav's ``make_omni_output`` returns
|
|
``{"model_outputs": [tensor_per_req, ...], "sr": [...]}``,
|
|
so each value is a ``list`` with ``len == num_reqs``.
|
|
"""
|
|
mm = model_output.multimodal_outputs
|
|
if not isinstance(mm, dict):
|
|
logger.warning(
|
|
"Unexpected multimodal_outputs type: %s; returning empty pooler_output",
|
|
type(mm).__name__ if mm is not None else "None",
|
|
)
|
|
return [None] * num_reqs
|
|
|
|
pooler: list[dict[str, Any] | None] = []
|
|
for i in range(num_reqs):
|
|
payload = {key: _materialize_generation_value(value, i, num_reqs) for key, value in mm.items()}
|
|
pooler.append(flatten_payload(payload) if isinstance(payload, dict) else payload)
|
|
return pooler
|
|
|
|
@staticmethod
|
|
def _build_pooler_output_from_cpu(
|
|
multimodal_outputs: dict[str, Any] | None,
|
|
num_reqs: int,
|
|
) -> list[dict[str, Any] | None]:
|
|
"""Build per-request payloads after an async whole-batch D2H copy."""
|
|
if not isinstance(multimodal_outputs, dict):
|
|
return [None] * num_reqs
|
|
|
|
pooler: list[dict[str, Any] | None] = []
|
|
for i in range(num_reqs):
|
|
payload = {
|
|
key: _materialize_generation_value(value, i, num_reqs, copy_cpu=False)
|
|
for key, value in multimodal_outputs.items()
|
|
}
|
|
pooler.append(flatten_payload(payload) if isinstance(payload, dict) else payload)
|
|
return pooler
|