ollama/x/models/qwen3_5/qwen3_5.go

// Package qwen3_5 provides the Qwen 3.5 text and MoE implementation for MLX.
package qwen3_5

import (
	"encoding/json"
	"fmt"
	"log/slog"
	"math"
	"slices"
	"strings"

	"github.com/ollama/ollama/x/mlxrunner/batch"
	"github.com/ollama/ollama/x/mlxrunner/cache"
	"github.com/ollama/ollama/x/mlxrunner/mlx"
	"github.com/ollama/ollama/x/mlxrunner/model"
	"github.com/ollama/ollama/x/mlxrunner/model/base"
	"github.com/ollama/ollama/x/models/nn"
	"github.com/ollama/ollama/x/tokenizer"
)

func init() {
	base.Register("Qwen3_5ForCausalLM", NewModel)
	base.Register("Qwen3_5ForConditionalGeneration", NewModel)
	base.Register("Qwen3NextForCausalLM", NewModel)
	base.Register("Qwen3NextForConditionalGeneration", NewModel)
}

// RopeParameters carries optional rope metadata embedded under rope_parameters.
type RopeParameters struct {
	Type                string  `json:"type"`
	RopeType            string  `json:"rope_type"`
	RopeTheta           float32 `json:"rope_theta"`
	PartialRotaryFactor float32 `json:"partial_rotary_factor"`
}

// Config holds Qwen 3.5 text config (top-level or nested text_config).
type Config struct {
	ModelType             string   `json:"model_type"`
	HiddenSize            int32    `json:"hidden_size"`
	IntermediateSize      int32    `json:"intermediate_size"`
	NumHiddenLayers       int32    `json:"num_hidden_layers"`
	NumAttentionHeads     int32    `json:"num_attention_heads"`
	NumKeyValueHeads      int32    `json:"num_key_value_heads"`
	HeadDim               int32    `json:"head_dim"`
	RMSNormEps            float32  `json:"rms_norm_eps"`
	VocabSize             int32    `json:"vocab_size"`
	MaxPositionEmbeddings int32    `json:"max_position_embeddings"`
	AttentionBias         bool     `json:"attention_bias"`
	TieWordEmbeddings     bool     `json:"tie_word_embeddings"`
	LayerTypes            []string `json:"layer_types"`

	FullAttentionInterval int32 `json:"full_attention_interval"`

	LinearNumValueHeads  int32 `json:"linear_num_value_heads"`
	LinearNumKeyHeads    int32 `json:"linear_num_key_heads"`
	LinearKeyHeadDim     int32 `json:"linear_key_head_dim"`
	LinearValueHeadDim   int32 `json:"linear_value_head_dim"`
	LinearConvKernelDim  int32 `json:"linear_conv_kernel_dim"`
	DecoderSparseStep    int32 `json:"decoder_sparse_step"`
	SharedExpertGateRank int32 `json:"-"`

	NumExperts                   int32   `json:"num_experts"`
	NumExpertsPerTok             int32   `json:"num_experts_per_tok"`
	SharedExpertIntermediateSize int32   `json:"shared_expert_intermediate_size"`
	MoeIntermediateSize          int32   `json:"moe_intermediate_size"`
	NormTopKProb                 bool    `json:"norm_topk_prob"`
	MLPOnlyLayers                []int32 `json:"mlp_only_layers"`

	RopeTheta           float32         `json:"rope_theta"`
	PartialRotaryFactor float32         `json:"partial_rotary_factor"`
	RopeScaling         map[string]any  `json:"rope_scaling"`
	RopeParameters      *RopeParameters `json:"rope_parameters"`

	// Quantization metadata.
	QuantGroupSize int                               `json:"-"`
	QuantBits      int                               `json:"-"`
	QuantMode      string                            `json:"-"`
	TensorQuant    map[string]*model.TensorQuantInfo `json:"-"`

	// Computed fields.
	Scale   float32 `json:"-"`
	RopeDim int32   `json:"-"`
}

// Model is the Qwen 3.5 model.
type Model struct {
	EmbedTokens nn.EmbeddingLayer
	Layers      []*Layer
	Norm        *nn.RMSNorm
	LMHead      nn.LinearLayer

	tok *tokenizer.Tokenizer
	*Config

	weightPrefix string
}

// Layer is a transformer decoder layer.
type Layer struct {
	InputNorm         *nn.RMSNorm
	PostAttentionNorm *nn.RMSNorm

	IsLinear bool
	FullAttn *FullAttention
	Linear   *GatedDeltaNet
	MLP      MLPBlock
}

// FullAttention is the full-attention branch used every N layers.
type FullAttention struct {
	QProj nn.LinearLayer
	KProj nn.LinearLayer
	VProj nn.LinearLayer
	OProj nn.LinearLayer

	QNorm *nn.RMSNorm
	KNorm *nn.RMSNorm
}

// GatedDeltaNet is the recurrent linear-attention branch.
type GatedDeltaNet struct {
	InProjQKV  nn.LinearLayer
	InProjZ    nn.LinearLayer
	InProjB    nn.LinearLayer
	InProjA    nn.LinearLayer
	InProjQKVZ nn.LinearLayer
	InProjBA   nn.LinearLayer
	OutProj    nn.LinearLayer

	Conv1D     *nn.Conv1d
	NormWeight *mlx.Array
	DtBias     *mlx.Array
	ALog       *mlx.Array
	AExp       *mlx.Array
}

// MLPBlock is the feed-forward interface for dense and MoE blocks.
type MLPBlock interface {
	Forward(x *mlx.Array, cfg *Config) *mlx.Array
}

// DenseMLP is SwiGLU feed-forward.
type DenseMLP struct {
	GateProj nn.LinearLayer
	UpProj   nn.LinearLayer
	DownProj nn.LinearLayer
}

// SparseMoE is Qwen3.5's sparse MoE with shared expert.
type SparseMoE struct {
	Gate             nn.LinearLayer
	SwitchMLP        *SwitchMLP
	SharedExpert     *DenseMLP
	SharedExpertGate nn.LinearLayer
}

// SwitchMLP executes selected expert MLPs. Gate and up are always held
// packed as one gate_up tensor; checkpoints that ship them separately are
// fused at load.
type SwitchMLP struct {
	GateUpWeight *mlx.Array
	DownWeight   *mlx.Array

	GateUpWeightQ, GateUpScales, GateUpBias *mlx.Array
	DownWeightQ, DownScales, DownBiases     *mlx.Array

	GateUpBits      int
	DownBits        int
	GateUpGroupSize int
	DownGroupSize   int
	GateUpMode      string
	DownMode        string

	// Tensor-level nvfp4 scales, [experts, 1, 1], applied to gather outputs.
	GateUpGlobalScale, DownGlobalScale *mlx.Array
}

type stackedExpertWeights struct {
	Weight *mlx.Array
	Scales *mlx.Array
	Biases *mlx.Array

	// Tensor-level nvfp4 scales as [experts, 1, 1], kept only for GatherQMM;
	// dequantizing paths fold them into the weights.
	GlobalScales *mlx.Array

	Bits      int
	GroupSize int
	Mode      string
}

