Use head_dim in config if exists (#2622)

2024-01-27 10:30:49 -08:00 · 2024-01-27 10:30:49 -08:00 · 220a47627b
commit 220a47627b
parent beb89f68b4
1 changed files with 2 additions and 0 deletions
--- a/vllm/config.py
+++ b/vllm/config.py
@ -213,6 +213,8 @@ class ModelConfig:
        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