[Fix] Align VoxCPM2 graph scheduler gate

Signed-off-by: Sy03 <1370724210@qq.com>

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -6,6 +6,7 @@ import pytest
 from vllm.v1.request import RequestStatus
 
 import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod
+import vllm_omni.model_executor.models.voxcpm2.scheduler as voxcpm2_scheduler_mod
 from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -51,14 +52,30 @@ class _MockRequest:
         return RequestStatus.is_finished(self.status)
 
 
-def _make_scheduler(*, enable_unified_decode_graph: bool = True) -> VoxCPM2OmniARAsyncScheduler:
+@pytest.fixture(autouse=True)
+def _mock_cuda_graph_platform(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True)
+
+
+def _make_scheduler(
+    *,
+    enable_unified_decode_graph: bool | None = True,
+    deterministic_cfm_noise: bool = False,
+    runtime_config_on_model_config: bool = False,
+) -> 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),
+    hf_config = SimpleNamespace()
+    model_config = SimpleNamespace(hf_config=hf_config)
+    if enable_unified_decode_graph is not None or deterministic_cfm_noise:
+        runtime_config = SimpleNamespace(
+            enable_unified_decode_graph=True if enable_unified_decode_graph is None else enable_unified_decode_graph,
+            deterministic_cfm_noise=deterministic_cfm_noise,
         )
-    )
+        if runtime_config_on_model_config:
+            model_config.voxcpm2_runtime_config = runtime_config
+        else:
+            hf_config.voxcpm2_runtime_config = runtime_config
+    sched.vllm_config = SimpleNamespace(model_config=model_config)
     return sched
 
 
@@ -86,6 +103,39 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None:
     assert not scheduler._should_defer_waiting_for_unified_decode_graph()
 
 
+def test_voxcpm2_unified_decode_graph_uses_model_runtime_defaults() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=None)
+    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_reads_model_config_runtime_config() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=True, runtime_config_on_model_config=True)
+    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_with_deterministic_noise() -> None:
+    scheduler = _make_scheduler(deterministic_cfm_noise=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False)
+    scheduler = _make_scheduler()
+    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")]

tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py

@@ -12,6 +12,7 @@ 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_RUNTIME_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "runtime_config.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"
@@ -74,6 +75,7 @@ def test_ar_runner_without_model_hook_stays_on_normal_path():
 
 def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
     source = TALKER.read_text()
+    runtime_config_source = VOXCPM2_RUNTIME_CONFIG.read_text()
     compact_source = "".join(source.split())
 
     assert "_voxcpm2_compile_without_inductor_cudagraphs" in source
@@ -86,8 +88,8 @@ def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
     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 "self._enable_unified_decode_graph=self._runtime_config.unified_decode_graph_available(" in compact_source
+    assert "andnotself.deterministic_cfm_noise" in "".join(runtime_config_source.split())
     assert "decode_tail_graph" not in source
 
     unified_source = source[source.index("def _forward_unified_decode") :]
@@ -183,7 +185,9 @@ def test_voxcpm2_scheduler_policy_stays_model_local():
     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 "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source
+    assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source
+    assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" 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