func parseConfig(configData []byte) (Config, error) {
	var rawTop map[string]json.RawMessage
	if err := json.Unmarshal(configData, &rawTop); err != nil {
		return Config{}, fmt.Errorf("parse config envelope: %w", err)
	}

	var cfg Config
	activeRaw := rawTop
	if textRaw, ok := rawTop["text_config"]; ok {
		if err := json.Unmarshal(textRaw, &cfg); err != nil {
			return Config{}, fmt.Errorf("parse text_config: %w", err)
		}
		if err := json.Unmarshal(textRaw, &activeRaw); err != nil {
			return Config{}, fmt.Errorf("parse text_config envelope: %w", err)
		}
	} else {
		if err := json.Unmarshal(configData, &cfg); err != nil {
			return Config{}, fmt.Errorf("parse config: %w", err)
		}
	}

	if cfg.HiddenSize <= 0 {
		return Config{}, fmt.Errorf("invalid hidden_size: %d", cfg.HiddenSize)
	}
	if cfg.NumHiddenLayers <= 0 {
		return Config{}, fmt.Errorf("invalid num_hidden_layers: %d", cfg.NumHiddenLayers)
	}
	if cfg.NumAttentionHeads <= 0 {
		return Config{}, fmt.Errorf("invalid num_attention_heads: %d", cfg.NumAttentionHeads)
	}
	if cfg.NumKeyValueHeads <= 0 {
		cfg.NumKeyValueHeads = cfg.NumAttentionHeads
	}
	if cfg.HeadDim <= 0 {
		if cfg.HiddenSize%cfg.NumAttentionHeads != 0 {
			return Config{}, fmt.Errorf("hidden_size (%d) must be divisible by num_attention_heads (%d)", cfg.HiddenSize, cfg.NumAttentionHeads)
		}
		cfg.HeadDim = cfg.HiddenSize / cfg.NumAttentionHeads
	}
	if cfg.HeadDim <= 0 {
		return Config{}, fmt.Errorf("invalid head_dim: %d", cfg.HeadDim)
	}

	if cfg.RMSNormEps == 0 {
		cfg.RMSNormEps = 1e-6
	}
	if cfg.LinearConvKernelDim <= 0 {
		cfg.LinearConvKernelDim = 4
	}
	if cfg.LinearNumKeyHeads <= 0 || cfg.LinearNumValueHeads <= 0 || cfg.LinearKeyHeadDim <= 0 || cfg.LinearValueHeadDim <= 0 {
		return Config{}, fmt.Errorf("invalid linear attention config (k_heads=%d v_heads=%d k_dim=%d v_dim=%d)",
			cfg.LinearNumKeyHeads, cfg.LinearNumValueHeads, cfg.LinearKeyHeadDim, cfg.LinearValueHeadDim)
	}
	if cfg.LinearNumValueHeads%cfg.LinearNumKeyHeads != 0 {
		return Config{}, fmt.Errorf("linear_num_value_heads (%d) must be divisible by linear_num_key_heads (%d)", cfg.LinearNumValueHeads, cfg.LinearNumKeyHeads)
	}

	if cfg.RopeParameters != nil {
		if cfg.RopeParameters.RopeTheta > 0 {
			cfg.RopeTheta = cfg.RopeParameters.RopeTheta
		}
		if cfg.RopeParameters.PartialRotaryFactor > 0 {
			cfg.PartialRotaryFactor = cfg.RopeParameters.PartialRotaryFactor
		}
	}
	if cfg.RopeTheta == 0 {
		cfg.RopeTheta = 100000.0
	}
	if cfg.PartialRotaryFactor == 0 {
		cfg.PartialRotaryFactor = 0.25
	}
	if cfg.PartialRotaryFactor < 0 {
		cfg.PartialRotaryFactor = 0.25
	}
	ropeDim := int32(float32(cfg.HeadDim) * cfg.PartialRotaryFactor)
	if ropeDim <= 0 {
		ropeDim = cfg.HeadDim
	}
	if ropeDim > cfg.HeadDim {
		ropeDim = cfg.HeadDim
	}
	cfg.RopeDim = ropeDim

	if cfg.FullAttentionInterval <= 0 {
		for i, lt := range cfg.LayerTypes {
			if strings.Contains(strings.ToLower(lt), "full") {
				cfg.FullAttentionInterval = int32(i + 1)
				break
			}
		}
		if cfg.FullAttentionInterval <= 0 {
			cfg.FullAttentionInterval = 4
		}
	}
	if cfg.FullAttentionInterval > cfg.NumHiddenLayers {
		cfg.FullAttentionInterval = cfg.NumHiddenLayers
	}

	if cfg.NumExperts > 0 {
		if cfg.NumExpertsPerTok <= 0 {
			cfg.NumExpertsPerTok = 1
		}
		if cfg.MoeIntermediateSize <= 0 {
			cfg.MoeIntermediateSize = cfg.IntermediateSize
		}
		if cfg.SharedExpertIntermediateSize <= 0 {
			cfg.SharedExpertIntermediateSize = cfg.IntermediateSize
		}
		if _, ok := activeRaw["norm_topk_prob"]; !ok {
			cfg.NormTopKProb = true
		}
		if cfg.DecoderSparseStep <= 0 {
			cfg.DecoderSparseStep = 1
		}
	}

	cfg.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim)))
	return cfg, nil
}

type tensorPathLayout struct {
	containerPrefix string
	modelPrefix     string
}

func (l tensorPathLayout) modelPath(suffix string) string {
	return l.containerPrefix + l.modelPrefix + suffix
}

func resolveTensorPathLayout(tensors map[string]*mlx.Array) tensorPathLayout {
	for _, layout := range []tensorPathLayout{
		{containerPrefix: "", modelPrefix: "model."},
		{containerPrefix: "language_model.", modelPrefix: "model."},
		{containerPrefix: "language_model.", modelPrefix: ""},
		{containerPrefix: "model.language_model.", modelPrefix: "model."},
		{containerPrefix: "model.language_model.", modelPrefix: ""},
	} {
		if tensors[layout.modelPath("embed_tokens.weight")] != nil {
			return layout
		}
	}

	return tensorPathLayout{modelPrefix: "model."}
}

func layerIsLinear(cfg *Config, layer int32) bool {
	if len(cfg.LayerTypes) == int(cfg.NumHiddenLayers) {
		t := strings.ToLower(cfg.LayerTypes[layer])
		return !strings.Contains(t, "full")
	}
	if cfg.FullAttentionInterval <= 0 {
		return true
	}
	return (layer+1)%cfg.FullAttentionInterval != 0
}

func layerUsesMoE(cfg *Config, layer int32) bool {
	if cfg.NumExperts <= 0 {
		return false
	}
	for _, l := range cfg.MLPOnlyLayers {
		if l == layer {
			return false
		}
	}
	if cfg.DecoderSparseStep <= 1 {
		return true
	}
	return (layer+1)%cfg.DecoderSparseStep == 0
}

