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:
Kharec committed 2026-10-15 17:52:37 +02:00
1 parent e80ad634ef
commit d6ee57130f
3 files changed
+74 -32

No files matched your search

+10 -5
View File
@@ -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