Merge branch 'main' into bugfix/qwen3-tts-seed-propagation

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

# Conflicts:
#	tests/entrypoints/openai_api/test_serving_speech.py

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/e2e/online_serving/test_qwen_image_layered_expansion.py

@@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model.
 Kept cases (maximum feature coverage, no cpu-offload):
   sp_001              : cache_dit + Ulysses-SP 2          (2×H100)
   cfg_parallel_002    : cache_dit + CFG-Parallel 2        (2×H100)
-  layers_guard_001    : layerwise offload + layers=3      (1×H100, issue #1969 guard)
+  layers_guard_001    : layerwise offload + layers=2      (1×H100, issue #1969 guard)
 
 Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload.
 """
@@ -116,8 +116,8 @@ LAYERS_GUARD_CASES = [
             model=MODEL,
             server_args=["--enable-layerwise-offload"],
         ),
-        3,
-        id="layers_guard_001_layers3",
+        2,
+        id="layers_guard_001_layers2",
         marks=SINGLE_CARD_FEATURE_MARKS,
     ),
 ]

tests/entrypoints/openai_api/test_image_server.py

@@ -1244,7 +1244,7 @@ def test_parameter_validation():
 
     # Invalid layers for layered models (must stay within the backend-supported range)
     with pytest.raises(ValueError):
-        ImageGenerationRequest(prompt="test", layers=2)
+        ImageGenerationRequest(prompt="test", layers=1)
 
     with pytest.raises(ValueError):
         ImageGenerationRequest(prompt="test", layers=11)
@@ -1663,13 +1663,13 @@ def test_image_edit_invalid_layers(async_omni_test_client):
         files=[("image", img_bytes)],
         data={
             "prompt": "test",
-            "layers": 2,
+            "layers": 1,
         },
     )
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
-    assert "layers must be between 3 and 10 inclusive" in detail
+    assert "layers must be between 2 and 10 inclusive" in detail
 
     # Test layers above the supported range
     response = async_omni_test_client.post(
@@ -1683,7 +1683,7 @@ def test_image_edit_invalid_layers(async_omni_test_client):
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
-    assert "layers must be between 3 and 10 inclusive" in detail
+    assert "layers must be between 2 and 10 inclusive" in detail
 
 
 def test_image_edit_resolution_and_size_conflict(async_omni_test_client):

tests/entrypoints/openai_api/test_serving_speech.py

@@ -39,6 +39,7 @@ from vllm_omni.entrypoints.openai.serving_speech import (
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
+from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
     build_fish_voice_clone_prompt_ids,
@@ -3364,6 +3365,40 @@ class TestTTSAsyncOffloading:
         assert stage0_params.extra_args["tts_local_seed"] == 42
         assert qwen3_tts_server.engine_client.default_sampling_params_list[0].extra_args is None
 
+    def test_prepare_speech_generation_uses_adapter_model_type_label(
+        self,
+        voxtral_server,
+        mocker: MockerFixture,
+    ):
+        """Adapter model_type should replace the legacy _tts_model_type label ladder."""
+        legacy_tts_model_type = "dummy_tts"
+        adapter_model_type = "adapter_dummy_tts"
+
+        class FakeAdapter:
+            def validate(self, request):
+                return None
+
+            async def build(self, request, sampling_params_list, has_inline_ref_audio):
+                return PreparedRequest(
+                    prompt={"prompt": request.input},
+                    tts_params={},
+                    model_type=adapter_model_type,
+                )
+
+        voxtral_server._tts_model_type = legacy_tts_model_type
+        mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter())
+        log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info")
+
+        asyncio.run(voxtral_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="hello")))
+
+        assert adapter_model_type != legacy_tts_model_type
+        assert any(
+            call.args
+            and call.args[0] == "TTS speech request %s: text=%r, model=%s"
+            and call.args[3] == adapter_model_type
+            for call in log_info.call_args_list
+        )
+
     def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture):
         """stream_format=sse should request delta-style multimodal outputs."""
         qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)

tests/helpers/stage_config.py

@@ -633,6 +633,18 @@ def get_deploy_config_path(rel_path: str) -> str:
     return str(_DEPLOY_DIR / rel_path)
 
 
+def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]:
+    """Return one stage entry from a deploy yaml by ``stage_id``."""
+    with open(get_deploy_config_path(rel_path), encoding="utf-8") as f:
+        cfg = yaml.safe_load(f) or {}
+
+    stage_key = "stages" if "stages" in cfg else "stage_args"
+    for stage in cfg.get(stage_key, []):
+        if stage.get("stage_id") == stage_id:
+            return stage
+    raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}")
+
+
 def _get_config_value_by_path(config_dict: dict, path: str) -> Any:
     """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``)."""
     current: Any = config_dict
@@ -707,6 +719,7 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st
 
 __all__ = [
     "get_deploy_config_path",
+    "get_deploy_config_stage",
     "modify_stage_config",
     "stage_config_path_for_run_level",
 ]

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

@@ -14,7 +14,7 @@ from unittest.mock import patch
 import pytest
 from transformers import PretrainedConfig, Qwen3OmniMoeConfig
 
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
 from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
 from vllm_omni.config.stage_config import (
@@ -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:
@@ -739,6 +740,15 @@ stages:
         assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}}
         assert "compilation_config" not in deploy.stages[0].engine_extras
 
+    def test_load_voxcpm2_deploy_config_preserves_engine_extras(self):
+        deploy_path = get_deploy_config_path("voxcpm2.yaml")
+        raw_stage = get_deploy_config_stage("voxcpm2.yaml", 0)
+        expected_runtime_config = raw_stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"]
+
+        deploy = load_deploy_config(deploy_path)
+        runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"]
+        assert runtime_config == expected_runtime_config
+
     def test_merge_pipeline_deploy(self):
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -800,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 = {
@@ -928,6 +958,54 @@ stages:
         assert "foo" in stage.engine_extras
         assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3}
 
+    def test_explicit_engine_extras_merge_with_flat_and_engine_args(self):
+        """Explicit engine_extras should be preserved with existing pass-through styles."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "engine_extras": {
+                        "hf_overrides": {
+                            "runtime_config": {
+                                "explicit_only": True,
+                                "shared": "explicit",
+                                "nested": {"a": 1},
+                            }
+                        }
+                    },
+                    "hf_overrides": {
+                        "runtime_config": {
+                            "flat_only": True,
+                            "shared": "flat",
+                            "nested": {"b": 2},
+                        }
+                    },
+                    "engine_args": {
+                        "hf_overrides": {
+                            "runtime_config": {
+                                "engine_args_only": True,
+                                "shared": "engine_args",
+                                "nested": {"c": 3},
+                            }
+                        }
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert deploy.stages[0].engine_extras["hf_overrides"] == {
+            "runtime_config": {
+                "engine_args_only": True,
+                "explicit_only": True,
+                "flat_only": True,
+                "nested": {"a": 1, "b": 2, "c": 3},
+                "shared": "engine_args",
+            }
+        }
+
     def test_deep_merge_does_not_mutate_inputs(self):
         """Merging engine_args must not mutate the base stage dict."""
         fake_config = {
@@ -1588,6 +1666,197 @@ class TestAuraOmniDeploy:
         assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav"
 
 
+class TestDeployCliOverrideFlow:
+    """Test deploy-YAML baselines overridden by CLI runtime overrides."""
+
+    def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path):
+        deploy_path = tmp_path / "diffusion_stage.yaml"
+        deploy_path.write_text(
+            """
+async_chunk: false
+stages:
+  - stage_id: 0
+    devices: "4,5"
+    parallel_config:
+      pipeline_parallel_size: 1
+      data_parallel_size: 1
+      tensor_parallel_size: 1
+      enable_expert_parallel: false
+      sequence_parallel_size: 1
+      ulysses_degree: 1
+      ring_degree: 1
+      cfg_parallel_size: 1
+      vae_patch_parallel_size: 1
+      use_hsdp: false
+      hsdp_shard_size: -1
+      hsdp_replicate_size: 1
+    engine_args:
+      cache_backend: cache_dit
+      diffusion_attention_backend: FLASH_ATTN
+      diffusion_kv_cache_dtype: auto
+      step_execution: false
+      vae_use_tiling: false
+      enable_cpu_offload: false
+      max_generated_image_size: 1048576
+      tts_max_instructions_length: 1000
+""",
+            encoding="utf-8",
+        )
+
+        pipeline = PipelineConfig(
+            model_type="test_diffusion",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="dit",
+                    execution_type=StageExecutionType.DIFFUSION,
+                    final_output=True,
+                    final_output_type="image",
+                ),
+            ),
+        )
+        deploy = load_deploy_config(deploy_path)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        stage = stages[0]
+
+        assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False
+        assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1
+        assert stage.yaml_engine_args["cache_backend"] == "cache_dit"
+        assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False
+        assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1
+        assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1
+        assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN"
+        assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto"
+        assert stage.yaml_engine_args["step_execution"] is False
+        assert stage.yaml_engine_args["vae_use_tiling"] is False
+        assert stage.yaml_engine_args["enable_cpu_offload"] is False
+        assert stage.yaml_engine_args["max_generated_image_size"] == 1048576
+        assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000
+
+        stage.runtime_overrides = StageConfigFactory._merge_cli_overrides(
+            stage,
+            {
+                "pipeline_parallel_size": 2,
+                "data_parallel_size": 3,
+                "tensor_parallel_size": 4,
+                "enable_expert_parallel": True,
+                "sequence_parallel_size": 24,
+                "ulysses_degree": 2,
+                "ring_degree": 4,
+                "cfg_parallel_size": 2,
+                "vae_patch_parallel_size": 2,
+                "cache_backend": "tea_cache",
+                "use_hsdp": True,
+                "hsdp_shard_size": 8,
+                "hsdp_replicate_size": 2,
+                "diffusion_attention_backend": "SAGE_ATTN",
+                "diffusion_kv_cache_dtype": "fp8",
+                "step_execution": True,
+                "vae_use_tiling": True,
+                "enable_cpu_offload": True,
+                "max_generated_image_size": 2097152,
+                "tts_max_instructions_length": 2000,
+            },
+        )
+
+        omega_config = stage.to_omegaconf()
+
+        assert omega_config.engine_args.cache_backend == "tea_cache"
+        assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN"
+        assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8"
+        assert omega_config.engine_args.step_execution is True
+        assert omega_config.engine_args.vae_use_tiling is True
+        assert omega_config.engine_args.enable_cpu_offload is True
+        assert omega_config.engine_args.max_generated_image_size == 2097152
+        assert omega_config.engine_args.tts_max_instructions_length == 2000
+        assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.data_parallel_size == 3
+        assert omega_config.engine_args.parallel_config.tensor_parallel_size == 4
+        assert omega_config.engine_args.parallel_config.enable_expert_parallel is True
+        assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24
+        assert omega_config.engine_args.parallel_config.ulysses_degree == 2
+        assert omega_config.engine_args.parallel_config.ring_degree == 4
+        assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.use_hsdp is True
+        assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8
+        assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2
+
+    def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path):
+        deploy_path = tmp_path / "llm_stage.yaml"
+        deploy_path.write_text(
+            """
+async_chunk: false
+stages:
+  - stage_id: 0
+    devices: "0,1"
+    engine_args:
+      tensor_parallel_size: 1
+      enable_expert_parallel: false
+      gpu_memory_utilization: 0.5
+      max_num_seqs: 16
+      max_num_batched_tokens: 1024
+      max_model_len: 4096
+      enforce_eager: false
+""",
+            encoding="utf-8",
+        )
+
+        pipeline = PipelineConfig(
+            model_type="test_llm",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="thinker",
+                    execution_type=StageExecutionType.LLM_AR,
+                    final_output=True,
+                    final_output_type="text",
+                ),
+            ),
+        )
+        deploy = load_deploy_config(deploy_path)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        stage = stages[0]
+
+        assert stage.yaml_engine_args["tensor_parallel_size"] == 1
+        assert stage.yaml_engine_args["enable_expert_parallel"] is False
+        assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5
+        assert stage.yaml_engine_args["max_num_seqs"] == 16
+        assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024
+        assert stage.yaml_engine_args["max_model_len"] == 4096
+        assert stage.yaml_engine_args["enforce_eager"] is False
+
+        stage.runtime_overrides = StageConfigFactory._merge_cli_overrides(
+            stage,
+            {
+                "tensor_parallel_size": 2,
+                "enable_expert_parallel": True,
+                "gpu_memory_utilization": 0.9,
+                "max_num_seqs": 32,
+                "max_num_batched_tokens": 2048,
+                "max_model_len": 8192,
+                "enforce_eager": True,
+            },
+        )
+
+        omega_config = stage.to_omegaconf()
+
+        assert omega_config.engine_args.tensor_parallel_size == 2
+        assert omega_config.engine_args.enable_expert_parallel is True
+        assert omega_config.engine_args.gpu_memory_utilization == 0.9
+        assert omega_config.engine_args.max_num_seqs == 32
+        assert omega_config.engine_args.max_num_batched_tokens == 2048
+        assert omega_config.engine_args.max_model_len == 8192
+        assert omega_config.engine_args.enforce_eager is True
+
+
 class TestSentinelDefaultPrecedence:
     """Caller-typed (non-None) values win over YAML; None values fall through
     to YAML / dataclass defaults (#3035)."""

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)