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

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