[Feature Request] Support for additional head_dim configurations (e.g., 64, 96)

[JiaxiangZheng]

The current repository appears to only support head_dim values of 128 and 80. This limitation restricts the use of other common model configurations. Could you consider adding support for additional head_dim values (e.g., 96, 64, etc.)?

Thanks!