squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Fix KV head calculation for MPT models when using GQA (#5142)

Browse Source
This commit is contained in:
Bruce Fontaine 2024-06-17 15:26:41 -07:00 committed by GitHub
parent e441bad674
commit a3e8a05d4c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/config.py
View File

@ -302,7 +302,11 @@ class ModelConfig:
return 1
# For DBRX and MPT
if self.hf_config.model_type in ["dbrx", "mpt"]:
if self.hf_config.model_type == "mpt":
if "kv_n_heads" in self.hf_config.attn_config:
return self.hf_config.attn_config["kv_n_heads"]
return self.hf_config.num_attention_heads
if self.hf_config.model_type == "dbrx":
return getattr(self.hf_config.attn_config, "kv_n_heads",
self.hf_config.num_attention_heads)