diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py
index bfd5324c5..dd9159769 100644
--- a/examples/offline_inference/qwen2_5_omni/end2end.py
+++ b/examples/offline_inference/qwen2_5_omni/end2end.py
@@ -322,7 +322,6 @@ def main(args):
     omni_llm = Omni(
         model=model_name,
         log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
         init_sleep_seconds=args.init_sleep_seconds,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py
index 7b13e10e8..4d1a49fb5 100644
--- a/examples/offline_inference/qwen3_omni/end2end.py
+++ b/examples/offline_inference/qwen3_omni/end2end.py
@@ -170,15 +170,9 @@ def main(args):
     else:
         query_result = query_func()
 
-    if not args.enable_stats:
-        log_file = None
-    else:
-        log_file = os.path.join(args.log_dir, f"omni_llm_pipeline_{args.query_type}")
-
     omni_llm = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
-        log_file=log_file,
         log_stats=args.enable_stats,
     )
 
diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
index 3033426aa..fec6169d8 100644
--- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
+++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml
@@ -81,7 +81,6 @@ stage_args:
       engine_output_type: audio  # Final output: audio waveform
       gpu_memory_utilization: 0.1
       distributed_executor_backend: "mp"
-      max_num_seqs: 1
       max_num_batched_tokens: 1000000
       hf_config_name: thinker_config
       load_format: dummy
diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py
index ce36b863d..bcc6d5bff 100644
--- a/tests/entrypoints/test_omni_llm.py
+++ b/tests/entrypoints/test_omni_llm.py
@@ -109,7 +109,6 @@ class _FakeStage:
         model: str,
         *,
         is_async: bool = False,
-        log_file=None,
         shm_threshold_bytes: int = 65536,
         ctx=None,
         batch_timeout: int = 10,
@@ -295,18 +294,6 @@ def _setup_ipc_mocks(monkeypatch):
 
 def _setup_log_mocks(monkeypatch):
     """Helper function to set up logging and stats mocks."""
-    # Mock init_stats_paths to return None files (no stats logging)
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.omni.init_stats_paths",
-        lambda enable_stats, log_file: (None, None),
-        raising=False,
-    )
-    # Mock configure_orchestrator_logger to do nothing
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.omni.configure_orchestrator_logger",
-        lambda logger, log_file: None,
-        raising=False,
-    )
     # Mock OrchestratorMetrics to be a simple class that doesn't require file operations
 
     class _FakeOrchestratorMetrics:
@@ -473,7 +460,6 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -485,13 +471,6 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    # Mock remove_old_logs to avoid bug (called before stage_list is created)
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
     # Mock load_stage_configs_from_model
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
@@ -542,7 +521,6 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -553,12 +531,6 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
         _fake_loader,
@@ -601,7 +573,6 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -612,12 +583,6 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
         _fake_loader,
@@ -706,7 +671,6 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -717,12 +681,6 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
         _fake_loader,
@@ -787,7 +745,6 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -798,12 +755,6 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
         _fake_loader,
@@ -862,7 +813,6 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -873,12 +823,6 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
         _fake_loader,
@@ -925,7 +869,6 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -936,12 +879,6 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
         _fake_loader,
@@ -1008,7 +945,6 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     for module_name in [
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
-        "vllm_omni.entrypoints.log_utils",
         "vllm_omni.entrypoints.omni_stage",
     ]:
         if module_name in sys.modules:
