Fix hstu

Author: xuzhao9

.github/workflows/pr.yaml

@@ -19,6 +19,9 @@ jobs:
     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
@@ -27,9 +30,3 @@ jobs:
       - name: Test Tritonbench operators
         run: |
           bash ./.ci/tritonbench/test-operators.sh
-      - name: Clean up Conda env
-        if: always()
-        run: |
-          . "${SETUP_SCRIPT}"
-          conda deactivate && conda deactivate
-          conda remove -n "${CONDA_ENV}" --all

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")