From b880b76c4308cd5a97cbef2ef86367d126d77f75 Mon Sep 17 00:00:00 2001
From: Jesse Gross <jesse@kernel.org>
Date: Sat, 1 Aug 2026 21:53:20 -0700
Subject: [PATCH] laguna: wire the DFlash target side

Add what a DFlash draft borrows from its target: the tapped layer
outputs, the raw embedding lookup, and the undecorated lm_head
projection. The laguna draft architecture (DFlashLagunaForCausalLM) is
registered here, alongside the only wired target.

Matched nvfp4 target+draft pairs, M5 Max, temp 0.8, repeat_penalty 1.1,
adaptive depth; decode tok/s:

                     prose   code   edit
  laguna-xs  plain   139.4  139.7  137.3
             DFlash  142.3  139.2  145.1
  laguna-s   plain    75.4   70.0   72.4
             DFlash   74.6   80.8  115.3
---
 x/models/dflash/dflash.go |  3 +++
 x/models/laguna/laguna.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/x/models/dflash/dflash.go b/x/models/dflash/dflash.go
index 07526e2c4..3283a9e7f 100644
--- a/x/models/dflash/dflash.go
+++ b/x/models/dflash/dflash.go
@@ -20,6 +20,9 @@ func init() {
 	base.RegisterDraft("DFlashDraftModel", func(root *model.Root, target base.Model) (base.DraftModel, error) {
 		return newModel(root, target, false)
 	})
+	base.RegisterDraft("DFlashLagunaForCausalLM", func(root *model.Root, target base.Model) (base.DraftModel, error) {
+		return newModel(root, target, true)
+	})
 }
 
 var _ base.BlockDraft = (*Model)(nil)
diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go
index 876d1752f..8940e93b9 100644
--- a/x/models/laguna/laguna.go
+++ b/x/models/laguna/laguna.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"math"
+	"slices"
 	"strings"
 
 	"github.com/ollama/ollama/x/mlxrunner/batch"
@@ -83,6 +84,9 @@ type Model struct {
 	Norm        *nn.RMSNorm
 	LMHead      nn.LinearLayer
 
+	// auxHiddenLayers are the tapped layers; empty means the final hidden.
+	auxHiddenLayers []int
+
 	tok *tokenizer.Tokenizer
 	*Config
 }
@@ -1395,6 +1399,7 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden
 
 func (m *Model) forward(b *batch.Batch, caches []cache.Cache, B, L int32) (hidden, auxHidden *mlx.Array) {
 	positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets))
+	var features []*mlx.Array
 	h := m.EmbedTokens.Forward(b.InputIDs)
 	for i, layer := range m.Layers {
 		var c cache.Cache
@@ -1402,11 +1407,34 @@ func (m *Model) forward(b *batch.Batch, caches []cache.Cache, B, L int32) (hidde
 			c = caches[i]
 		}
 		h = layer.Forward(h, b, c, positions, B, L, m.Config)
+		if slices.Contains(m.auxHiddenLayers, i) {
+			features = append(features, h)
+		}
 	}
 	out := m.Norm.Forward(h, m.RMSNormEps)
+	if features != nil {
+		return out, mlx.Concatenate(features, -1)
+	}
 	return out, out
 }
 
+// SetAuxHiddenLayers taps the listed layers' outputs, which Forward then
+// returns as the draft-conditioning state in place of the final hidden.
+func (m *Model) SetAuxHiddenLayers(layers []int) {
+	m.auxHiddenLayers = layers
+}
+
+// TokenEmbeddings is the raw lookup, for a draft that embeds with the
+// target's table.
+func (m *Model) TokenEmbeddings(ids *mlx.Array) *mlx.Array {
+	return m.EmbedTokens.Forward(ids)
+}
+
+// RawLogits matches Unembed: this head applies no output decoration.
+func (m *Model) RawLogits(hidden *mlx.Array) *mlx.Array {
+	return m.LMHead.Forward(hidden)
+}
+
 func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
 	return m.LMHead.Forward(x)
 }