@@ -1019,17 +955,6 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.log_utils.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
     monkeypatch.setattr(
         "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
         _fake_loader,
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 51b434714..dbf05a26b 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -4,8 +4,7 @@ import asyncio
 import multiprocessing as mp
 import os
 import time
-import uuid
-from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
+from collections.abc import AsyncGenerator, Iterable
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import asdict
 from pprint import pformat
@@ -13,19 +12,17 @@ from typing import Any
 
 from omegaconf import OmegaConf
 from vllm.config import VllmConfig
-from vllm.inputs import PromptType
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.plugins.io_processors import get_io_processor
 from vllm.sampling_params import SamplingParams
-from vllm.tokenizers import TokenizerLike, init_tokenizer_from_config
+from vllm.tokenizers import TokenizerLike
 from vllm.v1.engine.exceptions import EngineDeadError
 
 # Internal imports (our code)
 from vllm_omni.config import OmniModelConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.distributed.omni_connectors import (
     get_stage_connector_config,
     initialize_orchestrator_connectors,
@@ -77,9 +74,9 @@ class AsyncOmni:
 
     Args:
         *args: Variable length argument list.
-            - args[0] (str): Model name or path to load.
+            - args[0]: Model name or path to load.
         **kwargs: Arbitrary keyword arguments.
-            - model (str): Model name or path to load (if not in args).
+            - model: Model name or path to load (if not in args).
             - stage_configs_path: Optional path to YAML file containing stage
               configurations. If None, configurations are loaded from the model.
             - log_stats: Whether to enable statistics logging
@@ -88,6 +85,8 @@ class AsyncOmni:
               each stage process during initialization
             - shm_threshold_bytes: Threshold in bytes for using shared memory
               for IPC. Objects larger than this threshold will use shared memory.
+            - worker_backend: Backend for worker processes. Default is "multi_process".
+            - ray_address: Address of Ray cluster for Ray backend, if using Ray backend.
             - batch_timeout: Timeout in seconds for batching requests within a stage
             - init_timeout: Timeout in seconds for waiting for all stages to initialize
             - Additional keyword arguments passed to stage engines.
@@ -213,8 +212,8 @@ class AsyncOmni:
                 results.append(fut.result())
         results.sort(key=lambda x: x[0])
         self.stage_list = [st for _, st in results]
-
         self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
+
         self.output_modalities = [st.final_output_type for st in self.stage_list]
         logger.debug("[AsyncOrchestrator] Loaded %d stages", len(self.stage_list))
 
@@ -364,281 +363,6 @@ class AsyncOmni:
                 self.io_processor = None
                 self.model_config = None
 
-    async def _run_generation_async(
-        self,
-        prompts: PromptType | Sequence[PromptType] | OmniDiffusionRequest | Sequence[OmniDiffusionRequest],
-        sampling_params_list: Any | Sequence[Any] | None = None,
-    ) -> AsyncGenerator[OmniRequestOutput, None]:
-        """Asynchronously run pipeline generation."""
-        logger.debug("[AsyncOrchestrator] generate() called")
-        if sampling_params_list is None:
-            raise ValueError("sampling_params_list is required for pipelined generation")
-
-        # Normalize sampling_params_list to a list
-        if not isinstance(sampling_params_list, (list, tuple)):
-            sampling_params_list = [sampling_params_list]
-        else:
-            sampling_params_list = list(sampling_params_list)
-
-        if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
-
-        # Normalize prompts to a list for per-request iteration
-        if not isinstance(prompts, (list, tuple)):
-            request_prompts: list[PromptType] = [prompts]
-        else:
-            request_prompts = list(prompts)
-
-        # Create async output queue monitors for each stage
-        stage_monitors = []
-        for stage_id, stage in enumerate(self.stage_list):
-            monitor = asyncio.create_task(self._monitor_stage_outputs_async(stage_id, stage))
-            stage_monitors.append(monitor)
-
-        # Orchestrator keeps stage objects for input derivation
-        num_stages = len(self.stage_list)
-
-        # Generate globally unique request IDs and map them to original prompts
-        request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
-        request_id_to_prompt: dict[str, PromptType] = {rid: p for rid, p in zip(request_ids, request_prompts)}
-
-        # Track start time for each request for end-to-end timing
-        _req_start_ts: dict[str, float] = {}
-        _wall_start_ts: float = time.time()
-
-        # Determine final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
-        final_stage_id_to_prompt: dict[str, int] = {}
-        for rid, prompt in request_id_to_prompt.items():
-            if isinstance(prompt, dict):
-                prompt_modalities = prompt.get("modalities", None)
-            else:
-                prompt_modalities = None
-            final_stage_id_for_e2e = get_final_stage_id_for_e2e(
-                prompt_modalities, self.output_modalities, self.stage_list
-            )
-            final_stage_id_to_prompt[rid] = final_stage_id_for_e2e
-
-        # Metrics/aggregation helper
-        metrics = OrchestratorMetrics(
-            num_stages,
-            self._enable_stats,
-            self._stats_file,
-            self._overall_stats_file,
-            _wall_start_ts,
-        )
-
-        # Seed all requests into stage-0 queue
-        logger.debug("[AsyncOrchestrator] Seeding %d requests into stage-0", len(request_prompts))
-        # Mark first input time for stage-0
-        metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
-
-        for req_id, prompt in request_id_to_prompt.items():
-            sp0 = sampling_params_list[0]  # type: ignore[index]
-            task = {
-                "request_id": req_id,
-                "engine_inputs": prompt,
-                "sampling_params": sp0,
-            }
-            logger.info(f"task: {task}")
-            self.stage_list[0].submit(task)
-            _req_start_ts[req_id] = time.time()
-            logger.debug("[AsyncOrchestrator] Enqueued request %s to stage-0", req_id)
-
-        # For each stage, forward results in stage order; collect final results at the end
-        # We pipeline by continuously polling output queues
-        remaining_by_stage: list[int] = [len(request_prompts)] + [0] * (num_stages - 1)
-        completed_requests = 0
-        total_requests = len(request_prompts)
-
-        logger.debug(
-            "[AsyncOrchestrator] Entering scheduling loop: total_requests=%d, stages=%d",
-            total_requests,
-            num_stages,
-        )
-
-        while completed_requests < total_requests:
-            # Asynchronously wait for output from any stage
-            await asyncio.sleep(0.001)  # Brief sleep to avoid CPU overload
-
-            for stage_id, stage in enumerate(self.stage_list):
-                result = stage.try_collect()
-                if result is None:
-                    continue
-
-                req_id = result.get("request_id")
-                if "error" in result:
-                    logger.error(
-                        "Stage %s error on request %s: %s",
-                        stage_id,
-                        req_id,
-                        result["error"],
-                    )
-                    continue
-
-                if result.get("type") == "stage_ready":
-                    # If stage initialization is slower than expected, wait briefly and retry
-                    await asyncio.sleep(0.05)
-                    continue
-
-                engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-                # Mark last output time for this stage whenever we receive outputs
-                metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
-
-                try:
-                    _m = result.get("metrics")
-                    if _m is not None:
-                        metrics.on_stage_metrics(stage_id, req_id, _m)
-                except Exception as e:
-                    logger.exception(
-                        "[AsyncOrchestrator] Failed to process metrics for stage %s, req %s: %s",
-                        stage_id,
-                        req_id,
-                        e,
-                    )
-
-                logger.debug(
-                    "[AsyncOrchestrator] Stage-%s completed request %s; forwarding or finalizing",
-                    stage_id,
-                    req_id,
-                )
-                stage.set_engine_outputs(engine_outputs)
-
-                if getattr(stage, "final_output", False):
-                    # Handle diffusion outputs that already contain images
-                    if stage.final_output_type == "image":
-                        # Extract images from engine_outputs if it's an OmniRequestOutput
-                        images = []
-                        output_to_check = engine_outputs[0] if isinstance(engine_outputs, list) else engine_outputs
-                        if isinstance(output_to_check, OmniRequestOutput) and output_to_check.images:
-                            images = output_to_check.images
-                        elif hasattr(output_to_check, "images") and output_to_check.images:
-                            images = output_to_check.images
-                        final_output = OmniRequestOutput(
-                            stage_id=stage_id,
-                            final_output_type=stage.final_output_type,
-                            request_output=engine_outputs,
-                            images=images,
-                        )
-                    else:
-                        final_output = OmniRequestOutput(
-                            stage_id=stage_id,
-                            final_output_type=stage.final_output_type,  # type: ignore[attr-defined]
-                            request_output=engine_outputs,
-                        )
-                    yield final_output
-                    logger.debug(
-                        "[AsyncOrchestrator] Request %s finalized at stage-%s",
-                        req_id,
-                        stage_id,
-                    )
-
-                    # End-to-end timing and time-per-token for final output
-                    # (only once per request at the designated final stage)
-                    try:
-                        rid_key = str(req_id)
-                        if stage_id == final_stage_id_to_prompt[req_id] and rid_key not in metrics.e2e_done:
-                            metrics.on_finalize_request(
-                                stage_id,
-                                req_id,
-                                engine_outputs,
-                                _req_start_ts.get(req_id, _wall_start_ts),
-                            )
-                    except Exception as e:
-                        logger.exception(
-                            "[AsyncOrchestrator] Finalize request handling error for req %s at stage %s: %s",
-                            req_id,
-                            stage_id,
-                            e,
-                        )
-
-                next_stage_id = stage_id + 1
-                if next_stage_id <= final_stage_id_to_prompt[req_id]:
-                    next_stage: OmniStage = self.stage_list[next_stage_id]
-                    try:
-                        next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
-                    except Exception as e:
-                        logger.exception(
-                            "[AsyncOrchestrator] Process engine inputs error for req %s at stage %s: %s",
-                            req_id,
-                            next_stage_id,
-                            e,
-                        )
-                        continue
-                    sp_next = sampling_params_list[next_stage_id]  # type: ignore[index]
-
-                    # Check connector for this edge
-                    connector_key = (str(stage_id), str(next_stage_id))
-                    connector = self.connectors.get(connector_key)
-                    sent_via_connector = False
-                    if connector:
-                        sent_via_connector = try_send_via_connector(
-                            connector=connector,
-                            stage_id=stage_id,
-                            next_stage_id=next_stage_id,
-                            req_id=req_id,
-                            next_inputs=next_inputs,
-                            sampling_params=sp_next,
-                            original_prompt=request_id_to_prompt[req_id],
-                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
-                            metrics=metrics,
-                        )
-
-                    if not sent_via_connector:
-                        raise RuntimeError(
-                            f"[AsyncOrchestrator] Failed to send request {req_id} "
-                            f"to stage-{next_stage_id} via connector. Configure a connector "
-                            "for this edge or inspect connector logs for details."
-                        )
-                    logger.debug(
-                        "[AsyncOrchestrator] Forwarded request %s to stage-%s",
-                        req_id,
-                        next_stage_id,
-                    )
-                    remaining_by_stage[next_stage_id] += 1
-                else:
-                    completed_requests += 1
-                    logger.debug(
-                        "[AsyncOrchestrator] Request %s fully completed (%d/%d)",
-                        req_id,
-                        completed_requests,
-                        total_requests,
-                    )
-
-        logger.debug("[AsyncOrchestrator] All requests completed")
-
-        # Summarize and print statistics
-        try:
-            summary = metrics.build_and_log_summary(final_stage_id_to_prompt)
-            logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
-        except Exception as e:
-            logger.exception("[AsyncOrchestrator] Failed to build/log summary: %s", e)
-
-    async def _monitor_stage_outputs_async(self, stage_id: int, stage: OmniStage) -> None:
-        """Asynchronously monitor stage output queue."""
-        while True:
-            result = stage.try_collect()
-            if result is not None:
-                # Put result into async queue for main loop processing
-                pass
-            await asyncio.sleep(0.001)
-
-        # Here is a duplicated init for the first stage to get the tokenizer,
-        # input_processor, and io_processor
-        # TODO: Refactor the code to avoid the duplicated init, we need to pass
-        # the tokenizer, input_processor, and io_processor back here from the
-        # first stage.
-        vllm_config = self.stage_list[0].vllm_config
-        self.tokenizer = init_tokenizer_from_config(model_config=vllm_config.model_config)
-        self.input_processor = OmniInputProcessor(
-            vllm_config=vllm_config,
-            tokenizer=self.tokenizer,
-        )
-        self.model_config = vllm_config.model_config
-        self.io_processor = get_io_processor(
-            vllm_config,
-            self.model_config.io_processor_plugin,
-        )
-
     def _initialize_stages(
         self,
         model: str,
@@ -761,35 +485,31 @@ class AsyncOmni:
         except Exception as e:
             logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
 
-    async def generate(
-        self,
-        prompt: PromptType,
-        request_id: str,
-        sampling_params_list: SamplingParams | Sequence[SamplingParams] | None = None,
-        output_modalities: list[str] | None = None,
-        lora_request: LoRARequest | None = None,
-        trace_headers: Mapping[str, str] | None = None,
-        priority: int = 0,
-        data_parallel_rank: int | None = None,
-    ) -> AsyncGenerator[OmniRequestOutput, None]:
+    async def generate(self, *args: Any, **kwargs: dict[str, Any]) -> AsyncGenerator[OmniRequestOutput, None]:
         """Generate outputs for the given prompt asynchronously.
 
+        Coordinates multi-stage pipeline through YAML configuration.
+        Each stage will use AsyncOmniLLM or AsyncOmniDiffusion based on stage_type.
         Processes the prompt through all stages in the pipeline and yields
         outputs as they become available. Each stage uses its corresponding
         sampling parameters from the sampling_params_list.
 
         Args:
-            prompt: Prompt to process. Can be a text string, token IDs,
-                or multimodal prompt.
-            request_id: Unique identifier for this request
-            sampling_params_list: List of SamplingParams, one for each stage.
-                Must have the same length as the number of stages.
-                If None, uses default sampling params for each stage.
-            lora_request: Optional LoRA adapter request for this generation
-            trace_headers: Optional tracing headers for observability
-            priority: Request priority (higher values processed first)
-            data_parallel_rank: Optional data parallel rank for distributed
-                inference
+            *args: Arguments for generation.
+                - prompt: Prompt to process. Can be a text string, token IDs,
+                    or multimodal prompt.
+                - request_id: Unique identifier for this request
+                - sampling_params_list: List of SamplingParams, one for each stage.
+                    Must have the same length as the number of stages.
+                    If None, uses default sampling params for each stage.
+            **kwargs: Additional arguments for generation.
+                - prompt: Prompt to process. Can be a text string, token IDs,
+                    or multimodal prompt.
+                - request_id: Unique identifier for this request
+                - sampling_params_list: List of SamplingParams, one for each stage.
+                    Must have the same length as the number of stages.
+                    If None, uses default sampling params for each stage.
+                - output_modalities: Optional list of output modalities.
 
         Yields:
             OmniRequestOutput objects as they are produced by each stage.
@@ -803,18 +523,41 @@ class AsyncOmni:
         async with self._pause_cond:
             await self._pause_cond.wait_for(lambda: not self._paused)
 
-        logger.debug("[Orchestrator] generate() called")
+        logger.debug("[AsyncOrchestrator] generate() called")
 
         # Start output handler on the first call to generate()
         self._run_output_handler()
 
+        prompt = args[0] if args else kwargs.get("prompt")
+        request_id = args[1] if len(args) > 1 else kwargs.get("request_id")
+        sampling_params_list = args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
+        output_modalities = kwargs.get("output_modalities", None)
+        # TODO: lora_request, trace_headers, priority are not supported yet
+
         if sampling_params_list is None:
-            sampling_params_list = self.default_sampling_params_list
+            # For Omni LLM, the params are parsed via the yaml file. For the current version,
+            # diffusion params can parsed via the command line.
+            omni_params_kwargs = {
+                k: v for k, v in kwargs.items() if k not in ["prompt", "request_id", "output_modalities"]
+            }
+
+            per_stage_params: list[Any] = []
+            for stage_id, stage in enumerate(self.stage_list):
+                stage_type = getattr(stage, "stage_type", "llm")
+                if stage_type == "diffusion":
+                    default_dict = self.default_sampling_params_list[stage_id]
+                    # Merge user-provided kwargs
+                    merged = {**default_dict, **omni_params_kwargs}
+                    # Diffusion only needs to keep diff params, will be used via OmniDiffusionRequest
+                    per_stage_params.append(merged)
+                else:
+                    # LLM directly constructs SamplingParams, don't use the merged params
+                    per_stage_params.append(self.default_sampling_params_list[stage_id])
+
+            sampling_params_list = per_stage_params
+
         if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(
-                f"Expected {len(self.stage_list)} sampling params, \
-                got {len(sampling_params_list)}"
-            )
+            raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
 
         # Orchestrator keeps stage objects for input derivation
         num_stages = len(self.stage_list)
@@ -834,7 +577,7 @@ class AsyncOmni:
             _wall_start_ts,
         )
         # Seed stage-0 queue with all requests
-        logger.debug("[Orchestrator] Seeding request into stage-0")
+        logger.debug("[AsyncOrchestrator] Seeding request into stage-0")
         req_state = ClientRequestState(request_id)
         self.request_states[request_id] = req_state
 
@@ -849,9 +592,9 @@ class AsyncOmni:
         }
         self.stage_list[0].submit(task)
         _req_start_ts[request_id] = time.time()
-        logger.debug("[Orchestrator] Enqueued request %s to stage-0", request_id)
+        logger.debug("[AsyncOrchestrator] Enqueued request %s to stage-0", request_id)
 
-        logger.debug("[Orchestrator] Entering scheduling loop: stages=%d", num_stages)
+        logger.debug("[AsyncOrchestrator] Entering scheduling loop: stages=%d", num_stages)
         for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
             result = await req_state.queue.get()
             assert stage_id == req_state.stage_id
@@ -875,15 +618,13 @@ class AsyncOmni:
                     metrics.on_stage_metrics(stage_id, req_id, _m)
             except Exception as e:
                 logger.exception(
-                    "[Orchestrator] Failed to process metrics for stage %s, \
-                        req %s: %s",
+                    "[AsyncOrchestrator] Failed to process metrics for stage %s, req %s: %s",
                     stage_id,
                     req_id,
                     e,
                 )
             logger.debug(
-                "[Orchestrator] Stage-%s completed request %s; \
-                    forwarding or finalizing",
+                "[AsyncOrchestrator] Stage-%s completed request %s; forwarding or finalizing",
                 stage_id,
                 req_id,
             )
@@ -891,7 +632,7 @@ class AsyncOmni:
 
             if getattr(stage, "final_output", False):
                 logger.debug(
-                    "[Orchestrator] Request %s finalized at stage-%s",
+                    "[AsyncOrchestrator] Request %s finalized at stage-%s",
                     req_id,
                     stage_id,
                 )
@@ -909,8 +650,7 @@ class AsyncOmni:
                         )
                 except Exception as e:
                     logger.exception(
-                        "[Orchestrator] Finalize request handling error for \
-                            req %s at stage %s: %s",
+                        "[AsyncOrchestrator] Finalize request handling error for req %s at stage %s: %s",
                         req_id,
                         stage_id,
                         e,
@@ -969,27 +709,27 @@ class AsyncOmni:
                     # because continuing would cause the request to be silently dropped
                     # and the orchestrator to hang waiting for completion.
                     error_msg = (
-                        f"[Orchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
+                        f"[AsyncOrchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
                         "Configure a connector for this edge or inspect connector logs for details."
                     )
                     logger.error(error_msg)
                     raise RuntimeError(error_msg)
                 logger.debug(
-                    "[Orchestrator] Forwarded request %s to stage-%s",
+                    "[AsyncOrchestrator] Forwarded request %s to stage-%s",
                     req_id,
                     next_stage_id,
                 )
             else:
-                logger.debug("[Orchestrator] Request %s fully completed", req_id)
+                logger.debug("[AsyncOrchestrator] Request %s fully completed", req_id)
 
-        logger.debug("[Orchestrator] All requests completed")
+        logger.debug("[AsyncOrchestrator] All requests completed")
 
         # Summarize and print stats
         try:
             summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
             logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
         except Exception as e:
-            logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
+            logger.exception("[AsyncOrchestrator] Failed to build/log summary: %s", e)
         finally:
             self.request_states.pop(request_id, None)
 
diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py
index a35fa35f8..b57b063c1 100644
--- a/vllm_omni/entrypoints/omni.py
+++ b/vllm_omni/entrypoints/omni.py
@@ -10,11 +10,9 @@ from dataclasses import asdict
 from pprint import pformat
 from typing import Any
 
-import msgspec
 from omegaconf import OmegaConf
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
-from vllm.sampling_params import SamplingParams
 
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.distributed.omni_connectors import (
@@ -68,6 +66,15 @@ class Omni:
               configurations. If None, configurations are loaded from the model.
             - log_stats: Whether to enable statistics logging
               be written to files with stage-specific suffixes.
+            - init_sleep_seconds: Number of seconds to sleep between starting
+              each stage process during initialization
+            - shm_threshold_bytes: Threshold in bytes for using shared memory
+              for IPC. Objects larger than this threshold will use shared memory.
+            - worker_backend: Backend for worker processes. Default is "multi_process".
+            - ray_address: Address of Ray cluster for Ray backend, if using Ray backend.
+            - batch_timeout: Timeout in seconds for batching requests within a stage
+            - init_timeout: Timeout in seconds for waiting for all stages to initialize
+            - Additional keyword arguments passed to stage engines.
 
     Example:
         >>> omni = Omni(model="Qwen/Qwen2.5-Omni-7B")
@@ -173,6 +180,7 @@ class Omni:
                 results.append(fut.result())
         results.sort(key=lambda x: x[0])
         self.stage_list = [st for _, st in results]
+        self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
         self.output_modalities = [st.final_output_type for st in self.stage_list]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
@@ -300,20 +308,24 @@ class Omni:
             prompts = kwargs.get("prompt")
 
         if sampling_params_list is None:
-            omni_params_kwargs = {k: v for k, v in kwargs.items() if k not in ["prompts", "sampling_params_list"]}
+            # For Omni LLM, the params are parsed via the yaml file. For the current version,
+            # diffusion params can parsed via the command line.
+            omni_params_kwargs = {
+                k: v for k, v in kwargs.items() if k not in ["prompt", "request_id", "output_modalities"]
+            }
 
             per_stage_params: list[Any] = []
-            for stage in self.stage_list:
+            for stage_id, stage in enumerate(self.stage_list):
                 stage_type = getattr(stage, "stage_type", "llm")
-                default_dict = msgspec.to_builtins(getattr(stage, "default_sampling_params", {}))
-                # Merge user-provided kwargs
-                merged = {**default_dict, **omni_params_kwargs}
                 if stage_type == "diffusion":
+                    default_dict = self.default_sampling_params_list[stage_id]
+                    # Merge user-provided kwargs
+                    merged = {**default_dict, **omni_params_kwargs}
                     # Diffusion only needs to keep diff params, will be used via OmniDiffusionRequest
                     per_stage_params.append(merged)
                 else:
-                    # LLM directly constructs SamplingParams
-                    per_stage_params.append(SamplingParams(**merged))
+                    # LLM directly constructs SamplingParams, don't use the merged params
+                    per_stage_params.append(self.default_sampling_params_list[stage_id])
 
             sampling_params_list = per_stage_params
         return self._run_generation(prompts, sampling_params_list)