Remove obsolete RFC documents for automated rebase agent, output processor refactor, model runner clearance, long function decomposition, and model runner V2 integration. These files are no longer relevant to the current project direction.
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
This commit is contained in:
Taichang Zhou committed
1 parent
5511bc14c7
commit
2187db0e36
5 files changed
-1493
No files matched your search
@@ -1,196 +0,0 @@
|
||||
# RFC-001: Fully Automated Nightly Rebase via Multi-Agent System
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes building a fully automated agent system that performs a nightly rebase of vllm-omni onto the latest upstream vLLM commit. The primary challenge is that the codebase is large enough to exceed any single LLM context window, necessitating a **multi-agent architecture** that decomposes the rebase workflow into context-bounded subtasks.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current State
|
||||
|
||||
vllm-omni is an independent package that extends upstream vLLM (`vllm.v1.*`) with multimodal capabilities (text, image, video, audio, and diffusion models). Keeping vllm-omni in sync with upstream vLLM is critical because:
|
||||
|
||||
- **API drift**: vLLM frequently changes internal APIs (e.g., `GPUModelRunner`, `OutputProcessor`, `SchedulerOutput`). Each day of lag compounds merge conflict complexity.
|
||||
- **Manual burden**: Past rebase efforts (e.g., `75770c93 - Rebase to vllm v0.16.0`) required significant manual work, including fixing downstream references across workers, model runners, and platform adapters.
|
||||
|
||||
### Why Automation Is Difficult
|
||||
|
||||
| Challenge | Detail |
|
||||
|-----------|--------|
|
||||
| **Context window limits** | vllm-omni has 1,200+ lines in `gpu_model_runner.py`, 1,500+ in `omni_stage.py`, 800+ in `gpu_generation_model_runner.py`, and dozens of model-specific files. A single-agent approach cannot hold enough context to resolve all conflicts. |
|
||||
| **Semantic conflicts** | Many conflicts are not textual but semantic — upstream may rename a parameter, change a method signature, or restructure a class hierarchy. Blind merge tools fail silently. |
|
||||
| **Cross-file dependencies** | A change in `vllm.v1.worker.gpu_model_runner.GPUModelRunner` can cascade into `OmniGPUModelRunner`, `GPUARModelRunner`, `GPUGenerationModelRunner`, and all platform variants (NPU, XPU, ROCm). |
|
||||
| **Test validation** | After resolving conflicts, the result must be validated against CI (L1–L5 test levels, see `docs/contributing/ci/CI_5levels.md`). |
|
||||
|
||||
## 3. Proposed Design
|
||||
|
||||
### 3.1 Architecture Overview
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────────┐
|
||||
│ Orchestrator Agent │
|
||||
│ (Nightly cron trigger, git operations, conflict decomposition) │
|
||||
└───────┬──────────────┬──────────────┬──────────────┬────────────────┘
|
||||
│ │ │ │
|
||||
▼ ▼ ▼ ▼
|
||||
┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐
|
||||
│ Conflict │ │ Conflict │ │ Conflict │ │ Validation │
|
||||
│ Resolver │ │ Resolver │ │ Resolver │ │ Agent │
|
||||
│ Agent #1 │ │ Agent #2 │ │ Agent #N │ │ │
|
||||
│ (worker/) │ │ (engine/) │ │ (platform/) │ │ (CI runner) │
|
||||
└──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘
|
||||
```
|
||||
|
||||
### 3.2 Agent Roles
|
||||
|
||||
#### Orchestrator Agent
|
||||
|
||||
- **Trigger**: GitHub Actions cron schedule (e.g., `0 2 * * *` UTC)
|
||||
- **Responsibilities**:
|
||||
1. Fetch the latest upstream vLLM commit (`pip install vllm` or track a pinned ref)
|
||||
2. Attempt `git rebase` or dependency version bump on a fresh branch `auto-rebase/YYYY-MM-DD`
|
||||
3. If conflicts arise, decompose them by file/module into independent work units
|
||||
4. Dispatch work units to Conflict Resolver Agents
|
||||
5. Collect results, apply patches, and trigger Validation Agent
|
||||
6. If validation passes, open a PR; if not, open an issue with diagnostics
|
||||
|
||||
#### Conflict Resolver Agent (per-module)
|
||||
|
||||
- **Context boundary**: Each agent receives only the conflicting file(s) plus their direct upstream counterpart(s) and a module-level dependency graph
|
||||
- **Strategy**:
|
||||
1. Receive: `(omni_file, upstream_diff, conflict_markers, api_change_summary)`
|
||||
2. Analyze the upstream change (renamed method? new parameter? restructured class?)
|
||||
3. Apply the minimal adaptation to vllm-omni code
|
||||
4. Output: patch file + confidence score + human-review flags
|
||||
- **Module partitions** (based on current codebase structure):
|
||||
- `worker/` — `OmniGPUModelRunner`, `GPUARModelRunner`, `GPUGenerationModelRunner`
|
||||
- `engine/` — `OmniInputProcessor`, `MultimodalOutputProcessor`, `OmniEngineArgs`
|
||||
- `entrypoints/` — `omni_stage.py`, `omni_llm.py`, `async_omni_llm.py`
|
||||
- `model_executor/` — Model-specific code (Qwen3, MiMo, GLM, etc.)
|
||||
- `platforms/` — NPU, XPU, ROCm adapters
|
||||
- `diffusion/` — DiffusionModelRunner and related code
|
||||
- `core/sched/` — Scheduler adaptations
|
||||
|
||||
#### Validation Agent
|
||||
|
||||
- **Responsibilities**:
|
||||
1. Run lint checks (`ruff`, `pre-commit`)
|
||||
2. Run import-level smoke tests
|
||||
3. Trigger CI L1/L2 tests (unit tests, basic e2e)
|
||||
4. Report pass/fail with logs
|
||||
- **Escalation**: If validation fails, feed error context back to the relevant Conflict Resolver Agent for a retry (up to N attempts)
|
||||
|
||||
### 3.3 Context Management Strategy
|
||||
|
||||
The key innovation is **bounded context windows** per agent:
|
||||
|
||||
| Strategy | Description |
|
||||
|----------|-------------|
|
||||
| **File-level isolation** | Each Conflict Resolver Agent sees at most 2-3 files (~2-4K lines total) |
|
||||
| **API summary injection** | Instead of full upstream source, agents receive a structured summary of API changes (method signatures, parameter changes, class hierarchy diffs) |
|
||||
| **Dependency graph pruning** | Only relevant imports and call sites are included in context |
|
||||
| **Iterative refinement** | If a patch introduces cross-module issues, the Orchestrator re-dispatches with expanded context |
|
||||
|
||||
### 3.4 Workflow
|
||||
|
||||
```
|
||||
1. [Cron] Trigger nightly at 02:00 UTC
|
||||
2. [Orchestrator] git fetch upstream; attempt rebase/version bump
|
||||
3. [Orchestrator] If clean → skip to step 7
|
||||
4. [Orchestrator] Parse conflicts → group by module → create work units
|
||||
5. [Resolver Agents] Resolve conflicts in parallel (bounded context per agent)
|
||||
6. [Orchestrator] Collect patches → apply sequentially → check for cross-module issues
|
||||
7. [Validation Agent] Run lint + L1/L2 tests
|
||||
8. [Orchestrator] If pass → open PR "Auto-rebase to vLLM <commit>"
|
||||
9. [Orchestrator] If fail → retry up to 3 times with error feedback
|
||||
10. [Orchestrator] If still failing → open issue with diagnostics for human review
|
||||
```
|
||||
|
||||
### 3.5 GitHub Actions Integration
|
||||
|
||||
```yaml
|
||||
name: Nightly Auto-Rebase
|
||||
on:
|
||||
schedule:
|
||||
- cron: '0 2 * * *'
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
auto-rebase:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Run Orchestrator Agent
|
||||
run: python scripts/auto_rebase/orchestrator.py
|
||||
env:
|
||||
AGENT_API_KEY: ${{ secrets.AGENT_API_KEY }}
|
||||
- name: Open PR or Issue
|
||||
run: python scripts/auto_rebase/finalize.py
|
||||
```
|
||||
|
||||
## 4. Key Technical Decisions
|
||||
|
||||
### 4.1 Rebase vs. Dependency Bump
|
||||
|
||||
vllm-omni imports vLLM as a dependency rather than forking it. The "rebase" is actually:
|
||||
|
||||
1. Bumping the pinned vLLM version in `requirements/common.txt` or `setup.py`
|
||||
2. Adapting vllm-omni code to API changes in the new vLLM version
|
||||
3. Updating Dockerfiles (e.g., `Dockerfile.ci` base image `vllm/vllm-openai:v0.16.0`, `Dockerfile.rocm` notes)
|
||||
|
||||
### 4.2 Agent Backend
|
||||
|
||||
- Recommend using a code-specialized LLM with 128K+ context window
|
||||
- Each Conflict Resolver Agent should receive structured prompts with:
|
||||
- The specific vllm-omni file(s) being adapted
|
||||
- A diff of the upstream API change
|
||||
- The inheritance/import chain relevant to the file
|
||||
- Test files that cover the changed code
|
||||
|
||||
### 4.3 Rollback Safety
|
||||
|
||||
- All changes happen on a dedicated branch (`auto-rebase/YYYY-MM-DD`)
|
||||
- The PR requires at least one human approval before merge
|
||||
- Failed rebases are tracked in issues with full diagnostic context
|
||||
|
||||
## 5. Risks and Mitigations
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|------------|
|
||||
| Agent produces semantically wrong patches | Confidence scoring + human review for low-confidence patches |
|
||||
| Cross-module inconsistencies | Orchestrator performs a global consistency check after collecting all patches |
|
||||
| Upstream breaking changes too large | Fall back to opening a detailed issue rather than a broken PR |
|
||||
| API key / cost management | Rate-limit agent calls; cache upstream API summaries across runs |
|
||||
| Flaky CI tests | Distinguish test failures caused by rebase vs. pre-existing flakiness |
|
||||
|
||||
## 6. Success Metrics
|
||||
|
||||
- **Automation rate**: % of nightly rebases that produce a clean PR without human intervention
|
||||
- **Time to rebase**: Wall-clock time from cron trigger to PR creation (target: < 30 minutes)
|
||||
- **Conflict resolution accuracy**: % of agent-resolved conflicts that pass human review
|
||||
- **CI pass rate**: % of auto-rebase PRs that pass L1/L2 CI on first attempt
|
||||
|
||||
## 7. Implementation Plan
|
||||
|
||||
| Phase | Scope | Timeline |
|
||||
|-------|-------|----------|
|
||||
| Phase 1 | Orchestrator skeleton: git operations, conflict detection, branch management | 2 weeks |
|
||||
| Phase 2 | Conflict Resolver Agents: per-module resolution with bounded context | 3 weeks |
|
||||
| Phase 3 | Validation Agent: lint + smoke test integration | 1 week |
|
||||
| Phase 4 | GitHub Actions integration + PR/issue creation | 1 week |
|
||||
| Phase 5 | Monitoring, confidence tuning, feedback loops | Ongoing |
|
||||
|
||||
## 8. Open Questions
|
||||
|
||||
1. Should the system attempt to rebase daily or only when upstream has significant changes?
|
||||
2. What is the maximum acceptable cost per rebase run?
|
||||
3. Should platform-specific adapters (NPU, XPU, ROCm) be included in the nightly rebase or handled separately?
|
||||
4. How should we handle upstream changes that require new test coverage in vllm-omni?
|
||||
@@ -1,465 +0,0 @@
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes a comprehensive refactoring of the output processing pipeline in vllm-omni. The current `MultimodalOutputProcessor` and `OmniRequestState` classes have grown organically to handle multiple output modalities (text, image, audio, latents) but lack a clear, extensible architecture. A significant portion of the current code is **dead or unreachable**. We propose first simplifying the existing processor by removing dead code, then introducing a **Modality Registry**, a **Router abstraction**, and **composable output handlers** to support emerging multi-output models (e.g., DeepSeek with text+video+image).
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current Architecture
|
||||
|
||||
The output processing pipeline is concentrated in `vllm_omni/engine/output_processor.py` (469 lines) with two main classes:
|
||||
|
||||
- **`MultimodalOutputProcessor`** (extends `vllm.v1.engine.output_processor.OutputProcessor`): Routes `EngineCoreOutput` objects by `output_type` string to modality-specific handler methods.
|
||||
- **`OmniRequestState`** (extends `RequestState`): Accumulates multimodal tensors per-request during generation, with deferred concatenation.
|
||||
|
||||
### Problems
|
||||
|
||||
#### P0: Significant Dead / Unreachable Code
|
||||
|
||||
A thorough audit of the output processor reveals that **a large portion of the code is never executed** in any real production path. This dead code adds confusion, increases maintenance burden, and obscures the actual data flow. See Section 3 for the full dead code inventory.
|
||||
|
||||
#### P1: Rigid String-Based Routing
|
||||
|
||||
The `_route_and_normalize` method uses hardcoded string matching:
|
||||
|
||||
```python
|
||||
if output_type == "image":
|
||||
self._process_image_output(eco)
|
||||
elif output_type in ("text+image", "text,image", "image+text"):
|
||||
self._process_text_image_output(eco)
|
||||
elif output_type in ("latents", "latent"):
|
||||
self._process_latents_output(eco)
|
||||
elif output_type in ("audio", "speech"):
|
||||
self._process_audio_output(eco)
|
||||
elif output_type == "text":
|
||||
self._process_text_output(eco)
|
||||
```
|
||||
|
||||
Adding a new modality (e.g., `video`, `text+video`, `text+video+image`) requires modifying this central dispatch chain.
|
||||
|
||||
#### P2: No Composable Multi-Output Support
|
||||
|
||||
Models like DeepSeek may produce **text + video + image** simultaneously. The current architecture handles at most two modalities (e.g., `text+image`) via a single combined handler. There is no composition mechanism for arbitrary output type combinations.
|
||||
|
||||
#### P3: Inconsistent Key Extraction
|
||||
|
||||
Each `_process_*_output` method has its own hardcoded key lists:
|
||||
|
||||
- Image: `("image", "images", "pixel_values", "pixels")`
|
||||
- Audio: `("audio", "audios", "wav", "waveform", "audio_pcm", "pcm")`
|
||||
- Latents: `("latent", "latents", "z", "posterior")`
|
||||
|
||||
There is no standard contract between model output format and the output processor.
|
||||
|
||||
#### P4: Monolithic Accumulation Logic
|
||||
|
||||
`OmniRequestState.add_multimodal_tensor` handles all modalities in a single method with complex branching for dict/tensor/list inputs and nested dict merging (~65 lines). The special-case for audio tensor concatenation in `_consolidate_multimodal_tensors` (skipping `torch.cat` due to shape inconsistency) further indicates the need for per-modality accumulation strategies.
|
||||
|
||||
#### P5: Missing Modalities
|
||||
|
||||
- **Video output**: No dedicated `"video"` routing branch exists. Video generation would currently fall through to the heuristic fallback.
|
||||
- **Text+Video, Text+Video+Image**: No handling exists for these combinations, which future models (e.g., DeepSeek multimodal) will require.
|
||||
|
||||
## 3. Dead Code Audit and Simplification Plan
|
||||
|
||||
A line-by-line analysis of `output_processor.py` cross-referenced with all callers and stage configs reveals the following dead/unused code:
|
||||
|
||||
### 3.1 Dead Code Inventory
|
||||
|
||||
#### D1: `register_handler()` — Never Called (lines 281-292)
|
||||
|
||||
The `register_handler(modality, handler)` method is defined on `MultimodalOutputProcessor` and initializes `self.output_handlers: dict[str, Callable]`. However, **no code in the entire codebase calls `register_handler()`**. The `output_handlers` dict is always empty, making the custom handler dispatch path in `_route_and_normalize` (lines 373-379) unreachable.
|
||||
|
||||
**Evidence**: `rg "register_handler" vllm_omni/` only finds the definition and the docstring reference. Zero callers.
|
||||
|
||||
**Action**: Remove `register_handler()`, `self.output_handlers`, and the custom handler dispatch block.
|
||||
|
||||
#### D2: `_reqid_to_mm_type` — Written but Never Read (line 278, 346-350)
|
||||
|
||||
The `_reqid_to_mm_type: dict[str, str]` is populated in `process_outputs()`:
|
||||
|
||||
```python
|
||||
self._reqid_to_mm_type.clear()
|
||||
for eco in engine_core_outputs:
|
||||
mm_type = (self.engine_core_output_type or "").lower()
|
||||
if mm_type:
|
||||
self._reqid_to_mm_type[eco.request_id] = mm_type
|
||||
```
|
||||
|
||||
But this dict is **never read anywhere** — not in this class, not in any subclass, not in any caller. It is populated and then cleared on the next call.
|
||||
|
||||
**Action**: Remove `_reqid_to_mm_type` and the population code.
|
||||
|
||||
#### D3: `_process_text_image_output` — No Stage Config Uses "text+image" (lines 383-384, 406-424)
|
||||
|
||||
The routing branch `elif output_type in ("text+image", "text,image", "image+text")` and its handler `_process_text_image_output` are never triggered because:
|
||||
|
||||
- No stage config in the entire project uses `engine_output_type: text+image`, `text,image`, or `image+text`
|
||||
- `OmniEngineCoreOutput` does not have an `output_type` attribute, so `getattr(eco, "output_type", ...)` always falls back to `engine_core_output_type`
|
||||
- No code path ever sets `engine_core_output_type` to any of these composite strings
|
||||
|
||||
**Actual stage config values** (exhaustive search across all YAML files):
|
||||
- `latent` — Qwen3-TTS, Qwen3-Omni, Qwen2.5-Omni, MiMo-Audio
|
||||
- `audio` — Qwen3-TTS, Qwen3-Omni, Qwen2.5-Omni, MiMo-Audio
|
||||
- `text` — Bagel
|
||||
- `image` — Bagel
|
||||
- `token_ids` — GLM-Image
|
||||
|
||||
**Action**: Remove `_process_text_image_output` and its routing branch.
|
||||
|
||||
#### D4: `"speech"` Alias — Never Used (line 387)
|
||||
|
||||
The routing matches `("audio", "speech")`, but no stage config ever uses `engine_output_type: speech`. Only `"audio"` is used in practice.
|
||||
|
||||
**Action**: Remove `"speech"` from the match (minor, but reduces confusion).
|
||||
|
||||
#### D5: `_extract_from_multimodal_outputs` — Always Returns None (lines 457-469)
|
||||
|
||||
This method tries to read `eco.multimodal_outputs`:
|
||||
|
||||
```python
|
||||
mm = getattr(eco, "multimodal_outputs", None)
|
||||
```
|
||||
|
||||
However, `OmniEngineCoreOutput` (defined in `vllm_omni/engine/__init__.py:76-77`) only has:
|
||||
|
||||
```python
|
||||
class OmniEngineCoreOutput(EngineCoreOutput):
|
||||
pooling_output: dict[str, torch.Tensor] | None = None
|
||||
```
|
||||
|
||||
There is **no `multimodal_outputs` attribute** on `EngineCoreOutput` or `OmniEngineCoreOutput`. The `getattr` always returns `None`, making this method always return `None`.
|
||||
|
||||
This means the "extract from multimodal_outputs" fallback inside `_process_image_output`, `_process_latents_output`, `_process_audio_output`, and `_process_text_image_output` is **dead code**. These methods only execute when `eco.pooling_output is None`, and in that case they try `_extract_from_multimodal_outputs` which always returns `None` — so they effectively do nothing.
|
||||
|
||||
**Action**: Remove `_extract_from_multimodal_outputs` and the extraction fallback paths in all `_process_*_output` methods.
|
||||
|
||||
#### D6: `_process_pooling_output` — Broken for Actual Input Type (lines 446-455)
|
||||
|
||||
This method is hit as a fallback when `output_type` is unrecognized (e.g., `"token_ids"` for GLM-Image). It tries:
|
||||
|
||||
```python
|
||||
if not isinstance(eco.pooling_output, torch.Tensor):
|
||||
eco.pooling_output = torch.as_tensor(eco.pooling_output)
|
||||
```
|
||||
|
||||
But `OmniEngineCoreOutput.pooling_output` is typed as `dict[str, torch.Tensor] | None`. Calling `torch.as_tensor()` on a dict will raise a `TypeError`. This method is either dead or broken.
|
||||
|
||||
**Action**: Remove or fix `_process_pooling_output`. For `"token_ids"`, the data already flows correctly through `pooling_output` → `add_multimodal_tensor` in `process_outputs`.
|
||||
|
||||
#### D7: `_process_text_output` — No-Op (lines 442-444)
|
||||
|
||||
```python
|
||||
def _process_text_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""No-op; base processor will detokenize new_token_ids → text."""
|
||||
return
|
||||
```
|
||||
|
||||
This is a no-op method. It adds no value and can be replaced by a direct `pass` or `return` in the routing.
|
||||
|
||||
**Action**: Remove `_process_text_output` and replace the routing branch with a simple return/pass.
|
||||
|
||||
#### D8: Key Remapping Logic — Obscure and Partially Redundant (lines 64-69)
|
||||
|
||||
In `add_multimodal_tensor`, the code remaps keys `"model_outputs"` and `"hidden"` to `target_key` (which is `mm_type`):
|
||||
|
||||
```python
|
||||
if k == "model_outputs":
|
||||
k = target_key
|
||||
elif k == "hidden" and target_key != "hidden":
|
||||
k = target_key
|
||||
```
|
||||
|
||||
While this renaming does execute (AR runners produce `{"hidden": ...}` and generation runners produce `{"model_outputs": ...}`), downstream consumers like `serving_speech.py` still check for the **original key** `"model_outputs"`:
|
||||
|
||||
```python
|
||||
key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
|
||||
```
|
||||
|
||||
This means the renaming creates an inconsistency: the output processor renames `"model_outputs"` → `"audio"`, but the serving layer expects `"model_outputs"`. The system works only because the serving layer has a fallback, but the renaming adds confusion.
|
||||
|
||||
**Action**: Audit all downstream consumers and either commit to the renaming (update consumers) or remove it (keep original keys).
|
||||
|
||||
### 3.2 Effectively Unused Code Summary
|
||||
|
||||
| Code | Lines | Status | Reason |
|
||||
|------|-------|--------|--------|
|
||||
| `register_handler()` + `output_handlers` dict | 281-292, 277, 373-379 | **Dead** | Zero callers in entire codebase |
|
||||
| `_reqid_to_mm_type` dict | 278, 346-350 | **Dead** | Populated but never read |
|
||||
| `_process_text_image_output` | 383-384, 406-424 | **Dead** | No stage config uses "text+image"/"text,image"/"image+text" |
|
||||
| `"speech"` alias in routing | 387 | **Dead** | No config uses "speech", only "audio" |
|
||||
| `_extract_from_multimodal_outputs` | 457-469 | **Dead** | `eco.multimodal_outputs` attribute doesn't exist on `OmniEngineCoreOutput` |
|
||||
| Extraction fallbacks in `_process_*_output` | 399-440 (partial) | **Dead** | All call `_extract_from_multimodal_outputs` which always returns None |
|
||||
| `_process_pooling_output` | 446-455 | **Broken** | `torch.as_tensor()` on a dict raises TypeError |
|
||||
| `_process_text_output` | 442-444 | **No-op** | Empty method body |
|
||||
| Key remapping (`"model_outputs"` → mm_type) | 64-69 | **Inconsistent** | Renames keys that consumers still expect under original names |
|
||||
|
||||
**Estimated dead/broken code**: ~90 lines out of 469 (19% of the file).
|
||||
|
||||
### 3.3 What Actually Executes
|
||||
|
||||
After removing dead code, the **actual live data flow** is remarkably simple:
|
||||
|
||||
```
|
||||
1. process_outputs() is called with a list of EngineCoreOutput
|
||||
2. For each eco:
|
||||
a. engine_core_output_type determines mm_type (from stage config: "latent", "audio", "text", "image", "token_ids")
|
||||
b. _route_and_normalize() is called but effectively does nothing useful
|
||||
(all _process_*_output methods only act when pooling_output is None,
|
||||
then try _extract_from_multimodal_outputs which always returns None)
|
||||
c. If eco.pooling_output is not None AND req_state has a detokenizer:
|
||||
- pooling_output (a dict) is passed to req_state.add_multimodal_tensor()
|
||||
- eco.pooling_output is set to None (forcing text path in base processor)
|
||||
3. Base vLLM OutputProcessor.process_outputs() handles text detokenization
|
||||
4. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors
|
||||
5. _new_completion_output() attaches mm_accumulated to CompletionOutput
|
||||
```
|
||||
|
||||
The entire routing layer (`_route_and_normalize` and all `_process_*_output` methods) is effectively **a no-op** in production — the real work happens in the `pooling_output` capture at lines 355-361 and the accumulation in `OmniRequestState`.
|
||||
|
||||
## 4. Proposed Design
|
||||
|
||||
### 4.1 Core Abstractions
|
||||
|
||||
#### ModalityHandler (Protocol)
|
||||
|
||||
```python
|
||||
class ModalityHandler(Protocol):
|
||||
"""Handles extraction, normalization, and accumulation for a single output modality."""
|
||||
|
||||
modality: str # e.g., "image", "audio", "video", "latent"
|
||||
|
||||
def extract(self, eco: EngineCoreOutput) -> torch.Tensor | dict | None:
|
||||
"""Extract this modality's data from an EngineCoreOutput."""
|
||||
...
|
||||
|
||||
def accumulate(self, existing: Any, incoming: Any) -> Any:
|
||||
"""Accumulate incremental outputs (e.g., streaming audio chunks)."""
|
||||
...
|
||||
|
||||
def consolidate(self, accumulated: Any) -> Any:
|
||||
"""Finalize accumulated data (e.g., torch.cat for tensors)."""
|
||||
...
|
||||
|
||||
def validate(self, output: Any) -> bool:
|
||||
"""Optional validation of the final output."""
|
||||
...
|
||||
```
|
||||
|
||||
#### ModalityRegistry
|
||||
|
||||
```python
|
||||
class ModalityRegistry:
|
||||
"""Central registry mapping modality names to their handlers."""
|
||||
|
||||
def register(self, handler: ModalityHandler) -> None: ...
|
||||
def get(self, modality: str) -> ModalityHandler | None: ...
|
||||
def list_modalities(self) -> list[str]: ...
|
||||
```
|
||||
|
||||
#### OutputRouter
|
||||
|
||||
```python
|
||||
class OutputRouter:
|
||||
"""Routes EngineCoreOutput to one or more ModalityHandlers."""
|
||||
|
||||
def __init__(self, registry: ModalityRegistry): ...
|
||||
|
||||
def route(self, eco: EngineCoreOutput, output_type: str) -> list[tuple[str, ModalityHandler]]:
|
||||
"""Parse output_type and return list of (modality, handler) pairs.
|
||||
|
||||
Supports composite types like "text+video+image" by splitting
|
||||
on '+' or ',' and looking up each component in the registry.
|
||||
"""
|
||||
...
|
||||
```
|
||||
|
||||
### 4.2 Architecture
|
||||
|
||||
```
|
||||
┌────────────────────────────────────────────────────────────────────────┐
|
||||
│ MultimodalOutputProcessor (refactored) │
|
||||
│ │
|
||||
│ ┌──────────────┐ ┌──────────────┐ ┌───────────────────────┐ │
|
||||
│ │ OutputRouter │────▶│ ModalityReg │────▶│ ModalityHandler[] │ │
|
||||
│ │ │ │ │ │ - TextHandler │ │
|
||||
│ │ "text+video" │ │ "text" → H1 │ │ - ImageHandler │ │
|
||||
│ │ ↓ │ │ "video" → H2 │ │ - VideoHandler │ │
|
||||
│ │ ["text", │ │ "image" → H3 │ │ - AudioHandler │ │
|
||||
│ │ "video"] │ │ "audio" → H4 │ │ - LatentHandler │ │
|
||||
│ └──────────────┘ │ "video" → H5 │ │ - (custom handlers) │ │
|
||||
│ │ "latent"→ H6 │ └───────────────────────┘ │
|
||||
│ └──────────────┘ │
|
||||
│ │
|
||||
│ ┌─────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ OmniRequestState (refactored) │ │
|
||||
│ │ mm_accumulators: dict[str, ModalityAccumulator] │ │
|
||||
│ │ - Each accumulator uses its handler's accumulate/consolidate │ │
|
||||
│ └─────────────────────────────────────────────────────────────────┘ │
|
||||
└────────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### 4.3 Built-in Handlers
|
||||
|
||||
| Handler | Modality | Extract Keys | Accumulation | Consolidation |
|
||||
|---------|----------|-------------|--------------|---------------|
|
||||
| `TextHandler` | `text` | (none — handled by base vLLM detokenizer) | N/A | N/A |
|
||||
| `ImageHandler` | `image` | `image`, `images`, `pixel_values`, `pixels` | List append | `torch.cat(dim=0)` |
|
||||
| `VideoHandler` | `video` | `video`, `videos`, `frames`, `video_frames` | List append | `torch.cat(dim=0)` or `torch.stack` |
|
||||
| `AudioHandler` | `audio` | `audio`, `audios`, `wav`, `waveform`, `audio_pcm`, `pcm` | List append | `torch.cat(dim=-1)` (last dim for variable-length) |
|
||||
| `LatentHandler` | `latent` | `latent`, `latents`, `z`, `posterior` | List append | `torch.cat(dim=0)` |
|
||||
|
||||
### 4.4 Composite Output Types
|
||||
|
||||
The `OutputRouter` supports composite types by splitting on delimiters:
|
||||
|
||||
```python
|
||||
def route(self, eco, output_type):
|
||||
modalities = re.split(r'[+,]', output_type)
|
||||
handlers = []
|
||||
for m in modalities:
|
||||
m = m.strip().lower()
|
||||
handler = self.registry.get(m)
|
||||
if handler:
|
||||
handlers.append((m, handler))
|
||||
else:
|
||||
logger.warning("No handler for modality: %s", m)
|
||||
return handlers
|
||||
```
|
||||
|
||||
This naturally supports:
|
||||
- `"text+video+image"` → `[TextHandler, VideoHandler, ImageHandler]`
|
||||
- `"text+audio"` → `[TextHandler, AudioHandler]`
|
||||
- `"image"` → `[ImageHandler]`
|
||||
|
||||
### 4.5 Refactored OmniRequestState
|
||||
|
||||
```python
|
||||
class OmniRequestState(RequestState):
|
||||
def __init__(self, *args, registry: ModalityRegistry, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._registry = registry
|
||||
self._accumulators: dict[str, Any] = {}
|
||||
|
||||
def add_output(self, modality: str, payload: Any) -> None:
|
||||
handler = self._registry.get(modality)
|
||||
if handler is None:
|
||||
return
|
||||
existing = self._accumulators.get(modality)
|
||||
self._accumulators[modality] = handler.accumulate(existing, payload)
|
||||
|
||||
def consolidate(self) -> dict[str, Any]:
|
||||
result = {}
|
||||
for modality, data in self._accumulators.items():
|
||||
handler = self._registry.get(modality)
|
||||
if handler:
|
||||
result[modality] = handler.consolidate(data)
|
||||
else:
|
||||
result[modality] = data
|
||||
return result
|
||||
```
|
||||
|
||||
### 4.6 Model Output Contract
|
||||
|
||||
To standardize the interface between models and the output processor, we propose that models producing multimodal outputs should return an `OmniOutput` (already defined in `model_executor/models/output_templates.py`) with a well-defined `multimodal_outputs` dict:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class OmniOutput:
|
||||
text_hidden_states: torch.Tensor
|
||||
multimodal_outputs: dict[str, torch.Tensor | dict]
|
||||
# Keys should match registered modality names: "image", "video", "audio", etc.
|
||||
```
|
||||
|
||||
## 5. Migration Plan
|
||||
|
||||
### Phase 0: Simplify — Remove Dead Code and Clarify Live Paths
|
||||
|
||||
This phase is a prerequisite for any further refactoring. It reduces the file from 469 lines to ~300 and makes the actual data flow visible.
|
||||
|
||||
**Step 0.1**: Remove dead code (all items from Section 3.2):
|
||||
|
||||
- Delete `register_handler()`, `self.output_handlers`, and the custom handler dispatch block in `_route_and_normalize`
|
||||
- Delete `self._reqid_to_mm_type` and all code that populates it
|
||||
- Delete `_process_text_image_output` and its routing branch
|
||||
- Delete `_extract_from_multimodal_outputs` and all extraction fallback paths
|
||||
- Delete `_process_pooling_output`
|
||||
- Delete `_process_text_output`
|
||||
- Remove `"speech"` from the audio routing match
|
||||
|
||||
**Step 0.2**: Simplify `_route_and_normalize` to reflect what actually executes:
|
||||
|
||||
```python
|
||||
def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
|
||||
# Currently a no-op for all production paths.
|
||||
# Routing is actually handled by the pooling_output capture in process_outputs().
|
||||
# Kept as an extension point for future modality-specific normalization.
|
||||
pass
|
||||
```
|
||||
|
||||
Or, if we want to keep it as a dispatch point for future use, collapse into:
|
||||
|
||||
```python
|
||||
def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
|
||||
output_type = (self.engine_core_output_type or "").lower()
|
||||
if not output_type or output_type in ("text", "latent", "audio", "image", "token_ids"):
|
||||
return # All handled by pooling_output capture in process_outputs()
|
||||
```
|
||||
|
||||
**Step 0.3**: Audit and resolve the key remapping inconsistency (D8):
|
||||
|
||||
- Option A: Remove renaming, keep original keys (`"model_outputs"`, `"hidden"`), update the proposed ModalityHandler to standardize keys at extraction time
|
||||
- Option B: Commit to renaming, update `serving_speech.py` and other consumers to use mm_type keys
|
||||
|
||||
**Step 0.4**: Simplify `make_request_output` — the current override largely duplicates base class logic for streaming/DELTA handling. If the only additions are `_consolidate_multimodal_tensors()` and `_new_completion_output()`, consider hooking these without copying the full method.
|
||||
|
||||
**Estimated result**: ~300 lines (from 469), with clear separation between "what the code does" and "extension points for the future".
|
||||
|
||||
### Phase 1: Introduce Abstractions (Non-Breaking)
|
||||
|
||||
1. Create `ModalityHandler` protocol and `ModalityRegistry` class
|
||||
2. Implement built-in handlers (`TextHandler`, `ImageHandler`, `AudioHandler`, `LatentHandler`)
|
||||
3. Create `OutputRouter` with delimiter-based composite type parsing
|
||||
4. Place these in a new module: `vllm_omni/engine/output_handlers.py`
|
||||
|
||||
### Phase 2: Refactor MultimodalOutputProcessor
|
||||
|
||||
1. Replace the simplified `_route_and_normalize` with `OutputRouter.route()`
|
||||
2. Introduce handler-based dispatch for `pooling_output` extraction and normalization
|
||||
3. The existing `register_handler()` API concept is naturally replaced by `ModalityRegistry.register()`
|
||||
|
||||
### Phase 3: Refactor OmniRequestState
|
||||
|
||||
1. Replace `mm_accumulated` dict with per-modality accumulators
|
||||
2. Replace `add_multimodal_tensor` with `add_output(modality, payload)`
|
||||
3. Replace `_consolidate_multimodal_tensors` with handler-driven `consolidate()`
|
||||
|
||||
### Phase 4: Add New Modalities
|
||||
|
||||
1. Add `VideoHandler` for video output support
|
||||
2. Test composite types: `text+video`, `text+image+video`, etc.
|
||||
3. Integrate with upcoming DeepSeek model support
|
||||
|
||||
## 6. Backward Compatibility
|
||||
|
||||
- `register_handler()` was never called externally — removing it has zero backward compatibility impact
|
||||
- Existing stage configs with `engine_output_type` strings (e.g., `"audio"`, `"latent"`, `"image"`) continue to work unchanged
|
||||
- The `output_type` field on `EngineCoreOutput` is unchanged
|
||||
- Phase 0 (dead code removal) is a pure deletion with no behavioral change, since none of the removed paths were reachable
|
||||
|
||||
## 7. Testing Strategy
|
||||
|
||||
| Test Type | Phase | Scope |
|
||||
|-----------|-------|-------|
|
||||
| Regression (existing) | Phase 0 | Run `test_outputs.py` and all E2E tests after dead code removal — must pass unchanged |
|
||||
| Smoke test | Phase 0 | Verify Qwen3-TTS (audio), GLM-Image (token_ids), Bagel (text+image) pipelines still work |
|
||||
| Unit tests | Phase 1 | Each `ModalityHandler` (extract, accumulate, consolidate) |
|
||||
| Integration tests | Phase 2 | `OutputRouter` with composite types |
|
||||
| E2E tests | Phase 2-3 | Full pipeline with existing models (Qwen3-TTS, GLM-Image, etc.) |
|
||||
| New modality tests | Phase 4 | Video output handling, `text+video+image` combinations |
|
||||
|
||||
## 8. Open Questions
|
||||
|
||||
1. Should we support dynamic handler priority (e.g., if both `image` and `latent` are present, which takes precedence for `pooling_output`)?
|
||||
2. Should the modality registry be global (singleton) or per-engine-instance?
|
||||
3. How should error handling work when one handler in a composite route fails but others succeed?
|
||||
4. Should we support async accumulation for streaming scenarios?
|
||||
@@ -1,280 +0,0 @@
|
||||
# RFC-003: Model Runner Clearance — Extracting Model-Specific Logic
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes extracting model-specific logic currently embedded in the model runner classes (`OmniGPUModelRunner`, `GPUARModelRunner`, `GPUGenerationModelRunner`) into well-defined extension points and per-model adapter classes. The goal is to make the model runner a **model-agnostic execution framework** that delegates model-specific behavior through a clean plugin/adapter interface.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current State
|
||||
|
||||
`OmniGPUModelRunner` (1,292 lines in `vllm_omni/worker/gpu_model_runner.py`) is the core execution class that inherits from `vllm.v1.worker.gpu_model_runner.GPUModelRunner`. Over time, it has accumulated significant model-specific logic that should not live in a generic runner:
|
||||
|
||||
#### M-RoPE / Positional Encoding (Qwen2-VL, HunYuan-VL, GLM-Image)
|
||||
|
||||
- `_init_mrope_positions()` (lines 99-157): Extracts `image_grid_thw`, `video_grid_thw`, `second_per_grid_ts`, `audio_feature_lengths`, `use_audio_in_video` — all Qwen-specific metadata.
|
||||
- `_calc_mrope_positions()` (line 159): Delegates to upstream, then calls `_fixup_precomputed_mrope_decode_positions` for GLM-Image.
|
||||
- `_fixup_precomputed_mrope_decode_positions()` (lines 180-222): Entirely GLM-Image-specific — handles 2D spatial decode positions for image generation models with `precomputed_mrope_decode = True`.
|
||||
- XD-RoPE handling via `uses_xdrope_dim` for HunYuan-VL.
|
||||
|
||||
#### Talker MTP (Qwen3-TTS, Qwen3-Omni)
|
||||
|
||||
- `load_model()` (lines 71-97): Contains TTS-specific buffer allocation (`talker_mtp_input_ids`, `talker_mtp_inputs_embeds`, `last_talker_hidden`, `text_step`) with a `TODO` comment: "move this model specific logic to a separate class".
|
||||
- `_talker_mtp_forward()` (lines 1215-1246): Entirely Qwen3-TTS/Omni-specific — runs the talker MTP forward pass with CUDA graph wrapping.
|
||||
- `_preprocess()` (lines 1039-1213): Contains ~50 lines of talker MTP path (lines 1154-1204) interleaved with generic preprocessing.
|
||||
|
||||
#### MiMo-Audio
|
||||
|
||||
- `_maybe_attach_mimo_audio_req_infos()` (lines 1016-1037): Checks `self.model.__class__.__name__ == "MiMoAudioForConditionalGeneration"` — a direct class name check that couples the runner to a specific model.
|
||||
|
||||
#### Additional Information / Postprocess
|
||||
|
||||
- `_process_additional_information_updates()` (lines 937-969): Checks `self.model.has_postprocess` and calls `self.model.postprocess()` — a generic hook that's only used by specific models.
|
||||
- `_preprocess()` similarly checks `self.model.has_preprocess` and calls `self.model.preprocess()`.
|
||||
|
||||
### Problems
|
||||
|
||||
| Problem | Impact |
|
||||
|---------|--------|
|
||||
| **Tight coupling** | Adding a new model with unique positional encoding or pre/post-processing requires modifying the runner itself |
|
||||
| **Code readability** | The 1,292-line runner mixes framework logic with model-specific branches, making it hard to understand the core execution flow |
|
||||
| **Testing difficulty** | Model-specific logic cannot be unit-tested in isolation from the full runner |
|
||||
| **Rebase fragility** | Model-specific code in the runner is a frequent source of conflicts during upstream vLLM rebases |
|
||||
| **Platform duplication** | NPU, XPU, and ROCm runners must replicate or work around model-specific logic from the GPU runner |
|
||||
|
||||
## 3. Proposed Design
|
||||
|
||||
### 3.1 Model Runner Adapter Protocol
|
||||
|
||||
```python
|
||||
class ModelRunnerAdapter(Protocol):
|
||||
"""Extension point for model-specific behavior in the runner."""
|
||||
|
||||
def on_model_loaded(self, runner: "OmniGPUModelRunner", model: nn.Module) -> None:
|
||||
"""Called after model is loaded. Allocate model-specific buffers."""
|
||||
...
|
||||
|
||||
def init_request_positions(
|
||||
self, runner: "OmniGPUModelRunner", req_state: CachedRequestState
|
||||
) -> None:
|
||||
"""Initialize custom positional encoding for a new request (e.g., M-RoPE, XD-RoPE)."""
|
||||
...
|
||||
|
||||
def fixup_decode_positions(
|
||||
self, runner: "OmniGPUModelRunner", scheduler_output: SchedulerOutput
|
||||
) -> None:
|
||||
"""Apply model-specific position fixups for decode (e.g., GLM-Image 2D positions)."""
|
||||
...
|
||||
|
||||
def preprocess(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
input_ids: torch.Tensor,
|
||||
inputs_embeds: torch.Tensor | None,
|
||||
scheduler_output: SchedulerOutput,
|
||||
num_scheduled_tokens_np: np.ndarray,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None, dict]:
|
||||
"""Model-specific preprocessing before forward pass."""
|
||||
...
|
||||
|
||||
def postprocess(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
hidden_states: torch.Tensor,
|
||||
multimodal_outputs: Any,
|
||||
scheduler_output: SchedulerOutput,
|
||||
) -> None:
|
||||
"""Model-specific postprocessing after forward pass."""
|
||||
...
|
||||
|
||||
def extract_outputs(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
hidden_states: torch.Tensor | OmniOutput,
|
||||
) -> tuple[torch.Tensor, dict]:
|
||||
"""Extract text hidden states and multimodal outputs."""
|
||||
...
|
||||
|
||||
def enrich_request_info(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
req_id: str,
|
||||
req_state: CachedRequestState,
|
||||
req_infos: dict | None,
|
||||
) -> dict | None:
|
||||
"""Enrich per-request info before forward (e.g., MiMo-Audio mm_features)."""
|
||||
...
|
||||
```
|
||||
|
||||
### 3.2 Adapter Registry
|
||||
|
||||
```python
|
||||
class ModelAdapterRegistry:
|
||||
"""Maps model architecture names to their runner adapters."""
|
||||
|
||||
_adapters: dict[str, type[ModelRunnerAdapter]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, model_cls_name: str):
|
||||
def decorator(adapter_cls):
|
||||
cls._adapters[model_cls_name] = adapter_cls
|
||||
return adapter_cls
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def get_adapter(cls, model: nn.Module) -> ModelRunnerAdapter:
|
||||
name = model.__class__.__name__
|
||||
adapter_cls = cls._adapters.get(name, DefaultModelRunnerAdapter)
|
||||
return adapter_cls()
|
||||
```
|
||||
|
||||
### 3.3 Concrete Adapters
|
||||
|
||||
#### Default Adapter
|
||||
|
||||
```python
|
||||
class DefaultModelRunnerAdapter:
|
||||
"""No-op adapter for models without special requirements."""
|
||||
def on_model_loaded(self, runner, model): pass
|
||||
def init_request_positions(self, runner, req_state): pass
|
||||
def fixup_decode_positions(self, runner, scheduler_output): pass
|
||||
def preprocess(self, runner, input_ids, inputs_embeds, so, nst):
|
||||
return input_ids, inputs_embeds, {}
|
||||
def postprocess(self, runner, hs, mm, so): pass
|
||||
def extract_outputs(self, runner, hs):
|
||||
return (hs, {}) if isinstance(hs, torch.Tensor) else (hs[0], {})
|
||||
def enrich_request_info(self, runner, req_id, req_state, req_infos):
|
||||
return req_infos
|
||||
```
|
||||
|
||||
#### Qwen3-Omni / TTS Adapter
|
||||
|
||||
```python
|
||||
@ModelAdapterRegistry.register("Qwen3OmniMoeForConditionalGeneration")
|
||||
@ModelAdapterRegistry.register("Qwen3TTSForConditionalGeneration")
|
||||
class Qwen3TalkerAdapter(DefaultModelRunnerAdapter):
|
||||
def on_model_loaded(self, runner, model):
|
||||
# Allocate talker_mtp buffers, CUDAGraphWrapper, etc.
|
||||
...
|
||||
|
||||
def init_request_positions(self, runner, req_state):
|
||||
# M-RoPE with image_grid_thw, video_grid_thw, audio_feature_lengths
|
||||
...
|
||||
|
||||
def preprocess(self, runner, input_ids, inputs_embeds, so, nst):
|
||||
# Talker MTP forward for decode requests
|
||||
...
|
||||
```
|
||||
|
||||
#### GLM-Image Adapter
|
||||
|
||||
```python
|
||||
@ModelAdapterRegistry.register("GLMImageForConditionalGeneration")
|
||||
class GLMImageAdapter(DefaultModelRunnerAdapter):
|
||||
def fixup_decode_positions(self, runner, scheduler_output):
|
||||
# 2D spatial decode position fixup for precomputed_mrope_decode
|
||||
...
|
||||
```
|
||||
|
||||
#### MiMo-Audio Adapter
|
||||
|
||||
```python
|
||||
@ModelAdapterRegistry.register("MiMoAudioForConditionalGeneration")
|
||||
class MiMoAudioAdapter(DefaultModelRunnerAdapter):
|
||||
def enrich_request_info(self, runner, req_id, req_state, req_infos):
|
||||
# Attach mm_features and req_id
|
||||
...
|
||||
```
|
||||
|
||||
### 3.4 Refactored OmniGPUModelRunner
|
||||
|
||||
After extraction, the runner becomes much simpler:
|
||||
|
||||
```python
|
||||
class OmniGPUModelRunner(GPUModelRunner):
|
||||
def load_model(self, *args, **kwargs):
|
||||
super().load_model(*args, **kwargs)
|
||||
self._adapter = ModelAdapterRegistry.get_adapter(self.model)
|
||||
self._adapter.on_model_loaded(self, self.model)
|
||||
|
||||
def _update_states(self, scheduler_output):
|
||||
super()._update_states(scheduler_output)
|
||||
for req_id, req_state in ...:
|
||||
self._adapter.init_request_positions(self, req_state)
|
||||
|
||||
def _calc_mrope_positions(self, scheduler_output):
|
||||
super()._calc_mrope_positions(scheduler_output)
|
||||
self._adapter.fixup_decode_positions(self, scheduler_output)
|
||||
|
||||
def _preprocess(self, scheduler_output, num_input_tokens, ...):
|
||||
# Generic preprocessing
|
||||
...
|
||||
input_ids, inputs_embeds, extra = self._adapter.preprocess(...)
|
||||
return ...
|
||||
|
||||
def extract_multimodal_outputs(self, hidden_states):
|
||||
return self._adapter.extract_outputs(self, hidden_states)
|
||||
```
|
||||
|
||||
### 3.5 File Structure
|
||||
|
||||
```
|
||||
vllm_omni/worker/
|
||||
├── gpu_model_runner.py # Cleaned OmniGPUModelRunner (~600 lines)
|
||||
├── gpu_ar_model_runner.py # Unchanged
|
||||
├── gpu_generation_model_runner.py # Unchanged
|
||||
├── adapters/ # NEW: Model-specific adapters
|
||||
│ ├── __init__.py # ModelAdapterRegistry
|
||||
│ ├── base.py # ModelRunnerAdapter protocol + DefaultAdapter
|
||||
│ ├── qwen3_talker.py # Qwen3-Omni / TTS adapter
|
||||
│ ├── glm_image.py # GLM-Image adapter
|
||||
│ ├── mimo_audio.py # MiMo-Audio adapter
|
||||
│ └── hunyuan_vl.py # HunYuan-VL (XD-RoPE) adapter
|
||||
```
|
||||
|
||||
## 4. Migration Plan
|
||||
|
||||
| Phase | Scope | Risk |
|
||||
|-------|-------|------|
|
||||
| Phase 1 | Define `ModelRunnerAdapter` protocol and `ModelAdapterRegistry` | Low — additive only |
|
||||
| Phase 2 | Extract M-RoPE and positional encoding into adapters | Medium — touches core init path |
|
||||
| Phase 3 | Extract talker MTP logic into `Qwen3TalkerAdapter` | Medium — performance-critical path |
|
||||
| Phase 4 | Extract MiMo-Audio logic | Low — isolated code |
|
||||
| Phase 5 | Extract GLM-Image precomputed position fixup | Low — isolated code |
|
||||
| Phase 6 | Clean up `_preprocess` to use adapter dispatch | High — central hot path |
|
||||
| Phase 7 | Port platform adapters (NPU, XPU) to use the same adapter system | Medium |
|
||||
|
||||
Each phase should be a separate PR with targeted tests.
|
||||
|
||||
## 5. Testing Strategy
|
||||
|
||||
| Test | Description |
|
||||
|------|-------------|
|
||||
| Adapter unit tests | Test each adapter's methods independently with mock runner |
|
||||
| Registry tests | Verify correct adapter resolution by model class name |
|
||||
| Integration tests | Run existing E2E tests (Qwen3-TTS, GLM-Image, MiMo-Audio) |
|
||||
| Performance regression | Benchmark to ensure adapter indirection doesn't add measurable overhead |
|
||||
| Platform compatibility | Verify NPU/XPU runners work with the adapter system |
|
||||
|
||||
## 6. Benefits
|
||||
|
||||
| Benefit | Impact |
|
||||
|---------|--------|
|
||||
| **Code clarity** | Runner reduced from ~1,300 to ~600 lines, with model logic isolated |
|
||||
| **Extensibility** | New models only need an adapter class, not runner modifications |
|
||||
| **Testability** | Model-specific logic can be unit-tested independently |
|
||||
| **Rebase safety** | Upstream changes to `GPUModelRunner` are less likely to conflict with model-specific adapters |
|
||||
| **Platform parity** | NPU/XPU runners can share adapters or provide platform-specific overrides |
|
||||
|
||||
## 7. Open Questions
|
||||
|
||||
1. Should adapters be auto-discovered from the model class or explicitly configured in stage configs?
|
||||
2. Should the adapter have access to the full runner state or a restricted view (for safety)?
|
||||
3. How do we handle models that need multiple adapter behaviors composed together?
|
||||
4. Should `GPUARModelRunner` and `GPUGenerationModelRunner` also use the adapter pattern?
|
||||
@@ -1,275 +0,0 @@
|
||||
# RFC-004: Long Function Decomposition for Improved Readability
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes systematically decomposing the long functions in vllm-omni's engine and worker modules. Several functions span hundreds of lines, making them difficult to read, test, and maintain. We identify the worst offenders, propose decomposition strategies, and define guidelines to prevent future regression.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Identified Long Functions
|
||||
|
||||
A survey of the vllm-omni codebase reveals the following functions that significantly exceed reasonable length limits:
|
||||
|
||||
| Function | File | Lines | Description |
|
||||
|----------|------|-------|-------------|
|
||||
| `_stage_worker_async` | `entrypoints/omni_stage.py` | ~400 | Async stage worker: device setup, engine init, main loop, profiler handling, output processing |
|
||||
| `_stage_worker` | `entrypoints/omni_stage.py` | ~400 | Sync stage worker: near-identical structure to async version |
|
||||
| `handle_profiler_task_local` | `entrypoints/omni_stage.py` (nested) | ~270 | Profiler start/stop/dump handling for LLM and diffusion engines |
|
||||
| `_dummy_run` | `worker/gpu_model_runner.py` | ~320 | CUDA graph warm-up/capture with multiple execution paths |
|
||||
| `_preprocess` | `worker/gpu_model_runner.py` | ~175 | Input preparation with multiple branches for mm/embeds/text/preprocess |
|
||||
| `process_inputs` | `engine/input_processor.py` | ~193 | Multimodal input processing with tokenization, embeddings, LoRA |
|
||||
| `create_model_config` | `engine/arg_utils.py` | ~146 | Config construction (duplicated for sync/async) |
|
||||
| `_sequential_init_lock` | `entrypoints/omni_stage.py` | ~146 | Lock management for sequential GPU initialization |
|
||||
| `process_engine_inputs` | `entrypoints/omni_stage.py` (class method) | ~50 | Input wiring between stages |
|
||||
|
||||
### Problems Caused by Long Functions
|
||||
|
||||
| Problem | Impact |
|
||||
|---------|--------|
|
||||
| **Cognitive load** | Developers must hold 200-400 lines of context in their head to understand control flow |
|
||||
| **Testing difficulty** | Cannot unit-test internal logic phases; must test the entire function end-to-end |
|
||||
| **Merge conflicts** | Long functions increase the probability of merge conflicts during rebases |
|
||||
| **Code duplication** | `_stage_worker` and `_stage_worker_async` are near-duplicates (~800 lines combined) with minor async/await differences |
|
||||
| **Debugging difficulty** | Stack traces point to a single long function; hard to isolate which phase failed |
|
||||
|
||||
## 3. Decomposition Proposals
|
||||
|
||||
### 3.1 `_stage_worker` and `_stage_worker_async` (omni_stage.py)
|
||||
|
||||
These two functions share ~90% of their logic. The proposed decomposition:
|
||||
|
||||
#### Phase Extraction
|
||||
|
||||
```python
|
||||
# New helper functions (shared between sync and async):
|
||||
|
||||
def _init_stage_environment(stage_payload: dict) -> StageContext:
|
||||
"""Phase 1: Environment setup — ZMQ, device mapping, plugins."""
|
||||
...
|
||||
|
||||
def _init_stage_engine(
|
||||
model: str, stage_payload: dict, stage_ctx: StageContext
|
||||
) -> LLMEngine | OmniDiffusion:
|
||||
"""Phase 2: Engine initialization with sequential lock."""
|
||||
...
|
||||
|
||||
def _init_stage_connectors(
|
||||
stage_id: str, connectors_config: dict
|
||||
) -> dict[tuple[str, str], OmniConnectorBase]:
|
||||
"""Phase 3: Build OmniConnectors."""
|
||||
...
|
||||
|
||||
def _handle_profiler_task(
|
||||
task_type: OmniStageTaskType, engine: Any, stage_id: str, stage_type: str
|
||||
) -> dict:
|
||||
"""Phase 4: Profiler start/stop/dump (extracted from nested function)."""
|
||||
...
|
||||
|
||||
def _process_single_batch(
|
||||
batch: list[dict], engine: Any, stage_ctx: StageContext
|
||||
) -> list[dict]:
|
||||
"""Phase 5: Execute a single batch through the engine."""
|
||||
...
|
||||
|
||||
def _format_and_send_outputs(
|
||||
outputs: list, out_q: Queue, stage_ctx: StageContext
|
||||
) -> None:
|
||||
"""Phase 6: Format outputs and send via IPC."""
|
||||
...
|
||||
```
|
||||
|
||||
#### Sync/Async Unification
|
||||
|
||||
```python
|
||||
def _stage_worker(model, stage_payload, in_q, out_q, ...):
|
||||
ctx = _init_stage_environment(stage_payload)
|
||||
engine = _init_stage_engine(model, stage_payload, ctx)
|
||||
connectors = _init_stage_connectors(ctx.stage_id, ctx.connectors_config)
|
||||
|
||||
while True:
|
||||
batch = _collect_batch(in_q, ctx.max_batch_size, ctx.batch_timeout)
|
||||
if _is_shutdown(batch):
|
||||
break
|
||||
if _is_profiler_task(batch):
|
||||
result = _handle_profiler_task(batch.task_type, engine, ...)
|
||||
out_q.put(result)
|
||||
continue
|
||||
outputs = _process_single_batch(batch, engine, ctx)
|
||||
_format_and_send_outputs(outputs, out_q, ctx)
|
||||
|
||||
async def _stage_worker_async(model, stage_payload, in_q, out_q, ...):
|
||||
ctx = _init_stage_environment(stage_payload)
|
||||
engine = _init_stage_engine(model, stage_payload, ctx) # async variant
|
||||
connectors = _init_stage_connectors(ctx.stage_id, ctx.connectors_config)
|
||||
|
||||
while True:
|
||||
batch = await _collect_batch_async(in_q, ...)
|
||||
# ... same structure, async versions of batch processing
|
||||
```
|
||||
|
||||
### 3.2 `_dummy_run` (gpu_model_runner.py)
|
||||
|
||||
This function handles CUDA graph capture, warm-up, and profiling. Decompose into:
|
||||
|
||||
```python
|
||||
def _dummy_run(self, num_tokens, ...):
|
||||
num_scheduled_tokens, attn_metadata = self._prepare_dummy_batch(num_tokens, ...)
|
||||
hidden_states = self._execute_dummy_forward(
|
||||
num_scheduled_tokens, attn_metadata, ...
|
||||
)
|
||||
self._run_dummy_drafter(num_tokens, is_graph_capturing, ...)
|
||||
self._finalize_dummy_run(hidden_states, num_scheduled_tokens, ...)
|
||||
|
||||
def _prepare_dummy_batch(self, num_tokens, ...):
|
||||
"""Set up fake batch data, attention metadata, LoRA if needed."""
|
||||
...
|
||||
|
||||
def _execute_dummy_forward(self, num_scheduled_tokens, attn_metadata, ...):
|
||||
"""Run the model forward pass for CUDA graph capture or warm-up."""
|
||||
...
|
||||
|
||||
def _run_dummy_drafter(self, num_tokens, is_graph_capturing, ...):
|
||||
"""Run speculative decoding drafter dummy pass if configured."""
|
||||
...
|
||||
|
||||
def _finalize_dummy_run(self, hidden_states, num_scheduled_tokens, ...):
|
||||
"""Register NVTX hooks, run EPLB step, compute logit indices."""
|
||||
...
|
||||
```
|
||||
|
||||
### 3.3 `_preprocess` (gpu_model_runner.py)
|
||||
|
||||
The current function has 4 major branches. Decompose into:
|
||||
|
||||
```python
|
||||
def _preprocess(self, scheduler_output, num_input_tokens, intermediate_tensors=None):
|
||||
input_ids, inputs_embeds, model_kwargs, ec_connector_output = (
|
||||
self._prepare_model_inputs(scheduler_output, num_input_tokens)
|
||||
)
|
||||
positions = self._compute_positions(num_input_tokens)
|
||||
intermediate_tensors = self._prepare_intermediate_tensors(
|
||||
num_input_tokens, intermediate_tensors
|
||||
)
|
||||
self._apply_model_specific_preprocessing(
|
||||
scheduler_output, input_ids, inputs_embeds, model_kwargs
|
||||
)
|
||||
return input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ec_connector_output
|
||||
|
||||
def _prepare_model_inputs(self, scheduler_output, num_input_tokens):
|
||||
"""Branch: mm inputs / prompt embeds / has_preprocess / text-only."""
|
||||
...
|
||||
|
||||
def _compute_positions(self, num_input_tokens):
|
||||
"""Select M-RoPE / XD-RoPE / standard positions."""
|
||||
...
|
||||
|
||||
def _apply_model_specific_preprocessing(self, ...):
|
||||
"""Overlay prompt_embeds, run has_preprocess, talker MTP."""
|
||||
...
|
||||
```
|
||||
|
||||
### 3.4 `process_inputs` (input_processor.py)
|
||||
|
||||
Decompose into phases:
|
||||
|
||||
```python
|
||||
def process_inputs(self, request_id, prompt, params):
|
||||
preprocessed = self._preprocess_prompt(request_id, prompt, params)
|
||||
tokenized = self._tokenize_and_validate(preprocessed, params)
|
||||
mm_data = self._extract_multimodal_data(tokenized, params)
|
||||
prompt_embeds = self._extract_prompt_embeds(tokenized, params)
|
||||
additional_info = self._extract_additional_information(tokenized, params)
|
||||
return self._build_engine_core_request(
|
||||
request_id, tokenized, mm_data, prompt_embeds, additional_info, params
|
||||
)
|
||||
```
|
||||
|
||||
### 3.5 `create_model_config` (arg_utils.py)
|
||||
|
||||
This function is duplicated between `OmniEngineArgs` and `AsyncOmniEngineArgs` (~260 lines combined). Decompose and deduplicate:
|
||||
|
||||
```python
|
||||
class OmniEngineArgs:
|
||||
def create_model_config(self):
|
||||
return _create_model_config_common(self, is_async=False)
|
||||
|
||||
class AsyncOmniEngineArgs:
|
||||
def create_model_config(self):
|
||||
return _create_model_config_common(self, is_async=True)
|
||||
|
||||
def _create_model_config_common(args, is_async: bool):
|
||||
"""Shared config construction logic."""
|
||||
base_config = _build_base_config(args)
|
||||
_apply_omni_overrides(base_config, args)
|
||||
_validate_config(base_config)
|
||||
return base_config
|
||||
```
|
||||
|
||||
## 4. Guidelines for Future Code
|
||||
|
||||
To prevent regression, we propose the following guidelines:
|
||||
|
||||
### 4.1 Function Length Limits
|
||||
|
||||
| Severity | Line Count | Action |
|
||||
|----------|-----------|--------|
|
||||
| OK | < 50 lines | No action needed |
|
||||
| Warning | 50-100 lines | Consider decomposition |
|
||||
| Required | > 100 lines | Must decompose before merge |
|
||||
|
||||
### 4.2 Lint Rule
|
||||
|
||||
Add a custom lint check (or use `ruff` with `pylint.max-statements`) to flag functions exceeding the threshold:
|
||||
|
||||
```toml
|
||||
# pyproject.toml
|
||||
[tool.ruff.pylint]
|
||||
max-statements = 50
|
||||
```
|
||||
|
||||
### 4.3 Decomposition Principles
|
||||
|
||||
1. **Extract by phase**: Each phase of a multi-step function becomes its own method
|
||||
2. **Extract by branch**: If a function has N major branches, each branch becomes a helper
|
||||
3. **Extract by concern**: Separate I/O, computation, and error handling
|
||||
4. **Preserve hot paths**: Do not add unnecessary function call overhead in performance-critical paths (profile after decomposition)
|
||||
5. **Keep related logic together**: Extracted functions should be in the same module, not scattered
|
||||
|
||||
## 5. Implementation Plan
|
||||
|
||||
| Phase | Target | Functions | Est. Impact |
|
||||
|-------|--------|-----------|-------------|
|
||||
| Phase 1 | `omni_stage.py` | `_stage_worker`, `_stage_worker_async`, `handle_profiler_task_local` | ~800 lines → ~200 + shared helpers |
|
||||
| Phase 2 | `gpu_model_runner.py` | `_dummy_run`, `_preprocess` | ~500 lines → ~200 + helpers |
|
||||
| Phase 3 | `input_processor.py` | `process_inputs` | ~193 lines → ~80 + helpers |
|
||||
| Phase 4 | `arg_utils.py` | `create_model_config` (dedup) | ~260 lines → ~130 shared |
|
||||
| Phase 5 | Add lint rules | Configure ruff/pylint | Prevents regression |
|
||||
|
||||
## 6. Risk Mitigation
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|------------|
|
||||
| **Performance regression** | Benchmark critical paths (model runner `_preprocess`, `_dummy_run`) before and after |
|
||||
| **Behavioral changes** | Each decomposition should be a pure refactor — no logic changes |
|
||||
| **Merge conflicts** | Coordinate with ongoing PRs; complete in small incremental PRs |
|
||||
| **Variable scoping issues** | Use `dataclass` or `NamedTuple` for passing multi-value state between phases rather than long parameter lists |
|
||||
|
||||
## 7. Success Metrics
|
||||
|
||||
- No function in the codebase exceeds 100 lines (excluding comments/docstrings)
|
||||
- `_stage_worker` and `_stage_worker_async` share > 80% of logic via shared helpers
|
||||
- No performance regression in model runner benchmarks
|
||||
- Improved code review velocity for PRs touching refactored files
|
||||
|
||||
## 8. Open Questions
|
||||
|
||||
1. Should we introduce a `StageContext` dataclass to pass state between phases in `_stage_worker`, or use simpler parameter passing?
|
||||
2. Is the 100-line limit too aggressive for GPU kernel-related code paths where inlining matters?
|
||||
3. Should we use `ruff` max-statements or a custom pre-commit hook for enforcement?
|
||||
4. How do we handle the sync/async duplication — use a code generation approach or manual unification?
|
||||
@@ -1,277 +0,0 @@
|
||||
# RFC-005: Model Runner V2 Integration Roadmap
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
Upstream vLLM is developing a V2 model runner to replace the current V1 architecture (`vllm.v1.worker.gpu_model_runner.GPUModelRunner`). vllm-omni currently **explicitly disables** V2 model runner support with a fallback warning. This RFC outlines the adaptation work required to integrate the V2 model runner into vllm-omni, including identifying the "omni hooks" that must be ported, defining a phased migration plan, and establishing compatibility testing.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current State
|
||||
|
||||
Both `GPUARWorker` and `GPUGenerationWorker` contain identical guard code:
|
||||
|
||||
```python
|
||||
# vllm_omni/worker/gpu_ar_worker.py:96-99
|
||||
if self.use_v2_model_runner:
|
||||
logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.")
|
||||
self.use_v2_model_runner = False
|
||||
|
||||
# vllm_omni/worker/gpu_generation_worker.py:96-99
|
||||
if self.use_v2_model_runner:
|
||||
logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.")
|
||||
self.use_v2_model_runner = False
|
||||
```
|
||||
|
||||
The scheduler already has partial V2 awareness:
|
||||
|
||||
```python
|
||||
# vllm_omni/core/sched/output.py:44
|
||||
prefill_token_ids: Optional # prefill token IDs for v2 model runner
|
||||
```
|
||||
|
||||
`OmniNewRequestData.from_request()` conditionally populates `prefill_token_ids` based on `use_v2_model_runner`.
|
||||
|
||||
### Why V2 Matters
|
||||
|
||||
| Reason | Detail |
|
||||
|--------|--------|
|
||||
| **Upstream convergence** | vLLM is moving toward V2 as the default; V1 may eventually be deprecated |
|
||||
| **Performance** | V2 model runner introduces optimizations for attention metadata, micro-batching, and CUDA graph management |
|
||||
| **Feature access** | New upstream features may only be available in V2 |
|
||||
| **Maintenance burden** | Supporting V1 indefinitely means maintaining a growing delta from upstream |
|
||||
|
||||
### What Are "Omni Hooks"?
|
||||
|
||||
The warning message references "omni hooks" — the vllm-omni extensions to the model runner that V2 does not yet support. Based on code analysis, these are:
|
||||
|
||||
| Hook | Location | Description |
|
||||
|------|----------|-------------|
|
||||
| `additional_information` storage | `OmniGPUModelRunner.__init__` | Per-request `additional_information_cpu` attribute on `CachedRequestState` |
|
||||
| `_decode_and_store_request_payloads` | `OmniGPUModelRunner` | Decodes prompt_embeds and additional_information from scheduler output for new requests |
|
||||
| `_collect_additional_information_for_prefill` | `OmniGPUModelRunner` | Overlays prompt_embeds on inputs_embeds during prefill |
|
||||
| `_build_model_kwargs_extra` | `OmniGPUModelRunner` | Injects `runtime_additional_information` into model forward kwargs |
|
||||
| `_process_additional_information_updates` | `OmniGPUModelRunner` | Merges model postprocess outputs back into request state |
|
||||
| `_merge_additional_information_update` | `OmniGPUModelRunner` | CPU-side merge of update dicts with tensor offloading |
|
||||
| `extract_multimodal_outputs` | `OmniGPUModelRunner` | Splits `OmniOutput` into text hidden states + multimodal outputs |
|
||||
| `_model_forward` override | `OmniGPUModelRunner` | Injects extra kwargs and caches model output for multimodal |
|
||||
| M-RoPE extensions | `OmniGPUModelRunner` | Extended M-RoPE with audio, video, spatial metadata |
|
||||
| XD-RoPE | `OmniGPUModelRunner` | HunYuan-VL dimensional RoPE |
|
||||
| Talker MTP | `OmniGPUModelRunner` | Qwen3-TTS/Omni code predictor forward pass |
|
||||
| `has_preprocess` / `has_postprocess` | `OmniGPUModelRunner._preprocess` | Per-request model preprocessing/postprocessing |
|
||||
| `OmniModelRunnerOutput` | `outputs.py` | Extended `ModelRunnerOutput` with `multimodal_outputs` and `kv_extracted_req_ids` |
|
||||
|
||||
## 3. Analysis of V2 Model Runner Architecture
|
||||
|
||||
### 3.1 Key V2 Changes (Based on Upstream vLLM)
|
||||
|
||||
The V2 model runner refactors several areas that affect omni hooks:
|
||||
|
||||
| Area | V1 Behavior | V2 Behavior | Omni Impact |
|
||||
|------|------------|-------------|-------------|
|
||||
| **Input batch** | `GPUInputBatch` with direct attribute access | Restructured batch with new token layout | `_collect_additional_information_for_prefill` needs adaptation |
|
||||
| **Forward context** | `set_forward_context()` with CUDA graph mode | Enhanced forward context with batch descriptors | `_model_forward` override must align |
|
||||
| **Attention metadata** | Per-layer metadata builders | Restructured metadata pipeline | M-RoPE position injection path changes |
|
||||
| **Output handling** | `ModelRunnerOutput` returned from `execute_model` | Potentially different output structure | `OmniModelRunnerOutput` needs remapping |
|
||||
| **CUDA graph capture** | `_dummy_run` with manual graph management | V2 CUDA graph system | `_dummy_run` override may need rewrite |
|
||||
| **Scheduler output** | `SchedulerOutput` with `scheduled_new_reqs` | V2 scheduler output format | Request payload decoding path changes |
|
||||
|
||||
### 3.2 Hook Compatibility Assessment
|
||||
|
||||
| Hook | V2 Compatibility | Effort |
|
||||
|------|-----------------|--------|
|
||||
| `additional_information` storage | **Needs adaptation** — V2 may change `CachedRequestState` | Medium |
|
||||
| `_decode_and_store_request_payloads` | **Needs adaptation** — scheduler output format may differ | Medium |
|
||||
| `_collect_additional_information_for_prefill` | **Needs adaptation** — input batch layout changes | High |
|
||||
| `_build_model_kwargs_extra` | **Likely compatible** — model forward kwargs injection | Low |
|
||||
| `_model_forward` override | **Needs adaptation** — forward context API changes | Medium |
|
||||
| `extract_multimodal_outputs` | **Likely compatible** — output extraction is model-level | Low |
|
||||
| M-RoPE extensions | **Needs adaptation** — position computation pipeline changes | High |
|
||||
| Talker MTP | **Needs adaptation** — CUDA graph management changes | High |
|
||||
| `OmniModelRunnerOutput` | **Needs adaptation** — output class hierarchy may change | Medium |
|
||||
|
||||
## 4. Proposed Integration Plan
|
||||
|
||||
### 4.1 Phase 0: Upstream Tracking (Ongoing)
|
||||
|
||||
- Monitor upstream V2 model runner PRs and API stabilization
|
||||
- Maintain a mapping document of V1 → V2 API changes relevant to omni hooks
|
||||
- Track which omni hooks have upstream equivalents in V2
|
||||
|
||||
### 4.2 Phase 1: Abstract the Hook Interface
|
||||
|
||||
Before integrating V2, abstract the hook points so that both V1 and V2 can be supported:
|
||||
|
||||
```python
|
||||
class OmniModelRunnerHooks(Protocol):
|
||||
"""Interface for omni extensions to the model runner."""
|
||||
|
||||
def on_new_requests(self, scheduler_output) -> None:
|
||||
"""Decode and store payloads for new requests."""
|
||||
...
|
||||
|
||||
def prepare_prefill_embeds(self, num_scheduled_tokens_np) -> None:
|
||||
"""Overlay prompt_embeds for prefill."""
|
||||
...
|
||||
|
||||
def build_extra_kwargs(self) -> dict:
|
||||
"""Build extra model forward kwargs."""
|
||||
...
|
||||
|
||||
def process_model_output(self, hidden_states, mm_outputs, scheduler_output) -> None:
|
||||
"""Post-forward processing of model outputs."""
|
||||
...
|
||||
|
||||
def get_model_runner_output(self, base_output) -> ModelRunnerOutput:
|
||||
"""Wrap base output with omni-specific fields."""
|
||||
...
|
||||
```
|
||||
|
||||
This abstraction aligns with RFC-003 (Model Runner Clearance) and can be implemented concurrently.
|
||||
|
||||
### 4.3 Phase 2: V2 Runner Skeleton
|
||||
|
||||
Create a V2-compatible runner class:
|
||||
|
||||
```python
|
||||
class OmniGPUModelRunnerV2(GPUModelRunnerV2):
|
||||
"""V2 model runner with omni hooks."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._hooks = OmniModelRunnerHooksImpl(self)
|
||||
|
||||
def execute_model(self, scheduler_output, ...):
|
||||
self._hooks.on_new_requests(scheduler_output)
|
||||
self._hooks.prepare_prefill_embeds(...)
|
||||
result = super().execute_model(scheduler_output, ...)
|
||||
self._hooks.process_model_output(result, ...)
|
||||
return self._hooks.get_model_runner_output(result)
|
||||
```
|
||||
|
||||
### 4.4 Phase 3: Port Individual Hooks
|
||||
|
||||
Port each hook to V2, prioritized by complexity and importance:
|
||||
|
||||
| Priority | Hook | Rationale |
|
||||
|----------|------|-----------|
|
||||
| P0 | `_model_forward` override | Core execution path |
|
||||
| P0 | `extract_multimodal_outputs` | Required for any multimodal output |
|
||||
| P0 | `OmniModelRunnerOutput` | Required for output pipeline |
|
||||
| P1 | `additional_information` lifecycle | Required for multi-stage pipelines |
|
||||
| P1 | M-RoPE extensions | Required for Qwen2-VL, Qwen3-Omni |
|
||||
| P2 | Talker MTP | Required for Qwen3-TTS |
|
||||
| P2 | `has_preprocess` / `has_postprocess` | Required for custom models |
|
||||
| P3 | XD-RoPE | Only needed for HunYuan-VL |
|
||||
|
||||
### 4.5 Phase 4: Worker Updates
|
||||
|
||||
Update workers to conditionally use V2:
|
||||
|
||||
```python
|
||||
class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
|
||||
def init_device(self):
|
||||
...
|
||||
if self.use_v2_model_runner:
|
||||
self.model_runner = GPUARModelRunnerV2(self.vllm_config, self.device)
|
||||
else:
|
||||
self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
|
||||
```
|
||||
|
||||
### 4.6 Phase 5: Validation and Cutover
|
||||
|
||||
1. Run full test suite with `use_v2_model_runner=True`
|
||||
2. Benchmark V2 against V1 for key models (Qwen3-Omni, Qwen3-TTS, GLM-Image)
|
||||
3. Enable V2 as opt-in via environment variable
|
||||
4. Eventually make V2 the default when all hooks are ported
|
||||
|
||||
## 5. Compatibility Strategy
|
||||
|
||||
### 5.1 Dual Runner Support
|
||||
|
||||
During the transition, both V1 and V2 runners coexist:
|
||||
|
||||
```
|
||||
vllm_omni/worker/
|
||||
├── gpu_model_runner.py # V1: OmniGPUModelRunner (existing)
|
||||
├── gpu_model_runner_v2.py # V2: OmniGPUModelRunnerV2 (new)
|
||||
├── gpu_ar_model_runner.py # V1 AR runner
|
||||
├── gpu_ar_model_runner_v2.py # V2 AR runner (new)
|
||||
├── gpu_generation_model_runner.py # V1 generation runner
|
||||
├── gpu_generation_model_runner_v2.py # V2 generation runner (new)
|
||||
```
|
||||
|
||||
### 5.2 Feature Parity Matrix
|
||||
|
||||
Track feature parity with a compatibility matrix:
|
||||
|
||||
| Feature | V1 | V2 | Notes |
|
||||
|---------|----|----|-------|
|
||||
| Text generation | Yes | Phase 2 | |
|
||||
| Multimodal input (image/video/audio) | Yes | Phase 2 | |
|
||||
| Multimodal output | Yes | Phase 3 | |
|
||||
| Multi-stage pipeline | Yes | Phase 3 | |
|
||||
| M-RoPE | Yes | Phase 3 | |
|
||||
| Talker MTP | Yes | Phase 3 | |
|
||||
| CUDA graph capture | Yes | Phase 2 | |
|
||||
| Prompt embeds | Yes | Phase 3 | |
|
||||
| Diffusion models | Yes | Phase 4 | Separate runner |
|
||||
| NPU support | Yes | Phase 5 | Depends on `vllm_ascend` V2 |
|
||||
| XPU support | Yes | Phase 5 | |
|
||||
|
||||
## 6. Platform Considerations
|
||||
|
||||
### 6.1 NPU (Ascend)
|
||||
|
||||
NPU currently uses `vllm_ascend.worker.model_runner_v1.NPUModelRunner` as its base. The V2 integration depends on:
|
||||
- Whether `vllm_ascend` provides a V2 model runner
|
||||
- Whether NPU-specific changes can be layered on top of `OmniGPUModelRunnerV2`
|
||||
|
||||
### 6.2 XPU (Intel)
|
||||
|
||||
XPU runners (`XPUARModelRunner`, `XPUGenerationModelRunner`) inherit from the GPU runners. They should automatically benefit from V2 if the GPU V2 runner is correct, but may need platform-specific overrides for XPU memory management.
|
||||
|
||||
## 7. Testing Plan
|
||||
|
||||
| Level | Description |
|
||||
|-------|-------------|
|
||||
| Unit | Test each ported hook independently with mock V2 runner |
|
||||
| Integration | Run E2E tests for each model with `use_v2_model_runner=True` |
|
||||
| Benchmark | Compare V1 vs V2 throughput/latency for Qwen3-Omni, Qwen3-TTS |
|
||||
| CI gate | Add V2 runner to CI matrix (opt-in initially, required later) |
|
||||
| Platform | Test NPU/XPU with V2 runner (dependent on platform V2 support) |
|
||||
|
||||
## 8. Timeline
|
||||
|
||||
| Phase | Scope | Dependencies | Est. Duration |
|
||||
|-------|-------|-------------|---------------|
|
||||
| Phase 0 | Upstream tracking | None | Ongoing |
|
||||
| Phase 1 | Abstract hook interface | RFC-003 | 2 weeks |
|
||||
| Phase 2 | V2 runner skeleton + core hooks | Phase 1 + upstream V2 stabilization | 3 weeks |
|
||||
| Phase 3 | Port all hooks | Phase 2 | 4 weeks |
|
||||
| Phase 4 | Worker updates + opt-in flag | Phase 3 | 1 week |
|
||||
| Phase 5 | Validation + platform ports | Phase 4 | 3 weeks |
|
||||
| Cutover | V2 as default | Phase 5 + confidence from CI | TBD |
|
||||
|
||||
## 9. Risks
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|------------|
|
||||
| Upstream V2 API instability | Phase 0 tracking; abstract hook interface insulates from changes |
|
||||
| Performance regression in V2 | Benchmark gating before cutover |
|
||||
| Feature incompleteness | Feature parity matrix; V1 fallback always available |
|
||||
| Platform support gaps | NPU/XPU V2 treated as separate Phase 5 work |
|
||||
| Large merge conflicts during V2 port | Coordinate with rebase automation (RFC-001) |
|
||||
|
||||
## 10. Open Questions
|
||||
|
||||
1. Is upstream V2 model runner stable enough to begin Phase 2, or should we wait for an official release?
|
||||
2. Should we attempt to contribute omni hooks upstream to V2, or keep them in vllm-omni?
|
||||
3. How does V2 affect the `DiffusionModelRunner` (which has its own independent runner implementation)?
|
||||
4. Should the V2 integration be gated behind a configuration flag or environment variable during the transition?
|
||||
5. What is the minimum set of hooks needed for a useful V2 MVP (i.e., which models can work with a partial port)?
|
||||
Reference in new issue
Block a user