Merge pull request ollama#3466 from ollama/mxyng/head-kv

default head_kv to 1
galonge · Apr 3, 2024 · 4f895d6 · 4f895d6
2 parents 7d05a6e + 90f071c
commit 4f895d6
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/llm/ggml.go b/llm/ggml.go
@@ -148,15 +148,15 @@ func (kv KV) HeadCount() uint64 {
 }
 
 func (kv KV) HeadCountKV() uint64 {
-	return kv.u64(fmt.Sprintf("%s.attention.head_count_kv", kv.Architecture()))
+	if headCountKV := kv.u64(fmt.Sprintf("%s.attention.head_count_kv", kv.Architecture())); headCountKV > 0 {
+		return headCountKV
+	}
+
+	return 1
 }
 
 func (kv KV) GQA() uint64 {
-	if headCountKV := kv.HeadCountKV(); headCountKV > 0 {
-		return kv.HeadCount() / headCountKV
-	}
-
-	return 0
+	return kv.HeadCount() / kv.HeadCountKV()
 }
 
 func (kv KV) EmbeddingLength() uint64 {