[Refactor] Share offline profiler argument parser

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -21,6 +21,7 @@ from vllm.multimodal.media.audio import load_audio
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.profiler import add_profiler_config_arg
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
@@ -325,6 +326,7 @@ def main(args):
     else:
         query_result = query_func()
     args.quantization_config = quantization_config
+    # ``vars(args)`` below forwards profiler_config to Omni.
     omni_kwargs = vars(args).copy()
     # Override CLI --model with the derived model_name.
     omni_kwargs["model"] = model_name
@@ -378,11 +380,11 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.profiler_config is not None
     if profiler_enabled and hasattr(omni, "start_profile"):
         omni.start_profile(stages=[0])
     elif profiler_enabled:
-        print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
+        print("[Warn] --profiler-config was provided, but current engine does not support profiler controls.")
     omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
@@ -458,6 +460,7 @@ def parse_args():
         default=False,
         help="Enable writing detailed statistics (default: disabled)",
     )
+    add_profiler_config_arg(parser)
     parser.add_argument(
         "--stage-init-timeout",
         type=int,

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -14,6 +14,7 @@ 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
 from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+from vllm_omni.utils.profiler import add_profiler_config_arg
 
 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -45,6 +46,7 @@ def run_e2e():
         help="Override the deploy config path. If unset, auto-loads "
         "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
     )
+    add_profiler_config_arg(parser)
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
         "--prompt-text",
@@ -79,6 +81,7 @@ def run_e2e():
         deploy_config=args.deploy_config,
         tokenizer=args.tokenizer,
         log_stats=True,
+        profiler_config=args.profiler_config,
     )
 
     sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
@@ -149,8 +152,8 @@ def run_e2e():
 
     sampling_params_list = [gpt_sampling, s2mel_sampling]
 
-    # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
-    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
+    profiler_enabled = args.profiler_config is not None
+    if profiler_enabled:
         print("Starting profiler...")
         omni.start_profile()
 
@@ -158,7 +161,7 @@ def run_e2e():
     outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
 
     # Stop profiling and get results
-    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
+    if profiler_enabled:
         print("Stopping profiler...")
         profile_results = omni.stop_profile()
         print(f"Profile traces saved to: {profile_results}")

tests/diffusion/kernels/mot/test_mot_gemm.py

