[Perf][VoxCPM2] Add unified decode graph optimization (#4606)

Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod
+from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _MockQueue:
+    def __init__(self, items: list | None = None) -> None:
+        self._items = list(items or [])
+
+    def __bool__(self) -> bool:
+        return bool(self._items)
+
+    def __iter__(self):
+        return iter(self._items)
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+    def add_request(self, request) -> None:
+        self._items.append(request)
+
+    def prepend_requests(self, requests) -> None:
+        self._items = list(requests) + self._items
+
+
+class _MockRequest:
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        status: RequestStatus = RequestStatus.RUNNING,
+        num_prompt_tokens: int = 4,
+        num_computed_tokens: int = 4,
+        num_output_placeholders: int = 0,
+    ) -> None:
+        self.request_id = request_id
+        self.status = status
+        self.num_prompt_tokens = num_prompt_tokens
+        self.num_computed_tokens = num_computed_tokens
+        self.num_output_placeholders = num_output_placeholders
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+
+def _make_scheduler(*, enable_unified_decode_graph: bool = True) -> VoxCPM2OmniARAsyncScheduler:
+    sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler)
+    runtime_config = SimpleNamespace(enable_unified_decode_graph=enable_unified_decode_graph)
+    sched.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(voxcpm2_runtime_config=runtime_config),
+        )
+    )
+    return sched
+
+
+def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_decode_ready() -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("prefill-running", num_prompt_tokens=8, num_computed_tokens=4)]
+    scheduler.waiting = _MockQueue([_MockRequest("waiting", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=False)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
+    deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
+    original_waiting = _MockQueue([original_waiting_req])
+    scheduler.waiting = original_waiting
+    scheduler.policy = "fcfs"
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler._consume_pending_connector_output = lambda model_mode: None
+
+    monkeypatch.setattr(scheduler_mod, "create_request_queue", lambda _policy: _MockQueue())
+
+    def fake_upstream_schedule(self, throttle_prefills: bool = False):
+        assert self.waiting is not original_waiting
+        assert not self.waiting
+        self.waiting.add_request(deferred_by_upstream)
+        raise RuntimeError("stop before output wrapping")
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
+
+    with pytest.raises(RuntimeError, match="stop before output wrapping"):
+        scheduler.schedule()
+
+    assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]

tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Static regression checks for VoxCPM2 runner-assisted unified graph hooks.
+
+These tests intentionally avoid importing torch/vLLM so they can run in a
+lightweight local checkout. Runtime audio quality still requires a CUDA test.
+"""
+
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py"
+RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py"
+TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py"
+VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py"
+VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py"
+OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py"
+DEPLOY = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm2.yaml"
+
+
+def test_ar_runner_exposes_runner_assisted_full_metadata_hook():
+    source = RUNNER.read_text()
+
+    assert "Models without this hook keep the normal runner path" in source
+    assert "get_runner_assisted_full_attention_metadata_request" in source
+    assert "set_runner_assisted_full_attention_metadata_context" in source
+    assert "runner_assisted_full_attn" in source
+    assert "runner_assisted_full_attn_capture" in source
+    assert "pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn" in source
+    assert "for_cudagraph_capture=runner_assisted_full_attn_capture" in source
+    assert "cudagraph_runtime_mode=(" in source
+    assert "CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode" in source
+    refresh_source = source[source.index("def _refresh_runner_assisted_full_attention_metadata_buffers") :]
+    refresh_source = refresh_source[: refresh_source.index("def _set_runner_assisted_full_attention_metadata_context")]
+    assert "num_computed_tokens_cpu" in refresh_source
+    assert "num_scheduled_tokens_np" in refresh_source
+    assert "optimistic_seq_lens_cpu[:num_reqs]=1" not in "".join(refresh_source.split())
+    padding_source = source[
+        source.index("runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request") :
+    ]
+    padding_source = padding_source[: padding_source.index("ubatch_slices, ubatch_slices_padded")]
+    assert "runner_assisted_full_attn_request.num_reqs_padded" in padding_source
+    assert "runner_assisted_full_attn_request.for_cudagraph_capture" in padding_source
+    assert "num_reqs_padded * max_num_scheduled_tokens" in padding_source
+
+
+def test_ar_runner_without_model_hook_stays_on_normal_path():
+    source = RUNNER.read_text()
+    contract_source = RUNNER_ASSISTED_METADATA.read_text()
+    request_source = source[source.index("def _get_runner_assisted_full_attention_metadata_request") :]
+    request_source = request_source[
+        : request_source.index("def _refresh_runner_assisted_full_attention_metadata_buffers")
+    ]
+    compact_request_source = "".join(request_source.split())
+
+    assert "Models without this hook keep the normal runner path" in request_source
+    assert "class RunnerAssistedAttentionMetadataProvider(Protocol)" in contract_source
+    assert "class RunnerAssistedFullAttentionMetadataRequest(NamedTuple)" in contract_source
+    assert "tuple[int, bool]" not in contract_source
+    assert 'hook=getattr(self.model,"get_runner_assisted_full_attention_metadata_request",None)' in (
+        compact_request_source
+    )
+    assert "ifnotcallable(hook):returnNone" in compact_request_source
+    assert ".tolist()" not in compact_request_source
+    assert "isinstance(request,RunnerAssistedFullAttentionMetadataRequest)" in compact_request_source
+    assert "request.num_reqs_padded" in request_source
+    assert "request.for_cudagraph_capture" in request_source
+    assert "exceptException" not in compact_request_source
+
+    set_context_source = source[source.index("def _set_runner_assisted_full_attention_metadata_context") :]
+    set_context_source = set_context_source[: set_context_source.index("def _deferred_prefix_cache_mm_keys")]
+    assert "except Exception" not in set_context_source
+
+
+def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
+    source = TALKER.read_text()
+    compact_source = "".join(source.split())
+
+    assert "_voxcpm2_compile_without_inductor_cudagraphs" in source
+    assert "def _compile_without_inductor_cudagraphs" not in source
+    assert "_voxcpm2_unwrap_torch_compile" in source
+    assert "_orig_mod" in source
+    compile_source = source[source.index("def _voxcpm2_compile_without_inductor_cudagraphs") :]
+    compile_source = compile_source[: compile_source.index("def _setup_torch_compile")]
+    assert "torch._inductor.config" not in compile_source
+    assert '"triton.cudagraphs": False' in compile_source
+    assert '"triton.cudagraph_trees": False' in compile_source
+
+    assert "self._enable_unified_decode_graph=(use_cuda_graph" in compact_source
+    assert "andnotself._deterministic_cfm_noise" in compact_source
+    assert "decode_tail_graph" not in source
+
+    unified_source = source[source.index("def _forward_unified_decode") :]
+    unified_source = unified_source[: unified_source.index("# -------------------- vllm hooks")]
+    assert "except Exception:" not in unified_source
+    assert "_failed_unified_graph_sizes" not in source
+    assert "self._enable_unified_decode_graph = False" not in unified_source
+    assert "_forward_unified_decode_fallback" not in source
+    assert "capture_failed" not in compact_source
+    assert "record_fallback" not in source
+    assert "fallbacks" not in source
+    assert "**info_dict: Unpack[VoxCPM2PreprocessInput]" in source
+    assert "**info: Unpack[VoxCPM2PostprocessInput]" in source
+    assert "reference_audio: Any" not in source
+    assert "prompt_audio: Any" not in source
+    assert "def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike" in source
+    assert "class _PrefillInputs(NamedTuple)" in source
+    assert "-> _PrefillInputs" in source
+    assert 'inputs["audio_feat"]' not in source
+    assert 'inputs["text_token"]' not in source
+
+    capture_source = source[source.index("def _capture_unified_decode_graph") :]
+    capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")]
+    assert "_voxcpm2_compile_unified_capture_estimator" in capture_source
+    assert "_voxcpm2_compile_unified_capture_feat_encoder" in capture_source
+    graph_body = capture_source[capture_source.index("with torch.cuda.graph") :]
+    assert "g.cfm_noise.normal_()" not in graph_body
+    assert "g.cfm_noise.normal_()" in unified_source
+
+    assert "_pre_capture_unified_graphs" not in source
+    assert "_unified_graph_pre_capture_sizes" not in source
+    assert "unified_decode_graph_pre_capture_sizes" not in source
+
+
+def test_voxcpm2_batch_unified_graph_requires_runner_metadata_marker():
+    source = TALKER.read_text()
+    compact_source = "".join(source.split())
+
+    assert "enable_runner_assisted_unified_decode_graph" not in source
+    assert "allow_unified_decode_graph_batch_attention" not in source
+    assert "get_runner_assisted_full_attention_metadata_request" in source
+    assert "set_runner_assisted_full_attention_metadata_context" in source
+    assert "RunnerAssistedFullAttentionMetadataRequest" in source
+    assert "_runner_assisted_unified_decode_graph_active" in source
+    assert "runner_full_metadata_missing" in source
+    assert "_build_unified_graph_bucket_sizes" in source
+    assert "_select_unified_graph_bucket_size" in source
+
+    capture_source = source[source.index("def _capture_unified_decode_graph") :]
+    capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")]
+    assert "override_forward_context" in capture_source
+    assert "_nullify_volatile_metadata" in capture_source
+    assert "capture_context = override_forward_context(self._nullify_volatile_metadata(ctx))" in capture_source
+    assert "ifsize>1:continue" not in compact_source
+    forward_source = source[source.index("def _forward_unified_decode") :]
+    forward_source = forward_source[: forward_source.index("# -------------------- vllm hooks")]
+    assert "graph_size = self._select_unified_graph_bucket_size(num_reqs)" in forward_source
+    assert "self._unified_graphs[graph_size]" in forward_source
+    assert "g.input_embeds[num_reqs:graph_size].zero_()" not in forward_source
+    assert "padded rows as valid duplicates" in forward_source
+    assert "g.input_embeds[last : last + 1].expand" in forward_source
+    assert "ifnum_reqs>1andnotself._runner_assisted_unified_decode_graph_active" in compact_source
+    assert "andnotcfg.enable_runner_assisted_unified_decode_graph" not in compact_source
+    needs_source = source[source.index("def get_runner_assisted_full_attention_metadata_request") :]
+    needs_source = needs_source[: needs_source.index("def set_runner_assisted_full_attention_metadata_context")]
+    assert "_select_unified_graph_bucket_size(num_reqs)" in needs_source
+    assert "_should_use_decode_graph(num_reqs)" not in needs_source
+    assert "RunnerAssistedFullAttentionMetadataRequest(" in needs_source
+
+
+def test_voxcpm2_unified_skip_preserves_segmented_decode_graphs():
+    source = TALKER.read_text()
+    compact_source = "".join(source.split())
+    forward_source = source[source.index("def forward(") :]
+    forward_source = forward_source[: forward_source.index("# -------------------- prefill / decode helpers")]
+
+    assert "use_segmented_decode_graph" in forward_source
+    assert "unified_skip_reasonisnotNone" in compact_source
+    assert "notself._enable_unified_decode_graphorunified_skip_reasonisnotNone" in compact_source
+    assert "ifuse_segmented_decode_graph:" in compact_source
+    assert "anduse_segmented_decode_graph" in compact_source
+
+
+def test_voxcpm2_scheduler_policy_stays_model_local():
+    common_source = OMNI_AR_SCHEDULER.read_text()
+    voxcpm2_scheduler_source = VOXCPM2_SCHEDULER.read_text()
+    pipeline_source = VOXCPM2_PIPELINE.read_text()
+
+    assert "VoxCPM2TalkerForConditionalGeneration" not in common_source
+    assert "voxcpm2_runtime_config" not in common_source
+    assert "pure_decode_graph" not in common_source
+    assert "_schedule_with_optional_waiting_deferral" not in common_source
+    assert "def _should_defer_waiting_admission(self) -> bool:" in common_source
+
+    assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source
+    assert "voxcpm2_runtime_config" in voxcpm2_scheduler_source
+    assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source
+    assert "def schedule(" not in voxcpm2_scheduler_source
+    assert "create_request_queue" not in voxcpm2_scheduler_source
+    assert "original_waiting.prepend_requests(deferred_waiting)" not in voxcpm2_scheduler_source
+    assert "return self._should_defer_waiting_for_unified_decode_graph()" in voxcpm2_scheduler_source
+    assert "vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler" in pipeline_source
+
+
+def test_voxcpm2_deploy_defaults_to_full_unified_graph_only():
+    source = DEPLOY.read_text()
+
+    assert "max_num_seqs: 8" in source
+    assert "enable_unified_decode_graph: true" in source
+    assert "unified_decode_graph_max_batch_size: 8" in source
+    assert "unified_decode_graph_pre_capture_sizes" not in source
+    assert "enable_runner_assisted_unified_decode_graph" not in source
+    assert "allow_unified_decode_graph_batch_attention" not in source
+    assert "decode_graph_capture_policy" not in source

tests/test_config_factory.py

@@ -571,6 +571,7 @@ class TestStagePipelineConfig:
         assert s.final_output is False
         assert s.sampling_constraints == {}
         assert s.engine_output_type is None
+        assert s.scheduler_cls is None
 
 
 class TestPipelineConfigNew:
@@ -809,6 +810,26 @@ stages:
 
         assert stages[0].yaml_runtime["requires_multimodal_data"] is True
 
+    def test_merge_pipeline_deploy_preserves_pipeline_scheduler_cls(self):
+        scheduler_cls = "tests.fake.CustomScheduler"
+        pipeline = PipelineConfig(
+            model_type="test_scheduler",
+            model_arch="TestModel",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="ar",
+                    execution_type=StageExecutionType.LLM_AR,
+                    scheduler_cls=scheduler_cls,
+                ),
+            ),
+        )
+        deploy = DeployConfig(async_chunk=False, stages=[StageDeployConfig(stage_id=0)])
+
+        stages = merge_pipeline_deploy(pipeline, deploy)
+
+        assert stages[0].scheduler_cls == scheduler_cls
+
     def test_mixed_schema_preserves_flat_fields(self):
         """Ensure flat fields are not dropped when engine_args are present."""
         fake_config = {

tests/worker/test_gpu_ar_model_runner.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from types import SimpleNamespace
 
 import numpy as np
@@ -10,6 +11,7 @@ from vllm_omni.worker.gpu_ar_model_runner import (
     GPUARModelRunner,
     OmniAsyncGPUModelRunnerOutput,
 )
+from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -58,6 +60,101 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker():
     assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
 
 
+def test_runner_assisted_full_attention_metadata_request_is_opt_in():
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = object()
+    runner.scheduler_config = SimpleNamespace(max_num_seqs=16)
+
+    request = runner._get_runner_assisted_full_attention_metadata_request(
+        req_ids=["r1", "r2"],
+        num_reqs=2,
+        num_reqs_padded=4,
+        num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
+        num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32),
+        max_num_scheduled_tokens=1,
+    )
+
+    assert request is None
+
+
+def test_runner_assisted_full_attention_metadata_request_and_context_hooks():
+    calls = []
+
+    class Model:
+        def get_runner_assisted_full_attention_metadata_request(
+            self,
+            *,
+            req_ids: Sequence[str],
+            num_reqs: int,
+            num_scheduled_tokens: Sequence[int],
+            num_computed_tokens: Sequence[int],
+            max_num_scheduled_tokens: int,
+        ) -> RunnerAssistedFullAttentionMetadataRequest:
+            calls.append(
+                (
+                    "request",
+                    {
+                        "req_ids": list(req_ids),
+                        "num_reqs": num_reqs,
+                        "num_scheduled_tokens": [int(n) for n in num_scheduled_tokens],
+                        "num_computed_tokens": [int(n) for n in num_computed_tokens],
+                        "max_num_scheduled_tokens": max_num_scheduled_tokens,
+                    },
+                )
+            )
+            return RunnerAssistedFullAttentionMetadataRequest(
+                num_reqs_padded=12,
+                for_cudagraph_capture=True,
+            )
+
+        def set_runner_assisted_full_attention_metadata_context(
+            self,
+            *,
+            enabled: bool,
+            num_reqs: int = 0,
+        ) -> None:
+            calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs}))
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = Model()
+    runner.scheduler_config = SimpleNamespace(max_num_seqs=8)
+
+    request = runner._get_runner_assisted_full_attention_metadata_request(
+        req_ids=["r1", "r2"],
+        num_reqs=2,
+        num_reqs_padded=4,
+        num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
+        num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32),
+        max_num_scheduled_tokens=1,
+    )
+    context_enabled = runner._set_runner_assisted_full_attention_metadata_context(
+        enabled=True,
+        num_reqs=2,
+    )
+    context_disabled = runner._set_runner_assisted_full_attention_metadata_context(enabled=False)
+
+    assert request == RunnerAssistedFullAttentionMetadataRequest(
+        num_reqs_padded=8,
+        for_cudagraph_capture=True,
+    )
+    assert context_enabled
+    assert context_disabled
+    assert calls == [
+        (
+            "request",
+            {
+                "req_ids": ["r1", "r2"],
+                "num_reqs": 2,
+                "num_scheduled_tokens": [1, 1],
+                "num_computed_tokens": [5, 6],
+                "max_num_scheduled_tokens": 1,
+            },
+        ),
+        ("context", {"enabled": True, "num_reqs": 2}),
+        ("context", {"enabled": False, "num_reqs": 0}),
+    ]
+
+
 def test_omni_async_gpu_model_runner_output_builds_lazily_once():
     async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
     calls = []
@@ -357,6 +454,47 @@ def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
     assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]]
 
 
+def test_runner_assisted_full_attention_metadata_refresh_pads_buffers():
+    class QueryStartLoc:
+        def __init__(self):
+            self.np = np.full(5, -1, dtype=np.int32)
+            self.copied = False
+
+        def copy_to_gpu(self):
+            self.copied = True
+
+    class BlockTable:
+        def __init__(self):
+            self.commits = []
+
+        def commit_block_table(self, num_reqs_padded):
+            self.commits.append(num_reqs_padded)
+
+    runner = object.__new__(GPUARModelRunner)
+    block_table = BlockTable()
+    runner.input_batch = SimpleNamespace(
+        num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32),
+        block_table=block_table,
+    )
+    runner.optimistic_seq_lens_cpu = torch.zeros(4, dtype=torch.int32)
+    runner.seq_lens = torch.empty(4, dtype=torch.int32)
+    runner.query_pos = SimpleNamespace(np=np.empty(4, dtype=np.int32))
+    runner.query_start_loc = QueryStartLoc()
+    runner._get_cumsum_and_arange = lambda scheduled, _query_pos: np.cumsum(scheduled, dtype=np.int32)
+
+    runner._refresh_runner_assisted_full_attention_metadata_buffers(
+        num_reqs=2,
+        num_reqs_padded=4,
+        num_scheduled_tokens_np=np.array([1, 2], dtype=np.int32),
+    )
+
+    assert runner.optimistic_seq_lens_cpu.tolist() == [11, 22, 0, 0]
+    assert runner.seq_lens.tolist() == [11, 22, 0, 0]
+    assert runner.query_start_loc.np.tolist() == [0, 1, 3, 3, 3]
+    assert runner.query_start_loc.copied
+    assert block_table.commits == [4]
+
+
 @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
 def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
     runner = object.__new__(GPUARModelRunner)

vllm_omni/benchmarks/patch/patch.py

@@ -1264,7 +1264,7 @@ if "daily-omni" not in OPENAI_COMPATIBLE_BACKENDS:
 # Prevent import order from causing patch failures
 from vllm.benchmarks import serve
 from vllm.benchmarks.lib.ready_checker import wait_for_endpoint
-from vllm.benchmarks.serve import TaskType, _merge_overrides, calculate_metrics_for_embeddings, get_request
+from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request
 
 from vllm_omni.benchmarks.metrics.metrics import (
     MultiModalsBenchmarkMetrics,
@@ -1276,6 +1276,19 @@ from vllm_omni.benchmarks.metrics.metrics import (
 benchmark_old = serve.benchmark
 
 
+def _merge_overrides(base: dict | None, overrides: dict | None) -> dict | None:
+    """Merge benchmark extra_body with per-request overrides.
+
+    vLLM 0.24 removed the private helper from ``vllm.benchmarks.serve``.
+    Keep the same shallow-merge behavior here, with request overrides winning.
+    """
+    if not base and not overrides:
+        return None
+    merged = dict(base or {})
+    merged.update(overrides or {})
+    return merged
+
+
 async def benchmark(
     task_type: TaskType,
     endpoint_type: str,

vllm_omni/config/stage_config.py

@@ -228,6 +228,7 @@ class StagePipelineConfig:
     prompt_expand_func: str | None = None
     cfg_kv_collect_func: str | None = None
     omni_kv_config: dict[str, Any] | None = None
+    scheduler_cls: str | None = None
     # Model subdirectory indirections: for multi-component HF repos where the
     # stage's config/tokenizer lives in a subdirectory (e.g. GLM-Image's AR
     # config is in ``vision_language_encoder/``).  Consumed at stage-init time
@@ -896,7 +897,7 @@ def merge_pipeline_deploy(
                 final_output=ps.final_output,
                 final_output_type=ps.final_output_type,
                 worker_type=worker_type,
-                scheduler_cls=_scheduler_path(sched_cls),
+                scheduler_cls=ps.scheduler_cls or _scheduler_path(sched_cls),
                 hf_config_name=ps.hf_config_name,
                 is_comprehension=ps.owns_tokenizer,
                 yaml_engine_args=engine_args,

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStat
 from vllm.logger import init_logger
 from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMScheduler
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason
@@ -133,6 +134,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._omits_kv_transfer_cache[rid] = result
         return result
 
+    def _should_defer_waiting_admission(self) -> bool:
+        return False
+
     def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int]) -> bool:
         """
         Check triggers and process side effects (marking transfer).
@@ -222,9 +226,19 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 self.waiting, self.running, scheduler_requests=self.requests
             )
 
