refactor(model): extract attention primitives into model/attn

Snider · Snider · commit febe3a63447a · 2026-07-14T22:53:24.000+01:00
Move the shared fused-QKV split helpers (SplitContiguousQKV, SplitGroupedQKV,
SplitInterleavedQKV, SplitContiguousGateUp) and the RoPE geometry primitive
(RopeParams + RotaryDim) out of the flat model root into a dedicated model/attn
feature package, so any arch composes them off the shelf. model/attn is a leaf:
it imports only safetensors + core, never model, so there is no import cycle.

Consumers updated to attn.*: bloom, dbrx, falcon, glm4, mpt, phi, stablelm.
No engine/metal or engine/hip consumer references these symbols. RopeParams
error prefixes retargeted model.* -&gt; attn.* to match the new home.

Pure move, byte-identical: tests unmodified beyond the package/import rename.

The QK-norm / attention-scaling helpers named in the lane brief are NOT
standalone functions - they live inside the Arch type in model/arch.go and
resist a clean extraction, so they are left in place (flagged, not forced).

Co-Authored-By: Virgil &lt;virgil@lethean.io&gt;
diff --git a/go/model/attn/fused_qkv.go b/go/model/attn/fused_qkv.go
@@ -1,6 +1,6 @@
 // SPDX-Licence-Identifier: EUPL-1.2
 
-package model
+package attn
 
 import "dappco.re/go/inference/model/safetensors"
 
diff --git a/go/model/attn/fused_qkv_test.go b/go/model/attn/fused_qkv_test.go
@@ -1,6 +1,6 @@
 // SPDX-Licence-Identifier: EUPL-1.2
 
-package model
+package attn
 
 import (
 	"slices"
diff --git a/go/model/attn/rope.go b/go/model/attn/rope.go
@@ -1,6 +1,6 @@
 // SPDX-Licence-Identifier: EUPL-1.2
 
-package model
+package attn
 
 import (
 	"math"
@@ -19,22 +19,22 @@ type RopeParams struct {
 // RotaryDim resolves the even prefix width rotated by the RoPE op.
 func (p RopeParams) RotaryDim() (int, error) {
 	if p.HeadDim <= 0 {
-		return 0, core.NewError("model.RopeParams.RotaryDim: head dimension must be > 0")
+		return 0, core.NewError("attn.RopeParams.RotaryDim: head dimension must be > 0")
 	}
 	factor := p.PartialRotaryFactor
 	if factor == 0 {
 		factor = 1
 	}
 	if factor <= 0 || factor > 1 {
-		return 0, core.NewError("model.RopeParams.RotaryDim: partial rotary factor must be in (0,1]")
+		return 0, core.NewError("attn.RopeParams.RotaryDim: partial rotary factor must be in (0,1]")
 	}
 	resolved := float64(p.HeadDim) * float64(factor)
 	dim := int(math.Round(resolved))
 	if math.Abs(resolved-float64(dim)) > 1e-5 {
-		return 0, core.NewError("model.RopeParams.RotaryDim: factor must resolve to a whole dimension")
+		return 0, core.NewError("attn.RopeParams.RotaryDim: factor must resolve to a whole dimension")
 	}
 	if dim <= 0 || dim > p.HeadDim || dim%2 != 0 {
-		return 0, core.NewError("model.RopeParams.RotaryDim: resolved rotary dimension must be positive and even")
+		return 0, core.NewError("attn.RopeParams.RotaryDim: resolved rotary dimension must be positive and even")
 	}
 	return dim, nil
 }
diff --git a/go/model/attn/rope_bench_test.go b/go/model/attn/rope_bench_test.go
@@ -1,6 +1,6 @@
 // SPDX-Licence-Identifier: EUPL-1.2
 
-package model
+package attn
 
 import "testing"
 
diff --git a/go/model/attn/rope_example_test.go b/go/model/attn/rope_example_test.go
@@ -1,20 +1,20 @@
 // SPDX-Licence-Identifier: EUPL-1.2
 
-package model_test
+package attn_test
 
 import (
 	core "dappco.re/go"
-	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 )
 
 func ExampleRopeParams() {
-	dim, _ := (model.RopeParams{HeadDim: 80, PartialRotaryFactor: 0.4}).RotaryDim()
+	dim, _ := (attn.RopeParams{HeadDim: 80, PartialRotaryFactor: 0.4}).RotaryDim()
 	core.Println(dim)
 	// Output: 32
 }
 
 func ExampleRopeParams_RotaryDim() {
-	dim, _ := (model.RopeParams{HeadDim: 64}).RotaryDim()
+	dim, _ := (attn.RopeParams{HeadDim: 64}).RotaryDim()
 	core.Println(dim)
 	// Output: 64
 }
diff --git a/go/model/attn/rope_test.go b/go/model/attn/rope_test.go
@@ -1,6 +1,6 @@
 // SPDX-Licence-Identifier: EUPL-1.2
 
-package model
+package attn
 
 import "testing"
 
diff --git a/go/model/bloom/register.go b/go/model/bloom/register.go
@@ -5,6 +5,7 @@ package bloom
 import (
 	core "dappco.re/go"
 	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 	"dappco.re/go/inference/model/safetensors"
 )
 
@@ -25,7 +26,7 @@ func init() {
 				break
 			}
 			headDim := cfg.HiddenSize / cfg.NumAttentionHeads
-			tensors = model.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)
+			tensors = attn.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)
 		}
 		return tensors
 	}})
