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.

CHANGELOG.md

@@ -21,6 +21,9 @@ version is in the `VERSION` file.
   power cut just after a save could leave an empty script). An existing file keeps its
   permissions, and saving through a symlink replaces the file it points to instead of
   replacing the link.
+- Translation is faster and uses much less memory: the beam search no longer copies
+  every beam's decoder state for every candidate at every step, and reuses its
+  work buffers instead of allocating new ones per layer and step.
 
 ## [1.4] - 2026-10-15
 

internal/marian/decode.go

@@ -13,7 +13,9 @@ type beam struct {
 	tokens []int // starts with the decoder start token
 	score  float64
 	// Keys and values of the decoder's self-attention, one slice per layer
-	// holding len(tokens) rows: each step only computes its own row.
+	// holding len(tokens) rows: each step only computes its own row. When a
+	// beam is the only survivor of its parent it takes the parent's slices
+	// over instead of copying them (the parent is dead from that step on).
 	keys, values [][]float32
 }
 
@@ -27,6 +29,16 @@ type candidate struct {
 	score         float64
 }
 
+// decoderScratch holds the work buffers of a decoding step, shared by the
+// layers, the beams and the steps of a search: allocating them afresh for
+// every layer of every step costs more than the dot products they serve.
+type decoderScratch struct {
+	q, k, v   []float32
+	ctxv, tmp []float32
+	scores    []float32
+	hidden    []float32
+}
+
 // beamSearch generates the translation, as Hugging Face's generate() does
 // with the model's defaults: num_beams beams, length penalty 1, the
 // padding token never produced.
