vllm-omni/docs/design/rfcs/005-model-runner-v2-integration.md

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:

# 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:

# 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:

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:

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:

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)?