From 425040c7af134a3e034b25de6f73e6c38cd9f639 Mon Sep 17 00:00:00 2001
From: Zihao Ye <expye@outlook.com>
Date: Thu, 17 Oct 2024 17:19:42 -0700
Subject: [PATCH] bugfix: fix JIT compilation of prefill kernels (#536)

Some bugs were introduced in #534, this PR fix these issues.
---
 python/flashinfer/jit/attention.py | 3 ++-
 python/flashinfer/prefill.py       | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/python/flashinfer/jit/attention.py b/python/flashinfer/jit/attention.py
index 29e41da7..42781064 100644
--- 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
index bfc022b5..293073cb 100644
--- 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,