// NewModel creates a Qwen 3.5 model from a manifest root.
func NewModel(root *model.Root) (base.Model, error) {
	configData, err := root.Manifest.ReadConfig("config.json")
	if err != nil {
		return nil, fmt.Errorf("load config: %w", err)
	}

	cfg, err := parseConfig(configData)
	if err != nil {
		return nil, err
	}

	if qt := root.QuantType(); qt != "" {
		cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams(qt)
		if gs := root.GroupSize(); gs > 0 {
			cfg.QuantGroupSize = gs
		}
	} else {
		cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams("")
	}
	cfg.TensorQuant = root.AllTensorQuant()

	tokData, err := root.Manifest.ReadConfig("tokenizer.json")
	if err != nil {
		return nil, fmt.Errorf("load tokenizer config: %w", err)
	}

	tokConfig := &tokenizer.TokenizerConfig{ConfigJSON: configData}
	if genConfigData, err := root.Manifest.ReadConfig("generation_config.json"); err == nil {
		tokConfig.GenerationConfigJSON = genConfigData
	}
	if tokConfigData, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil {
		tokConfig.TokenizerConfigJSON = tokConfigData
	}

	tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig)
	if err != nil {
		return nil, fmt.Errorf("parse tokenizer: %w", err)
	}

	m := &Model{
		Layers: make([]*Layer, cfg.NumHiddenLayers),
		Config: &cfg,
		tok:    tok,
	}

	for i := range cfg.NumHiddenLayers {
		m.Layers[i] = &Layer{IsLinear: layerIsLinear(&cfg, i)}
	}

	return m, nil
}

func tensorAny(tensors map[string]*mlx.Array, keys ...string) (*mlx.Array, string) {
	for _, k := range keys {
		if v := tensors[k]; v != nil {
			return v, k
		}
	}
	return nil, ""
}

func tensorByBase(tensors map[string]*mlx.Array, base string) (*mlx.Array, string) {
	return tensorAny(tensors, base+".weight", base)
}

func supportsGatherQMM(mode string, bits int) bool {
	switch mode {
	case "affine":
		return bits == 4 || bits == 8
	case "mxfp8":
		return bits == 8
	case "nvfp4", "mxfp4":
		return bits == 4
	default:
		return false
	}
}

func freeTensorKeys(tensors map[string]*mlx.Array, keys ...string) {
	for _, k := range keys {
		if k == "" {
			continue
		}
		if t := tensors[k]; t != nil {
			delete(tensors, k)
		}
	}
}

func stackAndClone(parts []*mlx.Array) *mlx.Array {
	if len(parts) == 0 {
		return nil
	}
	stacked := mlx.Stack(parts, 0)
	cloned := stacked.Clone()
	mlx.Eval(cloned)
	return cloned
}

func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array {
	if w == nil || !w.Valid() || w.NumDims() != 3 {
		return w
	}
	t := mlx.Transpose(w, 0, 2, 1)
	cloned := t.Clone()
	mlx.Eval(cloned)
	return cloned
}

func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array {
	if a == nil || !a.Valid() || b == nil || !b.Valid() {
		return nil
	}
	out := mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone()
	mlx.Eval(out)
	return out
}

// tensorGlobalScale returns key's tensor-level nvfp4 scale; the import
// pipeline stores it as ModelOpt's weight_global_scale reciprocal.
func tensorGlobalScale(tensors map[string]*mlx.Array, key string) *mlx.Array {
	if gs := tensors[key+".global_scale"]; gs != nil {
		return gs
	}
	return tensors[key+".weight.global_scale"]
}

// expandExpertGlobalScale shapes scales as [experts, 1, 1] so Take with
// expert indices broadcasts against gather output rows.
func expandExpertGlobalScale(gs *mlx.Array, numExperts int) *mlx.Array {
	if gs == nil {
		return nil
	}
	if gs.Size() == 1 {
		gs = mlx.Add(mlx.Zeros(mlx.DTypeFloat32, numExperts, 1, 1), mlx.Reshape(gs, 1, 1, 1))
	} else {
		gs = mlx.Reshape(gs, int32(numExperts), 1, 1)
	}
	out := gs.Clone()
	mlx.Eval(out)
	return out
}

func foldExpertGlobalScale(w, gs *mlx.Array) *mlx.Array {
	if gs == nil {
		return w
	}
	return mlx.Mul(w, gs).AsType(w.DType())
}

// sameExpertGlobalScales reports whether two per-expert scale stacks carry
// identical values (nil means absent).
func sameExpertGlobalScales(a, b *mlx.Array) bool {
	if a == nil || b == nil {
		return a == b
	}
	if a.Size() != b.Size() {
		return false
	}
	av, bv := a.Floats(), b.Floats()
	for i := range av {
		if av[i] != bv[i] {
			return false
		}
	}
	return true
}

// fuseGateUpProjections joins gate and up stacks along the output dimension,
// which is exact for quantized stacks: groups run along the input dimension.
func fuseGateUpProjections(gate, up *stackedExpertWeights) *stackedExpertWeights {
	if gate == nil || up == nil {
		return nil
	}
	if gate.Weight == nil || gate.Weight.NumDims() != 3 || up.Weight == nil || up.Weight.NumDims() != 3 {
		return nil
	}
	if gate.Scales != nil && up.Scales != nil &&
		gate.Bits == up.Bits && gate.GroupSize == up.GroupSize && gate.Mode == up.Mode &&
		(gate.Biases == nil) == (up.Biases == nil) &&
		sameExpertGlobalScales(gate.GlobalScales, up.GlobalScales) {
		fused := &stackedExpertWeights{
			Weight:       fuseExpertStacks(gate.Weight, up.Weight, 1),
			Scales:       fuseExpertStacks(gate.Scales, up.Scales, 1),
			GlobalScales: gate.GlobalScales,
			Bits:         gate.Bits,
			GroupSize:    gate.GroupSize,
			Mode:         gate.Mode,
		}
		if gate.Biases != nil {
			fused.Biases = fuseExpertStacks(gate.Biases, up.Biases, 1)
		}
		return fused
	}
	if gate.Scales != nil || up.Scales != nil {
		slog.Warn("dequantizing gate/up experts: quantization parameters differ",
			"gate_mode", gate.Mode, "gate_bits", gate.Bits, "gate_group_size", gate.GroupSize,
			"up_mode", up.Mode, "up_bits", up.Bits, "up_group_size", up.GroupSize)
	}
	gateWeight := gate.Weight
	if gate.Scales != nil {
		gateWeight = mlx.Dequantize(gate.Weight, gate.Scales, gate.Biases, gate.GroupSize, gate.Bits, gate.Mode)
	}
	gateWeight = foldExpertGlobalScale(gateWeight, gate.GlobalScales)
	upWeight := up.Weight
	if up.Scales != nil {
		upWeight = mlx.Dequantize(up.Weight, up.Scales, up.Biases, up.GroupSize, up.Bits, up.Mode)
	}
	upWeight = foldExpertGlobalScale(upWeight, up.GlobalScales)
	return &stackedExpertWeights{
		Weight:    mlx.Concatenate([]*mlx.Array{gateWeight, upWeight}, 1),
		Bits:      gate.Bits,
		GroupSize: gate.GroupSize,
		Mode:      gate.Mode,
	}
}

