108 lines
4.3 KiB
Python
108 lines
4.3 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
"""Ban direct torch.cuda.* usage outside an allowlist.
|
||
|
|
|
||
|
|
Modeled on vllm/tools/pre_commit/check_torch_cuda.py. Call sites should use
|
||
|
|
current_omni_platform / OmniPlatform instead of CUDA-only torch.cuda APIs.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import sys
|
||
|
|
|
||
|
|
import regex as re
|
||
|
|
|
||
|
|
# Match `torch.cuda.xxx` device/memory/seed helpers that have portable Omni
|
||
|
|
# platform equivalents. CUDAGraph / is_available / stream capture are not in
|
||
|
|
# this list (same split as upstream vLLM).
|
||
|
|
_TORCH_CUDA_PATTERNS = [
|
||
|
|
r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|mem_get_info|set_device)\b",
|
||
|
|
# Match torch.cuda.device(...) including args. `device\()\b` misses
|
||
|
|
# `device()` (no word char after `)`) and `device(0)`.
|
||
|
|
r"\btorch\.cuda\.device\s*\(",
|
||
|
|
r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b",
|
||
|
|
r"\bwith\storch\.cuda\.device\b",
|
||
|
|
# Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally
|
||
|
|
r"\bcuda_device_count_stateless\(\)\b",
|
||
|
|
]
|
||
|
|
|
||
|
|
# Platform adapters must call native torch.cuda / torch.npu / torch.xpu APIs.
|
||
|
|
ALLOWED_PREFIXES = {
|
||
|
|
"vllm_omni/platforms/",
|
||
|
|
}
|
||
|
|
|
||
|
|
# Grandfathered call sites. Remove a path after migrating it to
|
||
|
|
# current_omni_platform / OmniPlatform. Do not add new files without review.
|
||
|
|
ALLOWED_FILES = {
|
||
|
|
"benchmarks/kernels/mot_linear_benchmarks.py",
|
||
|
|
"examples/online_serving/qwen2_5_omni/gradio_demo.py",
|
||
|
|
"examples/online_serving/qwen3_omni/gradio_demo.py",
|
||
|
|
"tests/dfx/reliability/helpers.py",
|
||
|
|
"tests/diffusion/distributed/test_cfg_parallel.py",
|
||
|
|
"tests/diffusion/distributed/test_pipeline_parallel.py",
|
||
|
|
"tests/e2e/accuracy/wan22_i2v/run_wan22_i2v_diffusers_cp.py",
|
||
|
|
"tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py",
|
||
|
|
"tests/e2e/features/fullduplex/test_personaplex_elastic.py",
|
||
|
|
"tools/configure_stage_memory.py",
|
||
|
|
"vllm_omni/diffusion/models/bagel/pipeline_bagel.py",
|
||
|
|
"vllm_omni/diffusion/models/lance/pipeline_lance.py",
|
||
|
|
"vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py",
|
||
|
|
"vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py",
|
||
|
|
"vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py",
|
||
|
|
"vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py",
|
||
|
|
"vllm_omni/experimental/ar_diffusion/runner.py",
|
||
|
|
"vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py",
|
||
|
|
"vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py",
|
||
|
|
"vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py",
|
||
|
|
"vllm_omni/worker/gpu_ar_model_runner.py",
|
||
|
|
}
|
||
|
|
|
||
|
|
_GENERIC_TIP = (
|
||
|
|
"Found torch.cuda API call. Use current_omni_platform / OmniPlatform "
|
||
|
|
"(empty_cache, synchronize, get_device_count, current_device, "
|
||
|
|
"get_device_memory, set_device) instead."
|
||
|
|
)
|
||
|
|
_SEED_TIP = (
|
||
|
|
"Found CUDA manual-seed API call. Use "
|
||
|
|
"torch.Generator(device=current_omni_platform.device_type).manual_seed(...) "
|
||
|
|
"instead."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _normalize_path(path: str) -> str:
|
||
|
|
return path.replace("\\", "/")
|
||
|
|
|
||
|
|
|
||
|
|
def _is_allowed(path: str) -> bool:
|
||
|
|
if any(path.startswith(prefix) for prefix in ALLOWED_PREFIXES):
|
||
|
|
return True
|
||
|
|
return path in ALLOWED_FILES
|
||
|
|
|
||
|
|
|
||
|
|
def scan_file(path: str) -> int:
|
||
|
|
with open(path, encoding="utf-8") as f:
|
||
|
|
content = f.read()
|
||
|
|
for pattern in _TORCH_CUDA_PATTERNS:
|
||
|
|
for match in re.finditer(pattern, content, re.MULTILINE):
|
||
|
|
line_start = content.rfind("\n", 0, match.start()) + 1
|
||
|
|
if "#" in content[line_start : match.start()]:
|
||
|
|
continue
|
||
|
|
line_num = content[: match.start() + 1].count("\n") + 1
|
||
|
|
matched_text = match.group(0)
|
||
|
|
tip = _SEED_TIP if "manual_seed" in matched_text else _GENERIC_TIP
|
||
|
|
print(f"{path}:{line_num}: \033[91merror:\033[0m {tip}")
|
||
|
|
return 1
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
returncode = 0
|
||
|
|
for filename in sys.argv[1:]:
|
||
|
|
path = _normalize_path(filename)
|
||
|
|
if _is_allowed(path):
|
||
|
|
continue
|
||
|
|
returncode |= scan_file(path)
|
||
|
|
return returncode
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|