[Engine] Add entrypoint class and stage management by Gaohan123 · Pull Request #13 · vllm-project/vllm-omni

Gaohan123 · 2025-10-22T06:46:59Z

Purpose

This PR implements Phase 1 features of #10 . It adds entrypoint class OmniLLM for vLLM-omni. And it also builds the stage management of multiple modules within a model.

Test Plan

Test Result

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/hsliuustc0106/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

gemini-code-assist · 2025-10-22T06:47:17Z

Summary of Changes

Hello @Gaohan123, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request lays the foundational infrastructure for multi-stage model execution within the vllm-omni framework. It introduces a new entrypoint class, OmniLLM, designed to orchestrate complex workflows by chaining together different model stages. This enables flexible configuration and dynamic management of multi-modal models, moving towards a more modular and extensible architecture for advanced generation capabilities.

Highlights

New Configuration Classes: Introduced OmniModelConfig to extend base model configurations with stage-specific parameters and OmniEngineArgs to provide command-line argument parsing for these new omni-specific engine settings.
Refactored OmniLLM Entrypoint: The OmniLLM class has been completely rewritten to manage a sequence of Stage objects, orchestrating multi-stage model execution. It now dynamically loads stage configurations and chains their outputs.
Stage Management System: Replaced the StageManager with a new Stage class, which encapsulates the configuration, engine, and input/output processing logic for individual stages within a multi-stage model pipeline.
Dynamic Stage Configuration Loading: Added utility functions to load stage configurations from YAML files based on the model type, enabling flexible and declarative definition of multi-stage model architectures.
Multi-modal Input/Output Handling: Defined OmniTokensPrompt to support multi-modal inputs (e.g., prompt embeddings, additional information) and OmniRequestOutput to capture stage-specific output details.
Example Stage Configuration: Included an example YAML configuration for the qwen2_5_omni model, demonstrating how to define multiple stages like 'thinker', 'talker', and 'code2wav' with their respective engine arguments and input/output types.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This PR introduces the OmniLLM class for managing multi-stage models, along with stage management and configuration loading. It also includes a StageLLM class and updates to configuration and input data structures. The review focuses on correctness and potential issues in the new implementation.

gemini-code-assist · 2025-10-22T06:48:38Z

-
-        return combined
+            stage = Stage(stage_config)
+            omni_llm = OmniLLM(model=model, **stage_config.engine_args)


It seems like you are creating an OmniLLM instance within the initialize_stages method, but the OmniLLM class is designed to manage multiple stages. This could lead to confusion and potential issues. Consider whether you should be instantiating StageLLM here instead.

gemini-code-assist · 2025-10-22T06:48:38Z

+@config
+@dataclass(config=ConfigDict(arbitrary_types_allowed=True))


Consider adding a docstring to explain what the @config decorator does and why it's being used here. This will improve readability and understanding of the code.

gemini-code-assist · 2025-10-22T06:48:38Z

