mirror of
https://github.com/ollama/ollama.git
synced 2026-10-20 18:41:17 +00:00
MLX: Qwen3.8 Flash Next support (#18032)
* MLX: Qwen3.8 Flash Next support * review comments
This commit is contained in:
Daniel Hiltgen
GitHub
parent
e2c6c7e894
commit
18ea9de0fe
22 changed files with 2406 additions and 2 deletions
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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"]}`,
|
||||
|
|
|
|||
|
|
@ -483,6 +483,7 @@ var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{
|
|||
"Qwen3_5MoeForConditionalGeneration": newQwen35ImportTransform,
|
||||
"Qwen3NextMoeForCausalLM": newQwen35ImportTransform,
|
||||
"Qwen3NextMoeForConditionalGeneration": newQwen35ImportTransform,
|
||||
"Qwen4ExpForConditionalGeneration": newQwen4ExpImportTransform,
|
||||
"Gemma4ForCausalLM": newGemma4ImportTransform,
|
||||
"Gemma4ForConditionalGeneration": newGemma4ImportTransform,
|
||||
"Gemma4UnifiedForCausalLM": newGemma4ImportTransform,
|
||||
|
|
|
|||
37
x/create/qwen4_exp.go
Normal file
37
x/create/qwen4_exp.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
52
x/create/qwen4_exp_test.go
Normal file
52
x/create/qwen4_exp_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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]).
|
||||
|
|
|
|||
203
x/models/qwen4_exp/blocks.go
Normal file
203
x/models/qwen4_exp/blocks.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
238
x/models/qwen4_exp/config.go
Normal file
238
x/models/qwen4_exp/config.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
132
x/models/qwen4_exp/config_test.go
Normal file
132
x/models/qwen4_exp/config_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
117
x/models/qwen4_exp/engram.go
Normal file
117
x/models/qwen4_exp/engram.go
Normal file
|
|
@ -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())
|
||||
}
|
||||
205
x/models/qwen4_exp/engram_cache.go
Normal file
205
x/models/qwen4_exp/engram_cache.go
Normal file
|
|
@ -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) }
|
||||
105
x/models/qwen4_exp/engram_test.go
Normal file
105
x/models/qwen4_exp/engram_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
50
x/models/qwen4_exp/hyper_connection.go
Normal file
50
x/models/qwen4_exp/hyper_connection.go
Normal file
|
|
@ -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
|
||||
}
|
||||
104
x/models/qwen4_exp/hyper_connection_test.go
Normal file
104
x/models/qwen4_exp/hyper_connection_test.go
Normal file
|
|
@ -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])
|
||||
}
|
||||
}
|
||||
}
|
||||
204
x/models/qwen4_exp/qsa.go
Normal file
204
x/models/qwen4_exp/qsa.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
125
x/models/qwen4_exp/qsa_test.go
Normal file
125
x/models/qwen4_exp/qsa_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
244
x/models/qwen4_exp/qwen4_exp.go
Normal file
244
x/models/qwen4_exp/qwen4_exp.go
Normal file
|
|
@ -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
|
||||
}
|
||||
462
x/models/qwen4_exp/weights.go
Normal file
462
x/models/qwen4_exp/weights.go
Normal file
|
|
@ -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
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue