bugfix: fix unittests not yield error in AOT mode (#657)

The unittests in AOT mode failed since
#629 because we didn't
use return instead of yield in warmup functions, this PR fixes the
issue.

Author: yzh119

tests/test_alibi.py

@@ -26,7 +26,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(

tests/test_batch_decode_kernels.py

@@ -24,7 +24,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(

tests/test_batch_prefill_kernels.py

@@ -24,7 +24,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_prefill_attention_func_args(

tests/test_block_sparse.py

@@ -26,7 +26,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(

tests/test_logits_cap.py

@@ -26,7 +26,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(

tests/test_non_contiguous_decode.py

@@ -8,7 +8,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(

tests/test_non_contiguous_prefill.py

@@ -24,7 +24,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_prefill_attention_func_args(

tests/test_shared_prefix_kernels.py

@@ -24,7 +24,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(

tests/test_sliding_window.py

@@ -24,7 +24,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(

tests/test_tensor_cores_decode.py

@@ -24,7 +24,7 @@
 @pytest.fixture(autouse=True, scope="module")
 def warmup_jit():
     if flashinfer.jit.has_prebuilt_ops:
-        return
+        yield
     try:
         flashinfer.jit.parallel_load_modules(
             jit_decode_attention_func_args(