@@ -1,95 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.layers.mot.ops import mot_gemm
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-_CONFIG = {
-    "BLOCK_SIZE_M": 16,
-    "BLOCK_SIZE_N": 64,
-    "BLOCK_SIZE_K": 32,
-    "GROUP_SIZE_M": 8,
-    "num_warps": 4,
-    "num_stages": 2,
-}
-
-
-class _RecordingKernel:
-    def __init__(self) -> None:
-        self.calls: list[tuple[tuple[object, ...], dict[str, object]]] = []
-
-    def __getitem__(self, grid):
-        del grid
-        return self._launch
-
-    def _launch(self, *args, **kwargs) -> None:
-        self.calls.append((args, kwargs))
-
-
-@pytest.mark.parametrize(
-    ("bias_text", "bias_vae"),
-    [(None, None), ("text", None), (None, "vae"), ("text", "vae")],
-    ids=["no-bias", "text-only", "vae-only", "both"],
-)
-def test_invoke_mot_gemm_accepts_independent_biases(
-    bias_text: str | None,
-    bias_vae: str | None,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    kernel = _RecordingKernel()
-    monkeypatch.setattr(mot_gemm, "mot_unified_gemm_kernel", kernel)
-
-    A = torch.ones((3, 4), dtype=torch.float16)
-    B_text = torch.ones((4, 5), dtype=torch.float16)
-    B_vae = torch.ones((4, 5), dtype=torch.float16)
-    C = torch.empty((3, 5), dtype=torch.float16)
-    text_indices = torch.tensor([0, 2], dtype=torch.long)
-    vae_indices = torch.tensor([1], dtype=torch.long)
-    text_bias = torch.ones(5, dtype=torch.float16) if bias_text else None
-    vae_bias = torch.ones(5, dtype=torch.float16) if bias_vae else None
-
-    mot_gemm.invoke_mot_gemm(
-        A=A,
-        B_text=B_text,
-        B_vae=B_vae,
-        C=C,
-        bias_text=text_bias,
-        bias_vae=vae_bias,
-        text_indices=text_indices,
-        vae_indices=vae_indices,
-        A_scale=None,
-        B_text_scale=None,
-        B_vae_scale=None,
-        use_fp8_w8a8=False,
-        use_int8_w8a8=False,
-        use_int8_w8a16=False,
-        use_int4_w4a16=False,
-        A_per_channel_quant=False,
-        B_per_channel_quant=False,
-        config=_CONFIG,
-    )
-
-    assert len(kernel.calls) == 1
-    args, kwargs = kernel.calls[0]
-    assert kwargs["HAS_BIAS_TEXT"] is (text_bias is not None)
-    assert kwargs["HAS_BIAS_VAE"] is (vae_bias is not None)
-    assert "HAS_BIAS" not in kwargs
-
-    launched_text_bias, launched_vae_bias = args[4:6]
-    if text_bias is None and vae_bias is None:
-        assert launched_text_bias == launched_vae_bias == 0
-    elif text_bias is None:
-        assert launched_text_bias is vae_bias
-        assert launched_vae_bias is vae_bias
-    elif vae_bias is None:
-        assert launched_text_bias is text_bias
-        assert launched_vae_bias is text_bias
-    else:
-        assert launched_text_bias is text_bias
-        assert launched_vae_bias is vae_bias

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -513,57 +513,6 @@ def test_mot_o_proj(
             )
 
 
-@pytest.mark.parametrize(
-    ("bias_text", "bias_vae"),
-    [(True, False), (False, True)],
-    ids=["text-only", "vae-only"],
-)
-def test_mot_row_one_sided_bias(bias_text: bool, bias_vae: bool):
-    """The fused row projection must support a bias on only one expert."""
-    dcfg = _parse_dtype("w16a16_bf16")
-    K = 32
-    N = 64
-    M = 5
-
-    with set_current_vllm_config(VllmConfig()):
-        text_linear = RowParallelLinear(
-            K,
-            N,
-            bias=bias_text,
-            input_is_parallel=True,
-            params_dtype=dcfg.torch_dtype,
-            disable_tp=True,
-        ).cuda()
-        vae_linear = RowParallelLinear(
-            K,
-            N,
-            bias=bias_vae,
-            input_is_parallel=True,
-            params_dtype=dcfg.torch_dtype,
-            disable_tp=True,
-        ).cuda()
-        mot_linear = MoTRowParallelLinear(
-            K,
-            N,
-            bias=bias_text,
-            vae_bias=bias_vae,
-            input_is_parallel=True,
-            params_dtype=dcfg.torch_dtype,
-            disable_tp=True,
-        ).cuda()
-
-        _sync_weights(text_linear, vae_linear, mot_linear)
-        text_idx = torch.tensor([0, 2], dtype=torch.long, device="cuda")
-        vae_idx = torch.tensor([1, 3, 4], dtype=torch.long, device="cuda")
-        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
-
-        with torch.no_grad():
-            ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear)
-            mot_out, _ = mot_linear(x, text_idx, vae_idx)
-
-        _check_and_report(ref, mot_out, f"Row one-sided bias text={bias_text} vae={bias_vae}")
-
-
 # =========================================================================
 #  Test: und-mode (text_indices=None) — falls back to standard forward
 # =========================================================================

tests/docs/test_offline_profiler_config.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_EXAMPLES = (
+    "examples/offline_inference/text_to_speech/cosyvoice3/end2end.py",
+    "examples/offline_inference/qwen2_5_omni/end2end.py",
+)
+
+
+@pytest.mark.parametrize("relative_path", _EXAMPLES)
+def test_offline_examples_use_typed_profiler_config(relative_path: str):
+    source = (_REPO_ROOT / relative_path).read_text(encoding="utf-8")
+
+    assert "add_profiler_config_arg(parser)" in source
+    assert "profiler_enabled = args.profiler_config is not None" in source
+    assert "VLLM_TORCH_PROFILER_DIR" not in source

vllm_omni/diffusion/layers/mot/ops/mot_gemm.py

