diff --git a/vllm_omni/model_executor/custom_process_mixin.py b/vllm_omni/model_executor/custom_process_mixin.py
new file mode 100644
index 000000000..9b2942086
--- /dev/null
+++ b/vllm_omni/model_executor/custom_process_mixin.py
@@ -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.")
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index 99a80dc82..e8536a500 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -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,
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
index 1a6cbae98..2b1759f71 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
@@ -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),
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
index 48fde0564..d68d1e65f 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
@@ -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,
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 162c66dd4..77b32c681 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -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 =====
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index 60dd271ab..9b326a652 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -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)