@@ -35,18 +47,32 @@ func (m *Model) beamSearch(ctx context.Context, enc []float32, srcLen int) ([]in
 	numBeams := max(m.cfg.NumBeams, 1)
 	crossK := make([][]float32, len(m.decoder))
 	crossV := make([][]float32, len(m.decoder))
+	ffn := 0
 	for i, l := range m.decoder {
 		crossK[i] = make([]float32, srcLen*d)
 		crossV[i] = make([]float32, srcLen*d)
 		l.cross.k.rows(crossK[i], enc, srcLen)
 		l.cross.v.rows(crossV[i], enc, srcLen)
+		ffn = max(ffn, l.fc1.out)
+	}
+	maxLen := min(m.cfg.MaxPositions-1, 2*srcLen+10)
+	scratch := &decoderScratch{
+		q: make([]float32, d), k: make([]float32, d), v: make([]float32, d),
+		ctxv: make([]float32, d), tmp: make([]float32, d),
+		scores: make([]float32, max(maxLen+1, srcLen)),
+		hidden: make([]float32, ffn),
+	}
+	// One vocabulary of logits per beam, reused at every step: the largest
+	// allocation of the search.
+	probs := make([][]float32, numBeams)
+	for i := range probs {
+		probs[i] = make([]float32, m.cfg.VocabSize)
 	}
 	beams := []beam{{
 		tokens: []int{m.cfg.DecoderStart},
 		keys:   make([][]float32, len(m.decoder)),
 		values: make([][]float32, len(m.decoder)),
 	}}
-	maxLen := min(m.cfg.MaxPositions-1, 2*srcLen+10)
 	var done []hypothesis
 	for len(beams[0].tokens) <= maxLen {
 		if err := ctx.Err(); err != nil {
@@ -54,9 +80,9 @@ func (m *Model) beamSearch(ctx context.Context, enc []float32, srcLen int) ([]in
 		}
 		hidden := make([][]float32, len(beams))
 		for i := range beams {
-			hidden[i] = m.decodeStep(&beams[i], crossK, crossV, srcLen)
+			hidden[i] = m.decodeStep(&beams[i], crossK, crossV, srcLen, scratch)
 		}
-		logProbs := m.logProbs(hidden)
+		logProbs := m.logProbs(hidden, probs[:len(beams)])
 		var cands []candidate
 		for i, lp := range logProbs {
 			for _, t := range topK(lp, 2*numBeams) {
@@ -72,29 +98,45 @@ func (m *Model) beamSearch(ctx context.Context, enc []float32, srcLen int) ([]in
 			}
 			return 0
 		})
-		var next []beam
+		// Choose the surviving beams first, then build them: a cache is
+		// copied at most once per step, and only when several survivors
+		// share its parent.
+		children := make([]int, len(beams))
+		chosen := make([]candidate, 0, numBeams)
 		for rank, c := range cands {
-			parent := beams[c.parent]
 			if c.token == m.cfg.EOSToken {
 				if rank < numBeams {
+					parent := beams[c.parent]
 					done = append(done, hypothesis{tokens: slices.Clone(parent.tokens), score: c.score / float64(len(parent.tokens))})
 				}
 				continue
 			}
+			chosen = append(chosen, c)
+			children[c.parent]++
+			if len(chosen) == numBeams {
+				break
+			}
+		}
+		next := make([]beam, 0, len(chosen))
+		for _, c := range chosen {
+			parent := &beams[c.parent]
 			child := beam{
 				tokens: append(slices.Clone(parent.tokens), c.token),
 				score:  c.score,
 				keys:   make([][]float32, len(parent.keys)),
 				values: make([][]float32, len(parent.values)),
 			}
+			takeOver := children[c.parent] == 1
 			for l := range parent.keys {
-				child.keys[l] = slices.Clone(parent.keys[l])
+				if takeOver {
-				child.values[l] = slices.Clone(parent.values[l])
+					child.keys[l] = parent.keys[l]
+					child.values[l] = parent.values[l]
+				} else {
+					child.keys[l] = slices.Clone(parent.keys[l])
+					child.values[l] = slices.Clone(parent.values[l])
+				}
 			}
 			next = append(next, child)
-			if len(next) == numBeams {
-				break
-			}
 		}
 		slices.SortStableFunc(done, func(a, b hypothesis) int {
 			switch {
@@ -126,22 +168,17 @@ func (m *Model) beamSearch(ctx context.Context, enc []float32, srcLen int) ([]in
 
 // decodeStep runs the decoder for the beam's last token and returns its
 // output vector.
-func (m *Model) decodeStep(b *beam, crossK, crossV [][]float32, srcLen int) []float32 {
+func (m *Model) decodeStep(b *beam, crossK, crossV [][]float32, srcLen int, s *decoderScratch) []float32 {
 	d := m.cfg.DModel
 	pos := len(b.tokens) - 1
 	tok := b.tokens[pos]
-	x := make([]float32, d)
+	x := make([]float32, d) // the beam's output vector, which lives past the step
 	emb := m.embed[tok*d : (tok+1)*d]
 	p := m.decPos[pos*d : (pos+1)*d]
 	for j := range x {
 		x[j] = emb[j]*m.embedScale + p[j]
 	}
-	q := make([]float32, d)
+	q, k, v, ctxv, tmp, scores := s.q, s.k, s.v, s.ctxv, s.tmp, s.scores
-	k := make([]float32, d)
-	v := make([]float32, d)
-	ctxv := make([]float32, d)
-	tmp := make([]float32, d)
-	scores := make([]float32, max(pos+1, srcLen))
 	for i, l := range m.decoder {
 		l.self.q.apply(q, x)
 		l.self.k.apply(k, x)
@@ -157,7 +194,7 @@ func (m *Model) decodeStep(b *beam, crossK, crossV [][]float32, srcLen int) []fl
 		l.cross.o.apply(tmp, ctxv)
 		addNorm(x, tmp, l.crossNorm, 1)
 
-		hidden := make([]float32, l.fc1.out)
+		hidden := s.hidden[:l.fc1.out]
 		l.fc1.apply(hidden, x)
 		for j := range hidden {
 			hidden[j] = m.activation(hidden[j])
@@ -169,14 +206,11 @@ func (m *Model) decodeStep(b *beam, crossK, crossV [][]float32, srcLen int) []fl
 }
 
 // logProbs projects each beam's vector on the vocabulary (the embedding
-// matrix, shared with the output layer) and returns log-softmax scores.
+// matrix, shared with the output layer) and returns log-softmax scores in
-// This is the largest product of a step, so it is split across CPUs.
+// out, one vocabulary per beam, reused from step to step. This is the
-func (m *Model) logProbs(hidden [][]float32) [][]float32 {
+// largest product of a step, so it is split across CPUs.
+func (m *Model) logProbs(hidden, out [][]float32) [][]float32 {
 	d, vocab := m.cfg.DModel, m.cfg.VocabSize
-	out := make([][]float32, len(hidden))
-	for i := range out {
-		out[i] = make([]float32, vocab)
-	}
 	workers := runtime.GOMAXPROCS(0)
 	chunk := (vocab + workers - 1) / workers
 	var wg sync.WaitGroup

internal/marian/model.go

@@ -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)
+		h := hidden[:n*l.fc1.out]
-		l.fc1.rows(hidden, x, n)
+		l.fc1.rows(h, x, n)
-		for i := range hidden {
+		for i := range h {
-			hidden[i] = m.activation(hidden[i])
+			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