Wang, Chuanqi 2025-06-17 12:42:36 +00:00 committed by PyTorch MergeBot
parent 1cce73b5f4
commit 11bb1ece50
3 changed files with 14 additions and 13 deletions

View File

@ -21,8 +21,11 @@ def read_triton_pin(device: str = "cuda") -> str:
return f.read().strip()
def read_triton_version() -> str:
with open(REPO_DIR / ".ci" / "docker" / "triton_version.txt") as f:
def read_triton_version(device: str = "cuda") -> str:
triton_version_file = "triton_version.txt"
if device == "xpu":
triton_version_file = "triton_xpu_version.txt"
with open(REPO_DIR / ".ci" / "docker" / triton_version_file) as f:
return f.read().strip()
@ -91,7 +94,7 @@ def build_triton(
patch_init_py(
triton_pythondir / "triton" / "__init__.py",
version=f"{version}",
expected_version=None,
expected_version=read_triton_version(device),
)
if device == "rocm":
@ -137,15 +140,19 @@ def main() -> None:
parser.add_argument("--py-version", type=str)
parser.add_argument("--commit-hash", type=str)
parser.add_argument("--with-clang-ldd", action="store_true")
parser.add_argument("--triton-version", type=str, default=read_triton_version())
parser.add_argument("--triton-version", type=str, default=None)
args = parser.parse_args()
triton_version = read_triton_version(args.device)
if args.triton_version:
triton_version = args.triton_version
build_triton(
device=args.device,
commit_hash=(
args.commit_hash if args.commit_hash else read_triton_pin(args.device)
),
version=args.triton_version,
version=triton_version,
py_version=args.py_version,
release=args.release,
with_clang_ldd=args.with_clang_ldd,

View File

@ -15,4 +15,4 @@ call conda run -n %PYTHON_PREFIX% pip install wheel pybind11 certifi cython cmak
dir "%VC_INSTALL_PATH%"
call "%VC_INSTALL_PATH%\VC\Auxiliary\Build\vcvarsall.bat" x64
call conda run -n %PYTHON_PREFIX% python .github/scripts/build_triton_wheel.py --device=%BUILD_DEVICE% %TRITON_VERSION% %RELEASE%
call conda run -n %PYTHON_PREFIX% python .github/scripts/build_triton_wheel.py --device=%BUILD_DEVICE% %RELEASE%

View File

@ -156,9 +156,8 @@ jobs:
fi
if [[ "${BUILD_DEVICE}" == xpu ]]; then
TRITON_VERSION=$(cat .ci/docker/triton_xpu_version.txt)
docker exec -t "${container_name}" bash -c "dnf install -y gcc-toolset-13-gcc-c++"
docker exec -t "${container_name}" bash -c "source /opt/rh/gcc-toolset-13/enable && ${PYTHON_EXECUTABLE} /pytorch/.github/scripts/build_triton_wheel.py --device=$BUILD_DEVICE --triton-version=$TRITON_VERSION $RELEASE"
docker exec -t "${container_name}" bash -c "source /opt/rh/gcc-toolset-13/enable && ${PYTHON_EXECUTABLE} /pytorch/.github/scripts/build_triton_wheel.py --device=$BUILD_DEVICE $RELEASE"
else
docker exec -t "${container_name}" bash -c "${PYTHON_EXECUTABLE} /pytorch/.github/scripts/build_triton_wheel.py --device=$BUILD_DEVICE $RELEASE $WITH_CLANG_LDD"
fi
@ -255,11 +254,6 @@ jobs:
if [[ "${IS_RELEASE_TAG}" == true ]]; then
export RELEASE="--release"
fi
export TRITON_VERSION=""
if [[ "${{ matrix.device }}" == "xpu" ]]; then
triton_version=$(cat .ci/docker/ci_commit_pins/triton-xpu.txt)
export TRITON_VERSION="--triton-version=${triton_version}"
fi
.github/scripts/windows/build_triton.bat
mkdir -p "${RUNNER_TEMP}/artifacts/"
mv ./*.whl "${RUNNER_TEMP}/artifacts/"