[docker] add meta triton env (#525)

Author: xuzhao9

.github/workflows/docker.yaml

@@ -18,6 +18,9 @@ on:
       nightly_date:
         description: "PyTorch nightly version"
         required: false
+      tags:
+        description: "Tags to attach to the docker image, split by comma"
+        required: false
 env:
   CONDA_ENV: "tritonbench"
   DOCKER_IMAGE: "ghcr.io/meta-pytorch/tritonbench:latest"
@@ -49,19 +52,24 @@ jobs:
           # and it is github.ref_name when triggered by workflow_dispatch
           branch_name=${{ github.head_ref || github.ref_name }}
           docker build . --build-arg TRITONBENCH_BRANCH="${branch_name}" --build-arg FORCE_DATE="${NIGHTLY_DATE}" \
-              -f tritonbench-nightly.dockerfile -t ghcr.io/meta-pytorch/tritonbench:latest
-          # Extract pytorch version from the docker
-          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/meta-pytorch/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
-          export DOCKER_TAG=$(awk '{match($0, /dev[0-9]+/, arr); print arr[0]}' <<< "${PYTORCH_VERSION}")
-          docker tag ghcr.io/meta-pytorch/tritonbench:latest ghcr.io/meta-pytorch/tritonbench:${DOCKER_TAG}
+              -f tritonbench-nightly.dockerfile -t ghcr.io/meta-pytorch/tritonbench:latest          
       - name: Push docker to remote
         if: github.event_name != 'pull_request'
         run: |
           # Extract pytorch version from the docker
           PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/meta-pytorch/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
           export DOCKER_TAG=$(awk '{match($0, /dev[0-9]+/, arr); print arr[0]}' <<< "${PYTORCH_VERSION}")
+          docker tag ghcr.io/meta-pytorch/tritonbench:latest ghcr.io/meta-pytorch/tritonbench:${DOCKER_TAG}
           docker push ghcr.io/meta-pytorch/tritonbench:${DOCKER_TAG}
           docker push ghcr.io/meta-pytorch/tritonbench:latest
+          set -x
+          if [[ -n "${{ github.event.inputs.tags }}" ]]; then
+            IFS=',' read -ra tags <<< "${{ github.event.inputs.tags }}"
+            for tag in "${tags[@]}"; do
+                docker tag ghcr.io/meta-pytorch/tritonbench:latest ghcr.io/meta-pytorch/tritonbench:${tag}
+                docker push ghcr.io/meta-pytorch/tritonbench:${tag}
+            done
+          fi
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true

benchmarks/run_config/bf16_blackwell_attentions_fwd_shape4.yaml

@@ -0,0 +1,7 @@
+bf16_flash_attention_fwd_shape4:
+  op:
+    blackwell_attentions
+  args:
+    --op blackwell_attentions --seq-len 8192 --batch 4 --n-heads 32 --d-head 128 --metrics tflops --rep 3000 --sleep 1.0 --only triton_tutorial_flash_dp_persistent_blackwell --simple-output --force
+  envs:
+    TRITON_ALWAYS_COMPILE=1

docker/tritonbench-nightly.dockerfile

@@ -4,6 +4,7 @@ FROM ${BASE_IMAGE}
 ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
 ENV CONDA_ENV=pytorch
 ENV CONDA_ENV_TRITON_MAIN=triton-main
+ENV CONDA_ENV_META_TRITON=meta-triton
 ENV SETUP_SCRIPT=/workspace/setup_instance.sh
 ARG OVERRIDE_GENCODE="-gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_90a,code=sm_90a"
 ARG OVERRIDE_GENCODE_CUDNN="-gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_90a,code=sm_90a"
@@ -105,6 +106,12 @@ RUN cd /workspace/tritonbench && \
         --repo triton-lang/triton --commit main --side single \
         --install-dir /workspace/triton-main
 
+# Build meta-triton conda env
+RUN cd /workspace/tritonbench && \
+    bash .ci/triton/install.sh --conda-env "${CONDA_ENV_META_TRITON}" \
+        --repo facebookexperimental/triton --commit ws-3.5 --side single \
+        --install-dir /workspace/meta-triton
+
 # Output setup script for inspection
 RUN cat "${SETUP_SCRIPT}"