diff --git a/python/flashinfer/page.py b/python/flashinfer/page.py
index 8efe2168..581d6207 100644
--- a/python/flashinfer/page.py
+++ b/python/flashinfer/page.py
@@ -14,7 +14,7 @@
 limitations under the License.
 """
 
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import torch
 import triton