@@ -238,7 +238,6 @@ def _get_mot_pointers(
     cur_stride_bk = stride_bk_vae
     cur_stride_bn = stride_bn_vae
     M_limit = M_vae
-    is_vae = 1
 
     # Text Path
     if pid < num_pid_m_text * num_pid_n:
@@ -252,7 +251,6 @@ def _get_mot_pointers(
         cur_stride_bk = stride_bk_text
         cur_stride_bn = stride_bn_text
         M_limit = M_text
-        is_vae = 0
 
     # 3. Calculate Grid coordinates(grouping)
     cur_num_pid_m = tl.cdiv(M_limit, BLOCK_SIZE_M)
@@ -285,7 +283,6 @@ def _get_mot_pointers(
         offs_n,
         n_mask,  # N-dim info
         M_limit,  # Boundary
-        is_vae,  # Selected expert kind: 0=text, 1=VAE
         cur_b_ptr,
         cur_bias_ptr,  # Selected Pointers
         cur_scale_b_ptr,  # Selected Scale Pointer
@@ -536,8 +533,7 @@ def mot_unified_gemm_kernel(
     # Quant Control
     # 0=None, 1=W8A8, 2=W8A16, 3=W4A16
     QUANT_TYPE: tl.constexpr,
-    HAS_BIAS_TEXT: tl.constexpr = False,
-    HAS_BIAS_VAE: tl.constexpr = False,
+    HAS_BIAS: tl.constexpr = False,
 ):
     pid = tl.program_id(axis=0)
 
@@ -552,7 +548,6 @@ def mot_unified_gemm_kernel(
         offs_n,
         n_mask,
         M_limit,
-        is_vae,
         cur_b_ptr,
         cur_bias_ptr,
         cur_scale_b_ptr,
@@ -680,12 +675,8 @@ def mot_unified_gemm_kernel(
     # -----------------------------------------------------------
 
     # Bias Add
-    if HAS_BIAS_TEXT or HAS_BIAS_VAE:
+    if HAS_BIAS:
         bias = tl.load(cur_bias_ptr + offs_n, mask=n_mask, other=0.0)
-        if HAS_BIAS_TEXT and not HAS_BIAS_VAE:
-            bias = tl.where(is_vae == 0, bias, 0.0)
-        elif HAS_BIAS_VAE and not HAS_BIAS_TEXT:
-            bias = tl.where(is_vae == 1, bias, 0.0)
         c = c + bias[None, :]  # reshape to (1, BLOCK_SIZE_N)
 
     # Cast C into the output dtype
@@ -841,8 +832,11 @@ def invoke_mot_gemm(
     STRIDE_BK_IS_1 = (B_text.stride(0) == 1) and (B_vae.stride(0) == 1)
     STRIDE_BN_IS_1 = (B_text.stride(1) == 1) and (B_vae.stride(1) == 1)
 
-    has_bias_text = bias_text is not None
-    has_bias_vae = bias_vae is not None
+    # bias check
+    assert (bias_text is None) == (bias_vae is None), (
+        "Bias must be provided for both Text and VAE simultaneously, or neither."
+    )
+    has_bias = bias_text is not None
 
     # --- 3. Grid Calculation ---
     def grid(META):
@@ -859,8 +853,7 @@ def invoke_mot_gemm(
             "ACCUMULATOR_DTYPE": ACCUMULATOR_DTYPE,
             "COMPUTE_DTYPE": COMPUTE_DTYPE,
             "OUTPUT_DTYPE": OUTPUT_DTYPE,
-            "HAS_BIAS_TEXT": has_bias_text,
-            "HAS_BIAS_VAE": has_bias_vae,
+            "HAS_BIAS": has_bias,
             "EVEN_K": EVEN_K,
             "EVEN_N": EVEN_N,
             "STRIDE_AK_IS_1": STRIDE_AK_IS_1,
@@ -873,17 +866,8 @@ def invoke_mot_gemm(
     p_a_scale = A_scale if A_scale is not None else 0
     p_b_text_scale = B_text_scale if B_text_scale is not None else 0
     p_b_vae_scale = B_vae_scale if B_vae_scale is not None else 0
-    # Triton requires both route-selected bias arguments to have a pointer
-    # type. Alias the missing route to the available tensor; its compile-time
-    # bias flag and runtime route mask keep it from being read for that route.
-    if bias_text is None:
-        p_bias_text = bias_vae if bias_vae is not None else 0
-    else:
-        p_bias_text = bias_text
-    if bias_vae is None:
-        p_bias_vae = bias_text if bias_text is not None else 0
-    else:
-        p_bias_vae = bias_vae
+    p_bias_text = bias_text if bias_text is not None else 0
+    p_bias_vae = bias_vae if bias_vae is not None else 0
 
     # Quantization granularity
     stride_scale_a = 1 if A_per_channel_quant else 0

vllm_omni/utils/profiler.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import json
+from typing import Any
+
+
+def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
+    """Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
+    return config
+
+
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    """Parse the JSON object passed to ``--profiler-config``."""
+    return parse_json_object(value, flag_name="--profiler-config")
+
+
+def add_profiler_config_arg(parser: argparse.ArgumentParser) -> None:
+    """Add the shared JSON profiler configuration option to ``parser``."""
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help=(
+            "JSON profiler config for torch/cuda profiling, e.g. "
+            '\'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.'
+        ),
+    )