func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuantized bool, bases ...string) *stackedExpertWeights {
	for _, base := range bases {
		w, key := tensorByBase(tensors, base)
		if w == nil {
			continue
		}

		scales := tensors[key+"_scale"]
		if scales == nil {
			return &stackedExpertWeights{Weight: w}
		}

		qbiases := tensors[key+"_qbias"]
		globalScale := expandExpertGlobalScale(tensorGlobalScale(tensors, key), w.Dim(0))
		groupSize, bits, mode := model.ResolveLinearQuantParams(
			cfg.QuantGroupSize,
			cfg.QuantBits,
			cfg.QuantMode,
			cfg.TensorQuant,
			key,
			w,
			scales,
		)
		if useQuantized && supportsGatherQMM(mode, bits) {
			return &stackedExpertWeights{
				Weight:       w,
				Scales:       scales,
				Biases:       qbiases,
				GlobalScales: globalScale,
				Bits:         bits,
				GroupSize:    groupSize,
				Mode:         mode,
			}
		}

		if useQuantized {
			slog.Warn("dequantizing expert weights: no gather kernel for format", "tensor", key, "mode", mode, "bits", bits)
		}
		return &stackedExpertWeights{
			Weight:    foldExpertGlobalScale(mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode), globalScale),
			Bits:      bits,
			GroupSize: groupSize,
			Mode:      mode,
		}
	}

	return nil
}

// expertProjection is one expert's tensors for a projection; the quant fields
// are zero when the expert ships unquantized.
type expertProjection struct {
	Weight, Scales, Biases, GlobalScale *mlx.Array
	Bits, GroupSize                     int
	Mode                                string
}

func collectPerExpertProjection(tensors map[string]*mlx.Array, cfg *Config, useQuantized bool, layerPrefix, proj string, numExperts int32) *stackedExpertWeights {
	experts := make([]expertProjection, 0, numExperts)
	consumedKeys := make([]string, 0, numExperts*4)

	for e := range numExperts {
		base := fmt.Sprintf("%s.mlp.experts.%d.%s", layerPrefix, e, proj)
		w, key := tensorByBase(tensors, base)
		if w == nil {
			continue
		}
		consumedKeys = append(consumedKeys, key)

		ex := expertProjection{Weight: w}
		if s := tensors[key+"_scale"]; s != nil {
			consumedKeys = append(consumedKeys, key+"_scale")
			ex.Scales = s
			if ex.Biases = tensors[key+"_qbias"]; ex.Biases != nil {
				consumedKeys = append(consumedKeys, key+"_qbias")
			}
			if ex.GlobalScale = tensorGlobalScale(tensors, key); ex.GlobalScale != nil {
				consumedKeys = append(consumedKeys, key+".global_scale", key+".weight.global_scale")
			}
			ex.GroupSize, ex.Bits, ex.Mode = model.ResolveLinearQuantParams(
				cfg.QuantGroupSize,
				cfg.QuantBits,
				cfg.QuantMode,
				cfg.TensorQuant,
				key,
				w,
				s,
			)
		}
		experts = append(experts, ex)
	}
	if len(experts) == 0 {
		return nil
	}

	// The gather path needs every expert quantized in one shared format;
	// otherwise dequantize them all so the stack is uniformly fp.
	first := experts[0]
	quantized := useQuantized
	for _, ex := range experts {
		if ex.Scales == nil || !supportsGatherQMM(ex.Mode, ex.Bits) ||
			ex.Bits != first.Bits || ex.GroupSize != first.GroupSize || ex.Mode != first.Mode ||
			(ex.Biases == nil) != (first.Biases == nil) {
			quantized = false
			break
		}
	}

	parts := make([]*mlx.Array, len(experts))
	out := &stackedExpertWeights{Mode: cfg.QuantMode}
	if quantized {
		out.Bits, out.GroupSize, out.Mode = first.Bits, first.GroupSize, first.Mode
		for i, ex := range experts {
			parts[i] = ex.Weight
		}
		out.Weight = stackAndClone(parts)
		for i, ex := range experts {
			parts[i] = ex.Scales
		}
		out.Scales = stackAndClone(parts)
		if first.Biases != nil {
			for i, ex := range experts {
				parts[i] = ex.Biases
			}
			out.Biases = stackAndClone(parts)
		}
		if slices.ContainsFunc(experts, func(ex expertProjection) bool { return ex.GlobalScale != nil }) {
			for i, ex := range experts {
				if ex.GlobalScale == nil {
					// An expert without a tensor-level scale is unscaled.
					parts[i] = mlx.FromValues([]float32{1}, 1)
				} else {
					parts[i] = mlx.Reshape(ex.GlobalScale.AsType(mlx.DTypeFloat32), 1)
				}
			}
			out.GlobalScales = expandExpertGlobalScale(stackAndClone(parts), len(experts))
		}
	} else {
		if useQuantized && slices.ContainsFunc(experts, func(ex expertProjection) bool { return ex.Scales != nil }) {
			slog.Warn("dequantizing expert weights: unsupported or mixed formats",
				"tensor", fmt.Sprintf("%s.mlp.experts.*.%s", layerPrefix, proj))
		}
		for i, ex := range experts {
			if ex.Scales == nil {
				parts[i] = ex.Weight
				continue
			}
			if out.Bits == 0 {
				out.Bits, out.GroupSize, out.Mode = ex.Bits, ex.GroupSize, ex.Mode
			}
			parts[i] = foldExpertGlobalScale(mlx.Dequantize(ex.Weight, ex.Scales, ex.Biases, ex.GroupSize, ex.Bits, ex.Mode), ex.GlobalScale)
		}
		out.Weight = stackAndClone(parts)
	}
	freeTensorKeys(tensors, consumedKeys...)
	return out
}

// combinedGateUpProjection loads the packed mlp.experts.gate_up_proj stack;
// nil when the checkpoint ships gate/up separately.
func combinedGateUpProjection(tensors map[string]*mlx.Array, cfg *Config, useQuantized bool, layerPrefix string) *stackedExpertWeights {
	gateUp, key := tensorAny(
		tensors,
		layerPrefix+".mlp.experts.gate_up_proj.weight",
		layerPrefix+".mlp.experts.gate_up_proj",
	)
	if gateUp == nil || gateUp.NumDims() != 3 {
		return nil
	}

	scales := tensors[key+"_scale"]
	if scales == nil {
		return &stackedExpertWeights{Weight: gateUp}
	}

	qbiases := tensors[key+"_qbias"]
	globalScale := expandExpertGlobalScale(tensorGlobalScale(tensors, key), gateUp.Dim(0))
	groupSize, bits, mode := model.ResolveLinearQuantParams(
		cfg.QuantGroupSize,
		cfg.QuantBits,
		cfg.QuantMode,
		cfg.TensorQuant,
		key,
		gateUp,
		scales,
	)
	if useQuantized && supportsGatherQMM(mode, bits) {
		return &stackedExpertWeights{
			Weight:       gateUp,
			Scales:       scales,
			Biases:       qbiases,
			GlobalScales: globalScale,
			Bits:         bits,
			GroupSize:    groupSize,
			Mode:         mode,
		}
	}
	if useQuantized {
		slog.Warn("dequantizing expert weights: no gather kernel for format", "tensor", key, "mode", mode, "bits", bits)
	}
	return &stackedExpertWeights{
		Weight:    foldExpertGlobalScale(mlx.Dequantize(gateUp, scales, qbiases, groupSize, bits, mode), globalScale),
		Bits:      bits,
		GroupSize: groupSize,
		Mode:      mode,
	}
}