vllm_omni/model_executor/models/voxcpm2/runtime_config.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass(frozen=True)
+class _VoxCPM2RuntimeConfig:
+    enable_profiling: bool = False
+    enable_nvtx_profile: bool = False
+    enable_loc_dit_layer_nvtx: bool = False
+    enable_loc_dit_fused_qkv: bool = True
+    enable_loc_dit_fused_mlp: bool = True
+    enable_loc_dit_zero_dt_cache: bool = True
+    enable_loc_dit_fast_rope: bool = False
+    enable_loc_dit_skip_qkv_contig: bool = True
+    enable_loc_dit_reduce_overhead_no_cg: bool = False
+    enable_loc_dit_fullgraph_no_cg: bool = False
+    cfg_cutoff_ratio: float = 1.0
+    decode_graph_capture_policy: str = "all"
+    enable_vae_cuda_graph: bool = False
+    enable_cfm_cuda_graph: bool = False
+    enable_cfm_prealloc_output: bool = False
+    enable_batched_cfm: bool = True
+    deterministic_cfm_noise: bool = False
+    deterministic_cfm_seed: int = 20260601
+    audio_emit_every: int = 1
+    vae_decode_every: int = 3
+    enable_delayed_audio_copy: bool = False
+    delayed_audio_copy_use_events: bool = False
+    coalesce_audio_d2h: bool = True
+    enable_batched_vae_decode: bool = True
+    enable_batched_fsq_fusion: bool = False
+    batched_fsq_fusion_max_batch: int = 32
+    enable_batched_prefill_tail: bool = False
+    enable_unified_decode_graph: bool = True
+    unified_decode_graph_max_batch_size: int = 64
+
+    @classmethod
+    def from_vllm_config(cls, vllm_config: Any) -> _VoxCPM2RuntimeConfig:
+        model_config = vllm_config.model_config
+        raw = getattr(model_config, "voxcpm2_runtime_config", None)
+        if raw is None:
+            raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None)
+        if raw is None:
+            return cls()
+        if hasattr(raw, "to_dict"):
+            raw = raw.to_dict()
+        elif not isinstance(raw, dict) and hasattr(raw, "__dict__"):
+            raw = vars(raw)
+        if not isinstance(raw, dict):
+            logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw)
+            return cls()
+
+        fields = {field.name: field for field in dataclasses.fields(cls)}
+        values: dict[str, Any] = {}
+        for key, value in raw.items():
+            if key not in fields:
+                logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key)
+                continue
+            default = getattr(cls(), key)
+            values[key] = cls._coerce_value(key, value, default)
+        return cls(**values)._normalized()
+
+    def _normalized(self) -> _VoxCPM2RuntimeConfig:
+        cfg = self
+        if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph:
+            logger.warning(
+                "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; "
+                "disabling CFM CUDA Graph and keeping batched CFM."
+            )
+            cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False)
+        if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph:
+            logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.")
+            cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False)
+        if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy:
+            logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.")
+            cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False)
+        if cfg.enable_batched_vae_decode and (
+            cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph
+        ):
+            logger.warning(
+                "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, "
+                "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode."
+            )
+            cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False)
+        if cfg.enable_unified_decode_graph:
+            if cfg.enable_cfm_cuda_graph:
+                logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.")
+                cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False)
+            if not cfg.enable_batched_cfm:
+                logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.")
+                cfg = dataclasses.replace(cfg, enable_batched_cfm=True)
+        return cfg
+
+    def unified_decode_graph_available(self, *, use_cuda_graph: bool) -> bool:
+        return bool(use_cuda_graph and self.enable_unified_decode_graph and not self.deterministic_cfm_noise)
+
+    @staticmethod
+    def _coerce_value(key: str, value: Any, default: Any) -> Any:
+        if isinstance(default, bool):
+            if isinstance(value, str):
+                return value.strip().lower() in ("1", "true", "yes", "on")
+            return bool(value)
+        if isinstance(default, int) and not isinstance(default, bool):
+            value = int(value)
+            if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}:
+                return max(1, value)
+            return value
+        if isinstance(default, float):
+            if key == "cfg_cutoff_ratio":
+                return min(1.0, max(0.0, float(value)))
+            return float(value)
+        if isinstance(default, str):
+            return str(value)
+        return value

vllm_omni/model_executor/models/voxcpm2/scheduler.py

@@ -6,6 +6,9 @@ from __future__ import annotations
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler
+from vllm_omni.platforms import current_omni_platform
+
+from .runtime_config import _VoxCPM2RuntimeConfig
 
 
 class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
