fix: copy decoder caches once per surviving beam, not per candidate
Each beam-search step deep-cloned the keys and values of every layer for every candidate it kept, so the cost grew quadratically with the sentence length and produced garbage faster than the unrolled dot() could win. Survivors are now chosen first; a beam whose parent has a single survivor takes its slices over (the parent is dead), and only a shared cache is cloned. The step's other buffers (query/key/value, hidden, logits) are reused across layers, beams and steps instead of being allocated afresh each time. TestRealModel confirms the output is unchanged.
This commit is contained in:
1 parent
e80ad634ef
commit
d6ee57130f
3 files changed
+74
-32
No files matched your search
@@ -276,6 +276,11 @@ func (m *Model) encode(ids []int) []float32 {
|
||||
ctxv := make([]float32, n*d)
|
||||
tmp := make([]float32, n*d)
|
||||
scores := make([]float32, n)
|
||||
ffn := 0 // the hidden layer of every encoder layer, allocated once
|
||||
for _, l := range m.encoder {
|
||||
ffn = max(ffn, l.fc1.out)
|
||||
}
|
||||
hidden := make([]float32, n*ffn)
|
||||
for _, l := range m.encoder {
|
||||
l.self.q.rows(q, x, n)
|
||||
l.self.k.rows(k, x, n)
|
||||
@@ -285,12 +290,12 @@ func (m *Model) encode(ids []int) []float32 {
|
||||
}
|
||||
l.self.o.rows(tmp, ctxv, n)
|
||||
addNorm(x, tmp, l.selfNorm, n)
|
||||
hidden := make([]float32, n*l.fc1.out)
|
||||
l.fc1.rows(hidden, x, n)
|
||||
for i := range hidden {
|
||||
hidden[i] = m.activation(hidden[i])
|
||||
h := hidden[:n*l.fc1.out]
|
||||
l.fc1.rows(h, x, n)
|
||||
for i := range h {
|
||||
h[i] = m.activation(h[i])
|
||||
}
|
||||
l.fc2.rows(tmp, hidden, n)
|
||||
l.fc2.rows(tmp, h, n)
|
||||
addNorm(x, tmp, l.finalLayerNorm, n)
|
||||
}
|
||||
return x
|
||||
|
||||
Reference in new issue
Block a user