Fix hstu

Author: xuzhao9

.github/workflows/pr.yaml

@@ -19,11 +19,9 @@ jobs:
     steps:
       - name: Checkout Tritonbench
         uses: actions/checkout@v3
-      - name: Tune Nvidia GPU
-        run: |
-          sudo nvidia-smi -pm 1
-          sudo ldconfig
-          nvidia-smi
+        with:
+          # no need to checkout submodules recursively
+          submodules: true
       - 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