Skip to content

Commit a5ac343

Browse files
xuzhao9facebook-github-bot
authored andcommitted
Fix loading ThunderKitten libraries (#22)
Summary: As mentioned in #17 Pull Request resolved: #22 Test Plan: ``` (base) ➜ tritonbench git:(xz9/fix-tk-load) python -c "from tritonbench.utils.loader import load_library; load_library('tk/tk_attn_h100_fwd.so')" (base) ➜ tritonbench git:(xz9/fix-tk-load) echo $? 0 ``` Reviewed By: FindHao Differential Revision: D65147532 Pulled By: xuzhao9 fbshipit-source-id: 669b2aa4db9b0581a2dbf709e9b577dbbe5e670e
1 parent 3f337fa commit a5ac343

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

Diff for: tritonbench/operators/flash_attention/operator.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@
9999
# colfax Flash Attention V2 for Hopper
100100
torch.ops.load_library("//ai_codesign/gen_ai/cutlass-kernels:fmha_forward_lib")
101101
else:
102-
from userbenchmark.triton.loader import load_library
102+
from tritonbench.utils.loader import load_library
103103

104104
load_library("cutlass_kernels/fmha_forward_lib.so")
105105
colfax_cutlass_fmha = torch.ops.cutlass.fmha_forward

Diff for: tritonbench/utils/loader.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,5 @@ def load_library(library_path: str):
55
import torch
66

77
prefix, _delimiter, so_file = library_path.partition("/")
8-
so_full_path = REPO_PATH.joinpath(prefix, ".data", so_file).resolve()
8+
so_full_path = REPO_PATH.joinpath("utils", prefix, ".data", so_file).resolve()
99
torch.ops.load_library(str(so_full_path))

0 commit comments

Comments
 (0)