Merge branch 'main' into feat/diffusion-prefetch-protection

docs/design/qwen3_omni_tts_performance_optimization.md

@@ -411,6 +411,8 @@ Notes:
 - `runtime.max_batch_size` controls stage-level batching.
 - Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
 - Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
+- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
+ than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
 
 #### 2) Enable async chunk
 

docs/mkdocs/hooks/generate_argparse.py

@@ -119,9 +119,10 @@ def extract_omni_serve_subparser_init():
                     exec_globals = {
                         "_FlexibleArgumentParser": _FlexibleArgumentParser,
                         "FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "TrackingArgumentParser": _FlexibleArgumentParser,  # only needs add_argument, not tracking
+                        "TrackingNamespace": None,  # type annotation only
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
                         "_ensure_vllm_platform": lambda: None,  # no-op for doc
-                        "nullify_stage_engine_defaults": lambda parser: None,  # no-op for doc
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -17,6 +17,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 
 The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
 automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
+The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
+Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
+backend.
 Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
 
 To explicitly utilize a custom deployment YAML, specify the configuration path:
@@ -72,6 +75,12 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
 ```
 
+To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
+```bash
+VLLM_USE_FLASHINFER_MOE_FP16=1 \
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
 For the stage-based CLI, you usually do **not** need `--stage-overrides` for
 that kind of change. Since each command launches one stage, just pass the knob
 directly on that stage command:

examples/offline_inference/bagel/README.md

@@ -195,6 +195,101 @@ stages:
     devices: "0,1"
 ```
 
+### VAE Patch Parallelism
+
+[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
+
+**Bagel-specific notes:**
+
+- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
+- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
+- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
+- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
+
+**Single-stage via deploy YAML** (recommended for `end2end.py`):
+
+```yaml
+pipeline: bagel_single_stage
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_batched_tokens: 32768
+    max_num_seqs: 1
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    devices: "0,1"
+    vae_use_tiling: true
+    parallel_config:
+      tensor_parallel_size: 2
+      vae_patch_parallel_size: 2
+    default_sampling_params:
+      seed: 52
+```
+
+```bash
+cd examples/offline_inference/bagel
+
+CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
+    --model /path/to/BAGEL-7B-MoT \
+    --deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
+    --modality text2img \
+    --prompts "A cute cat" \
+    --steps 10 \
+    --output ./out_vae_pp
+```
+
+**Single-stage via `Omni` kwargs** (same flags as online serving):
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+omni = Omni(
+    model="ByteDance-Seed/BAGEL-7B-MoT",
+    deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
+    tensor_parallel_size=2,
+    vae_patch_parallel_size=2,
+    vae_use_tiling=True,
+)
+# Then call omni.generate(...) as in end2end.py
+```
+
+**Two-stage (VAE PP on DiT only):**
+
+```yaml
+stages:
+  - stage_id: 0
+    devices: "0"
+    # AR Thinker — no vae_patch_parallel here
+
+  - stage_id: 1
+    devices: "0,1"
+    vae_use_tiling: true
+    parallel_config:
+      tensor_parallel_size: 2
+      vae_patch_parallel_size: 2
+```
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+    --deploy-config /path/to/bagel_vae_pp.yaml \
+    --modality text2img \
+    --prompts "A cute cat"
+```
+
+**Startup log checks:**
+
+```text
+INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
+```
+
+| Setting | Role |
+| :------ | :--- |
+| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
+| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
+| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
+
 #### Hybrid Sharded Data Parallel (HSDP)
 
 For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.

examples/offline_inference/bagel/end2end.py

@@ -121,9 +121,6 @@ def parse_args():
         help="Temperature for text generation sampling (default: 0.3).",
     )
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
     return args
 
@@ -154,6 +151,7 @@ def main():
     from vllm_omni.entrypoints.omni import Omni
 
     omni_kwargs = vars(args).copy()
+    omni_kwargs.pop("model", None)
     deploy_config = args.deploy_config
     if args.think and deploy_config is None:
         deploy_config = "vllm_omni/deploy/bagel_think.yaml"
@@ -164,9 +162,7 @@ def main():
     if args.quantization:
         omni_kwargs["quantization_config"] = args.quantization
 
-    # Override CLI --model with the derived model_name.
-    omni_kwargs["model"] = model_name
-    omni = Omni(**omni_kwargs)
+    omni = Omni(model=model_name, **omni_kwargs)
 
     formatted_prompts = []
     for p in prompts:

examples/offline_inference/dynin_omni/end2end.py

@@ -18,8 +18,6 @@ import numpy as np
 import torch
 from PIL import Image
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
 TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
 
 TASK_DEFAULT_RUNTIME = {
@@ -972,8 +970,6 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
     parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
 
     parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/hunyuan_image3/end2end.py

@@ -107,9 +107,6 @@ def parse_args():
         ),
     )
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -19,7 +19,6 @@ from vllm import SamplingParams
 from vllm.multimodal.image import convert_image_mode
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 DEFAULT_SYSTEM = "You are a helpful assistant."
 DEFAULT_QUESTION = "Please summarize the content of this image."
@@ -49,7 +48,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -29,7 +29,6 @@ from PIL import Image
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
@@ -118,7 +117,6 @@ def parse_args() -> argparse.Namespace:
     )
     p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
     p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
-    nullify_stage_engine_defaults(p)
     args = p.parse_args()
     if not args.prompt:
         args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]

examples/offline_inference/mimo_audio/end2end.py

@@ -21,11 +21,10 @@ from message_convert import (
     to_prompt,
 )
 from vllm import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniTokensPrompt
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -343,7 +342,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model-name",
         "-m",
@@ -439,7 +438,6 @@ def parse_args():
         "vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/ming_flash_omni/end2end.py

@@ -17,14 +17,13 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 import vllm_omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 
 # Imports the processor also registers itself
 from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor  # noqa: F401
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
@@ -397,7 +396,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
+    parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
     parser.add_argument(
         "--query-type",
         "-q",
@@ -496,7 +495,6 @@ def parse_args():
         help="Output directory for results.",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -19,10 +19,9 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
 from vllm.sampling_params import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -432,7 +431,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model",
         type=str,
@@ -554,7 +553,6 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -19,10 +19,9 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -422,7 +421,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model",
         type=str,
@@ -561,7 +560,6 @@ def parse_args():
         help="Model dtype (auto, half, float16, bfloat16, float, float32).",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -39,10 +39,9 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -383,7 +382,9 @@ async def run_all(args):
     print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
     async_omni = None
     try:
-        async_omni = AsyncOmni(**vars(args))
+        # ``from_cli_args`` forwards only explicitly-passed CLI args so
+        # argparse defaults do not silently override deploy YAML values.
+        async_omni = AsyncOmni.from_cli_args(args, model=args.model)
 
         # Use default sampling params from stage config (they are pre-configured
         # in the YAML for each stage).
@@ -453,7 +454,7 @@ async def run_all(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(
+    parser = TrackingArgumentParser(
         description=(
             "Offline inference with async_chunk enabled via AsyncOmni. "
             "Downstream stages start before upstream stages finish, "
@@ -589,7 +590,6 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/sensenova_u1/end2end.py

@@ -29,7 +29,6 @@ import os
 
 from PIL import Image
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -168,7 +167,6 @@ def parse_args():
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -307,9 +307,6 @@ def parse_args() -> argparse.Namespace:
         help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
     )
     current_omni_platform.pre_register_and_update(parser)
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -10,7 +10,6 @@ import soundfile as sf
 from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
@@ -65,7 +64,6 @@ def run_e2e():
         required=True,
         help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
     )
-    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
     # Ensure tokenizer directory exists
     if not os.path.exists(args.tokenizer):

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -26,16 +26,14 @@ import torch
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
 from vllm_omni import AsyncOmni, Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     build_fish_text_only_prompt_ids,
     estimate_fish_voice_clone_prompt_len_from_normalized,
     normalize_fish_voice_clone_texts,
 )
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -211,7 +209,7 @@ async def main_streaming(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(
+    parser = TrackingArgumentParser(
         description="Fish Speech S2 Pro offline inference with vLLM Omni",
     )
     parser.add_argument(
@@ -267,7 +265,6 @@ def parse_args():
         default=False,
         help="Stream audio chunks as they arrive via AsyncOmni.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py

@@ -30,10 +30,9 @@ from pathlib import Path
 import numpy as np
 import soundfile as sf
 import torch
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SAMPLE_RATE = 24_000
 DEFAULT_TEXTS = (
@@ -43,7 +42,7 @@ DEFAULT_TEXTS = (
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
+    parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
     parser.add_argument(
         "--model",
         type=str,
@@ -88,7 +87,6 @@ def parse_args():
         default=None,
         help="Transcript of the reference clip. Required when --ref-audio is set.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py

@@ -5,11 +5,11 @@ from typing import Any
 
 import soundfile as sf
 import torch
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
@@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
+    parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
     parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
     parser.add_argument(
         "--deploy-config",
@@ -89,7 +89,6 @@ def parse_args():
     parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
     parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py

@@ -38,7 +38,8 @@ from pathlib import Path
 import soundfile as sf
 import torch
 from vllm import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 # Prevent multiprocessing from re-importing CUDA in the wrong context.
 os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
@@ -155,7 +156,7 @@ def main(args) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
+    parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference")
     parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
     parser.add_argument(
         "--ref-audio",

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -21,7 +21,6 @@ import os
 import numpy as np
 import soundfile as sf
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -86,7 +85,6 @@ def run_e2e():
         default=None,
         help="Random seed for generation",
     )
-    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
 
     if not os.path.exists(args.stage_config):

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

@@ -15,10 +15,8 @@ import torch
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
 from vllm_omni import AsyncOmni, Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -435,7 +433,7 @@ async def main_streaming(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",
@@ -548,7 +546,6 @@ def parse_args():
         help="Number of prompts per batch (default: 1, sequential).",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/voxcpm2/end2end.py

@@ -12,17 +12,16 @@ from pathlib import Path
 
 import soundfile as sf
 import torch
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 REPO_ROOT = Path(__file__).resolve().parents[4]
 SAMPLE_RATE = 48_000
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference")
+    parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference")
     parser.add_argument(
         "--model",
         type=str,
@@ -60,7 +59,6 @@ def parse_args():
         default=None,
         help="Optional transcript of --ref-audio (enables continuation mode).",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/voxtral_tts/end2end.py

@@ -27,11 +27,10 @@ except ImportError:
         "pip install -e ."
     )
 from vllm import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import AsyncOmni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -230,7 +229,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
 
 
 def parse_args() -> Namespace:
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
     parser.add_argument(
         "--model",
         type=str,
@@ -313,7 +312,6 @@ def parse_args() -> Namespace:
         default=None,
         help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/online_serving/bagel/README.md

@@ -68,6 +68,67 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size
 
 Or set `tensor_parallel_size` per stage in a custom deploy YAML.
 
+### VAE Patch Parallelism
+
+[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck.
+
+**Scope for Bagel:**
+
+| Topology | VAE patch parallel |
+| :------- | :----------------- |
+| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) |
+| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated |
+
+**Requirements:**
+
+- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline).
+- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage).
+- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup.
+
+VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`).
+
+**Online serving (single-stage, 2 GPUs):**
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
+    --tensor-parallel-size 2 \
+    --vae-patch-parallel-size 2 \
+    --vae-use-tiling
+```
+
+**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example:
+
+```yaml
+stages:
+  - stage_id: 0
+    devices: "0"
+    # Thinker (AR) — no VAE patch parallel here
+
+  - stage_id: 1
+    devices: "0,1"
+    vae_use_tiling: true
+    parallel_config:
+      tensor_parallel_size: 2
+      vae_patch_parallel_size: 2
+```
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --deploy-config /path/to/bagel_vae_pp.yaml
+```
+
+**Verify it is active** (check server logs at startup):
+
+```text
+INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
+```
+
+| CLI flag | Default | Description |
+| :------- | :------ | :---------- |
+| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). |
+| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). |
+
 #### Hybrid Sharded Data Parallel (HSDP)
 
 For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
