kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][chore] cleanup build script (#7865)

Browse Source 
Signed-off-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com>
This commit is contained in:
Yuan Tong 2025-09-24 21:15:01 +08:00 committed by GitHub
parent 60101eb8a5
commit 51bef1beb0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 34 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
scripts/build_wheel.py
View File

@ -1,5 +1,5 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
@ -17,6 +17,7 @@
import os import os
import platform import platform
import re import re
import shutil
import sys import sys
import sysconfig import sysconfig
import tempfile import tempfile
@ -154,7 +155,7 @@ def setup_venv(project_dir: Path, requirements_file: Path,
warnings.warn( warnings.warn(
f"Using the NVIDIA PyTorch container with PyPI distributed PyTorch may lead to compatibility issues.\n" f"Using the NVIDIA PyTorch container with PyPI distributed PyTorch may lead to compatibility issues.\n"
f"If you encounter any problems, please delete the environment at `{venv_prefix}` so that " f"If you encounter any problems, please delete the environment at `{venv_prefix}` so that "
f"`build_wheel.py` can recreate the virtual environment correctly." f"`build_wheel.py` can recreate a virtual environment using container-provided PyTorch installation."
) )
print("^^^^^^^^^^ IMPORTANT WARNING ^^^^^^^^^^", file=sys.stderr) print("^^^^^^^^^^ IMPORTANT WARNING ^^^^^^^^^^", file=sys.stderr)
input("Press Ctrl+C to stop, any key to continue...\n") input("Press Ctrl+C to stop, any key to continue...\n")
@ -187,7 +188,7 @@ def setup_venv(project_dir: Path, requirements_file: Path,
f"Incompatible NVIDIA PyTorch container detected. " f"Incompatible NVIDIA PyTorch container detected. "
f"The container provides PyTorch version {version_installed}, " f"The container provides PyTorch version {version_installed}, "
f"but current revision requires {version_required}. " f"but current revision requires {version_required}. "
f"Please recreate your container using image specified in .devcontainer/docker-compose.yml. " f"Please recreate your container using image specified in jenkins/current_image_tags.properties. "
f"NOTE: Please don't try install PyTorch using pip. " f"NOTE: Please don't try install PyTorch using pip. "
f"Using the NVIDIA PyTorch container with PyPI distributed PyTorch may lead to compatibility issues." f"Using the NVIDIA PyTorch container with PyPI distributed PyTorch may lead to compatibility issues."
) )
@ -253,7 +254,6 @@ def generate_fmha_cu(project_dir, venv_python):
fmha_v2_cu_dir.mkdir(parents=True, exist_ok=True) fmha_v2_cu_dir.mkdir(parents=True, exist_ok=True)
fmha_v2_dir = project_dir / "cpp/kernels/fmha_v2" fmha_v2_dir = project_dir / "cpp/kernels/fmha_v2"
os.chdir(fmha_v2_dir)
env = os.environ.copy() env = os.environ.copy()
env.update({ env.update({
@ -267,19 +267,41 @@ def generate_fmha_cu(project_dir, venv_python):
"GENERATE_CU_TRTLLM": "true" "GENERATE_CU_TRTLLM": "true"
}) })
build_run("rm -rf generated") shutil.rmtree(fmha_v2_dir / "generated", ignore_errors=True)
build_run("rm -rf temp") shutil.rmtree(fmha_v2_dir / "temp", ignore_errors=True)
build_run("rm -rf obj") shutil.rmtree(fmha_v2_dir / "obj", ignore_errors=True)
build_run("python3 setup.py", env=env) build_run("python3 setup.py", env=env, cwd=fmha_v2_dir)
# Only touches generated source files if content is updated
def move_if_updated(src, dst):
with open(src, "rb") as f:
new_content = f.read()
try:
with open(dst, "rb") as f:
old_content = f.read()
except FileNotFoundError:
old_content = None
if old_content != new_content:
shutil.move(src, dst)
# Copy generated header file when cu path is active and cubins are deleted. # Copy generated header file when cu path is active and cubins are deleted.
cubin_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin" cubin_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin"
build_run(f"mv generated/fmha_cubin.h {cubin_dir}") move_if_updated(fmha_v2_dir / "generated/fmha_cubin.h",
cubin_dir / "fmha_cubin.h")
generated_files = set()
for cu_file in (fmha_v2_dir / "generated").glob("*sm*.cu"): for cu_file in (fmha_v2_dir / "generated").glob("*sm*.cu"):
build_run(f"mv {cu_file} {fmha_v2_cu_dir}") dst_file = fmha_v2_cu_dir / os.path.basename(cu_file)
move_if_updated(cu_file, dst_file)
generated_files.add(str(dst_file.resolve()))
os.chdir(project_dir) # Remove extra files
for root, _, files in os.walk(fmha_v2_cu_dir):
for file in files:
file_path = os.path.realpath(os.path.join(root, file))
if file_path not in generated_files:
os.remove(file_path)
def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]) -> str: def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]) -> str:
@ -606,10 +628,7 @@ def main(*,
source_dir = get_source_dir() source_dir = get_source_dir()
fmha_v2_cu_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu" fmha_v2_cu_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu"
if clean or generate_fmha: if clean or generate_fmha or not fmha_v2_cu_dir.exists():
build_run(f"rm -rf {fmha_v2_cu_dir}")
generate_fmha_cu(project_dir, venv_python)
elif not fmha_v2_cu_dir.exists():
generate_fmha_cu(project_dir, venv_python) generate_fmha_cu(project_dir, venv_python)
with working_directory(build_dir): with working_directory(build_dir):