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

Merge pull request #7657 from ollama/mxyng/sync

Browse source 
fix(mllama): sync backend between batches
This commit is contained in:
Michael Yang 2024-11-14 09:40:04 -08:00 committed by GitHub
commit 549c2bdfcf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 11 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
llama/llama.go
View file

@ -600,6 +600,10 @@ func (c *Context) SetCrossAttention(state bool) {
C.llama_set_cross_attention(c.c, C.bool(state)) C.llama_set_cross_attention(c.c, C.bool(state))
} }
func (c *Context) Synchronize() {
C.llama_synchronize(c.c)
}
// sampling // sampling
// TODO: this is a temporary wrapper to allow calling C++ code from CGo // TODO: this is a temporary wrapper to allow calling C++ code from CGo
type SamplingContext struct { type SamplingContext struct {

7
llama/runner/runner.go
View file

@ -409,6 +409,13 @@ func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch)
return return
} }
if crossAttention {
// synchronize state to ensure the cross attention batch is complete.
// needed specifically for multi-GPU systems otherwise an inflight
// task may be incorrectly invalidated causing a crash
s.lc.Synchronize()
}
for i, seq := range s.seqs { for i, seq := range s.seqs {
if seq == nil { if seq == nil {
continue continue