@@ -323,6 +384,8 @@ python openai_chat_client.py \
 | `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
 | `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
 | `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
+| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) |
+| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) |
 | `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
 | `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
 

examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py

@@ -12,7 +12,8 @@ from typing import Any
 import requests
 from openai import OpenAI
 from vllm.assets.audio import AudioAsset
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 # Modify OpenAI's API key and API base to use vLLM's API server.
 openai_api_key = "EMPTY"
@@ -493,7 +494,7 @@ def run_multimodal_generation(args) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

@@ -6,7 +6,8 @@ from typing import NamedTuple
 import requests
 from openai import OpenAI
 from vllm.assets.audio import AudioAsset
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -478,7 +479,7 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",

tests/comfyui/test_comfyui_integration.py

@@ -30,12 +30,12 @@ from PIL import Image
 from pytest_mock import MockerFixture
 from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -444,7 +444,7 @@ def mock_async_omni(
 @pytest.fixture
 def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
     """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     cmd = OmniServeCommand()
     cmd.subparser_init(subparsers)

tests/e2e/online_serving/test_bagel_expansion.py

@@ -9,6 +9,8 @@ Coverage:
 - Ulysses-SP
 - Ring-Attention
 - Layerwise Offloading
+- Hybrid Sharded Data Parallel
+- Tensor Parallelism + VAE Patch Parallelism
 
 assert_diffusion_response validates successful generation and the expected
 512x512 resolution.
@@ -38,12 +40,28 @@ BAGEL_HSDP_2_DEPLOY = modify_stage_config(
     BAGEL_CI_DEPLOY,
     updates={"stages": {0: {"devices": "0"}, 1: {"devices": "0,1"}}},
 )
+BAGEL_TP_VAE_PP_2_DEPLOY = modify_stage_config(
+    BAGEL_CI_DEPLOY,
+    updates={
+        "stages": {
+            0: {"devices": "0"},
+            1: {
+                "devices": "0,1",
+                "parallel_config": {
+                    "tensor_parallel_size": 2,
+                    "vae_patch_parallel_size": 2,
+                },
+            },
+        },
+    },
+)
 
 
 def _get_diffusion_feature_cases(model: str):
     """Return L4 diffusion feature cases for Bagel.
     TeaCache, Cache-DiT, CFG-Parallel,
-    Ulysses-SP, Ring-Attention, Layerwise Offloading.
+    Ulysses-SP, Ring-Attention, Layerwise Offloading,
+    Hybrid Sharded Data Parallel, Tensor Parallelism, VAE Patch Parallelism.
     """
 
     return [
@@ -135,6 +153,15 @@ def _get_diffusion_feature_cases(model: str):
             id="parallel_hsdp_2",
             marks=HSDP_2_FEATURE_MARKS,
         ),
+        # Tensor Parallelism (TP) + VAE Patch Parallelism (size=2)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                stage_config_path=BAGEL_TP_VAE_PP_2_DEPLOY,
+            ),
+            id="tp_vae_patch_parallel_2",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        ),
     ]
 
 
@@ -157,6 +184,7 @@ def test_bagel(
     - Ring-Attention (degree=2)
     - Layerwise Offloading
     - Hybrid Sharded Data Parallel (size=2)
+    - Tensor Parallelism (TP) + VAE Patch Parallelism (size=2)
 
     Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py,
     which checks output dimensions and basic correctness.

tests/engine/test_async_omni_engine_stage_init.py

@@ -640,38 +640,6 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
     assert captured_timeout == 302
 
 
-def test_async_omni_engine_reads_tokenizer_from_engine_args(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    class DummyThread:
-        def __init__(self, *args, **kwargs):
-            pass
-
-        def start(self):
-            pass
-
-        def is_alive(self):
-            return False
-
-    fake_engine_args = types.SimpleNamespace(
-        _explicit_fields=frozenset({"tokenizer"}),
-        explicit_kwargs=lambda: {"tokenizer": "/tokenizer/from-engine-args"},
-    )
-
-    monkeypatch.setattr(engine_mod.threading, "Thread", DummyThread)
-    monkeypatch.setattr(
-        AsyncOmniEngine,
-        "_resolve_stage_configs",
-        lambda self, model, kwargs: ("dummy-config", [types.SimpleNamespace(engine_args={})]),
-    )
-    monkeypatch.setattr(AsyncOmniEngine, "_wait_for_orchestrator_init", lambda *_, **__: None)
-
-    engine = AsyncOmniEngine("dummy-model", engine_args=fake_engine_args)
-
-    assert engine.tokenizer == "/tokenizer/from-engine-args"
-    engine.shutdown()
-
-
 def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path):
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 

tests/engine/test_single_stage_mode.py

@@ -403,39 +403,6 @@ class TestSingleStageModeDetection:
         assert engine.single_stage_mode is True
         assert engine._single_stage_id_filter is None
 
-    def test_engine_args_create_only_forwards_explicit_fields(self, mocker: MockerFixture):
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-
-        captured: dict[str, Any] = {}
-
-        def fake_resolve(self, model: str, kwargs: dict[str, Any]):
-            captured.update(kwargs)
-            return "/fake/path", [_make_stage_cfg(0)]
-
-        mocker.patch.object(AsyncOmniEngine, "_resolve_stage_configs", fake_resolve)
-        mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator")
-        mock_thread_cls = mocker.patch("threading.Thread")
-        mock_future_cls = mocker.patch("concurrent.futures.Future")
-        mock_future = mocker.Mock()
-        mock_future.result.return_value = mocker.Mock()
-        mock_future_cls.return_value = mock_future
-        mock_thread = mocker.Mock()
-        mock_thread.is_alive.return_value = False
-        mock_thread_cls.return_value = mock_thread
-
-        ea = OmniEngineArgs.create(model="ignored", gpu_memory_utilization=0.5)
-        AsyncOmniEngine(model="fake-model", engine_args=ea)
-
-        assert captured["gpu_memory_utilization"] == 0.5
-        assert "model" not in captured
-        assert "max_num_seqs" not in captured
-
-    def test_bare_engine_args_rejected(self, mocker: MockerFixture):
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-
-        with pytest.raises(TypeError, match="OmniEngineArgs.create"):
-            self._make_engine_no_thread(mocker, engine_args=OmniEngineArgs(model="fake-model"))
-
     def test_master_address_and_port_stored(self, mocker: MockerFixture):
         engine = self._make_engine_no_thread(
             mocker,

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -4,12 +4,11 @@
 from types import SimpleNamespace
 
 import pytest
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.config.stage_config import deploy_override_field_names
 from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-from vllm_omni.entrypoints.cli.serve import OmniServeCommand, _create_default_diffusion_stage_cfg
+from vllm_omni.entrypoints.cli.serve import OmniServeCommand
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -34,15 +33,6 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["model_stage"] == "diffusion"
 
 
-def test_default_stage_config_ignores_none_deploy_overrides():
-    """Ensure nullified deploy override defaults do not alter diffusion defaults."""
-    baseline = AsyncOmniEngine._create_default_diffusion_stage_cfg({})[0]
-    nullified_overrides = {name: None for name in deploy_override_field_names()}
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(nullified_overrides)[0]
-
-    assert stage_cfg == baseline
-
-
 def test_default_cache_config_used_when_missing():
     """Ensure default cache_config is synthesized when only backend is given."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -221,7 +211,7 @@ def test_default_stage_config_whitelist_none_fallback():
 
 def test_serve_cli_accepts_ulysses_mode():
     """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -237,7 +227,8 @@ def test_serve_cli_accepts_ulysses_mode():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
 
     assert args.ulysses_mode == "advanced_uaa"
@@ -247,7 +238,7 @@ def test_serve_cli_accepts_ulysses_mode():
 
 def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     """Ensure diffusion serve CLI exposes the profiler switch."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -260,7 +251,8 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
 
     assert args.enable_diffusion_pipeline_profiler is True
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
@@ -268,7 +260,7 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
 
 def test_serve_cli_accepts_diffusion_attention_backend():
     """Ensure diffusion serve CLI exposes the shorthand backend flag."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -282,7 +274,8 @@ def test_serve_cli_accepts_diffusion_attention_backend():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
     diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
 
     assert args.diffusion_attention_backend == "FLASH_ATTN"
@@ -293,7 +286,7 @@ def test_serve_cli_accepts_diffusion_attention_backend():
 
 def test_serve_cli_accepts_additional_config():
     """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
 
@@ -307,7 +300,8 @@ def test_serve_cli_accepts_additional_config():
         ]
     )
 
-    stage_cfg = _create_default_diffusion_stage_cfg(args)[0]
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))[0]
+
     engine_args = stage_cfg["engine_args"]
 
     assert args.additional_config == {"torchair_graph_config": {"enabled": True}}

tests/entrypoints/test_omni_entrypoints.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import argparse
 import queue
 from collections.abc import Callable
 from types import SimpleNamespace
@@ -176,83 +175,6 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine)
     )
 
 
-def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-    from vllm_omni.entrypoints.omni import Omni
-
-    captured: dict[str, Any] = {}
-
-    def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine:
-        captured.update(kwargs)
-        return FakeAsyncOmniEngine()
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine)
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--batch-timeout", type=int, default=10)
-    nullify_stage_engine_defaults(parser)
-    args = parser.parse_args([])
-    args.model = "fake-model"
-
-    Omni(**vars(args))
-
-    assert captured["gpu_memory_utilization"] is None
-    assert captured["batch_timeout"] == 10
-    assert "_cli_explicit_keys" not in captured
-
-
-def test_from_cli_args_warns_and_forwards_without_internal_keys(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    captured: dict[str, Any] = {}
-
-    def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine:
-        captured.update(kwargs)
-        return FakeAsyncOmniEngine()
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine)
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
-
-    args = argparse.Namespace(model="fake-model", gpu_memory_utilization=0.9, _cli_explicit_keys={"model"})
-    with pytest.deprecated_call(match="from_cli_args"):
-        Omni.from_cli_args(args)
-
-    assert captured["gpu_memory_utilization"] == 0.9
-    assert "_cli_explicit_keys" not in captured
-
-
-def test_deprecated_from_cli_args_preserves_legacy_parser_nulling(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    from vllm_omni.entrypoints.omni import Omni
-
-    captured: dict[str, Any] = {}
-
-    def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine:
-        captured.update(kwargs)
-        return FakeAsyncOmniEngine()
-
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine)
-    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
-    monkeypatch.setattr("sys.argv", ["prog"])
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--batch-timeout", type=int, default=10)
-    args = parser.parse_args([])
-    args.model = "fake-model"
-
-    with pytest.deprecated_call(match="from_cli_args"):
-        Omni.from_cli_args(args, parser=parser)
-
-    assert captured["gpu_memory_utilization"] is None
-    assert captured["batch_timeout"] == 10
-
-
 def _make_base():
     from vllm_omni.entrypoints.omni_base import OmniBase
 

tests/entrypoints/test_serve.py

@@ -4,65 +4,56 @@ from __future__ import annotations
 
 import argparse
 from types import SimpleNamespace
-from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def test_serve_parser_accepts_no_async_chunk() -> None:
-    """``--no-async-chunk`` should parse after deploy-overriding parser
-    defaults are nullified."""
-    try:
-        from vllm.utils.argparse_utils import FlexibleArgumentParser
-    except Exception as exc:
-        pytest.skip(f"Cannot build parser in this environment: {exc}")
-
-    root = FlexibleArgumentParser()
-    subparsers = root.add_subparsers(dest="subcommand")
+def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
+    """``--no-async-chunk`` should parse to ``async_chunk=False`` and mark the
+    shared deploy-level dest as explicitly provided by the user."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
     cmd = OmniServeCommand()
     cmd.subparser_init(subparsers)
 
     argv = ["serve", "fake-model", "--omni", "--no-async-chunk"]
-    args = root.parse_args(argv)
-
+    args = parser.parse_args(argv)
     assert args.async_chunk is False
 
-
-# ---------------------------------------------------------------------------
-# run_headless validation
-# ---------------------------------------------------------------------------
+    explicit = args.get_explicit_kwargs_dict()
+    assert args.get_explicit_kwargs_dict()
+    assert not explicit["async_chunk"]
 
 
-def _make_headless_args(**overrides: Any) -> argparse.Namespace:
-    """Build an argparse.Namespace shaped like the headless CLI passes in.
-
-    Defaults pass every validation gate so individual tests can mutate just
-    the field they're exercising.
-    """
-    defaults = dict(
-        model="fake-model",
-        stage_id=0,
-        replica_id=0,
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_replica_address=None,
-        omni_dp_size_local=1,
-        api_server_count=0,
-        worker_backend="multi_process",
-        stage_configs_path=None,
-        deploy_config=None,
-        log_stats=False,
-        disable_log_stats=False,
-        stage_init_timeout=600,
-        tokenizer=None,
+def _make_headless_args(**kwargs) -> TrackingNamespace:
+    defaults = {
+        "model": "fake-model",
+        "stage_id": 0,
+        "replica_id": 0,
+        "omni_master_address": "127.0.0.1",
+        "omni_master_port": 26000,
+        "omni_replica_address": None,
+        "omni_dp_size_local": 1,
+        "worker_backend": "multi_process",
+        "stage_configs_path": None,
+        "deploy_config": None,
+        "log_stats": False,
+        "disable_log_stats": False,
+        "stage_init_timeout": 600,
+        "tokenizer": None,
+    }
+    ns_kwargs = {**defaults, **kwargs}
+    ns = argparse.Namespace(**ns_kwargs)
+    return TrackingNamespace(
+        unfiltered_ns=ns,
+        explicit_keys=frozenset(ns.__dict__.keys()),
     )
-    defaults.update(overrides)
-    return argparse.Namespace(**defaults)
 
 
 def test_run_headless_requires_stage_id() -> None:
@@ -83,12 +74,6 @@ def test_run_headless_requires_master_port() -> None:
         run_headless(args)
 
 
-def test_run_headless_rejects_multi_api_server_count() -> None:
-    args = _make_headless_args(api_server_count=2)
-    with pytest.raises(ValueError, match="api_server_count can't be set"):
-        run_headless(args)
-
-
 def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
     args = _make_headless_args(worker_backend="ray")
     with pytest.raises(ValueError, match="worker_backend=multi_process"):

tests/test_arg_utils.py

@@ -5,303 +5,41 @@ hold for the orchestrator/engine/server CLI flag partition."""
 
 from __future__ import annotations
 
-import logging
-from dataclasses import dataclass, fields
+from dataclasses import fields
 
 import pytest
 
 from vllm_omni.engine.arg_utils import (
     SHARED_FIELDS,
-    derive_server_dests_from_vllm_parser,
+    OmniEngineArgs,
     internal_blacklist_keys,
-    orchestrator_args_from_argparse,
     orchestrator_field_names,
-    split_kwargs,
 )
 
-# ---------------------------------------------------------------------------
-# Fake engine class for unit testing — avoids pulling in the full vllm
-# EngineArgs and its heavy __post_init__ at test time.
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class _FakeEngineArgs:
-    """Stand-in for OmniEngineArgs with a representative subset of fields."""
-
-    model: str = ""
-    stage_id: int = 0
-    max_num_seqs: int = 64
-    gpu_memory_utilization: float = 0.9
-    async_chunk: bool = False  # also in OrchestratorArgs → shared
-    log_stats: bool = False  # also in OrchestratorArgs → shared
-    stage_configs_path: str | None = None
-
-
-# ============================================================================
-# Invariant 1 — OrchestratorArgs and engine must not ambiguously overlap.
-# ============================================================================
-
-
-def test_no_ambiguous_overlap_with_fake_engine():
-    """OrchestratorArgs ∩ engine fields must be ⊆ SHARED_FIELDS."""
-    orch = orchestrator_field_names()
-    engine = {f.name for f in fields(_FakeEngineArgs)}
-    overlap = orch & engine
-    unexpected = overlap - SHARED_FIELDS
-    assert not unexpected, (
-        f"Fields declared in both OrchestratorArgs and the engine class "
-        f"but not in SHARED_FIELDS: {sorted(unexpected)}. These cause "
-        f"double-routing — either remove the duplicate declaration or add "
-        f"to SHARED_FIELDS if sharing is intentional."
-    )
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_no_ambiguous_overlap_with_real_engine():
-    """Same check, but against the real OmniEngineArgs."""
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
+    """OrchestratorArgs ∩ OmniEngineArgs must be ⊆ SHARED_FIELDS ∪ orchestrator-captured.
+
+    Fields that the orchestrator intentionally captures to prevent uniform
+    propagation (async_chunk, tokenizer) are excluded — they exist in both
+    classes by design but flow through DeployConfig, not kwargs forwarding.
+    """
+    # Fields on both classes by design: orchestrator captures them to prevent
+    # uniform per-stage propagation; redistributed via DeployConfig.
+    orchestrator_captured = {"async_chunk", "tokenizer"}
 
     orch = orchestrator_field_names()
     engine = {f.name for f in fields(OmniEngineArgs)}
     overlap = orch & engine
-    unexpected = overlap - SHARED_FIELDS
+    unexpected = overlap - SHARED_FIELDS - orchestrator_captured
     assert not unexpected, (
-        f"Real OmniEngineArgs has ambiguous overlap with OrchestratorArgs: "
+        f"OmniEngineArgs has ambiguous overlap with OrchestratorArgs: "
         f"{sorted(unexpected)}. Update SHARED_FIELDS or remove duplication."
     )
 
 
-# ============================================================================
-# Invariant 2 — split_kwargs partitions correctly.
-# ============================================================================
-
-
-def test_split_orchestrator_only():
-    """Pure orchestrator fields go to OrchestratorArgs, not engine_kwargs."""
-    raw = {"stage_init_timeout": 500, "worker_backend": "ray"}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert orch.stage_init_timeout == 500
-    assert orch.worker_backend == "ray"
-    assert "stage_init_timeout" not in engine
-    assert "worker_backend" not in engine
-
-
-def test_split_engine_only():
-    """Pure engine fields go to engine_kwargs, not OrchestratorArgs."""
-    raw = {"max_num_seqs": 128, "gpu_memory_utilization": 0.85}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert engine["max_num_seqs"] == 128
-    assert engine["gpu_memory_utilization"] == 0.85
-    # These fields don't exist on OrchestratorArgs at all.
-
-
-def test_split_shared_fields_go_to_both():
-    """Fields in SHARED_FIELDS are copied to both buckets."""
-    raw = {"model": "Qwen/Qwen2.5-Omni-7B", "log_stats": True}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert orch.log_stats is True
-    assert engine["model"] == "Qwen/Qwen2.5-Omni-7B"
-    assert engine["log_stats"] is True
-
-
-def test_split_drops_unclassified():
-    """Unclassified fields (uvicorn/server) are dropped silently."""
-    raw = {
-        "max_num_seqs": 64,  # engine
-        "host": "0.0.0.0",  # unclassified (server)
-        "port": 8091,  # unclassified (server)
-        "ssl_keyfile": "key.pem",  # unclassified (server)
-    }
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert engine == {"max_num_seqs": 64}
-    assert "host" not in engine
-    assert "port" not in engine
-    assert "ssl_keyfile" not in engine
-
-
-def test_split_mixed_real_world():
-    """End-to-end: raw CLI kwargs with all three classes present."""
-    raw = {
-        # orchestrator
-        "stage_init_timeout": 400,
-        "deploy_config": "/tmp/deploy.yaml",
-        "worker_backend": "multi_process",
-        "async_chunk": True,
-        # engine
-        "max_num_seqs": 32,
-        "gpu_memory_utilization": 0.8,
-        # shared
-        "model": "Qwen/Qwen3-Omni",
-        "log_stats": False,
-        # server / unclassified
-        "host": "0.0.0.0",
-        "port": 8091,
-        "api_key": "secret",
-        # None values
-        "ray_address": None,
-    }
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-
-    # Orchestrator side
-    assert orch.stage_init_timeout == 400
-    assert orch.deploy_config == "/tmp/deploy.yaml"
-    assert orch.worker_backend == "multi_process"
-    assert orch.async_chunk is True
-    assert orch.log_stats is False  # shared, read from raw
-    assert orch.ray_address is None  # default preserved
-
-    # Engine side
-    assert engine["max_num_seqs"] == 32
-    assert engine["gpu_memory_utilization"] == 0.8
-    assert engine["model"] == "Qwen/Qwen3-Omni"
-    assert engine["log_stats"] is False
-    assert "host" not in engine
-    assert "port" not in engine
-    assert "api_key" not in engine
-    # orchestrator-only keys never reach engine
-    assert "stage_init_timeout" not in engine
-    assert "deploy_config" not in engine
-    # async_chunk is in SHARED_FIELDS so it flows to both
-    assert "async_chunk" in engine
-
-
-# ============================================================================
-# Invariant 3 — user-typed unclassifiable flags warn (don't fail silently).
-# ============================================================================
-
-
-def test_user_typed_unclassified_warns(monkeypatch):
-    """If the user types a flag we can't route, warn — don't silently drop."""
-    import io
-
-    buf = io.StringIO()
-    vllm_logger = logging.getLogger("vllm")
-    orig_handlers = list(vllm_logger.handlers)
-    vllm_logger.handlers.clear()
-    handler = logging.StreamHandler(buf)
-    handler.setLevel(logging.WARNING)
-    vllm_logger.addHandler(handler)
-    try:
-        raw = {"bogus_flag": "value", "max_num_seqs": 64}
-        split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed={"bogus_flag"})
-        assert "bogus_flag" in buf.getvalue(), f"Expected 'bogus_flag' in log output, got: {buf.getvalue()!r}"
-    finally:
-        vllm_logger.handlers.clear()
-        for h in orig_handlers:
-            vllm_logger.addHandler(h)
-
-
-def test_unclassified_without_user_typed_silent():
-    """Without user_typed, unclassified keys drop silently (argparse defaults
-    for server flags shouldn't spam logs on every launch)."""
-    import io
-
-    vllm_logger = logging.getLogger("vllm")
-    orig_handlers = list(vllm_logger.handlers)
-    vllm_logger.handlers.clear()
-    buf = io.StringIO()
-    handler = logging.StreamHandler(buf)
-    handler.setLevel(logging.WARNING)
-    vllm_logger.addHandler(handler)
-    try:
-        raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64}
-        split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed=None)
-        assert "host" not in buf.getvalue()
-    finally:
-        vllm_logger.handlers.clear()
-        for h in orig_handlers:
-            vllm_logger.addHandler(h)
-
-
-# ============================================================================
-# Invariant 4 — CLI flag classification completeness.
-# Catches new flags added without updating OrchestratorArgs or OmniEngineArgs.
-# ============================================================================
-
-
-def test_all_omni_cli_flags_classified():
-    """Every vllm-omni-added CLI flag must be classifiable.
-
-    Runs ``OmniServeCommand.subparser_init`` and checks that every new
-    argument (compared to vllm's base parser) is either:
-      - a field on OrchestratorArgs, OR
-      - a field on OmniEngineArgs, OR
-      - in SHARED_FIELDS
-    """
-    try:
-        from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-        from vllm_omni.entrypoints.cli.serve import OmniServeCommand
-    except Exception as exc:
-        pytest.skip(f"Cannot build parser in this environment: {exc}")
-
-    # Build the serve parser
-    root = FlexibleArgumentParser()
-    subparsers = root.add_subparsers()
-    cmd = OmniServeCommand()
-    try:
-        parser = cmd.subparser_init(subparsers)
-    except Exception as exc:
-        pytest.skip(f"subparser_init failed (dev env issue): {exc}")
-
-    all_dests = {a.dest for a in parser._actions if a.dest and a.dest not in {"help", "model_tag"}}
-
-    orch = orchestrator_field_names()
-    engine = {f.name for f in fields(OmniEngineArgs)}
-    server_derived = derive_server_dests_from_vllm_parser()
-
-    unclassified = all_dests - orch - engine - SHARED_FIELDS - server_derived
-    # Some argparse-internal dests (suppressed, private) may not match —
-    # filter those out.
-    unclassified = {d for d in unclassified if not d.startswith("_")}
-
-    assert not unclassified, (
-        f"These CLI flags are not classified as "
-        f"orchestrator/engine/shared/server: {sorted(unclassified)}. "
-        f"Add them to OrchestratorArgs (if consumed by orchestrator), "
-        f"OmniEngineArgs (if consumed by per-stage engine), or the known-server "
-        f"allowlist (if they're vllm frontend flags). "
-        f"If intentional (e.g. a new CLI-only flag that doesn't map to either "
-        f"dataclass), add it to a KNOWN_UNROUTED allowlist."
-    )
-
-
-# ============================================================================
-# argparse interop (Phase 3).
-# ============================================================================
-
-
-def test_orchestrator_args_from_argparse():
-    """Can build OrchestratorArgs from an argparse.Namespace."""
-    import argparse
-
-    ns = argparse.Namespace(
-        stage_init_timeout=500,
-        deploy_config="/tmp/x.yaml",
-        max_num_seqs=64,  # engine field — ignored
-        host="0.0.0.0",  # server field — ignored
-    )
-    orch = orchestrator_args_from_argparse(ns)
-    assert orch.stage_init_timeout == 500
-    assert orch.deploy_config == "/tmp/x.yaml"
-    assert orch.worker_backend == "multi_process"  # default
-
-
-def test_derive_server_dests_returns_frozenset():
-    """Server-dest derivation returns a frozenset (possibly empty)."""
-    result = derive_server_dests_from_vllm_parser()
-    assert isinstance(result, frozenset)
-
-
-# ============================================================================
-# internal_blacklist_keys — single source of truth for per-stage forwarding.
-# ============================================================================
-
-
 def test_internal_blacklist_keys_derived_from_orchestrator():
     """Blacklist is exactly OrchestratorArgs fields minus SHARED_FIELDS.
 
@@ -311,214 +49,3 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
     """
     blacklist = internal_blacklist_keys()
     assert blacklist == orchestrator_field_names() - SHARED_FIELDS
-    # Spot-check expected entries
-    assert "stage_init_timeout" in blacklist
-    assert "deploy_config" in blacklist
-    assert "async_chunk" not in blacklist  # now in SHARED_FIELDS
-    # Shared fields must NOT appear — they flow to both orchestrator and engine
-    assert "model" not in blacklist
-    assert "log_stats" not in blacklist
-
-
-# ============================================================================
-# Boundary value analysis — edge cases around split_kwargs.
-# ============================================================================
-
-
-def test_split_empty_kwargs():
-    """Empty kwargs yields default OrchestratorArgs and empty engine dict."""
-    orch, engine = split_kwargs({}, engine_cls=_FakeEngineArgs)
-    assert orch.stage_init_timeout == 300  # dataclass default
-    assert orch.worker_backend == "multi_process"  # dataclass default
-    assert engine == {}
-
-
-def test_split_all_none_values_preserved_on_orchestrator():
-    """None values for orchestrator fields are kept (represents 'not set')."""
-    raw = {"ray_address": None, "deploy_config": None, "max_num_seqs": None}
-    orch, engine = split_kwargs(raw, engine_cls=_FakeEngineArgs)
-    assert orch.ray_address is None
-    assert orch.deploy_config is None
-    # Engine-side None still passes through; caller decides semantics downstream.
-    assert engine.get("max_num_seqs") is None
-
-
-def test_split_user_typed_with_empty_kwargs_no_warn(caplog):
-    """user_typed non-empty but kwargs empty — no warnings emitted."""
-    with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"):
-        split_kwargs({}, engine_cls=_FakeEngineArgs, user_typed={"nothing"})
-    assert not caplog.records
-
-
-def test_ambiguous_field_strict_raises():
-    """strict=True raises ValueError on overlap outside SHARED_FIELDS."""
-
-    # deploy_config is on OrchestratorArgs; declaring it on the engine class
-    # too (without adding to SHARED_FIELDS) creates an ambiguous route.
-    @dataclass
-    class _AmbiguousEngine:
-        deploy_config: str | None = None
-
-    with pytest.raises(ValueError, match="both OrchestratorArgs and"):
-        split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=True)
-
-
-def test_ambiguous_field_non_strict_routes_to_orchestrator():
-    """strict=False logs ERROR but routes the ambiguous field to orchestrator."""
-
-    @dataclass
-    class _AmbiguousEngine:
-        deploy_config: str | None = None
-
-    import io
-
-    vllm_logger = logging.getLogger("vllm")
-    orig_handlers = list(vllm_logger.handlers)
-    vllm_logger.handlers.clear()
-    buf = io.StringIO()
-    handler = logging.StreamHandler(buf)
-    handler.setLevel(logging.ERROR)
-    vllm_logger.addHandler(handler)
-    try:
-        orch, engine = split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=False)
-        assert orch.deploy_config == "x"
-        assert "deploy_config" not in engine
-        assert "both OrchestratorArgs" in buf.getvalue()
-    finally:
-        vllm_logger.handlers.clear()
-        for h in orig_handlers:
-            vllm_logger.addHandler(h)
-
-
-# Sentinel-default precedence invariants (#3035)
-
-
-def _build_full_serve_parser():
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-    try:
-        from vllm.entrypoints.openai.cli_args import make_arg_parser
-    except ImportError:
-        pytest.skip("vllm parser not importable")
-    return make_arg_parser(FlexibleArgumentParser())
-
-
-def test_nullify_stage_engine_defaults_resets_inherited_defaults():
-    import argparse
-
-    from vllm_omni.config.stage_config import deploy_override_field_names
-    from vllm_omni.engine.arg_utils import (
-        nullify_stage_engine_defaults,
-    )
-
-    parser = _build_full_serve_parser()
-    nullify_stage_engine_defaults(parser)
-
-    override_dests = deploy_override_field_names()
-    offenders = [
-        (a.dest, a.default)
-        for a in parser._actions
-        if a.dest not in ("help", "version")
-        and a.option_strings
-        and a.dest in override_dests
-        and a.default is not None
-        and a.default is not argparse.SUPPRESS
-    ]
-    assert not offenders, f"Stage flags with non-None defaults after nullify: {offenders}"
-
-
-def test_non_override_flags_keep_real_defaults_after_nullify():
-    import argparse
-
-    from vllm_omni.config.stage_config import deploy_override_field_names
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--batch-timeout", type=int, default=10, help="Batch timeout.")
-    parser.add_argument("--max-num-seqs", type=int, default=64, help="Max num seqs.")
-    nullify_stage_engine_defaults(parser)
-
-    assert "batch_timeout" not in deploy_override_field_names()
-
-    batch_timeout = next(a for a in parser._actions if a.dest == "batch_timeout")
-    max_num_seqs = next(a for a in parser._actions if a.dest == "max_num_seqs")
-    assert batch_timeout.default == 10
-    assert max_num_seqs.default is None
-
-
-def test_help_text_preserves_default_after_nullify():
-    # Real defaults must stay visible in --help even though parser stores None.
-    import argparse
-
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--max-num-seqs", type=int, default=42, help="Example knob.")
-    nullify_stage_engine_defaults(parser)
-
-    action = next(a for a in parser._actions if a.dest == "max_num_seqs")
-    assert action.default is None
-    assert "(default: 42)" in action.help
-
-
-_OMNIENGINEARGS_USER_INPUT_FIELDS = frozenset(
-    {
-        "model_arch",
-        "engine_output_type",
-        "hf_config_name",
-        "custom_process_next_stage_input_func",
-        "subtalker_sampling_params",
-        "omni_kv_config",
-        "quantization_config",
-        "worker_type",
-        "task_type",
-        "worker_cls",
-        "omni_master_address",
-        "omni_master_port",
-        "stage_configs_path",
-        "output_modalities",
-        "custom_pipeline_args",
-    }
-)
-
-
-def test_omniengineargs_user_input_fields_default_to_none():
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
-
-    offenders = [
-        (f.name, f.default)
-        for f in fields(OmniEngineArgs)
-        if f.name in _OMNIENGINEARGS_USER_INPUT_FIELDS
-        and f.default is not dataclasses.MISSING
-        and f.default is not None
-    ]
-    assert not offenders, f"User-input fields with non-None defaults: {offenders}"
-
-
-def test_omniengineargs_create_tracks_explicit_fields():
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
-
-    ea = OmniEngineArgs.create(model="x", gpu_memory_utilization=0.5)
-    assert ea._explicit_fields == frozenset({"model", "gpu_memory_utilization"})
-    assert ea.explicit_kwargs() == {"model": "x", "gpu_memory_utilization": 0.5}
-
-
-def test_omniengineargs_bare_constructor_has_no_explicit_tracking():
-    try:
-        from vllm_omni.engine.arg_utils import OmniEngineArgs
-    except Exception as exc:
-        pytest.skip(f"OmniEngineArgs not importable: {exc}")
-
-    ea = OmniEngineArgs(model="x")
-    assert not hasattr(ea, "_explicit_fields")
-    assert "model" in ea.explicit_kwargs()
-
-
-# dataclasses already imported via ``from dataclasses import dataclass, fields``
-import dataclasses  # noqa: E402  -- needed for MISSING sentinel above

tests/test_config_factory.py

@@ -29,7 +29,6 @@ from vllm_omni.config.stage_config import (
     _deep_merge_stage,
     _resolve_scheduler,
     build_stage_runtime_overrides,
-    deploy_override_field_names,
     load_deploy_config,
     merge_pipeline_deploy,
     register_pipeline,
@@ -39,6 +38,8 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_black
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class TestStageType:
     """Tests for StageType enum."""
@@ -156,23 +157,6 @@ class TestStageConfig:
         omega_config = config.to_omegaconf()
         assert omega_config.engine_args.max_num_seqs == 32
 
-    def test_to_omegaconf_omits_none_deploy_overrides_for_engine_args(self):
-        """None deploy overrides must fall through to EngineArgs defaults."""
-
-        config = StageConfig(
-            stage_id=0,
-            model_stage="thinker",
-            runtime_overrides={name: None for name in deploy_override_field_names()},
-        )
-
-        omega_config = config.to_omegaconf()
-        engine_args = dict(omega_config.engine_args)
-
-        assert "devices" not in engine_args
-        assert "max_batch_size" not in engine_args
-        for name in deploy_override_field_names() - {"devices"}:
-            assert name not in engine_args
-
     def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self):
         config = StageConfig(
             stage_id=1,
@@ -1018,64 +1002,7 @@ class TestPipelineRegistry:
 
 
 class TestDeployConfigLoading:
-    def test_deploy_override_fields_include_deploy_schema_fields(self):
-        expected_fields = {
-            "active_stream_window",
-            "async_chunk",
-            # StageDeployConfig: stage placement and runtime fields.
-            "devices",
-            # StageDeployConfig: vLLM EngineArgs fields.
-            "async_scheduling",
-            "compilation_config",
-            "config_format",
-            "custom_voice_dir",
-            "data_parallel_size",
-            "devices",
-            "disable_hybrid_kv_cache_manager",
-            "enable_flashinfer_autotune",
-            "enforce_eager",
-            "gpu_memory_utilization",
-            "load_format",
-            "max_model_len",
-            "max_num_batched_tokens",
-            "max_num_seqs",
-            "mm_processor_cache_gb",
-            "profiler_config",
-            "skip_mm_profiling",
-            "subtalker_sampling_params",
-            "tensor_parallel_size",
-            "tokenizer_mode",
-            # StageDeployConfig: diffusion parallel_config deploy override fields.
-            "cfg_parallel_size",
-            "enable_expert_parallel",
-            "hsdp_replicate_size",
-            "hsdp_shard_size",
-            "ring_degree",
-            "sequence_parallel_size",
-            "ulysses_degree",
-            "ulysses_mode",
-            "use_hsdp",
-            "vae_patch_parallel_size",
-            # DeployConfig: pipeline-wide engine settings.
-            "data_parallel_size",
-            "distributed_executor_backend",
-            "dtype",
-            "enable_chunked_prefill",
-            "enable_prefix_caching",
-            "pipeline_parallel_size",
-            "quantization",
-            "trust_remote_code",
-        }
-
-        actual_fields = deploy_override_field_names()
-        assert expected_fields == actual_fields, (
-            f"added={actual_fields - expected_fields}, removed={expected_fields - actual_fields}"
-        )
-
     def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path):
-        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml"
         custom_voice_dir = tmp_path / "voices"
         deploy_path.write_text(
@@ -1913,16 +1840,6 @@ class TestSentinelDefaultPrecedence:
         named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD]
         assert "cli_explicit_keys" not in {p.name for p in named}
 
-    def test_cli_explicit_keys_kwarg_emits_deprecation(self):
-        with warnings.catch_warnings(record=True) as w:
-            warnings.simplefilter("always")
-            StageConfigFactory._create_from_registry(
-                "qwen3_omni_moe",
-                cli_overrides={},
-                cli_explicit_keys={"max_num_seqs"},
-            )
-            assert any(issubclass(x.category, DeprecationWarning) for x in w)
-
     def test_async_chunk_dispatches_processors(self):
         """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end
         processors based on ``deploy.async_chunk``, without needing a

tests/utils/test_tracking_parser.py

@@ -0,0 +1,545 @@
+"""Tests for TrackingArgumentParser and related utilities."""
+
+import argparse
+import json
+from unittest.mock import patch
+
+import pytest
+import yaml
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.config.stage_config import (
+    _PIPELINE_REGISTRY,
+    DeployConfig,
+    PipelineConfig,
+    StageConfigFactory,
+    StageDeployConfig,
+    StagePipelineConfig,
+)
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
+
+### Fake pipeline/deploy config for integration tests
+
+_TEST_MODEL = "_test_tracking"
+
+_TEST_PIPELINE = PipelineConfig(
+    model_type=_TEST_MODEL,
+    stages=(
+        StagePipelineConfig(stage_id=0, model_stage="stage0"),
+        StagePipelineConfig(stage_id=1, model_stage="stage1"),
+        StagePipelineConfig(stage_id=2, model_stage="stage2"),
+    ),
+)
+
+_TEST_DEPLOY = DeployConfig(
+    async_chunk=False,
+    stages=[
+        StageDeployConfig(stage_id=0, gpu_memory_utilization=0.5, max_num_seqs=32),
+        StageDeployConfig(stage_id=1, gpu_memory_utilization=0.6, max_num_seqs=64),
+        StageDeployConfig(stage_id=2, gpu_memory_utilization=0.1, max_num_seqs=16),
+    ],
+)
+
+
+@pytest.fixture()
+def mock_stages(monkeypatch):
+    """Register a fake pipeline and mock deploy YAML loading."""
+    monkeypatch.setitem(_PIPELINE_REGISTRY._loaded, _TEST_MODEL, _TEST_PIPELINE)
+    monkeypatch.setattr(
+        "vllm_omni.config.stage_config.load_deploy_config",
+        lambda _path: _TEST_DEPLOY,
+    )
+    return __file__
+
+
+### Tests for TrackingNamespace
+def test_tracking_namespaces_cant_be_nested():
+    """Ensure tracking namespaces explode if we try to nest them."""
+    track_ns = TrackingNamespace(
+        unfiltered_ns=argparse.Namespace(foo="bar"),
+        explicit_keys=frozenset(),
+    )
+
+    with pytest.raises(ValueError):
+        TrackingNamespace(
+            unfiltered_ns=track_ns,
+            explicit_keys=frozenset(),
+        )
+
+
+def test_tracking_namespaces_init():
+    """Check simple initialization for tracking namespaces."""
+    unfiltered_ns = argparse.Namespace(foo="bar")
+    tracked_ns = TrackingNamespace(
+        unfiltered_ns=unfiltered_ns,
+        explicit_keys=frozenset({"foo"}),
+    )
+    assert tracked_ns.foo == "bar"
+    assert tracked_ns.explicit_keys == frozenset({"foo"})
+
+
+def test_tracking_filtering():
+    """Ensure tracking namespaces are filterable."""
+    unfiltered_ns = argparse.Namespace(foo="bar", baz="foobar")
+    tracked_ns = TrackingNamespace(
+        unfiltered_ns=unfiltered_ns,
+        explicit_keys=frozenset({"foo"}),
+    )
+    assert tracked_ns.foo == "bar"
+    assert tracked_ns.baz == "foobar"
+    assert tracked_ns.explicit_keys == frozenset({"foo"})
+    # baz gets dropped because it's not marked in explicit_keys
+    assert tracked_ns.get_explicit_kwargs_dict() == {"foo": "bar"}
+
+
+def test_setattr_writes_through_to_unfiltered_ns():
+    """Ensure mutating an attribute on TrackingNamespace forwards to
+    get_explicit_kwargs_dict() and vars()."""
+    unfiltered_ns = argparse.Namespace(model="original")
+    tracked_ns = TrackingNamespace(
+        unfiltered_ns=unfiltered_ns,
+        explicit_keys=frozenset({"model"}),
+    )
+    assert tracked_ns.model == "original"
+    assert tracked_ns.get_explicit_kwargs_dict() == {"model": "original"}
+    assert vars(tracked_ns) == {"model": "original"}
+
+    # Ensure if we update the namespace, it's forwarded correctly
+    tracked_ns.model = "updated"
+    assert tracked_ns.model == "updated"
+    assert tracked_ns.get_explicit_kwargs_dict() == {"model": "updated"}
+    assert vars(tracked_ns) == {"model": "updated"}
+
+
+### Tests for simple cases (no nested parsers or groups)
+def test_vars_does_not_expose_internals():
+    """Ensure vars on a TrackingArgumentParser is identical to running on the real one."""
+    tracking = TrackingArgumentParser()
+    flexible = FlexibleArgumentParser()
+    namespaces = []
+
+    # For both parser, register two args, but only pass one
+    for parser in [tracking, flexible]:
+        parser.add_argument("--foo", type=int, default=42)
+        parser.add_argument("--bar", type=str, default="x")
+        namespaces.append(parser.parse_args(["--foo", "100"]))
+
+    # Ensure that vars are the same result, since it isn't filtered
+    tracked_ns, flexible_ns = namespaces
+    tracked_kwargs = vars(tracked_ns)
+    real_kwargs = vars(flexible_ns)
+    assert tracked_kwargs == real_kwargs == {"foo": 100, "bar": "x"}
+    assert tracked_ns.explicit_keys == {"foo"}
+
+
+def test_default_not_detected():
+    """Ensure omitted defaults aren't in explicit keys and take defaults."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == 42
+
+
+@pytest.mark.parametrize("val", ["42", "100"])
+def test_explicit_value_equal_to_default(val):
+    """Ensure explicit keys correctly handles passed values."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns = p.parse_args(["--foo", val])
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == int(val)
+
+
+def test_equals_syntax():
+    """Ensure equals syntax is handled correctly wrt explicit keys."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns = p.parse_args(["--foo=100"])
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == 100
+
+
+def test_explicit_none_via_store_const():
+    """Ensure passing `None` to the namespace, it isn't filtered from explicit_keys."""
+    parser = TrackingArgumentParser()
+    parser.add_argument("--foo", action="store_const", const=None, default="Something else")
+    ns = parser.parse_args(["--foo"])
+    # User explicitly passed --foo, so it should be in the explicit keys even though it's None
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo is None
+
+
+def test_multiple_args_mixed():
+    """Ensure that explicit keys are correct when some are passed and others aren't."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=1)
+    p.add_argument("--bar", type=str, default="x")
+    p.add_argument("--baz", type=float, default=0.5)
+    ns = p.parse_args(["--foo", "10", "--baz", "0.9"])
+    assert ns.explicit_keys == {"foo", "baz"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == "x"
+
+
+def test_store_true_default():
+    """Ensure that store true is handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--verbose", action="store_true")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose is False
+
+
+def test_store_true_explicit():
+    """Ensure that store true is handled correctly when passed."""
+    p = TrackingArgumentParser()
+    p.add_argument("--verbose", action="store_true")
+    ns = p.parse_args(["--verbose"])
+    assert ns.explicit_keys == {"verbose"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose is True
+
+
+def test_store_false_default():
+    """Ensure that store false is handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--disable-x", action="store_false", dest="enable_x", default=True)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.enable_x is True
+
+
+def test_store_false_explicit():
+    """Ensure that store false is handled correctly when passed."""
+    p = TrackingArgumentParser()
+    p.add_argument("--disable-x", action="store_false")
+    ns = p.parse_args(["--disable-x"])
+    assert ns.explicit_keys == {"disable_x"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.disable_x is False
+
+
+def test_dest_is_reflected_in_explicit_keys():
+    """Ensure that explicit keys use dest correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, dest="bar")
+    ns = p.parse_args(["--foo", "100"])
+    assert ns.explicit_keys == {"bar"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == 100
+
+
+def test_boolean_optional_action_default():
+    """Check that boolean optional actions are handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--flag", action=argparse.BooleanOptionalAction)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flag is None
+
+
+def test_boolean_optional_action_positive():
+    """Check that boolean optional actions are handled correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--flag", action=argparse.BooleanOptionalAction, default=None)
+    ns = p.parse_args(["--flag"])
+    assert ns.explicit_keys == {"flag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flag is True
+
+
+def test_no_option_strings_are_handled():
+    """Ensure --no-<feature> sets <feature> in the explicit keys correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--flag", action=argparse.BooleanOptionalAction, default=None)
+    ns = p.parse_args(["--no-flag"])
+    assert ns.explicit_keys == {"flag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flag is False
+
+
+def test_json_type():
+    """Ensure json type is handled correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--cfg", type=json.loads, default="{}")
+    ns = p.parse_args(["--cfg", '{"a": 1}'])
+    assert "cfg" in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.cfg == {"a": 1}
+
+
+def test_choices():
+    """Ensure choices are handled correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--mode", choices=["fast", "slow"], default="fast")
+    ns = p.parse_args(["--mode", "slow"])
+    assert "mode" in ns.explicit_keys
+    assert ns is not None
+    assert ns.mode == "slow"
+
+
+def test_explicit_positional_arg():
+    """Ensure positional args are handled correctly when provided."""
+    p = TrackingArgumentParser()
+    p.add_argument("name", nargs="?", default=None)
+    ns = p.parse_args(["hello"])
+    assert "name" in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.name == "hello"
+
+
+def test_omitted_positional_arg():
+    """Ensure positional args are handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("name", nargs="?", default=None)
+    ns = p.parse_args([])
+    assert "name" not in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.name is None
+
+
+def test_explicit_nargs():
+    """Ensure that variable num args are handled correctly when provided."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", nargs="*", default=None)
+    ns = p.parse_args(["--items", "a", "b"])
+    assert "items" in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items == ["a", "b"]
+
+
+def test_omitted_nargs():
+    """Ensure that variable num args are handled correctly when omitted."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", nargs="*", default=None)
+    ns = p.parse_args([])
+    assert "items" not in ns.explicit_keys
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items is None
+
+
+def test_parse_known_args_tracking():
+    """Ensure parse_known_args is also trackable"""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int, default=42)
+    ns, remaining = p.parse_known_args(["--foo", "10", "--unknown", "val"])
+    assert ns.explicit_keys == {"foo"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.foo == 10
+    assert remaining == ["--unknown", "val"]
+
+
+### Tests for group handling
+def test_group_arg_default():
+    """Ensure that that groups with defaults are handled correctly."""
+    p = TrackingArgumentParser()
+    g = p.add_argument_group("TestGroup")
+    g.add_argument("--bar", type=str, default="baz")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == "baz"
+
+
+def test_group_arg_explicit():
+    """Ensure that that groups are handled correctly."""
+    p = TrackingArgumentParser()
+    g = p.add_argument_group("TestGroup")
+    g.add_argument("--bar", type=str, default="baz")
+    ns = p.parse_args(["--bar", "qux"])
+    assert ns.explicit_keys == {"bar"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.bar == "qux"
+
+
+def test_multiple_groups():
+    """Test multiple groups behave correctly."""
+    p = TrackingArgumentParser()
+    g1 = p.add_argument_group("Group1")
+    g2 = p.add_argument_group("Group2")
+    g1.add_argument("--a", type=int, default=1)
+    g2.add_argument("--b", type=int, default=2)
+    ns = p.parse_args(["--b", "20"])
+    assert ns.explicit_keys == {"b"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.a == 1
+    assert ns.b == 20
+
+
+def test_omitted_mutually_exclusive_group():
+    """Ensure that that mutually exclusive groups with defaults are handled correctly."""
+    p = TrackingArgumentParser()
+    grp = p.add_mutually_exclusive_group()
+    grp.add_argument("--json", action="store_true")
+    grp.add_argument("--text", action="store_true")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+
+
+def test_mutually_exclusive_group():
+    """Ensure that that mutually exclusive groups are handled correctly."""
+    p = TrackingArgumentParser()
+    grp = p.add_mutually_exclusive_group()
+    grp.add_argument("--json", action="store_true")
+    grp.add_argument("--text", action="store_true")
+    ns = p.parse_args(["--json"])
+    assert ns.explicit_keys == {"json"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.json is True
+    assert ns.text is False
+
+
+### Tests for subparser handling
+def test_subparser_explicit_detection():
+    """Ensure that subparsers are handled correctly."""
+    p = TrackingArgumentParser()
+    sub = p.add_subparsers()
+    child = sub.add_parser("foo")
+    child.add_argument("--bar", type=str)
+    ns = p.parse_args(["foo", "--bar", "baz"])
+    assert isinstance(child, TrackingArgumentParser)
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.explicit_keys == {"bar"}
+    assert ns.bar == "baz"
+
+
+def test_subparser_group_args():
+    """Ensure that subparsers with groups are handled correctly."""
+    p = TrackingArgumentParser()
+    sub = p.add_subparsers()
+    child = sub.add_parser("foo")
+    g = child.add_argument_group("Config")
+    g.add_argument("--port", type=int, default=8000)
+    g.add_argument("--host", type=str, default="localhost")
+    ns = p.parse_args(["foo", "--port", "9000"])
+    assert ns.explicit_keys == {"port"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.host == "localhost"
+
+
+### Tests for specific behaviors against FlexibleArgumentParser
+def test_config_file_args_detected(tmp_path):
+    """Ensure config is handled correctly for vLLM's FlexibleArgumentParser."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int)
+    p.add_argument("--bar", type=int)
+    cfg = tmp_path / "test.yaml"
+    cfg.write_text(yaml.dump({"foo": 100}))
+    ns = p.parse_args(["--config", str(cfg)])
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.explicit_keys == {"foo"}
+    assert ns.foo == 100
+
+
+def test_cli_overrides_config(tmp_path):
+    """Ensure tracking parser handles config vs cli overrides correctly."""
+    p = TrackingArgumentParser()
+    p.add_argument("--foo", type=int)
+    cfg = tmp_path / "test.yaml"
+    cfg.write_text(yaml.dump({"foo": 100}))
+    ns = p.parse_args(["--config", str(cfg), "--foo", "200"])
+    assert isinstance(ns, TrackingNamespace)
+    assert "foo" in ns.explicit_keys
+    assert ns.foo == 200
+
+
+### Integration tests for arg resolution through StageConfigFactory
+def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages):
+    """Explicitly passed CLI values reach runtime_overrides on all stages."""
+    p = TrackingArgumentParser()
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    ns = p.parse_args(["--max-num-seqs", "999"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    for stage in stages:
+        assert stage.runtime_overrides.get("max_num_seqs") == 999
+
+
+def test_omitted_default_not_in_runtime_overrides(mock_stages):
+    """Omitted defaults are overridden by deploy config values"""
+    p = TrackingArgumentParser()
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    ns = p.parse_args([])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    for stage in stages:
+        assert stage.runtime_overrides == {}
+
+
+def test_config_file_args_reach_runtime_overrides(mock_stages):
+    """Args from --config YAML must be treated as explicitly passed and
+    flow through to runtime_overrides."""
+    p = TrackingArgumentParser()
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    with patch(
+        "vllm.utils.argparse_utils.FlexibleArgumentParser.load_config_file", return_value=["--max-num-seqs", "999"]
+    ):
+        ns = p.parse_args(["--config", "fake.yaml"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    for stage in stages:
+        assert stage.runtime_overrides.get("max_num_seqs") == 999
+        assert "gpu_memory_utilization" not in stage.runtime_overrides
+
+
+def test_per_stage_override_routes_correctly(mock_stages):
+    """Ensure stage_<N>_<key> only affects the targeted stage."""
+    p = TrackingArgumentParser()
+    p.add_argument("--stage-0-gpu-memory-utilization", type=float)
+    ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+    assert stages[0].runtime_overrides == {"gpu_memory_utilization": 0.42}
+    assert stages[1].runtime_overrides == {}
+    assert stages[2].runtime_overrides == {}
+
+
+def test_explicit_args_omitted_from_yaml(mock_stages):
+    """Ensure only passed args end up in runtime overrides (regardless of whether
+    they are defined in the yaml config)."""
+    p = TrackingArgumentParser()
+    p.add_argument("--enforce-eager", action="store_true")
+    p.add_argument("--max-num-seqs", type=int, default=64)
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    # NOTE: enforce eager is not set in the mock config yaml.
+    ns = p.parse_args(["--enforce-eager"])
+
+    explicit_kwargs = ns.get_explicit_kwargs_dict()
+    stages = StageConfigFactory._create_from_registry(
+        _TEST_MODEL,
+        explicit_kwargs,
+        deploy_config_path=mock_stages,
+    )
+
+    for stage in stages:
+        assert stage.runtime_overrides == {"enforce_eager": True}

vllm_omni/config/stage_config.py

@@ -30,15 +30,6 @@ def get_pipeline_path(model_dir: str, filename: str) -> Path:
 logger = init_logger(__name__)
 
 
-def _warn_deprecated_kwargs(kwargs: dict[str, Any]) -> None:
-    if "cli_explicit_keys" in kwargs:
-        warnings.warn(
-            "cli_explicit_keys= is deprecated and ignored. Remove the kwarg.",
-            DeprecationWarning,
-            stacklevel=3,
-        )
-
-
 _STAGE_OVERRIDE_PATTERN = re.compile(r"^stage_(\d+)_(.+)$")
 
 
@@ -837,15 +828,6 @@ _PIPELINE_WIDE_ENGINE_FIELDS: tuple[str, ...] = (
 )
 
 
-def deploy_override_field_names() -> frozenset[str]:
-    """Return deploy-schema fields whose CLI defaults must not override YAML."""
-    return (
-        frozenset(_STAGE_DEPLOY_FIELDS)
-        | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS)
-        | frozenset({"async_chunk", "devices"})
-    )
-
-
 def _build_engine_args(
     ps: StagePipelineConfig,
     ds: StageDeployConfig | None,
@@ -1169,14 +1151,11 @@ class StageConfigFactory:
         model: str,
         cli_overrides: dict[str, Any] | None = None,
         deploy_config_path: str | None = None,
-        **deprecated_kwargs: Any,
     ) -> list[StageConfig] | None:
         """Load pipeline + deploy config, merge with CLI overrides.
 
         Checks _PIPELINE_REGISTRY first (new path), falls back to legacy YAML.
         """
-        _warn_deprecated_kwargs(deprecated_kwargs)
-
         if cli_overrides is None:
             cli_overrides = {}
 
@@ -1256,15 +1235,12 @@ class StageConfigFactory:
         model_type: str,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None = None,
-        **deprecated_kwargs: Any,
     ) -> list[StageConfig]:
         """Create StageConfigs from pipeline registry + deploy YAML.
 
         Precedence: caller-typed (non-None) value > deploy YAML >
         StageDeployConfig dataclass default.
         """
-        _warn_deprecated_kwargs(deprecated_kwargs)
-
         # Resolve deploy config path
         if deploy_config_path is None:
             deploy_path = _DEPLOY_DIR / f"{model_type}.yaml"

vllm_omni/diffusion/models/bagel/autoencoder.py

@@ -14,6 +14,16 @@ from dataclasses import dataclass
 import torch
 from einops import rearrange
 from torch import Tensor, nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
 
 
 @dataclass
@@ -322,3 +332,229 @@ class AutoEncoder(nn.Module):
 
     def forward(self, x: Tensor) -> Tensor:
         return self.decode(self.encode(x))
+
+
+class DistributedAutoEncoder(AutoEncoder, DistributedVaeMixin):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__(params)
+        self.init_distributed()
+
+        self.use_tiling = False
+        self.tile_sample_min_height = 512
+        self.tile_sample_min_width = 512
+        self.tile_sample_stride_height = 448
+        self.tile_sample_stride_width = 448
+        self.spatial_compression_ratio = params.downsample
+
+    def blend_v(self, above: Tensor, current: Tensor, blend_extent: int) -> Tensor:
+        blend_extent = min(above.shape[-2], current.shape[-2], blend_extent)
+        if blend_extent <= 0:
+            return current
+
+        for y in range(blend_extent):
+            alpha = y / blend_extent
+            current[:, :, y, :] = above[:, :, -blend_extent + y, :] * (1 - alpha) + current[:, :, y, :] * alpha
+        return current
+
+    def blend_h(self, left: Tensor, current: Tensor, blend_extent: int) -> Tensor:
+        blend_extent = min(left.shape[-1], current.shape[-1], blend_extent)
+        if blend_extent <= 0:
+            return current
+
+        for x in range(blend_extent):
+            alpha = x / blend_extent
+            current[:, :, :, x] = left[:, :, :, -blend_extent + x] * (1 - alpha) + current[:, :, :, x] * alpha
+        return current
+
+    def decode_tile_split(self, z: Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, latent_height, latent_width = z.shape
+
+        sample_height = latent_height * self.spatial_compression_ratio
+        sample_width = latent_width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        tiletask_list = []
+        for i in range(0, latent_height, tile_latent_stride_height):
+            for j in range(0, latent_width, tile_latent_stride_width):
+                tile = z[
+                    :,
+                    :,
+                    i : i + tile_latent_min_height,
+                    j : j + tile_latent_min_width,
+                ]
+                tiletask_list.append(
+                    TileTask(
+                        tile_id=len(tiletask_list),
+                        grid_coord=(i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        tensor=tile,
+                        workload=tile.shape[-2] * tile.shape[-1],
+                    )
+                )
+        grid_spec = GridSpec(
+            split_dims=(2, 3),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "sample_height": sample_height,
+                "sample_width": sample_width,
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+            },
+            output_dtype=z.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def decode_tile_exec(self, task: TileTask) -> Tensor:
+        tile = task.tensor
+        tile = tile / self.scale_factor + self.shift_factor
+        return self.decoder(tile)
+
+    def decode_tile_merge(self, coord_tensor_map: dict[tuple[int, ...], Tensor], grid_spec: GridSpec) -> Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        : self.tile_sample_stride_height,
+                        : self.tile_sample_stride_width,
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=-2)[
+            :,
+            :,
+            : grid_spec.tile_spec["sample_height"],
+            : grid_spec.tile_spec["sample_width"],
+        ]
+        return dec
+
+    def decode(self, z: Tensor) -> Tensor:
+        if not self.is_distributed_enabled():
+            return super().decode(z)
+
+        logger.debug("Bagel VAE decode running with distributed executor")
+        return self.distributed_executor.execute(
+            z,
+            DistributedOperator(
+                split=self.decode_tile_split,
+                exec=self.decode_tile_exec,
+                merge=self.decode_tile_merge,
+            ),
+            broadcast_result=True,
+        )
+
+    def encode_tile_split(self, x: Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, height, width = x.shape
+
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        tiletask_list = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            for j in range(0, width, self.tile_sample_stride_width):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_height,
+                    j : j + self.tile_sample_min_width,
+                ]
+                tiletask_list.append(
+                    TileTask(
+                        tile_id=len(tiletask_list),
+                        grid_coord=(i // self.tile_sample_stride_height, j // self.tile_sample_stride_width),
+                        tensor=tile,
+                        workload=tile.shape[-2] * tile.shape[-1],
+                    )
+                )
+        grid_spec = GridSpec(
+            split_dims=(2, 3),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "latent_height": latent_height,
+                "latent_width": latent_width,
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+                "tile_latent_stride_height": tile_latent_stride_height,
+                "tile_latent_stride_width": tile_latent_stride_width,
+            },
+            output_dtype=next(self.parameters()).dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> Tensor:
+        tile = task.tensor
+        z = self.reg(self.encoder(tile))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def encode_tile_merge(self, coord_tensor_map: dict[tuple[int, ...], Tensor], grid_spec: GridSpec) -> Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_latent_stride_height"],
+                        : grid_spec.tile_spec["tile_latent_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=-2)[
+            :,
+            :,
+            : grid_spec.tile_spec["latent_height"],
+            : grid_spec.tile_spec["latent_width"],
+        ]
+        return enc
+
+    def encode(self, x: Tensor) -> Tensor:
+        if not self.is_distributed_enabled():
+            return super().encode(x)
+        logger.debug("Bagel VAE encode running with distributed executor")
+        return self.distributed_executor.execute(
+            x,
+            DistributedOperator(
+                split=self.encode_tile_split,
+                exec=self.encode_tile_exec,
+                merge=self.encode_tile_merge,
+            ),
+            broadcast_result=True,
+        )

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -32,7 +32,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
-from .autoencoder import AutoEncoder, AutoEncoderParams
+from .autoencoder import AutoEncoder, AutoEncoderParams, DistributedAutoEncoder
 from .bagel_transformer import Bagel, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM
 
 logger = init_logger(__name__)
@@ -253,7 +253,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         )
         self.transformer = self.language_model.model
         ae_params: AutoEncoderParams = default_ae_params()
-        self.vae = AutoEncoder(ae_params)
+        self.vae = DistributedAutoEncoder(ae_params)
 
         self.bagel = Bagel(
             language_model=self.language_model,

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -43,7 +43,6 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
-from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 logger = init_logger(__name__)
@@ -538,8 +537,8 @@ class QwenImageCrossAttention(nn.Module):
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -576,8 +575,8 @@ class QwenImageCrossAttention(nn.Module):
             prefix=_join_prefix(prefix, "to_out"),
         )
 
-        self.norm_added_q = RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -992,7 +991,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -402,12 +402,12 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
             truncation=True,
             add_special_tokens=True,
             return_tensors="pt",
-        ).to(self.device)
+        )
         text_input_ids = text_inputs.input_ids
         untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
 
         if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
             logger.warning(
                 "The following part of your input was truncated because `max_sequence_length` is set to "
                 f" {max_sequence_length} tokens: {removed_text}"

vllm_omni/engine/arg_utils.py

@@ -1,5 +1,4 @@
 import argparse
-import dataclasses
 import json
 import os
 import tempfile
@@ -191,30 +190,6 @@ class OmniEngineArgs(EngineArgs):
         load_omni_general_plugins()
         super().__post_init__()
 
-    @classmethod
-    def from_cli_args(cls, args: argparse.Namespace) -> "OmniEngineArgs":
-        attrs = [attr.name for attr in dataclasses.fields(cls)]
-        engine_args = cls(**{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)})
-        engine_args._explicit_fields = frozenset(
-            attr for attr in attrs if hasattr(args, attr) and getattr(args, attr) is not None
-        )
-        return engine_args
-
-    @classmethod
-    def create(cls, **explicit: Any) -> "OmniEngineArgs":
-        """Tracks caller-set fields for ``Omni(..., engine_args=ea)``."""
-        ea = cls(**explicit)
-        ea._explicit_fields = frozenset(explicit.keys())
-        return ea
-
-    def explicit_kwargs(self) -> dict[str, Any]:
-        explicit = getattr(self, "_explicit_fields", None)
-        if explicit is None:
-            return {
-                f.name: getattr(self, f.name) for f in dataclasses.fields(self) if getattr(self, f.name) is not None
-            }
-        return {k: getattr(self, k) for k in explicit}
-
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):
             return True
@@ -420,8 +395,7 @@ class OrchestratorArgs:
           forwarding).
 
     Fields that BOTH orchestrator and engine genuinely need (e.g. ``model``,
-    ``log_stats``) should be listed in ``SHARED_FIELDS`` below; ``split_kwargs``
-    will copy them to both buckets.
+    ``log_stats``) should be listed in ``SHARED_FIELDS`` below.
     """
 
     # === Lifecycle ===
@@ -530,163 +504,3 @@ def internal_blacklist_keys() -> frozenset[str]:
     dataclass — this function updates automatically.
     """
     return orchestrator_field_names() - SHARED_FIELDS
-
-
-def split_kwargs(
-    kwargs: dict[str, Any],
-    *,
-    engine_cls: type | None = None,
-    user_typed: set[str] | None = None,
-    strict: bool = False,
-) -> tuple[OrchestratorArgs, dict[str, Any]]:
-    """Partition CLI kwargs into (orchestrator, engine) buckets.
-
-    Args:
-        kwargs: Raw dict, typically ``vars(args)``.
-        engine_cls: Engine dataclass used to whitelist-filter the engine
-            bucket. Defaults to ``OmniEngineArgs``. Pass a custom class
-            for testing.
-        user_typed: Keys the user actually typed on the command line. Used
-            to warn when a user-typed flag is unclassifiable.
-        strict: If True, raise ``ValueError`` on ambiguous (double-classified
-            but not in ``SHARED_FIELDS``) fields. Default False to keep the
-            rollout non-breaking; flip to True in tests and CI.
-
-    Returns:
-        ``(orchestrator_args, engine_kwargs)``. ``engine_kwargs`` has already
-        been whitelist-filtered against ``engine_cls`` — safe to pass directly
-        to ``engine_cls(**engine_kwargs)``.
-    """
-    if engine_cls is None:
-        engine_cls = OmniEngineArgs
-
-    orch_fields = orchestrator_field_names()
-    engine_fields = {f.name for f in fields(engine_cls)}
-
-    orch_kwargs: dict[str, Any] = {}
-    engine_candidate: dict[str, Any] = {}
-    shared_values: dict[str, Any] = {}
-    unclassified: dict[str, Any] = {}
-
-    for key, value in kwargs.items():
-        in_orch = key in orch_fields
-        in_engine = key in engine_fields
-        is_shared = key in SHARED_FIELDS
-
-        if is_shared:
-            shared_values[key] = value
-        elif in_orch and in_engine:
-            # Declared in both but not marked shared → ambiguous.
-            msg = (
-                f"Field {key!r} is defined on both OrchestratorArgs and "
-                f"{engine_cls.__name__} but is not in SHARED_FIELDS. "
-                f"This causes double-routing. Either remove the duplicate or "
-                f"add {key!r} to SHARED_FIELDS if the sharing is intentional."
-            )
-            if strict:
-                raise ValueError(msg)
-            logger.error(msg)
-            # Default: treat as orchestrator-only to preserve existing behavior.
-            orch_kwargs[key] = value
-        elif in_orch:
-            orch_kwargs[key] = value
-        elif in_engine:
-            engine_candidate[key] = value
-        else:
-            unclassified[key] = value
-
-    # Warn on user-typed but unclassifiable flags so we don't silently drop
-    # something the user cared about (fixes the class of bug that spawned #873).
-    if unclassified and user_typed:
-        user_typed_unknown = sorted(k for k in unclassified if k in user_typed)
-        if user_typed_unknown:
-            logger.warning(
-                "CLI flags not consumed by vllm-omni and dropped before "
-                "per-stage engine construction: %s. If these are vllm "
-                "frontend/uvicorn flags (host, port, ssl_*, api_key, …) this "
-                "is expected; otherwise check your spelling.",
-                user_typed_unknown,
-            )
-
-    # Engine bucket: shared + engine-only. We do NOT pass through unclassified
-    # fields — that's exactly the server/uvicorn noise we want to shed.
-    engine_kwargs = {**shared_values, **engine_candidate}
-
-    # Construct the orchestrator dataclass. Shared fields that OrchestratorArgs
-    # also declares get copied into its constructor.
-    orch_init: dict[str, Any] = dict(orch_kwargs)
-    for key, value in shared_values.items():
-        if key in orch_fields:
-            orch_init[key] = value
-    orch_args = OrchestratorArgs(**orch_init)
-
-    return orch_args, engine_kwargs
-
-
-def derive_server_dests_from_vllm_parser() -> frozenset[str]:
-    """Derive the set of argparse dests that belong to vllm's frontend/server.
-
-    Returns every dest registered by ``make_arg_parser`` that is NOT a field
-    of ``OmniEngineArgs`` and NOT a field of ``OrchestratorArgs``. Useful for
-    CI tests to assert all CLI flags are classifiable without maintaining
-    a hardcoded server list.
-
-    Returns empty frozenset if vllm's parser cannot be built (e.g. in a
-    minimal test environment).
-    """
-    try:
-        from vllm.entrypoints.openai.cli_args import make_arg_parser
-        from vllm.utils.argparse_utils import FlexibleArgumentParser
-    except ImportError:
-        logger.debug("Cannot import vllm parser — server-dest derivation skipped")
-        return frozenset()
-
-    try:
-        parser = make_arg_parser(FlexibleArgumentParser())
-        all_dests = {a.dest for a in parser._actions if a.dest and a.dest != "help"}
-    except Exception as exc:
-        logger.debug("Failed to build vllm parser: %s", exc)
-        return frozenset()
-
-    engine_fields = {f.name for f in fields(OmniEngineArgs)}
-    orch_fields = orchestrator_field_names()
-
-    return frozenset(all_dests - engine_fields - orch_fields - SHARED_FIELDS)
-
-
-def orchestrator_args_from_argparse(args: Any) -> OrchestratorArgs:
-    """Build an ``OrchestratorArgs`` from an ``argparse.Namespace``.
-
-    Only copies attributes that exist on the namespace — missing fields fall
-    back to the dataclass default. Useful when the full parser is already
-    built and ``vars(args)`` would include noise.
-    """
-    kwargs: dict[str, Any] = {}
-    for f in fields(OrchestratorArgs):
-        if hasattr(args, f.name):
-            value = getattr(args, f.name)
-            if value is not None or f.default is None:
-                kwargs[f.name] = value
-    return OrchestratorArgs(**kwargs)
-
-
-def nullify_stage_engine_defaults(parser: argparse.ArgumentParser) -> None:
-    """Reset stage-level engine flag defaults to ``None``; preserve real
-    default in help text. Only deploy-YAML override fields are touched.
-    Idempotent."""
-    from vllm_omni.config.stage_config import deploy_override_field_names
-
-    override_dests = deploy_override_field_names()
-
-    for action in parser._actions:
-        if action.dest in ("help", "version") or not action.option_strings:
-            continue
-        if action.dest not in override_dests:
-            continue
-        if action.default is None or action.default is argparse.SUPPRESS:
-            continue
-        if action.help and "(default:" not in action.help and "%(default)" not in action.help:
-            action.help = f"{action.help} (default: {action.default})"
-        action.default = None
-
-    parser._omni_nullified = True  # type: ignore[attr-defined]

vllm_omni/engine/async_omni_engine.py

@@ -21,7 +21,7 @@ import weakref
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from contextlib import ExitStack
 from dataclasses import asdict, dataclass
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import Any, Literal, cast
 
 import janus
 import torch
@@ -113,9 +113,6 @@ from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 from vllm_omni.platforms import current_omni_platform
 
-if TYPE_CHECKING:
-    from vllm_omni.engine.arg_utils import OmniEngineArgs
-
 logger = init_logger(__name__)
 
 _STARTUP_POLL_INTERVAL_S = 1.0
@@ -269,16 +266,17 @@ class AsyncOmniEngine:
     def __init__(
         self,
         model: str,
-        engine_args: OmniEngineArgs | None = None,
         stage_init_timeout: int = 300,
         init_timeout: int = 600,
         diffusion_batch_size: int = 1,
         single_stage_mode: bool = False,
         transfer_emitter: Any = None,
         log_stats: bool = False,
+        tokenizer: str | None = None,
         **kwargs: Any,
     ) -> None:
         self.model = model
+        self.tokenizer = tokenizer
         self.diffusion_batch_size = diffusion_batch_size
         startup_timeout = int(init_timeout)
         # Forwarded into Orchestrator so its _forward_to_next_stage path can
@@ -294,22 +292,6 @@ class AsyncOmniEngine:
 
         logger.info(f"[AsyncOmniEngine] Initializing with model {model}")
 
-        # Merge tracked engine_args fields into kwargs; explicit kwargs take priority.
-        if engine_args is not None:
-            if not hasattr(engine_args, "_explicit_fields"):
-                raise TypeError(
-                    "engine_args=OmniEngineArgs(...) is ambiguous under "
-                    "sentinel-default precedence. Use "
-                    "OmniEngineArgs.create(**explicit) or pass explicit kwargs "
-                    "directly."
-                )
-            ea_dict = engine_args.explicit_kwargs()
-            # Remove model since it is passed as a positional arg already.
-            ea_dict.pop("model", None)
-            kwargs = {**ea_dict, **kwargs}
-
-        self.tokenizer: str | None = kwargs.get("tokenizer")
-
         # ------------------------------------------------------------------ #
         # Single-stage mode detection                                        #
         # ------------------------------------------------------------------ #
@@ -2072,7 +2054,6 @@ class AsyncOmniEngine:
         stage_configs_path = kwargs.get("stage_configs_path", None)
         deploy_config_path = kwargs.pop("deploy_config", None)
         stage_overrides_json = kwargs.pop("stage_overrides", None)
-        kwargs.pop("_cli_explicit_keys", None)
         explicit_stage_configs = kwargs.pop("stage_configs", None)
         if explicit_stage_configs is not None:
             logger.warning(

vllm_omni/engine/stage_init_utils.py

@@ -482,6 +482,15 @@ def prepare_engine_environment() -> None:
         pass
 
 
+def _maybe_set_qwen3_omni_moe_env(engine_args_dict: dict[str, Any]) -> None:
+    if (
+        engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration"
+        and "VLLM_USE_FLASHINFER_MOE_FP16" not in os.environ
+    ):
+        os.environ["VLLM_USE_FLASHINFER_MOE_FP16"] = "0"
+        logger.info("[stage_init] Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni stage")
+
+
 def split_devices_for_replicas(
     devices_str: str | None,
     num_replicas: int,
@@ -762,6 +771,10 @@ def build_engine_args_dict(
     default_sp = _to_dict(getattr(stage_config, "default_sampling_params", {}))
     engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args"))
 
+    # TODO: Remove this after the performance regression is fixed
+    # Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni to avoid performance regression
+    _maybe_set_qwen3_omni_moe_env(engine_args_dict)
+
     return engine_args_dict
 
 

vllm_omni/entrypoints/cli/benchmark/main.py

@@ -9,7 +9,7 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase
 
 if typing.TYPE_CHECKING:
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
+    from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 
 class OmniBenchmarkSubcommand(CLISubcommand):
@@ -25,7 +25,7 @@ class OmniBenchmarkSubcommand(CLISubcommand):
     def validate(self, args: argparse.Namespace) -> None:
         pass
 
-    def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+    def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgumentParser:
         bench_parser = subparsers.add_parser(
             self.name, description=self.help, usage=f"vllm {self.name} <bench_type> [options]"
         )

vllm_omni/entrypoints/cli/main.py

@@ -24,10 +24,10 @@ def main():
             os.environ["VLLM_LOGGING_COLOR"] = "1"
 
         from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
-        from vllm.utils.argparse_utils import FlexibleArgumentParser
 
         import vllm_omni.entrypoints.cli.benchmark.main
         import vllm_omni.entrypoints.cli.serve
+        from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
         CMD_MODULES = [
             vllm_omni.entrypoints.cli.serve,
@@ -40,7 +40,7 @@ def main():
 
         _ensure_vllm_platform()
 
-        parser = FlexibleArgumentParser(
+        parser = TrackingArgumentParser(
             description="vLLM OMNI CLI",
             epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
         )

vllm_omni/entrypoints/cli/serve.py

@@ -19,11 +19,10 @@ from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.cli.logo import log_logo
 from vllm_omni.entrypoints.openai.api_server import omni_run_server
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 logger = init_logger(__name__)
 
@@ -84,10 +83,10 @@ class OmniServeCommand(CLISubcommand):
     name = "serve"
     # Parser stashed at subparser_init so ``cmd`` can resolve each user-typed
     # flag to its real ``dest`` via the parser's action table.
-    _parser: FlexibleArgumentParser | None = None
+    _parser: TrackingArgumentParser
 
     @staticmethod
-    def cmd(args: argparse.Namespace) -> None:
+    def cmd(args: TrackingNamespace) -> None:
         if not os.environ.get("VLLM_DISABLE_LOG_LOGO"):
             os.environ["VLLM_DISABLE_LOG_LOGO"] = "1"
             log_logo()
@@ -176,7 +175,7 @@ class OmniServeCommand(CLISubcommand):
             return
         validate_parsed_serve_args(args)
 
-    def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+    def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgumentParser:
         serve_parser = subparsers.add_parser(
             self.name,
             description=DESCRIPTION,
@@ -682,22 +681,10 @@ class OmniServeCommand(CLISubcommand):
         # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError.
         type(self)._parser = serve_parser
 
-        nullify_stage_engine_defaults(serve_parser)
         return serve_parser
 
 
-def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[str, Any]]:
-    """Create default diffusion stage configuration.
-
-    Uses AsyncOmniEngine's implementation which doesn't have OmegaConf
-    compatibility issues.
-    """
-    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-
-    return AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))
-
-
-def run_headless(args: argparse.Namespace) -> None:
+def run_headless(args: TrackingNamespace) -> None:
     """Run a single stage in headless mode.
 
     Honors ``--omni-dp-size-local``: launches that many replicas locally for
@@ -738,43 +725,39 @@ def run_headless(args: argparse.Namespace) -> None:
     stage_id: int | None = args.stage_id
     omni_master_address: str | None = args.omni_master_address
     omni_master_port: int | None = args.omni_master_port
+    worker_backend: str | None = args.worker_backend
+    stage_configs_path: str | None = args.stage_configs_path
     omni_replica_address: str | None = getattr(args, "omni_replica_address", None)
     omni_dp_size_local: int = max(1, int(getattr(args, "omni_dp_size_local", 1) or 1))
 
+    if not model:
+        raise ValueError("Failed to pass model from kwargs")
     if stage_id is None:
         raise ValueError("--stage-id is required in headless mode")
+    if omni_master_address is None or omni_master_port is None:
+        raise ValueError("--omni-master-address and --omni-master-port are required in headless mode")
+    if worker_backend != "multi_process":
+        raise ValueError("headless mode requires worker_backend=multi_process")
+
+    # Filter down to a dict of things explicitly requested by the user
+    args_dict = args.get_explicit_kwargs_dict()
 
     # ``--replica-id`` is deprecated and ignored — replica ids are
     # auto-assigned by ``OmniMasterServer`` so headless processes carry
     # no knowledge of their per-replica id at launch time. Warn (don't
     # error) when the operator still supplies it so existing launchers
     # keep working with a single log line.
-    explicit_cli_keys: set[str] = getattr(args, "_cli_explicit_keys", set()) or set()
-    if "replica_id" in explicit_cli_keys:
+    if "replica_id" in args_dict:
         logger.warning(
             "[Headless] --replica-id is deprecated and ignored "
             "(supplied value: %s). Replica ids are auto-assigned by the "
             "master server.",
             args.replica_id,
         )
-    if omni_master_address is None or omni_master_port is None:
-        raise ValueError("--omni-master-address and --omni-master-port are required in headless mode")
-    api_server_count = args.api_server_count or 0
-    if api_server_count > 1:
-        raise ValueError("api_server_count can't be set in headless mode")
-    if args.worker_backend != "multi_process":
-        raise ValueError("headless mode requires worker_backend=multi_process")
 
-    args_dict = vars(args).copy()
-    args_dict.pop("_cli_explicit_keys", None)
-    # Forward ``--deploy-config`` so the headless reads the same YAML the
-    # head was launched with — otherwise ``load_and_resolve_stage_configs``
-    # falls back to ``vllm_omni/deploy/<model>.yaml`` and the headless's
-    # view of ``stage.runtime.devices`` diverges from the head's, breaking
-    # the per-replica device split.
     config_path, stage_configs = load_and_resolve_stage_configs(
         model,
-        args_dict.get("stage_configs_path"),
+        stage_configs_path,
         args_dict,
         deploy_config_path=args_dict.get("deploy_config"),
     )

vllm_omni/entrypoints/omni_base.py

@@ -1,13 +1,10 @@
 from __future__ import annotations
 
-import argparse
 import os
-import sys
 import time
-import warnings
 import weakref
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Literal
+from typing import Any, Literal
 
 import huggingface_hub
 from vllm.logger import init_logger
@@ -29,9 +26,7 @@ from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetric
 from vllm_omni.metrics.transfer import OmniTransferMetrics
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.outputs import OmniRequestOutput
-
-if TYPE_CHECKING:
-    from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.utils.tracking_parser import TrackingNamespace
 
 logger = init_logger(__name__)
 
@@ -105,48 +100,44 @@ class OmniBase(PDDisaggregationMixin):
     @classmethod
     def from_cli_args(
         cls,
-        args: argparse.Namespace,
-        *,
-        parser: argparse.ArgumentParser | None = None,
-        **overrides: Any,
+        args: TrackingNamespace,
+        model: str | None = None,
     ) -> OmniBase:
-        """Deprecated argparse builder.
-
-        Build from argparse. If ``parser`` is passed and not yet nullified,
-        un-typed engine fields are reset to ``None``. New callers should
-        nullify deploy-overriding parser defaults with
-        ``nullify_stage_engine_defaults(parser)`` and construct Omni/AsyncOmni
-        directly.
+        """Build from a TrackingNamespace parsed by TrackingArgumentParser.
+        Only args that are explicitly passed to parse_args are forwarded.
         """
-        warnings.warn(
-            "`from_cli_args()` is deprecated. Nullify deploy-overriding parser defaults "
-            "with `nullify_stage_engine_defaults(parser)` and construct Omni/AsyncOmni "
-            "directly from `vars(args)`.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        kwargs: dict[str, Any] = {k: v for k, v in vars(args).items() if not k.startswith("_")}
+        if not isinstance(args, TrackingNamespace):
+            raise TypeError(
+                f"expected args to be of type TrackingNamespace, got {type(args)}. "
+                "Hint: did you parse your args with TrackingArgumentParser?"
+            )
 
-        if parser is not None and not getattr(parser, "_omni_nullified", False):
-            from vllm_omni.config.stage_config import deploy_override_field_names
-            from vllm_omni.entrypoints.utils import detect_explicit_cli_keys
+        explicit_kwargs = args.get_explicit_kwargs_dict()
+        args_model = explicit_kwargs.pop("model", None) or args.model
+        if model is not None and args_model is not None and model != args_model:
+            raise ValueError(
+                f"explicit model kwarg and args.model were both provided, but do not match [{model} != {args_model}]"
+            )
 
-            explicit = detect_explicit_cli_keys(sys.argv[1:], parser) or set()
-            override_dests = deploy_override_field_names()
-            for key in list(kwargs):
-                if key in override_dests and key not in explicit:
-                    kwargs[key] = None
+        if model is None and args_model is None:
+            raise ValueError(
+                "model must be explicitly passed as a parsed arg in the TrackingNamespace or directly provided."
+            )
 
-        kwargs.update(overrides)
-        return cls(**kwargs)
+        resolved_model = model or args_model
+        return cls(model=resolved_model, **explicit_kwargs)
 
     def __init__(
         self,
         model: str,
         **kwargs: Any,
     ) -> None:
-        engine_args: OmniEngineArgs | None = kwargs.pop("engine_args", None)
-
+        if "engine_args" in kwargs:
+            logger.warning(
+                "engine_args were passed as a kwarg to an Omni instance; this is not supported. "
+                "You should instead, pass the keyword arguments used to initialize the engine args "
+                "directly to this object's initializer."
+            )
         stage_init_timeout = kwargs.pop("stage_init_timeout", 300)
         init_timeout = kwargs.pop("init_timeout", 600)
         log_stats = kwargs.pop("log_stats", False)
@@ -180,7 +171,6 @@ class OmniBase(PDDisaggregationMixin):
         st = time.time()
         self.engine = AsyncOmniEngine(
             model=model,
-            engine_args=engine_args,
             init_timeout=init_timeout,
             stage_init_timeout=stage_init_timeout,
             diffusion_batch_size=diffusion_batch_size,

vllm_omni/entrypoints/openai/api_server.py

@@ -136,6 +136,7 @@ from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_reques
 from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference
 from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 logger = init_logger(__name__)
 router = APIRouter()
@@ -510,7 +511,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
 @asynccontextmanager
 async def build_async_omni(
-    args: Namespace,
+    args: TrackingNamespace,
     *,
     disable_frontend_multiprocessing: bool | None = None,
     client_config: dict[str, Any] | None = None,
@@ -550,7 +551,7 @@ async def build_async_omni(
 
 @asynccontextmanager
 async def build_async_omni_from_stage_config(
-    args: Namespace,
+    args: TrackingNamespace,
     *,
     disable_frontend_multiprocessing: bool = False,
 ) -> AsyncIterator[EngineClient]:
@@ -604,9 +605,9 @@ async def build_async_omni_from_stage_config(
             logger.debug("Pre-loading transformers_modules failed: %s", e)
 
     try:
-        kwargs = vars(args).copy()
-        kwargs.pop("model", None)
-        async_omni = AsyncOmni(model=args.model, **kwargs)
+        kwargs = args.get_explicit_kwargs_dict()
+        model = kwargs.pop("model", None) or args.model
+        async_omni = AsyncOmni(model=model, **kwargs)
 
         # # Don't keep the dummy data in memory
         # await async_llm.reset_mm_cache()
@@ -3040,25 +3041,23 @@ async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request):
 
 
 if __name__ == "__main__":
-    import argparse
     import asyncio
 
     from vllm.entrypoints.openai.cli_args import make_arg_parser
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    parser = argparse.ArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server")
+    parser = TrackingArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server")
     parser = make_arg_parser(parser)
     registered_flags = set()
     for action in parser._actions:
         registered_flags.update(action.option_strings)
+
+    # FIXME - this is broken on `main`. `make_arg_parser` does not handle omni engine args properly.
     if "--omni" not in registered_flags:
         parser.add_argument("--omni", action="store_true", default=False, help="Enable vLLM-Omni mode.")
     if "--enable-sleep-mode" not in registered_flags:
         parser.add_argument(
             "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode."
         )
-    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
     if not hasattr(args, "model_tag"):
         setattr(args, "model_tag", args.model)

vllm_omni/entrypoints/utils.py

@@ -1,4 +1,3 @@
-import argparse
 import os
 import types
 from collections import Counter
@@ -24,81 +23,11 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent
 logger = init_logger(__name__)
 
 
-def _warn_deprecated_explicit_keys(kwargs: dict[str, Any]) -> None:
-    if "cli_explicit_keys" in kwargs:
-        import warnings
-
-        warnings.warn(
-            "cli_explicit_keys= is deprecated and ignored. Remove the kwarg.",
-            DeprecationWarning,
-            stacklevel=3,
-        )
-
-
 _DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = {
     "GlmImagePipeline": "glm_image",
 }
 
 
-def detect_explicit_cli_keys(
-    argv: list[str],
-    parser: argparse.ArgumentParser | None = None,
-) -> set[str]:
-    """Walk ``argv`` and return the set of ``dest`` attribute names the user
-    explicitly provided (e.g. ``--max-num-seqs 64`` → ``max_num_seqs``).
-
-    Used to distinguish user-typed CLI args from argparse default values so
-    deploy YAMLs are not silently overridden by parser defaults. Shared
-    across online (``vllm serve``) and offline (scripts, examples, tests,
-    CI) entry points — offline callers that parse CLI args via argparse
-    should invoke this on ``sys.argv[1:]`` and pass the result through to
-    ``AsyncOmni`` / ``Omni`` via the ``_cli_explicit_keys`` kwarg.
-
-    When ``parser`` is provided, each token is looked up in the parser's
-    action table to find its real ``dest``. This correctly handles flags
-    with ``dest=`` overrides, alias flags (e.g. ``--usp`` /
-    ``--ulysses-degree`` both mapping to ``ulysses_degree``), and
-    ``--disable-foo`` / ``store_false`` patterns that map to a differently
-    named dest. Callers with access to an ``argparse.ArgumentParser`` should
-    always pass it.
-
-    When ``parser`` is ``None``, a name-based heuristic is used as a
-    fallback (hyphens → underscores, plus a ``no_`` prefix strip for
-    ``argparse.BooleanOptionalAction``). This is correct for simple flags
-    but silently misidentifies ``--disable-X``-style flags and explicit
-    ``dest=`` overrides, so prefer the parser-aware form.
-    """
-    if parser is not None:
-        dest_map: dict[str, str] = {}
-        for action in parser._actions:
-            for opt in action.option_strings:
-                dest_map[opt] = action.dest
-        explicit: set[str] = set()
-        for tok in argv:
-            if not tok.startswith("--"):
-                continue
-            flag = tok.split("=", 1)[0]
-            dest = dest_map.get(flag)
-            if dest is not None:
-                explicit.add(dest)
-        return explicit
-
-    # Fallback: name-based heuristic (legacy path for callers without a parser).
-    explicit = set()
-    for tok in argv:
-        if not tok.startswith("--"):
-            continue
-        name = tok[2:].split("=", 1)[0]
-        if not name:
-            continue
-        attr = name.replace("-", "_")
-        explicit.add(attr)
-        # BooleanOptionalAction: --no-foo records as dest `foo`, not `no_foo`.
-        if attr.startswith("no_"):
-            explicit.add(attr[3:])
-    return explicit
-
-
 def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None:
     """Inject connector configuration into stage engine arguments."""
     # Prepare omni_kv_config dict
@@ -383,7 +312,6 @@ def load_stage_configs_from_model(
     base_engine_args: dict | None = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-    **deprecated_kwargs: Any,
 ) -> list:
     """Load stage configurations from model's default config file.
 
@@ -402,8 +330,6 @@ def load_stage_configs_from_model(
     Returns:
         List of stage configuration dictionaries
     """
-    _warn_deprecated_explicit_keys(deprecated_kwargs)
-
     if base_engine_args is None:
         base_engine_args = {}
 
@@ -577,7 +503,6 @@ def load_and_resolve_stage_configs(
     default_stage_cfg_factory: Any = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-    **deprecated_kwargs: Any,
 ) -> tuple[str, list]:
     """Load stage configurations from model or YAML file with fallback to defaults.
 
@@ -621,8 +546,6 @@ def load_and_resolve_stage_configs(
                 stage_configs_path,
             )
 
-    _warn_deprecated_explicit_keys(deprecated_kwargs)
-
     if deploy_config_path is not None:
         config_path = deploy_config_path
         stage_configs = load_stage_configs_from_model(

vllm_omni/utils/tracking_parser.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+from typing import Any
+
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+UNSET = object()
+
+
+class TrackingNamespace(argparse.Namespace):
+    """Proxy that wraps an argparse namespace with explicit keys, which
+    can be filtered down to a dict containing only explicitly passed values.
+    """
+
+    def __init__(self, unfiltered_ns: argparse.Namespace, explicit_keys: frozenset[str]) -> None:
+        # We never have nested tracking namespaces, but explicitly guard
+        # against them to prevent bad behavior with nested __dict__ overrides.
+        if isinstance(unfiltered_ns, TrackingNamespace):
+            raise ValueError("Tracking namespaces cannot be nested")
+
+        self.unfiltered_ns = unfiltered_ns
+        self.explicit_keys = explicit_keys
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name in ("unfiltered_ns", "explicit_keys"):
+            object.__setattr__(self, name, value)
+        else:
+            setattr(self.unfiltered_ns, name, value)
+
+    def get_explicit_kwargs_dict(self):
+        """Return a dict containing only the explicitly passed key-value pairs."""
+        return {k: v for k, v in vars(self.unfiltered_ns).items() if k in self.explicit_keys}
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.unfiltered_ns, name)
+
+    @property
+    def __dict__(self):
+        # NOTE: We do this so that vars() etc forward directly into the encapsulated namespace,
+        # which makes this class a drop-in replacement for the original namespace, while also
+        # ensuring that updates to the encapsulated namespace are correctly reflected.
+        return self.unfiltered_ns.__dict__
+
+
+class TrackingGroup:
+    """Proxy that wraps an argument group and its corresponding shadow group."""
+
+    def __init__(
+        self,
+        real_group: argparse._ArgumentGroup,
+        shadow_group: argparse._ArgumentGroup,
+    ):
+        self._real = real_group
+        self._shadow = shadow_group
+
+    def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action:
+        """Add an argument to the real group and to the shadow group."""
+        action = self._real.add_argument(*args, **kwargs)
+        default_kwargs = {**kwargs, "default": UNSET}
+        self._shadow.add_argument(*args, **default_kwargs)
+        return action
+
+    def __getattr__(self, name: str) -> Any:
+        # Any attribute access is forwarded to the real argument group.
+        return getattr(self._real, name)
+
+
+class TrackingSubparsers:
+    """Proxy that wraps a subparser and its corresponding shadow subparser."""
+
+    def __init__(
+        self,
+        real_sub: argparse._SubParsersAction,
+        shadow_sub: argparse._SubParsersAction,
+    ):
+        self._real = real_sub
+        self._shadow = shadow_sub
+
+    def add_parser(self, name, *args, **kwargs):
+        """Add a parser to the encapsulated real parser and its shadow."""
+        real_parser = self._real.add_parser(name, *args, **kwargs)
+        # real_parser is a TrackingArgumentParser with its own _shadow.
+        # Reuse that shadow as the parent shadow's child — so when
+        # real_parser.add_argument() mirrors to real_parser._shadow,
+        # the parent's shadow sees it too.
+        self._shadow._name_parser_map[name] = real_parser._shadow
+        return real_parser
+
+    def __getattr__(self, name: str) -> Any:
+        # Any attribute access is forwarded to the real subparser.
+        return getattr(self._real, name)
+
+
+class TrackingArgumentParser(FlexibleArgumentParser):
+    """Drop-in replacement for FlexibleArgumentParser, which tracks keys that
+    were explicitly passed as args on the parser namespace.
+
+    Unfortunately, Argparse does not provide an easy way of doing this without
+    depending on a lot of internal attributes, so we implement it by instead
+    using a 'shadow' parser, which is essentially a clone of the parser, where
+    defaults are overridden to `None`. By comparing the parser against its
+    shadow, we can tell which values were passed in a non-destructive manner.
+    """
+
+    def __init__(self, *args, **kwargs):
+        # NOTE: We have to define the shadow parser before calling init,
+        # with add_help=False, since otherwise init will call add_argument
+        # and delegate to the override on this class and cause problems.
+        shadow_kwargs = {**kwargs, "add_help": False}
+        self._shadow = FlexibleArgumentParser(*args, **shadow_kwargs)
+        super().__init__(*args, **kwargs)
+
+    def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action:
+        """Add an arg to the parser & the shadow, where the latter has UNSET for the default."""
+        action = super().add_argument(*args, **kwargs)
+        shadow_kwargs = {**kwargs, "default": UNSET}
+        self._shadow.add_argument(*args, **shadow_kwargs)
+        return action
+
+    def add_argument_group(self, *args, **kwargs) -> TrackingGroup:
+        real_group = super().add_argument_group(*args, **kwargs)
+        shadow_group = self._shadow.add_argument_group(*args, **kwargs)
+        return TrackingGroup(real_group, shadow_group)
+
+    def add_mutually_exclusive_group(self, *args, **kwargs) -> TrackingGroup:
+        real_group = super().add_mutually_exclusive_group(*args, **kwargs)
+        shadow_group: argparse._MutuallyExclusiveGroup = self._shadow.add_mutually_exclusive_group(*args, **kwargs)
+        return TrackingGroup(real_group, shadow_group)
+
+    def add_subparsers(self, *args, **kwargs) -> TrackingSubparsers:
+        real_sub = super().add_subparsers(*args, **kwargs)
+        shadow_sub = self._shadow.add_subparsers(*args, **kwargs)
+        return TrackingSubparsers(real_sub, shadow_sub)
+
+    def build_tracking_namespace(self, real_ns: argparse.Namespace, shadow_ns: argparse.Namespace) -> TrackingNamespace:
+        """Build a tracking namespace for the real / shadow namespaces."""
+        explicit_keys = frozenset(k for k, v in vars(shadow_ns).items() if v is not UNSET)
+        return TrackingNamespace(real_ns, explicit_keys)
+
+    def parse_args(
+        self,
+        args: list[str] | None = None,
+        namespace: argparse.Namespace | None = None,
+    ) -> TrackingNamespace:
+        """Parse the args on the real/shadow parser."""
+        # Only the real parser should use the namespace if one is,
+        # given since shadow parser will set its own defaults to None.
+        real_ns = super().parse_args(args, namespace)
+        shadow_ns = self._shadow.parse_args(args)
+        if real_ns is None or shadow_ns is None:
+            raise ValueError("Parse args created empty namespaces")
+
+        # If this is called through parse_known_args on self, we will already
+        # get a TrackingNamespace back, which will already have set the explicit
+        # keys through build_tracking_namespace, so no need to do it again.
+        if isinstance(real_ns, TrackingNamespace):
+            return real_ns
+
+        return self.build_tracking_namespace(real_ns, shadow_ns)
+
+    def parse_known_args(
+        self,
+        args: list[str] | None = None,
+        namespace: argparse.Namespace | None = None,
+    ) -> tuple[TrackingNamespace, list[str]]:
+        """Parse the known args on the real/shadow parser."""
+        real_ns, remaining = super().parse_known_args(args, namespace)
+        shadow_ns, _ = self._shadow.parse_known_args(args)
+        tracked_ns = self.build_tracking_namespace(real_ns, shadow_ns)
+
+        return tracked_ns, remaining