Skip to content

Commit 6bb969b

Browse files
authored
fix: minor errors in cubin loader (#1295)
<!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. -->
1 parent 927a41e commit 6bb969b

File tree

1 file changed

+5
-10
lines changed

1 file changed

+5
-10
lines changed

flashinfer/decode.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1788,15 +1788,6 @@ def run(
17881788
run_return_lse = functools.partialmethod(run, return_lse=True)
17891789

17901790

1791-
# todo(Yingyi): update the params list
1792-
@functools.cache
1793-
def get_trtllm_fmha_gen_module():
1794-
mod = trtllm_fmha_gen_module()
1795-
op = mod.build_and_load()
1796-
setup_cubin_loader(mod.get_library_path())
1797-
return op
1798-
1799-
18001791
class TrtllmGenDecodeModule:
18011792
def _paged_run(
18021793
self,
@@ -1841,9 +1832,13 @@ def _plan(self, *args, **kwargs):
18411832
def __init__(self):
18421833
self._mod = trtllm_fmha_gen_module()
18431834
self._op = self._mod.build_and_load()
1844-
from flashinfer.jit.cubin_loader import setup_cubin_loader
1835+
from flashinfer.jit.cubin_loader import (
1836+
setup_cubin_loader,
1837+
setup_metainfo_loader,
1838+
)
18451839

18461840
setup_cubin_loader(self._mod.get_library_path())
1841+
setup_metainfo_loader(self._mod.get_library_path())
18471842

18481843

18491844
@functools.cache

0 commit comments

Comments
 (0)