search

pytorch-labs / tritonbench

Tritonbench is a collection of PyTorch custom operators with example inputs to measure their performance.
BSD 3-Clause "New" or "Revised" License
20 stars 3 forks source link

Fix loading ThunderKitten libraries #22

Closed xuzhao9 closed 3 weeks ago

xuzhao9 commented 3 weeks ago

As mentioned in https://github.com/pytorch-labs/tritonbench/issues/17

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
facebook-github-bot commented 3 weeks ago

@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

facebook-github-bot commented 3 weeks ago

@xuzhao9 merged this pull request in pytorch-labs/tritonbench@a5ac343e40a074f01dbb1913a01f924e19dc2e9b.