baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Merge pull request #5121 from ollama/mxyng/deepseekv2

Browse source 
deepseek v2 graph
This commit is contained in:
Michael Yang 2024-06-18 16:30:58 -07:00 committed by GitHub
commit 21adf8b6d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 11 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
llm/ggml.go
View file

@ -367,6 +367,17 @@ func (llm GGML) GraphSize(context, batch uint64) (partialOffload, fullOffload ui
4*batch*(vocab+2*embedding), 4*batch*(vocab+2*embedding),
fullOffload, fullOffload,
) )
case "deepseek2":
keys := uint64(llm.KV()["deepseek2.attention.key_length"].(uint32))
fullOffload = max(
4*batch*(3*embedding+vocab),
4*batch*(3*embedding+2+context*(1+headsKV)+2*keys*headsKV),
)
partialOffload = max(
4*batch*(3*embedding+vocab)+embedding*vocab*105/128,
4*batch*(2*embedding+1+2*keys*headsKV+context+context*headsKV)+4*keys*context*headsKV+embedding*keys*headsKV*9/16,
)
} }
return return