misc: remove head dimension 64 from AOT (flashinfer-ai#782)

This PR removes `head_dim=64` from pre-built AOT wheels, `head_dim=64`
will be supported using JIT in the future.

Author: yzh119

setup.py

@@ -27,10 +27,12 @@
 root = Path(__file__).parent.resolve()
 gen_dir = root / "csrc" / "generated"
 
-head_dims = os.environ.get("FLASHINFER_HEAD_DIMS", "64,128,256").split(",")
+head_dims = os.environ.get("FLASHINFER_HEAD_DIMS", "128,256").split(",")
 head_dims = list(map(int, head_dims))
-SM90_ALLOWED_HEAD_DIMS = {(64, 64), (128, 128), (256, 256), (192, 128)}
-head_dims_sm90 = list(SM90_ALLOWED_HEAD_DIMS)  # No support for custom head dims for SM90
+SM90_ALLOWED_HEAD_DIMS = {(128, 128), (256, 256), (192, 128)}
+head_dims_sm90 = list(
+    SM90_ALLOWED_HEAD_DIMS
+)  # No support for custom head dims for SM90
 
 mask_modes = [0, 1, 2]