[Core] customize process between stages, remove the need of "request_ids" in model forward, rebase to 0.12.0 (#458)
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
This commit is contained in:
Peiqi Yin authored and GitHub committed
1 parent
341457a121
commit
e9e6b3711c
6 files changed
+466
-606
No files matched your search
@@ -0,0 +1,44 @@
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class CustomProcessMixin:
|
||||
"""
|
||||
Mixin class for all stages in the Omni model.
|
||||
"""
|
||||
|
||||
def set_custom_preprocess(self, preprocess_fn: Callable) -> None:
|
||||
"""
|
||||
Set a preprocess function for the stage.
|
||||
Args:
|
||||
preprocess_fn: The preprocess function to register.
|
||||
"""
|
||||
self.preprocess = preprocess_fn
|
||||
|
||||
def set_custom_postprocess(self, postprocess_fn: Callable) -> None:
|
||||
"""
|
||||
Set a postprocess function for the stage.
|
||||
Args:
|
||||
postprocess_fn: The postprocess function to register.
|
||||
"""
|
||||
self.postprocess = postprocess_fn
|
||||
|
||||
def preprocess(
|
||||
self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **input_dict: object
|
||||
) -> tuple[torch.Tensor, torch.Tensor, dict]:
|
||||
"""
|
||||
Process the input_ids and input_embeds for the given input_dict.
|
||||
Returns the processed input_ids, input_embeds, and the input_dict.
|
||||
If the stage don't applicable, return the original input_ids, input_embeds, and an empty dict.
|
||||
"""
|
||||
raise NotImplementedError("Preprocess is not implemented for this stage.")
|
||||
|
||||
def postprocess(self, model_output, **info_dict: object):
|
||||
"""
|
||||
Postprocess the model output.
|
||||
Returns the postprocessed model output and the save dictionary.
|
||||
Args:
|
||||
model_output: The model output to postprocess.
|
||||
"""
|
||||
raise NotImplementedError("Postprocess is not implemented for this stage.")
|
||||
@@ -30,6 +30,7 @@ from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
|
||||
@@ -49,10 +50,11 @@ logger = init_logger(__name__)
|
||||
dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
|
||||
)
|
||||
class Qwen2_5OmniForConditionalGeneration(
|
||||
nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
|
||||
nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE, Qwen2_5OmniConditionalGenerationMixin, CustomProcessMixin
|
||||
):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self.has_preprocess = False
|
||||
self.have_multimodal_outputs = True
|
||||
config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
@@ -83,6 +85,9 @@ class Qwen2_5OmniForConditionalGeneration(
|
||||
self.token2wav = None
|
||||
|
||||
elif self.model_stage == "talker":
|
||||
# register the process function for the talker stage
|
||||
self.has_preprocess = True
|
||||
self.set_custom_preprocess(self.talker_preprocess)
|
||||
self.thinker = None
|
||||
# Initialize talker model wrapper (handles projection + LM)
|
||||
self.talker = init_vllm_registered_model(
|
||||
@@ -267,151 +272,14 @@ 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.
|
||||
|
||||
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,
|
||||
)
|
||||
additional_information = {}
|
||||
# mock data for profile
|
||||
if input_ids is None:
|
||||
input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
|
||||
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.embed_input_ids(input_ids)
|
||||
|
||||
# ------- Request-scoped additional information (no cross-request concat) -------
|
||||
request_ids: list[str] | None = kwargs.get("request_ids") # ordered
|
||||
request_token_spans: list[tuple[int, int]] | None = kwargs.get("request_token_spans")
|
||||
addi_by_req: dict | None = 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]
|
||||
# TODO(Peiqi): temporal hack here to support voice_type.
|
||||
if not hasattr(self, "voice_type"):
|
||||
self.voice_type = voice_type
|
||||
|
||||
with torch.inference_mode():
|
||||
talker_hidden = self.talker(
|
||||
@@ -419,10 +287,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,
|
||||
@@ -431,7 +295,7 @@ class Qwen2_5OmniForConditionalGeneration(
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=talker_hidden,
|
||||
multimodal_outputs=multimodal_outputs,
|
||||
multimodal_outputs=None,
|
||||
)
|
||||
|
||||
if self.model_stage == "code2wav":
|
||||
@@ -743,6 +607,74 @@ class Qwen2_5OmniForConditionalGeneration(
|
||||
return talker_hf_config.tts_text_start_token_id
|
||||
return self.tts_text_spk_token_ids[voice_type]
|
||||
|
||||
def talker_preprocess(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
input_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 input_embeds is None and input_ids is not None:
|
||||
input_embeds = self.talker.embed_input_ids(input_ids)
|
||||
|
||||
span_len = input_ids.shape[0]
|
||||
if span_len > 1:
|
||||
# prefill
|
||||
return self.thinker_to_talker_process(input_ids, input_embeds, **info_dict)
|
||||
else:
|
||||
# decode
|
||||
return self.thinker_to_talker_decode_one_step(input_ids, input_embeds, **info_dict)
|
||||
|
||||
def thinker_to_talker_process(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
input_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=input_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=input_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 = []
|
||||
|
||||
# TODO(Peiqi): add voice_type support
|
||||
req_input_ids, req_embeds = self._thinker_to_talker_prefill(
|
||||
voice_type=self.voice_type,
|
||||
output_prompt_embeds=thinker_result.to(input_embeds.dtype).to(self._module_device(self.model)),
|
||||
output_token_ids=thinker_output_token_ids,
|
||||
thinker_prompt_embeds=prompt_embeds.to(input_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,
|
||||
@@ -788,6 +720,38 @@ class Qwen2_5OmniForConditionalGeneration(
|
||||
)
|
||||
return prompt_token_ids_processed, prompt_embeds
|
||||
|
||||
def thinker_to_talker_decode_one_step(self, input_ids, input_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(input_embeds.dtype).to(self._module_device(self.model)),
|
||||
output_token_ids=one_id,
|
||||
)
|
||||
input_embeds[0] = one_embed[0]
|
||||
return input_ids[0:1], input_embeds[0:1], update_dict
|
||||
|
||||
def _thinker_to_talker_decode_one_step(
|
||||
self,
|
||||
output_prompt_embeds,
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
# Copyright 2025 The Qwen team.
|
||||
"""Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav)."""
|
||||
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterable
|
||||
from functools import cached_property
|
||||
|
||||
@@ -31,8 +30,9 @@ from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
|
||||
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape
|
||||
|
||||
# Special token IDs for Qwen3 Omni MoE
|
||||
# Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
|
||||
@@ -65,7 +65,7 @@ logger = init_logger(__name__)
|
||||
dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder,
|
||||
)
|
||||
class Qwen3OmniMoeForConditionalGeneration(
|
||||
nn.Module, SupportsMultiModal, SupportsPP, Qwen3OmniMoeConditionalGenerationMixin
|
||||
nn.Module, SupportsMultiModal, SupportsPP, Qwen3OmniMoeConditionalGenerationMixin, CustomProcessMixin
|
||||
):
|
||||
"""
|
||||
Unified Qwen3 Omni MoE model combining thinker, talker, and code2wav.
|
||||
@@ -82,6 +82,8 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self.have_multimodal_outputs = True
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
config: Qwen3OmniMoeConfig = vllm_config.model_config.hf_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
|
||||
@@ -121,6 +123,10 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
self.talker = None
|
||||
self.code2wav = None
|
||||
elif self.model_stage == "talker":
|
||||
self.has_preprocess = True
|
||||
self.has_postprocess = True
|
||||
self.set_custom_preprocess(self.talker_preprocess)
|
||||
self.set_custom_postprocess(self.talker_postprocess)
|
||||
self.thinker = None
|
||||
# Initialize talker model (text embeddings → codec codes)
|
||||
# Create a new vllm_config with talker_config as the hf_config
|
||||
@@ -359,19 +365,8 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
|
||||
# ========== Stage 2.1: Talker ==========
|
||||
elif 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.
|
||||
|
||||
if input_ids is None and additional_information is None:
|
||||
if input_ids is None:
|
||||
input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
|
||||
additional_information = {}
|
||||
self.code_predictor_hidden_pool = torch.zeros_like(inputs_embeds)
|
||||
is_profile = True
|
||||
else:
|
||||
@@ -381,274 +376,39 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
if inputs_embeds is None and input_ids is not None:
|
||||
inputs_embeds = self.talker.embed_input_ids(input_ids)
|
||||
|
||||
# ------- Request-scoped additional information (no cross-request concat) -------
|
||||
request_ids: list[str] | None = kwargs.get("request_ids") # ordered
|
||||
request_token_spans: list[tuple[int, int]] | None = kwargs.get("request_token_spans")
|
||||
addi_by_req: dict | None = 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., code_predictor_hidden_per_request)
|
||||
update_by_req_id: dict[str, dict] = {}
|
||||
text_step = torch.zeros(
|
||||
1,
|
||||
self.talker_config.text_config.hidden_size,
|
||||
device=self._module_device(self.talker),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
last_talker_hidden = torch.zeros(
|
||||
1,
|
||||
1,
|
||||
self.talker_config.text_config.hidden_size,
|
||||
device=self._module_device(self.talker),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
batched_talker_inputs = defaultdict(dict)
|
||||
if is_profile:
|
||||
batched_talker_inputs["0"] = {
|
||||
"input_ids": input_ids,
|
||||
"positions": positions,
|
||||
"inputs_embeds": inputs_embeds.clone(),
|
||||
"text_step": text_step,
|
||||
"last_talker_hidden": last_talker_hidden,
|
||||
}
|
||||
|
||||
# ------- 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
|
||||
batched_talker_inputs[rid] = {
|
||||
"input_ids": input_ids[s:e],
|
||||
"positions": positions[s:e] if len(positions.shape) == 1 else positions[0, s:e],
|
||||
"inputs_embeds": inputs_embeds[s:e].clone(),
|
||||
"text_step": text_step,
|
||||
"last_talker_hidden": last_talker_hidden,
|
||||
}
|
||||
info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
|
||||
# Read thinker outputs for prefill
|
||||
thinker_sequence_embeds = info.get("thinker_embeddings").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
) # Tensor [P,H]
|
||||
thinker_hidden_states = info.get("thinker_hidden_states").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
) # Tensor [K,H]
|
||||
thinker_sequences = (
|
||||
info.get("thinker_sequences")
|
||||
if info.get("thinker_sequences") is None
|
||||
else torch.as_tensor(info.get("thinker_sequences"), device=self._module_device(self.talker))
|
||||
)
|
||||
thinker_chatml_ids = (
|
||||
info.get("thinker_input_ids")
|
||||
if info.get("thinker_input_ids") is None
|
||||
else torch.as_tensor(info.get("thinker_input_ids"), device=self._module_device(self.talker))
|
||||
)
|
||||
# Optional thinker-side TTS token embeddings (to be projected by talker)
|
||||
tts_bos_thinker = info.get("tts_bos_embed").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
)
|
||||
tts_eos_thinker = info.get("tts_eos_embed").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
)
|
||||
tts_pad_thinker = info.get("tts_pad_embed").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
if thinker_sequence_embeds is None or thinker_hidden_states is None:
|
||||
raise ValueError(
|
||||
f"additional_information_by_req_id[{rid}] must include "
|
||||
f"'thinker_embeddings' and 'thinker_hidden_states' for talker prefill."
|
||||
)
|
||||
|
||||
# Normalize to tensors
|
||||
if not isinstance(thinker_sequence_embeds, torch.Tensor):
|
||||
thinker_sequence_embeds = torch.as_tensor(
|
||||
thinker_sequence_embeds, device=self._module_device(self.talker)
|
||||
)
|
||||
if not isinstance(thinker_hidden_states, torch.Tensor):
|
||||
thinker_hidden_states = torch.as_tensor(
|
||||
thinker_hidden_states, device=self._module_device(self.talker)
|
||||
)
|
||||
|
||||
if isinstance(thinker_chatml_ids, torch.Tensor) or isinstance(thinker_chatml_ids, list):
|
||||
ids_chatml = (
|
||||
thinker_chatml_ids
|
||||
if isinstance(thinker_chatml_ids, torch.Tensor)
|
||||
else torch.as_tensor(thinker_chatml_ids, device=self._module_device(self.talker))
|
||||
)
|
||||
if ids_chatml.ndim == 1:
|
||||
ids_chatml = ids_chatml.unsqueeze(0)
|
||||
else:
|
||||
# Fallback: create dummy ids if not provided
|
||||
ids_chatml = torch.zeros(
|
||||
(1, thinker_sequence_embeds.shape[1]),
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
)
|
||||
thinker_sequences = ids_chatml
|
||||
|
||||
speaker_id = self._get_text_spk_token_id(voice_type)
|
||||
req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill(
|
||||
thinker_embed=thinker_sequence_embeds.to(self._module_device(self.talker)),
|
||||
thinker_hidden=thinker_hidden_states.to(self._module_device(self.talker)),
|
||||
multimodal_mask=None,
|
||||
input_ids=ids_chatml.to(self._module_device(self.talker)),
|
||||
thinker_result_ids=thinker_sequences.to(self._module_device(self.talker)),
|
||||
speaker_id=speaker_id,
|
||||
tts_bos_thinker=tts_bos_thinker,
|
||||
tts_eos_thinker=tts_eos_thinker,
|
||||
tts_pad_thinker=tts_pad_thinker,
|
||||
)
|
||||
# req_embeds: [S,D], req_input_ids: [S] → map onto the current span
|
||||
seg_len = min(span_len, req_embeds.shape[0])
|
||||
# Force-align dtype/device/memory layout to avoid mismatches from precision or layout differences
|
||||
inputs_embeds[s : s + seg_len] = (
|
||||
req_embeds[:seg_len].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype).contiguous()
|
||||
)
|
||||
|
||||
if isinstance(req_input_ids, torch.Tensor) and req_input_ids.shape[0] >= seg_len:
|
||||
input_ids[s : s + seg_len] = req_input_ids[:seg_len].to(input_ids.device)
|
||||
# Queue tailing_text_hidden for decode (drop first for next steps),
|
||||
# similar to Qwen2.5's thinker_reply_part_per_request
|
||||
batched_talker_inputs[rid]["inputs_embeds"] = inputs_embeds[s : s + seg_len].clone()
|
||||
try:
|
||||
if isinstance(trailing_text_hidden, torch.Tensor) and trailing_text_hidden.numel() > 0:
|
||||
if trailing_text_hidden.ndim == 2:
|
||||
rem_tail = trailing_text_hidden
|
||||
elif trailing_text_hidden.ndim == 1:
|
||||
rem_tail = torch.zeros(
|
||||
0,
|
||||
trailing_text_hidden.shape[0],
|
||||
dtype=trailing_text_hidden.dtype,
|
||||
device=trailing_text_hidden.device,
|
||||
)
|
||||
else:
|
||||
# compatible with old shape [1,S,D]
|
||||
rem_tail = trailing_text_hidden.squeeze(0)
|
||||
if rem_tail.shape[0] > 0:
|
||||
update_by_req_id.setdefault(rid, {})["tailing_text_hidden_per_request"] = (
|
||||
rem_tail.detach().to("cpu").contiguous()
|
||||
)
|
||||
# Also persist projected tts_pad for decode fallback if needed
|
||||
if isinstance(tts_pad_thinker, torch.Tensor):
|
||||
pad_in = tts_pad_thinker
|
||||
if pad_in.ndim == 2:
|
||||
pad_in = pad_in.unsqueeze(0)
|
||||
if pad_in.ndim == 1:
|
||||
pad_in = pad_in.view(1, 1, -1)
|
||||
pad_proj = self.talker.text_projection(pad_in.to(self._module_device(self.talker)))
|
||||
update_by_req_id.setdefault(rid, {})["tts_pad_embed_projected"] = (
|
||||
pad_proj.detach().to("cpu").contiguous()
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ------- 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
|
||||
batched_talker_inputs[rid] = {
|
||||
"input_ids": input_ids[s:e],
|
||||
"positions": positions[s:e] if len(positions.shape) == 1 else positions[0, s:e],
|
||||
"inputs_embeds": inputs_embeds[s:e].clone(),
|
||||
"text_step": text_step,
|
||||
"last_talker_hidden": last_talker_hidden,
|
||||
}
|
||||
try:
|
||||
q_tail = None
|
||||
if isinstance(rid, str):
|
||||
q_tail = runtime_addi_by_req.get(rid, {}).get("tailing_text_hidden_per_request")
|
||||
if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:
|
||||
use_vec = q_tail[0:1, :]
|
||||
new_q_tail = (
|
||||
q_tail[1:, :].detach().to("cpu").contiguous()
|
||||
if q_tail.shape[1] > 1
|
||||
else self.tts_pad_embed.to(inputs_embeds.device, dtype=inputs_embeds.dtype)
|
||||
)
|
||||
text_step = use_vec.to(inputs_embeds.device, dtype=inputs_embeds.dtype)
|
||||
update_by_req_id.setdefault(rid, {})["tailing_text_hidden_per_request"] = new_q_tail
|
||||
else:
|
||||
text_step = self.tts_pad_embed.to(inputs_embeds.device, dtype=inputs_embeds.dtype)
|
||||
|
||||
if isinstance(rid, str):
|
||||
last_talker_hidden = (
|
||||
runtime_addi_by_req.get(rid, {})
|
||||
.get("last_talker_hidden")
|
||||
.to(inputs_embeds.device, dtype=inputs_embeds.dtype)
|
||||
)
|
||||
last_talker_hidden = last_talker_hidden.reshape(
|
||||
*last_talker_hidden.shape[-2:]
|
||||
) # [1, hidden_size]
|
||||
batched_talker_inputs[rid]["last_talker_hidden"] = last_talker_hidden
|
||||
batched_talker_inputs[rid]["text_step"] = text_step
|
||||
except Exception as e:
|
||||
logger.error(f"Error in decode: {e}")
|
||||
|
||||
if not is_profile and input_ids.shape[0] > 1:
|
||||
self.input_ids_list = input_ids.tolist()
|
||||
|
||||
last_talker_hidden = None
|
||||
# TODO(Peiqi): temporal hack here to support voice_type.
|
||||
if not hasattr(self, "voice_type"):
|
||||
self.voice_type = voice_type
|
||||
|
||||
# Run talker forward
|
||||
with torch.inference_mode():
|
||||
talker_hidden, batched_code_predictor_codes = self.talker(
|
||||
batched_talker_inputs=batched_talker_inputs,
|
||||
talker_hidden = self.talker.forward(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
inputs_embeds=inputs_embeds,
|
||||
)
|
||||
if not is_profile:
|
||||
for rid in request_ids:
|
||||
idx_req = request_ids.index(rid)
|
||||
s, e = request_token_spans[idx_req]
|
||||
span_len = int(e) - int(s)
|
||||
update_by_req_id.setdefault(rid, {})["last_talker_hidden"] = (
|
||||
talker_hidden[e - 1 : e, :].detach().to("cpu").contiguous()
|
||||
)
|
||||
|
||||
# merge the code_predictor_codes from the info_dict list into a single tensor
|
||||
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 not is_profile and multimodal_outputs is not None and batched_code_predictor_codes is not None:
|
||||
multimodal_outputs.update({"code_predictor_codes": batched_code_predictor_codes})
|
||||
elif not is_profile and multimodal_outputs is None and batched_code_predictor_codes is not None:
|
||||
multimodal_outputs = {"code_predictor_codes": batched_code_predictor_codes}
|
||||
if not is_profile:
|
||||
# Here is the only place to use runtime_additional_information. After MTP in the
|
||||
# preprocess function, the code_predictor_codes are stored in the info_dict list.
|
||||
# We need to merge the tensors from different requests into a single tensor.
|
||||
# In the future, we may allow user to custom an aggregated function.
|
||||
info_dicts = kwargs.get("runtime_additional_information")
|
||||
code_predictor_codes = [info.get("code_predictor_codes") for info in info_dicts]
|
||||
multimodal_outputs = {"code_predictor_codes": torch.cat(code_predictor_codes, dim=0)}
|
||||
|
||||
next_token_id = None
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=talker_hidden,
|
||||
multimodal_outputs=multimodal_outputs,
|
||||
next_token_id=next_token_id,
|
||||
)
|
||||
return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=multimodal_outputs)
|
||||
|
||||
# ========== Stage 3: Code2Wav ==========
|
||||
elif self.model_stage == "code2wav":
|
||||
# Extract codec codes from input
|
||||
codes = []
|
||||
if input_ids is not None:
|
||||
# addi_by_req = kwargs.get("additional_information_by_req_id")
|
||||
codes.append(input_ids.reshape(1, 16, -1))
|
||||
|
||||
else:
|
||||
@@ -801,6 +561,183 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
return self.tts_text_spk_token_ids[self.default_tts_text_spk_type]
|
||||
return self.tts_text_spk_token_ids[voice_type]
|
||||
|
||||
def talker_postprocess(self, hidden_states: torch.Tensor, **info_dict: object):
|
||||
"""
|
||||
Postprocess the talker hidden states.
|
||||
"""
|
||||
update_dict = {}
|
||||
update_dict["last_talker_hidden"] = hidden_states[-1, :].detach().to("cpu").contiguous()
|
||||
return update_dict
|
||||
|
||||
def talker_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
|
||||
"""
|
||||
Preprocess talker embeds. Noted that we set the MTP here.
|
||||
"""
|
||||
# Ensure we have base embeddings when only ids are provided
|
||||
if input_embeds is None and input_ids is not None:
|
||||
input_embeds = self.talker.embed_input_ids(input_ids)
|
||||
|
||||
text_step = torch.zeros(
|
||||
1,
|
||||
self.talker_config.text_config.hidden_size,
|
||||
device=self._module_device(self.talker),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
last_talker_hidden = torch.zeros(
|
||||
1,
|
||||
1,
|
||||
self.talker_config.text_config.hidden_size,
|
||||
device=self._module_device(self.talker),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
span_len = input_ids.shape[0]
|
||||
if span_len > 1:
|
||||
# prefill
|
||||
input_ids, input_embeds, update_dict = self.talker_preprocess_prefill(input_ids, input_embeds, **info_dict)
|
||||
else:
|
||||
last_talker_hidden, text_step, update_dict = self.talker_preprocess_decode(
|
||||
input_ids, input_embeds, **info_dict
|
||||
)
|
||||
|
||||
# execute talker MTP
|
||||
input_embeds, code_predictor_codes = self.talker_mtp(input_ids, input_embeds, last_talker_hidden, text_step)
|
||||
update_dict["code_predictor_codes"] = code_predictor_codes
|
||||
|
||||
return input_ids, input_embeds, update_dict
|
||||
|
||||
def talker_mtp(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
input_embeds: torch.Tensor,
|
||||
last_talker_hidden: torch.Tensor,
|
||||
text_step: torch.Tensor,
|
||||
):
|
||||
# TODO(Peiqi): not support intermediate_tensors now
|
||||
input_ids = safe_tensor_reshape(input_ids, (1, -1))
|
||||
inputs_embeds = safe_tensor_reshape(input_embeds, (-1, self.talker_config.text_config.hidden_size))
|
||||
text_step = safe_tensor_reshape(text_step, (1, -1))
|
||||
last_talker_hidden = safe_tensor_reshape(last_talker_hidden, (1, 1, self.talker_config.text_config.hidden_size))
|
||||
# for profiling
|
||||
if inputs_embeds.shape[-1] == 2048:
|
||||
inputs_embeds = self.text_projection(inputs_embeds)
|
||||
if inputs_embeds.shape[0] == 1:
|
||||
code_predictor_codes, summed_embeddings = self.talker.code_predictor_forward(
|
||||
input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden
|
||||
)
|
||||
inputs_embeds = summed_embeddings.clone()
|
||||
else:
|
||||
code_predictor_codes = torch.zeros((0, self.talker.num_code_groups), dtype=torch.long)
|
||||
inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size)
|
||||
return inputs_embeds, code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous()
|
||||
|
||||
def talker_preprocess_prefill(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
|
||||
# Containers to return per-request updates (e.g., code_predictor_hidden_per_request)
|
||||
update_dict: dict[str, dict] = {}
|
||||
# TODO(Peiqi): add voice_type support
|
||||
voice_type = self.voice_type
|
||||
|
||||
# Read thinker outputs for prefill
|
||||
thinker_sequence_embeds = info_dict.get("thinker_embeddings").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
) # Tensor [P,H]
|
||||
thinker_hidden_states = info_dict.get("thinker_hidden_states").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
) # Tensor [K,H]
|
||||
thinker_sequences = (
|
||||
info_dict.get("thinker_sequences")
|
||||
if info_dict.get("thinker_sequences") is None
|
||||
else torch.as_tensor(info_dict.get("thinker_sequences"), device=self._module_device(self.talker))
|
||||
)
|
||||
thinker_chatml_ids = (
|
||||
info_dict.get("thinker_input_ids")
|
||||
if info_dict.get("thinker_input_ids") is None
|
||||
else torch.as_tensor(info_dict.get("thinker_input_ids"), device=self._module_device(self.talker))
|
||||
)
|
||||
# Optional thinker-side TTS token embeddings (to be projected by talker)
|
||||
tts_bos_thinker = info_dict.get("tts_bos_embed").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
)
|
||||
tts_eos_thinker = info_dict.get("tts_eos_embed").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
)
|
||||
tts_pad_thinker = info_dict.get("tts_pad_embed").to(
|
||||
device=self._module_device(self.talker), dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
if thinker_sequence_embeds is None or thinker_hidden_states is None:
|
||||
raise ValueError(
|
||||
"additional_information_by_req_id must include "
|
||||
"'thinker_embeddings' and 'thinker_hidden_states' for talker prefill."
|
||||
)
|
||||
|
||||
# Normalize to tensors
|
||||
if not isinstance(thinker_sequence_embeds, torch.Tensor):
|
||||
thinker_sequence_embeds = torch.as_tensor(thinker_sequence_embeds, device=self._module_device(self.talker))
|
||||
if not isinstance(thinker_hidden_states, torch.Tensor):
|
||||
thinker_hidden_states = torch.as_tensor(thinker_hidden_states, device=self._module_device(self.talker))
|
||||
|
||||
if isinstance(thinker_chatml_ids, torch.Tensor) or isinstance(thinker_chatml_ids, list):
|
||||
ids_chatml = (
|
||||
thinker_chatml_ids
|
||||
if isinstance(thinker_chatml_ids, torch.Tensor)
|
||||
else torch.as_tensor(thinker_chatml_ids, device=self._module_device(self.talker))
|
||||
)
|
||||
if ids_chatml.ndim == 1:
|
||||
ids_chatml = ids_chatml.unsqueeze(0)
|
||||
else:
|
||||
# Fallback: create dummy ids if not provided
|
||||
ids_chatml = torch.zeros(
|
||||
(1, thinker_sequence_embeds.shape[1]),
|
||||
dtype=torch.long,
|
||||
device=self._module_device(self.talker),
|
||||
)
|
||||
thinker_sequences = ids_chatml
|
||||
|
||||
speaker_id = self._get_text_spk_token_id(voice_type)
|
||||
req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill(
|
||||
thinker_embed=thinker_sequence_embeds.to(self._module_device(self.talker)),
|
||||
thinker_hidden=thinker_hidden_states.to(self._module_device(self.talker)),
|
||||
multimodal_mask=None,
|
||||
input_ids=ids_chatml.to(self._module_device(self.talker)),
|
||||
thinker_result_ids=thinker_sequences.to(self._module_device(self.talker)),
|
||||
speaker_id=speaker_id,
|
||||
tts_bos_thinker=tts_bos_thinker,
|
||||
tts_eos_thinker=tts_eos_thinker,
|
||||
tts_pad_thinker=tts_pad_thinker,
|
||||
)
|
||||
|
||||
# Queue tailing_text_hidden for decode (drop first for next steps),
|
||||
try:
|
||||
if isinstance(trailing_text_hidden, torch.Tensor) and trailing_text_hidden.numel() > 0:
|
||||
if trailing_text_hidden.ndim == 2:
|
||||
rem_tail = trailing_text_hidden
|
||||
elif trailing_text_hidden.ndim == 1:
|
||||
rem_tail = torch.zeros(
|
||||
0,
|
||||
trailing_text_hidden.shape[0],
|
||||
dtype=trailing_text_hidden.dtype,
|
||||
device=trailing_text_hidden.device,
|
||||
)
|
||||
else:
|
||||
# compatible with old shape [1,S,D]
|
||||
rem_tail = trailing_text_hidden.squeeze(0)
|
||||
if rem_tail.shape[0] > 0:
|
||||
update_dict["tailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous()
|
||||
# Also persist projected tts_pad for decode fallback if needed
|
||||
if isinstance(tts_pad_thinker, torch.Tensor):
|
||||
pad_in = tts_pad_thinker
|
||||
if pad_in.ndim == 2:
|
||||
pad_in = pad_in.unsqueeze(0)
|
||||
if pad_in.ndim == 1:
|
||||
pad_in = pad_in.view(1, 1, -1)
|
||||
pad_proj = self.talker.text_projection(pad_in.to(self._module_device(self.talker)))
|
||||
update_dict["tts_pad_embed_projected"] = pad_proj.detach().to("cpu").contiguous()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return req_input_ids, req_embeds, update_dict
|
||||
|
||||
def _thinker_to_talker_prefill(
|
||||
self,
|
||||
thinker_embed: torch.Tensor,
|
||||
@@ -905,6 +842,30 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
|
||||
return talker_input_id, talker_input_embed, trailing_text_hidden_all
|
||||
|
||||
def talker_preprocess_decode(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
|
||||
update_dict: dict[str, dict] = {}
|
||||
|
||||
try:
|
||||
q_tail = info_dict.get("tailing_text_hidden")
|
||||
if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:
|
||||
use_vec = q_tail[0:1, :]
|
||||
new_q_tail = (
|
||||
q_tail[1:, :].detach().to("cpu").contiguous()
|
||||
if q_tail.shape[1] > 1
|
||||
else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
|
||||
)
|
||||
text_step = use_vec.to(input_embeds.device, dtype=input_embeds.dtype)
|
||||
update_dict["tailing_text_hidden"] = new_q_tail
|
||||
else:
|
||||
text_step = self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
|
||||
|
||||
last_talker_hidden = info_dict.get("last_talker_hidden").to(input_embeds.device, dtype=input_embeds.dtype)
|
||||
last_talker_hidden = last_talker_hidden.reshape(*last_talker_hidden.shape[-2:]) # [1, hidden_size]
|
||||
except Exception as e:
|
||||
logger.error(f"Error in decode: {e}")
|
||||
|
||||
return last_talker_hidden, text_step, update_dict
|
||||
|
||||
def _get_talker_user_parts(self, im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed):
|
||||
user_talker_part = torch.empty(
|
||||
(segment_end_index - im_start_index, self.config.talker_config.text_config.hidden_size),
|
||||
|
||||
@@ -48,7 +48,6 @@ try:
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
flash_attn = None
|
||||
|
||||
from vllm_omni.model_executor.models.utils import safe_tensor_reshape
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -368,75 +367,22 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batched_talker_inputs: dict,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
inputs_embeds: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
"""Forward pass through the talker model.
|
||||
|
||||
For inference, the talker receives inputs_embeds that should already be
|
||||
projected to talker's hidden dimension. If receiving raw thinker outputs,
|
||||
use project_thinker_outputs() first.
|
||||
"""
|
||||
# If intermediate_tensors is provided (pipeline parallel),
|
||||
# inputs_embeds should be None
|
||||
batched_input_ids = []
|
||||
batched_positions = []
|
||||
batched_input_embeds = []
|
||||
batched_code_predictor_codes = []
|
||||
if intermediate_tensors is not None:
|
||||
inputs_embeds = None
|
||||
for rid in batched_talker_inputs:
|
||||
input_ids = safe_tensor_reshape(batched_talker_inputs[rid]["input_ids"], (1, -1))
|
||||
positions = batched_talker_inputs[rid]["positions"]
|
||||
inputs_embeds = safe_tensor_reshape(
|
||||
batched_talker_inputs[rid]["inputs_embeds"], (-1, self.config.text_config.hidden_size)
|
||||
)
|
||||
text_step = safe_tensor_reshape(batched_talker_inputs[rid]["text_step"], (1, -1))
|
||||
last_talker_hidden = safe_tensor_reshape(
|
||||
batched_talker_inputs[rid]["last_talker_hidden"], (1, 1, self.config.text_config.hidden_size)
|
||||
)
|
||||
# for profiling
|
||||
if inputs_embeds.shape[-1] == 2048:
|
||||
inputs_embeds = self.text_projection(inputs_embeds)
|
||||
if inputs_embeds.shape[0] == 1:
|
||||
code_predictor_codes, summed_embeddings = self.code_predictor_forward(
|
||||
input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden
|
||||
)
|
||||
inputs_embeds = summed_embeddings.clone()
|
||||
else:
|
||||
code_predictor_codes = torch.zeros((0, self.num_code_groups), dtype=torch.long)
|
||||
batched_input_ids.append(input_ids.reshape(-1))
|
||||
batched_positions.append(positions)
|
||||
batched_input_embeds.append((inputs_embeds + text_step).reshape(-1, self.config.text_config.hidden_size))
|
||||
batched_code_predictor_codes.append(code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous())
|
||||
try:
|
||||
talker_input_ids = torch.cat(batched_input_ids, dim=0)
|
||||
talker_positions = torch.cat(batched_positions, dim=0)
|
||||
talker_input_embeds = torch.cat(batched_input_embeds, dim=0)
|
||||
except Exception as e:
|
||||
print(f"Error in talker_input_embeds: {e}")
|
||||
print(f"talker_input_embeds shape: {[tensor.shape for tensor in batched_input_embeds]}")
|
||||
print(f"talker_input_ids shape: {[tensor.shape for tensor in batched_input_ids]}")
|
||||
print(f"talker_positions shape: {[tensor.shape for tensor in batched_positions]}")
|
||||
raise e
|
||||
"""Forward pass through the talker model."""
|
||||
talker_hidden_states, _ = self.language_model.model(
|
||||
talker_input_ids,
|
||||
talker_positions,
|
||||
input_ids,
|
||||
positions,
|
||||
intermediate_tensors,
|
||||
inputs_embeds=talker_input_embeds,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Pass talker hidden states to code predictor
|
||||
# Returns: (residual_codes, summed_embeddings=sum of all layer embeddings at each position)
|
||||
|
||||
# Return both talker hidden states and code predictor results
|
||||
# code_predictor_codes: [batch, num_code_groups, seq_len]
|
||||
# summed_embeddings: [batch, seq_len, hidden_size]
|
||||
# - Sum of all layer embeddings at each position (like Transformers)
|
||||
|
||||
return talker_hidden_states, batched_code_predictor_codes
|
||||
return talker_hidden_states
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
|
||||
@@ -374,8 +374,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
with record_function_or_nullcontext("gpu_model_runner: eplb"):
|
||||
self.eplb_step()
|
||||
|
||||
self._process_additional_information_updates(multimodal_outputs)
|
||||
|
||||
hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
|
||||
num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
|
||||
if num_scheduled_tokens_np is None:
|
||||
@@ -385,6 +383,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
dtype=np.int32,
|
||||
)
|
||||
|
||||
self._process_additional_information_updates(hidden_states, multimodal_outputs, num_scheduled_tokens_np)
|
||||
|
||||
pooler_output: list[dict[str, object]] = []
|
||||
for rid in req_ids_output_copy:
|
||||
idx = req_id_to_index_output_copy[rid]
|
||||
@@ -449,24 +449,3 @@ class GPUARModelRunner(OmniGPUModelRunner):
|
||||
)
|
||||
|
||||
return async_output
|
||||
|
||||
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):
|
||||
merged[k] = [
|
||||
(item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v
|
||||
]
|
||||
else:
|
||||
merged[k] = v
|
||||
setattr(req_state, "additional_information_cpu", merged)
|
||||
|
||||
# ===== Helper functions extracted for clarity and reuse =====
|
||||
@@ -662,25 +662,13 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
req_token_spans.append((start_offset, start_offset + sched_tokens))
|
||||
return req_token_spans
|
||||
|
||||
def _build_model_kwargs_extra(
|
||||
self,
|
||||
per_req_additional_information: dict[str, dict] | None,
|
||||
num_scheduled_tokens_np,
|
||||
) -> dict:
|
||||
def _build_model_kwargs_extra(self) -> dict:
|
||||
"""Build extra keyword arguments passed to the model for this step, including:
|
||||
- additional_information_by_req_id: per-request additional information provided by preprocess for this step
|
||||
- runtime_additional_information: per-request additional information stored in request state
|
||||
- request_ids: the request id list in batch order
|
||||
- request_token_spans: token spans per request within the flattened sequence
|
||||
"""
|
||||
model_kwargs_extra: dict[str, object] = {}
|
||||
if per_req_additional_information:
|
||||
model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
|
||||
try:
|
||||
model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information()
|
||||
model_kwargs_extra["request_ids"] = list(self.input_batch.req_ids)
|
||||
if num_scheduled_tokens_np is not None:
|
||||
model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np)
|
||||
except Exception as e:
|
||||
logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}")
|
||||
import traceback
|
||||
@@ -688,49 +676,29 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
traceback.print_exc()
|
||||
return model_kwargs_extra
|
||||
|
||||
def _process_additional_information_updates(self, multimodal_outputs: object) -> None:
|
||||
def _process_additional_information_updates(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
multimodal_outputs: object,
|
||||
num_scheduled_tokens_np: np.ndarray,
|
||||
) -> None:
|
||||
"""Process model-provided per-request additional_information updates and merge into request state."""
|
||||
try:
|
||||
if multimodal_outputs is None:
|
||||
logger.debug("[OMNI] multimodal_outputs is None, no updates to process")
|
||||
return
|
||||
if not isinstance(multimodal_outputs, dict):
|
||||
logger.debug(f"[OMNI] multimodal_outputs is not dict: {type(multimodal_outputs)}")
|
||||
return
|
||||
if (
|
||||
"additional_information_update" not in multimodal_outputs
|
||||
and "additional_information_update_by_req_id" not in multimodal_outputs
|
||||
):
|
||||
logger.debug(f"[OMNI] multimodal_outputs has no update keys: {list(multimodal_outputs.keys())}")
|
||||
return
|
||||
updates_list = multimodal_outputs.get("additional_information_update")
|
||||
if isinstance(updates_list, list):
|
||||
logger.debug(f"[OMNI] Processing additional_information_update list with {len(updates_list)} items")
|
||||
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)
|
||||
updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
|
||||
if isinstance(updates_map, dict):
|
||||
logger.debug(
|
||||
f"[OMNI] Processing additional_information_update_by_req_id for {list(updates_map.keys())}"
|
||||
)
|
||||
for req_id, upd in updates_map.items():
|
||||
if not isinstance(upd, dict):
|
||||
continue
|
||||
if req_id not in self.requests:
|
||||
logger.warning(f"[OMNI] req_id {req_id} not found in requests")
|
||||
continue
|
||||
self._merge_additional_information_update(req_id, upd)
|
||||
|
||||
# execute the custom postprocess function
|
||||
# TODO(Peiqi): do we have a more elegant way to do this?
|
||||
if hasattr(self.model, "has_postprocess") and self.model.has_postprocess:
|
||||
for req_index, req_id in enumerate(self.input_batch.req_ids):
|
||||
req_state = self.requests.get(req_id)
|
||||
if req_state:
|
||||
info = getattr(req_state, "additional_information_cpu", {})
|
||||
if "thinker_reply_part_per_request" in info:
|
||||
q = info["thinker_reply_part_per_request"]
|
||||
if hasattr(q, "shape"):
|
||||
logger.debug(f"[OMNI] After update, req={req_id} queue shape: {q.shape}")
|
||||
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
|
||||
# only consider to store data into update dict.
|
||||
hidden_states_slice = hidden_states[s:e]
|
||||
update_dict = self.model.postprocess(hidden_states_slice, **req_infos)
|
||||
self._merge_additional_information_update(req_id, update_dict)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error merging for requests:{self.input_batch.req_ids} "
|
||||
@@ -741,46 +709,21 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
def _build_mm_inputs_and_overlays(
|
||||
self,
|
||||
scheduler_output: "SchedulerOutput",
|
||||
num_scheduled_tokens: int,
|
||||
num_scheduled_tokens_np: np.ndarray,
|
||||
) -> dict[str, dict]:
|
||||
"""Run MM encoder, fill scheduled inputs_embeds, and overlay prompt embeds;
|
||||
collect per-request additional_information for the prefill portion."""
|
||||
# Run the multimodal encoder if any.
|
||||
self._execute_mm_encoder(scheduler_output)
|
||||
mm_embeds = self._gather_mm_embeddings(scheduler_output)
|
||||
# Always use embeddings as inputs for MM models
|
||||
inputs_embeds_scheduled = self.model.get_input_embeddings(
|
||||
input_ids=self.input_ids.gpu[:num_scheduled_tokens],
|
||||
multimodal_embeddings=mm_embeds or None,
|
||||
)
|
||||
self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
|
||||
# Reset per-step additional information collector (deprecated concat path)
|
||||
if hasattr(self, "_forward_additional_information"):
|
||||
self._forward_additional_information = None
|
||||
# Overlay and collect per-request info
|
||||
return self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
|
||||
|
||||
def _collect_additional_information_for_prefill(
|
||||
self,
|
||||
num_scheduled_tokens_np: np.ndarray,
|
||||
) -> dict[str, dict]:
|
||||
"""Overlay per-request prompt_embeds for the prefill portion and collect
|
||||
additional_information slices for this step. Returns a map req_id -> dict."""
|
||||
per_req_additional_information: dict[str, dict] = {}
|
||||
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)
|
||||
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
||||
overlay_len = min(sched_tokens, prompt_remaining)
|
||||
if overlay_len <= 0 and not (isinstance(addi_cpu, dict) and addi_cpu):
|
||||
if overlay_len <= 0:
|
||||
continue
|
||||
if overlay_len > 0 and pe_cpu is not None:
|
||||
src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
|
||||
@@ -788,18 +731,6 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
)
|
||||
start_offset = int(self.query_start_loc.cpu[req_index])
|
||||
self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
|
||||
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):
|
||||
req_info[k] = v.detach().to("cpu").contiguous()
|
||||
elif isinstance(v, list):
|
||||
req_info[k] = v
|
||||
else:
|
||||
req_info[k] = v
|
||||
if req_info:
|
||||
per_req_additional_information[req_id] = req_info
|
||||
return per_req_additional_information
|
||||
|
||||
def _preprocess(
|
||||
self,
|
||||
@@ -909,17 +840,36 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
)
|
||||
self._omni_num_scheduled_tokens_np = num_scheduled_tokens_np
|
||||
|
||||
# Always collect additional_information, not only in prefill.
|
||||
# _collect_additional_information_for_prefill handles prompt_embeds overlay during prefill,
|
||||
# but even in decode we need per_req_additional_information so that
|
||||
# _build_model_kwargs_extra can pass runtime_additional_information correctly.
|
||||
per_req_additional_information: dict[str, dict] = {}
|
||||
# Note: only prefill need collect additional_information for now.
|
||||
# Decode don't need per_req_additional_information anymore.
|
||||
if inputs_embeds is not None:
|
||||
# Prefill: overlay prompt_embeds and collect additional_information
|
||||
per_req_additional_information = self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
|
||||
# Save for _model_forward regardless of prefill; runtime_additional_information
|
||||
# will be fetched from request state inside _build_model_kwargs_extra
|
||||
self._omni_per_req_additional_information = per_req_additional_information
|
||||
self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
|
||||
|
||||
if hasattr(self.model, "has_preprocess") and self.model.has_preprocess:
|
||||
# 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.preprocess(
|
||||
input_ids=input_ids[s:e], input_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 (
|
||||
input_ids,
|
||||
@@ -939,10 +889,7 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
**model_kwargs: dict[str, Any],
|
||||
):
|
||||
"""Inject omni-specific kwargs into forward and cache model output"""
|
||||
model_kwargs_extra = self._build_model_kwargs_extra(
|
||||
self._omni_per_req_additional_information,
|
||||
self._omni_num_scheduled_tokens_np,
|
||||
)
|
||||
model_kwargs_extra = self._build_model_kwargs_extra()
|
||||
|
||||
runtime_info = model_kwargs_extra.get("runtime_additional_information", [])
|
||||
if runtime_info:
|
||||
@@ -961,3 +908,22 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
# Cache model output so later sample_tokens can consume multimodal results.
|
||||
self._omni_last_model_output = model_output
|
||||
return model_output
|
||||
|
||||
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):
|
||||
merged[k] = [
|
||||
(item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v
|
||||
]
|
||||
else:
|
||||
merged[k] = v
|
||||
setattr(req_state, "additional_information_cpu", merged)
|
||||
Reference in new issue
Block a user