anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

allow omni model custom pre-process function

Browse Source
This commit is contained in:
yinpeiqi committed 2025-11-27 14:20:00 +07:00
1 parent 99c46e1adc
commit 423f918e79
2 files changed
+140 -234
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+108 -152
View File
@@ -60,6 +60,7 @@ class Qwen2_5OmniForConditionalGeneration(
):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.has_process = False
self.have_multimodal_outputs = True
config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
multimodal_config = vllm_config.model_config.multimodal_config
@@ -90,6 +91,8 @@ class Qwen2_5OmniForConditionalGeneration(
self.token2wav = None
elif self.model_stage == "talker":
# temporal use this prepresation
self.has_process = True
self.thinker = None
# Initialize talker model wrapper (handles projection + LM)
self.talker = init_vllm_registered_model(
@@ -236,13 +239,14 @@ class Qwen2_5OmniForConditionalGeneration(
if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
inputs_embeds = inputs_embeds.to(thinker_dev)
# Run thinker
thinker_output = self.thinker(
input_ids=input_ids,
positions=positions[0],
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**kwargs,
)
with torch.inference_mode():
thinker_output = self.thinker(
input_ids=input_ids,
positions=positions[0],
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**kwargs,
)
if isinstance(thinker_output, tuple):
embeds, text_hidden_states = thinker_output
@@ -257,151 +261,13 @@ class Qwen2_5OmniForConditionalGeneration(
# 2) Talker (if codec not provided)
if self.model_stage == "talker":
# Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
# Rules:
# - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
# - Decode segments consist of a single non-special token.
# - If additional_information is provided
# (can be a list split by request or a concatenated tensor plus a list of shapes),
# then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
# - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them;
# otherwise, keep the original embedding.
# mock data for profile
if input_ids is None and additional_information is None:
input_ids = torch.zeros(
inputs_embeds.shape[0],
dtype=torch.long,
device=inputs_embeds.device,
inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
)
additional_information = {}
self.thinker_reply_part = torch.zeros_like(inputs_embeds)
is_profile = True
else:
is_profile = False
# Ensure we have base embeddings when only ids are provided
if inputs_embeds is None and input_ids is not None:
inputs_embeds = self.talker.get_input_embeddings(input_ids)
# ------- Request-scoped additional information (no cross-request concat) -------
request_ids: Optional[list[str]] = kwargs.get("request_ids") # ordered
request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans")
addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
runtime_addi = kwargs.get("runtime_additional_information")
# Normalize runtime_addi into a mapping by request_id for convenience
runtime_addi_by_req: dict[str, dict] = {}
if (
isinstance(request_ids, list)
and isinstance(runtime_addi, list)
and len(runtime_addi) == len(request_ids)
):
for i, rid in enumerate(request_ids):
if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
runtime_addi_by_req[rid] = runtime_addi[i]
elif isinstance(request_ids, list) and isinstance(runtime_addi, dict):
for rid in request_ids:
if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict):
runtime_addi_by_req[rid] = runtime_addi[rid]
# Containers to return per-request updates (e.g., thinker_reply_part_per_request)
update_by_req_id: dict[str, dict] = {}
# ------- Prefill: span_len > 1 -------
if (
not is_profile
and isinstance(request_ids, list)
and isinstance(request_token_spans, list)
and isinstance(addi_by_req, dict)
):
for idx_req, rid in enumerate(request_ids):
s, e = request_token_spans[idx_req]
span_len = int(e) - int(s)
if span_len <= 1:
continue
info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
if not isinstance(info, dict):
info = {}
pe = info.get("prompt_embeds") # Tensor [P,H]
tr = info.get("thinker_result") # Tensor [K,H]
ptoks = info.get("prompt_token_ids") # list[int]
otoks = info.get("thinker_output_token_ids") # list[int]
if not isinstance(pe, torch.Tensor):
pe = torch.zeros(
0,
self.talker.config.hidden_size,
dtype=inputs_embeds.dtype,
device=self._module_device(self.model),
)
if not isinstance(tr, torch.Tensor):
tr = torch.zeros(
0,
self.talker.config.hidden_size,
dtype=inputs_embeds.dtype,
device=self._module_device(self.model),
)
if not isinstance(ptoks, (list, torch.Tensor)):
ptoks = []
if not isinstance(otoks, (list, torch.Tensor)):
otoks = []
req_input_ids, req_embeds = self._thinker_to_talker_prefill(
voice_type=voice_type,
output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
output_token_ids=otoks,
thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
prompt_token_ids=ptoks,
)
seg_len = min(span_len, req_embeds.shape[0])
inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
input_ids[s : s + seg_len] = req_input_ids
# Prepare per-request reply queue for subsequent decode: drop first row
if tr.ndim == 2 and tr.shape[0] > 0:
update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = (
tr[1:].detach().to("cpu").contiguous()
)
# ------- Decode: span_len == 1 -------
if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
for idx_req, rid in enumerate(request_ids):
s, e = request_token_spans[idx_req]
if (int(e) - int(s)) != 1:
continue
# choose step vector in priority order
step_vec = None
# A) runtime queue
q = None
if isinstance(rid, str):
q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
if isinstance(q, torch.Tensor) and q.numel() > 0:
step_vec = q[0:1]
new_q = q[1:].detach().to("cpu").contiguous()
update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
else:
# B) per-request provided decode vector (optional)
info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
if isinstance(dv, torch.Tensor) and dv.numel() > 0:
step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
elif (
hasattr(self, "thinker_reply_part")
and isinstance(self.thinker_reply_part, torch.Tensor)
and self.thinker_reply_part.numel() > 0
):
# C) fallback shared pool
step_vec = self.thinker_reply_part[0:1]
self.thinker_reply_part = self.thinker_reply_part[1:]
if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
one_id = input_ids[s : s + 1]
_, one_embed = self._thinker_to_talker_decode_one_step(
output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
output_token_ids=one_id,
)
inputs_embeds[s] = one_embed[0]
with torch.inference_mode():
talker_hidden = self.talker(
@@ -409,10 +275,6 @@ class Qwen2_5OmniForConditionalGeneration(
positions=positions[0],
inputs_embeds=inputs_embeds,
)
multimodal_outputs: dict = None
# Return updates if any
if update_by_req_id:
multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
if sampling_metadata is not None:
# the padding token id is set to text model's pad token id,
@@ -421,7 +283,7 @@ class Qwen2_5OmniForConditionalGeneration(
return OmniOutput(
text_hidden_states=talker_hidden,
multimodal_outputs=multimodal_outputs,
multimodal_outputs=None,
)
if self.model_stage == "code2wav":
@@ -733,6 +595,68 @@ class Qwen2_5OmniForConditionalGeneration(
return talker_hf_config.tts_text_start_token_id
return self.tts_text_spk_token_ids[voice_type]
def process_fn(self,
input_ids: torch.Tensor,
inputs_embeds: torch.Tensor,
**info_dict: object,
):
# Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
# Rules:
# - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
# - Decode segments consist of a single non-special token.
# - If additional_information is provided (can be a list split by request or a
# concatenated tensor plus a list of shapes), then for each request, reconstruct
# the thinker→talker input embeddings for the Prefill segments;
# - For Decode segments, if per-request auxiliary decode embeddings are provided (optional),
# add them; otherwise, keep the original embedding.
# Ensure we have base embeddings when only ids are provided
if inputs_embeds is None and input_ids is not None:
inputs_embeds = self.talker.get_input_embeddings(input_ids)
span_len = input_ids.shape[0]
if span_len > 1:
# prefill
return self.thinker_to_talker_process(input_ids, inputs_embeds, **info_dict)
else:
# decode
return self.thinker_to_talker_decode_one_step(input_ids, inputs_embeds, **info_dict)
def thinker_to_talker_process(self,
input_ids: torch.Tensor,
inputs_embeds: torch.Tensor,
**info_dict: object,
):
update_dict = {}
prompt_embeds = info_dict.get("prompt_embeds") # Tensor [P,H]
thinker_result = info_dict.get("thinker_result") # Tensor [K,H]
prompt_token_ids = info_dict.get("prompt_token_ids") # list[int]
thinker_output_token_ids = info_dict.get("thinker_output_token_ids") # list[int]
if not isinstance(prompt_embeds, torch.Tensor):
prompt_embeds = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
if not isinstance(thinker_result, torch.Tensor):
thinker_result = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
if not isinstance(prompt_token_ids, (list, torch.Tensor)):
prompt_token_ids = []
if not isinstance(thinker_output_token_ids, (list, torch.Tensor)):
thinker_output_token_ids = []
req_input_ids, req_embeds = self._thinker_to_talker_prefill(
voice_type=info_dict.get("voice_type"),
output_prompt_embeds=thinker_result.to(inputs_embeds.dtype).to(self._module_device(self.model)),
output_token_ids=thinker_output_token_ids,
thinker_prompt_embeds=prompt_embeds.to(inputs_embeds.dtype).to(self._module_device(self.model)),
prompt_token_ids=prompt_token_ids,
)
if thinker_result.ndim == 2 and thinker_result.shape[0] > 0:
update_dict["thinker_reply_part"] = thinker_result[1:].detach().to("cpu").contiguous()
return req_input_ids, req_embeds, update_dict
def _thinker_to_talker_prefill(
self,
voice_type: str,
@@ -778,6 +702,38 @@ class Qwen2_5OmniForConditionalGeneration(
)
return prompt_token_ids_processed, prompt_embeds
def thinker_to_talker_decode_one_step(self,
input_ids,
inputs_embeds,
**info_dict
):
update_dict = {}
# choose step vector in priority order
step_vec = None
q = info_dict.get("thinker_reply_part", None)
if isinstance(q, torch.Tensor) and q.numel() > 0:
step_vec = q[0:1]
new_q = q[1:].detach().to("cpu").contiguous()
update_dict["thinker_reply_part"] = new_q
else:
# B) per-request provided decode vector (optional)
dv = info_dict.get("decode_output_prompt_embeds") if isinstance(info_dict, dict) else None
if isinstance(dv, torch.Tensor) and dv.numel() > 0:
step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
# C) fallback shared pool
step_vec = self.thinker_reply_part[0:1]
self.thinker_reply_part = self.thinker_reply_part[1:]
if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
one_id = input_ids[0:1]
_, one_embed = self._thinker_to_talker_decode_one_step(
output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
output_token_ids=one_id,
)
inputs_embeds[0] = one_embed[0]
return input_ids[0:1], inputs_embeds[0:1], update_dict
def _thinker_to_talker_decode_one_step(
self,
output_prompt_embeds,
vllm_omni/worker/gpu_ar_model_runner.py
+32 -82
View File
@@ -79,7 +79,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
# _prepare_inputs may reorder the batch, so we must gather multi
# modal outputs after that to ensure the correct order
per_req_additional_information: dict[str, dict] | None = None
if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
# Run the multimodal encoder if any.
self._execute_mm_encoder(scheduler_output)
@@ -100,14 +99,12 @@ class GPUARModelRunner(OmniGPUModelRunner):
if hasattr(self, "_forward_additional_information"):
self._forward_additional_information = None
# New: per-request additional information for this step
per_req_additional_information = {}
# Overlay custom prompt_embeds per request for the prompt portion;
# collect additional_information (tensor/list) for prefill portion only
for req_index, req_id in enumerate(self.input_batch.req_ids):
req_state = self.requests[req_id]
pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
addi_cpu = getattr(req_state, "additional_information_cpu", None)
num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
prompt_len = len(req_state.prompt_token_ids)
prompt_remaining = max(0, prompt_len - num_computed_tokens)
@@ -121,31 +118,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
)
start_offset = int(self.query_start_loc.cpu[req_index])
self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
# Build per-request additional information (no cross-request concat)
if addi_cpu is not None and isinstance(addi_cpu, dict):
req_info: dict[str, object] = {}
for k, v in addi_cpu.items():
if isinstance(v, torch.Tensor):
# For prefill tokens, pass only the scheduled slice;
# for decode or no scheduled tokens, pass whole tensor
if overlay_len > 0:
try:
seg = (
v[num_computed_tokens : num_computed_tokens + overlay_len]
.detach()
.to("cpu")
.contiguous()
)
except Exception:
seg = v.detach().to("cpu").contiguous()
req_info[k] = seg
else:
req_info[k] = v.detach().to("cpu").contiguous()
elif isinstance(v, list):
req_info[k] = v
else:
req_info[k] = v
per_req_additional_information[req_id] = req_info
input_ids = self.input_ids.gpu[:num_input_tokens]
inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
@@ -200,6 +172,34 @@ class GPUARModelRunner(OmniGPUModelRunner):
encoder_inputs = self._extract_encoder_inputs(scheduler_output)
model_kwargs.update(encoder_inputs)
if self.model.has_process:
# Overlay custom prompt_embeds per request for the prompt portion;
# collect additional_information (tensor/list) for prefill portion only
for req_index, req_id in enumerate(self.input_batch.req_ids):
req_state = self.requests.get(req_id)
req_infos = (getattr(req_state, "additional_information_cpu", None)
if req_state is not None else None)
start_offset = int(self.query_start_loc.cpu[req_index])
sched_tokens = int(num_scheduled_tokens_np[req_index])
s, e = start_offset, start_offset + sched_tokens
span_len = int(e) - int(s)
# call the custom process function
req_input_ids, req_embeds, update_dict = self.model.process_fn(
input_ids=input_ids[s:e],
inputs_embeds=inputs_embeds[s:e],
**req_infos
)
# TODO(Peiqi): the merge stage could move out from the critical path
self._merge_additional_information_update(req_id, update_dict)
# update the inputs_embeds and input_ids
seg_len = min(span_len, req_embeds.shape[0])
inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
input_ids[s : s + seg_len] = req_input_ids
return (
num_scheduled_tokens,
num_input_tokens,
@@ -209,7 +209,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
positions,
intermediate_tensors,
model_kwargs,
per_req_additional_information,
)
@torch.inference_mode()
@@ -309,7 +308,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
positions,
intermediate_tensors,
model_kwargs,
per_req_additional_information,
) = self._preprocess(
scheduler_output,
num_scheduled_tokens_np,
@@ -344,29 +342,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
record_function_or_nullcontext("Forward"),
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
):
# TODO(Peiqi): delete this
model_kwargs_extra = {}
# Pass per-request additional information map for this step (no concat)
if per_req_additional_information:
model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
# Always pass per-request runtime additional_information (persisted in request state)
try:
per_req_runtime_info = []
for req_id in self.input_batch.req_ids:
req_state = self.requests.get(req_id)
info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
per_req_runtime_info.append(info if isinstance(info, dict) else {})
model_kwargs_extra["runtime_additional_information"] = per_req_runtime_info
model_kwargs_extra["request_ids"] = self.input_batch.req_ids
# Pass each request's token span within the flattened sequence for this step,
# enabling the model to map decode/prefill by request
req_token_spans = []
for req_index in range(len(self.input_batch.req_ids)):
start_offset = int(self.query_start_loc.cpu[req_index])
sched_tokens = int(num_scheduled_tokens_np[req_index])
req_token_spans.append((start_offset, start_offset + sched_tokens))
model_kwargs_extra["request_token_spans"] = req_token_spans
except Exception:
pass
model_output = self.model(
input_ids=input_ids,
positions=positions,
@@ -389,36 +367,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
aux_hidden_states = None
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(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.
try:
if isinstance(multimodal_outputs, dict) and (
"additional_information_update" in multimodal_outputs
or "additional_information_update_by_req_id" in multimodal_outputs
):
# Option A: list[dict] in batch order
updates_list = multimodal_outputs.get("additional_information_update")
if isinstance(updates_list, list):
for idx, upd in enumerate(updates_list):
if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
continue
req_id = self.input_batch.req_ids[idx]
self._merge_additional_information_update(req_id, upd)
# Option B: dict[str, dict] keyed by req_id
updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
if isinstance(updates_map, dict):
for req_id, upd in updates_map.items():
if not isinstance(upd, dict):
continue
if req_id not in self.requests:
continue
self._merge_additional_information_update(req_id, upd)
except Exception as e:
logger.error(
f"Error merging for requests:{self.input_batch.req_ids} additional \
information update: {e}, with the multimodal_outputs as {multimodal_outputs}"
)
# TODO(Peiqi): move the merge stage out here
if not self.broadcast_pp_output:
# Common case.
if not get_pp_group().is_last_rank: