Add pr test to run operators on h100

.github/workflows/docker.yaml

@@ -21,6 +21,7 @@ jobs:
   build-push-docker:
     if: ${{ github.repository_owner == 'pytorch-labs' }}
     runs-on: 32-core-ubuntu
+    environment: docker-s3-upload
     steps:
       - name: Checkout
         uses: actions/checkout@v3

.github/workflows/pr.yaml

@@ -0,0 +1,32 @@
+name: TritonBench PR Test
+on:
+  pull_request:
+    paths:
+      - .ci/*
+      - tritonbench/*
+      - .github/workflows/pr.yaml
+
+jobs:
+  h100-pytorch-test:
+    # Don't run on forked repos
+    if: github.repository_owner == 'pytorch-labs'
+    runs-on: [gcp-h100-runner]
+    timeout-minutes: 240
+    environment: docker-s3-upload
+    env:
+      CONDA_ENV: "pytorch"
+      SETUP_SCRIPT: "/workspace/setup_instance.sh"
+    steps:
+      - name: Checkout Tritonbench
+        uses: actions/checkout@v3
+        with:
+          # no need to checkout submodules recursively
+          submodules: true
+      - name: Tune Nvidia GPU
+        run: |
+          sudo nvidia-smi -pm 1
+          sudo ldconfig
+          nvidia-smi
+      - name: Test Tritonbench operators
+        run: |
+          bash ./.ci/tritonbench/test-operators.sh

tritonbench/operators/ragged_attention/hstu.py

@@ -4,7 +4,7 @@
 
 try:
     # Internal Import
-    from hammer.generative_recommenders.ops.triton.triton_ragged_hstu_attention import (
+    from hammer.oss.generative_recommenders.ops.triton.triton_ragged_hstu_attention import (
         _ragged_hstu_attn_fwd,
         _ragged_hstu_attn_fwd_persistent,
     )

tritonbench/operators/sum/operator.py

@@ -4,8 +4,6 @@
 import os
 from typing import Callable, Generator, List, Optional, Tuple
 
-import matplotlib.pyplot as plt
-
 import torch
 import triton
 import triton.language as tl

tritonbench/utils/path_utils.py

@@ -3,7 +3,7 @@
 
 from pathlib import Path
 
-REPO_PATH = Path(os.path.abspath(__file__)).parent.parent
+REPO_PATH = Path(os.path.abspath(__file__)).parent.parent.parent
 SUBMODULE_PATH = REPO_PATH.joinpath("submodules")