diff --git a/go/model/dbrx/register.go b/go/model/dbrx/register.go
@@ -5,6 +5,7 @@ package dbrx
 import (
 	core "dappco.re/go"
 	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 	"dappco.re/go/inference/model/composed"
 	"dappco.re/go/inference/model/safetensors"
 )
@@ -33,7 +34,7 @@ func NormalizeWeights(in map[string]safetensors.Tensor, cfg Config) map[string]s
 		alias(dst+"input_layernorm.weight", src+"norm_attn_norm.norm_1.weight")
 		alias(dst+"post_attention_layernorm.weight", src+"norm_attn_norm.norm_2.weight")
 		alias(dst+"self_attn.o_proj.weight", src+"norm_attn_norm.attn.out_proj.weight")
-		out = model.SplitContiguousQKV(out, src+"norm_attn_norm.attn.Wqkv", dst+"self_attn.q_proj", dst+"self_attn.k_proj", dst+"self_attn.v_proj", cfg.Heads*headDim, cfg.Attention.KVHeads*headDim)
+		out = attn.SplitContiguousQKV(out, src+"norm_attn_norm.attn.Wqkv", dst+"self_attn.q_proj", dst+"self_attn.k_proj", dst+"self_attn.v_proj", cfg.Heads*headDim, cfg.Attention.KVHeads*headDim)
 		alias(dst+"mlp.gate.weight", src+"ffn.router.layer.weight")
 		for _, part := range []struct{ source, target string }{{"w1", "gate_proj.weight"}, {"v1", "up_proj.weight"}, {"w2", "down_proj.weight"}} {
 			tensor, ok := in[src+"ffn.experts.mlp."+part.source]
diff --git a/go/model/falcon/register.go b/go/model/falcon/register.go
@@ -5,6 +5,7 @@ package falcon
 import (
 	core "dappco.re/go"
 	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 	"dappco.re/go/inference/model/safetensors"
 )
 
@@ -27,11 +28,11 @@ func init() {
 			}
 			headDim := cfg.HiddenSize / cfg.NumAttentionHeads
 			if cfg.NewDecoderArchitecture && cfg.NumKVHeads > 0 {
-				tensors = model.SplitGroupedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, cfg.NumKVHeads, headDim)
+				tensors = attn.SplitGroupedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, cfg.NumKVHeads, headDim)
 			} else if cfg.MultiQuery {
-				tensors = model.SplitContiguousQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads*headDim, headDim)
+				tensors = attn.SplitContiguousQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads*headDim, headDim)
 			} else {
-				tensors = model.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)
+				tensors = attn.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)
 			}
 		}
 		return tensors
diff --git a/go/model/glm4/glm4.go b/go/model/glm4/glm4.go
@@ -8,6 +8,7 @@ import (
 
 	core "dappco.re/go"
 	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 	"dappco.re/go/inference/model/safetensors"
 )
 
@@ -32,7 +33,7 @@ func (c *Config) Arch() (model.Arch, error) {
 	if c.ModelType != "glm4" || c.HiddenSize <= 0 || c.IntermediateSize <= 0 || c.NumHiddenLayers <= 0 || c.NumAttentionHeads <= 0 || c.NumKeyValueHeads <= 0 || c.HeadDim <= 0 || c.VocabSize <= 0 || c.RMSNormEps <= 0 || c.RopeTheta <= 0 || c.NumAttentionHeads%c.NumKeyValueHeads != 0 {
 		return model.Arch{}, core.NewError("glm4.Config.Arch: invalid dense GLM-4 declaration")
 	}
-	rd, err := (model.RopeParams{HeadDim: c.HeadDim, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()
+	rd, err := (attn.RopeParams{HeadDim: c.HeadDim, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()
 	if err != nil {
 		return model.Arch{}, core.E("glm4.Config.Arch", "partial rotary", err)
 	}
@@ -56,7 +57,7 @@ func init() {
 			if _, ok := out[p+"gate_up_proj.weight"]; !ok {
 				break
 			}
-			out = model.SplitContiguousGateUp(out, p+"gate_up_proj", p+"gate_proj", p+"up_proj")
+			out = attn.SplitContiguousGateUp(out, p+"gate_up_proj", p+"gate_proj", p+"up_proj")
 		}
 		return out
 	}})
diff --git a/go/model/mpt/register.go b/go/model/mpt/register.go
@@ -5,6 +5,7 @@ package mpt
 import (
 	core "dappco.re/go"
 	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 	"dappco.re/go/inference/model/safetensors"
 )
 
@@ -21,7 +22,7 @@ func init() {
 		rows := c.DModel
 		for i := 0; i < c.NLayers; i++ {
 			p := core.Sprintf("transformer.blocks.%d.attn.", i)
-			t = model.SplitContiguousQKV(t, p+"Wqkv.weight", p+"q_proj.weight", p+"k_proj.weight", p+"v_proj.weight", rows, rows)
+			t = attn.SplitContiguousQKV(t, p+"Wqkv.weight", p+"q_proj.weight", p+"k_proj.weight", p+"v_proj.weight", rows, rows)
 		}
 		return t
 	}})