// splitLastAxisHalves views the two halves of a's last axis.
func splitLastAxisHalves(a *mlx.Array) (lo, hi *mlx.Array) {
	dims := a.Dims()
	nd := len(dims)
	beg := make([]int32, nd)
	end := make([]int32, nd)
	for i, d := range dims {
		end[i] = int32(d)
	}
	mid := int32(dims[nd-1]) / 2
	endLo := append([]int32(nil), end...)
	endLo[nd-1] = mid
	begHi := append([]int32(nil), beg...)
	begHi[nd-1] = mid
	return mlx.SliceStartStop(a, beg, endLo), mlx.SliceStartStop(a, begHi, end)
}

// loadSwitchMLP assembles a layer's routed experts from any supported
// checkpoint layout.
func loadSwitchMLP(tensors map[string]*mlx.Array, cfg *Config, useQuantized bool, layerPrefix string) (*SwitchMLP, error) {
	gateUpW := combinedGateUpProjection(tensors, cfg, useQuantized, layerPrefix)
	downW := loadStackedProjection(tensors, cfg, useQuantized,
		layerPrefix+".mlp.switch_mlp.down_proj",
		layerPrefix+".mlp.experts.down_proj",
	)
	if gateUpW == nil {
		gateW := loadStackedProjection(tensors, cfg, useQuantized,
			layerPrefix+".mlp.switch_mlp.gate_proj",
			layerPrefix+".mlp.experts.gate_proj",
		)
		upW := loadStackedProjection(tensors, cfg, useQuantized,
			layerPrefix+".mlp.switch_mlp.up_proj",
			layerPrefix+".mlp.experts.up_proj",
		)
		if gateW == nil {
			gateW = collectPerExpertProjection(tensors, cfg, useQuantized, layerPrefix, "gate_proj", cfg.NumExperts)
		}
		if upW == nil {
			upW = collectPerExpertProjection(tensors, cfg, useQuantized, layerPrefix, "up_proj", cfg.NumExperts)
		}
		gateUpW = fuseGateUpProjections(gateW, upW)
	}
	if downW == nil {
		downW = collectPerExpertProjection(tensors, cfg, useQuantized, layerPrefix, "down_proj", cfg.NumExperts)
	}
	if gateUpW == nil || downW == nil {
		return nil, fmt.Errorf("missing switch expert weights")
	}

	switchMLP := &SwitchMLP{}
	if gateUpW.Scales != nil {
		switchMLP.GateUpWeightQ = gateUpW.Weight
		switchMLP.GateUpScales = gateUpW.Scales
		switchMLP.GateUpBias = gateUpW.Biases
		switchMLP.GateUpBits = gateUpW.Bits
		switchMLP.GateUpGroupSize = gateUpW.GroupSize
		switchMLP.GateUpMode = gateUpW.Mode
		switchMLP.GateUpGlobalScale = gateUpW.GlobalScales
	} else {
		switchMLP.GateUpWeight = transposeExpertWeightForGatherMM(gateUpW.Weight)
	}
	if downW.Scales != nil {
		switchMLP.DownWeightQ = downW.Weight
		switchMLP.DownScales = downW.Scales
		switchMLP.DownBiases = downW.Biases
		switchMLP.DownBits = downW.Bits
		switchMLP.DownGroupSize = downW.GroupSize
		switchMLP.DownMode = downW.Mode
		switchMLP.DownGlobalScale = downW.GlobalScales
	} else {
		switchMLP.DownWeight = transposeExpertWeightForGatherMM(downW.Weight)
	}
	return switchMLP, nil
}

func sanitizeConvWeight(w *mlx.Array) *mlx.Array {
	if w == nil {
		return nil
	}
	if w.NumDims() == 3 {
		if w.Dim(1) == 1 {
			return mlx.Squeeze(w, 1)
		}
		if w.Dim(2) == 1 {
			return mlx.Squeeze(w, 2)
		}
	}
	return w
}

func shouldShiftNormKey(key string) bool {
	for _, suffix := range []string{
		".input_layernorm.weight",
		".post_attention_layernorm.weight",
		"model.norm.weight",
		".self_attn.q_norm.weight",
		".self_attn.k_norm.weight",
	} {
		if strings.HasSuffix(key, suffix) {
			return true
		}
	}
	return false
}

func maybeShiftNormWeight(key string, w *mlx.Array, shouldShift bool) *mlx.Array {
	if !shouldShift || w == nil || w.NumDims() != 1 || !shouldShiftNormKey(key) {
		return w
	}
	return mlx.AddScalar(w, 1.0)
}

