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.
flashinfer-ai · Dec 13, 2024 · 6dfc9d8 · 6dfc9d8
1 parent 4c15777
commit 6dfc9d8
Show file tree

Hide file tree

Showing 10 changed files with 10 additions and 10 deletions.
diff --git a/tests/test_alibi.py b/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(

diff --git a/tests/test_batch_decode_kernels.py b/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(

diff --git a/tests/test_batch_prefill_kernels.py b/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(

diff --git a/tests/test_block_sparse.py b/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(

diff --git a/tests/test_logits_cap.py b/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(

diff --git a/tests/test_non_contiguous_decode.py b/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(

diff --git a/tests/test_non_contiguous_prefill.py b/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(

diff --git a/tests/test_shared_prefix_kernels.py b/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(

diff --git a/tests/test_sliding_window.py b/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(

diff --git a/tests/test_tensor_cores_decode.py b/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(