diff --git a/python/flashinfer/page.py b/python/flashinfer/page.py
index d7944a782..876e98229 100644
--- a/python/flashinfer/page.py
+++ b/python/flashinfer/page.py
@@ -27,6 +27,8 @@
     else:
         raise e
 
+from .utils import check_kv_layout, TensorLayout
+
 
 def append_paged_kv_cache(
     append_key: torch.Tensor,