CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-08-27 04:06:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 27

model: add Laguna MLX support (#17237)

Browse source 
* model: add Laguna MLX support

Add Laguna XS 2, XS 2.1, and S 2.1 support to the MLX model and create paths.

Read the source config to apply one quantization policy across dense and routed MoE layers. Keep the tied output head and router at source precision, quantize supported attention and expert projections, selectively promote sensitive expert down projections, and emit per-tensor metadata for mixed quantization blobs.

Correct dense expert loading, BF16 source-layout handling, expert global-scale shapes and dtypes, routing-score scaling, and mixed-precision expert dispatch. Gate/up and down projections select quantized or dense execution independently so promoted BF16 down projections do not force quantized gate/up weights through the dense fallback.

Optimize the forward pass with compatible gate/up fusion, sorted standard GatherMM and GatherQMM operations for larger prefills, model-local mlx.Compile closures for elementwise MoE work, and cache-backed 512-token prefill chunks. This keeps the implementation on maintained MLX operations without custom kernels.

Add focused tests for Laguna configuration variants, quantization policy and metadata, dense and routed expert loading, mixed-precision dispatch, compiled-versus-eager parity, fused projections, routing, and prefill chunking.

* review comments and S 2.1 performance fixes

Address renderer/parser selection and mixed-precision expert quantization review feedback.

Keep Laguna weights resident on Metal to prevent repeated paging of its large, sparsely accessed expert buffers. Scope this policy to Laguna GPU execution.

Remove obsolete 512-token prefill chunking now that the runner's 2048-token path is faster.

* review comments addressed

* fix create
This commit is contained in:
Daniel Hiltgen 2026-07-24 18:24:53 -07:00 • committed by GitHub
parent 132e0ca25d
commit 64ee2f9847
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 1449 additions and 124 deletions
Download patch file Download diff file Expand all files Collapse all files

24
x/create/blockfp8.go
View file

@ -71,18 +71,18 @@ func planBlockFP8(inv Inventory, target string, policy quantizePolicy) ([]BlobSp
}
for _, gp := range sortedKeys(groups) {
spec, err := planExpertGroup(gp, groups[gp], "", policy)
groupSpecs, err := planExpertGroup(gp, groups[gp], "", policy)
if err != nil {
return nil, err
}
specs = append(specs, spec)
specs = append(specs, groupSpecs...)
}
for _, gp := range sortedKeys(fp8Groups) {
spec, err := planFP8ExpertGroup(gp, fp8Groups[gp], inv, target, policy)
groupSpecs, err := planFP8ExpertGroup(gp, fp8Groups[gp], inv, target, policy)
if err != nil {
return nil, err
}
specs = append(specs, spec)
specs = append(specs, groupSpecs...)
}
return specs, nil
}
@ -92,7 +92,7 @@ func planBlockFP8(inv Inventory, target string, policy quantizePolicy) ([]BlobSp
// dequantized from FP8 with their block scales, and quantized per the policy.
// The stacking, decode, and quantize all run on the MLX writer thread; the
// planner only groups and orders the source weights and their scale companions.
func planFP8ExpertGroup(groupPrefix string, tensors []SourceTensor, inv Inventory, target string, policy quantizePolicy) (BlobSpec, error) {
func planFP8ExpertGroup(groupPrefix string, tensors []SourceTensor, inv Inventory, target string, policy quantizePolicy) ([]BlobSpec, error) {
type expert struct {
idx int
weight SourceTensor
@ -102,16 +102,16 @@ func planFP8ExpertGroup(groupPrefix string, tensors []SourceTensor, inv Inventor
for _, t := range tensors {
idx, proj, err := parseExpertTensor(groupPrefix, t.Name)
if err != nil {
return BlobSpec{}, err
return nil, err
}
scaleName, ok := fp8ScaleFor(inv, t.Name)
if !ok {
return BlobSpec{}, fmt.Errorf("fp8 expert weight %q has no scale companion", t.Name)
return nil, fmt.Errorf("fp8 expert weight %q has no scale companion", t.Name)
}
byProj[proj] = append(byProj[proj], expert{idx: idx, weight: t, scale: inv.Tensors[scaleName]})
}
spec := BlobSpec{Name: groupPrefix}
var tensorSpecs []TensorSpec
for _, proj := range sortedKeys(byProj) {
experts := byProj[proj]
sort.Slice(experts, func(i, j int) bool { return experts[i].idx < experts[j].idx })
@ -124,11 +124,11 @@ func planFP8ExpertGroup(groupPrefix string, tensors []SourceTensor, inv Inventor
scales := make([]SourceTensor, 0, len(experts))
for _, e := range experts {
if e.weight.Dtype != base.Dtype || !slices.Equal(e.weight.Shape, base.Shape) {
return BlobSpec{}, fmt.Errorf("fp8 expert group %s projection %s has mismatched weight layout (%s %v vs %s %v)",
return nil, fmt.Errorf("fp8 expert group %s projection %s has mismatched weight layout (%s %v vs %s %v)",
groupPrefix, proj, base.Dtype, base.Shape, e.weight.Dtype, e.weight.Shape)
}
if e.scale.Dtype != baseScale.Dtype || !slices.Equal(e.scale.Shape, baseScale.Shape) {
return BlobSpec{}, fmt.Errorf("fp8 expert group %s projection %s has mismatched scale layout (%s %v vs %s %v)",
return nil, fmt.Errorf("fp8 expert group %s projection %s has mismatched scale layout (%s %v vs %s %v)",
groupPrefix, proj, baseScale.Dtype, baseScale.Shape, e.scale.Dtype, e.scale.Shape)
}
sources = append(sources, e.weight)
@ -138,7 +138,7 @@ func planFP8ExpertGroup(groupPrefix string, tensors []SourceTensor, inv Inventor
stackedName := groupPrefix + "." + proj + ".weight"
stackedShape := append([]int32{int32(len(experts))}, base.Shape...)
spec.Tensors = append(spec.Tensors, TensorSpec{
tensorSpecs = append(tensorSpecs, TensorSpec{
Name: stackedName,
Sources: sources,
Transform: TransformDecodeStackFP8,
@ -147,7 +147,7 @@ func planFP8ExpertGroup(groupPrefix string, tensors []SourceTensor, inv Inventor
OutShape: stackedShape,
})
}
return spec, nil
return homogeneousExpertBlobs(groupPrefix, tensorSpecs), nil
}
// isFP8Weight reports whether name is an F8_E4M3 weight with a block-scale

25
x/create/client/create.go
View file

@ -593,6 +593,23 @@ func isQwen35Family(s string) bool {
return strings.Contains(s, "qwen3_5") || strings.Contains(s, "qwen3next")
}
func lagunaRendererParserName(modelDir string) string {
const poolsideV1Marker = "laguna_glm_thinking_v8"
if strings.Contains(readChatTemplate(modelDir), poolsideV1Marker) {
return "poolside-v1"
}
// Poolside's tokenizer config includes the standalone template by name
// rather than embedding it, so inspect that file as well.
if data, err := os.ReadFile(filepath.Join(modelDir, "chat_template.jinja")); err == nil &&
strings.Contains(string(data), poolsideV1Marker) {
return "poolside-v1"
}
return "laguna"
}
// getParserName returns the parser name for a model based on its architecture.
// This reads the config.json from the model directory and determines the appropriate parser.
func getParserName(modelDir string) string {
@ -614,7 +631,7 @@ func getParserName(modelDir string) string {
for _, arch := range cfg.Architectures {
archLower := strings.ToLower(arch)
if strings.Contains(archLower, "laguna") {
return "laguna"
return lagunaRendererParserName(modelDir)
}
if strings.Contains(archLower, "cohere2moe") || strings.Contains(archLower, "cohere2_moe") {
return "cohere"
@ -640,7 +657,7 @@ func getParserName(modelDir string) string {
if cfg.ModelType != "" {
typeLower := strings.ToLower(cfg.ModelType)
if strings.Contains(typeLower, "laguna") {
return "laguna"
return lagunaRendererParserName(modelDir)
}
if strings.Contains(typeLower, "cohere2_moe") {
return "cohere"
@ -686,7 +703,7 @@ func getRendererName(modelDir string) string {
for _, arch := range cfg.Architectures {
archLower := strings.ToLower(arch)
if strings.Contains(archLower, "laguna") {
return "laguna"
return lagunaRendererParserName(modelDir)
}
if strings.Contains(archLower, "cohere2moe") || strings.Contains(archLower, "cohere2_moe") {
return "cohere"
@ -712,7 +729,7 @@ func getRendererName(modelDir string) string {
if cfg.ModelType != "" {
typeLower := strings.ToLower(cfg.ModelType)
if strings.Contains(typeLower, "laguna") {
return "laguna"
return lagunaRendererParserName(modelDir)
}
if strings.Contains(typeLower, "cohere2_moe") {
return "cohere"

46
x/create/client/create_test.go
View file

@ -650,6 +650,11 @@ func TestInferSafetensorsCapabilitiesFromParser(t *testing.T) {
parserName: "laguna",
want: []string{"completion", "tools", "thinking"},
},
{
name: "poolside tools and thinking",
parserName: "poolside-v1",
want: []string{"completion", "tools", "thinking"},
},
{
name: "functiongemma tools only",
parserName: "functiongemma",
@ -797,3 +802,44 @@ func TestGetRendererName(t *testing.T) {
})
}
}
func TestGetLagunaRendererParserName(t *testing.T) {
tests := []struct {
name string
chatTemplate string
want string
}{
{
name: "v5",
chatTemplate: `{#- Iteration on laguna_glm_thinking_v5/chat_template.jinja -#}`,
want: "laguna",
},
{
name: "v8",
chatTemplate: `{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#}`,
want: "poolside-v1",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{"architectures":["LagunaForCausalLM"],"model_type":"laguna"}`), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(dir, "tokenizer_config.json"), []byte(`{"chat_template":"{% include 'chat_template.jinja' %}"}`), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(dir, "chat_template.jinja"), []byte(tt.chatTemplate), 0o644); err != nil {
t.Fatal(err)
}
if got := getParserName(dir); got != tt.want {
t.Errorf("getParserName() = %q, want %q", got, tt.want)
}
if got := getRendererName(dir); got != tt.want {
t.Errorf("getRendererName() = %q, want %q", got, tt.want)
}
})
}
}

160
x/create/laguna.go
View file

@ -2,22 +2,164 @@ package create
import (
"encoding/json"
"fmt"
"strings"
)
type lagunaImportTransform struct{}
func newLagunaImportTransform(json.RawMessage) (quantizePolicy, error) {
return lagunaImportTransform{}, nil
type lagunaImportTransform struct {
denseMLPLayers map[int]bool
numLayers int
}
func (lagunaImportTransform) quantizationType(name string, shape []int32, quantize string) string {
if !lagunaIsHFRoutedExpertWeight(name) {
type lagunaConfig struct {
NumHiddenLayers int `json:"num_hidden_layers"`
MLPOnlyLayers []int `json:"mlp_only_layers"`
MLPLayerTypes []string `json:"mlp_layer_types"`
}
func newLagunaImportTransform(rawConfig json.RawMessage) (quantizePolicy, error) {
var cfg lagunaConfig
if len(rawConfig) > 0 {
if err := json.Unmarshal(rawConfig, &cfg); err != nil {
return nil, fmt.Errorf("laguna: parse config.json: %w", err)
}
}
denseLayers := make(map[int]bool)
for i, typ := range cfg.MLPLayerTypes {
if typ == "dense" {
denseLayers[i] = true
}
}
for _, layer := range cfg.MLPOnlyLayers {
denseLayers[layer] = true
}
if len(denseLayers) == 0 {
denseLayers[0] = true
}
numLayers := cfg.NumHiddenLayers
if numLayers == 0 {
numLayers = len(cfg.MLPLayerTypes)
}
if numLayers == 0 {
numLayers = 40
}
return lagunaImportTransform{
denseMLPLayers: denseLayers,
numLayers: numLayers,
}, nil
}
func (t lagunaImportTransform) quantizationType(name string, shape []int32, quantize string) string {
base := normalizeQuantType(quantize)
if !lagunaFPQuant(base) {
return GetTensorQuantization(name, shape, quantize)
}
switch {
case isEmbedTokensWeight(name) || strings.HasSuffix(name, "lm_head.weight"):
// Laguna has separate embedding and output weights. Both are large,
// quality-sensitive tensors, so keep them at 8-bit for FP quants.
return promoteEmbedding(shape, base)
case strings.HasSuffix(name, ".mlp.gate.weight"):
return ""
case base == "mxfp8" && lagunaAttentionProjection(name):
return ""
case lagunaAttentionProjection(name):
return lagunaQuantizationType(name, shape, base)
case lagunaDenseMLPProjection(name) && t.denseMLPLayers[layerIndex(name)]:
return lagunaQuantizationType(name, shape, base)
case base == "mxfp8" && lagunaRoutedExpertDownProjection(name):
if lagunaPromoteExpertDown(layerIndex(name), t.numLayers) {
return ""
}
return lagunaQuantizationType(name, shape, base)
case lagunaSharedExpertDownProjection(name):
return lagunaSensitiveType(lagunaPromoteExpertDown(layerIndex(name), t.numLayers), name, shape, base)
case lagunaSharedExpertProjection(name):
return lagunaQuantizationType(name, shape, base)
case lagunaRoutedExpertProjection(name):
return lagunaQuantizationType(name, shape, base)
default:
return ""
}
return GetTensorQuantization(name, shape, quantize)
}
func lagunaIsHFRoutedExpertWeight(name string) bool {
return strings.HasSuffix(name, ".weight") && strings.Contains(name, ".mlp.experts.")
func lagunaFPQuant(quantize string) bool {
return quantize == "nvfp4" || quantize == "mxfp4" || quantize == "mxfp8"
}
func lagunaAttentionProjection(name string) bool {
return strings.Contains(name, ".self_attn.q_proj.weight") ||
strings.Contains(name, ".self_attn.k_proj.weight") ||
strings.Contains(name, ".self_attn.v_proj.weight") ||
strings.Contains(name, ".self_attn.o_proj.weight") ||
strings.Contains(name, ".self_attn.g_proj.weight")
}
func lagunaDenseMLPProjection(name string) bool {
return strings.Contains(name, ".mlp.gate_proj.weight") ||
strings.Contains(name, ".mlp.up_proj.weight") ||
strings.Contains(name, ".mlp.down_proj.weight")
}
func lagunaRoutedExpertProjection(name string) bool {
if !lagunaMLPProjectionWeight(name) {
return false
}
return strings.Contains(name, ".mlp.experts.")
}
func lagunaRoutedExpertDownProjection(name string) bool {
return strings.Contains(name, ".mlp.experts.") && strings.HasSuffix(name, ".down_proj.weight")
}
func lagunaSharedExpertProjection(name string) bool {
if !lagunaMLPProjectionWeight(name) {
return false
}
return strings.Contains(name, ".mlp.shared_expert.")
}
func lagunaSharedExpertDownProjection(name string) bool {
return strings.Contains(name, ".mlp.shared_expert.down_proj.weight")
}
// Laguna XS 2 and 2.1 are sensitive to fully quantizing expert down
// projections. Keep the same cadence for both: use higher precision on the
// input-side layers, final layers, and a sparse cadence early in the residual
// stream. For 4-bit fp quants that higher precision is mxfp8. For mxfp8, the
// shared expert down projections stay at source precision because the tensor
// class is small; routed expert down projections use the cadence to avoid
// pushing the model too close to bf16 size.
func lagunaPromoteExpertDown(layerIdx, numLayers int) bool {
return useMoreBitsWithMiddleEnd(layerIdx, numLayers, numLayers/2-4)
}
func lagunaMLPProjectionWeight(name string) bool {
return strings.HasSuffix(name, ".gate_proj.weight") ||
strings.HasSuffix(name, ".up_proj.weight") ||
strings.HasSuffix(name, ".down_proj.weight")
}
func lagunaQuantizationType(name string, shape []int32, quantize string) string {
q := GetTensorQuantization(name, shape, quantize)
// Laguna's architecture policy decides which sensitive tensors to
// promote. Undo the generic policy's blanket 4-to-8-bit promotion here.
if q != quantize && q == eightBit(quantize) {
return quantize
}
return q
}
func lagunaSensitiveType(promote bool, name string, shape []int32, quantize string) string {
if quantize == "mxfp8" {
return ""
}
if promote {
return GetTensorQuantization(name, shape, quantize)
}
return lagunaQuantizationType(name, shape, quantize)
}

448
x/create/laguna_test.go Normal file
View file

@ -0,0 +1,448 @@
package create
import (
"encoding/json"
"strconv"
"testing"
)
func TestLagunaImportTransformRegistration(t *testing.T) {
inv := Inventory{
Config: sourceModelConfig{Architectures: []string{"LagunaForCausalLM"}},
RawConfig: json.RawMessage(`{"num_hidden_layers":40,"mlp_layer_types":["dense","sparse","sparse"]}`),
}
policy, err := newTensorImportTransform(inv)
if err != nil {
t.Fatalf("newTensorImportTransform() error = %v", err)
}
transform, ok := policy.(lagunaImportTransform)
if !ok {
t.Fatalf("newTensorImportTransform() = %T, want lagunaImportTransform", policy)
}
if !transform.denseMLPLayers[0] || transform.denseMLPLayers[1] {
t.Fatalf("denseMLPLayers = %v, want only layer 0", transform.denseMLPLayers)
}
if transform.numLayers != 40 {
t.Fatalf("numLayers = %d, want 40", transform.numLayers)
}
}
func TestLagunaImportTransformSameRecipeAcrossConfigs(t *testing.T) {
configs := map[string]json.RawMessage{
"laguna xs.2": json.RawMessage(`{
"num_hidden_layers": 40,
"mlp_layer_types": ["dense", "sparse", "sparse"]
}`),
"laguna xs 2.1": json.RawMessage(`{
"num_hidden_layers": 40,
"mlp_only_layers": [0],
"gating_types": ["per_head"]
}`),
}
for name, rawConfig := range configs {
t.Run(name, func(t *testing.T) {
policy, err := newLagunaImportTransform(rawConfig)
if err != nil {
t.Fatal(err)
}
testLagunaQuantizationRecipe(t, policy)
})
}
}
func TestLagunaPlanQuantizesEmbeddingAndHeadAtMXFP8(t *testing.T) {
policy, err := newLagunaImportTransform(json.RawMessage(`{
"num_hidden_layers": 40,
"mlp_only_layers": [0]
}`))
if err != nil {
t.Fatal(err)
}
const (
embedding = "model.embed_tokens.weight"
head = "lm_head.weight"
)
inv := Inventory{Tensors: map[string]SourceTensor{
embedding: {
Name: embedding,
Dtype: "BF16",
Shape: []int32{100352, 2048},
},
head: {
Name: head,
Dtype: "BF16",
Shape: []int32{100352, 2048},
},
}}
specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "nvfp4"}, policy)
if err != nil {
t.Fatal(err)
}
for _, name := range []string{embedding, head} {
if got := quantizeForPlannedTensor(specs, name); got != "mxfp8" {
t.Errorf("planned quantization for %s = %q, want mxfp8", name, got)
}
}
}
func TestLagunaPlanExpertGroupUsesStackedDownProjectionPolicy(t *testing.T) {
policy, err := newLagunaImportTransform(json.RawMessage(`{
"num_hidden_layers": 40,
"mlp_layer_types": ["dense", "sparse"]
}`))
if err != nil {
t.Fatal(err)
}
inv := Inventory{Tensors: map[string]SourceTensor{}}
for _, layer := range []int{1, 5} {
for expert := range 2 {
for _, projection := range []string{"gate_proj", "down_proj"} {
name := "model.layers." + strconv.Itoa(layer) + ".mlp.experts." + strconv.Itoa(expert) + "." + projection + ".weight"
inv.Tensors[name] = SourceTensor{
Name: name,
Dtype: "BF16",
Shape: []int32{2048, 512},
}
}
}
}
specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "mxfp8"}, policy)
if err != nil {
t.Fatal(err)
}
tests := map[string]string{
"model.layers.1.mlp.experts.down_proj.weight": "",
"model.layers.1.mlp.experts.gate_proj.weight": "mxfp8",
"model.layers.5.mlp.experts.down_proj.weight": "mxfp8",
}
for tensor, want := range tests {
if got := quantizeForPlannedTensor(specs, tensor); got != want {
t.Fatalf("planned quantization for %s = %q, want %q", tensor, got, want)
}
}
if _, ok := specByName(specs, "model.layers.1.mlp.experts"); ok {
t.Fatal("mixed layer 1 expert projections should use separate blobs")
}
for _, name := range []string{
"model.layers.1.mlp.experts.down_proj.weight",
"model.layers.1.mlp.experts.gate_proj.weight",
} {
spec, ok := specByName(specs, name)
if !ok || len(spec.Tensors) != 1 {
t.Fatalf("missing homogeneous blob %s; got %v", name, specNames(specs))
}
}
spec, ok := specByName(specs, "model.layers.5.mlp.experts")
if !ok || len(spec.Tensors) != 2 {
t.Fatalf("uniform layer 5 projections should share one blob; got %v", specNames(specs))
}
}
func quantizeForPlannedTensor(specs []BlobSpec, name string) string {
for _, spec := range specs {
for _, ts := range spec.Tensors {
if ts.Name == name {
return ts.Quantize
}
}
}
return "<missing>"
}
func testLagunaQuantizationRecipe(t *testing.T, policy quantizePolicy) {
t.Helper()
tests := []struct {
name string
tensor string
shape []int32
quantize string
want string
}{
{
name: "attention q projection uses requested fp4",
tensor: "model.layers.1.self_attn.q_proj.weight",
shape: []int32{8192, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "attention v projection uses requested fp4 before promotion layer",
tensor: "model.layers.0.self_attn.v_proj.weight",
shape: []int32{1024, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "attention v projection uses requested fp4 on layer 4",
tensor: "model.layers.4.self_attn.v_proj.weight",
shape: []int32{1024, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "attention v projection uses requested fp4 after promotion layer",
tensor: "model.layers.5.self_attn.v_proj.weight",
shape: []int32{1024, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "attention k projection uses requested fp4 on input layer",
tensor: "model.layers.0.self_attn.k_proj.weight",
shape: []int32{1024, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "attention k projection uses requested fp4 past layer 0",
tensor: "model.layers.4.self_attn.k_proj.weight",
shape: []int32{1024, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "attention k projection stays source precision for mxfp8",
tensor: "model.layers.4.self_attn.k_proj.weight",
shape: []int32{1024, 2048},
quantize: "mxfp8",
want: "",
},
{
name: "attention v projection stays source precision for mxfp8",
tensor: "model.layers.4.self_attn.v_proj.weight",
shape: []int32{1024, 2048},
quantize: "mxfp8",
want: "",
},
{
name: "attention q projection stays source precision for mxfp8",
tensor: "model.layers.4.self_attn.q_proj.weight",
shape: []int32{8192, 2048},
quantize: "mxfp8",
want: "",
},
{
name: "attention o projection stays source precision for mxfp8",
tensor: "model.layers.4.self_attn.o_proj.weight",
shape: []int32{2048, 8192},
quantize: "mxfp8",
want: "",
},
{
name: "attention gate projection stays source precision for mxfp8",
tensor: "model.layers.4.self_attn.g_proj.weight",
shape: []int32{64, 2048},
quantize: "mxfp8",
want: "",
},
{
name: "attention gate projection uses requested fp4",
tensor: "model.layers.1.self_attn.g_proj.weight",
shape: []int32{64, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "dense gate projection uses requested fp4",
tensor: "model.layers.0.mlp.gate_proj.weight",
shape: []int32{8192, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "dense down projection uses requested fp4",
tensor: "model.layers.0.mlp.down_proj.weight",
shape: []int32{2048, 8192},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "unsupported dense projection in sparse layer stays source precision",
tensor: "model.layers.1.mlp.down_proj.weight",
shape: []int32{2048, 8192},
quantize: "nvfp4",
want: "",
},
{
name: "routed expert gate uses requested fp4",
tensor: "model.layers.1.mlp.experts.gate_proj.weight",
shape: []int32{256, 512, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "routed expert down uses requested fp4 on cadence layer",
tensor: "model.layers.1.mlp.experts.down_proj.weight",
shape: []int32{256, 2048, 512},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "routed expert down uses requested fp4 on later layer",
tensor: "model.layers.5.mlp.experts.down_proj.weight",
shape: []int32{256, 2048, 512},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "shared expert gate uses requested fp4",
tensor: "model.layers.1.mlp.shared_expert.gate_proj.weight",
shape: []int32{512, 2048},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "shared expert down promotes to mxfp8 on input-side layer",
tensor: "model.layers.1.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "mxfp8",
},
{
name: "shared expert down uses requested fp4 before middle cadence",
tensor: "model.layers.5.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "shared expert down promotes to mxfp8 on first selected middle layer",
tensor: "model.layers.7.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "mxfp8",
},
{
name: "shared expert down promotes to mxfp8 on early middle layer",
tensor: "model.layers.10.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "mxfp8",
},
{
name: "shared expert down promotes to mxfp8 on last selected middle layer",
tensor: "model.layers.13.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "mxfp8",
},
{
name: "shared expert down uses requested fp4 after selected middle layers",
tensor: "model.layers.16.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "shared expert down uses requested fp4 on late middle layer",
tensor: "model.layers.19.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "nvfp4",
},
{
name: "shared expert down promotes to mxfp8 on final layers",
tensor: "model.layers.39.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "nvfp4",
want: "mxfp8",
},
{
name: "shared expert down stays source precision for mxfp8 on selected layer",
tensor: "model.layers.1.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "mxfp8",
want: "",
},
{
name: "shared expert down stays source precision for mxfp8 off selected layers",
tensor: "model.layers.5.mlp.shared_expert.down_proj.weight",
shape: []int32{2048, 512},
quantize: "mxfp8",
want: "",
},
{
name: "routed expert down stays source precision for mxfp8 on selected layer",
tensor: "model.layers.1.mlp.experts.down_proj.weight",
shape: []int32{256, 2048, 512},
quantize: "mxfp8",
want: "",
},
{
name: "per-expert routed down stays source precision for mxfp8 on selected layer",
tensor: "model.layers.1.mlp.experts.0.down_proj.weight",
shape: []int32{2048, 512},
quantize: "mxfp8",
want: "",
},
{
name: "routed expert down uses mxfp8 off selected layers",
tensor: "model.layers.5.mlp.experts.down_proj.weight",
shape: []int32{256, 2048, 512},
quantize: "mxfp8",
want: "mxfp8",
},
{
name: "per-expert routed down uses mxfp8 off selected layers",
tensor: "model.layers.5.mlp.experts.0.down_proj.weight",
shape: []int32{2048, 512},
quantize: "mxfp8",
want: "mxfp8",
},
{
name: "router gate stays source precision",
tensor: "model.layers.1.mlp.gate.weight",
shape: []int32{256, 2048},
quantize: "nvfp4",
want: "",
},
{
name: "embedding promotes to mxfp8 for nvfp4",
tensor: "model.embed_tokens.weight",
shape: []int32{100352, 2048},
quantize: "nvfp4",
want: "mxfp8",
},
{
name: "embedding uses requested mxfp8",
tensor: "model.embed_tokens.weight",
shape: []int32{100352, 2048},
quantize: "mxfp8",
want: "mxfp8",
},
{
name: "lm head promotes to mxfp8 for nvfp4",
tensor: "lm_head.weight",
shape: []int32{100352, 2048},
quantize: "nvfp4",
want: "mxfp8",
},
{
name: "lm head uses requested mxfp8",
tensor: "lm_head.weight",
shape: []int32{100352, 2048},
quantize: "mxfp8",
want: "mxfp8",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := policy.quantizationType(tt.tensor, tt.shape, tt.quantize); got != tt.want {
t.Fatalf("quantizationType(%q, %v, %q) = %q, want %q", tt.tensor, tt.shape, tt.quantize, got, tt.want)
}
})
}
}

43
x/create/plan.go
View file

@ -160,20 +160,20 @@ func planFloat(inv Inventory, quantize string, policy quantizePolicy) ([]BlobSpe
}
for _, gp := range sortedKeys(groups) {
spec, err := planExpertGroup(gp, groups[gp], quantize, policy)
groupSpecs, err := planExpertGroup(gp, groups[gp], quantize, policy)
if err != nil {
return nil, err
}
specs = append(specs, spec)
specs = append(specs, groupSpecs...)
}
return specs, nil
}
// planExpertGroup packs a layer's per-expert weights into one blob: the experts
// of each projection are stacked into a single [experts, out, in] tensor and
// quantized per the policy. Output tensor names keep the source's ".experts."
// path; only the per-expert index is dropped.
func planExpertGroup(groupPrefix string, tensors []SourceTensor, quantize string, policy quantizePolicy) (BlobSpec, error) {
// planExpertGroup stacks each projection's per-expert weights into an
// [experts, out, in] tensor. Uniform projections share one blob; mixed
// precisions use one blob per projection so safetensors quantization metadata
// always describes the entire blob.
func planExpertGroup(groupPrefix string, tensors []SourceTensor, quantize string, policy quantizePolicy) ([]BlobSpec, error) {
type expert struct {
idx int
t SourceTensor
@ -182,12 +182,12 @@ func planExpertGroup(groupPrefix string, tensors []SourceTensor, quantize string
for _, t := range tensors {
idx, proj, err := parseExpertTensor(groupPrefix, t.Name)
if err != nil {
return BlobSpec{}, err
return nil, err
}
byProj[proj] = append(byProj[proj], expert{idx: idx, t: t})
}
spec := BlobSpec{Name: groupPrefix}
var tensorSpecs []TensorSpec
for _, proj := range sortedKeys(byProj) {
experts := byProj[proj]
sort.Slice(experts, func(i, j int) bool { return experts[i].idx < experts[j].idx })
@ -196,7 +196,7 @@ func planExpertGroup(groupPrefix string, tensors []SourceTensor, quantize string
sources := make([]SourceTensor, len(experts))
for i, e := range experts {
if e.t.Dtype != base.Dtype || !slices.Equal(e.t.Shape, base.Shape) {
return BlobSpec{}, fmt.Errorf("expert group %s projection %s has mismatched expert layout (%s %v vs %s %v)",
return nil, fmt.Errorf("expert group %s projection %s has mismatched expert layout (%s %v vs %s %v)",
groupPrefix, proj, base.Dtype, base.Shape, e.t.Dtype, e.t.Shape)
}
sources[i] = e.t
@ -208,7 +208,7 @@ func planExpertGroup(groupPrefix string, tensors []SourceTensor, quantize string
if quantize != "" {
q = policy.quantizationType(stackedName, stackedShape, quantize)
}
spec.Tensors = append(spec.Tensors, TensorSpec{
tensorSpecs = append(tensorSpecs, TensorSpec{
Name: stackedName,
Sources: sources,
Transform: TransformStackExperts,
@ -217,7 +217,26 @@ func planExpertGroup(groupPrefix string, tensors []SourceTensor, quantize string
OutShape: stackedShape,
})
}
return spec, nil
return homogeneousExpertBlobs(groupPrefix, tensorSpecs), nil
}
func homogeneousExpertBlobs(groupPrefix string, tensors []TensorSpec) []BlobSpec {
if len(tensors) == 0 {
return nil
}
quantize := tensors[0].Quantize
for _, tensor := range tensors[1:] {
if tensor.Quantize != quantize {
blobs := make([]BlobSpec, len(tensors))
for i, tensor := range tensors {
blobs[i] = BlobSpec{Name: tensor.Name, Tensors: []TensorSpec{tensor}}
}
return blobs
}
}
return []BlobSpec{{Name: groupPrefix, Tensors: tensors}}
}
// parseExpertTensor splits a per-expert weight name of the form

24
x/create/plan_test.go
View file

@ -105,7 +105,8 @@ func TestPlanFloat(t *testing.T) {
func TestPlanFloatExpertGroup(t *testing.T) {
// A float MoE layer that ships per-expert tensors: two experts, two
// projections. Each projection is stacked into one [experts, out, in]
// tensor in a single packed blob; the routing gate and norm stay plain.
// tensor. Different quantization types use separate homogeneous blobs; the
// routing gate and norm stay plain.
inv := newInventory(sourceModelConfig{}, map[string]string{
"model.layers.0.mlp.experts.0.gate_proj.weight": "BF16",
"model.layers.0.mlp.experts.1.gate_proj.weight": "BF16",
@ -120,17 +121,17 @@ func TestPlanFloatExpertGroup(t *testing.T) {
t.Fatalf("Plan() error = %v", err)
}
group, ok := specByName(specs, "model.layers.0.mlp.experts")
gateBlob, ok := specByName(specs, "model.layers.0.mlp.experts.gate_proj.weight")
if !ok {
t.Fatalf("missing packed expert blob; got %v", specNames(specs))
t.Fatalf("missing stacked gate projection blob; got %v", specNames(specs))
}
if len(group.Tensors) != 2 {
t.Fatalf("packed blob has %d tensors, want 2 (gate_proj, down_proj)", len(group.Tensors))
if len(gateBlob.Tensors) != 1 {
t.Fatalf("gate projection blob has %d tensors, want 1", len(gateBlob.Tensors))
}
gate, ok := inputByOutput(group, "model.layers.0.mlp.experts.gate_proj.weight")
gate, ok := inputByOutput(gateBlob, "model.layers.0.mlp.experts.gate_proj.weight")
if !ok {
t.Fatal("packed blob missing stacked gate_proj")
t.Fatal("gate projection blob missing stacked gate_proj")
}
if gate.Transform != TransformStackExperts || len(gate.Sources) != 2 {
t.Errorf("gate_proj = %+v, want stack of 2 experts", gate)
@ -146,10 +147,17 @@ func TestPlanFloatExpertGroup(t *testing.T) {
t.Errorf("gate_proj quantize = %q, want int4", gate.Quantize)
}
down, _ := inputByOutput(group, "model.layers.0.mlp.experts.down_proj.weight")
downBlob, ok := specByName(specs, "model.layers.0.mlp.experts.down_proj.weight")
if !ok || len(downBlob.Tensors) != 1 {
t.Fatalf("missing homogeneous stacked down projection blob; got %v", specNames(specs))
}
down, _ := inputByOutput(downBlob, "model.layers.0.mlp.experts.down_proj.weight")
if down.Quantize != "int8" {
t.Errorf("down_proj quantize = %q, want int8 (promoted)", down.Quantize)
}
if _, ok := specByName(specs, "model.layers.0.mlp.experts"); ok {
t.Error("mixed expert projections must not share a blob")
}
// Routing gate and norm are not expert tensors; they stay as their own blobs.
for _, name := range []string{"model.layers.0.mlp.gate.weight", "model.layers.0.input_layernorm.weight"} {

17
x/create/quantpolicy.go
View file

@ -38,9 +38,20 @@ func layerIndex(name string) int {
// input grounding and final output refinement), plus every 3rd layer in between
// to limit error accumulation through the residual stream.
func useMoreBits(layerIdx, numLayers int) bool {
return layerIdx < numLayers/8 ||
layerIdx >= 7*numLayers/8 ||
(layerIdx-numLayers/8)%3 == 2
return useMoreBitsWithMiddleEnd(layerIdx, numLayers, 7*numLayers/8)
}
// useMoreBitsWithMiddleEnd applies the standard early/late promotion and
// limits the every-third-layer cadence to layers before middleEnd.
func useMoreBitsWithMiddleEnd(layerIdx, numLayers, middleEnd int) bool {
if layerIdx < 0 || numLayers <= 0 {
return false
}
first := numLayers / 8
last := 7 * numLayers / 8
return layerIdx < first ||
layerIdx >= last ||
(layerIdx >= first && layerIdx < middleEnd && (layerIdx-first)%3 == 2)
}
// eightBit returns the 8-bit quantization type in base's family: int8 for the

475
x/models/laguna/laguna.go
View file

@ -114,10 +114,17 @@ type MLPBlock interface {
Forward(x *mlx.Array, cfg *Config) *mlx.Array
}
type MLPBlockAdder interface {
ForwardAdd(x, residual *mlx.Array, cfg *Config) *mlx.Array
}
type DenseMLP struct {
GateProj nn.LinearLayer
UpProj nn.LinearLayer
DownProj nn.LinearLayer
GateProj nn.LinearLayer
UpProj nn.LinearLayer
DownProj nn.LinearLayer
GateUpProj nn.LinearLayer
GateUpGateScale *mlx.Array
GateUpUpScale *mlx.Array
}
type SparseMoE struct {
@ -125,6 +132,7 @@ type SparseMoE struct {
SwitchMLP *SwitchMLP
SharedExpert *DenseMLP
EScoreCorrectionBias *mlx.Array
RoutedScale *mlx.Array
}
type SwitchMLP struct {
@ -132,6 +140,11 @@ type SwitchMLP struct {
GateWeight *mlx.Array
UpWeight *mlx.Array
DownWeight *mlx.Array
// Source-layout expert weights are stored as [experts, out, in], matching
// the published tensors. GatherMM transposes them lazily so load does not
// materialize huge BF16 expert tensors.
GateUpWeightsSourceLayout bool
DownWeightSourceLayout bool
GateUpWeightQ, GateUpScales, GateUpBiases *mlx.Array
GateWeightQ, GateScales, GateBiases *mlx.Array
@ -143,7 +156,6 @@ type SwitchMLP struct {
GateUpBits, GateBits, UpBits, DownBits int
GateUpGroupSize, GateGroupSize, UpGroupSize, DownGroupSize int
GateUpMode, GateMode, UpMode, DownMode string
UseQuantized, UseFusedGateUp bool
}
type stackedExpertWeights struct {
@ -588,6 +600,55 @@ func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array {
return t
}
func transposeExpertWeightViewForGatherMM(w *mlx.Array) *mlx.Array {
if w == nil || !w.Valid() || w.NumDims() != 3 {
return w
}
return mlx.Transpose(w, 0, 2, 1)
}
func denseExpertWeight(w *stackedExpertWeights) *mlx.Array {
if w == nil {
return nil
}
weight := w.Weight
if w.Scales != nil {
weight = mlx.Dequantize(w.Weight, w.Scales, w.Biases, w.GroupSize, w.Bits, w.Mode)
if w.GlobalScales != nil {
scale := w.GlobalScales
if scale.DType() != weight.DType() {
scale = scale.AsType(weight.DType())
}
if !(scale.NumDims() == 0 || (scale.NumDims() == 1 && scale.Dim(0) == 1)) {
scale = mlx.ExpandDims(mlx.ExpandDims(scale, -1), -1)
}
weight = mlx.Mul(weight, scale)
}
}
return weight
}
func denseExpertWeightForGatherMM(w *stackedExpertWeights) *mlx.Array {
weight := denseExpertWeight(w)
if weight == nil {
return nil
}
return transposeExpertWeightForGatherMM(weight)
}
func denseExpertWeightSupportsSourceLayout(w *stackedExpertWeights) bool {
return w != nil && w.Weight != nil && w.Weight.Valid() && w.Scales == nil && w.Weight.DType() == mlx.DTypeBFloat16
}
func denseExpertWeightsSupportSourceLayout(weights ...*stackedExpertWeights) bool {
for _, w := range weights {
if !denseExpertWeightSupportsSourceLayout(w) {
return false
}
}
return true
}
func canFuseQuantizedGateUp(gateW, upW *stackedExpertWeights) bool {
if gateW == nil || upW == nil || gateW.Scales == nil || upW.Scales == nil {
return false
@ -604,6 +665,95 @@ func canFuseQuantizedGateUp(gateW, upW *stackedExpertWeights) bool {
return gateW.Weight.NumDims() == 3 && upW.Weight.NumDims() == 3
}
func canFuseDenseQuantizedLinears(a, b *nn.QuantizedLinear) bool {
if a == nil || b == nil || a.Scales == nil || b.Scales == nil {
return false
}
if a.QBiases != nil || b.QBiases != nil ||
a.Bias != nil || b.Bias != nil {
return false
}
if !isScalarGlobalScale(a.GlobalScale) || !isScalarGlobalScale(b.GlobalScale) {
return false
}
if a.Bits != b.Bits || a.GroupSize != b.GroupSize || a.Mode != b.Mode {
return false
}
if a.Weight.NumDims() != 2 || b.Weight.NumDims() != 2 ||
a.Scales.NumDims() != 2 || b.Scales.NumDims() != 2 {
return false
}
if a.Weight.Dim(1) != b.Weight.Dim(1) || a.Scales.Dim(1) != b.Scales.Dim(1) {
return false
}
return (a.Mode == "nvfp4" && a.Bits == 4 && a.GroupSize == 16) ||
(a.Mode == "mxfp8" && a.Bits == 8 && a.GroupSize == 32)
}
func isScalarGlobalScale(scale *mlx.Array) bool {
if scale == nil {
return true
}
return scale.NumDims() == 0 || (scale.NumDims() == 1 && scale.Dim(0) == 1)
}
func fuseDenseQuantizedLinears(a, b nn.LinearLayer) nn.LinearLayer {
aq, ok := a.(*nn.QuantizedLinear)
if !ok {
return nil
}
bq, ok := b.(*nn.QuantizedLinear)
if !ok {
return nil
}
if !canFuseDenseQuantizedLinears(aq, bq) {
return nil
}
return &nn.QuantizedLinear{
Weight: fuseExpertStacks(aq.Weight, bq.Weight, 0),
Scales: fuseExpertStacks(aq.Scales, bq.Scales, 0),
GroupSize: aq.GroupSize,
Bits: aq.Bits,
Mode: aq.Mode,
}
}
// fuseDenseGateUp fuses only the quantized weights and scales. Scalar global
// scales from the original projections stay on DenseMLP and are applied after
// the fused output is split back into gate/up halves.
func fuseDenseGateUp(gate, up nn.LinearLayer) nn.LinearLayer {
return fuseDenseQuantizedLinears(gate, up)
}
func linearGlobalScale(l nn.LinearLayer) *mlx.Array {
if q, ok := l.(*nn.QuantizedLinear); ok {
return q.GlobalScale
}
return nil
}
func applyDenseGlobalScale(x, globalScale *mlx.Array) *mlx.Array {
if globalScale == nil {
return x
}
return mlx.Mul(x, globalScale).AsType(x.DType())
}
func splitLastDim(x *mlx.Array, first int32) (*mlx.Array, *mlx.Array) {
dims := x.Dims()
starts := make([]int32, len(dims))
leftStops := make([]int32, len(dims))
rightStarts := make([]int32, len(dims))
rightStops := make([]int32, len(dims))
for i, dim := range dims {
leftStops[i] = int32(dim)
rightStops[i] = int32(dim)
}
leftStops[len(leftStops)-1] = first
rightStarts[len(rightStarts)-1] = first
return mlx.SliceStartStop(x, starts, leftStops), mlx.SliceStartStop(x, rightStarts, rightStops)
}
func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array {
if a == nil || !a.Valid() || b == nil || !b.Valid() {
return nil
@ -613,6 +763,77 @@ func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array {
return out
}
func applyExpertGlobalScale(x, globalScale, idx *mlx.Array) *mlx.Array {
if globalScale == nil {
return x
}
if globalScale.NumDims() == 0 || (globalScale.NumDims() == 1 && globalScale.Dim(0) == 1) {
return mlx.Mul(x, globalScale).AsType(x.DType())
}
scale := mlx.ExpandDims(mlx.ExpandDims(mlx.Take(globalScale, idx, 0), -1), -1)
return mlx.Mul(x, scale).AsType(x.DType())
}
var lagunaSwiGLUGatheredGateScale = mlx.Compile(
"LagunaSwiGLUGatheredGateScale",
func(in ...*mlx.Array) []*mlx.Array {
gate := applyExpertGlobalScale(in[0], in[2], in[3])
return []*mlx.Array{mlx.SwiGLU(gate, in[1])}
},
)
func newLagunaSigmoidTopKRouter(name string, normalize, scaleScores bool) mlx.CompileFunc {
return mlx.Compile(name, func(in ...*mlx.Array) []*mlx.Array {
gates, bias := in[0].AsType(mlx.DTypeFloat32), in[1]
probs, neg := mlx.SigmoidRouter(gates, bias)
inds := mlx.Argpartition(neg, 7, -1)
inds = mlx.SliceStartStop(inds, []int32{0, 0}, []int32{int32(gates.Dim(0)), 8})
scores := mlx.TakeAlongAxis(probs, inds, -1)
if normalize {
scores = mlx.Div(scores, mlx.Sum(scores, -1, true))
}
if scaleScores {
scores = scaleScoresByExpert(scores, inds, in[2])
scores = scaleScoresByExpert(scores, inds, in[3])
}
return []*mlx.Array{scores, inds}
})
}
var (
lagunaSigmoidTopK8 = newLagunaSigmoidTopKRouter(
"LagunaSigmoidTopK8", false, false)
lagunaSigmoidTopK8Normalized = newLagunaSigmoidTopKRouter(
"LagunaSigmoidTopK8Normalized", true, false)
lagunaSigmoidTopK8Scaled = newLagunaSigmoidTopKRouter(
"LagunaSigmoidTopK8Scaled", false, true)
lagunaSigmoidTopK8ScaledNormalized = newLagunaSigmoidTopKRouter(
"LagunaSigmoidTopK8ScaledNormalized", true, true)
)
func lagunaWeightedSum(expert, scores, scale *mlx.Array) *mlx.Array {
weighted := mlx.Mul(expert, mlx.ExpandDims(scores.AsType(expert.DType()), -1))
weighted = mlx.Sum(weighted, 2, false)
return mlx.Mul(weighted, scale.AsType(weighted.DType()))
}
var (
lagunaMoEWeightedSumAdd = mlx.Compile(
"LagunaMoEWeightedSumAdd",
func(in ...*mlx.Array) []*mlx.Array {
y := lagunaWeightedSum(in[0], in[1], in[2]).AsType(in[3].DType())
return []*mlx.Array{mlx.Add(y, in[3])}
},
)
lagunaMoEWeightedSumAdd2 = mlx.Compile(
"LagunaMoEWeightedSumAdd2",
func(in ...*mlx.Array) []*mlx.Array {
y := lagunaWeightedSum(in[0], in[1], in[2]).AsType(in[3].DType())
return []*mlx.Array{mlx.Add(mlx.Add(y, in[3]), in[4])}
},
)
)
func combinedTensorGlobalScale(tensors map[string]*mlx.Array, key string) (*mlx.Array, []string) {
var names []string
weightGlobal := tensors[key+".global_scale"]
@ -715,26 +936,43 @@ func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuanti
if w == nil {
continue
}
consumedKeys := []string{key}
s := tensors[key+"_scale"]
if s == nil {
s = tensors[key+".scale"]
if s != nil {
consumedKeys = append(consumedKeys, key+"_scale")
}
if s == nil {
s = tensors[key+".scale"]
if s != nil {
consumedKeys = append(consumedKeys, key+".scale")
}
}
if s == nil {
freeTensorKeys(tensors, consumedKeys...)
return &stackedExpertWeights{Weight: w}
}
qb := tensors[key+"_qbias"]
if qb != nil {
consumedKeys = append(consumedKeys, key+"_qbias")
}
if qb == nil {
qb = tensors[key+".bias"]
if qb != nil {
consumedKeys = append(consumedKeys, key+".bias")
}
}
globalScale, _ := combinedTensorGlobalScale(tensors, key)
globalScale, globalScaleKeys := combinedTensorGlobalScale(tensors, key)
consumedKeys = append(consumedKeys, globalScaleKeys...)
gs, b, m := model.ResolveLinearQuantParams(cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant, key, w, s)
if useQuantized && supportsGatherQMM(m, b) {
freeTensorKeys(tensors, consumedKeys...)
return &stackedExpertWeights{Weight: w, Scales: s, Biases: qb, GlobalScales: globalScale, Bits: b, GroupSize: gs, Mode: m}
}
deq := mlx.Dequantize(w, s, qb, gs, b, m)
if globalScale != nil {
deq = mlx.Mul(deq, globalScale)
}
freeTensorKeys(tensors, consumedKeys...)
return &stackedExpertWeights{Weight: deq, GlobalScales: globalScale, Bits: b, GroupSize: gs, Mode: m}
}
return nil
@ -744,6 +982,7 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
prefix := resolveWeightPrefix(tensors)
cfg := m.Config
linears := model.NewLinearFactory(tensors, cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant)
routedScale := mlx.FromValue(cfg.MoeRoutedScalingFactor)
m.EmbedTokens = model.MakeEmbeddingLayer(tensors, prefix+"embed_tokens", cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant)
if m.EmbedTokens == nil {
@ -814,7 +1053,10 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
}
if layerUsesMoE(cfg, i) {
moe := &SparseMoE{Gate: linears.Make(layerPrefix + ".mlp.gate")}
moe := &SparseMoE{
Gate: linears.Make(layerPrefix + ".mlp.gate"),
RoutedScale: routedScale,
}
if moe.Gate == nil {
return fmt.Errorf("layer %d: missing moe gate", i)
}
@ -849,13 +1091,8 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
return fmt.Errorf("layer %d: missing moe expert weights", i)
}
sw := &SwitchMLP{}
if gateW.Scales != nil && upW.Scales != nil && downW.Scales != nil {
sw.UseQuantized = true
sw.DownWeightQ, sw.DownScales, sw.DownBiases = downW.Weight, downW.Scales, downW.Biases
sw.DownGlobalScale = downW.GlobalScales
sw.DownBits, sw.DownGroupSize, sw.DownMode = downW.Bits, downW.GroupSize, downW.Mode
if gateW.Scales != nil && upW.Scales != nil {
if canFuseQuantizedGateUp(gateW, upW) {
sw.UseFusedGateUp = true
sw.GateUpWeightQ = fuseExpertStacks(gateW.Weight, upW.Weight, 1)
sw.GateUpScales = fuseExpertStacks(gateW.Scales, upW.Scales, 1)
sw.GateUpBiases = fuseExpertStacks(gateW.Biases, upW.Biases, 1)
@ -869,27 +1106,55 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
sw.UpBits, sw.UpGroupSize, sw.UpMode = upW.Bits, upW.GroupSize, upW.Mode
}
} else {
sw.GateWeight = transposeExpertWeightForGatherMM(gateW.Weight)
sw.UpWeight = transposeExpertWeightForGatherMM(upW.Weight)
sw.DownWeight = transposeExpertWeightForGatherMM(downW.Weight)
sw.GateUpWeight = fuseExpertStacks(sw.GateWeight, sw.UpWeight, 2)
sw.UseFusedGateUp = sw.GateUpWeight != nil
sw.GateUpWeightsSourceLayout = denseExpertWeightsSupportSourceLayout(gateW, upW)
if sw.GateUpWeightsSourceLayout {
sw.GateWeight = denseExpertWeight(gateW)
sw.UpWeight = denseExpertWeight(upW)
// Avoid pre-fusing source-layout BF16 gate/up weights: the
// full-size concatenate can time out during model load.
} else {
sw.GateWeight = denseExpertWeightForGatherMM(gateW)
sw.UpWeight = denseExpertWeightForGatherMM(upW)
sw.GateUpWeight = fuseExpertStacks(sw.GateWeight, sw.UpWeight, 2)
}
}
if downW.Scales != nil {
sw.DownWeightQ, sw.DownScales, sw.DownBiases = downW.Weight, downW.Scales, downW.Biases
sw.DownGlobalScale = downW.GlobalScales
sw.DownBits, sw.DownGroupSize, sw.DownMode = downW.Bits, downW.GroupSize, downW.Mode
} else {
sw.DownWeightSourceLayout = denseExpertWeightSupportsSourceLayout(downW)
if sw.DownWeightSourceLayout {
sw.DownWeight = denseExpertWeight(downW)
} else {
sw.DownWeight = denseExpertWeightForGatherMM(downW)
}
}
moe.SwitchMLP = sw
sharedGate := linears.Make(layerPrefix + ".mlp.shared_expert.gate_proj")
sharedUp := linears.Make(layerPrefix + ".mlp.shared_expert.up_proj")
moe.SharedExpert = &DenseMLP{
GateProj: linears.Make(layerPrefix + ".mlp.shared_expert.gate_proj"),
UpProj: linears.Make(layerPrefix + ".mlp.shared_expert.up_proj"),
DownProj: linears.Make(layerPrefix + ".mlp.shared_expert.down_proj"),
GateProj: sharedGate,
UpProj: sharedUp,
DownProj: linears.Make(layerPrefix + ".mlp.shared_expert.down_proj"),
GateUpProj: fuseDenseGateUp(sharedGate, sharedUp),
GateUpGateScale: linearGlobalScale(sharedGate),
GateUpUpScale: linearGlobalScale(sharedUp),
}
if moe.SharedExpert.GateProj == nil || moe.SharedExpert.UpProj == nil || moe.SharedExpert.DownProj == nil {
return fmt.Errorf("layer %d: missing shared expert weights", i)
}
layer.MLP = moe
} else {
gate := linears.Make(layerPrefix + ".mlp.gate_proj")
up := linears.Make(layerPrefix + ".mlp.up_proj")
mlp := &DenseMLP{
GateProj: linears.Make(layerPrefix + ".mlp.gate_proj"),
UpProj: linears.Make(layerPrefix + ".mlp.up_proj"),
DownProj: linears.Make(layerPrefix + ".mlp.down_proj"),
GateProj: gate,
UpProj: up,
DownProj: linears.Make(layerPrefix + ".mlp.down_proj"),
GateUpProj: fuseDenseGateUp(gate, up),
GateUpGateScale: linearGlobalScale(gate),
GateUpUpScale: linearGlobalScale(up),
}
if mlp.GateProj == nil || mlp.UpProj == nil || mlp.DownProj == nil {
return fmt.Errorf("layer %d: missing dense mlp projections", i)
@ -941,9 +1206,23 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio
}
func (m *DenseMLP) Forward(x *mlx.Array, _ *Config) *mlx.Array {
if m.GateUpProj != nil {
gateUp := m.GateUpProj.Forward(x)
gate, up := splitLastDim(gateUp, int32(gateUp.Dim(len(gateUp.Dims())-1))/2)
gate = applyDenseGlobalScale(gate, m.GateUpGateScale)
up = applyDenseGlobalScale(up, m.GateUpUpScale)
return m.DownProj.Forward(mlx.SwiGLU(gate, up))
}
return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x)))
}
func weightForGatherMM(w *mlx.Array, sourceLayout bool) *mlx.Array {
if sourceLayout {
return transposeExpertWeightViewForGatherMM(w)
}
return w
}
func (s *SwitchMLP) Forward(x *mlx.Array, indices *mlx.Array, cfg *Config) *mlx.Array {
dims := x.Dims()
B, L := int32(dims[0]), int32(dims[1])
@ -964,44 +1243,41 @@ func (s *SwitchMLP) Forward(x *mlx.Array, indices *mlx.Array, cfg *Config) *mlx.
idxFlat = mlx.Reshape(mlx.Take(idxAll, order, 0), n, 1)
}
var gate, up, hidden, down *mlx.Array
if s.UseQuantized {
if s.UseFusedGateUp {
gateUp := mlx.GatherQMM(xFlat, s.GateUpWeightQ, s.GateUpScales, s.GateUpBiases, nil, idxFlat, true, s.GateUpGroupSize, s.GateUpBits, s.GateUpMode, doSort)
guDims := gateUp.Dims()
mid := int32(guDims[len(guDims)-1] / 2)
gate = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, 0}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), mid})
up = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, mid}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), int32(guDims[len(guDims)-1])})
hidden = mlx.SwiGLU(gate, up)
var gate, up, hidden *mlx.Array
switch {
case s.GateUpWeightQ != nil:
gateUp := mlx.GatherQMM(xFlat, s.GateUpWeightQ, s.GateUpScales, s.GateUpBiases, nil, idxFlat, true, s.GateUpGroupSize, s.GateUpBits, s.GateUpMode, doSort)
guDims := gateUp.Dims()
mid := int32(guDims[len(guDims)-1] / 2)
gate = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, 0}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), mid})
up = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, mid}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), int32(guDims[len(guDims)-1])})
hidden = mlx.SwiGLU(gate, up)
case s.GateWeightQ != nil && s.UpWeightQ != nil:
gate = mlx.GatherQMM(xFlat, s.GateWeightQ, s.GateScales, s.GateBiases, nil, idxFlat, true, s.GateGroupSize, s.GateBits, s.GateMode, doSort)
up = mlx.GatherQMM(xFlat, s.UpWeightQ, s.UpScales, s.UpBiases, nil, idxFlat, true, s.UpGroupSize, s.UpBits, s.UpMode, doSort)
if s.GateGlobalScale != nil {
hidden = lagunaSwiGLUGatheredGateScale(gate, up, s.GateGlobalScale, idxFlat)[0]
} else {
gate = mlx.GatherQMM(xFlat, s.GateWeightQ, s.GateScales, s.GateBiases, nil, idxFlat, true, s.GateGroupSize, s.GateBits, s.GateMode, doSort)
if s.GateGlobalScale != nil {
gate = mlx.Mul(gate, mlx.Take(s.GateGlobalScale, idxFlat, 0))
}
up = mlx.GatherQMM(xFlat, s.UpWeightQ, s.UpScales, s.UpBiases, nil, idxFlat, true, s.UpGroupSize, s.UpBits, s.UpMode, doSort)
if s.UpGlobalScale != nil {
up = mlx.Mul(up, mlx.Take(s.UpGlobalScale, idxFlat, 0))
}
hidden = mlx.SwiGLU(gate, up)
}
case s.GateUpWeight != nil:
gateUp := mlx.GatherMM(xFlat, s.GateUpWeight, nil, idxFlat, doSort)
guDims := gateUp.Dims()
mid := int32(guDims[len(guDims)-1] / 2)
gate = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, 0}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), mid})
up = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, mid}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), int32(guDims[len(guDims)-1])})
hidden = mlx.SwiGLU(gate, up)
default:
gate = mlx.GatherMM(xFlat, weightForGatherMM(s.GateWeight, s.GateUpWeightsSourceLayout), nil, idxFlat, doSort)
up = mlx.GatherMM(xFlat, weightForGatherMM(s.UpWeight, s.GateUpWeightsSourceLayout), nil, idxFlat, doSort)
hidden = mlx.SwiGLU(gate, up)
}
var down *mlx.Array
if s.DownWeightQ != nil {
down = mlx.GatherQMM(hidden, s.DownWeightQ, s.DownScales, s.DownBiases, nil, idxFlat, true, s.DownGroupSize, s.DownBits, s.DownMode, doSort)
if s.DownGlobalScale != nil {
down = mlx.Mul(down, mlx.Take(s.DownGlobalScale, idxFlat, 0))
}
} else {
if s.UseFusedGateUp && s.GateUpWeight != nil {
gateUp := mlx.GatherMM(xFlat, s.GateUpWeight, nil, idxFlat, doSort)
guDims := gateUp.Dims()
mid := int32(guDims[len(guDims)-1] / 2)
gate = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, 0}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), mid})
up = mlx.SliceStartStop(gateUp, []int32{0, 0, 0, mid}, []int32{int32(guDims[0]), int32(guDims[1]), int32(guDims[2]), int32(guDims[len(guDims)-1])})
hidden = mlx.SwiGLU(gate, up)
} else {
gate = mlx.GatherMM(xFlat, s.GateWeight, nil, idxFlat, doSort)
up = mlx.GatherMM(xFlat, s.UpWeight, nil, idxFlat, doSort)
hidden = mlx.SwiGLU(gate, up)
}
down = mlx.GatherMM(hidden, s.DownWeight, nil, idxFlat, doSort)
down = mlx.GatherMM(hidden, weightForGatherMM(s.DownWeight, s.DownWeightSourceLayout), nil, idxFlat, doSort)
}
if doSort {
down = mlx.Reshape(mlx.Take(mlx.Squeeze(mlx.Squeeze(down, 2), 1), invOrder, 0), B*L, topK, cfg.HiddenSize)
@ -1011,12 +1287,45 @@ func (s *SwitchMLP) Forward(x *mlx.Array, indices *mlx.Array, cfg *Config) *mlx.
return mlx.Reshape(down, B, L, topK, cfg.HiddenSize)
}
func (m *SparseMoE) route(xFlat *mlx.Array, cfg *Config) (scores, inds *mlx.Array) {
gates := m.Gate.Forward(xFlat).AsType(mlx.DTypeFloat32)
func scaleScoresByExpert(scores, inds, globalScale *mlx.Array) *mlx.Array {
if globalScale == nil {
return scores
}
scale := globalScale
if scale.DType() != scores.DType() {
scale = scale.AsType(scores.DType())
}
if scale.NumDims() == 0 || (scale.NumDims() == 1 && scale.Dim(0) == 1) {
return mlx.Mul(scores, scale)
}
return mlx.Mul(scores, mlx.Take(scale, inds, 0))
}
func (m *SparseMoE) route(xFlat *mlx.Array, cfg *Config) (scores, inds *mlx.Array, scalesFolded bool) {
gates := m.Gate.Forward(xFlat)
var probs, neg *mlx.Array
if m.EScoreCorrectionBias != nil && cfg.NumExpertsPerTok == 8 {
normalize := cfg.NormTopKProb
if m.SwitchMLP != nil && m.SwitchMLP.UpGlobalScale != nil && m.SwitchMLP.DownGlobalScale != nil {
fn := lagunaSigmoidTopK8Scaled
if normalize {
fn = lagunaSigmoidTopK8ScaledNormalized
}
out := fn(gates, m.EScoreCorrectionBias, m.SwitchMLP.UpGlobalScale, m.SwitchMLP.DownGlobalScale)
return out[0], out[1], true
}
fn := lagunaSigmoidTopK8
if normalize {
fn = lagunaSigmoidTopK8Normalized
}
out := fn(gates, m.EScoreCorrectionBias)
return out[0], out[1], false
}
if m.EScoreCorrectionBias != nil {
gates = gates.AsType(mlx.DTypeFloat32)
probs, neg = mlx.SigmoidRouter(gates, m.EScoreCorrectionBias)
} else {
gates = gates.AsType(mlx.DTypeFloat32)
probs = mlx.Sigmoid(gates)
neg = mlx.Neg(probs)
}
@ -1026,37 +1335,65 @@ func (m *SparseMoE) route(xFlat *mlx.Array, cfg *Config) (scores, inds *mlx.Arra
if cfg.NormTopKProb && cfg.NumExpertsPerTok > 1 {
scores = mlx.Div(scores, mlx.Sum(scores, -1, true))
}
return scores, inds
return scores, inds, false
}
func (m *SparseMoE) Forward(x *mlx.Array, cfg *Config) *mlx.Array {
return m.forward(x, nil, cfg)
}
func (m *SparseMoE) ForwardAdd(x, residual *mlx.Array, cfg *Config) *mlx.Array {
return m.forward(x, residual, cfg)
}
func (m *SparseMoE) forward(x, residual *mlx.Array, cfg *Config) *mlx.Array {
dims := x.Dims()
B, L := int32(dims[0]), int32(dims[1])
BL := B * L
shared := m.SharedExpert.Forward(x, cfg)
xFlat := mlx.Reshape(x, BL, cfg.HiddenSize)
scores, inds := m.route(xFlat, cfg)
scores = scores.AsType(x.DType())
scores, inds, scalesFolded := m.route(xFlat, cfg)
if !scalesFolded {
scores = scaleScoresByExpert(scores, inds, m.SwitchMLP.UpGlobalScale)
scores = scaleScoresByExpert(scores, inds, m.SwitchMLP.DownGlobalScale)
}
expertOut := m.SwitchMLP.Forward(x, inds, cfg)
routed := mlx.Sum(mlx.Mul(expertOut, mlx.ExpandDims(mlx.Reshape(scores, B, L, cfg.NumExpertsPerTok), -1)), 2, false)
if cfg.MoeRoutedScalingFactor != 1 {
routed = mlx.MulScalar(routed, cfg.MoeRoutedScalingFactor)
scoreDims := mlx.Reshape(scores, B, L, cfg.NumExpertsPerTok)
if residual != nil {
return moeWeightedSumAdd2(expertOut, scoreDims, m.RoutedScale, shared, residual)
}
return mlx.Add(routed, shared)
return moeWeightedSumAdd(expertOut, scoreDims, m.RoutedScale, shared)
}
func moeWeightedSumAdd(expertOut, scores, scale, shared *mlx.Array) *mlx.Array {
return lagunaMoEWeightedSumAdd(expertOut, scores, scale, shared)[0]
}
func moeWeightedSumAdd2(expertOut, scores, scale, addA, addB *mlx.Array) *mlx.Array {
return lagunaMoEWeightedSumAdd2(expertOut, scores, scale, addA, addB)[0]
}
func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array {
r := l.Attention.Forward(l.InputNorm.Forward(x, cfg.RMSNormEps), b, c, positions, B, L, l, cfg)
xn := l.InputNorm.Forward(x, cfg.RMSNormEps)
r := l.Attention.Forward(xn, b, c, positions, B, L, l, cfg)
h := mlx.Add(x, r)
r = l.MLP.Forward(l.PostAttentionNorm.Forward(h, cfg.RMSNormEps), cfg)
mn := l.PostAttentionNorm.Forward(h, cfg.RMSNormEps)
if mlp, ok := l.MLP.(MLPBlockAdder); ok {
return mlp.ForwardAdd(mn, h, cfg)
}
r = l.MLP.Forward(mn, cfg)
return mlx.Add(h, r)
}
func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array {
dims := b.InputIDs.Dims()
B, L := int32(dims[0]), int32(dims[1])
return m.forward(b, caches, B, L)
}
func (m *Model) forward(b *batch.Batch, caches []cache.Cache, B, L int32) *mlx.Array {
positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets))
h := m.EmbedTokens.Forward(b.InputIDs)
for i, layer := range m.Layers {

311
x/models/laguna/laguna_test.go
View file

@ -2,6 +2,8 @@ package laguna
import (
"math"
"slices"
"strings"
"testing"
"github.com/ollama/ollama/x/mlxrunner/batch"
@ -319,9 +321,6 @@ func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) {
if !ok {
t.Fatalf("layer 1 MLP type = %T, want *SparseMoE", m.Layers[1].MLP)
}
if !moe.SwitchMLP.UseFusedGateUp {
t.Fatal("expected dense SwitchMLP to fuse gate/up expert weights")
}
if moe.SwitchMLP.GateUpWeight == nil {
t.Fatal("expected fused GateUpWeight to be populated")
}
@ -330,6 +329,136 @@ func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) {
}
}
func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) {
skipIfNoMLX(t)
cfg, err := parseConfig([]byte(`{
"model_type": "laguna",
"hidden_size": 8,
"intermediate_size": 12,
"moe_intermediate_size": 4,
"shared_expert_intermediate_size": 4,
"num_hidden_layers": 2,
"num_attention_heads": 2,
"num_attention_heads_per_layer": [2, 2],
"num_key_value_heads": 1,
"head_dim": 4,
"vocab_size": 16,
"max_position_embeddings": 64,
"layer_types": ["full_attention", "sliding_attention"],
"sliding_window": 2,
"mlp_only_layers": [0],
"decoder_sparse_step": 1,
"num_experts": 2,
"num_experts_per_tok": 1,
"norm_topk_prob": false,
"moe_routed_scaling_factor": 2.5,
"gating": "per-head",
"rms_norm_eps": 1e-5
}`))
if err != nil {
t.Fatal(err)
}
tensors := tinyLagunaTensors()
for key, tensor := range tensors {
if strings.Contains(key, ".mlp.experts.") && strings.HasSuffix(key, ".weight") {
tensors[key] = tensor.AsType(mlx.DTypeBFloat16)
}
}
m := &Model{
Config: &cfg,
Layers: []*Layer{
{LayerIdx: 0, IsSliding: false},
{LayerIdx: 1, IsSliding: true},
},
}
if err := m.LoadWeights(tensors); err != nil {
t.Fatalf("LoadWeights failed: %v", err)
}
moe, ok := m.Layers[1].MLP.(*SparseMoE)
if !ok {
t.Fatalf("layer 1 MLP type = %T, want *SparseMoE", m.Layers[1].MLP)
}
if !moe.SwitchMLP.GateUpWeightsSourceLayout || !moe.SwitchMLP.DownWeightSourceLayout {
t.Fatal("expected BF16 dense SwitchMLP to keep source-layout expert weights")
}
if moe.SwitchMLP.GateUpWeight != nil {
t.Fatal("expected BF16 source-layout SwitchMLP to avoid pre-fused gate/up weights")
}
if got, want := moe.SwitchMLP.GateWeight.Dims(), []int{2, 4, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] {
t.Fatalf("GateWeight dims = %v, want %v", got, want)
}
}
func TestTinyLagunaLoadWeightsKeepsMixedExpertPrecision(t *testing.T) {
skipIfNoMLX(t)
cfg := &Config{
HiddenSize: 8,
IntermediateSize: 12,
MoeIntermediateSize: 4,
SharedExpertIntermediate: 4,
NumHiddenLayers: 2,
NumAttentionHeads: 2,
NumAttentionHeadsPerLayer: []int32{2, 2},
NumKeyValueHeads: 1,
HeadDim: 4,
VocabSize: 16,
LayerTypes: []string{"full_attention", "sliding_attention"},
MLPOnlyLayers: []int32{0},
DecoderSparseStep: 1,
NumExperts: 2,
NumExpertsPerTok: 1,
MoeRoutedScalingFactor: 2.5,
RMSNormEps: 1e-5,
QuantGroupSize: 4,
QuantBits: 4,
QuantMode: "affine",
}
tensors := tinyLagunaTensors()
for expert := range 2 {
prefix := "model.layers.1.mlp.experts." + string(rune('0'+expert))
for _, proj := range []string{"gate_proj", "up_proj"} {
key := prefix + "." + proj + ".weight"
weight, scales, biases := mlx.Quantize(tensors[key], cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode)
tensors[key] = weight
tensors[key+"_scale"] = scales
tensors[key+"_qbias"] = biases
}
downKey := prefix + ".down_proj.weight"
tensors[downKey] = tensors[downKey].AsType(mlx.DTypeBFloat16)
}
m := &Model{
Config: cfg,
Layers: []*Layer{
{LayerIdx: 0, IsSliding: false},
{LayerIdx: 1, IsSliding: true},
},
}
if err := m.LoadWeights(tensors); err != nil {
t.Fatalf("LoadWeights failed: %v", err)
}
moe, ok := m.Layers[1].MLP.(*SparseMoE)
if !ok {
t.Fatalf("layer 1 MLP type = %T, want *SparseMoE", m.Layers[1].MLP)
}
hasFusedGateUp := moe.SwitchMLP.GateUpWeightQ != nil && moe.SwitchMLP.GateUpScales != nil
hasSeparateGateUp := moe.SwitchMLP.GateWeightQ != nil && moe.SwitchMLP.GateScales != nil &&
moe.SwitchMLP.UpWeightQ != nil && moe.SwitchMLP.UpScales != nil
if !hasFusedGateUp && !hasSeparateGateUp {
t.Fatal("expected quantized gate/up expert weights")
}
if moe.SwitchMLP.GateUpWeight != nil || moe.SwitchMLP.GateWeight != nil || moe.SwitchMLP.UpWeight != nil {
t.Fatal("quantized gate/up expert weights fell back to dense")
}
if moe.SwitchMLP.DownWeight == nil || moe.SwitchMLP.DownWeightQ != nil || !moe.SwitchMLP.DownWeightSourceLayout {
t.Fatal("expected BF16 down expert weights to retain source layout")
}
}
func TestSparseMoERouteBiasAffectsSelectionNotRoutingWeights(t *testing.T) {
skipIfNoMLX(t)
cfg := &Config{
@ -345,7 +474,10 @@ func TestSparseMoERouteBiasAffectsSelectionNotRoutingWeights(t *testing.T) {
}
xFlat := mlx.FromValues([]float32{1}, 1, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16)
scores, inds := moe.route(xFlat, cfg)
scores, inds, scalesFolded := moe.route(xFlat, cfg)
if scalesFolded {
t.Fatal("route folded scales without expert projection scales")
}
scores = scores.AsType(mlx.DTypeFloat32)
inds = inds.AsType(mlx.DTypeInt32)
mlx.Eval(scores, inds)
@ -368,6 +500,108 @@ func TestSparseMoERouteBiasAffectsSelectionNotRoutingWeights(t *testing.T) {
}
}
func TestLagunaSigmoidTopK8CompiledMatchesEager(t *testing.T) {
skipIfNoMLX(t)
gates := make([]float32, 2*16)
for i := range gates {
gates[i] = float32((i%13)-6) * 0.2
}
bias := make([]float32, 16)
scaleA := make([]float32, 16)
scaleB := make([]float32, 16)
for i := range bias {
bias[i] = float32((i%5)-2) * 0.03
scaleA[i] = 0.5 + float32(i)*0.01
scaleB[i] = 0.75 + float32(i)*0.005
}
gatesArray := mlx.FromValues(gates, 2, 16).AsType(mlx.DTypeBFloat16)
biasArray := mlx.FromValues(bias, 16)
scaleAArray := mlx.FromValues(scaleA, 16)
scaleBArray := mlx.FromValues(scaleB, 16)
got := lagunaSigmoidTopK8ScaledNormalized(gatesArray, biasArray, scaleAArray, scaleBArray)
probs, neg := mlx.SigmoidRouter(gatesArray.AsType(mlx.DTypeFloat32), biasArray)
wantIndices := mlx.Argpartition(neg, 7, -1)
wantIndices = mlx.SliceStartStop(wantIndices, []int32{0, 0}, []int32{2, 8})
wantScores := mlx.TakeAlongAxis(probs, wantIndices, -1)
wantScores = mlx.Div(wantScores, mlx.Sum(wantScores, -1, true))
wantScores = scaleScoresByExpert(wantScores, wantIndices, scaleAArray)
wantScores = scaleScoresByExpert(wantScores, wantIndices, scaleBArray)
gotScores := got[0].AsType(mlx.DTypeFloat32)
gotIndices := got[1].AsType(mlx.DTypeInt32)
wantScores = wantScores.AsType(mlx.DTypeFloat32)
wantIndices = wantIndices.AsType(mlx.DTypeInt32)
mlx.Eval(gotScores, gotIndices, wantScores, wantIndices)
assertFloatSlicesClose(t, gotScores.Floats(), wantScores.Floats(), 1e-6)
if got, want := gotIndices.Ints(), wantIndices.Ints(); !slices.Equal(got, want) {
t.Fatalf("indices = %v, want %v", got, want)
}
}
func TestLagunaSwiGLUGatheredGateScaleCompiledMatchesEager(t *testing.T) {
skipIfNoMLX(t)
gateValues := make([]float32, 2*8*4)
upValues := make([]float32, len(gateValues))
for i := range gateValues {
gateValues[i] = float32((i%17)-8) * 0.04
upValues[i] = float32((i%11)-5) * 0.03
}
scaleValues := make([]float32, 16)
for i := range scaleValues {
scaleValues[i] = 0.5 + float32(i)*0.025
}
indices := mlx.FromValues([]int32{
0, 3, 6, 9, 12, 15, 2, 5,
1, 4, 7, 10, 13, 14, 8, 11,
}, 2, 8)
gate := mlx.FromValues(gateValues, 2, 8, 1, 4).AsType(mlx.DTypeBFloat16)
up := mlx.FromValues(upValues, 2, 8, 1, 4).AsType(mlx.DTypeBFloat16)
scales := mlx.FromValues(scaleValues, 16)
got := lagunaSwiGLUGatheredGateScale(gate, up, scales, indices)[0]
want := mlx.SwiGLU(applyExpertGlobalScale(gate, scales, indices), up)
got = got.AsType(mlx.DTypeFloat32)
want = want.AsType(mlx.DTypeFloat32)
mlx.Eval(got, want)
assertFloatSlicesClose(t, got.Floats(), want.Floats(), 1e-6)
}
func TestLagunaMoEWeightedSumCompiledMatchesEager(t *testing.T) {
skipIfNoMLX(t)
expertValues := make([]float32, 1*2*8*4)
scoreValues := make([]float32, 1*2*8)
for i := range expertValues {
expertValues[i] = float32((i%19)-9) * 0.02
}
for i := range scoreValues {
scoreValues[i] = float32(i+1) / 20
}
addAValues := []float32{0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8}
addBValues := []float32{-0.8, 0.7, -0.6, 0.5, -0.4, 0.3, -0.2, 0.1}
expert := mlx.FromValues(expertValues, 1, 2, 8, 4).AsType(mlx.DTypeBFloat16)
scores := mlx.FromValues(scoreValues, 1, 2, 8)
scale := mlx.FromValue(float32(2.5))
addA := mlx.FromValues(addAValues, 1, 2, 4).AsType(mlx.DTypeBFloat16)
addB := mlx.FromValues(addBValues, 1, 2, 4).AsType(mlx.DTypeBFloat16)
weighted := mlx.Mul(expert, mlx.ExpandDims(scores.AsType(expert.DType()), -1))
weighted = mlx.Mul(mlx.Sum(weighted, 2, false), scale.AsType(expert.DType()))
wantAdd := mlx.Add(weighted.AsType(addA.DType()), addA)
wantAdd2 := mlx.Add(wantAdd, addB)
gotAdd := lagunaMoEWeightedSumAdd(expert, scores, scale, addA)[0]
gotAdd2 := lagunaMoEWeightedSumAdd2(expert, scores, scale, addA, addB)[0]
gotAdd = gotAdd.AsType(mlx.DTypeFloat32)
gotAdd2 = gotAdd2.AsType(mlx.DTypeFloat32)
wantAdd = wantAdd.AsType(mlx.DTypeFloat32)
wantAdd2 = wantAdd2.AsType(mlx.DTypeFloat32)
mlx.Eval(gotAdd, gotAdd2, wantAdd, wantAdd2)
assertFloatSlicesClose(t, gotAdd.Floats(), wantAdd.Floats(), 1e-6)
assertFloatSlicesClose(t, gotAdd2.Floats(), wantAdd2.Floats(), 1e-6)
}
func TestSwitchMLPFusedGateUpMatchesSeparate(t *testing.T) {
skipIfNoMLX(t)
cfg := &Config{HiddenSize: 4, NumExpertsPerTok: 2}
@ -391,9 +625,8 @@ func TestSwitchMLPFusedGateUpMatchesSeparate(t *testing.T) {
DownWeight: makePatternExpertWeight(2, 3, 4, 0.013),
}
fused := &SwitchMLP{
GateUpWeight: fuseExpertStacks(separate.GateWeight, separate.UpWeight, 2),
DownWeight: separate.DownWeight,
UseFusedGateUp: true,
GateUpWeight: fuseExpertStacks(separate.GateWeight, separate.UpWeight, 2),
DownWeight: separate.DownWeight,
}
gotSeparate := separate.Forward(x, indices, cfg)
@ -406,6 +639,70 @@ func TestSwitchMLPFusedGateUpMatchesSeparate(t *testing.T) {
assertFloatSlicesClose(t, gotFusedF32.Floats(), gotSeparateF32.Floats(), 1e-5)
}
func TestSwitchMLPMixedQuantizedGateUpDenseDownMatchesDense(t *testing.T) {
skipIfNoMLX(t)
cfg := &Config{HiddenSize: 32, NumExpertsPerTok: 2}
x := makePatternExpertWeight(1, 2, int(cfg.HiddenSize), 0.013)
indices := mlx.FromValues([]int32{0, 1, 1, 0}, 2, int(cfg.NumExpertsPerTok))
gateWeight := makePatternExpertWeight(2, 32, 32, 0.011)
upWeight := makePatternExpertWeight(2, 32, 32, 0.017)
downWeight := makePatternExpertWeight(2, 32, 32, 0.013)
gateQ, gateScales, gateBiases := mlx.Quantize(gateWeight, 32, 8, "mxfp8")
upQ, upScales, upBiases := mlx.Quantize(upWeight, 32, 8, "mxfp8")
mlx.Eval(gateQ, gateScales, upQ, upScales)
mixed := &SwitchMLP{
GateUpWeightQ: fuseExpertStacks(gateQ, upQ, 1),
GateUpScales: fuseExpertStacks(gateScales, upScales, 1),
GateUpBiases: fuseExpertStacks(gateBiases, upBiases, 1),
GateUpBits: 8,
GateUpGroupSize: 32,
GateUpMode: "mxfp8",
DownWeight: downWeight,
DownWeightSourceLayout: true,
}
dense := &SwitchMLP{
GateWeight: gateWeight,
UpWeight: upWeight,
DownWeight: downWeight,
GateUpWeightsSourceLayout: true,
DownWeightSourceLayout: true,
}
got := mixed.Forward(x, indices, cfg).AsType(mlx.DTypeFloat32)
want := dense.Forward(x, indices, cfg).AsType(mlx.DTypeFloat32)
mlx.Eval(got, want)
assertFloatSlicesClose(t, got.Floats(), want.Floats(), 0.02)
}
func TestDenseExpertWeightForGatherMMDequantizesQuantizedWeight(t *testing.T) {
skipIfNoMLX(t)
weight := makePatternExpertWeight(2, 4, 32, 0.011)
qweight, scales, qbiases := mlx.Quantize(weight, 32, 8, "mxfp8")
mlx.Eval(qweight, scales)
got := denseExpertWeightForGatherMM(&stackedExpertWeights{
Weight: qweight,
Scales: scales,
Biases: qbiases,
GroupSize: 32,
Bits: 8,
Mode: "mxfp8",
})
mlx.Eval(got)
if got == nil {
t.Fatal("denseExpertWeightForGatherMM returned nil")
}
if dims := got.Dims(); len(dims) != 3 || dims[0] != 2 || dims[1] != 32 || dims[2] != 4 {
t.Fatalf("dense expert dims = %v, want [2 32 4]", dims)
}
if got.DType() == mlx.DTypeUint32 {
t.Fatal("dense expert fallback kept packed U32 weight")
}
}
func TestCombinedTensorGlobalScaleIgnoresInputGlobalScale(t *testing.T) {
skipIfNoMLX(t)
tensors := map[string]*mlx.Array{