367 lines
17 KiB
Python
367 lines
17 KiB
Python
"""AR GPU Model Runner for vLLM-Omni.
|
|
|
|
Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
|
|
and also outputs sampled tokens.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
from vllm.forward_context import BatchDescriptor
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
from vllm.v1.outputs import AsyncModelRunnerOutput
|
|
from vllm.v1.structured_output.utils import apply_grammar_bitmask
|
|
from vllm.v1.utils import record_function_or_nullcontext
|
|
from vllm.v1.worker.gpu_model_runner import (
|
|
EMPTY_MODEL_RUNNER_OUTPUT,
|
|
AsyncGPUModelRunnerOutput,
|
|
IntermediateTensors,
|
|
get_pp_group,
|
|
get_tp_group,
|
|
has_kv_transfer_group,
|
|
set_forward_context,
|
|
)
|
|
from vllm.v1.worker.utils import is_residual_scattered_for_sp
|
|
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class GPUARModelRunner(OmniGPUModelRunner):
|
|
"""Autoregressive GPU model runner that returns hidden states per request.
|
|
|
|
This runner follows the same preparation and forward path as GPUModelRunner
|
|
(inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
|
|
and additionally performs lightweight sampling so that sampled tokens are
|
|
available in outputs. Hidden representations are taken at the same indices
|
|
that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
|
|
# each model stage has their own hidden size
|
|
self.hidden_size = self.model_config.hf_text_config.hidden_size
|
|
self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
|
|
with record_function_or_nullcontext("Preprocess"):
|
|
with self.synchronize_input_prep():
|
|
# Update persistent batch states.
|
|
self._update_states(scheduler_output)
|
|
|
|
# Decode per-request prompt_embeds / additional_information payloads
|
|
# (if present) into CPU tensors
|
|
self._decode_and_store_request_payloads(scheduler_output)
|
|
|
|
if not scheduler_output.total_num_scheduled_tokens:
|
|
if not has_kv_transfer_group():
|
|
# Return empty ModelRunnerOutput if no work to do.
|
|
return EMPTY_MODEL_RUNNER_OUTPUT
|
|
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
|
|
if self.cache_config.kv_sharing_fast_prefill:
|
|
assert not self.input_batch.num_prompt_logprobs, (
|
|
"--kv-sharing-fast-prefill produces incorrect "
|
|
"logprobs for prompt tokens, tokens, please disable "
|
|
"it when the requests need prompt logprobs"
|
|
)
|
|
|
|
# Prepare the decoder inputs.
|
|
(
|
|
attn_metadata,
|
|
logits_indices,
|
|
spec_decode_metadata,
|
|
num_scheduled_tokens_np,
|
|
spec_decode_common_attn_metadata,
|
|
max_query_len,
|
|
ubatch_slices,
|
|
num_tokens_after_padding,
|
|
) = self._prepare_inputs(scheduler_output)
|
|
|
|
(
|
|
num_scheduled_tokens,
|
|
num_input_tokens,
|
|
num_tokens_across_dp,
|
|
input_ids,
|
|
inputs_embeds,
|
|
positions,
|
|
intermediate_tensors,
|
|
model_kwargs,
|
|
per_req_additional_information,
|
|
) = self._preprocess(
|
|
scheduler_output,
|
|
num_scheduled_tokens_np,
|
|
intermediate_tensors,
|
|
ubatch_slices,
|
|
num_tokens_after_padding,
|
|
)
|
|
|
|
uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
|
|
num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
|
|
)
|
|
batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=uniform_decode)
|
|
cudagraph_runtime_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(batch_descriptor)
|
|
|
|
# This is currently to get around the assert in the DPMetadata
|
|
# where it wants `num_tokens_across_dp` to align with `num_tokens`
|
|
if ubatch_slices is not None:
|
|
num_input_tokens = ubatch_slices[0].num_tokens
|
|
|
|
# Run the model.
|
|
# Use persistent buffers for CUDA graphs.
|
|
with (
|
|
set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_input_tokens,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
|
batch_descriptor=batch_descriptor,
|
|
ubatch_slices=ubatch_slices,
|
|
),
|
|
record_function_or_nullcontext("Forward"),
|
|
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
|
|
):
|
|
model_kwargs_extra = self._build_model_kwargs_extra(per_req_additional_information, num_scheduled_tokens_np)
|
|
model_output = self.model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**model_kwargs,
|
|
sampling_metadata=self.input_batch.sampling_metadata,
|
|
logits_index=logits_indices,
|
|
sampler=self.sampler,
|
|
**model_kwargs_extra,
|
|
)
|
|
|
|
with record_function_or_nullcontext("Postprocess"):
|
|
if self.use_aux_hidden_state_outputs:
|
|
# True when EAGLE 3 is used.
|
|
hidden_states, aux_hidden_states = model_output
|
|
else:
|
|
# Common case.
|
|
hidden_states = model_output
|
|
aux_hidden_states = None
|
|
|
|
# hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
|
|
# hidden_states
|
|
# )
|
|
multimodal_outputs = model_output.multimodal_outputs
|
|
hidden_states = model_output.text_hidden_states
|
|
# The model side may return per-request additional_information updates (model-agnostic channel).
|
|
# Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
|
|
# the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
|
|
self._process_additional_information_updates(multimodal_outputs)
|
|
if not self.broadcast_pp_output:
|
|
# Common case.
|
|
if not get_pp_group().is_last_rank:
|
|
# Return the intermediate tensors.
|
|
assert isinstance(hidden_states, IntermediateTensors)
|
|
hidden_states.kv_connector_output = kv_connector_output
|
|
return hidden_states
|
|
|
|
if self.is_pooling_model:
|
|
# Return the pooling output.
|
|
output = self._pool(hidden_states, num_scheduled_tokens, num_scheduled_tokens_np)
|
|
output.kv_connector_output = kv_connector_output
|
|
return output
|
|
|
|
sample_hidden_states = hidden_states[logits_indices]
|
|
logits = self.model.compute_logits(sample_hidden_states)
|
|
else:
|
|
# Rare case.
|
|
assert not self.is_pooling_model
|
|
|
|
if not get_pp_group().is_last_rank:
|
|
all_gather_tensors = {
|
|
"residual": not is_residual_scattered_for_sp(self.vllm_config, num_input_tokens)
|
|
}
|
|
get_pp_group().send_tensor_dict(
|
|
hidden_states.tensors,
|
|
all_gather_group=get_tp_group(),
|
|
all_gather_tensors=all_gather_tensors,
|
|
)
|
|
logits = None
|
|
else:
|
|
sample_hidden_states = hidden_states[logits_indices]
|
|
logits = self.model.compute_logits(sample_hidden_states)
|
|
|
|
model_output_broadcast_data = {}
|
|
if logits is not None:
|
|
model_output_broadcast_data["logits"] = logits.contiguous()
|
|
|
|
model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
|
|
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
|
|
)
|
|
assert model_output_broadcast_data is not None
|
|
logits = model_output_broadcast_data["logits"]
|
|
|
|
# Apply structured output bitmasks if present
|
|
if scheduler_output.grammar_bitmask is not None:
|
|
apply_grammar_bitmask(scheduler_output, self.input_batch, logits, self.device)
|
|
|
|
with record_function_or_nullcontext("Sample"):
|
|
sampler_output = self._sample(logits, spec_decode_metadata)
|
|
|
|
def propose_draft_token_ids(sampled_token_ids):
|
|
assert spec_decode_common_attn_metadata is not None
|
|
with record_function_or_nullcontext("Draft"):
|
|
self._draft_token_ids = self.propose_draft_token_ids(
|
|
scheduler_output,
|
|
sampled_token_ids,
|
|
self.input_batch.sampling_metadata,
|
|
hidden_states,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
)
|
|
|
|
use_padded_batch_for_eagle = (
|
|
self.speculative_config
|
|
and self.speculative_config.use_eagle()
|
|
and not self.speculative_config.disable_padded_drafter_batch
|
|
)
|
|
effective_drafter_max_model_len = self.max_model_len
|
|
if effective_drafter_max_model_len is None:
|
|
effective_drafter_max_model_len = self.model_config.max_model_len
|
|
if (
|
|
self.speculative_config
|
|
and self.speculative_config.draft_model_config is not None
|
|
and self.speculative_config.draft_model_config.max_model_len is not None
|
|
):
|
|
effective_drafter_max_model_len = self.speculative_config.draft_model_config.max_model_len
|
|
input_fits_in_drafter = spec_decode_common_attn_metadata and (
|
|
spec_decode_common_attn_metadata.seq_lens.max() + self.speculative_config.num_speculative_tokens
|
|
<= effective_drafter_max_model_len
|
|
)
|
|
if use_padded_batch_for_eagle and input_fits_in_drafter:
|
|
# EAGLE speculative decoding can use the GPU sampled tokens
|
|
# as inputs, and does not need to wait for bookkeeping to finish.
|
|
propose_draft_token_ids(sampler_output.sampled_token_ids)
|
|
|
|
with record_function_or_nullcontext("Bookkeep"):
|
|
(
|
|
num_nans_in_logits,
|
|
logprobs_lists,
|
|
valid_sampled_token_ids,
|
|
prompt_logprobs_dict,
|
|
req_ids_output_copy,
|
|
req_id_to_index_output_copy,
|
|
invalid_req_indices,
|
|
) = self._bookkeeping_sync(
|
|
scheduler_output,
|
|
sampler_output,
|
|
logits,
|
|
hidden_states,
|
|
num_scheduled_tokens,
|
|
)
|
|
|
|
if self.speculative_config and not use_padded_batch_for_eagle and input_fits_in_drafter:
|
|
# ngram and other speculative decoding methods use the sampled
|
|
# tokens on the CPU, so they are run after bookkeeping.
|
|
propose_draft_token_ids(valid_sampled_token_ids)
|
|
|
|
with record_function_or_nullcontext("EPLB"):
|
|
self.eplb_step()
|
|
|
|
# Convert to per-request tensors on CPU
|
|
hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
|
|
# pooler_output: list[torch.Tensor | None] = []
|
|
pooler_output: list[torch.Tensor | None] = []
|
|
prev_logits_index = 0
|
|
for rid, logits_index in zip(req_ids_output_copy, logits_indices):
|
|
# Base payload: hidden slice for this request in this iteration
|
|
hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1]
|
|
payload: dict[str, object] = {"hidden": hidden_slice}
|
|
# Merge multimodal_outputs if present
|
|
if isinstance(multimodal_outputs, dict) and multimodal_outputs:
|
|
mm_payload: dict[str, object] = {}
|
|
for k, v in multimodal_outputs.items():
|
|
try:
|
|
# Case 1: tensor aligned on token dimension
|
|
if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
|
|
mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
|
|
elif isinstance(v, torch.Tensor) and v.shape[0] != hidden_states_cpu.shape[0]:
|
|
logger.error(
|
|
f"Error in merge multimodal outputs: Tensor dimension mismatch, \
|
|
{v.shape} != {hidden_states_cpu.shape} for {k}"
|
|
)
|
|
# Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers)
|
|
elif isinstance(v, dict):
|
|
sub_dict: dict[str, torch.Tensor] = {}
|
|
for sk, sv in v.items():
|
|
if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
|
|
sub_dict[str(sk)] = (
|
|
sv.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
|
|
)
|
|
if sub_dict:
|
|
mm_payload[k] = sub_dict
|
|
elif isinstance(v, list):
|
|
element = v[0]
|
|
if isinstance(element, torch.Tensor):
|
|
element = element.detach().to("cpu").contiguous()
|
|
multimodal_outputs[k] = v[1:] if len(v) > 1 else v
|
|
mm_payload[k] = element
|
|
except Exception as e:
|
|
# Best-effort; skip malformed entries
|
|
logger.error(f"Error in merge multimodal outputs: {e}")
|
|
if mm_payload:
|
|
payload.update(mm_payload)
|
|
pooler_output.append(payload) # type: ignore[arg-type]
|
|
prev_logits_index = logits_index + 1
|
|
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 self.vllm_config.model_config.engine_output_type != "text" else None),
|
|
kv_connector_output=kv_connector_output,
|
|
num_nans_in_logits=num_nans_in_logits,
|
|
)
|
|
|
|
if not self.use_async_scheduling:
|
|
return output
|
|
|
|
return AsyncGPUModelRunnerOutput(
|
|
model_runner_output=output,
|
|
sampled_token_ids=sampler_output.sampled_token_ids,
|
|
invalid_req_indices=invalid_req_indices,
|
|
async_output_copy_stream=self.async_output_copy_stream,
|
|
)
|
|
|
|
def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
|
|
req_state = self.requests.get(req_id)
|
|
if req_state is None:
|
|
return
|
|
existing = getattr(req_state, "additional_information_cpu", {})
|
|
if not isinstance(existing, dict):
|
|
existing = {}
|
|
merged = dict(existing)
|
|
for k, v in upd.items():
|
|
if isinstance(v, torch.Tensor):
|
|
merged[k] = v.detach().to("cpu").contiguous()
|
|
elif isinstance(v, list):
|
|
new_list = []
|
|
for item in v:
|
|
if isinstance(item, torch.Tensor):
|
|
new_list.append(item.detach().to("cpu").contiguous())
|
|
else:
|
|
new_list.append(item)
|
|
merged[k] = new_list
|
|
else:
|
|
merged[k] = v
|
|
setattr(req_state, "additional_information_cpu", merged)
|
|
|
|
# ===== Helper functions extracted for clarity and reuse =====
|