diff --git a/go/model/phi/config.go b/go/model/phi/config.go
@@ -8,6 +8,7 @@ import (
 
 	core "dappco.re/go"
 	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 	"dappco.re/go/inference/model/safetensors"
 )
 
@@ -57,7 +58,7 @@ func (c *Config) Arch() (model.Arch, error) {
 	if c.ModelType == "phi3" && kvHeads != c.NumAttentionHeads {
 		return model.Arch{}, core.NewError("phi.Config.Arch: Phi-3 fused qkv requires equal query and key/value head counts")
 	}
-	rotaryDim, err := (model.RopeParams{HeadDim: headDim, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()
+	rotaryDim, err := (attn.RopeParams{HeadDim: headDim, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()
 	if err != nil {
 		return model.Arch{}, core.E("phi.Config.Arch", "resolve partial rotary", err)
 	}
diff --git a/go/model/stablelm/config.go b/go/model/stablelm/config.go
@@ -8,6 +8,7 @@ import (
 
 	core "dappco.re/go"
 	"dappco.re/go/inference/model"
+	"dappco.re/go/inference/model/attn"
 	"dappco.re/go/inference/model/safetensors"
 )
 
@@ -44,7 +45,7 @@ func (c *Config) Arch() (model.Arch, error) {
 	if kv <= 0 || c.NumAttentionHeads%kv != 0 {
 		return model.Arch{}, core.E("stablelm.Config.Arch", "invalid key/value heads", nil)
 	}
-	rd, err := (model.RopeParams{HeadDim: hd, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()
+	rd, err := (attn.RopeParams{HeadDim: hd, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()
 	if err != nil {
 		return model.Arch{}, core.E("stablelm.Config.Arch", "resolve partial rotary", err)
 	}

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`// SPDX-Licence-Identifier: EUPL-1.2`
`2`	`2`
`3`		`-package model`
	`3`	`+package attn`
`4`	`4`
`5`	`5`	`import (`
`6`	`6`	`"math"`
`@@ -19,22 +19,22 @@ type RopeParams struct {`
`19`	`19`	`// RotaryDim resolves the even prefix width rotated by the RoPE op.`
`20`	`20`	`func (p RopeParams) RotaryDim() (int, error) {`
`21`	`21`	`if p.HeadDim <= 0 {`
`22`		`- return 0, core.NewError("model.RopeParams.RotaryDim: head dimension must be > 0")`
	`22`	`+ return 0, core.NewError("attn.RopeParams.RotaryDim: head dimension must be > 0")`
`23`	`23`	`}`
`24`	`24`	`factor := p.PartialRotaryFactor`
`25`	`25`	`if factor == 0 {`
`26`	`26`	`factor = 1`
`27`	`27`	`}`
`28`	`28`	`if factor <= 0 \|\| factor > 1 {`
`29`		`- return 0, core.NewError("model.RopeParams.RotaryDim: partial rotary factor must be in (0,1]")`
	`29`	`+ return 0, core.NewError("attn.RopeParams.RotaryDim: partial rotary factor must be in (0,1]")`
`30`	`30`	`}`
`31`	`31`	`resolved := float64(p.HeadDim) * float64(factor)`
`32`	`32`	`dim := int(math.Round(resolved))`
`33`	`33`	`if math.Abs(resolved-float64(dim)) > 1e-5 {`
`34`		`- return 0, core.NewError("model.RopeParams.RotaryDim: factor must resolve to a whole dimension")`
	`34`	`+ return 0, core.NewError("attn.RopeParams.RotaryDim: factor must resolve to a whole dimension")`
`35`	`35`	`}`
`36`	`36`	`if dim <= 0 \|\| dim > p.HeadDim \|\| dim%2 != 0 {`
`37`		`- return 0, core.NewError("model.RopeParams.RotaryDim: resolved rotary dimension must be positive and even")`
	`37`	`+ return 0, core.NewError("attn.RopeParams.RotaryDim: resolved rotary dimension must be positive and even")`
`38`	`38`	`}`
`39`	`39`	`return dim, nil`
`40`	`40`	`}`
Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@ package bloom`
`5`	`5`	`import (`
`6`	`6`	`core "dappco.re/go"`
`7`	`7`	`"dappco.re/go/inference/model"`
	`8`	`+ "dappco.re/go/inference/model/attn"`
`8`	`9`	`"dappco.re/go/inference/model/safetensors"`
`9`	`10`	`)`
`10`	`11`
`@@ -25,7 +26,7 @@ func init() {`
`25`	`26`	`break`
`26`	`27`	`}`
`27`	`28`	`headDim := cfg.HiddenSize / cfg.NumAttentionHeads`
`28`		`- tensors = model.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)`
	`29`	`+ tensors = attn.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)`
`29`	`30`	`}`
`30`	`31`	`return tensors`
`31`	`32`	`}})`
Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@ package falcon`
`5`	`5`	`import (`
`6`	`6`	`core "dappco.re/go"`
`7`	`7`	`"dappco.re/go/inference/model"`
	`8`	`+ "dappco.re/go/inference/model/attn"`
`8`	`9`	`"dappco.re/go/inference/model/safetensors"`
`9`	`10`	`)`
`10`	`11`
`@@ -27,11 +28,11 @@ func init() {`
`27`	`28`	`}`
`28`	`29`	`headDim := cfg.HiddenSize / cfg.NumAttentionHeads`
`29`	`30`	`if cfg.NewDecoderArchitecture && cfg.NumKVHeads > 0 {`
`30`		`- tensors = model.SplitGroupedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, cfg.NumKVHeads, headDim)`
	`31`	`+ tensors = attn.SplitGroupedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, cfg.NumKVHeads, headDim)`
`31`	`32`	`} else if cfg.MultiQuery {`
`32`		`- tensors = model.SplitContiguousQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads*headDim, headDim)`
	`33`	`+ tensors = attn.SplitContiguousQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads*headDim, headDim)`
`33`	`34`	`} else {`
`34`		`- tensors = model.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)`
	`35`	`+ tensors = attn.SplitInterleavedQKV(tensors, p+"query_key_value", p+"query", p+"key", p+"value", cfg.NumAttentionHeads, headDim)`
`35`	`36`	`}`
`36`	`37`	`}`
`37`	`38`	`return tensors`
Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,7 @@ import (`
`8`	`8`
`9`	`9`	`core "dappco.re/go"`
`10`	`10`	`"dappco.re/go/inference/model"`
	`11`	`+ "dappco.re/go/inference/model/attn"`
`11`	`12`	`"dappco.re/go/inference/model/safetensors"`
`12`	`13`	`)`
`13`	`14`
`@@ -32,7 +33,7 @@ func (c *Config) Arch() (model.Arch, error) {`
`32`	`33`	`if c.ModelType != "glm4" \|\| c.HiddenSize <= 0 \|\| c.IntermediateSize <= 0 \|\| c.NumHiddenLayers <= 0 \|\| c.NumAttentionHeads <= 0 \|\| c.NumKeyValueHeads <= 0 \|\| c.HeadDim <= 0 \|\| c.VocabSize <= 0 \|\| c.RMSNormEps <= 0 \|\| c.RopeTheta <= 0 \|\| c.NumAttentionHeads%c.NumKeyValueHeads != 0 {`
`33`	`34`	`return model.Arch{}, core.NewError("glm4.Config.Arch: invalid dense GLM-4 declaration")`
`34`	`35`	`}`
`35`		`- rd, err := (model.RopeParams{HeadDim: c.HeadDim, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()`
	`36`	`+ rd, err := (attn.RopeParams{HeadDim: c.HeadDim, PartialRotaryFactor: c.PartialRotaryFactor}).RotaryDim()`
`36`	`37`	`if err != nil {`
`37`	`38`	`return model.Arch{}, core.E("glm4.Config.Arch", "partial rotary", err)`
`38`	`39`	`}`
`@@ -56,7 +57,7 @@ func init() {`
`56`	`57`	`if _, ok := out[p+"gate_up_proj.weight"]; !ok {`
`57`	`58`	`break`
`58`	`59`	`}`
`59`		`- out = model.SplitContiguousGateUp(out, p+"gate_up_proj", p+"gate_proj", p+"up_proj")`
	`60`	`+ out = attn.SplitContiguousGateUp(out, p+"gate_up_proj", p+"gate_proj", p+"up_proj")`
`60`	`61`	`}`
`61`	`62`	`return out`
`62`	`63`	`}})`