diff --git a/x/create/client/create.go b/x/create/client/create.go index e050cc35f..ae35caa20 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -568,11 +568,11 @@ func chatTemplateHasThinkingSupport(chatTemplate string) bool { } func alwaysSupportsThinking(architectures []string, modelType string) bool { - if isQwen35Family(modelType) { + if isQwen35Family(modelType) || isQwen4Family(modelType) { return true } for _, arch := range architectures { - if isQwen35Family(arch) { + if isQwen35Family(arch) || isQwen4Family(arch) { return true } } @@ -584,6 +584,12 @@ func isQwen35Family(s string) bool { return strings.Contains(s, "qwen3_5") || strings.Contains(s, "qwen3next") } +func isQwen4Family(s string) bool { + s = strings.ToLower(s) + return strings.Contains(s, "qwen4exp") || + strings.Contains(s, "qwen4_exp") +} + func qwen35RendererName(modelDir string) string { template := readChatTemplate(modelDir) if strings.Contains(template, "resolved_reasoning_effort") && @@ -676,6 +682,8 @@ func parserNameForIdentifier(modelDir, s string) string { return "deepseek3" case strings.Contains(s, "gemma4"): return "gemma4" + case isQwen4Family(s): + return "qwen3.5" case isQwen35Family(s): return "qwen3.5" case strings.Contains(s, "qwen3"): @@ -739,6 +747,8 @@ func rendererNameForIdentifier(modelDir, s string) string { return "glm-4.7" case strings.Contains(s, "deepseek"): return "deepseek3" + case isQwen4Family(s): + return "qwen3.8" case isQwen35Family(s): return qwen35RendererName(modelDir) case strings.Contains(s, "qwen3"): diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 678fe37b7..607f6b7d4 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -642,6 +642,11 @@ func TestDetectCapabilities(t *testing.T) { configJSON: `{"architectures": ["Qwen3NextForCausalLM"]}`, want: modelCapabilities{thinking: true}, }, + { + name: "qwen4 always thinks", + configJSON: `{"architectures": ["Qwen4ExpForConditionalGeneration"], "model_type": "qwen4_exp"}`, + want: modelCapabilities{vision: false, thinking: true}, + }, { name: "vision config", configJSON: `{"architectures": ["Gemma4ForConditionalGeneration"], "vision_config": {}}`, @@ -786,6 +791,11 @@ func TestGetParserName(t *testing.T) { configJSON: `{"architectures": ["Qwen3_5ForConditionalGeneration"]}`, want: "qwen3.5", }, + { + name: "qwen4 model", + configJSON: `{"architectures": ["Qwen4ExpForConditionalGeneration"], "model_type": "qwen4_exp"}`, + want: "qwen3.5", + }, { name: "deepseek model", configJSON: `{"architectures": ["DeepseekV3ForCausalLM"]}`, @@ -868,6 +878,11 @@ func TestGetRendererName(t *testing.T) { configJSON: `{"architectures": ["Qwen3_5ForConditionalGeneration"]}`, want: "qwen3.5", }, + { + name: "qwen4 model", + configJSON: `{"architectures": ["Qwen4ExpForConditionalGeneration"], "model_type": "qwen4_exp"}`, + want: "qwen3.8", + }, { name: "qwen3.8 embedded template", configJSON: `{"architectures": ["Qwen3_5ForConditionalGeneration"]}`, diff --git a/x/create/create.go b/x/create/create.go index 1c6603c4a..84d17dd34 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -483,6 +483,7 @@ var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ "Qwen3_5MoeForConditionalGeneration": newQwen35ImportTransform, "Qwen3NextMoeForCausalLM": newQwen35ImportTransform, "Qwen3NextMoeForConditionalGeneration": newQwen35ImportTransform, + "Qwen4ExpForConditionalGeneration": newQwen4ExpImportTransform, "Gemma4ForCausalLM": newGemma4ImportTransform, "Gemma4ForConditionalGeneration": newGemma4ImportTransform, "Gemma4UnifiedForCausalLM": newGemma4ImportTransform, diff --git a/x/create/qwen4_exp.go b/x/create/qwen4_exp.go new file mode 100644 index 000000000..4113332fc --- /dev/null +++ b/x/create/qwen4_exp.go @@ -0,0 +1,37 @@ +package create + +import ( + "encoding/json" + "strings" +) + +type qwen4ExpImportTransform struct{} + +func newQwen4ExpImportTransform(_ json.RawMessage) (quantizePolicy, error) { + return qwen4ExpImportTransform{}, nil +} + +func (qwen4ExpImportTransform) quantizationType(name string, shape []int32, quantize string) string { + // PLE is a hashed n-gram feature table rather than the model's token + // embedding. It dominates the checkpoint size, and its lookup path supports + // row-wise dequantization, so quantize its shards at the requested type. + if qwen4ExpIsPLEEmbeddingShard(name) { + quantize = normalizeQuantType(quantize) + if len(shape) == 2 && isAligned(shape, quantize) { + return quantize + } + return "" + } + + // Match the conservative Qwen linear-attention policy: these low-rank + // projections are sensitive and negligible in size. + if qwen35IsLowRankProjection(name) { + return "" + } + return GetTensorQuantization(name, shape, quantize) +} + +func qwen4ExpIsPLEEmbeddingShard(name string) bool { + return strings.Contains(name, ".ple.ple_embedding.ngram_embedding.shard_") && + strings.HasSuffix(name, ".weight") +} diff --git a/x/create/qwen4_exp_test.go b/x/create/qwen4_exp_test.go new file mode 100644 index 000000000..c5159eff6 --- /dev/null +++ b/x/create/qwen4_exp_test.go @@ -0,0 +1,52 @@ +package create + +import "testing" + +func TestQwen4ExpQuantizationType(t *testing.T) { + policy := qwen4ExpImportTransform{} + + tests := []struct { + name string + tensor string + shape []int32 + quantize string + want string + }{ + { + name: "PLE shard", + tensor: "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight", + shape: []int32{156250, 2560}, + quantize: "nvfp4", + want: "nvfp4", + }, + { + name: "token embedding", + tensor: "model.language_model.embed_tokens.weight", + shape: []int32{248320, 2560}, + quantize: "nvfp4", + want: "", + }, + { + name: "low rank projection", + tensor: "model.language_model.layers.0.linear_attn.in_proj_a.weight", + shape: []int32{24, 5120}, + quantize: "nvfp4", + want: "", + }, + { + name: "ordinary projection", + tensor: "model.language_model.layers.0.linear_attn.out_proj.weight", + shape: []int32{5120, 3072}, + quantize: "nvfp4", + want: "nvfp4", + }, + } + + 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/mlxrunner/imports.go b/x/mlxrunner/imports.go index 5f2cc4cb9..a155de98b 100644 --- a/x/mlxrunner/imports.go +++ b/x/mlxrunner/imports.go @@ -13,4 +13,5 @@ import ( _ "github.com/ollama/ollama/x/models/qwen3" _ "github.com/ollama/ollama/x/models/qwen3_5" _ "github.com/ollama/ollama/x/models/qwen3_5_moe" + _ "github.com/ollama/ollama/x/models/qwen4_exp" ) diff --git a/x/mlxrunner/mlx/ops.go b/x/mlxrunner/mlx/ops.go index 08ce26763..060075aa1 100644 --- a/x/mlxrunner/mlx/ops.go +++ b/x/mlxrunner/mlx/ops.go @@ -49,6 +49,12 @@ func (t *Array) AsType(dtype DType) *Array { return out } +func (t *Array) BitwiseXor(other *Array) *Array { + out := New("BITWISE_XOR") + C.mlx_bitwise_xor(&out.ctx, t.ctx, other.ctx, DefaultStream().ctx) + return out +} + func (t *Array) AsStrided(shape []int, strides []int, offset int) *Array { cShape := make([]C.int, len(shape)) for i, s := range shape { @@ -226,12 +232,24 @@ func (t *Array) RightShift(other *Array) *Array { return out } +func (t *Array) Remainder(other *Array) *Array { + out := New("REMAINDER") + C.mlx_remainder(&out.ctx, t.ctx, other.ctx, DefaultStream().ctx) + return out +} + func (t *Array) Sigmoid() *Array { out := New("SIGMOID") C.mlx_sigmoid(&out.ctx, t.ctx, DefaultStream().ctx) return out } +func (t *Array) Sign() *Array { + out := New("SIGN") + C.mlx_sign(&out.ctx, t.ctx, DefaultStream().ctx) + return out +} + func (t *Array) Sqrt() *Array { out := New("SQRT") C.mlx_sqrt(&out.ctx, t.ctx, DefaultStream().ctx) diff --git a/x/mlxrunner/mlx/ops_extra.go b/x/mlxrunner/mlx/ops_extra.go index 585a55dc7..79d78c78a 100644 --- a/x/mlxrunner/mlx/ops_extra.go +++ b/x/mlxrunner/mlx/ops_extra.go @@ -113,6 +113,12 @@ func GatherQMM(x, w, scales *Array, biases, lhsIndices, rhsIndices *Array, trans // Missing tensor ops +func Arange(start, stop, step float64, dtype DType) *Array { + out := New("ARANGE") + C.mlx_arange(&out.ctx, C.double(start), C.double(stop), C.double(step), C.mlx_dtype(dtype), DefaultStream().ctx) + return out +} + func Tile(a *Array, reps []int32) *Array { cReps := make([]C.int, len(reps)) for i, r := range reps { diff --git a/x/models/qwen3_5/vision.go b/x/models/qwen3_5/vision.go index bc8107134..fe1b45a1e 100644 --- a/x/models/qwen3_5/vision.go +++ b/x/models/qwen3_5/vision.go @@ -92,6 +92,70 @@ type visionLayout struct { delta int32 } +// VisionAdapter exposes the shared Qwen3.5/Qwen 4 vision tower without +// coupling another text architecture to this package's decoder internals. +// The publisher uses the same tensor layout, image preprocessing, and MRoPE +// layout for both families. +type VisionAdapter struct { + // Model is exported so mlx.Collect traverses and pins every tower weight. + // An unexported wrapper field is invisible to the reflection collector. + Model *Model +} + +type VisionAdapterConfig struct { + HiddenSize int32 + RopeDim int32 + RopeTheta float32 + MropeSection []int32 + QuantGroup int + QuantBits int + QuantMode string + TensorQuant map[string]*model.TensorQuantInfo +} + +// NewVisionAdapter returns nil when the config has no vision tower. +func NewVisionAdapter(configData []byte, tensors map[string]*mlx.Array, cfg VisionAdapterConfig) (*VisionAdapter, error) { + mm, err := parseMultimodalConfig(configData) + if err != nil { + return nil, err + } + if mm.VisionConfig == nil { + return nil, nil + } + text := &Config{ + HiddenSize: cfg.HiddenSize, + RopeDim: cfg.RopeDim, + RopeTheta: cfg.RopeTheta, + MropeSection: cfg.MropeSection, + QuantGroupSize: cfg.QuantGroup, + QuantBits: cfg.QuantBits, + QuantMode: cfg.QuantMode, + TensorQuant: cfg.TensorQuant, + } + m := &Model{Config: text, MM: mm, Vision: mm.VisionConfig} + linears := model.NewLinearFactory(tensors, cfg.QuantGroup, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant) + if err := m.loadVisionWeights(tensors, linears); err != nil { + return nil, err + } + return &VisionAdapter{Model: m}, nil +} + +func (a *VisionAdapter) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { + return a.Model.PrepareMedia(segments) +} + +func (a *VisionAdapter) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { + return a.Model.EncodeMedia(item, data) +} + +func (a *VisionAdapter) ScatterMedia(hidden *mlx.Array, b *batch.Batch, delta int) *mlx.Array { + return a.Model.scatterMedia(hidden, b, delta) +} + +func (a *VisionAdapter) RopePositions(b *batch.Batch, length int32) []int32 { + return ropePositions(b, length) +} + // visionWeightRoot locates the tower's tensor prefix: the import pipeline // renames the reference's model.visual to vision_tower. func visionWeightRoot(tensors map[string]*mlx.Array) string { diff --git a/x/models/qwen3_5/vision_test.go b/x/models/qwen3_5/vision_test.go index b1d3938e8..e96307f6c 100644 --- a/x/models/qwen3_5/vision_test.go +++ b/x/models/qwen3_5/vision_test.go @@ -7,9 +7,20 @@ import ( "strings" "testing" + "github.com/ollama/ollama/x/internal/mlxtest" + "github.com/ollama/ollama/x/mlxrunner/mlx" "github.com/ollama/ollama/x/mlxrunner/model/base" ) +func TestVisionAdapterWeightsAreCollectable(t *testing.T) { + mlxtest.Setup(t) + weight := mlx.FromValue(float32(1)) + adapter := &VisionAdapter{Model: &Model{VisionTower: &VisionTower{PosEmbed: weight}}} + if got := mlx.Collect(adapter); len(got) != 1 || got[0] != weight { + t.Fatalf("mlx.Collect(adapter) = %v, want the tower weight", got) + } +} + func TestSmartResize(t *testing.T) { // Goldens from the reference smart_resize at the family's processor // bounds (factor 32, pixels in [65536, 16777216]). diff --git a/x/models/qwen4_exp/blocks.go b/x/models/qwen4_exp/blocks.go new file mode 100644 index 000000000..76bcd5b55 --- /dev/null +++ b/x/models/qwen4_exp/blocks.go @@ -0,0 +1,203 @@ +package qwen4_exp + +import ( + "math" + + "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/models/nn" +) + +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c, side cache.Cache, positions, ropePositions *mlx.Array, cfg *Config) *mlx.Array { + branch, state := l.AttentionConnection.Prepare(x, cfg) + if l.Linear != nil { + branch = l.Linear.Forward(branch, b, c, cfg) + } else { + branch = l.Attention.Forward(branch, b, c, side, positions, ropePositions, cfg) + } + x = l.AttentionConnection.Inject(state, branch, cfg) + + branch, state = l.MLPConnection.Prepare(x, cfg) + branch = l.MoE.Forward(branch, cfg) + return l.MLPConnection.Inject(state, branch, cfg) +} + +func (a *linearAttention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, cfg *Config) *mlx.Array { + dims := x.Dims() + B, L := int32(dims[0]), int32(dims[1]) + valueDim := cfg.LinearNumValueHeads * cfg.LinearValueHeadDim + + qkv := a.QKV.Forward(x) + z := mlx.Reshape(a.Z.Forward(x), B, L, cfg.LinearNumValueHeads, cfg.LinearValueHeadDim) + ba := mlx.Concatenate([]*mlx.Array{a.B.Forward(x), a.A.Forward(x)}, -1) + + convTail := cfg.LinearConvKernelDim - 1 + options := []nn.RecurrentOption{nn.WithConvSiLU()} + var recurrent *cache.RecurrentCache + if value, ok := c.(*cache.RecurrentCache); ok { + recurrent = value + options = append(options, nn.WithRecurrentHistory(recurrent.Get(b, x.DType()))) + if splits := recurrent.SnapshotSplits(int(L)); len(splits) > 0 { + options = append(options, nn.WithSnapshotSplits(splits)) + } + } else { + options = append(options, 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, convState := nn.CausalConv1D(b, qkv, a.Conv, int(convTail), options...) + out, deltaState := nn.GatedDelta(b, convOut, ba, a.DtBias, a.AExp, options...) + outType := out.DType() + out = mlx.RMSNormFn(out, a.NormWeight, cfg.RMSNormEps) + // Qwen 4 uses a sigmoid Gated DeltaNet output gate rather than the SiLU + // gate used by Qwen3.5. + gate := mlx.Sigmoid(z.AsType(mlx.DTypeFloat32)) + out = mlx.Mul(out.AsType(mlx.DTypeFloat32), gate).AsType(outType) + out = a.Out.Forward(mlx.Reshape(out, B, L, valueDim)) + if recurrent != nil { + recurrent.Put(b, convState, deltaState) + } + return out +} + +func (a *fullAttention) Forward(x *mlx.Array, b *batch.Batch, c, side cache.Cache, positions, ropePositions *mlx.Array, cfg *Config) *mlx.Array { + dims := x.Dims() + B, L := int32(dims[0]), int32(dims[1]) + + qg := mlx.Reshape(a.QProj.Forward(x), 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 := mlx.Reshape(a.KProj.Forward(x), B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + v := mlx.Reshape(a.VProj.Forward(x), 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) + cos, sin := qsaRopeCosSin(cfg, ropePositions) + q = applyQwenRoPE(q, cos, sin, cfg.RopeDim) + k = applyQwenRoPE(k, cos, sin, cfg.RopeDim) + + projected := a.Indexer.QKProj.Forward(x) + indexQ := mlx.SliceStartStop(projected, + []int32{0, 0, 0}, + []int32{B, L, cfg.IndexerNumHeads * cfg.IndexerHeadDim}, + ) + rawIndexK := mlx.SliceStartStop(projected, + []int32{0, 0, cfg.IndexerNumHeads * cfg.IndexerHeadDim}, + []int32{B, L, (cfg.IndexerNumHeads + cfg.IndexerKVHeads) * cfg.IndexerHeadDim}, + ) + indexQ = mlx.Reshape(indexQ, B, L, cfg.IndexerNumHeads, cfg.IndexerHeadDim) + indexQ = a.Indexer.QNorm.Forward(indexQ, cfg.RMSNormEps) + indexQ = mlx.Transpose(indexQ, 0, 2, 1, 3) + indexQ = applyQwenRoPE(indexQ, cos, sin, cfg.RopeDim) + rawIndexK = mlx.Transpose( + mlx.Reshape(rawIndexK, B, L, cfg.IndexerKVHeads, cfg.IndexerHeadDim), + 0, 2, 1, 3, + ) + + var mainHistory, indexHistory *nn.KVHistory + if c != nil { + mainHistory = c.(cache.Attention).Update(b, k, v) + if side == nil { + panic("QSA side cache is missing") + } + indexHistory = side.(cache.Attention).Update(b, rawIndexK, ropePositions) + } + + var out *mlx.Array + keyLength := L + for _, offset := range b.SeqOffsets { + keyLength = max(keyLength, offset+L) + } + if mainHistory == nil { + out = nn.ScaledDotProductAttention(b, q, cfg.Scale, + nn.WithKV(k, v, b.SeqQueryLens), nn.WithMask(nn.CausalMask())) + } else if keyLength <= cfg.IndexerBudget { + // QSA selects the entire visible history within the indexer budget. + out = nn.ScaledDotProductAttention(b, q, cfg.Scale, + nn.WithKVHistory(mainHistory), nn.WithMask(nn.CausalMask())) + } else { + compressed := qsaCompressedKeys(indexHistory.K(), indexHistory.V(), a.Indexer, cfg) + scores := mlx.Matmul(indexQ.AsType(mlx.DTypeFloat32), mlx.Transpose(compressed.AsType(mlx.DTypeFloat32), 0, 1, 3, 2)) + scores = mlx.DivScalar(mlx.Sum(mlx.ReLU(scores), 1, false), float32(math.Sqrt(float64(cfg.IndexerHeadDim)))) + logical, valid := qsaLogicalIndices(scores, b, keyLength, cfg) + out = qsaSparseAttention(q, mainHistory, logical, valid, cfg) + } + out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) + out = mlx.Mul(out, mlx.Sigmoid(gate)) + return a.OProj.Forward(out) +} + +func splitLastAxis(a *mlx.Array) (*mlx.Array, *mlx.Array) { + dims := a.Dims() + start := make([]int32, len(dims)) + stop := make([]int32, len(dims)) + for i, dim := range dims { + stop[i] = int32(dim) + } + mid := int32(dims[len(dims)-1]) / 2 + leftStop := append([]int32(nil), stop...) + leftStop[len(leftStop)-1] = mid + rightStart := append([]int32(nil), start...) + rightStart[len(rightStart)-1] = mid + return mlx.SliceStartStop(a, start, leftStop), mlx.SliceStartStop(a, rightStart, stop) +} + +func (m *sparseMoE) Forward(x *mlx.Array, cfg *Config) *mlx.Array { + dims := x.Dims() + B, L := int32(dims[0]), int32(dims[1]) + topK := cfg.NumExpertsPerTok + + probs := mlx.SoftmaxAxis(m.Gate.Forward(x), -1, true) + indices := mlx.Argpartition(mlx.Neg(probs), int(topK)-1, -1) + shape := indices.Dims() + indices = mlx.SliceStartStop(indices, []int32{0, 0, 0}, []int32{int32(shape[0]), int32(shape[1]), topK}) + scores := mlx.TakeAlongAxis(probs, indices, -1) + scores = mlx.Div(scores, mlx.Sum(scores, -1, true)) + + xFlat := mlx.Reshape(mlx.ExpandDims(mlx.ExpandDims(x, -2), -2), B*L, 1, 1, cfg.HiddenSize) + indexFlat := mlx.Reshape(indices, B*L, topK) + doSort := B*L >= 64 + var inverse *mlx.Array + if doSort { + all := mlx.Flatten(indexFlat) + order := mlx.Argsort(all, 0) + inverse = mlx.Argsort(order, 0) + xFlat = mlx.ExpandDims(mlx.Take(mlx.Squeeze(xFlat, 1), mlx.FloorDivideScalar(order, topK), 0), 1) + indexFlat = mlx.Reshape(mlx.Take(all, order, 0), B*L*topK, 1) + } + + var gateUp *mlx.Array + if m.GateUpScales != nil { + gateUp = mlx.GatherQMM(xFlat, m.GateUpExperts, m.GateUpScales, m.GateUpBiases, + nil, indexFlat, true, m.GateUpGroup, m.GateUpBits, m.GateUpMode, doSort) + } else { + gateUp = mlx.GatherMM(xFlat, m.GateUpExperts, nil, indexFlat, doSort) + } + gate, up := splitLastAxis(gateUp) + hidden := mlx.SwiGLU(gate, up) + var down *mlx.Array + if m.DownScales != nil { + down = mlx.GatherQMM(hidden, m.DownExperts, m.DownScales, m.DownBiases, + nil, indexFlat, true, m.DownGroup, m.DownBits, m.DownMode, doSort) + } else { + down = mlx.GatherMM(hidden, m.DownExperts, nil, indexFlat, doSort) + } + if doSort { + down = mlx.Reshape(mlx.Take(mlx.Squeeze(mlx.Squeeze(down, 2), 1), inverse, 0), B*L, topK, cfg.HiddenSize) + } else { + down = mlx.Squeeze(down, 2) + } + differentiated := mlx.Reshape(down, B, L, topK, cfg.HiddenSize) + out := mlx.Sum(mlx.Mul(differentiated, mlx.ExpandDims(scores, -1)), 2, false) + + shared := m.SharedDownProj.Forward(mlx.SwiGLU(m.SharedGateProj.Forward(x), m.SharedUpProj.Forward(x))) + shared = mlx.Mul(shared, mlx.Sigmoid(m.SharedGate.Forward(x))) + return mlx.Add(out, shared) +} diff --git a/x/models/qwen4_exp/config.go b/x/models/qwen4_exp/config.go new file mode 100644 index 000000000..a8db1685d --- /dev/null +++ b/x/models/qwen4_exp/config.go @@ -0,0 +1,238 @@ +package qwen4_exp + +import ( + "encoding/json" + "fmt" + "math" + "strings" +) + +type ropeParameters struct { + RopeTheta float32 `json:"rope_theta"` + PartialRotaryFactor float32 `json:"partial_rotary_factor"` + MRoPEInterleaved bool `json:"mrope_interleaved"` + MRoPESection []int32 `json:"mrope_section"` +} + +type mtpConfig struct { + Hybrid bool `json:"hybrid"` + LayerTypes []string `json:"layer_types"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + UseHiddenStateFromLayer *int32 `json:"mtp_use_hidden_state_from_layer"` + RopeTheta float32 `json:"rope_theta"` +} + +// Config is the text portion of a Qwen 4 checkpoint config. +// Keep the publisher's field names represented explicitly: this schema is an +// artifact contract and must not silently reinterpret unknown shapes as +// Qwen3.5 defaults. +type Config struct { + ModelType string `json:"model_type"` + HiddenSize int32 `json:"hidden_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"` + EOSTokenID int64 `json:"eos_token_id"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + LayerTypes []string `json:"layer_types"` + + FullAttentionInterval int32 `json:"full_attention_interval"` + OutputGateType string `json:"output_gate_type"` + AttentionBias bool `json:"attention_bias"` + AttentionDropout float32 `json:"attention_dropout"` + HiddenAct string `json:"hidden_act"` + MambaSSMDType string `json:"mamba_ssm_dtype"` + TieWordEmbeddings bool `json:"tie_word_embeddings"` + UseCache bool `json:"use_cache"` + + 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"` + + NumExperts int32 `json:"num_experts"` + NumExpertsPerTok int32 `json:"num_experts_per_tok"` + MoeIntermediateSize int32 `json:"moe_intermediate_size"` + SharedExpertIntermediateSize int32 `json:"shared_expert_intermediate_size"` + + HCCount int32 `json:"hc_count"` + HCLowRank int32 `json:"hc_lowrank"` + + IndexerBudget int32 `json:"indexer_budget"` + IndexerCompressRatio int32 `json:"indexer_compress_ratio"` + IndexerHeadDim int32 `json:"indexer_head_dim"` + IndexerKVHeads int32 `json:"indexer_kv_heads"` + IndexerNumHeads int32 `json:"indexer_n_heads"` + + PLEConvKernelSize int32 `json:"ple_conv_kernel_size"` + PLEEmbedDim int32 `json:"ple_embed_dim"` + PLELayerIDs []int32 `json:"ple_layer_ids"` + + HeadsPerNGram int32 `json:"heads_per_ngram"` + MakeNGramVocabSizeDivisible int32 `json:"make_ngram_vocab_size_divisible_by"` + NGramSize int32 `json:"ngram_size"` + NGramVocabSizeBase int64 `json:"ngram_vocab_size_base"` + SplitNGramParts int32 `json:"split_ngram_parts"` + + MTPNumHiddenLayers int32 `json:"mtp_num_hidden_layers"` + MTPUseDedicatedEmbeddings bool `json:"mtp_use_dedicated_embeddings"` + MTP mtpConfig `json:"mtp"` + + PartialRotaryFactor float32 `json:"partial_rotary_factor"` + RopeParameters ropeParameters `json:"rope_parameters"` + + Scale float32 `json:"-"` + RopeDim int32 `json:"-"` +} + +func parseConfig(data []byte) (Config, error) { + var envelope struct { + TextConfig json.RawMessage `json:"text_config"` + } + if err := json.Unmarshal(data, &envelope); err != nil { + return Config{}, fmt.Errorf("parse config envelope: %w", err) + } + + active := data + if len(envelope.TextConfig) > 0 && string(envelope.TextConfig) != "null" { + active = envelope.TextConfig + } + + var cfg Config + if err := json.Unmarshal(active, &cfg); err != nil { + return Config{}, fmt.Errorf("parse text config: %w", err) + } + if err := cfg.validate(); err != nil { + return Config{}, err + } + return cfg, nil +} + +func (cfg *Config) validate() error { + if cfg.ModelType != "qwen4_exp_text" { + return fmt.Errorf("unsupported model_type %q", cfg.ModelType) + } + if cfg.HiddenSize <= 0 || cfg.NumHiddenLayers <= 0 { + return fmt.Errorf("invalid model dimensions: hidden_size=%d num_hidden_layers=%d", cfg.HiddenSize, cfg.NumHiddenLayers) + } + if cfg.EOSTokenID < 0 || cfg.EOSTokenID >= int64(cfg.VocabSize) { + return fmt.Errorf("invalid eos_token_id %d for vocab_size %d", cfg.EOSTokenID, cfg.VocabSize) + } + if cfg.NumAttentionHeads <= 0 || cfg.NumKeyValueHeads <= 0 || cfg.HeadDim <= 0 { + return fmt.Errorf("invalid attention dimensions: heads=%d kv_heads=%d head_dim=%d", cfg.NumAttentionHeads, cfg.NumKeyValueHeads, cfg.HeadDim) + } + if cfg.NumAttentionHeads%cfg.NumKeyValueHeads != 0 { + return fmt.Errorf("num_attention_heads (%d) must be divisible by num_key_value_heads (%d)", cfg.NumAttentionHeads, cfg.NumKeyValueHeads) + } + if cfg.MaxPositionEmbeddings <= 0 { + return fmt.Errorf("invalid max_position_embeddings %d", cfg.MaxPositionEmbeddings) + } + if len(cfg.LayerTypes) != int(cfg.NumHiddenLayers) { + return fmt.Errorf("layer_types has %d entries, want %d", len(cfg.LayerTypes), cfg.NumHiddenLayers) + } + for i, kind := range cfg.LayerTypes { + switch kind { + case "linear_attention", "full_attention": + default: + return fmt.Errorf("layer_types[%d] has unsupported value %q", i, kind) + } + } + + if cfg.LinearNumKeyHeads <= 0 || cfg.LinearNumValueHeads <= 0 || cfg.LinearKeyHeadDim <= 0 || cfg.LinearValueHeadDim <= 0 || cfg.LinearConvKernelDim <= 0 { + return fmt.Errorf("invalid linear attention dimensions") + } + if cfg.LinearNumValueHeads%cfg.LinearNumKeyHeads != 0 { + return fmt.Errorf("linear_num_value_heads (%d) must be divisible by linear_num_key_heads (%d)", cfg.LinearNumValueHeads, cfg.LinearNumKeyHeads) + } + if cfg.OutputGateType != "sigmoid" { + return fmt.Errorf("unsupported output_gate_type %q", cfg.OutputGateType) + } + if cfg.AttentionBias || cfg.AttentionDropout != 0 { + return fmt.Errorf("unsupported attention bias/dropout: bias=%v dropout=%v", cfg.AttentionBias, cfg.AttentionDropout) + } + if cfg.HiddenAct != "silu" || cfg.MambaSSMDType != "float32" { + return fmt.Errorf("unsupported activation/recurrent dtype: hidden_act=%q mamba_ssm_dtype=%q", cfg.HiddenAct, cfg.MambaSSMDType) + } + if cfg.TieWordEmbeddings || !cfg.UseCache { + return fmt.Errorf("unsupported embedding/cache contract: tied=%v use_cache=%v", cfg.TieWordEmbeddings, cfg.UseCache) + } + + if cfg.NumExperts <= 0 || cfg.NumExpertsPerTok <= 0 || cfg.NumExpertsPerTok > cfg.NumExperts { + return fmt.Errorf("invalid MoE dimensions: experts=%d experts_per_tok=%d", cfg.NumExperts, cfg.NumExpertsPerTok) + } + if cfg.MoeIntermediateSize <= 0 || cfg.SharedExpertIntermediateSize <= 0 { + return fmt.Errorf("invalid MoE intermediate dimensions") + } + + if cfg.HCCount <= 0 || cfg.HCLowRank <= 0 { + return fmt.Errorf("invalid hyper-connection dimensions: count=%d lowrank=%d", cfg.HCCount, cfg.HCLowRank) + } + if cfg.IndexerBudget <= 0 || cfg.IndexerCompressRatio <= 0 || cfg.IndexerHeadDim <= 0 || cfg.IndexerKVHeads <= 0 || cfg.IndexerNumHeads <= 0 { + return fmt.Errorf("invalid sparse indexer dimensions") + } + if cfg.IndexerBudget < cfg.IndexerCompressRatio { + return fmt.Errorf("indexer_budget (%d) must be at least indexer_compress_ratio (%d)", cfg.IndexerBudget, cfg.IndexerCompressRatio) + } + if cfg.IndexerKVHeads != 1 && cfg.IndexerKVHeads != cfg.IndexerNumHeads { + return fmt.Errorf("indexer_kv_heads (%d) must be 1 or equal indexer_n_heads (%d)", cfg.IndexerKVHeads, cfg.IndexerNumHeads) + } + if cfg.PLEEmbedDim != cfg.HiddenSize || cfg.PLEConvKernelSize <= 0 || len(cfg.PLELayerIDs) == 0 { + return fmt.Errorf("invalid PLE dimensions: embed_dim=%d conv_kernel=%d layer_ids=%v", cfg.PLEEmbedDim, cfg.PLEConvKernelSize, cfg.PLELayerIDs) + } + pleLayers := make(map[int32]struct{}, len(cfg.PLELayerIDs)) + for _, layer := range cfg.PLELayerIDs { + if layer < 1 || layer > cfg.NumHiddenLayers { + return fmt.Errorf("PLE layer ID %d is outside [1, %d]", layer, cfg.NumHiddenLayers) + } + if _, ok := pleLayers[layer]; ok { + return fmt.Errorf("duplicate PLE layer ID %d", layer) + } + pleLayers[layer] = struct{}{} + } + if cfg.NGramSize <= 0 || cfg.HeadsPerNGram <= 0 || cfg.SplitNGramParts <= 0 || cfg.NGramVocabSizeBase <= 0 { + return fmt.Errorf("invalid n-gram embedding dimensions") + } + if cfg.MTPNumHiddenLayers != 1 || cfg.MTP.NumHiddenLayers != 1 || len(cfg.MTP.LayerTypes) != 1 || cfg.MTP.LayerTypes[0] != "full_attention" { + return fmt.Errorf("unsupported MTP layout") + } + if cfg.MTPUseDedicatedEmbeddings || !cfg.MTP.Hybrid || cfg.MTP.UseHiddenStateFromLayer != nil { + return fmt.Errorf("unsupported MTP embedding/hidden-state contract") + } + + if cfg.RMSNormEps == 0 { + cfg.RMSNormEps = 1e-6 + } + if cfg.RopeParameters.RopeTheta <= 0 { + return fmt.Errorf("invalid rope_theta: %v", cfg.RopeParameters.RopeTheta) + } + if cfg.RopeParameters.PartialRotaryFactor > 0 { + cfg.PartialRotaryFactor = cfg.RopeParameters.PartialRotaryFactor + } + cfg.RopeDim = int32(float32(cfg.HeadDim) * cfg.PartialRotaryFactor) + if cfg.RopeDim <= 0 || cfg.RopeDim > cfg.HeadDim || cfg.RopeDim > cfg.IndexerHeadDim || cfg.RopeDim%2 != 0 { + return fmt.Errorf("invalid rotary dimension %d from head_dim=%d partial_rotary_factor=%v", cfg.RopeDim, cfg.HeadDim, cfg.PartialRotaryFactor) + } + if !cfg.RopeParameters.MRoPEInterleaved || len(cfg.RopeParameters.MRoPESection) != 3 { + return fmt.Errorf("unsupported MRoPE layout: interleaved=%v sections=%v", cfg.RopeParameters.MRoPEInterleaved, cfg.RopeParameters.MRoPESection) + } + var ropeSections int32 + for _, section := range cfg.RopeParameters.MRoPESection { + ropeSections += section + } + if ropeSections != cfg.RopeDim/2 { + return fmt.Errorf("MRoPE sections total %d, want rotary half-dimension %d", ropeSections, cfg.RopeDim/2) + } + if cfg.MTP.RopeTheta != cfg.RopeParameters.RopeTheta { + return fmt.Errorf("MTP rope_theta %v does not match text rope_theta %v", cfg.MTP.RopeTheta, cfg.RopeParameters.RopeTheta) + } + cfg.Scale = float32(1 / math.Sqrt(float64(cfg.HeadDim))) + return nil +} + +func (cfg *Config) layerIsLinear(layer int) bool { + return strings.EqualFold(cfg.LayerTypes[layer], "linear_attention") +} diff --git a/x/models/qwen4_exp/config_test.go b/x/models/qwen4_exp/config_test.go new file mode 100644 index 000000000..0d42bb6fb --- /dev/null +++ b/x/models/qwen4_exp/config_test.go @@ -0,0 +1,132 @@ +package qwen4_exp + +import ( + "math" + "testing" +) + +const testConfig = `{ + "text_config": { + "model_type": "qwen4_exp_text", + "hidden_size": 2560, + "num_hidden_layers": 4, + "num_attention_heads": 24, + "num_key_value_heads": 2, + "head_dim": 256, + "rms_norm_eps": 0.000001, + "vocab_size": 248320, + "eos_token_id": 248044, + "max_position_embeddings": 262144, + "layer_types": ["linear_attention", "linear_attention", "linear_attention", "full_attention"], + "full_attention_interval": 4, + "output_gate_type": "sigmoid", + "attention_bias": false, + "attention_dropout": 0.0, + "hidden_act": "silu", + "mamba_ssm_dtype": "float32", + "tie_word_embeddings": false, + "use_cache": true, + "linear_num_value_heads": 48, + "linear_num_key_heads": 16, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "num_experts": 512, + "num_experts_per_tok": 10, + "moe_intermediate_size": 640, + "shared_expert_intermediate_size": 640, + "hc_count": 4, + "hc_lowrank": 320, + "indexer_budget": 2048, + "indexer_compress_ratio": 4, + "indexer_head_dim": 128, + "indexer_kv_heads": 1, + "indexer_n_heads": 4, + "ple_conv_kernel_size": 4, + "ple_embed_dim": 2560, + "ple_layer_ids": [2], + "heads_per_ngram": 8, + "make_ngram_vocab_size_divisible_by": 128, + "ngram_size": 3, + "ngram_vocab_size_base": 20000000, + "split_ngram_parts": 128, + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "mtp": { + "hybrid": true, + "layer_types": ["full_attention"], + "num_hidden_layers": 1, + "rope_theta": 10000000 + }, + "partial_rotary_factor": 0.25, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [11, 11, 10], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + } + } +}` + +func TestParseConfig(t *testing.T) { + cfg, err := parseConfig([]byte(testConfig)) + if err != nil { + t.Fatal(err) + } + + if cfg.ModelType != "qwen4_exp_text" || cfg.HiddenSize != 2560 || cfg.HCCount != 4 { + t.Fatalf("unexpected config identity: model_type=%q hidden=%d hc_count=%d", cfg.ModelType, cfg.HiddenSize, cfg.HCCount) + } + if !cfg.layerIsLinear(0) || cfg.layerIsLinear(3) { + t.Fatalf("unexpected layer classification: %v", cfg.LayerTypes) + } + if cfg.RopeDim != 64 { + t.Fatalf("rope dimension = %d, want 64", cfg.RopeDim) + } + wantScale := float32(1 / math.Sqrt(256)) + if cfg.Scale != wantScale { + t.Fatalf("attention scale = %v, want %v", cfg.Scale, wantScale) + } +} + +func TestParseConfigRejectsUnknownSemantics(t *testing.T) { + for _, replacement := range []struct { + name string + old string + new string + }{ + {name: "gate", old: `"output_gate_type": "sigmoid"`, new: `"output_gate_type": "swish"`}, + {name: "layer", old: `"full_attention"]`, new: `"sliding_attention"]`}, + {name: "mtp", old: "\"num_hidden_layers\": 1,\n \"rope_theta\"", new: "\"num_hidden_layers\": 2,\n \"rope_theta\""}, + {name: "attention head grouping", old: `"num_attention_heads": 24`, new: `"num_attention_heads": 23`}, + {name: "linear convolution", old: `"linear_conv_kernel_dim": 4`, new: `"linear_conv_kernel_dim": 0`}, + {name: "indexer budget", old: `"indexer_budget": 2048`, new: `"indexer_budget": 2`}, + {name: "indexer head broadcast", old: `"indexer_kv_heads": 1`, new: `"indexer_kv_heads": 2`}, + {name: "PLE layer range", old: `"ple_layer_ids": [2]`, new: `"ple_layer_ids": [0]`}, + {name: "duplicate PLE layer", old: `"ple_layer_ids": [2]`, new: `"ple_layer_ids": [2, 2]`}, + {name: "dedicated MTP embeddings", old: `"mtp_use_dedicated_embeddings": false`, new: `"mtp_use_dedicated_embeddings": true`}, + {name: "non-hybrid MTP", old: `"hybrid": true`, new: `"hybrid": false`}, + {name: "MTP hidden-state selection", old: `"hybrid": true,`, new: `"hybrid": true, "mtp_use_hidden_state_from_layer": 1,`}, + {name: "MTP rope", old: "\"rope_theta\": 10000000\n },\n \"partial_rotary_factor\"", new: "\"rope_theta\": 10000001\n },\n \"partial_rotary_factor\""}, + {name: "non-interleaved MRoPE", old: `"mrope_interleaved": true`, new: `"mrope_interleaved": false`}, + {name: "MRoPE sections", old: `"mrope_section": [11, 11, 10]`, new: `"mrope_section": [10, 11, 10]`}, + {name: "odd rotary dimension", old: "\"partial_rotary_factor\": 0.25,\n \"rope_theta\"", new: "\"partial_rotary_factor\": 0.25390625,\n \"rope_theta\""}, + } { + t.Run(replacement.name, func(t *testing.T) { + data := []byte(replaceOnce(testConfig, replacement.old, replacement.new)) + if _, err := parseConfig(data); err == nil { + t.Fatal("parseConfig accepted unsupported semantics") + } + }) + } +} + +func replaceOnce(s, old, new string) string { + for i := 0; i+len(old) <= len(s); i++ { + if s[i:i+len(old)] == old { + return s[:i] + new + s[i+len(old):] + } + } + return s +} diff --git a/x/models/qwen4_exp/engram.go b/x/models/qwen4_exp/engram.go new file mode 100644 index 000000000..1e6f5d8c9 --- /dev/null +++ b/x/models/qwen4_exp/engram.go @@ -0,0 +1,117 @@ +package qwen4_exp + +import ( + "math" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +func (n *streamRMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array { + originalShape := x.Dims() + if originalShape[len(originalShape)-1] == n.Weight.Size() { + shape := append([]int(nil), originalShape[:len(originalShape)-1]...) + shape = append(shape, n.Weight.Dim(0), n.Weight.Dim(1)) + x = x.Reshape(shape...) + } + dtype := x.DType() + values := x.AsType(mlx.DTypeFloat32) + meanSquare := mlx.Mean(mlx.Mul(values, values), -1, true) + invRMS := mlx.Div(mlx.NewScalarArray(1), mlx.AddScalar(meanSquare, eps).Sqrt()) + normalized := mlx.Mul(mlx.Mul(values, invRMS), n.Weight.AsType(mlx.DTypeFloat32)).AsType(dtype) + return normalized.Reshape(originalShape...) +} + +func (p *PLE) hashes(b *batch.Batch, history *mlx.Array, cfg *Config) (*mlx.Array, *mlx.Array) { + tokens := b.InputIDs.AsType(mlx.DTypeInt64) + joined := mlx.Concatenate([]*mlx.Array{history, tokens}, 1) + length := b.InputIDs.Dim(1) + shifts := make([]*mlx.Array, cfg.NGramSize) + eos := mlx.FromValues([]int64{cfg.EOSTokenID}, 1) + for k := range cfg.NGramSize { + start := int(cfg.NGramSize - 1 - k) + shifts[k] = mlx.SliceStartStop(joined, + []int32{0, int32(start)}, + []int32{int32(joined.Dim(0)), int32(start + length)}, + ) + for previous := int32(1); previous <= k; previous++ { + boundary := mlx.SliceStartStop(joined, + []int32{0, cfg.NGramSize - 1 - previous}, + []int32{int32(joined.Dim(0)), cfg.NGramSize - 1 - previous + int32(length)}, + ) + shifts[k] = mlx.Where(boundary.Equal(eos), eos, shifts[k]) + } + } + + hashes := make([]*mlx.Array, 0, (cfg.NGramSize-1)*cfg.HeadsPerNGram) + head := 0 + for ngram := int32(2); ngram <= cfg.NGramSize; ngram++ { + multiplier := mlx.Take(p.LayerMultipliers, mlx.FromValue(0), 0) + mix := mlx.Mul(shifts[0], multiplier) + for k := int32(1); k < ngram; k++ { + multiplier = mlx.Take(p.LayerMultipliers, mlx.FromValue(int(k)), 0) + mix = mix.BitwiseXor(mlx.Mul(shifts[k], multiplier)) + } + for range cfg.HeadsPerNGram { + modulus := mlx.Take(p.HeadVocabSizes, mlx.FromValue(head), 0) + hashes = append(hashes, mix.Remainder(modulus)) + head++ + } + } + return mlx.Add(mlx.Stack(hashes, 2), p.HeadOffsets), tokens +} + +func (p *PLE) lookup(globalIDs *mlx.Array) *mlx.Array { + rows := mlx.FromValues([]int64{int64(p.ShardRows)}, 1) + shardIDs := globalIDs.FloorDivide(rows) + localIDs := globalIDs.Remainder(rows) + zeroID := mlx.FromValues([]int64{0}, 1) + + // The checkpoint partitions one logical table into equal contiguous row + // shards. This portable fallback touches every shard, then masks and sums + // the selected rows. It prioritizes reference bring-up over performance; + // production support should replace it with a fused split-table lookup. + var result *mlx.Array + for i, embedding := range p.EmbeddingShards { + mask := shardIDs.Equal(mlx.FromValues([]int64{int64(i)}, 1)) + indices := mlx.Where(mask, localIDs, zeroID) + values := embedding.Forward(indices) + values = mlx.Mul(values, mlx.ExpandDims(mask, -1).AsType(values.DType())) + if result == nil { + result = values + } else { + result = mlx.Add(result, values) + } + } + return result +} + +// Forward implements the published Engram dataflow against the checkpoint's +// fused four-stream projections and split embedding table. +func (p *PLE) Forward(hidden *mlx.Array, b *batch.Batch, tokens *engramCache, cfg *Config) *mlx.Array { + dims := hidden.Dims() + B, L := int32(dims[0]), int32(dims[1]) + streams := mlx.Reshape(hidden, B, L, cfg.HCCount, cfg.HiddenSize) + + hashes, inputIDs := p.hashes(b, tokens.get(b), cfg) + embeddings := p.lookup(hashes) + embeddings = mlx.Reshape(embeddings, B, L, cfg.HiddenSize) + + key := mlx.Reshape(p.KeyProj.Forward(embeddings), B, L, cfg.HCCount, cfg.HiddenSize) + key = p.NormKey.Forward(key, cfg.RMSNormEps) + query := p.NormQuery.Forward(streams, cfg.RMSNormEps) + gate := mlx.DivScalar(mlx.Sum(mlx.Mul(key, query), -1, false), float32(math.Sqrt(float64(cfg.HiddenSize)))) + gate = mlx.Mul(mlx.Maximum(gate.Abs(), mlx.NewScalarArray(1e-6)).Sqrt(), gate.Sign()) + gate = mlx.ExpandDims(mlx.Sigmoid(gate), -1) + + value := mlx.ExpandDims(p.ValueProj.Forward(embeddings), 2) + value = mlx.Mul(gate, value) + convInput := p.NormConv.Forward(value, cfg.RMSNormEps) + convInput = mlx.Reshape(convInput, B, L, cfg.HCCount*cfg.HiddenSize) + convHistory := tokens.getConvHistory(b, convInput.DType()) + conv := p.Conv.Forward(mlx.Concatenate([]*mlx.Array{convHistory, convInput}, 1)) + conv = mlx.Reshape(mlx.SiLU(conv), B, L, cfg.HCCount, cfg.HiddenSize) + + tokens.put(b, inputIDs, convInput) + return mlx.Reshape(mlx.Add(value, conv), B, L, cfg.HCCount*cfg.HiddenSize).AsType(hidden.DType()) +} diff --git a/x/models/qwen4_exp/engram_cache.go b/x/models/qwen4_exp/engram_cache.go new file mode 100644 index 000000000..b836a46d3 --- /dev/null +++ b/x/models/qwen4_exp/engram_cache.go @@ -0,0 +1,205 @@ +package qwen4_exp + +import ( + "fmt" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/cache" + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +// engramCache keeps the two raw token IDs and nine normalized +// convolution inputs preceding the current chunk. Both Engram operations span +// forward boundaries, so this state follows the model cache lifecycle. +type engramCache struct { + history *mlx.Array + convHistory *mlx.Array + offset int + eosID int64 + width int + convTail int + convDim int + + scheduled []int + captured []cache.Snapshot +} + +type engramSnapshot struct { + history *mlx.Array + convHistory *mlx.Array + offset int +} + +func newEngramCache(width, convTail, convDim int, eosID int64) *engramCache { + return &engramCache{width: width, convTail: convTail, convDim: convDim, eosID: eosID} +} + +func (c *engramCache) setHistory(value *mlx.Array) { + value = value.Clone() + mlx.Pin(value) + mlx.Unpin(c.history) + c.history = value +} + +func (c *engramCache) setConvHistory(value *mlx.Array) { + value = value.Clone() + mlx.Pin(value) + mlx.Unpin(c.convHistory) + c.convHistory = value +} + +func (c *engramCache) get(b *batch.Batch) *mlx.Array { + if c.history == nil { + batchSize := b.InputIDs.Dim(0) + values := make([]int64, batchSize*c.width) + for i := range values { + values[i] = c.eosID + } + c.setHistory(mlx.FromValues(values, batchSize, c.width)) + } + return c.history +} + +func (c *engramCache) getConvHistory(b *batch.Batch, dtype mlx.DType) *mlx.Array { + if c.convHistory == nil { + c.setConvHistory(mlx.Zeros(dtype, b.InputIDs.Dim(0), c.convTail, c.convDim)) + } + return c.convHistory +} + +func tailAfter(history, values *mlx.Array, count, width int) *mlx.Array { + if count > 0 && count < values.Dim(1) { + start := make([]int32, values.NumDims()) + stop := make([]int32, values.NumDims()) + for i, dim := range values.Dims() { + stop[i] = int32(dim) + } + stop[1] = int32(count) + values = mlx.SliceStartStop(values, start, stop) + } + joined := history + if count > 0 { + joined = mlx.Concatenate([]*mlx.Array{history, values}, 1) + } + dims := joined.Dims() + start := make([]int32, len(dims)) + stop := make([]int32, len(dims)) + for i, dim := range dims { + stop[i] = int32(dim) + } + start[1] = int32(joined.Dim(1) - width) + return mlx.SliceStartStop(joined, start, stop) +} + +func (c *engramCache) put(b *batch.Batch, inputIDs, convInput *mlx.Array) { + if b.InputIDs.Dim(0) != 1 || len(b.SeqQueryLens) != 1 { + panic("Engram cache requires a single sequence") + } + history := c.get(b) + convHistory := c.getConvHistory(b, convInput.DType()) + length := int(b.SeqQueryLens[0]) + start, end := c.offset, c.offset+length + for i, boundary := range c.scheduled { + if c.captured[i] != nil || boundary < start || boundary > end { + continue + } + c.captured[i] = newEngramSnapshot( + tailAfter(history, inputIDs, boundary-start, c.width), + tailAfter(convHistory, convInput, boundary-start, c.convTail), + boundary, + ) + } + c.setHistory(tailAfter(history, inputIDs, length, c.width)) + c.setConvHistory(tailAfter(convHistory, convInput, length, c.convTail)) + c.offset = end +} + +func (c *engramCache) State() []*mlx.Array { + state := []*mlx.Array{c.history, c.convHistory} + for _, snapshot := range c.captured { + if snapshot != nil { + value := snapshot.(*engramSnapshot) + state = append(state, value.history, value.convHistory) + } + } + return state +} + +func (c *engramCache) Free() { + mlx.Unpin(c.history, c.convHistory) + c.history = nil + c.convHistory = nil + c.offset = 0 + c.scheduled = nil + c.captured = nil +} + +func (c *engramCache) Offset() int { return c.offset } + +func (c *engramCache) PrepareSnapshots(offsets []int) { + if c.captured != nil { + panic("Engram cache: previous snapshot schedule was not drained") + } + for i, offset := range offsets { + if offset < c.offset || (i > 0 && offset <= offsets[i-1]) { + panic(fmt.Sprintf("Engram cache: invalid snapshot offsets %v at current offset %d", offsets, c.offset)) + } + } + c.scheduled = append([]int(nil), offsets...) + c.captured = make([]cache.Snapshot, len(offsets)) + for i, offset := range offsets { + if offset == c.offset && c.history != nil && c.convHistory != nil { + c.captured[i] = newEngramSnapshot(c.history, c.convHistory, c.offset) + } + } +} + +func (c *engramCache) TakeSnapshots() []cache.Snapshot { + result := c.captured + c.scheduled = nil + c.captured = nil + return result +} + +func (c *engramCache) Snapshot(int) cache.Snapshot { + if c.history == nil || c.convHistory == nil { + return nil + } + return newEngramSnapshot(c.history, c.convHistory, c.offset) +} + +func (c *engramCache) Restore(snapshot cache.Snapshot, target int) bool { + if snapshot == nil { + return target == c.offset + } + value, ok := snapshot.(*engramSnapshot) + if !ok || value.offset != target { + return false + } + c.setHistory(value.history) + c.setConvHistory(value.convHistory) + c.offset = target + return true +} + +func (c *engramCache) Merge(parent, child cache.Snapshot) cache.Snapshot { + if parent != nil { + parent.Close() + } + return child +} + +func (c *engramCache) Split(snapshot cache.Snapshot, _ int) (cache.Snapshot, cache.Snapshot) { + return nil, snapshot +} + +func newEngramSnapshot(history, convHistory *mlx.Array, offset int) *engramSnapshot { + snapshot := &engramSnapshot{history: history.Clone(), convHistory: convHistory.Clone(), offset: offset} + mlx.Pin(snapshot.history, snapshot.convHistory) + mlx.AsyncEval(snapshot.history, snapshot.convHistory) + return snapshot +} + +func (s *engramSnapshot) Size() int { return s.history.NumBytes() + s.convHistory.NumBytes() } +func (s *engramSnapshot) SetMaterializeHook(func(int)) {} +func (s *engramSnapshot) Close() { mlx.Unpin(s.history, s.convHistory) } diff --git a/x/models/qwen4_exp/engram_test.go b/x/models/qwen4_exp/engram_test.go new file mode 100644 index 000000000..44a3f5f5b --- /dev/null +++ b/x/models/qwen4_exp/engram_test.go @@ -0,0 +1,105 @@ +package qwen4_exp + +import ( + "slices" + "testing" + + "github.com/ollama/ollama/x/internal/mlxtest" + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +func TestEngramHashes(t *testing.T) { + mlxtest.Setup(t) + + p := &PLE{ + LayerMultipliers: mlx.FromValues([]int64{3, 5, 7}, 3), + HeadVocabSizes: mlx.FromValues([]int64{11, 13, 17, 19}, 4), + HeadOffsets: mlx.FromValues([]int64{0, 11, 24, 41}, 4), + } + cfg := &Config{NGramSize: 3, HeadsPerNGram: 2, EOSTokenID: 9} + b := &batch.Batch{InputIDs: mlx.FromValues([]int32{1, 9, 2, 3}, 1, 4)} + history := mlx.FromValues([]int64{7, 8}, 1, 2) + + got, _ := p.hashes(b, history, cfg) + got = got.AsType(mlx.DTypeInt32) + mlx.Eval(got) + values := got.Ints() + want := []int32{10, 15, 33, 48, 8, 15, 28, 41, 10, 15, 27, 42, 3, 14, 33, 44} + if !slices.Equal(values, want) { + t.Fatalf("hashes = %v, want %v", values, want) + } +} + +func TestEngramCacheCarriesChunkHistory(t *testing.T) { + mlxtest.Setup(t) + c := newEngramCache(2, 3, 1, 9) + t.Cleanup(c.Free) + b := &batch.Batch{ + InputIDs: mlx.FromValues([]int32{1, 2}, 1, 2), + SeqQueryLens: []int32{2}, + } + c.put(b, mlx.FromValues([]int64{1, 2}, 1, 2), mlx.FromValues([]float32{10, 20}, 1, 2, 1)) + b.InputIDs = mlx.FromValues([]int32{3, 4}, 1, 2) + c.put(b, mlx.FromValues([]int64{3, 4}, 1, 2), mlx.FromValues([]float32{30, 40}, 1, 2, 1)) + + history := c.history.AsType(mlx.DTypeInt32) + mlx.Eval(history, c.convHistory) + if got, want := history.Ints(), []int32{3, 4}; !slices.Equal(got, want) { + t.Fatalf("token history = %v, want %v", got, want) + } + if got, want := c.convHistory.Floats(), []float32{20, 30, 40}; !slices.Equal(got, want) { + t.Fatalf("conv history = %v, want %v", got, want) + } +} + +func TestEngramCacheRestoresScheduledSnapshot(t *testing.T) { + mlxtest.Setup(t) + c := newEngramCache(2, 3, 1, 9) + t.Cleanup(c.Free) + b := &batch.Batch{ + InputIDs: mlx.FromValues([]int32{1, 2}, 1, 2), + SeqQueryLens: []int32{2}, + } + c.put(b, mlx.FromValues([]int64{1, 2}, 1, 2), mlx.FromValues([]float32{10, 20}, 1, 2, 1)) + + c.PrepareSnapshots([]int{3, 4}) + b.InputIDs = mlx.FromValues([]int32{3, 4}, 1, 2) + inputIDs := mlx.FromValues([]int64{3, 4}, 1, 2) + convInput := mlx.FromValues([]float32{30, 40}, 1, 2, 1) + c.put(b, inputIDs, convInput) + snapshots := c.TakeSnapshots() + if len(snapshots) != 2 || snapshots[0] == nil || snapshots[1] == nil { + t.Fatalf("TakeSnapshots() = %v, want two captured snapshots", snapshots) + } + for _, snapshot := range snapshots { + t.Cleanup(snapshot.Close) + } + first := snapshots[0].(*engramSnapshot) + mlx.Eval(first.history, first.convHistory) + if got, want := first.convHistory.Floats(), []float32{10, 20, 30}; !slices.Equal(got, want) { + t.Fatalf("captured convolution history = %v, want %v", got, want) + } + + b.InputIDs = mlx.FromValues([]int32{5, 6}, 1, 2) + c.put(b, mlx.FromValues([]int64{5, 6}, 1, 2), mlx.FromValues([]float32{50, 60}, 1, 2, 1)) + if !c.Restore(snapshots[0], 3) { + t.Fatal("Restore(snapshot, 3) failed") + } + history := c.history.AsType(mlx.DTypeInt32) + mlx.Eval(history, c.convHistory) + if got, want := history.Ints(), []int32{2, 3}; !slices.Equal(got, want) { + t.Fatalf("restored token history = %v, want %v", got, want) + } + if got, want := c.convHistory.Floats(), []float32{10, 20, 30}; !slices.Equal(got, want) { + t.Fatalf("restored convolution history = %v, want %v", got, want) + } + + parent, child := c.Split(snapshots[1], 3) + if parent != nil || child != snapshots[1] { + t.Fatalf("Split(snapshot) = (%v, %v), want (nil, snapshot)", parent, child) + } + if merged := c.Merge(nil, child); merged != child { + t.Fatalf("Merge(nil, child) = %v, want child", merged) + } +} diff --git a/x/models/qwen4_exp/hyper_connection.go b/x/models/qwen4_exp/hyper_connection.go new file mode 100644 index 000000000..85c1c3421 --- /dev/null +++ b/x/models/qwen4_exp/hyper_connection.go @@ -0,0 +1,50 @@ +package qwen4_exp + +import "github.com/ollama/ollama/x/mlxrunner/mlx" + +type hyperConnectionState struct { + residual *mlx.Array + normed *mlx.Array +} + +func expandStreams(x *mlx.Array, cfg *Config) *mlx.Array { + dims := x.Dims() + B, L := int32(dims[0]), int32(dims[1]) + expanded := mlx.Tile(mlx.ExpandDims(x, -2), []int32{1, 1, cfg.HCCount, 1}) + return mlx.Reshape(expanded, B, L, cfg.HCCount*cfg.HiddenSize) +} + +// Prepare reduces the residual streams into one branch input and preserves +// the state needed to inject the branch output. +func (h *hyperConnection) Prepare(residual *mlx.Array, cfg *Config) (*mlx.Array, hyperConnectionState) { + dims := residual.Dims() + B, L := int32(dims[0]), int32(dims[1]) + normed := h.Norm.Forward(residual, cfg.RMSNormEps) + mixHidden := mlx.SiLU(mlx.DivScalar(h.InputMixDown.Forward(normed), float32(cfg.HCCount))) + mix := h.InputMixUp.Forward(mixHidden) + mix = mlx.Reshape(mix, B, L, cfg.HCCount, cfg.HiddenSize) + streams := mlx.Reshape(normed, B, L, cfg.HCCount, cfg.HiddenSize) + branch := mlx.DivScalar(mlx.Sum(mlx.Mul(streams, mlx.Sigmoid(mix)), 2, false), float32(cfg.HCCount)) + return branch, hyperConnectionState{residual: residual, normed: normed} +} + +// Inject broadcasts one branch result back into the residual streams. The +// public HC post-connection uses 2*sigmoid so an untrained coefficient is one; +// block_inject_weight has exactly one dynamic coefficient per stream. +func (h *hyperConnection) Inject(state hyperConnectionState, branch *mlx.Array, cfg *Config) *mlx.Array { + dims := branch.Dims() + B, L := int32(dims[0]), int32(dims[1]) + weight := mlx.MulScalar(mlx.Sigmoid(mlx.DivScalar(h.BlockInject.Forward(state.normed), float32(cfg.HCCount))), 2) + weight = mlx.ExpandDims(weight, -1) + branch = mlx.ExpandDims(branch, -2) + update := mlx.Mul(weight, branch) + residual := mlx.Reshape(state.residual, B, L, cfg.HCCount, cfg.HiddenSize) + return mlx.Reshape(mlx.Add(residual, update), B, L, cfg.HCCount*cfg.HiddenSize) +} + +// Reduce applies the trained input mixer as the final stream reduction. The +// final checkpoint mixer has no block injection tensor, matching this role. +func (h *hyperConnection) Reduce(residual *mlx.Array, cfg *Config) *mlx.Array { + branch, _ := h.Prepare(residual, cfg) + return branch +} diff --git a/x/models/qwen4_exp/hyper_connection_test.go b/x/models/qwen4_exp/hyper_connection_test.go new file mode 100644 index 000000000..58657b8a6 --- /dev/null +++ b/x/models/qwen4_exp/hyper_connection_test.go @@ -0,0 +1,104 @@ +package qwen4_exp + +import ( + "math" + "testing" + + "github.com/ollama/ollama/x/internal/mlxtest" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" +) + +func TestHyperConnectionMatchesReferenceFormula(t *testing.T) { + mlxtest.Setup(t) + + cfg := &Config{HCCount: 2, HiddenSize: 2, HCLowRank: 2, RMSNormEps: 1e-6} + normWeight := []float32{1.1, 0.9, 1.2, 0.8} + downWeight := [][]float32{{0.2, -0.3, 0.4, 0.1}, {-0.1, 0.5, 0.2, -0.4}} + upWeight := [][]float32{{0.3, -0.2}, {0.1, 0.4}, {-0.5, 0.2}, {0.25, 0.15}} + injectWeight := [][]float32{{0.2, -0.1, 0.3, 0.4}, {-0.3, 0.2, 0.1, 0.5}} + input := []float32{1, 2, 3, 4} + + h := &hyperConnection{ + Norm: &streamRMSNorm{Weight: mlx.FromValues(normWeight, 2, 2)}, + InputMixDown: nn.NewLinear(matrix(downWeight), nil), + InputMixUp: nn.NewLinear(matrix(upWeight), nil), + BlockInject: nn.NewLinear(matrix(injectWeight), nil), + } + residual := mlx.FromValues(input, 1, 1, 4) + branch, state := h.Prepare(residual, cfg) + got := h.Inject(state, branch, cfg).AsType(mlx.DTypeFloat32) + reduced := h.Reduce(residual, cfg).AsType(mlx.DTypeFloat32) + mlx.Eval(got, reduced) + + normed := append([]float32(nil), input...) + for stream := range int(cfg.HCCount) { + start := stream * int(cfg.HiddenSize) + var square float64 + for i := range int(cfg.HiddenSize) { + x := float64(input[start+i]) + square += x * x + } + invRMS := 1 / math.Sqrt(square/float64(cfg.HiddenSize)+float64(cfg.RMSNormEps)) + for i := range int(cfg.HiddenSize) { + normed[start+i] = float32(float64(input[start+i]) * invRMS * float64(normWeight[start+i])) + } + } + + down := matvec(downWeight, normed) + for i := range down { + down[i] /= float32(cfg.HCCount) + down[i] *= 1 / (1 + float32(math.Exp(float64(-down[i])))) + } + mix := matvec(upWeight, down) + wantBranch := make([]float32, cfg.HiddenSize) + for stream := range int(cfg.HCCount) { + for i := range int(cfg.HiddenSize) { + j := stream*int(cfg.HiddenSize) + i + gate := 1 / (1 + float32(math.Exp(float64(-mix[j])))) + wantBranch[i] += gate * normed[j] / float32(cfg.HCCount) + } + } + + injection := matvec(injectWeight, normed) + want := append([]float32(nil), input...) + for stream := range int(cfg.HCCount) { + weight := 2 / (1 + float32(math.Exp(float64(-injection[stream]/float32(cfg.HCCount))))) + for i := range int(cfg.HiddenSize) { + want[stream*int(cfg.HiddenSize)+i] += weight * wantBranch[i] + } + } + + assertClose(t, "mixed branch", reduced.Floats(), wantBranch) + assertClose(t, "injected streams", got.Floats(), want) +} + +func matrix(rows [][]float32) *mlx.Array { + values := make([]float32, 0, len(rows)*len(rows[0])) + for _, row := range rows { + values = append(values, row...) + } + return mlx.FromValues(values, len(rows), len(rows[0])) +} + +func matvec(weight [][]float32, input []float32) []float32 { + output := make([]float32, len(weight)) + for i, row := range weight { + for j, value := range row { + output[i] += value * input[j] + } + } + return output +} + +func assertClose(t *testing.T, name string, got, want []float32) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("%s length = %d, want %d", name, len(got), len(want)) + } + for i := range want { + if math.Abs(float64(got[i]-want[i])) > 1e-5 { + t.Fatalf("%s[%d] = %v, want %v", name, i, got[i], want[i]) + } + } +} diff --git a/x/models/qwen4_exp/qsa.go b/x/models/qwen4_exp/qsa.go new file mode 100644 index 000000000..f70040082 --- /dev/null +++ b/x/models/qwen4_exp/qsa.go @@ -0,0 +1,204 @@ +package qwen4_exp + +import ( + "math" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" +) + +// canonicalRopePositionRows returns exact [B, 1, L, 3] position rows. MRoPE +// input is the channel-major [B][3][L] layout used by the vision path. +func canonicalRopePositionRows(b *batch.Batch, L int32, mrope []int32) *mlx.Array { + B := len(b.SeqOffsets) + rows := make([]int32, B*int(L)*3) + for r, base := range b.SeqOffsets { + for i := range int(L) { + row := (r*int(L) + i) * 3 + if mrope == nil { + p := base + int32(i) + rows[row], rows[row+1], rows[row+2] = p, p, p + continue + } + channel := mrope[r*3*int(L):] + rows[row] = channel[i] + rows[row+1] = channel[int(L)+i] + rows[row+2] = channel[2*int(L)+i] + } + } + return mlx.FromValues(rows, B, 1, int(L), 3) +} + +// qsaRopeCosSin builds the main attention's exact 1D/MRoPE tables from +// arbitrary cached positions. positions is [B, 1, L, 3]. +func qsaRopeCosSin(cfg *Config, positions *mlx.Array) (cos, sin *mlx.Array) { + half := int(cfg.RopeDim) / 2 + selector := make([]int32, half) + sections := cfg.RopeParameters.MRoPESection + if len(sections) == 3 && cfg.RopeParameters.MRoPEInterleaved { + for channel := int32(1); channel <= 2; channel++ { + limit := min(int(sections[channel])*3, half) + for i := int(channel); i < limit; i += 3 { + selector[i] = channel + } + } + } + + invFreq := make([]float32, half) + for i := range invFreq { + invFreq[i] = float32(1 / math.Pow(float64(cfg.RopeParameters.RopeTheta), float64(2*i)/float64(cfg.RopeDim))) + } + selected := mlx.Take(positions, mlx.FromValues(selector, half), 3).AsType(mlx.DTypeFloat32) + angles := mlx.Mul(selected, mlx.FromValues(invFreq, 1, 1, 1, half)) + embedding := mlx.Concatenate([]*mlx.Array{angles, angles}, -1) + return mlx.Cos(embedding), mlx.Sin(embedding) +} + +func applyQwenRoPE(x, cos, sin *mlx.Array, ropeDim int32) *mlx.Array { + headDim := int32(x.Dim(3)) + rot := mlx.SliceStartStop(x, + []int32{0, 0, 0, 0}, + []int32{int32(x.Dim(0)), int32(x.Dim(1)), int32(x.Dim(2)), ropeDim}, + ) + half := ropeDim / 2 + left := mlx.SliceStartStop(rot, + []int32{0, 0, 0, 0}, + []int32{int32(rot.Dim(0)), int32(rot.Dim(1)), int32(rot.Dim(2)), half}, + ) + right := mlx.SliceStartStop(rot, + []int32{0, 0, 0, half}, + []int32{int32(rot.Dim(0)), int32(rot.Dim(1)), int32(rot.Dim(2)), ropeDim}, + ) + rotated := mlx.Concatenate([]*mlx.Array{mlx.Neg(right), left}, -1) + out := mlx.Add( + mlx.Mul(rot, cos.AsType(x.DType())), + mlx.Mul(rotated, sin.AsType(x.DType())), + ) + if ropeDim == headDim { + return out + } + pass := mlx.SliceStartStop(x, + []int32{0, 0, 0, ropeDim}, + []int32{int32(x.Dim(0)), int32(x.Dim(1)), int32(x.Dim(2)), headDim}, + ) + return mlx.Concatenate([]*mlx.Array{out, pass}, -1) +} + +func qsaCompressedKeys(rawKeys, positions *mlx.Array, indexer *attentionIndexer, cfg *Config) *mlx.Array { + ratio := cfg.IndexerCompressRatio + groups := int32(rawKeys.Dim(2)) / ratio + complete := groups * ratio + rawKeys = mlx.SliceStartStop(rawKeys, + []int32{0, 0, 0, 0}, + []int32{int32(rawKeys.Dim(0)), int32(rawKeys.Dim(1)), complete, int32(rawKeys.Dim(3))}, + ) + rawKeys = mlx.Reshape(rawKeys, int32(rawKeys.Dim(0)), int32(rawKeys.Dim(1)), groups, ratio, cfg.IndexerHeadDim) + pooled := mlx.Mean(rawKeys.AsType(mlx.DTypeFloat32), 3, false).AsType(rawKeys.DType()) + pooled = indexer.KNorm.Forward(pooled, cfg.RMSNormEps) + + positions = mlx.SliceStartStop(positions, + []int32{0, 0, 0, 0}, + []int32{int32(positions.Dim(0)), int32(positions.Dim(1)), complete, int32(positions.Dim(3))}, + ) + positions = mlx.Reshape(positions, int32(positions.Dim(0)), int32(positions.Dim(1)), groups, ratio, 3) + positions = mlx.Squeeze(mlx.SliceStartStop(positions, + []int32{0, 0, 0, 0, 0}, + []int32{int32(positions.Dim(0)), int32(positions.Dim(1)), groups, 1, 3}, + ), 3) + cos, sin := qsaRopeCosSin(cfg, positions) + return applyQwenRoPE(pooled, cos, sin, cfg.RopeDim) +} + +// qsaLogicalIndices reproduces the reference's compressed-block top-k and +// incomplete-group tail. Invalid rows are kept as zero indices with a mask. +func qsaLogicalIndices(scores *mlx.Array, b *batch.Batch, keyLength int32, cfg *Config) (indices, valid *mlx.Array) { + B, L := int32(scores.Dim(0)), int32(scores.Dim(1)) + ratio := cfg.IndexerCompressRatio + blocks := keyLength / ratio + blockIDs := mlx.Reshape(mlx.Arange(0, float64(blocks), 1, mlx.DTypeInt32), 1, 1, blocks) + + visibleValues := make([]int32, int(B*L)) + for r, offset := range b.SeqOffsets { + for i := range int(L) { + visibleValues[r*int(L)+i] = offset + int32(i) + 1 + } + } + visibleTokens := mlx.FromValues(visibleValues, int(B), int(L), 1) + visibleBlocks := mlx.FloorDivideScalar(visibleTokens, ratio) + blockValid := blockIDs.Less(visibleBlocks) + + selectedCount := min(blocks, cfg.IndexerBudget/ratio) + if blocks > selectedCount { + fill := mlx.AddScalar(mlx.Zeros(scores.DType(), int(B), int(L), int(blocks)), -float32(math.MaxFloat32)) + masked := mlx.Where(blockValid, scores, fill) + selected := mlx.Argpartition(mlx.Neg(masked), int(selectedCount)-1, -1) + indices = mlx.SliceStartStop(selected, + []int32{0, 0, 0}, + []int32{B, L, selectedCount}, + ) + } else { + indices = mlx.Tile(blockIDs, []int32{B, L, 1}) + } + blockValid = indices.Less(visibleBlocks) + + offsets := mlx.Reshape(mlx.Arange(0, float64(ratio), 1, mlx.DTypeInt32), 1, 1, 1, ratio) + expanded := mlx.Add(mlx.MulScalar(mlx.ExpandDims(indices, -1), float32(ratio)), offsets) + expandedValid := mlx.Mul( + mlx.ExpandDims(blockValid.AsType(mlx.DTypeInt32), -1), + expanded.Less(mlx.ExpandDims(visibleTokens, -1)).AsType(mlx.DTypeInt32), + ) + expanded = mlx.Reshape(expanded, B, L, selectedCount*ratio) + expandedValid = mlx.Reshape(expandedValid, B, L, selectedCount*ratio) + + // The current partial compression group is always attended causally. + tailWidth := ratio - 1 + tailOffsets := mlx.Reshape(mlx.Arange(0, float64(tailWidth), 1, mlx.DTypeInt32), 1, 1, tailWidth) + tailStart := mlx.MulScalar(mlx.FloorDivideScalar(visibleTokens, ratio), float32(ratio)) + tail := mlx.Add(tailStart, tailOffsets) + tailCount := mlx.Sub(visibleTokens, tailStart) + tailValid := tailOffsets.Less(tailCount).AsType(mlx.DTypeInt32) + + indices = mlx.Concatenate([]*mlx.Array{expanded, tail}, -1) + validInt := mlx.Concatenate([]*mlx.Array{expandedValid, tailValid}, -1) + valid = validInt.Greater(mlx.Zeros(mlx.DTypeInt32, int(B), int(L), validInt.Dim(2))) + indices = mlx.Where(valid, indices, mlx.Zeros(mlx.DTypeInt32, int(B), int(L), indices.Dim(2))) + return indices, valid +} + +func qsaSparseAttention(q *mlx.Array, history *nn.KVHistory, indices, valid *mlx.Array, cfg *Config) *mlx.Array { + outputType := q.DType() + B, queryHeads, L, D := int32(q.Dim(0)), int32(q.Dim(1)), int32(q.Dim(2)), int32(q.Dim(3)) + kvHeads := cfg.NumKeyValueHeads + repeats := queryHeads / kvHeads + + k := qsaGatherHistory(history.K(), indices) + v := qsaGatherHistory(history.V(), indices) + + qr := mlx.Reshape(q, B, kvHeads, repeats, L, 1, D) + kr := mlx.Transpose(mlx.ExpandDims(k, 2), 0, 1, 2, 3, 5, 4) + scores := mlx.Squeeze(mlx.Matmul(qr.AsType(mlx.DTypeFloat32), kr.AsType(mlx.DTypeFloat32)), 4) + scores = mlx.MulScalar(scores, cfg.Scale) + mask := mlx.ExpandDims(mlx.ExpandDims(valid, 1), 1) + fill := mlx.AddScalar(mlx.Zeros(scores.DType(), scores.Dims()...), -float32(math.MaxFloat32)) + scores = mlx.Where(mask, scores, fill) + probs := mlx.SoftmaxAxis(scores, -1, true) + + vr := mlx.ExpandDims(v.AsType(mlx.DTypeFloat32), 2) + out := mlx.Matmul(mlx.ExpandDims(probs, 4), vr) + out = mlx.Squeeze(out, 4) + return mlx.Reshape(out, B, queryHeads, L, D).AsType(outputType) +} + +// qsaGatherHistory selects each batch row's logical token indices from a +// [B, H, K, D] cache history and returns [B, H, L, S, D]. +func qsaGatherHistory(history, indices *mlx.Array) *mlx.Array { + B, H, K, D := int32(history.Dim(0)), int32(history.Dim(1)), int32(history.Dim(2)), int32(history.Dim(3)) + offsets := make([]int32, B) + for i := range offsets { + offsets[i] = int32(i) * K + } + logical := mlx.Add(indices, mlx.FromValues(offsets, int(B), 1, 1)) + flattened := mlx.Reshape(mlx.Transpose(history, 1, 0, 2, 3), H, B*K, D) + return mlx.Transpose(mlx.Take(flattened, logical, 1), 1, 0, 2, 3, 4) +} diff --git a/x/models/qwen4_exp/qsa_test.go b/x/models/qwen4_exp/qsa_test.go new file mode 100644 index 000000000..9bf367c9a --- /dev/null +++ b/x/models/qwen4_exp/qsa_test.go @@ -0,0 +1,125 @@ +package qwen4_exp + +import ( + "math" + "slices" + "testing" + + "github.com/ollama/ollama/x/internal/mlxtest" + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" +) + +func TestQSASelectsCompressedBlocksAndCausalTail(t *testing.T) { + mlxtest.Setup(t) + cfg := &Config{IndexerBudget: 8, IndexerCompressRatio: 4} + scores := mlx.FromValues([]float32{0.1, 4, 2, 3}, 1, 1, 4) + b := &batch.Batch{SeqOffsets: []int32{16}} + + indices, valid := qsaLogicalIndices(scores, b, 16, cfg) + indices = indices.AsType(mlx.DTypeInt32) + valid = valid.AsType(mlx.DTypeInt32) + mlx.Eval(indices, valid) + + values, mask := indices.Ints(), valid.Ints() + selected := make([]int32, 0, len(values)) + for i, value := range values { + if mask[i] != 0 { + selected = append(selected, value) + } + } + slices.Sort(selected) + want := []int32{4, 5, 6, 7, 12, 13, 14, 15, 16} + if !slices.Equal(selected, want) { + t.Fatalf("selected indices = %v, want %v", selected, want) + } +} + +func TestQSASelectionMasksFutureBlocks(t *testing.T) { + mlxtest.Setup(t) + cfg := &Config{IndexerBudget: 8, IndexerCompressRatio: 4} + // Only block 0 and token 4 are visible. Give every future block a much + // larger score so the test fails if selection sees cached-but-causal-junk. + scores := mlx.FromValues([]float32{0.1, 100, 90, 80, 70}, 1, 1, 5) + b := &batch.Batch{SeqOffsets: []int32{4}} + + indices, valid := qsaLogicalIndices(scores, b, 20, cfg) + indices = indices.AsType(mlx.DTypeInt32) + valid = valid.AsType(mlx.DTypeInt32) + mlx.Eval(indices, valid) + + var selected []int32 + for i, value := range indices.Ints() { + if valid.Ints()[i] != 0 { + selected = append(selected, value) + } + } + slices.Sort(selected) + if want := []int32{0, 1, 2, 3, 4}; !slices.Equal(selected, want) { + t.Fatalf("selected indices = %v, want %v", selected, want) + } +} + +func TestQSASparseAttentionMatchesReference(t *testing.T) { + mlxtest.Setup(t) + cfg := &Config{NumKeyValueHeads: 1, Scale: 1} + q := mlx.FromValues([]float32{1, 0}, 1, 1, 1, 2) + k := mlx.FromValues([]float32{1, 0, 0, 1, 2, 0}, 1, 1, 3, 2) + v := mlx.FromValues([]float32{10, 1, 20, 2, 30, 3}, 1, 1, 3, 2) + indices := mlx.FromValues([]int32{2, 0}, 1, 1, 2) + valid := mlx.FromValues([]bool{true, true}, 1, 1, 2) + + out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg) + out = out.AsType(mlx.DTypeFloat32) + mlx.Eval(out) + got := out.Floats() + p2 := float32(math.Exp(2) / (math.Exp(2) + math.Exp(1))) + want := []float32{p2*30 + (1-p2)*10, p2*3 + (1-p2)*1} + for i := range want { + if math.Abs(float64(got[i]-want[i])) > 1e-4 { + t.Fatalf("sparse attention[%d] = %v, want %v", i, got[i], want[i]) + } + } +} + +func TestQSASparseAttentionIgnoresInvalidRows(t *testing.T) { + mlxtest.Setup(t) + cfg := &Config{NumKeyValueHeads: 1, Scale: 1} + q := mlx.FromValues([]float32{1, 0}, 1, 1, 1, 2) + // Row 0 is intentionally dominant junk. Only row 1 is logically valid. + k := mlx.FromValues([]float32{100, 0, 1, 0}, 1, 1, 2, 2) + v := mlx.FromValues([]float32{999, 999, 7, 3}, 1, 1, 2, 2) + indices := mlx.FromValues([]int32{0, 1}, 1, 1, 2) + valid := mlx.FromValues([]bool{false, true}, 1, 1, 2) + + out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg) + out = out.AsType(mlx.DTypeFloat32) + mlx.Eval(out) + if got, want := out.Floats(), []float32{7, 3}; !slices.Equal(got, want) { + t.Fatalf("sparse attention = %v, want %v", got, want) + } +} + +func TestQSASparseAttentionKeepsBatchRowsIndependent(t *testing.T) { + mlxtest.Setup(t) + cfg := &Config{NumKeyValueHeads: 1, Scale: 1} + q := mlx.FromValues([]float32{1, 0, 1, 0}, 2, 1, 1, 2) + k := mlx.FromValues([]float32{ + 1, 0, 0, 1, + 1, 0, 0, 1, + }, 2, 1, 2, 2) + v := mlx.FromValues([]float32{ + 10, 1, 20, 2, + 30, 3, 40, 4, + }, 2, 1, 2, 2) + indices := mlx.FromValues([]int32{0, 1}, 2, 1, 1) + valid := mlx.FromValues([]bool{true, true}, 2, 1, 1) + + out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg) + out = out.AsType(mlx.DTypeFloat32) + mlx.Eval(out) + if got, want := out.Floats(), []float32{10, 1, 40, 4}; !slices.Equal(got, want) { + t.Fatalf("sparse attention = %v, want %v", got, want) + } +} diff --git a/x/models/qwen4_exp/qwen4_exp.go b/x/models/qwen4_exp/qwen4_exp.go new file mode 100644 index 000000000..069ca2eba --- /dev/null +++ b/x/models/qwen4_exp/qwen4_exp.go @@ -0,0 +1,244 @@ +// Package qwen4_exp implements the Qwen 4 MLX model family. +package qwen4_exp + +import ( + "fmt" + + "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/models/qwen3_5" + "github.com/ollama/ollama/x/tokenizer" +) + +func init() { + base.Register("Qwen4ExpForConditionalGeneration", NewModel) +} + +var ( + _ base.Model = (*Model)(nil) + _ base.SelfDraft = (*Model)(nil) + _ base.DraftModel = (*mtpDraft)(nil) + _ base.MediaModel = (*Model)(nil) +) + +// Model carries the validated architecture and tokenizer. +type Model struct { + EmbedTokens nn.EmbeddingLayer + Layers []*Layer + OutputMixer *hyperConnection + LMHead nn.LinearLayer + MTP *MTPHead + Vision *qwen3_5.VisionAdapter + + *Config + tok *tokenizer.Tokenizer + quantBits int + quantGroup int + quantMode string + tensorQuant map[string]*model.TensorQuantInfo + configData []byte +} + +// NewModel validates the publisher config before constructing runtime state. +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 + } + + tokenizerData, err := root.Manifest.ReadConfig("tokenizer.json") + if err != nil { + return nil, fmt.Errorf("load tokenizer: %w", err) + } + tokenizerConfig := &tokenizer.TokenizerConfig{ConfigJSON: configData} + if data, err := root.Manifest.ReadConfig("generation_config.json"); err == nil { + tokenizerConfig.GenerationConfigJSON = data + } + if data, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil { + tokenizerConfig.TokenizerConfigJSON = data + } + tok, err := tokenizer.LoadFromBytesWithConfig(tokenizerData, tokenizerConfig) + if err != nil { + return nil, fmt.Errorf("parse tokenizer: %w", err) + } + group, bits, mode := model.QuantizationParams(root.QuantType()) + if root.GroupSize() > 0 { + group = root.GroupSize() + } + return &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + tok: tok, + quantBits: bits, + quantGroup: group, + quantMode: mode, + tensorQuant: root.AllTensorQuant(), + configData: configData, + }, nil +} + +func (m *Model) NewCaches() []cache.Cache { + fullLayers := 0 + for _, layer := range m.Layers { + if layer.Attention != nil { + fullLayers++ + } + } + caches := make([]cache.Cache, 0, len(m.Layers)+fullLayers+1) + convTail := m.LinearConvKernelDim - 1 + convDim := 2*m.LinearNumKeyHeads*m.LinearKeyHeadDim + m.LinearNumValueHeads*m.LinearValueHeadDim + for i := range m.Layers { + if m.layerIsLinear(i) { + caches = append(caches, cache.NewRecurrentCache(convTail, convDim, m.LinearNumValueHeads, m.LinearValueHeadDim, m.LinearKeyHeadDim)) + } else { + caches = append(caches, cache.NewKVCache()) + } + } + for range fullLayers { + // Full-history raw QSA keys plus exact 3-axis positions. This is the + // simple semantic implementation; a compressed streaming cache can + // replace it without changing the attention contract. + caches = append(caches, cache.NewKVCache()) + } + caches = append(caches, newEngramCache( + int(m.NGramSize-1), + int((m.PLEConvKernelSize-1)*m.NGramSize), + int(m.HCCount*m.HiddenSize), + m.EOSTokenID, + )) + return caches +} + +func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (*mlx.Array, *mlx.Array) { + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + L := int32(b.InputIDs.Dim(1)) + var mrope []int32 + if m.Vision != nil { + mrope = m.Vision.RopePositions(b, L) + } + ropePositions := canonicalRopePositionRows(b, L, mrope) + embedding := m.EmbedTokens.Forward(b.InputIDs) + if len(b.Media) > 0 { + embedding = m.Vision.ScatterMedia(embedding, b, 0) + } + hidden := expandStreams(embedding, m.Config) + + var tokens *engramCache + pleCacheIndex := len(m.Layers) + for _, layer := range m.Layers { + if layer.Attention != nil { + pleCacheIndex++ + } + } + if len(caches) > pleCacheIndex { + tokens, _ = caches[pleCacheIndex].(*engramCache) + } + if tokens == nil { + tokens = newEngramCache( + int(m.NGramSize-1), + int((m.PLEConvKernelSize-1)*m.NGramSize), + int(m.HCCount*m.HiddenSize), + m.EOSTokenID, + ) + } + + for i, layer := range m.Layers { + if layer.PLE != nil { + hidden = mlx.Add(hidden, layer.PLE.Forward(hidden, b, tokens, m.Config)) + } + var c cache.Cache + if i < len(caches) { + c = caches[i] + } + var side cache.Cache + if layer.Attention != nil && layer.Attention.SideCacheIndex < len(caches) { + side = caches[layer.Attention.SideCacheIndex] + } + hidden = layer.Forward(hidden, b, c, side, positions, ropePositions, m.Config) + } + multiHidden := hidden + hidden = m.OutputMixer.Reduce(multiHidden, m.Config) + return hidden, multiHidden +} + +func (m *Model) Unembed(hidden *mlx.Array) *mlx.Array { + return m.LMHead.Forward(hidden) +} + +func (m *Model) Tokenizer() *tokenizer.Tokenizer { return m.tok } + +func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { + if m.Vision == nil { + return nil, fmt.Errorf("this model does not support media input") + } + return m.Vision.PrepareMedia(segments) +} + +func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { + return m.Vision.EncodeMedia(item, data) +} + +func (m *Model) MaxContextLength() int { + return int(m.MaxPositionEmbeddings) +} + +// mtpDraft views the checkpoint's inline predictor through DraftModel. +type mtpDraft Model + +func (m *Model) SelfDraft() base.DraftModel { + if m.MTP == nil { + return nil + } + return (*mtpDraft)(m) +} + +func (m *mtpDraft) NewCaches() []cache.Cache { + if m.MTP == nil { + return nil + } + return []cache.Cache{cache.NewKVCache(), cache.NewKVCache()} +} + +func (m *mtpDraft) LoadWeights(map[string]*mlx.Array) error { return nil } + +func (m *mtpDraft) Unembed(x *mlx.Array) *mlx.Array { return (*Model)(m).Unembed(x) } + +func (m *mtpDraft) Forward(b *batch.Batch, _, draftCaches []cache.Cache) (hidden, auxHidden *mlx.Array) { + B, L := int32(b.InputIDs.Dim(0)), int32(b.InputIDs.Dim(1)) + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + var mrope []int32 + if m.Vision != nil { + mrope = m.Vision.RopePositions(b, L) + } + ropePositions := canonicalRopePositionRows(b, L, mrope) + + embedding := m.EmbedTokens.Forward(b.InputIDs) + // The reference MTP is text-only: visual information reaches it through + // the target's multi-stream hidden state, while this branch embeds raw IDs. + embedding = m.MTP.FCEmbedding.Forward(m.MTP.EmbeddingNorm.Forward(embedding, m.RMSNormEps)) + + backbone := m.MTP.HiddenNorm.Forward(b.Hidden, m.RMSNormEps) + backbone = mlx.Reshape(backbone, B, L, m.HCCount, m.HiddenSize) + backbone = m.MTP.FCHidden.Forward(backbone) + backbone = mlx.Add(backbone, mlx.ExpandDims(embedding, -2)) + backbone = mlx.Reshape(backbone, B, L, m.HCCount*m.HiddenSize) + + var main, side cache.Cache + if len(draftCaches) > 0 { + main = draftCaches[0] + } + if len(draftCaches) > 1 { + side = draftCaches[1] + } + multi := m.MTP.Layer.Forward(backbone, b, main, side, positions, ropePositions, m.Config) + sample := m.MTP.Mixer.Reduce(multi, m.Config) + return sample, multi +} diff --git a/x/models/qwen4_exp/weights.go b/x/models/qwen4_exp/weights.go new file mode 100644 index 000000000..eece909c9 --- /dev/null +++ b/x/models/qwen4_exp/weights.go @@ -0,0 +1,462 @@ +package qwen4_exp + +import ( + "fmt" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + "github.com/ollama/ollama/x/models/nn" + "github.com/ollama/ollama/x/models/qwen3_5" +) + +type hyperConnection struct { + Norm *streamRMSNorm + InputMixDown nn.LinearLayer + InputMixUp nn.LinearLayer + BlockInject nn.LinearLayer +} + +type linearAttention struct { + QKV nn.LinearLayer + Z nn.LinearLayer + B nn.LinearLayer + A nn.LinearLayer + Out nn.LinearLayer + + Conv *nn.Conv1d + NormWeight *mlx.Array + DtBias *mlx.Array + ALog *mlx.Array + AExp *mlx.Array +} + +type attentionIndexer struct { + QKProj nn.LinearLayer + QNorm *nn.RMSNorm + KNorm *nn.RMSNorm +} + +type fullAttention struct { + QProj nn.LinearLayer + KProj nn.LinearLayer + VProj nn.LinearLayer + OProj nn.LinearLayer + QNorm *nn.RMSNorm + KNorm *nn.RMSNorm + + Indexer *attentionIndexer + + SideCacheIndex int +} + +type sparseMoE struct { + Gate nn.LinearLayer + GateUpExperts *mlx.Array + GateUpScales *mlx.Array + GateUpBiases *mlx.Array + DownExperts *mlx.Array + DownScales *mlx.Array + DownBiases *mlx.Array + GateUpGroup int + GateUpBits int + GateUpMode string + DownGroup int + DownBits int + DownMode string + SharedGate nn.LinearLayer + SharedGateProj nn.LinearLayer + SharedUpProj nn.LinearLayer + SharedDownProj nn.LinearLayer +} + +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 loadPackedExperts(tensors map[string]*mlx.Array, base string, m *Model) (weight, scales, biases *mlx.Array, group, bits int, mode string, err error) { + var key string + for _, candidate := range []string{base, base + ".weight"} { + if weight = tensors[candidate]; weight != nil { + key = candidate + break + } + } + if weight == nil { + return nil, nil, nil, 0, 0, "", fmt.Errorf("missing packed expert tensor %q", base) + } + scales = tensors[key+"_scale"] + if scales == nil { + return mlx.Transpose(weight, 0, 2, 1), nil, nil, 0, 0, "", nil + } + biases = tensors[key+"_qbias"] + group, bits, mode = model.ResolveLinearQuantParams(m.quantGroup, m.quantBits, m.quantMode, m.tensorQuant, key, weight, scales) + if !supportsGatherQMM(mode, bits) { + return nil, nil, nil, 0, 0, "", fmt.Errorf("packed expert tensor %q uses unsupported quantization mode=%q bits=%d", key, mode, bits) + } + return weight, scales, biases, group, bits, mode, nil +} + +// PLE is the model's n-gram embedding block. The checkpoint stores one logical +// table in split vocabulary shards. +type PLE struct { + Conv *nn.Conv1d + KeyProj nn.LinearLayer + ValueProj nn.LinearLayer + NormConv *streamRMSNorm + NormKey *streamRMSNorm + NormQuery *streamRMSNorm + + LayerMultipliers *mlx.Array + HeadOffsets *mlx.Array + HeadVocabSizes *mlx.Array + EmbeddingShards []nn.EmbeddingLayer + ShardRows int +} + +type streamRMSNorm struct { + Weight *mlx.Array +} + +type Layer struct { + AttentionConnection *hyperConnection + MLPConnection *hyperConnection + Linear *linearAttention + Attention *fullAttention + MoE *sparseMoE + PLE *PLE +} + +type MTPHead struct { + EmbeddingNorm *nn.RMSNorm + HiddenNorm *nn.RMSNorm + FCEmbedding nn.LinearLayer + FCHidden nn.LinearLayer + Mixer *hyperConnection + Layer *Layer +} + +func requiredArray(tensors map[string]*mlx.Array, names ...string) (*mlx.Array, error) { + for _, name := range names { + if value := tensors[name]; value != nil { + return value, nil + } + } + return nil, fmt.Errorf("missing tensor %q", names[0]) +} + +func requiredLinear(linears model.LinearFactory, name string) (nn.LinearLayer, error) { + value := linears.Make(name) + if value == nil { + return nil, fmt.Errorf("missing linear %q", name) + } + return value, nil +} + +func loadNorm(tensors map[string]*mlx.Array, name string, eps float32) (*nn.RMSNorm, error) { + weight, err := requiredArray(tensors, name+".weight", name) + if err != nil { + return nil, err + } + // Qwen 4 RMSNorm weights are stored zero-centered. + return nn.NewRMSNorm(mlx.AddScalar(weight, 1), eps), nil +} + +func loadStreamNorm(tensors map[string]*mlx.Array, name string, cfg *Config) (*streamRMSNorm, error) { + weight, err := requiredArray(tensors, name+".weight", name) + if err != nil { + return nil, err + } + if weight.Size() != int(cfg.HCCount*cfg.HiddenSize) { + return nil, fmt.Errorf("%s has %d elements, want %d", name, weight.Size(), cfg.HCCount*cfg.HiddenSize) + } + weight = mlx.Reshape(mlx.AddScalar(weight, 1), cfg.HCCount, cfg.HiddenSize) + return &streamRMSNorm{Weight: weight}, nil +} + +func loadHyperConnection(linears model.LinearFactory, tensors map[string]*mlx.Array, prefix string, withInjection bool, cfg *Config) (*hyperConnection, error) { + h := &hyperConnection{} + var err error + if h.Norm, err = loadStreamNorm(tensors, prefix+".hc_norm", cfg); err != nil { + return nil, err + } + if h.InputMixDown, err = requiredLinear(linears, prefix+".input_mix_weight_down"); err != nil { + return nil, err + } + if h.InputMixUp, err = requiredLinear(linears, prefix+".input_mix_weight_up"); err != nil { + return nil, err + } + if withInjection { + if h.BlockInject, err = requiredLinear(linears, prefix+".block_inject_weight"); err != nil { + return nil, err + } + } + return h, nil +} + +func loadLinearAttention(linears model.LinearFactory, tensors map[string]*mlx.Array, prefix string) (*linearAttention, error) { + a := &linearAttention{} + var err error + for name, dst := range map[string]*nn.LinearLayer{ + "in_proj_qkv": &a.QKV, + "in_proj_z": &a.Z, + "in_proj_b": &a.B, + "in_proj_a": &a.A, + "out_proj": &a.Out, + } { + *dst, err = requiredLinear(linears, prefix+"."+name) + if err != nil { + return nil, err + } + } + convWeight, err := requiredArray(tensors, prefix+".conv1d.weight", prefix+".conv1d") + if err != nil { + return nil, err + } + if convWeight.NumDims() == 3 && convWeight.Dim(1) == 1 { + convWeight = mlx.Reshape(convWeight, int32(convWeight.Dim(0)), int32(convWeight.Dim(2))) + } + if convWeight.NumDims() != 2 { + return nil, fmt.Errorf("%s.conv1d weight must reduce to 2D, got %dD", prefix, convWeight.NumDims()) + } + a.Conv = nn.NewConv1d(mlx.ExpandDims(convWeight, 2), nil, 1, 0, 1, int32(convWeight.Dim(0))) + if a.NormWeight, err = requiredArray(tensors, prefix+".norm.weight", prefix+".norm"); err != nil { + return nil, err + } + if a.DtBias, err = requiredArray(tensors, prefix+".dt_bias", prefix+".dt_proj"); err != nil { + return nil, err + } + if a.ALog, err = requiredArray(tensors, prefix+".A_log", prefix+".a_log"); err != nil { + return nil, err + } + a.AExp = mlx.Exp(a.ALog.AsType(mlx.DTypeFloat32)) + return a, nil +} + +func loadFullAttention(linears model.LinearFactory, tensors map[string]*mlx.Array, prefix string, eps float32) (*fullAttention, error) { + a := &fullAttention{} + var err error + for name, dst := range map[string]*nn.LinearLayer{ + "q_proj": &a.QProj, + "k_proj": &a.KProj, + "v_proj": &a.VProj, + "o_proj": &a.OProj, + } { + *dst, err = requiredLinear(linears, prefix+"."+name) + if err != nil { + return nil, err + } + } + if a.QNorm, err = loadNorm(tensors, prefix+".q_norm", eps); err != nil { + return nil, err + } + if a.KNorm, err = loadNorm(tensors, prefix+".k_norm", eps); err != nil { + return nil, err + } + + idx := &attentionIndexer{} + if idx.QKProj, err = requiredLinear(linears, prefix+".indexer.index_qk_proj"); err != nil { + return nil, err + } + if idx.QNorm, err = loadNorm(tensors, prefix+".indexer.q_layernorm", eps); err != nil { + return nil, err + } + if idx.KNorm, err = loadNorm(tensors, prefix+".indexer.k_layernorm", eps); err != nil { + return nil, err + } + a.Indexer = idx + return a, nil +} + +func (m *Model) loadSparseMoE(linears model.LinearFactory, tensors map[string]*mlx.Array, prefix string) (*sparseMoE, error) { + moe := &sparseMoE{} + var err error + if moe.Gate, err = requiredLinear(linears, prefix+".gate"); err != nil { + return nil, err + } + if moe.GateUpExperts, moe.GateUpScales, moe.GateUpBiases, moe.GateUpGroup, moe.GateUpBits, moe.GateUpMode, err = loadPackedExperts(tensors, prefix+".experts.gate_up_proj", m); err != nil { + return nil, err + } + if moe.DownExperts, moe.DownScales, moe.DownBiases, moe.DownGroup, moe.DownBits, moe.DownMode, err = loadPackedExperts(tensors, prefix+".experts.down_proj", m); err != nil { + return nil, err + } + for name, dst := range map[string]*nn.LinearLayer{ + "shared_expert_gate": &moe.SharedGate, + "shared_expert.gate_proj": &moe.SharedGateProj, + "shared_expert.up_proj": &moe.SharedUpProj, + "shared_expert.down_proj": &moe.SharedDownProj, + } { + *dst, err = requiredLinear(linears, prefix+"."+name) + if err != nil { + return nil, err + } + } + return moe, nil +} + +func (m *Model) loadPLE(linears model.LinearFactory, tensors map[string]*mlx.Array, prefix string) (*PLE, error) { + p := &PLE{} + var err error + convWeight, err := requiredArray(tensors, prefix+".conv1d.weight", prefix+".conv1d") + if err != nil { + return nil, err + } + if convWeight.NumDims() != 3 || convWeight.Dim(1) != 1 || convWeight.Dim(2) != int(m.PLEConvKernelSize) { + return nil, fmt.Errorf("%s.conv1d has shape %v, want [%d, 1, %d]", prefix, convWeight.Dims(), m.HCCount*m.HiddenSize, m.PLEConvKernelSize) + } + convWeight = mlx.Transpose(convWeight, 0, 2, 1) + p.Conv = nn.NewConv1d(convWeight, nil, 1, 0, m.NGramSize, m.HCCount*m.HiddenSize) + if p.KeyProj, err = requiredLinear(linears, prefix+".key_proj"); err != nil { + return nil, err + } + if p.ValueProj, err = requiredLinear(linears, prefix+".value_proj"); err != nil { + return nil, err + } + if p.NormConv, err = loadStreamNorm(tensors, prefix+".norm_conv", m.Config); err != nil { + return nil, err + } + if p.NormKey, err = loadStreamNorm(tensors, prefix+".norm_key", m.Config); err != nil { + return nil, err + } + if p.NormQuery, err = loadStreamNorm(tensors, prefix+".norm_query", m.Config); err != nil { + return nil, err + } + if p.LayerMultipliers, err = requiredArray(tensors, prefix+".ple_embedding.layer_multipliers"); err != nil { + return nil, err + } + if p.HeadOffsets, err = requiredArray(tensors, prefix+".ple_embedding.ngram_heads_offsets"); err != nil { + return nil, err + } + if p.HeadVocabSizes, err = requiredArray(tensors, prefix+".ple_embedding.ngram_heads_vocab_sizes"); err != nil { + return nil, err + } + + p.EmbeddingShards = make([]nn.EmbeddingLayer, m.SplitNGramParts) + for i := range m.SplitNGramParts { + path := fmt.Sprintf("%s.ple_embedding.ngram_embedding.shard_%d", prefix, i) + weight, err := requiredArray(tensors, path+".weight") + if err != nil { + return nil, err + } + if i == 0 { + p.ShardRows = weight.Dim(0) + } else if weight.Dim(0) != p.ShardRows { + return nil, fmt.Errorf("%s has %d rows, want %d", path, weight.Dim(0), p.ShardRows) + } + p.EmbeddingShards[i] = model.MakeEmbeddingLayer(tensors, path, m.quantGroup, m.quantBits, m.quantMode, m.tensorQuant) + if p.EmbeddingShards[i] == nil { + return nil, fmt.Errorf("missing embedding %q", path) + } + } + return p, nil +} + +func (m *Model) loadLayer(linears model.LinearFactory, tensors map[string]*mlx.Array, prefix, kind string, layer int) (*Layer, error) { + l := &Layer{} + var err error + if l.AttentionConnection, err = loadHyperConnection(linears, tensors, prefix+".attn_hyper_connection", true, m.Config); err != nil { + return nil, err + } + if l.MLPConnection, err = loadHyperConnection(linears, tensors, prefix+".mlp_hyper_connection", true, m.Config); err != nil { + return nil, err + } + if kind == "linear_attention" { + if l.Linear, err = loadLinearAttention(linears, tensors, prefix+".linear_attn"); err != nil { + return nil, err + } + } else if l.Attention, err = loadFullAttention(linears, tensors, prefix+".self_attn", m.RMSNormEps); err != nil { + return nil, err + } + if l.MoE, err = m.loadSparseMoE(linears, tensors, prefix+".mlp"); err != nil { + return nil, err + } + + for _, id := range m.PLELayerIDs { + // PLE layer IDs are one-based and identify the layer receiving the + // additive n-gram features before its attention branch. + if int(id-1) == layer { + l.PLE, err = m.loadPLE(linears, tensors, prefix+".ple") + if err != nil { + return nil, err + } + } + } + return l, nil +} + +func (m *Model) loadMTP(linears model.LinearFactory, tensors map[string]*mlx.Array) (*MTPHead, error) { + head := &MTPHead{} + var err error + if head.EmbeddingNorm, err = loadNorm(tensors, "mtp.pre_fc_norm_embedding", m.RMSNormEps); err != nil { + return nil, err + } + if head.HiddenNorm, err = loadNorm(tensors, "mtp.pre_fc_norm_hidden", m.RMSNormEps); err != nil { + return nil, err + } + if head.FCEmbedding, err = requiredLinear(linears, "mtp.fc_embedding"); err != nil { + return nil, err + } + if head.FCHidden, err = requiredLinear(linears, "mtp.fc_hidden"); err != nil { + return nil, err + } + if head.Mixer, err = loadHyperConnection(linears, tensors, "mtp.hyper_connection_mixer", false, m.Config); err != nil { + return nil, err + } + if head.Layer, err = m.loadLayer(linears, tensors, "mtp.layers.0", m.Config.MTP.LayerTypes[0], -1); err != nil { + return nil, err + } + return head, nil +} + +func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { + linears := model.NewLinearFactory(tensors, m.quantGroup, m.quantBits, m.quantMode, m.tensorQuant) + var err error + m.EmbedTokens = model.MakeEmbeddingLayer(tensors, "model.language_model.embed_tokens", m.quantGroup, m.quantBits, m.quantMode, m.tensorQuant) + if m.EmbedTokens == nil { + return fmt.Errorf("missing embedding weight") + } + if m.LMHead, err = requiredLinear(linears, "lm_head"); err != nil { + return err + } + if m.OutputMixer, err = loadHyperConnection(linears, tensors, "model.language_model.hyper_connection_mixer", false, m.Config); err != nil { + return err + } + + sideCacheIndex := len(m.Layers) + for i := range m.Layers { + prefix := fmt.Sprintf("model.language_model.layers.%d", i) + if m.Layers[i], err = m.loadLayer(linears, tensors, prefix, m.LayerTypes[i], i); err != nil { + return fmt.Errorf("load layer %d: %w", i, err) + } + if m.Layers[i].Attention != nil { + m.Layers[i].Attention.SideCacheIndex = sideCacheIndex + sideCacheIndex++ + } + } + if m.MTP, err = m.loadMTP(linears, tensors); err != nil { + return fmt.Errorf("load MTP: %w", err) + } + m.Vision, err = qwen3_5.NewVisionAdapter(m.configData, tensors, qwen3_5.VisionAdapterConfig{ + HiddenSize: m.HiddenSize, + RopeDim: m.RopeDim, + RopeTheta: m.RopeParameters.RopeTheta, + MropeSection: m.RopeParameters.MRoPESection, + QuantGroup: m.quantGroup, + QuantBits: m.quantBits, + QuantMode: m.quantMode, + TensorQuant: m.tensorQuant, + }) + if err != nil { + return fmt.Errorf("load vision tower: %w", err) + } + + return nil +}