Use head_dim in config if exists (vllm-project#2622)

Author: xiangxu-google

vllm/config.py

@@ -248,6 +248,8 @@ def get_hidden_size(self) -> int:
         return self.hf_config.hidden_size
 
     def get_head_size(self) -> int:
+        if hasattr(self.hf_config, "head_dim"):
+            return self.hf_config.head_dim
         # FIXME(woosuk): This may not be true for all models.
         return self.hf_config.hidden_size // self.hf_config.num_attention_heads