[infra] Fix docker workflow (#343)

Author: xuzhao9

.github/workflows/docker-rocm.yaml

@@ -28,7 +28,7 @@ jobs:
         uses: docker/login-action@v2
         with:
           registry: ghcr.io
-          username: pytorch-labs
+          username: meta-pytorch
           password: ${{ secrets.TRITONBENCH_ACCESS_TOKEN }}
       - name: Build TritonBench nightly docker
         run: |
@@ -39,19 +39,19 @@ 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-rocm-nightly.dockerfile -t ghcr.io/pytorch-labs/tritonbench:rocm-latest
+              -f tritonbench-rocm-nightly.dockerfile -t ghcr.io/meta-pytorch/tritonbench:rocm-latest
           # Extract pytorch version from the docker
-          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/pytorch-labs/tritonbench:rocm-latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
+          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/meta-pytorch/tritonbench:rocm-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/pytorch-labs/tritonbench:rocm-latest ghcr.io/pytorch-labs/tritonbench:rocm-${DOCKER_TAG}
+          docker tag ghcr.io/meta-pytorch/tritonbench:rocm-latest ghcr.io/meta-pytorch/tritonbench:rocm-${DOCKER_TAG}
       - 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/pytorch-labs/tritonbench:rocm-latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
+          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/meta-pytorch/tritonbench:rocm-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 push ghcr.io/pytorch-labs/tritonbench:rocm-${DOCKER_TAG}
-          docker push ghcr.io/pytorch-labs/tritonbench:rocm-latest
+          docker push ghcr.io/meta-pytorch/tritonbench:rocm-${DOCKER_TAG}
+          docker push ghcr.io/meta-pytorch/tritonbench:rocm-latest
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true

.github/workflows/docker.yaml

@@ -20,7 +20,7 @@ on:
         required: false
 env:
   CONDA_ENV: "tritonbench"
-  DOCKER_IMAGE: "ghcr.io/pytorch-labs/tritonbench:latest"
+  DOCKER_IMAGE: "ghcr.io/meta-pytorch/tritonbench:latest"
   SETUP_SCRIPT: "/workspace/setup_instance.sh"
 
 jobs:
@@ -38,7 +38,7 @@ jobs:
         uses: docker/login-action@v2
         with:
           registry: ghcr.io
-          username: pytorch-labs
+          username: meta-pytorch
           password: ${{ secrets.TRITONBENCH_ACCESS_TOKEN }}
       - name: Build TritonBench nightly docker
         run: |
@@ -49,19 +49,19 @@ 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/pytorch-labs/tritonbench:latest
+              -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/pytorch-labs/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
+          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/pytorch-labs/tritonbench:latest ghcr.io/pytorch-labs/tritonbench:${DOCKER_TAG}
+          docker tag ghcr.io/meta-pytorch/tritonbench:latest ghcr.io/meta-pytorch/tritonbench:${DOCKER_TAG}
       - 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/pytorch-labs/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
+          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 push ghcr.io/pytorch-labs/tritonbench:${DOCKER_TAG}
-          docker push ghcr.io/pytorch-labs/tritonbench:latest
+          docker push ghcr.io/meta-pytorch/tritonbench:${DOCKER_TAG}
+          docker push ghcr.io/meta-pytorch/tritonbench:latest
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true

docker/infra/README.md

@@ -6,6 +6,9 @@ The Infra is a Kubernetes cluster built on top of Google Cloud Platform.
 ## Step 1: Create the cluster and install the ARC Controller
 
 ```
+# login ghcr.io so that remote can pull the image
+docker login ghcr.io
+
 # Get credentials for the cluster so that kubectl could use it
 gcloud container clusters get-credentials --location us-central1 tritonbench-h100-cluster
 

docker/infra/values.yaml

@@ -223,7 +223,7 @@ template:
     containers:
     - name: runner
       # image: ghcr.io/actions/actions-runner:latest
-      image: ghcr.io/pytorch-labs/tritonbench:latest
+      image: ghcr.io/meta-pytorch/tritonbench:latest
       command: ["sh", "-c", "sudo cp -r /usr/bin/nvidia/* /usr/bin; sudo cp -r /usr/lib/x86_64-linux-gnu/nvidia/* /usr/lib/x86_64-linux-gnu; bash /home/runner/run.sh"]
       securityContext:
         privileged: true

run.py

@@ -15,6 +15,7 @@
 
 from tritonbench.operators import load_opbench_by_name
 from tritonbench.operators_collection import list_operators_by_collection
+from tritonbench.utils.ab_test import compare_ab_results, run_ab_test
 from tritonbench.utils.env_utils import is_fbcode
 from tritonbench.utils.gpu_utils import gpu_lockdown
 from tritonbench.utils.list_operator_details import list_operator_details
@@ -23,7 +24,6 @@
 
 from tritonbench.utils.triton_op import BenchmarkOperatorResult
 from tritonbench.utils.tritonparse_utils import tritonparse_init, tritonparse_parse
-from tritonbench.utils.ab_test import run_ab_test, compare_ab_results
 
 try:
     if is_fbcode():
@@ -34,8 +34,6 @@
     usage_report_logger = lambda *args, **kwargs: None
 
 
-
-
 def _run(args: argparse.Namespace, extra_args: List[str]) -> BenchmarkOperatorResult:
     if is_loader_op(args.op):
         Opbench = get_op_loader_bench_cls_by_name(args.op)
@@ -132,23 +130,26 @@ def run(args: List[str] = []):
     # Check if A/B testing mode is enabled
     if args.side_a is not None and args.side_b is not None:
         # A/B testing mode - only support single operator
-        assert len(ops) == 1, "A/B testing validation should have caught multiple operators"
+        assert (
+            len(ops) == 1
+        ), "A/B testing validation should have caught multiple operators"
         op = ops[0]
         args.op = op
-        
+
         print("[A/B Testing Mode Enabled]")
         print(f"Operator: {op}")
         print()
-        
+
         with gpu_lockdown(args.gpu_lockdown):
             try:
                 result_a, result_b = run_ab_test(args, extra_args, _run)
-                
+
                 from tritonbench.utils.ab_test import parse_ab_config
+
                 config_a_args = parse_ab_config(args.side_a)
                 config_b_args = parse_ab_config(args.side_b)
                 compare_ab_results(result_a, result_b, config_a_args, config_b_args)
-                
+
             except Exception as e:
                 print(f"A/B test failed: {e}")
                 if not args.bypass_fail:
@@ -166,7 +167,7 @@ def run(args: List[str] = []):
                     run_in_task(op)
                 else:
                     _run(args, extra_args)
-                    
+
     tritonparse_parse(args.tritonparse)
 
 

tools/python_utils.py

@@ -8,12 +8,12 @@
 DEFAULT_PYTHON_VERSION = "3.12"
 
 PYTHON_VERSION_MAP = {
-   "3.11": {
+    "3.11": {
         "pytorch_url": "cp311",
-   },
-   "3.12": {
+    },
+    "3.12": {
         "pytorch_url": "cp312",
-   },
+    },
 }
 REPO_DIR = Path(__file__).parent.parent
 

tritonbench/operators/addmm/operator.py

@@ -11,7 +11,12 @@
 except ModuleNotFoundError:
     from .hstu import triton_addmm
 
-from tritonbench.operators.gemm.stream_k import streamk_matmul
+try:
+    from tritonbench.operators.gemm.stream_k import streamk_matmul
+except ImportError:
+    streamk_matmul = None
+
+from tritonbench.operators.gemm import stream_k
 from tritonbench.utils.triton_op import (
     BenchmarkOperator,
     BenchmarkOperatorMetrics,
@@ -94,7 +99,7 @@ def __init__(
     def triton_addmm(self, a, mat1, mat2) -> Callable:
         return lambda: triton_addmm(a, mat1, mat2)
 
-    @register_benchmark()
+    @register_benchmark(enabled=bool(streamk_matmul))
     def streamk_addmm(self, a, mat1, mat2) -> Callable:
         return lambda: streamk_matmul(mat1, mat2, bias=a)
 

tritonbench/operators/fp8_gemm/fp8_gemm.py

@@ -65,13 +65,17 @@ def args(m, n, k):
             )
             return (a, b)
 
-        if hasattr(self, 'external_shapes') and self.external_shapes:  # Check for external shapes loaded from input-loader
+        if (
+            hasattr(self, "external_shapes") and self.external_shapes
+        ):  # Check for external shapes loaded from input-loader
             for shape in self.external_shapes:
                 if len(shape) == 3:
                     m, n, k = shape
                     yield args(m, n, k)
                 else:
-                    logger.warning(f"Skipping invalid shape: {shape}, expected [M, N, K]")
+                    logger.warning(
+                        f"Skipping invalid shape: {shape}, expected [M, N, K]"
+                    )
         elif self.extra_args.llama:
             for m, n, k, _bias in llama_shapes():
                 yield args(m, n, k)

tritonbench/utils/triton_utils.py

@@ -1,9 +1,10 @@
 # utils to identify triton versions
 
-import triton.language as tl
 import functools
 import importlib.util
 
+import triton.language as tl
+
 
 class AsyncTaskContext:
     """Context manager that dispatches to tl.async_task if available, otherwise no-op."""