// LoadWeights assigns tensors to model fields.
func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
	layout := resolveTensorPathLayout(tensors)
	m.weightPrefix = layout.containerPrefix
	prefix := m.weightPrefix
	modelPrefix := layout.containerPrefix + layout.modelPrefix
	cfg := m.Config

	linears := model.NewLinearFactory(tensors, cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant)

	shouldShiftNormWeights := false
	mtpKeys := make([]string, 0)
	for name, t := range tensors {
		if strings.Contains(name, "mtp.") {
			shouldShiftNormWeights = true
			mtpKeys = append(mtpKeys, name)
			continue
		}
		if !shouldShiftNormWeights && strings.Contains(name, ".linear_attn.conv1d.weight") && t != nil && t.NumDims() == 3 && t.Dim(2) != 1 {
			shouldShiftNormWeights = true
		}
	}
	if len(mtpKeys) > 0 {
		freeTensorKeys(tensors, mtpKeys...)
	}

	embedTokens := model.MakeEmbeddingLayer(tensors, modelPrefix+"embed_tokens", cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant)
	if embedTokens == nil {
		return fmt.Errorf("missing embedding weight: %sembed_tokens.weight", modelPrefix)
	}
	m.EmbedTokens = embedTokens

	normKey := modelPrefix + "norm.weight"
	normWeight := maybeShiftNormWeight(normKey, tensors[normKey], shouldShiftNormWeights)
	if normWeight == nil {
		return fmt.Errorf("missing final norm weight: %snorm.weight", modelPrefix)
	}
	m.Norm = nn.NewRMSNorm(normWeight, cfg.RMSNormEps)

	if cfg.TieWordEmbeddings {
		m.LMHead = m.EmbedTokens.AsLinear()
	} else if lmHead := linears.Make(prefix + "lm_head"); lmHead != nil {
		m.LMHead = lmHead
	} else if lmHead := linears.Make("lm_head"); lmHead != nil {
		m.LMHead = lmHead
	} else {
		m.LMHead = m.EmbedTokens.AsLinear()
	}

	useQuantizedExperts := supportsGatherQMM(cfg.QuantMode, cfg.QuantBits)
	if !useQuantizedExperts && cfg.TensorQuant != nil {
		for _, tq := range cfg.TensorQuant {
			if tq == nil {
				continue
			}
			_, bits, mode := model.QuantizationParams(tq.QuantType)
			if supportsGatherQMM(mode, bits) {
				useQuantizedExperts = true
				break
			}
		}
	}

	for i := range cfg.NumHiddenLayers {
		layerPrefix := fmt.Sprintf("%slayers.%d", modelPrefix, i)
		layer := &Layer{IsLinear: layerIsLinear(cfg, i)}

		if w := maybeShiftNormWeight(layerPrefix+".input_layernorm.weight", tensors[layerPrefix+".input_layernorm.weight"], shouldShiftNormWeights); w != nil {
			layer.InputNorm = nn.NewRMSNorm(w, cfg.RMSNormEps)
		}
		if w := maybeShiftNormWeight(layerPrefix+".post_attention_layernorm.weight", tensors[layerPrefix+".post_attention_layernorm.weight"], shouldShiftNormWeights); w != nil {
			layer.PostAttentionNorm = nn.NewRMSNorm(w, cfg.RMSNormEps)
		}
		if layer.InputNorm == nil || layer.PostAttentionNorm == nil {
			return fmt.Errorf("layer %d: missing layer norms", i)
		}

		if layer.IsLinear {
			lin := &GatedDeltaNet{}
			lin.InProjQKV = linears.Make(layerPrefix + ".linear_attn.in_proj_qkv")
			lin.InProjZ = linears.Make(layerPrefix + ".linear_attn.in_proj_z")
			lin.InProjB = linears.Make(layerPrefix + ".linear_attn.in_proj_b")
			lin.InProjA = linears.Make(layerPrefix + ".linear_attn.in_proj_a")
			lin.InProjQKVZ = linears.Make(layerPrefix + ".linear_attn.in_proj_qkvz")
			lin.InProjBA = linears.Make(layerPrefix + ".linear_attn.in_proj_ba")
			lin.OutProj = linears.Make(layerPrefix + ".linear_attn.out_proj")

			convWeight := sanitizeConvWeight(tensors[layerPrefix+".linear_attn.conv1d.weight"])
			if convWeight == nil {
				convWeight = sanitizeConvWeight(tensors[layerPrefix+".linear_attn.conv1d"])
			}
			lin.NormWeight, _ = tensorAny(tensors,
				layerPrefix+".linear_attn.norm.weight",
				layerPrefix+".linear_attn.norm",
			)
			lin.DtBias, _ = tensorAny(tensors,
				layerPrefix+".linear_attn.dt_bias",
				layerPrefix+".linear_attn.dt_proj",
			)
			lin.ALog, _ = tensorAny(tensors,
				layerPrefix+".linear_attn.A_log",
				layerPrefix+".linear_attn.a_log",
			)
			if lin.ALog != nil {
				lin.AExp = mlx.Exp(lin.ALog.AsType(mlx.DTypeFloat32))
			}

			hasSplit := lin.InProjQKV != nil && lin.InProjZ != nil && lin.InProjB != nil && lin.InProjA != nil
			hasCombined := lin.InProjQKVZ != nil && lin.InProjBA != nil
			if (!hasSplit && !hasCombined) || lin.OutProj == nil {
				return fmt.Errorf("layer %d: missing linear attention projections", i)
			}
			if convWeight == nil || lin.NormWeight == nil || lin.DtBias == nil || lin.ALog == nil || lin.AExp == nil {
				return fmt.Errorf("layer %d: missing linear attention state tensors", i)
			}
			if convWeight.NumDims() != 2 {
				return fmt.Errorf("layer %d: conv1d weight must be 2D after sanitization, got %dD", i, convWeight.NumDims())
			}
			// MLX grouped conv expects [Cout, K, Cin/groups]; for depthwise
			// conv (groups=C) the [C, K] kernel becomes [C, K, 1].
			lin.Conv1D = nn.NewConv1d(mlx.ExpandDims(convWeight, 2), nil, 1, 0, 1, int32(convWeight.Dim(0)))

			layer.Linear = lin
		} else {
			attn := &FullAttention{}
			attn.QProj = linears.Make(layerPrefix + ".self_attn.q_proj")
			attn.KProj = linears.Make(layerPrefix + ".self_attn.k_proj")
			attn.VProj = linears.Make(layerPrefix + ".self_attn.v_proj")
			attn.OProj = linears.Make(layerPrefix + ".self_attn.o_proj")

			if w := maybeShiftNormWeight(layerPrefix+".self_attn.q_norm.weight", tensors[layerPrefix+".self_attn.q_norm.weight"], shouldShiftNormWeights); w != nil {
				attn.QNorm = nn.NewRMSNorm(w, cfg.RMSNormEps)
			}
			if w := maybeShiftNormWeight(layerPrefix+".self_attn.k_norm.weight", tensors[layerPrefix+".self_attn.k_norm.weight"], shouldShiftNormWeights); w != nil {
				attn.KNorm = nn.NewRMSNorm(w, cfg.RMSNormEps)
			}

			if attn.QProj == nil || attn.KProj == nil || attn.VProj == nil || attn.OProj == nil {
				return fmt.Errorf("layer %d: missing full attention projections", i)
			}
			if attn.QNorm == nil || attn.KNorm == nil {
				return fmt.Errorf("layer %d: missing full attention q/k norms", i)
			}
			layer.FullAttn = attn
		}

		if layerUsesMoE(cfg, i) {
			moe := &SparseMoE{}
			moe.Gate = linears.Make(layerPrefix + ".mlp.gate")
			if moe.Gate == nil {
				return fmt.Errorf("layer %d: missing moe gate", i)
			}

			switchMLP, err := loadSwitchMLP(tensors, cfg, useQuantizedExperts, layerPrefix)
			if err != nil {
				return fmt.Errorf("layer %d: %w", i, err)
			}
			moe.SwitchMLP = switchMLP

			sharedGateProj := linears.Make(layerPrefix + ".mlp.shared_expert.gate_proj")
			sharedUpProj := linears.Make(layerPrefix + ".mlp.shared_expert.up_proj")
			sharedDownProj := linears.Make(layerPrefix + ".mlp.shared_expert.down_proj")
			if sharedGateProj != nil && sharedUpProj != nil && sharedDownProj != nil {
				moe.SharedExpert = &DenseMLP{
					GateProj: sharedGateProj,
					UpProj:   sharedUpProj,
					DownProj: sharedDownProj,
				}
				moe.SharedExpertGate = linears.Make(layerPrefix + ".mlp.shared_expert_gate")
			}

			layer.MLP = moe
		} else {
			mlp := &DenseMLP{
				GateProj: linears.Make(layerPrefix + ".mlp.gate_proj"),
				UpProj:   linears.Make(layerPrefix + ".mlp.up_proj"),
				DownProj: linears.Make(layerPrefix + ".mlp.down_proj"),
			}
			if mlp.GateProj == nil || mlp.UpProj == nil || mlp.DownProj == nil {
				return fmt.Errorf("layer %d: missing dense mlp projections", i)
			}
			layer.MLP = mlp
		}

		m.Layers[i] = layer
	}

	return nil
}

func softplus(x *mlx.Array) *mlx.Array {
	return mlx.Logaddexp(x, mlx.Zeros(x.DType(), x.Dims()...))
}