+        parser.add_argument("--model-stage", type=str, default=EngineArgs.model_stage, 
+        help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.")


The help message for --model-stage uses EngineArgs.model_stage as the default value, but it should reference OmniEngineArgs.model_stage to be accurate.

Suggested change

parser.add_argument("--model-stage", type=str, default=EngineArgs.model_stage,

help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.")

parser.add_argument("--model-stage", type=str, default=OmniEngineArgs.model_stage,

help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.")

gemini-code-assist · 2025-10-22T06:48:38Z

+                engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)
+            else:
+                engine_inputs = prompts
+            engine_outputs = self._run_generation(stage, sampling_params_list[stage_id], engine_inputs)


Consider adding a check to ensure that sampling_params_list has the same length as self.stage_list. If they don't match, it could lead to indexing errors or incorrect behavior.

gemini-code-assist · 2025-10-22T06:48:38Z

+            if len(self.engine_input_source) == 0:
+                raise ValueError("engine_input_source is empty")


Raising a ValueError if engine_input_source is empty might be too strict. There might be valid cases where a stage doesn't require an input source. Consider using a different approach, such as skipping the input processing or providing a default input.

hsliuustc0106 · 2025-10-22T14:40:20Z

 """
 Configuration module for vLLM-omni.
 """
+from vllm.config import ModelConfig


import order should be reconfigrated, maybe we need to add precommit

fixed. Later I think we can uniformly build a precommit and submit a PR to recheck all existing codes.

hsliuustc0106 · 2025-10-22T14:42:36Z

+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.AR_gpu_worker.ARGPUWorker


the worker file and class name should be renamed

hsliuustc0106 · 2025-10-22T14:42:54Z

+      engine_output_type: latent
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+


del blank here

hsliuustc0106 · 2025-10-22T14:48:26Z

-            raise IndexError(f"Stage config {stage_id} not found. Available stages: 0-{len(self.stage_configs)-1}")
-
-        return self.stage_configs[stage_id]
+class Stage:


Stage -> OmniStage
does it only process one stage?

if so, the stage_manager.py file name should be changed

Yes it only process one. Personally I think the name "Stage" is simple and clear. vLLM main doesn't have this concept. Also I modify stage_manager.py to stage.py

hsliuustc0106 · 2025-10-22T14:51:23Z

@@ -0,0 +1,32 @@
+from vllm.inputs import TextPrompt


imports should be reordered

hsliuustc0106 · 2025-10-22T14:58:35Z

+        return engine_outputs
+
+
+class StageLLM(LLM): 


in this file, OmniLLM, StageLLM, LLM, these relateionships between them are not clear

After thinking, OmniLLM is the main entry class, which keeps consistent usage habit with LLM in vLLM main branch. And StageLLM inherited from LLM to process each stage with one engine. Specifically, One OmniLLM includes multiple StageLLMs, which inherits from LLM.

Gaohan123 · 2025-10-24T02:24:47Z

Fix necessary reviews from #23 , other reviews for future development

hsliuustc0106 · 2025-10-24T02:38:03Z

lgtm
approve

[Engine] Add entrypoint class and stage management

talker debug

Add entrypoint class and stage management

20cfd15

Gaohan123 requested review from fake0fan, hsliuustc0106 and tzhouam October 22, 2025 06:46

gemini-code-assist Bot reviewed Oct 22, 2025

View reviewed changes

Gaohan123 added 2 commits October 22, 2025 17:46

fix some reviews from gemini

6111c28

supplement some imports

87d6414

Gaohan123 changed the title ~~[Core] Add entrypoint class and stage management~~ [Engine] Add entrypoint class and stage management Oct 22, 2025

hsliuustc0106 reviewed Oct 22, 2025

View reviewed changes

fix for some reviews

98a85d4

Gaohan123 requested a review from hsliuustc0106 October 23, 2025 02:32

Copilot AI mentioned this pull request Oct 24, 2025

[Review] Add comprehensive expert review of PR #13 - Entrypoint class and stage management #23

Closed

3 tasks

Gaohan123 added 2 commits October 24, 2025 10:17

fix for reviews of copilot

2c08c04

fix file name

61a215f

hsliuustc0106 merged commit 15dbe97 into vllm-project:main Oct 24, 2025

Gaohan123 deleted the entrypoint_stage branch November 1, 2025 02:27

princepride pushed a commit to princepride/vllm-omni that referenced this pull request Jan 10, 2026

Merge pull request vllm-project#13 from Gaohan123/entrypoint_stage

1ce4e71

[Engine] Add entrypoint class and stage management

R2-Y pushed a commit to R2-Y/vllm-omni that referenced this pull request Jan 20, 2026

Merge pull request vllm-project#13 from R2-Y/why_async_chunk

6ce85ba

talker debug

		@config
		@dataclass(config=ConfigDict(arbitrary_types_allowed=True))

		parser.add_argument("--model-stage", type=str, default=EngineArgs.model_stage,
		help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.")

		if len(self.engine_input_source) == 0:
		raise ValueError("engine_input_source is empty")

Conversation

Gaohan123 commented Oct 22, 2025

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist Bot commented Oct 22, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Oct 22, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 22, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 22, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 22, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 22, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Gaohan123 commented Oct 24, 2025

Uh oh!

hsliuustc0106 commented Oct 24, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants