Update pytorch-triton version (pytorch#2775)

Summary:
- Update pytorch-triton version to follow PyTorch pytorch/pytorch#126098

Pull Request resolved: pytorch#2775

Reviewed By: brad-mengchi

Differential Revision: D58959021

Pulled By: q10

fbshipit-source-id: 575d42dc07020778ef9f201ea1a234fb92be27f2

.github/scripts/utils_triton.bash

@@ -15,7 +15,7 @@
 # Triton Setup Functions
 ################################################################################
 
-install_triton_gitmodule () {
+install_triton_git_repo () {
   local env_name="$1"
   local triton_version="$2"
   if [ "$env_name" == "" ]; then
@@ -26,7 +26,7 @@ install_triton_gitmodule () {
     return 1
   else
     echo "################################################################################"
-    echo "# Build + Install Triton (gitmodule)"
+    echo "# Build + Install Triton (git repo)"
     echo "#"
     echo "# [$(date --utc +%FT%T.%3NZ)] + ${FUNCNAME[0]} ${*}"
     echo "################################################################################"
@@ -39,20 +39,23 @@ install_triton_gitmodule () {
   local env_prefix=$(env_name_or_prefix "${env_name}")
 
   echo "[BUILD] Checking out triton ..."
-  cd ../third_party/triton/python || return 1
+  cd ~/                                               || return 1
+  git clone https://github.com/triton-lang/triton.git || return 1
+  cd triton/python                                    || return 1
+
   if [ "$triton_version" != "" ]; then
-    (print_exec git checkout "${triton_version}") || return 1
+    (print_exec git reset --hard "${triton_version}") || return 1
   fi
 
-  echo "[BUILD] Installing Triton from gitmodule ..."
+  echo "[BUILD] Installing Triton from git repo ..."
   # shellcheck disable=SC2086
   (exec_with_retries 3 conda run --no-capture-output ${env_prefix} python -m pip install -e .) || return 1
 
   # shellcheck disable=SC2086
   (test_python_import_package "${env_name}" triton) || return 1
 
   cd - || return 1
-  echo "[INSTALL] Successfully installed Triton ${triton_version} from gitmodule"
+  echo "[INSTALL] Successfully installed Triton ${triton_version} from git repo"
 }
 
 install_triton_pip () {
@@ -72,8 +75,21 @@ install_triton_pip () {
   fi
 
   echo "[BUILD] Installing Triton from PIP ..."
+  # NOTE: Install pytorch-triton from nightly with commit SHA that follows the
+  # version tracked in FB internal (/third-party/tp2/triton/2.0.x/METADATA.bzl)
+  # as close as possible.
+  #
+  # https://download.pytorch.org/whl/nightly/pytorch-triton/
+  # https://download.pytorch.org/whl/nightly/pytorch-triton-rocm/
+  #
+  # This needs to be manually updated to follow PyTorch's triton pinning
+  # updates:
+  #
+  # https://github.com/pytorch/pytorch/blob/main/.ci/docker/ci_commit_pins/triton.txt
+  # https://github.com/pytorch/pytorch/pull/126098
+  #
   # shellcheck disable=SC2086
-  install_from_pytorch_pip "${env_name}" pytorch-triton nightly/3.0.0+45fff310c8 || return 1
+  install_from_pytorch_pip "${env_name}" pytorch-triton nightly/3.0.0+dedb7bdf33 || return 1
 
   # shellcheck disable=SC2086
   (test_python_import_package "${env_name}" triton) || return 1

.gitmodules

@@ -13,6 +13,3 @@
 [submodule "third_party/cutlass"]
 	path = third_party/cutlass
 	url = https://github.com/NVIDIA/cutlass.git
-[submodule "third_party/triton"]
-	path = third_party/triton
-	url = https://github.com/openai/triton.git

fbgemm_gpu/docs/src/fbgemm_gpu-development/BuildInstructions.rst

@@ -423,6 +423,29 @@ For the CUDA variant of PyTorch, verify that at the minimum ``cuda_cmake_macros.
   conda_prefix=$(conda run -n ${env_name} printenv CONDA_PREFIX)
   find "${conda_prefix}" -name cuda_cmake_macros.h
 
+Install PyTorch-Triton
+~~~~~~~~~~~~~~~~~~~~~~
+
+This section is only applicable to building the experimental FBGEMM_GPU
+Triton-GEMM module.  Triton should be installed via the ``pytorch-triton``,
+which generally comes installing ``torch``, but can also be installed manually:
+
+.. code:: sh
+
+  # pytorch-triton repos:
+  # https://download.pytorch.org/whl/nightly/pytorch-triton/
+  # https://download.pytorch.org/whl/nightly/pytorch-triton-rocm/
+
+  # The version SHA should follow the one pinned in PyTorch
+  # https://github.com/pytorch/pytorch/blob/main/.ci/docker/ci_commit_pins/triton.txt
+  conda run -n ${env_name} pip install --pre pytorch-triton==3.0.0+dedb7bdf33 --index-url https://download.pytorch.org/whl/nightly/
+
+Verify the PyTorch-Triton installation with an ``import`` test:
+
+.. code:: sh
+
+  # Ensure that the package loads properly
+  conda run -n ${env_name} python -c "import triton"
 
 Build the FBGEMM_GPU Package
 ----------------------------