bugfix: fix JIT compilation of prefill kernels (#536)

Some bugs were introduced in #534, this PR fix these issues.
flashinfer-ai · Oct 18, 2024 · 425040c · 425040c
1 parent de25d76
commit 425040c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/python/flashinfer/jit/attention.py b/python/flashinfer/jit/attention.py
@@ -269,8 +269,9 @@ def gen_batch_prefill_cu(*args) -> Tuple[str, pathlib.Path]:
         os.makedirs(gen_directory)
     uri = get_batch_prefill_uri(*args)
     file_name = f"{uri}.cu"
-    path = gen_directory / file_name,
+    path = gen_directory / file_name
     write_if_different(
+        path,
         get_batch_prefill_cu_str(*args),
     )
     return uri, path

diff --git a/python/flashinfer/prefill.py b/python/flashinfer/prefill.py
@@ -55,7 +55,7 @@ def compile_single_prefill_module(
     *args,
     verbose: bool = False,
 ):
-    uri, path = get_single_prefill_uri(*args)
+    uri, path = gen_single_prefill_cu(*args)
     return load_cuda_ops(
         uri, [path],
         verbose=verbose,
@@ -66,7 +66,7 @@ def compile_batch_prefill_module(
     *args,
     verbose: bool = False,
 ):
-    uri, path = get_batch_prefill_uri(*args)
+    uri, path = gen_batch_prefill_cu(*args)
     return load_cuda_ops(
         uri, [path],
         verbose=verbose,