misc: add page ops to shared-prefix kernel unittest warmup function (#630)

yzh119 · web-flow · commit 6ec4db344cbc · 2024-11-23T19:13:01.000-08:00
In #629 , page ops are missing in shared-prefix kernel unittest's warmup function
diff --git a/tests/test_shared_prefix_kernels.py b/tests/test_shared_prefix_kernels.py
@@ -44,6 +44,7 @@ def warmup_jit():
                 [False],  # use_logits_soft_caps
                 [False],  # allow_fp16_qk_reductions
             )
+            + [flashinfer.page.get_page_module, []]
         )
     except Exception as e:
         # abort the test session if warmup fails

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,7 @@ def warmup_jit():`
`44`	`44`	`[False], # use_logits_soft_caps`
`45`	`45`	`[False], # allow_fp16_qk_reductions`
`46`	`46`	`)`
	`47`	`+ + [flashinfer.page.get_page_module, []]`
`47`	`48`	`)`
`48`	`49`	`except Exception as e:`
`49`	`50`	`# abort the test session if warmup fails`