func splitQKVZBA(mixedQKVZ, mixedBA *mlx.Array, cfg *Config, B, L int32) (q, k, v, z, beta, alpha *mlx.Array) {
	nk := cfg.LinearNumKeyHeads
	nv := cfg.LinearNumValueHeads
	dk := cfg.LinearKeyHeadDim
	dv := cfg.LinearValueHeadDim
	vPerK := nv / nk

	mixedQKVZ = mlx.Reshape(mixedQKVZ, B, L, nk, 2*dk+2*vPerK*dv)
	q = mlx.SliceStartStop(mixedQKVZ, []int32{0, 0, 0, 0}, []int32{B, L, nk, dk})
	k = mlx.SliceStartStop(mixedQKVZ, []int32{0, 0, 0, dk}, []int32{B, L, nk, 2 * dk})
	v = mlx.SliceStartStop(mixedQKVZ, []int32{0, 0, 0, 2 * dk}, []int32{B, L, nk, 2*dk + vPerK*dv})
	z = mlx.SliceStartStop(mixedQKVZ, []int32{0, 0, 0, 2*dk + vPerK*dv}, []int32{B, L, nk, 2*dk + 2*vPerK*dv})

	v = mlx.Reshape(v, B, L, nv, dv)
	z = mlx.Reshape(z, B, L, nv, dv)

	mixedBA = mlx.Reshape(mixedBA, B, L, nk, 2*vPerK)
	beta = mlx.SliceStartStop(mixedBA, []int32{0, 0, 0, 0}, []int32{B, L, nk, vPerK})
	alpha = mlx.SliceStartStop(mixedBA, []int32{0, 0, 0, vPerK}, []int32{B, L, nk, 2 * vPerK})
	beta = mlx.Reshape(beta, B, L, nv)
	alpha = mlx.Reshape(alpha, B, L, nv)

	return q, k, v, z, beta, alpha
}

func (a *FullAttention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array {
	qg := a.QProj.Forward(x)
	qg = mlx.Reshape(qg, B, L, cfg.NumAttentionHeads, cfg.HeadDim*2)
	q := mlx.SliceStartStop(qg, []int32{0, 0, 0, 0}, []int32{B, L, cfg.NumAttentionHeads, cfg.HeadDim})
	gate := mlx.SliceStartStop(qg, []int32{0, 0, 0, cfg.HeadDim}, []int32{B, L, cfg.NumAttentionHeads, cfg.HeadDim * 2})
	gate = mlx.Reshape(gate, B, L, cfg.NumAttentionHeads*cfg.HeadDim)

	k := a.KProj.Forward(x)
	v := a.VProj.Forward(x)
	k = mlx.Reshape(k, B, L, cfg.NumKeyValueHeads, cfg.HeadDim)
	v = mlx.Reshape(v, B, L, cfg.NumKeyValueHeads, cfg.HeadDim)

	q = a.QNorm.Forward(q, cfg.RMSNormEps)
	k = a.KNorm.Forward(k, cfg.RMSNormEps)

	q = mlx.Transpose(q, 0, 2, 1, 3)
	k = mlx.Transpose(k, 0, 2, 1, 3)
	v = mlx.Transpose(v, 0, 2, 1, 3)

	q = mlx.RoPEWithBase(q, int(cfg.RopeDim), false, cfg.RopeTheta, 1.0, positions)
	k = mlx.RoPEWithBase(k, int(cfg.RopeDim), false, cfg.RopeTheta, 1.0, positions)

	var kv nn.SDPAOption
	if c != nil {
		history := c.(cache.Attention).Update(b, k, v)
		kv = nn.WithKVHistory(history)
	} else {
		kv = nn.WithKV(k, v, b.SeqQueryLens)
	}
	out := nn.ScaledDotProductAttention(b, q, cfg.Scale, kv, nn.WithMask(nn.CausalMask()))
	out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim)
	gateSigmoid := mlx.Sigmoid(gate)
	out = mlx.Mul(out, gateSigmoid)
	out = a.OProj.Forward(out)
	return out
}

func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B, L int32, cfg *Config) *mlx.Array {
	var qkv, z, beta, alpha *mlx.Array
	useSplitProj := g.InProjQKV != nil && g.InProjZ != nil && g.InProjB != nil && g.InProjA != nil
	if useSplitProj {
		qkv = g.InProjQKV.Forward(x)
		z = g.InProjZ.Forward(x)
		z = mlx.Reshape(z, B, L, cfg.LinearNumValueHeads, cfg.LinearValueHeadDim)
		beta = g.InProjB.Forward(x)
		alpha = g.InProjA.Forward(x)
	} else {
		mixedQKVZ := g.InProjQKVZ.Forward(x)
		mixedBA := g.InProjBA.Forward(x)
		var q, k, v *mlx.Array
		q, k, v, z, beta, alpha = splitQKVZBA(mixedQKVZ, mixedBA, cfg, B, L)
		qkv = mlx.Concatenate([]*mlx.Array{
			mlx.Reshape(q, B, L, cfg.LinearNumKeyHeads*cfg.LinearKeyHeadDim),
			mlx.Reshape(k, B, L, cfg.LinearNumKeyHeads*cfg.LinearKeyHeadDim),
			mlx.Reshape(v, B, L, cfg.LinearNumValueHeads*cfg.LinearValueHeadDim),
		}, -1)
	}
	convTail := cfg.LinearConvKernelDim - 1
	var rc *cache.RecurrentCache
	opts := make([]nn.RecurrentOption, 0, 2)
	if typed, ok := c.(*cache.RecurrentCache); ok {
		rc = typed
		opts = append(opts, nn.WithRecurrentHistory(rc.Get(b, x.DType())))
		// When the cache has scheduled per-token snapshots, segment the
		// recurrent kernels at the interior offsets so each boundary state
		// can be captured.
		if splits := rc.SnapshotSplits(int(L)); len(splits) > 0 {
			opts = append(opts, nn.WithSnapshotSplits(splits))
		}
	} else {
		opts = append(opts, nn.WithRecurrentState(
			mlx.Zeros(x.DType(), int(B), int(convTail), qkv.Dim(2)),
			mlx.Zeros(mlx.DTypeFloat32, int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)),
		))
	}

	convOut, convStates := nn.CausalConv1D(b, qkv, g.Conv1D, int(convTail), opts...)
	convOut = mlx.SiLU(convOut)

	keyDim := cfg.LinearNumKeyHeads * cfg.LinearKeyHeadDim
	valueDim := cfg.LinearNumValueHeads * cfg.LinearValueHeadDim
	q := mlx.SliceStartStop(convOut, []int32{0, 0, 0}, []int32{B, L, keyDim})
	k := mlx.SliceStartStop(convOut, []int32{0, 0, keyDim}, []int32{B, L, 2 * keyDim})
	v := mlx.SliceStartStop(convOut, []int32{0, 0, 2 * keyDim}, []int32{B, L, 2*keyDim + valueDim})
	q = mlx.Reshape(q, B, L, cfg.LinearNumKeyHeads, cfg.LinearKeyHeadDim)
	k = mlx.Reshape(k, B, L, cfg.LinearNumKeyHeads, cfg.LinearKeyHeadDim)
	v = mlx.Reshape(v, B, L, cfg.LinearNumValueHeads, cfg.LinearValueHeadDim)
	invScale := float32(1.0 / math.Sqrt(float64(cfg.LinearKeyHeadDim)))
	q = mlx.MulScalar(mlx.RMSNormFn(q, nil, 1e-6), invScale*invScale)
	k = mlx.MulScalar(mlx.RMSNormFn(k, nil, 1e-6), invScale)

	gDecay := softplus(mlx.Add(alpha, g.DtBias))
	gDecay = mlx.Mul(gDecay, g.AExp)
	gDecay = mlx.Exp(mlx.MulScalar(gDecay, -1))
	gDecay = gDecay.AsType(alpha.DType())

	betaGate := mlx.Sigmoid(beta)

	out, deltaStates := nn.GatedDelta(b, q, k, v, gDecay, betaGate, opts...)
	outDType := out.DType()
	out = mlx.RMSNormFn(out, g.NormWeight, cfg.RMSNormEps)
	out = mlx.Mul(out.AsType(mlx.DTypeFloat32), mlx.SiLU(z.AsType(mlx.DTypeFloat32))).AsType(outDType)
	out = mlx.Reshape(out, B, L, valueDim)
	out = g.OutProj.Forward(out)
	if rc != nil {
		rc.Put(b, convStates, deltaStates)
	}
	return out
}