+        original_waiting = None
+        if self._should_defer_waiting_admission():
+            original_waiting = self.waiting
+            self.waiting = create_request_queue(self.policy)
+
         try:
             scheduler_output = super().schedule(throttle_prefills)
         finally:
+            if original_waiting is not None:
+                deferred_waiting = list(self.waiting)
+                if deferred_waiting:
+                    original_waiting.prepend_requests(deferred_waiting)
+                self.waiting = original_waiting
             if self.chunk_transfer_adapter:
                 # Add request waiting for chunk to the waiting and running queue
                 self.chunk_transfer_adapter.restore_queues(

vllm_omni/deploy/voxcpm2.yaml

@@ -8,16 +8,17 @@ dtype: bfloat16
 
 stages:
   - stage_id: 0
-    max_num_seqs: 32
+    max_num_seqs: 8
     gpu_memory_utilization: 0.9
     # Right-size the KV cache instead of letting it fill gpu_memory_utilization.
-    # max_num_seqs(32) * max_model_len(4096) ~= 131k tokens; 6 GiB holds ~175k
-    # (42x concurrency, above the 32-seq scheduler cap, so no throughput loss).
+    # max_num_seqs(8) * max_model_len(4096) ~= 32k tokens; 6 GiB leaves headroom
+    # for longer live requests while keeping VoxCPM2 TTS decode batches small
+    # enough to avoid the high ITL observed with large admission caps.
     # The remaining VRAM stays free for the VoxCPM2 diffusion side-path
     # (LocDiT / AudioVAE / code2wav + batch CUDA-graph pools), which allocates
     # outside vLLM's KV accounting during inference. Keeps peak ~13 GiB on any
     # card and avoids per-rebase OOM tuning on small CI GPUs.
-    # Assumes prefix caching off and max_num_seqs*max_model_len <= ~175k tokens;
+    # Assumes prefix caching off and typical live decode concurrency <= ~8;
     # raise this cap if either grows or enable_prefix_caching is turned on.
     kv_cache_memory_bytes: 6442450944  # 6 GiB
     enforce_eager: true
@@ -27,18 +28,20 @@ stages:
     max_num_batched_tokens: 4096
     max_model_len: 4096
     devices: "0"
-    engine_extras:
-      hf_overrides:
-        voxcpm2_runtime_config:
-          audio_emit_every: 1
-          vae_decode_every: 3
-          enable_batched_cfm: true
-          enable_batched_vae_decode: true
-          coalesce_audio_d2h: true
-          enable_loc_dit_fused_qkv: true
-          enable_loc_dit_fused_mlp: true
-          enable_loc_dit_zero_dt_cache: true
-          enable_loc_dit_skip_qkv_contig: true
+    hf_overrides:
+      voxcpm2_runtime_config:
+        audio_emit_every: 1
+        vae_decode_every: 3
+        enable_batched_cfm: true
+        enable_batched_vae_decode: true
+        coalesce_audio_d2h: true
+        enable_loc_dit_fused_qkv: true
+        enable_loc_dit_fused_mlp: true
+        enable_loc_dit_zero_dt_cache: true
+        enable_loc_dit_skip_qkv_contig: true
+        enable_loc_dit_reduce_overhead_no_cg: true
+        enable_unified_decode_graph: true
+        unified_decode_graph_max_batch_size: 8
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0

vllm_omni/model_executor/models/bagel/bagel.py

@@ -587,9 +587,9 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
         input_ids: torch.Tensor | None,
         positions: torch.Tensor | None,
         inputs_embeds: torch.Tensor | None,
-        req_ids: list[str],
-        num_computed_tokens: list[int],
-        num_scheduled_tokens: list[int],
+        req_ids: Sequence[str],
+        num_computed_tokens: Sequence[int],
+        num_scheduled_tokens: Sequence[int],
         input_ids_buffer: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Restore input_ids so _adjust_positions_for_img2img can locate
@@ -599,7 +599,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
             input_ids = input_ids_buffer
         return input_ids, positions
 
-    def flush_pending_metadata(self, req_ids: list[str]) -> None:
+    def flush_pending_metadata(self, req_ids: Sequence[str]) -> None:
         """Map pending metadata (batch order) to req_ids after forward().
 
         Guard: if a request already has metadata with ``image_shape``

vllm_omni/model_executor/models/voxcpm2/pipeline.py

@@ -26,6 +26,7 @@ VOXCPM2_PIPELINE = PipelineConfig(
             final_output_type="audio",
             owns_tokenizer=True,
             engine_output_type="audio",
+            scheduler_cls="vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler",
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [1],

vllm_omni/model_executor/models/voxcpm2/scheduler.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler
+
+
+class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
+    """VoxCPM2 scheduler variant for full unified decode graph serving.
+
+    VoxCPM2's full unified decode graph only applies to pure decode batches.
+    When a decode-ready request is already running, this scheduler defers new
+    waiting admissions for the current tick so the decode batch can stay on the
+    VoxCPM2 graph path. This is a model-local serving policy, not a generic AR
+    scheduler rule.
+    """
+
+    def _unified_decode_graph_enabled(self) -> bool:
+        hf_config = getattr(self.vllm_config.model_config, "hf_config", None)
+        runtime_config = getattr(hf_config, "voxcpm2_runtime_config", None)
+        if isinstance(runtime_config, dict):
+            return bool(runtime_config.get("enable_unified_decode_graph", False))
+        return bool(getattr(runtime_config, "enable_unified_decode_graph", False))
+
+    def _should_defer_waiting_for_unified_decode_graph(self) -> bool:
+        if not self._unified_decode_graph_enabled():
+            return False
+        if not self.waiting or not self.running:
+            return False
+
+        for request in self.running:
+            if getattr(request, "status", None) != RequestStatus.RUNNING or request.is_finished():
+                continue
+            if self._get_confirmed_num_computed_tokens(request) >= request.num_prompt_tokens:
+                return True
+        return False
+
+    def _should_defer_waiting_admission(self) -> bool:
+        return self._should_defer_waiting_for_unified_decode_graph()

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -621,8 +621,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     positions=positions,
                     inputs_embeds=inputs_embeds,
                     req_ids=req_ids[:num_reqs],
-                    num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)],
-                    num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)],
+                    num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs],
+                    num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
                     input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
                 )
             #  -------------------------------------- Omni-new -------------------------------------------------
@@ -690,8 +690,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
         with record_function_or_nullcontext("post process"):
             #  -------------------------------------- Omni-new -------------------------------------------------
             # [Omni] Map pending ropes metadata to req_ids.
-            if hasattr(self.model, "flush_pending_metadata"):
-                self.model.flush_pending_metadata(list(req_ids))
+            flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None)
+            if callable(flush_pending_metadata):
+                flush_pending_metadata(req_ids[:num_reqs])
 
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
 

vllm_omni/worker/gpu_ar_model_runner.py

@@ -8,7 +8,7 @@ from __future__ import annotations
 
 import gc
 import threading
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Mapping, Sequence
 from contextlib import nullcontext
 from copy import copy
 from dataclasses import replace
@@ -46,6 +46,7 @@ from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 
 logger = init_logger(__name__)
 
@@ -588,6 +589,93 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         model = getattr(self, "model", None)
         return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))
 
+    def _get_runner_assisted_full_attention_metadata_request(
+        self,
+        *,
+        req_ids: Sequence[str],
+        num_reqs: int,
+        num_reqs_padded: int,
+        num_scheduled_tokens_np: np.ndarray,
+        num_computed_tokens_cpu: torch.Tensor | np.ndarray,
+        max_num_scheduled_tokens: int,
+    ) -> RunnerAssistedFullAttentionMetadataRequest | None:
+        # Models without this hook keep the normal runner path. VoxCPM2 uses
+        # it to ask the runner for padded FULL attention metadata while keeping
+        # graph policy and lifecycle in the model layer.
+        hook = getattr(self.model, "get_runner_assisted_full_attention_metadata_request", None)
+        if not callable(hook):
+            return None
+        num_computed_tokens = num_computed_tokens_cpu
+        if hasattr(num_computed_tokens, "numpy"):
+            num_computed_tokens = num_computed_tokens.numpy()
+        request = hook(
+            req_ids=req_ids,
+            num_reqs=num_reqs,
+            num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
+            num_computed_tokens=num_computed_tokens,
+            max_num_scheduled_tokens=max_num_scheduled_tokens,
+        )
+        if request is None:
+            return None
+        if not isinstance(request, RunnerAssistedFullAttentionMetadataRequest):
+            raise TypeError(
+                "runner-assisted full attention metadata hook must return "
+                "RunnerAssistedFullAttentionMetadataRequest or None, got "
+                f"{type(request).__name__}"
+            )
+        padded_num_reqs = max(
+            num_reqs_padded,
+            min(int(request.num_reqs_padded), self.scheduler_config.max_num_seqs),
+        )
+        return RunnerAssistedFullAttentionMetadataRequest(
+            num_reqs_padded=padded_num_reqs,
+            for_cudagraph_capture=bool(request.for_cudagraph_capture),
+        )
+
+    def _refresh_runner_assisted_full_attention_metadata_buffers(
+        self,
+        *,
+        num_reqs: int,
+        num_reqs_padded: int,
+        num_scheduled_tokens_np: np.ndarray,
+    ) -> None:
+        num_computed_tokens = self.input_batch.num_computed_tokens_cpu[:num_reqs]
+        if hasattr(num_computed_tokens, "numpy"):
+            num_computed_tokens = num_computed_tokens.numpy()
+        np.add(
+            num_computed_tokens,
+            num_scheduled_tokens_np,
+            out=self.optimistic_seq_lens_cpu[:num_reqs].numpy(),
+        )
+        self.optimistic_seq_lens_cpu[num_reqs:num_reqs_padded].fill_(0)
+        self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True)
+
+        cum_num_tokens = self._get_cumsum_and_arange(
+            num_scheduled_tokens_np,
+            self.query_pos.np,
+        )
+        self.query_start_loc.np[0] = 0
+        self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+        self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1])
+        self.query_start_loc.copy_to_gpu()
+
+        self.input_batch.block_table.commit_block_table(num_reqs_padded)
+
+    def _set_runner_assisted_full_attention_metadata_context(
+        self,
+        *,
+        enabled: bool,
+        num_reqs: int = 0,
+    ) -> bool:
+        hook = getattr(self.model, "set_runner_assisted_full_attention_metadata_context", None)
+        if not callable(hook):
+            return False
+        hook(
+            enabled=enabled,
+            num_reqs=num_reqs,
+        )
+        return True
+
     def _deferred_prefix_cache_mm_keys(self) -> set[str]:
         """Model-declared multimodal keys whose prefix-cache writes are deferred."""
         model = getattr(self, "model", None)
@@ -1028,6 +1116,26 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             )
             num_tokens_padded = batch_desc.num_tokens
             num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
+            num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu[:num_reqs]
+
+            runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request(
+                req_ids=req_ids[:num_reqs],
+                num_reqs=num_reqs,
+                num_reqs_padded=num_reqs_padded,
+                num_scheduled_tokens_np=num_scheduled_tokens_np,
+                num_computed_tokens_cpu=num_computed_tokens_cpu,
+                max_num_scheduled_tokens=max_num_scheduled_tokens,
+            )
+            runner_assisted_full_attn_capture = False
+            if runner_assisted_full_attn_request is not None:
+                num_reqs_padded = runner_assisted_full_attn_request.num_reqs_padded
+                runner_assisted_full_attn_capture = runner_assisted_full_attn_request.for_cudagraph_capture
+                num_tokens_padded = max(
+                    num_tokens_padded,
+                    num_reqs_padded * max_num_scheduled_tokens,
+                )
+            runner_assisted_full_attn = runner_assisted_full_attn_request is not None
+
             ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
                 should_ubatch,
                 num_scheduled_tokens_np,
@@ -1035,8 +1143,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 num_reqs_padded,
                 self.parallel_config.num_ubatches,
             )
-
-            pad_attn = cudagraph_mode == CUDAGraphMode.FULL
+            pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn
 
             use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
             ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
@@ -1058,11 +1165,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 num_tokens_unpadded=num_tokens_unpadded,
                 ubatch_slices=ubatch_slices_padded,
             )
+            if runner_assisted_full_attn:
+                self._refresh_runner_assisted_full_attention_metadata_buffers(
+                    num_reqs=num_reqs,
+                    num_reqs_padded=num_reqs_padded,
+                    num_scheduled_tokens_np=num_scheduled_tokens_np,
+                )
 
+            attn_num_reqs = num_reqs_padded if runner_assisted_full_attn else num_reqs
             attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
                 num_tokens=num_tokens_unpadded,
                 num_tokens_padded=num_tokens_padded if pad_attn else None,
-                num_reqs=num_reqs,
+                num_reqs=attn_num_reqs,
                 num_reqs_padded=num_reqs_padded if pad_attn else None,
                 max_query_len=max_num_scheduled_tokens,
                 ubatch_slices=ubatch_slices_attn,
@@ -1071,6 +1185,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
                 cascade_attn_prefix_lens=cascade_attn_prefix_lens,
                 slot_mappings=slot_mappings_by_group,
+                for_cudagraph_capture=runner_assisted_full_attn_capture,
             )
             self._maybe_attach_attention_metadata_extensions(
                 attn_metadata=attn_metadata,
@@ -1079,6 +1194,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 max_query_len=max_num_scheduled_tokens,
                 pad_attn=pad_attn,
                 num_scheduled_tokens_np=num_scheduled_tokens_np,
+                for_cudagraph_capture=runner_assisted_full_attn_capture,
             )
 
             (
@@ -1092,14 +1208,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
 
         # Let the model adjust inputs before forward (e.g. restore input_ids
         # for multimodal position detection, fix decode position offsets).
-        if hasattr(self.model, "prepare_runner_inputs"):
-            input_ids, positions = self.model.prepare_runner_inputs(
+        prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None)
+        if callable(prepare_runner_inputs):
+            input_ids, positions = prepare_runner_inputs(
                 input_ids=input_ids,
                 positions=positions,
                 inputs_embeds=inputs_embeds,
                 req_ids=req_ids[:num_reqs],
-                num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)],
-                num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)],
+                num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs],
+                num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
                 input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
             )
 
@@ -1108,46 +1225,60 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # with CUDA graph capture.
         if self.calculate_kv_scales:
             cudagraph_mode = CUDAGraphMode.NONE
+            runner_assisted_full_attn = False
             # Mark KV scales as calculated after the first forward pass
             self.calculate_kv_scales = False
 
+        runner_assisted_context_enabled = False
+        if runner_assisted_full_attn:
+            runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context(
+                enabled=True,
+                num_reqs=num_reqs,
+            )
+            runner_assisted_full_attn = runner_assisted_context_enabled
+
         # Run the model.
         # Use persistent buffers for CUDA graphs.
         # When spec decode is enabled, defer connector finalization
         # (wait_for_save + clear metadata) until after draft model runs.
         defer_kv_connector_finalize = self.speculative_config is not None
-        with (
-            nullcontext(),
-            set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_padded,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_mode,
-                batch_descriptor=batch_desc,
-                ubatch_slices=ubatch_slices_padded,
-                slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
-            ),
-            record_function_or_nullcontext("gpu_model_runner: forward"),
-            self.maybe_get_kv_connector_output(
-                scheduler_output,
-                defer_finalize=defer_kv_connector_finalize,
-            ) as kv_connector_output,
-        ):
-            model_output = self._model_forward(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                **model_kwargs,
-                sampling_metadata=self.input_batch.sampling_metadata,
-                logits_index=logits_indices,
-                sampler=self.sampler,
-            )
+        try:
+            with (
+                nullcontext(),
+                set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_padded,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    cudagraph_runtime_mode=(CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode),
+                    batch_descriptor=batch_desc,
+                    ubatch_slices=ubatch_slices_padded,
+                    slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
+                ),
+                record_function_or_nullcontext("gpu_model_runner: forward"),
+                self.maybe_get_kv_connector_output(
+                    scheduler_output,
+                    defer_finalize=defer_kv_connector_finalize,
+                ) as kv_connector_output,
+            ):
+                model_output = self._model_forward(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **model_kwargs,
+                    sampling_metadata=self.input_batch.sampling_metadata,
+                    logits_index=logits_indices,
+                    sampler=self.sampler,
+                )
 
-            # [Omni] Map pending ropes metadata to req_ids.
-            if hasattr(self.model, "flush_pending_metadata"):
-                self.model.flush_pending_metadata(list(req_ids))
+                # [Omni] Map pending ropes metadata to req_ids.
+                flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None)
+                if callable(flush_pending_metadata):
+                    flush_pending_metadata(req_ids[:num_reqs])
+        finally:
+            if runner_assisted_context_enabled:
+                self._set_runner_assisted_full_attention_metadata_context(enabled=False)
 
         with record_function_or_nullcontext("gpu_model_runner: postprocess"):
             if self.use_aux_hidden_state_outputs:

vllm_omni/worker/runner_assisted_metadata.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import NamedTuple, Protocol
+
+
+class RunnerAssistedFullAttentionMetadataRequest(NamedTuple):
+    num_reqs_padded: int
+    for_cudagraph_capture: bool
+
+
+class RunnerAssistedAttentionMetadataProvider(Protocol):
+    """Optional model hook for runner-built full attention metadata."""
+
+    def get_runner_assisted_full_attention_metadata_request(
+        self,
+        *,
+        req_ids: Sequence[str],
+        num_reqs: int,
+        num_scheduled_tokens: Sequence[int],
+        num_computed_tokens: Sequence[int],
+        max_num_scheduled_tokens: int,
+    ) -> RunnerAssistedFullAttentionMetadataRequest | None: ...
+
+    def set_runner_assisted_full_attention_metadata_context(
+        self,
+        *,
+        enabled: bool,
+        num_reqs: int = 0,
+    ) -> None: ...