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

hongxiayang · Jan 27, 2024 · 1b669dd · 1b669dd
1 parent 17f61af
commit 1b669dd
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/vllm/config.py b/vllm/config.py
@@ -213,6 +213,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