From 64ee2f9847ccaedd8f05a139c30f086e9e0abe73 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Fri, 24 Jul 2026 18:24:53 -0700 Subject: [PATCH] model: add Laguna MLX support (#17237) * 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 --- x/create/blockfp8.go | 24 +- x/create/client/create.go | 25 +- x/create/client/create_test.go | 46 ++++ x/create/laguna.go | 160 ++++++++++- x/create/laguna_test.go | 448 +++++++++++++++++++++++++++++++ x/create/plan.go | 43 ++- x/create/plan_test.go | 24 +- x/create/quantpolicy.go | 17 +- x/models/laguna/laguna.go | 475 ++++++++++++++++++++++++++++----- x/models/laguna/laguna_test.go | 311 ++++++++++++++++++++- 10 files changed, 1449 insertions(+), 124 deletions(-) create mode 100644 x/create/laguna_test.go diff --git a/x/create/blockfp8.go b/x/create/blockfp8.go index 359423e9e..1a9c041c2 100644 --- a/x/create/blockfp8.go +++ b/x/create/blockfp8.go @@ -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 diff --git a/x/create/client/create.go b/x/create/client/create.go index b25ceb952..c4f61755d 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -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" diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index e86ec6a80..e7e5d8b82 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -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) + } + }) + } +} diff --git a/x/create/laguna.go b/x/create/laguna.go index 0e5b4de61..6cfcb6520 100644 --- a/x/create/laguna.go +++ b/x/create/laguna.go @@ -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) } diff --git a/x/create/laguna_test.go b/x/create/laguna_test.go new file mode 100644 index 000000000..c78f9f13d --- /dev/null +++ b/x/create/laguna_test.go @@ -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 "" +} + +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) + } + }) + } +} diff --git a/x/create/plan.go b/x/create/plan.go index d61dcb394..d754859aa 100644 --- a/x/create/plan.go +++ b/x/create/plan.go @@ -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 diff --git a/x/create/plan_test.go b/x/create/plan_test.go index ccc2cab2d..72ad33cf8 100644 --- a/x/create/plan_test.go +++ b/x/create/plan_test.go @@ -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"} { diff --git a/x/create/quantpolicy.go b/x/create/quantpolicy.go index 2fc000663..8b53aa104 100644 --- a/x/create/quantpolicy.go +++ b/x/create/quantpolicy.go @@ -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 diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go index e78edb719..224546d88 100644 --- a/x/models/laguna/laguna.go +++ b/x/models/laguna/laguna.go @@ -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 { diff --git a/x/models/laguna/laguna_test.go b/x/models/laguna/laguna_test.go index ad81d1ebd..8560b1a32 100644 --- a/x/models/laguna/laguna_test.go +++ b/x/models/laguna/laguna_test.go @@ -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{