misc: allow head_dim=64 for sm90 AOT (#783)

Keeps the default behavior of #782, i.e., build AOT without
`head_dim=64`. But gives an option to specifically enable `head_dim=64`.

`head_dim=64` is used by some small models like
[Qwen2.5-0.5B](https://huggingface.co/Qwen/Qwen2.5-0.5B/blob/main/config.json).

Author: abcdabcd987

setup.py

@@ -29,10 +29,11 @@
 
 head_dims = os.environ.get("FLASHINFER_HEAD_DIMS", "128,256").split(",")
 head_dims = list(map(int, head_dims))
-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
+SM90_ALLOWED_HEAD_DIMS = {(64, 64), (128, 128), (256, 256), (192, 128)}
+head_dims_sm90 = [(d, d) for d in head_dims if (d, d) in SM90_ALLOWED_HEAD_DIMS]
+head_dims_sm90.extend(
+    [(k, v) for k, v in SM90_ALLOWED_HEAD_DIMS if k != v]
+)  # Always enable (192,128)
 
 mask_modes = [0, 1, 2]