Fix hstu

pytorch-labs · Oct 29, 2024 · 4eb6547 · 4eb6547
1 parent a6a8aed
commit 4eb6547
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 8 deletions.
diff --git a/.github/workflows/pr.yaml b/.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

diff --git a/tritonbench/operators/ragged_attention/hstu.py b/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,
     )

diff --git a/tritonbench/operators/sum/operator.py b/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