func (m *DenseMLP) Forward(x *mlx.Array, _ *Config) *mlx.Array {
	return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x)))
}

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])
	topK := cfg.NumExpertsPerTok

	xExpanded := mlx.ExpandDims(mlx.ExpandDims(x, -2), -2)
	xFlat := mlx.Reshape(xExpanded, B*L, 1, 1, cfg.HiddenSize)
	idxFlat := mlx.Reshape(indices, B*L, topK)

	doSort := B*L >= 64
	var invOrder *mlx.Array
	n := B * L * topK

	if doSort {
		idxAll := mlx.Flatten(idxFlat)
		order := mlx.Argsort(idxAll, 0)
		invOrder = mlx.Argsort(order, 0)
		xFlat = mlx.ExpandDims(mlx.Take(mlx.Squeeze(xFlat, 1), mlx.FloorDivideScalar(order, topK), 0), 1)
		idxFlat = mlx.Reshape(mlx.Take(idxAll, order, 0), n, 1)
	}

	var gateUp, down *mlx.Array
	if s.GateUpWeightQ != nil {
		gateUp = mlx.GatherQMM(xFlat, s.GateUpWeightQ, s.GateUpScales, s.GateUpBias,
			nil, idxFlat, true, s.GateUpGroupSize, s.GateUpBits, s.GateUpMode, doSort)
	} else {
		gateUp = mlx.GatherMM(xFlat, s.GateUpWeight, nil, idxFlat, doSort)
	}
	if s.GateUpGlobalScale != nil {
		gateUp = mlx.Mul(gateUp, mlx.Take(s.GateUpGlobalScale, idxFlat, 0)).AsType(xFlat.DType())
	}
	gate, up := splitLastAxisHalves(gateUp)
	hidden := mlx.SwiGLU(gate, up)
	if s.DownWeightQ != nil {
		down = mlx.GatherQMM(hidden, s.DownWeightQ, s.DownScales, s.DownBiases,
			nil, idxFlat, true, s.DownGroupSize, s.DownBits, s.DownMode, doSort)
	} else {
		down = mlx.GatherMM(hidden, s.DownWeight, nil, idxFlat, doSort)
	}
	if s.DownGlobalScale != nil {
		down = mlx.Mul(down, mlx.Take(s.DownGlobalScale, idxFlat, 0)).AsType(hidden.DType())
	}

	if doSort {
		down = mlx.Reshape(mlx.Take(mlx.Squeeze(mlx.Squeeze(down, 2), 1), invOrder, 0), B*L, topK, cfg.HiddenSize)
	} else {
		down = mlx.Squeeze(down, 2)
	}

	return mlx.Reshape(down, B, L, topK, cfg.HiddenSize)
}

func (m *SparseMoE) Forward(x *mlx.Array, cfg *Config) *mlx.Array {
	dims := x.Dims()
	B, L := int32(dims[0]), int32(dims[1])

	probs := mlx.SoftmaxAxis(m.Gate.Forward(x), -1, true)
	neg := mlx.Neg(probs)
	inds := mlx.Argpartition(neg, int(cfg.NumExpertsPerTok)-1, -1)
	shape := inds.Dims()
	inds = mlx.SliceStartStop(inds, []int32{0, 0, 0}, []int32{int32(shape[0]), int32(shape[1]), cfg.NumExpertsPerTok})

	scores := mlx.TakeAlongAxis(probs, inds, -1)
	if cfg.NormTopKProb && cfg.NumExpertsPerTok > 1 {
		sumScores := mlx.Sum(scores, -1, true)
		scores = mlx.Div(scores, sumScores)
	}

	expertOut := m.SwitchMLP.Forward(x, inds, cfg)
	y := mlx.Sum(mlx.Mul(expertOut, mlx.ExpandDims(scores, -1)), 2, false)

	if m.SharedExpert != nil {
		shared := m.SharedExpert.Forward(x, cfg)
		if m.SharedExpertGate != nil {
			shared = mlx.Mul(shared, mlx.Sigmoid(m.SharedExpertGate.Forward(x)))
		}
		y = mlx.Add(y, shared)
	}

	return mlx.Reshape(y, B, L, cfg.HiddenSize)
}

func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array {
	var r *mlx.Array
	normed := l.InputNorm.Forward(x, cfg.RMSNormEps)
	if l.IsLinear {
		r = l.Linear.Forward(normed, b, c, B, L, cfg)
	} else {
		r = l.FullAttn.Forward(normed, b, c, positions, B, L, cfg)
	}
	h := mlx.Add(x, r)
	r = l.MLP.Forward(l.PostAttentionNorm.Forward(h, cfg.RMSNormEps), 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])
	positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets))

	h := m.EmbedTokens.Forward(b.InputIDs)
	for i, layer := range m.Layers {
		var c cache.Cache
		if caches != nil && i < len(caches) {
			c = caches[i]
		}
		h = layer.Forward(h, b, c, positions, B, L, m.Config)
	}
	out := m.Norm.Forward(h, m.RMSNormEps)
	return out
}

func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
	return m.LMHead.Forward(x)
}

func (m *Model) NumLayers() int {
	return len(m.Layers)
}

func (m *Model) MaxContextLength() int {
	return int(m.MaxPositionEmbeddings)
}

func (m *Model) Tokenizer() *tokenizer.Tokenizer {
	return m.tok
}

func (m *Model) NewCaches() []cache.Cache {
	caches := make([]cache.Cache, len(m.Layers))
	convTail := m.LinearConvKernelDim - 1
	convDim := 2*m.LinearNumKeyHeads*m.LinearKeyHeadDim + m.LinearNumValueHeads*m.LinearValueHeadDim
	for i, layer := range m.Layers {
		if layer.IsLinear {
			caches[i] = cache.NewRecurrentCache(convTail, convDim, m.LinearNumValueHeads, m.LinearValueHeadDim, m.LinearKeyHeadDim)
		} else {
			caches[i] = cache.NewKVCache()
		}
	}
	return caches
}