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,