@@ -19,11 +22,8 @@ class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
     """
 
     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))
+        runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)
+        return runtime_config.unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())
 
     def _should_defer_waiting_for_unified_decode_graph(self) -> bool:
         if not self._unified_decode_graph_enabled():

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py

@@ -46,6 +46,7 @@ from vllm_omni.utils.speaker_cache import (
 from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 
 from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM
+from .runtime_config import _VoxCPM2RuntimeConfig
 from .voxcpm2_import_utils import import_voxcpm2_core
 
 logger = init_logger(__name__)
@@ -104,115 +105,6 @@ class _DecodeResidualMeta(NamedTuple):
     new_lm_hidden: torch.Tensor
 
 
-@dataclasses.dataclass(frozen=True)
-class _VoxCPM2RuntimeConfig:
-    enable_profiling: bool = False
-    enable_nvtx_profile: bool = False
-    enable_loc_dit_layer_nvtx: bool = False
-    enable_loc_dit_fused_qkv: bool = True
-    enable_loc_dit_fused_mlp: bool = True
-    enable_loc_dit_zero_dt_cache: bool = True
-    enable_loc_dit_fast_rope: bool = False
-    enable_loc_dit_skip_qkv_contig: bool = True
-    enable_loc_dit_reduce_overhead_no_cg: bool = False
-    enable_loc_dit_fullgraph_no_cg: bool = False
-    cfg_cutoff_ratio: float = 1.0
-    decode_graph_capture_policy: str = "all"
-    enable_vae_cuda_graph: bool = False
-    enable_cfm_cuda_graph: bool = False
-    enable_cfm_prealloc_output: bool = False
-    enable_batched_cfm: bool = True
-    deterministic_cfm_noise: bool = False
-    deterministic_cfm_seed: int = 20260601
-    audio_emit_every: int = 1
-    vae_decode_every: int = 3
-    enable_delayed_audio_copy: bool = False
-    delayed_audio_copy_use_events: bool = False
-    coalesce_audio_d2h: bool = True
-    enable_batched_vae_decode: bool = True
-    enable_batched_fsq_fusion: bool = False
-    batched_fsq_fusion_max_batch: int = 32
-    enable_batched_prefill_tail: bool = False
-    enable_unified_decode_graph: bool = True
-    unified_decode_graph_max_batch_size: int = 64
-
-    @classmethod
-    def from_vllm_config(cls, vllm_config: VllmConfig) -> _VoxCPM2RuntimeConfig:
-        model_config = vllm_config.model_config
-        raw = getattr(model_config, "voxcpm2_runtime_config", None)
-        if raw is None:
-            raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None)
-        if raw is None:
-            return cls()
-        if hasattr(raw, "to_dict"):
-            raw = raw.to_dict()
-        elif not isinstance(raw, dict) and hasattr(raw, "__dict__"):
-            raw = vars(raw)
-        if not isinstance(raw, dict):
-            logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw)
-            return cls()
-
-        fields = {field.name: field for field in dataclasses.fields(cls)}
-        values: dict[str, Any] = {}
-        for key, value in raw.items():
-            if key not in fields:
-                logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key)
-                continue
-            default = getattr(cls(), key)
-            values[key] = cls._coerce_value(key, value, default)
-        return cls(**values)._normalized()
-
-    def _normalized(self) -> _VoxCPM2RuntimeConfig:
-        cfg = self
-        if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph:
-            logger.warning(
-                "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; "
-                "disabling CFM CUDA Graph and keeping batched CFM."
-            )
-            cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False)
-        if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph:
-            logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.")
-            cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False)
-        if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy:
-            logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.")
-            cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False)
-        if cfg.enable_batched_vae_decode and (
-            cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph
-        ):
-            logger.warning(
-                "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, "
-                "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode."
-            )
-            cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False)
-        if cfg.enable_unified_decode_graph:
-            if cfg.enable_cfm_cuda_graph:
-                logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.")
-                cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False)
-            if not cfg.enable_batched_cfm:
-                logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.")
-                cfg = dataclasses.replace(cfg, enable_batched_cfm=True)
-        return cfg
-
-    @staticmethod
-    def _coerce_value(key: str, value: Any, default: Any) -> Any:
-        if isinstance(default, bool):
-            if isinstance(value, str):
-                return value.strip().lower() in ("1", "true", "yes", "on")
-            return bool(value)
-        if isinstance(default, int) and not isinstance(default, bool):
-            value = int(value)
-            if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}:
-                return max(1, value)
-            return value
-        if isinstance(default, float):
-            if key == "cfg_cutoff_ratio":
-                return min(1.0, max(0.0, float(value)))
-            return float(value)
-        if isinstance(default, str):
-            return str(value)
-        return value
-
-
 def is_cjk_char(c: str) -> bool:
     """Check if a character is a CJK ideograph."""
     cp = ord(c)
@@ -1028,8 +920,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         self._cuda_graph_warmup_steps = 0
         self._cuda_graph_warmup_threshold = 3
         self._unified_graphs: dict[int, _CapturedUnifiedDecodeGraph] = {}
-        self._enable_unified_decode_graph = (
-            use_cuda_graph and self._runtime_config.enable_unified_decode_graph and not self._deterministic_cfm_noise
+        self._enable_unified_decode_graph = self._runtime_config.unified_decode_graph_available(
+            use_cuda_graph=use_cuda_graph
         )
         self._unified_decode_graph_max_batch_size = self._runtime_config.unified_decode_graph_max_batch_size
         self._unified_graph_bucket_sizes: frozenset[int] = self._build_unified_graph_bucket_sizes(