From 36914807bc9515bb90df92e1d6d541e5ae19f162 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Thu, 27 Aug 2026 11:08:49 -0700 Subject: [PATCH 1/2] MLX: reduce Qwen3.8 Flash Next memory use Preserve BF16 for accuracy-sensitive QSA projections and the non-expert MTP path while using MXFP8 for other eligible dense and control weights. Memory-map the PLE safetensors and transfer only selected packed rows into MLX. Pass compressed QSA block selections directly to a fused indexed-attention Metal kernel for long prompts, with the graph path retained for short or unsupported shapes. --- create/qwen4_exp.go | 41 ++- create/qwen4_exp_test.go | 53 +++- mlx/indexed_attention.go | 272 ++++++++++++++++++ mlx/indexed_attention_test.go | 59 ++++ mlxrunner/model/qwen4_exp/blocks.go | 4 +- mlxrunner/model/qwen4_exp/engram.go | 4 + mlxrunner/model/qwen4_exp/ple_host.go | 251 ++++++++++++++++ mlxrunner/model/qwen4_exp/ple_host_test.go | 90 ++++++ .../model/qwen4_exp/ple_mapping_other.go | 42 +++ mlxrunner/model/qwen4_exp/ple_mapping_unix.go | 55 ++++ mlxrunner/model/qwen4_exp/qsa.go | 81 ++---- mlxrunner/model/qwen4_exp/qsa_test.go | 70 ++--- mlxrunner/model/qwen4_exp/qwen4_exp.go | 6 + mlxrunner/model/qwen4_exp/weights.go | 7 + 14 files changed, 924 insertions(+), 111 deletions(-) create mode 100644 mlx/indexed_attention.go create mode 100644 mlx/indexed_attention_test.go create mode 100644 mlxrunner/model/qwen4_exp/ple_host.go create mode 100644 mlxrunner/model/qwen4_exp/ple_host_test.go create mode 100644 mlxrunner/model/qwen4_exp/ple_mapping_other.go create mode 100644 mlxrunner/model/qwen4_exp/ple_mapping_unix.go diff --git a/create/qwen4_exp.go b/create/qwen4_exp.go index 4113332fc..d6d09c5f4 100644 --- a/create/qwen4_exp.go +++ b/create/qwen4_exp.go @@ -12,11 +12,12 @@ func newQwen4ExpImportTransform(_ json.RawMessage) (quantizePolicy, error) { } func (qwen4ExpImportTransform) quantizationType(name string, shape []int32, quantize string) string { + quantize = normalizeQuantType(quantize) + // 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. + // embedding. Its lookup path dequantizes selected rows, so quantize its + // shards at the requested type. if qwen4ExpIsPLEEmbeddingShard(name) { - quantize = normalizeQuantType(quantize) if len(shape) == 2 && isAligned(shape, quantize) { return quantize } @@ -28,6 +29,25 @@ func (qwen4ExpImportTransform) quantizationType(name string, shape []int32, quan if qwen35IsLowRankProjection(name) { return "" } + + if quantize == "nvfp4" { + // Keep QSA Q/K/V and the indexer projection on the validated BF16 path. + // Keep the whole draft-side control path at source precision as well. + if qwen4ExpIsSensitiveQSA(name) || + (strings.HasPrefix(name, "mtp.") && !qwen4ExpIsRoutedExpert(name)) { + return "" + } + if isEmbedTokensWeight(name) { + return promoteEmbedding(shape, quantize) + } + + // Routed experts dominate the remaining model size and have validated NVFP4 + // kernels. Use MXFP8 for the other eligible target weights to avoid the + // long-context quality loss seen with blanket NVFP4. + if !qwen4ExpIsRoutedExpert(name) { + return GetTensorQuantization(name, shape, "mxfp8") + } + } return GetTensorQuantization(name, shape, quantize) } @@ -35,3 +55,18 @@ func qwen4ExpIsPLEEmbeddingShard(name string) bool { return strings.Contains(name, ".ple.ple_embedding.ngram_embedding.shard_") && strings.HasSuffix(name, ".weight") } + +func qwen4ExpIsRoutedExpert(name string) bool { + return (strings.HasPrefix(name, "model.language_model.layers.") || + strings.HasPrefix(name, "mtp.layers.")) && + strings.Contains(name, ".mlp.experts.") +} + +func qwen4ExpIsSensitiveQSA(name string) bool { + return strings.HasPrefix(name, "model.language_model.layers.") && + strings.Contains(name, ".self_attn.") && + (strings.HasSuffix(name, ".q_proj.weight") || + strings.HasSuffix(name, ".k_proj.weight") || + strings.HasSuffix(name, ".v_proj.weight") || + strings.HasSuffix(name, ".indexer.index_qk_proj.weight")) +} diff --git a/create/qwen4_exp_test.go b/create/qwen4_exp_test.go index c5159eff6..200710676 100644 --- a/create/qwen4_exp_test.go +++ b/create/qwen4_exp_test.go @@ -13,18 +13,25 @@ func TestQwen4ExpQuantizationType(t *testing.T) { want string }{ { - name: "PLE shard", + name: "PLE shard uses requested NVFP4", tensor: "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight", shape: []int32{156250, 2560}, quantize: "nvfp4", want: "nvfp4", }, + { + name: "PLE shard preserves explicit MXFP8 request", + tensor: "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight", + shape: []int32{156250, 2560}, + quantize: "mxfp8", + want: "mxfp8", + }, { name: "token embedding", tensor: "model.language_model.embed_tokens.weight", shape: []int32{248320, 2560}, quantize: "nvfp4", - want: "", + want: "mxfp8", }, { name: "low rank projection", @@ -38,8 +45,50 @@ func TestQwen4ExpQuantizationType(t *testing.T) { tensor: "model.language_model.layers.0.linear_attn.out_proj.weight", shape: []int32{5120, 3072}, quantize: "nvfp4", + want: "mxfp8", + }, + { + name: "QSA query projection", + tensor: "model.language_model.layers.3.self_attn.q_proj.weight", + shape: []int32{12288, 2560}, + quantize: "nvfp4", + want: "", + }, + { + name: "MTP routed expert", + tensor: "mtp.layers.0.mlp.experts.0.down_proj.weight", + shape: []int32{2560, 1024}, + quantize: "nvfp4", want: "nvfp4", }, + { + name: "MTP control path", + tensor: "mtp.layers.0.self_attn.q_proj.weight", + shape: []int32{12288, 2560}, + quantize: "nvfp4", + want: "", + }, + { + name: "main routed expert", + tensor: "model.language_model.layers.0.mlp.experts.0.down_proj.weight", + shape: []int32{2560, 1024}, + quantize: "nvfp4", + want: "nvfp4", + }, + { + name: "QSA query projection mxfp8", + tensor: "model.language_model.layers.3.self_attn.q_proj.weight", + shape: []int32{12288, 2560}, + quantize: "mxfp8", + want: "mxfp8", + }, + { + name: "hyperconnection projection", + tensor: "model.language_model.layers.0.attn_hyper_connection.block_inject_weight.weight", + shape: []int32{2560, 2560}, + quantize: "nvfp4", + want: "mxfp8", + }, } for _, tt := range tests { diff --git a/mlx/indexed_attention.go b/mlx/indexed_attention.go new file mode 100644 index 000000000..477e70076 --- /dev/null +++ b/mlx/indexed_attention.go @@ -0,0 +1,272 @@ +package mlx + +import ( + "fmt" + "math" +) + +const indexedBlockAttentionMetalSource = ` +uint row = threadgroup_position_in_grid.y; +uint simd_group = simdgroup_index_in_threadgroup; +uint lane = thread_index_in_simdgroup; + +constexpr int SIMDGroups = 32; +constexpr int SIMDWidth = 32; +constexpr int ValuesPerLane = D / SIMDWidth; + +thread float query[ValuesPerLane]; +thread float values[ValuesPerLane]; +threadgroup float outputs[SIMDGroups * SIMDWidth]; +threadgroup float max_scores[SIMDGroups]; +threadgroup float sum_scores[SIMDGroups]; + +int query_index = int(row % QueryLength); +int batch_head = int(row / QueryLength); +int batch = batch_head / QueryHeads; +int query_head = batch_head - batch * QueryHeads; +int kv_head = query_head / (QueryHeads / KVHeads); +int query_end = query_ends[batch * QueryLength + query_index]; +int key_count = int(key_length); + +auto query_ptr = q + ((batch_head * QueryLength + query_index) * D) + lane * ValuesPerLane; +auto output_ptr = out + ((batch_head * QueryLength + query_index) * D) + simd_group * ValuesPerLane; + +for (int i = 0; i < ValuesPerLane; ++i) { + query[i] = float(scale) * static_cast(query_ptr[i]); + values[i] = 0.0f; +} + +int blocks_offset = (batch * QueryLength + query_index) * TopKBlocks; +float max_score = -3.402823466e+38F; +float sum_score = 0.0f; + +// Each SIMD group owns an interleaved subset of the selected tokens and +// accumulates an independent online-softmax state. +for (int selected = int(simd_group); selected < TopKBlocks * BlockSize; selected += SIMDGroups) { + int block_slot = selected / BlockSize; + int block_offset = selected - block_slot * BlockSize; + int block = block_indices[blocks_offset + block_slot]; + int key_index = block * BlockSize + block_offset; + bool valid = block >= 0 && key_index < key_count && key_index < query_end; + + float score = -INFINITY; + if (valid) { + auto key_ptr = k + (((batch * KVHeads + kv_head) * key_count + key_index) * D) + lane * ValuesPerLane; + score = 0.0f; + for (int i = 0; i < ValuesPerLane; ++i) { + score += query[i] * static_cast(key_ptr[i]); + } + score = simd_sum(score); + } + + float next_max = metal::max(max_score, score); + float previous_factor = fast::exp(max_score - next_max); + float weight = valid ? fast::exp(score - next_max) : 0.0f; + max_score = next_max; + sum_score = sum_score * previous_factor + weight; + + if (valid) { + auto value_ptr = v + (((batch * KVHeads + kv_head) * key_count + key_index) * D) + lane * ValuesPerLane; + for (int i = 0; i < ValuesPerLane; ++i) { + values[i] = values[i] * previous_factor + weight * static_cast(value_ptr[i]); + } + } else { + for (int i = 0; i < ValuesPerLane; ++i) { + values[i] *= previous_factor; + } + } +} + +// QSA always attends the incomplete compression block after the selected +// complete blocks. At most BlockSize-1 SIMD groups have a valid tail token. +int tail_start = (query_end / BlockSize) * BlockSize; +int tail_index = tail_start + int(simd_group); +bool valid_tail = tail_index < query_end && tail_index < key_count; +if (valid_tail) { + auto key_ptr = k + (((batch * KVHeads + kv_head) * key_count + tail_index) * D) + lane * ValuesPerLane; + float score = 0.0f; + for (int i = 0; i < ValuesPerLane; ++i) { + score += query[i] * static_cast(key_ptr[i]); + } + score = simd_sum(score); + + float next_max = metal::max(max_score, score); + float previous_factor = fast::exp(max_score - next_max); + float weight = fast::exp(score - next_max); + max_score = next_max; + sum_score = sum_score * previous_factor + weight; + + auto value_ptr = v + (((batch * KVHeads + kv_head) * key_count + tail_index) * D) + lane * ValuesPerLane; + for (int i = 0; i < ValuesPerLane; ++i) { + values[i] = values[i] * previous_factor + weight * static_cast(value_ptr[i]); + } +} + +if (lane == 0) { + max_scores[simd_group] = max_score; + sum_scores[simd_group] = sum_score; +} +threadgroup_barrier(mem_flags::mem_threadgroup); + +max_score = max_scores[lane]; +float row_max = simd_max(max_score); +float row_factor = fast::exp(max_score - row_max); +float row_sum = simd_sum(sum_scores[lane] * row_factor); + +// Transpose the lane-by-SIMD-group accumulators through threadgroup memory. +// Each SIMD group then reduces and writes one contiguous output segment. +for (int i = 0; i < ValuesPerLane; ++i) { + outputs[lane * SIMDWidth + simd_group] = values[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + values[i] = simd_sum(outputs[simd_group * SIMDWidth + lane] * row_factor); + values[i] /= row_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +if (lane == 0) { + for (int i = 0; i < ValuesPerLane; ++i) { + output_ptr[i] = static_cast(values[i]); + } +} +` + +var indexedBlockAttention = &gpuKernel{ + name: "indexed_block_scaled_dot_product_attention", + inputs: []string{"q", "k", "v", "block_indices", "query_ends", "scale", "key_length", "block_size"}, + outputs: []string{"out"}, + metal: gpuSource{ + source: indexedBlockAttentionMetalSource, + header: "#include \nusing namespace metal;\n", + }, + fallback: func(launch gpuLaunch) []*Array { + in := launch.inputs + return []*Array{indexedBlockScaledDotProductAttentionGraph( + in[0], in[1], in[2], in[3], in[4], int(in[7].Int()), in[5].Float(), + )} + }, +} + +// The graph path is faster for small query batches. At 1024 tokens the +// block-indexed kernel pulls ahead and avoids the graph's context-dependent +// intermediate growth. +const indexedBlockAttentionMinQueryLength = 1024 + +func indexedTokenScaledDotProductAttentionGraph(q, k, v, indices, valid *Array, scale float32) *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, K := int32(k.Dim(1)), int32(k.Dim(2)) + repeats := queryHeads / kvHeads + + offsets := make([]int32, B) + for i := range offsets { + offsets[i] = int32(i) * K + } + logical := Add(indices, FromValues(offsets, int(B), 1, 1)) + gather := func(history *Array) *Array { + flattened := Reshape(Transpose(history, 1, 0, 2, 3), kvHeads, B*K, D) + return Transpose(Take(flattened, logical, 1), 1, 0, 2, 3, 4) + } + selectedK, selectedV := gather(k), gather(v) + + qr := Reshape(q, B, kvHeads, repeats, L, 1, D) + kr := Transpose(ExpandDims(selectedK, 2), 0, 1, 2, 3, 5, 4) + scores := Squeeze(Matmul(qr.AsType(DTypeFloat32), kr.AsType(DTypeFloat32)), 4) + scores = MulScalar(scores, scale) + mask := ExpandDims(ExpandDims(valid, 1), 1) + fill := AddScalar(Zeros(scores.DType(), scores.Dims()...), -float32(math.MaxFloat32)) + scores = Where(mask, scores, fill) + probs := SoftmaxAxis(scores, -1, true) + + vr := ExpandDims(selectedV.AsType(DTypeFloat32), 2) + out := Matmul(ExpandDims(probs, 4), vr) + out = Squeeze(out, 4) + return Reshape(out, B, queryHeads, L, D).AsType(outputType) +} + +func indexedBlockScaledDotProductAttentionGraph(q, k, v, blockIndices, queryEnds *Array, blockSize int, scale float32) *Array { + B, L, topK := int32(blockIndices.Dim(0)), int32(blockIndices.Dim(1)), int32(blockIndices.Dim(2)) + keyLength := k.Dim(2) + + offsets := Reshape(Arange(0, float64(blockSize), 1, DTypeInt32), 1, 1, 1, int32(blockSize)) + indices := Add(MulScalar(ExpandDims(blockIndices, -1), float32(blockSize)), offsets) + validBlocks := ExpandDims(blockIndices.Greater(FromValue(-1)), -1) + queryLimit := ExpandDims(ExpandDims(queryEnds, -1), -1) + valid := Mul(validBlocks.AsType(DTypeInt32), indices.Less(queryLimit).AsType(DTypeInt32)) + valid = Mul(valid, indices.Less(FromValue(keyLength)).AsType(DTypeInt32)) + indices = Reshape(indices, B, L, topK*int32(blockSize)) + valid = Reshape(valid, B, L, topK*int32(blockSize)) + + tailWidth := blockSize - 1 + if tailWidth > 0 { + tailOffsets := Reshape(Arange(0, float64(tailWidth), 1, DTypeInt32), 1, 1, int32(tailWidth)) + tailStart := MulScalar(ExpandDims(FloorDivideScalar(queryEnds, int32(blockSize)), -1), float32(blockSize)) + tail := Add(tailStart, tailOffsets) + tailValid := Mul(tail.Less(ExpandDims(queryEnds, -1)).AsType(DTypeInt32), tail.Less(FromValue(keyLength)).AsType(DTypeInt32)) + indices = Concatenate([]*Array{indices, tail}, -1) + valid = Concatenate([]*Array{valid, tailValid}, -1) + } + + validMask := valid.Greater(FromValue(0)) + indices = Where(validMask, indices, Zeros(DTypeInt32, indices.Dims()...)).AsType(DTypeInt32) + return indexedTokenScaledDotProductAttentionGraph(q, k, v, indices, validMask, scale) +} + +// IndexedBlockScaledDotProductAttention attends to a per-query list of +// complete K/V blocks followed by that query's incomplete causal block. +// q is [B, QH, L, D], k and v are [B, KVH, K, D], blockIndices is +// [B, L, TopK], and queryEnds is [B, L]. Negative block indices are ignored. +// A finite initial online-softmax maximum keeps SIMD groups with no valid +// selected token from evaluating -Inf - -Inf. +func IndexedBlockScaledDotProductAttention(q, k, v, blockIndices, queryEnds *Array, blockSize int, scale float32) *Array { + if q == nil || k == nil || v == nil || blockIndices == nil || queryEnds == nil || + q.NumDims() != 4 || k.NumDims() != 4 || v.NumDims() != 4 || + blockIndices.NumDims() != 3 || queryEnds.NumDims() != 2 || blockSize <= 0 { + panic("mlx.IndexedBlockScaledDotProductAttention: invalid inputs") + } + B, queryHeads, L, D := q.Dim(0), q.Dim(1), q.Dim(2), q.Dim(3) + kvHeads, keyLength := k.Dim(1), k.Dim(2) + if k.Dim(0) != B || v.Dim(0) != B || v.Dim(1) != kvHeads || v.Dim(2) != keyLength || + k.Dim(3) != D || v.Dim(3) != D || queryHeads%kvHeads != 0 || + blockIndices.Dim(0) != B || blockIndices.Dim(1) != L || + queryEnds.Dim(0) != B || queryEnds.Dim(1) != L { + panic(fmt.Sprintf( + "mlx.IndexedBlockScaledDotProductAttention: incompatible shapes q=%v k=%v v=%v blocks=%v queryEnds=%v", + q.Dims(), k.Dims(), v.Dims(), blockIndices.Dims(), queryEnds.Dims(), + )) + } + if !MetalIsAvailable() || L < indexedBlockAttentionMinQueryLength || blockSize > 32 || D%32 != 0 || + q.DType() != k.DType() || q.DType() != v.DType() || + (q.DType() != DTypeBFloat16 && q.DType() != DTypeFloat16) || + blockIndices.DType() != DTypeInt32 || queryEnds.DType() != DTypeInt32 { + return indexedBlockScaledDotProductAttentionGraph(q, k, v, blockIndices, queryEnds, blockSize, scale) + } + return indexedBlockScaledDotProductAttentionMetal(q, k, v, blockIndices, queryEnds, blockSize, scale) +} + +func indexedBlockScaledDotProductAttentionMetal(q, k, v, blockIndices, queryEnds *Array, blockSize int, scale float32) *Array { + B, queryHeads, L, D := q.Dim(0), q.Dim(1), q.Dim(2), q.Dim(3) + kvHeads, keyLength, topK := k.Dim(1), k.Dim(2), blockIndices.Dim(2) + outs := indexedBlockAttention.run(gpuLaunch{ + dtypes: []gpuDTypeArg{{"InT", q.DType()}}, + ints: []gpuIntArg{ + {"D", D}, + {"QueryLength", L}, + {"QueryHeads", queryHeads}, + {"KVHeads", kvHeads}, + {"BlockSize", blockSize}, + {"TopKBlocks", topK}, + }, + outputs: []gpuOutputSpec{{ + "INDEXED_BLOCK_SCALED_DOT_PRODUCT_ATTENTION", + []int32{int32(B), int32(queryHeads), int32(L), int32(D)}, + q.DType(), + }}, + grid: [3]int{1024, B * queryHeads * L, 1}, + threadGroup: [3]int{1024, 1, 1}, + inputs: []*Array{ + q, k, v, blockIndices, queryEnds, + FromValue(scale), FromValue(keyLength), FromValue(blockSize), + }, + }) + return outs[0] +} diff --git a/mlx/indexed_attention_test.go b/mlx/indexed_attention_test.go new file mode 100644 index 000000000..6abf972d4 --- /dev/null +++ b/mlx/indexed_attention_test.go @@ -0,0 +1,59 @@ +package mlx + +import ( + "math" + "testing" + + "github.com/ollama/ollama/mlx/mlxthread/mlxthreadtest" +) + +func TestIndexedBlockScaledDotProductAttentionMatchesGraph(t *testing.T) { + withMLXThread(t, func(t *mlxthreadtest.T) { + Scoped(func() { + const ( + batch = 2 + queryHeads = 4 + kvHeads = 2 + queryLength = indexedBlockAttentionMinQueryLength + keyLength = 64 + headDim = 32 + blockSize = 4 + topK = 16 + ) + + q := patternArray(DTypeBFloat16, []int{batch, queryHeads, queryLength, headDim}, 0.1, 0.03, 7, 19) + k := patternArray(DTypeBFloat16, []int{batch, kvHeads, keyLength, headDim}, -0.05, 0.02, 5, 17) + v := patternArray(DTypeBFloat16, []int{batch, kvHeads, keyLength, headDim}, 0.2, 0.04, 3, 13) + blockValues := make([]int32, batch*queryLength*topK) + queryEndValues := make([]int32, batch*queryLength) + for b := range batch { + for row := range queryLength { + i := (b*queryLength + row) * topK + for j := range topK { + blockValues[i+j] = -1 + if j%2 == 0 { + blockValues[i+j] = int32(j / 2) + } + } + queryEndValues[b*queryLength+row] = keyLength + } + } + blocks := FromValues(blockValues, batch, queryLength, topK) + queryEnds := FromValues(queryEndValues, batch, queryLength) + const scale = float32(0.176776695) + got := IndexedBlockScaledDotProductAttention(q, k, v, blocks, queryEnds, blockSize, scale).AsType(DTypeFloat32) + want := indexedBlockScaledDotProductAttentionGraph(q, k, v, blocks, queryEnds, blockSize, scale).AsType(DTypeFloat32) + Eval(got, want) + + gotValues, wantValues := got.Floats(), want.Floats() + for i := range wantValues { + if math.IsNaN(float64(gotValues[i])) || math.IsInf(float64(gotValues[i]), 0) { + t.Fatalf("out[%d] = %v, want finite output", i, gotValues[i]) + } + if diff := float32(math.Abs(float64(gotValues[i] - wantValues[i]))); diff > 0.005 { + t.Fatalf("out[%d] = %v, want %v (diff %v)", i, gotValues[i], wantValues[i], diff) + } + } + }) + }) +} diff --git a/mlxrunner/model/qwen4_exp/blocks.go b/mlxrunner/model/qwen4_exp/blocks.go index 760c36751..1ad2511cc 100644 --- a/mlxrunner/model/qwen4_exp/blocks.go +++ b/mlxrunner/model/qwen4_exp/blocks.go @@ -126,8 +126,8 @@ func (a *fullAttention) Forward(x *mlx.Array, b *batch.Batch, c, side cache.Cach 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) + blocks, queryEnds := qsaLogicalBlocks(scores, b, keyLength, cfg) + out = qsaSparseAttention(q, mainHistory, blocks, queryEnds, cfg) } out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) out = mlx.Mul(out, mlx.Sigmoid(gate)) diff --git a/mlxrunner/model/qwen4_exp/engram.go b/mlxrunner/model/qwen4_exp/engram.go index ddb787304..d441e61db 100644 --- a/mlxrunner/model/qwen4_exp/engram.go +++ b/mlxrunner/model/qwen4_exp/engram.go @@ -62,6 +62,10 @@ func (p *PLE) hashes(b *batch.Batch, history *mlx.Array, cfg *Config) (*mlx.Arra } func (p *PLE) lookup(globalIDs *mlx.Array) *mlx.Array { + if p.hostEmbedding != nil { + return p.hostEmbedding.lookup(globalIDs) + } + rows := mlx.FromValues([]int64{int64(p.ShardRows)}, 1) shardIDs := globalIDs.FloorDivide(rows) localIDs := globalIDs.Remainder(rows) diff --git a/mlxrunner/model/qwen4_exp/ple_host.go b/mlxrunner/model/qwen4_exp/ple_host.go new file mode 100644 index 000000000..e9d1793a8 --- /dev/null +++ b/mlxrunner/model/qwen4_exp/ple_host.go @@ -0,0 +1,251 @@ +package qwen4_exp + +import ( + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "log/slog" + "sort" + "strconv" + "strings" + + "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/mlx" +) + +var errHostPLEUnsupported = errors.New("host-backed PLE requires standard NVFP4") + +type safetensorEntry struct { + DType string `json:"dtype"` + Shape []int64 `json:"shape"` + DataOffsets [2]int64 `json:"data_offsets"` +} + +type hostPLEShard struct { + mapping *hostMapping + + weightStart int64 + scaleStart int64 + rows int + weightColumns int + scaleColumns int +} + +type hostPLETable struct { + shards []*hostPLEShard + rows int + rowWidth int +} + +func openHostPLETables(m *manifest.Manifest, cfg *Config) (map[string]*hostPLETable, error) { + headCount := (cfg.NGramSize - 1) * cfg.HeadsPerNGram + if headCount <= 0 || cfg.PLEEmbedDim%headCount != 0 { + return nil, nil + } + rowWidth := int(cfg.PLEEmbedDim / headCount) + + layers := make(map[string]manifest.Layer) + for _, layer := range m.TensorLayers() { + layers[layer.Name] = layer + } + + tables := make(map[string]*hostPLETable, len(cfg.PLELayerIDs)) + rowsPerShard := 0 + for _, id := range cfg.PLELayerIDs { + prefix := fmt.Sprintf("model.language_model.layers.%d.ple", id-1) + table, err := openHostPLETable(layers, prefix, int(cfg.SplitNGramParts), rowWidth) + if errors.Is(err, errHostPLEUnsupported) { + closeHostPLETables(tables) + return nil, nil + } + if err != nil { + closeHostPLETables(tables) + return nil, fmt.Errorf("%s: %w", prefix, err) + } + tables[prefix] = table + rowsPerShard = table.rows + } + slog.Info("enabled host-backed Qwen PLE", "layers", len(tables), "shards_per_layer", cfg.SplitNGramParts, "rows_per_shard", rowsPerShard) + return tables, nil +} + +func openHostPLETable(layers map[string]manifest.Layer, prefix string, shardCount, rowWidth int) (*hostPLETable, error) { + table := &hostPLETable{ + shards: make([]*hostPLEShard, shardCount), + rowWidth: rowWidth, + } + for i := range shardCount { + name := fmt.Sprintf("%s.ple_embedding.ngram_embedding.shard_%d.weight", prefix, i) + layer, ok := layers[name] + if !ok { + table.close() + return nil, fmt.Errorf("missing tensor layer %q", name) + } + path, err := manifest.BlobsPath(layer.Digest) + if err != nil { + table.close() + return nil, fmt.Errorf("resolve %s: %w", name, err) + } + shard, err := openHostPLEShard(path, name, rowWidth) + if err != nil { + table.close() + return nil, fmt.Errorf("open %s: %w", name, err) + } + if i == 0 { + table.rows = shard.rows + } else if shard.rows != table.rows { + shard.mapping.close() + table.close() + return nil, fmt.Errorf("shard %d has %d rows, want %d", i, shard.rows, table.rows) + } + table.shards[i] = shard + } + return table, nil +} + +func openHostPLEShard(path, name string, rowWidth int) (*hostPLEShard, error) { + mapping, err := openHostMapping(path) + if err != nil { + return nil, err + } + fail := func(err error) (*hostPLEShard, error) { + mapping.close() + return nil, err + } + + prefix, err := mapping.bytes(0, 8) + if err != nil { + return fail(err) + } + headerSize := int64(binary.LittleEndian.Uint64(prefix)) + headerData, err := mapping.bytes(8, 8+headerSize) + if err != nil { + return fail(err) + } + var header map[string]json.RawMessage + if err := json.Unmarshal(headerData, &header); err != nil { + return fail(fmt.Errorf("parse safetensors header: %w", err)) + } + var metadata map[string]string + if err := json.Unmarshal(header["__metadata__"], &metadata); err != nil { + return fail(errHostPLEUnsupported) + } + if !strings.EqualFold(metadata["quant_type"], "nvfp4") || metadata["group_size"] != strconv.Itoa(16) { + return fail(errHostPLEUnsupported) + } + + var weight, scale safetensorEntry + if err := json.Unmarshal(header[name], &weight); err != nil { + return fail(fmt.Errorf("parse weight header: %w", err)) + } + if err := json.Unmarshal(header[name+".scale"], &scale); err != nil { + return fail(errHostPLEUnsupported) + } + if weight.DType != "U32" || len(weight.Shape) != 2 || weight.Shape[0] <= 0 || weight.Shape[1]*8 != int64(rowWidth) { + return fail(fmt.Errorf("weight has dtype=%s shape=%v, want U32 [rows,%d]", weight.DType, weight.Shape, rowWidth/8)) + } + if scale.DType != "U8" || len(scale.Shape) != 2 || scale.Shape[0] != weight.Shape[0] || scale.Shape[1]*16 != int64(rowWidth) { + return fail(fmt.Errorf("scale has dtype=%s shape=%v, want U8 [%d,%d]", scale.DType, scale.Shape, weight.Shape[0], rowWidth/16)) + } + + dataStart := 8 + headerSize + weightStart := dataStart + weight.DataOffsets[0] + weightEnd := dataStart + weight.DataOffsets[1] + scaleStart := dataStart + scale.DataOffsets[0] + scaleEnd := dataStart + scale.DataOffsets[1] + if _, err := mapping.bytes(weightStart, weightEnd); err != nil { + return fail(err) + } + if _, err := mapping.bytes(scaleStart, scaleEnd); err != nil { + return fail(err) + } + if weightEnd-weightStart != weight.Shape[0]*weight.Shape[1]*4 || scaleEnd-scaleStart != scale.Shape[0]*scale.Shape[1] { + return fail(fmt.Errorf("tensor byte sizes do not match their shapes")) + } + + return &hostPLEShard{ + mapping: mapping, + weightStart: weightStart, + scaleStart: scaleStart, + rows: int(weight.Shape[0]), + weightColumns: int(weight.Shape[1]), + scaleColumns: int(scale.Shape[1]), + }, nil +} + +func closeHostPLETables(tables map[string]*hostPLETable) { + for _, table := range tables { + table.close() + } +} + +func (t *hostPLETable) close() { + for _, shard := range t.shards { + if shard != nil { + shard.mapping.close() + } + } +} + +func (t *hostPLETable) lookup(globalIDs *mlx.Array) *mlx.Array { + shape := globalIDs.Dims() + ids := globalIDs.AsType(mlx.DTypeInt32).Ints() + order := make([]int, len(ids)) + maxID := len(t.shards) * t.rows + for i, id := range ids { + if id < 0 || int(id) >= maxID { + panic(fmt.Sprintf("Qwen PLE row ID %d outside [0,%d)", id, maxID)) + } + order[i] = i + } + sort.Slice(order, func(i, j int) bool { + return ids[order[i]] < ids[order[j]] + }) + + weightColumns := t.shards[0].weightColumns + scaleColumns := t.shards[0].scaleColumns + weights := make([]uint32, len(ids)*weightColumns) + scales := make([]uint8, len(ids)*scaleColumns) + previousID := int32(-1) + previousOutput := -1 + for _, output := range order { + id := ids[output] + weightOutput := weights[output*weightColumns : (output+1)*weightColumns] + scaleOutput := scales[output*scaleColumns : (output+1)*scaleColumns] + if id == previousID { + copy(weightOutput, weights[previousOutput*weightColumns:(previousOutput+1)*weightColumns]) + copy(scaleOutput, scales[previousOutput*scaleColumns:(previousOutput+1)*scaleColumns]) + continue + } + + shard := t.shards[int(id)/t.rows] + row := int64(int(id) % t.rows) + weightBytes, err := shard.mapping.bytes( + shard.weightStart+row*int64(weightColumns*4), + shard.weightStart+(row+1)*int64(weightColumns*4), + ) + if err != nil { + panic(fmt.Sprintf("read Qwen PLE weight row %d: %v", id, err)) + } + for i := range weightColumns { + weightOutput[i] = binary.LittleEndian.Uint32(weightBytes[i*4:]) + } + scaleBytes, err := shard.mapping.bytes( + shard.scaleStart+row*int64(scaleColumns), + shard.scaleStart+(row+1)*int64(scaleColumns), + ) + if err != nil { + panic(fmt.Sprintf("read Qwen PLE scale row %d: %v", id, err)) + } + copy(scaleOutput, scaleBytes) + previousID = id + previousOutput = output + } + + weightShape := append(append([]int(nil), shape...), weightColumns) + scaleShape := append(append([]int(nil), shape...), scaleColumns) + weight := mlx.FromValues(weights, weightShape...) + scale := mlx.FromValues(scales, scaleShape...) + return mlx.Dequantize(weight, scale, nil, 16, 4, "nvfp4", nil) +} diff --git a/mlxrunner/model/qwen4_exp/ple_host_test.go b/mlxrunner/model/qwen4_exp/ple_host_test.go new file mode 100644 index 000000000..2c3380471 --- /dev/null +++ b/mlxrunner/model/qwen4_exp/ple_host_test.go @@ -0,0 +1,90 @@ +package qwen4_exp + +import ( + "bytes" + "encoding/binary" + "fmt" + "io" + "slices" + "testing" + + "github.com/ollama/ollama/fs/safetensors" + "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlx/mlxtest" + "github.com/ollama/ollama/mlxrunner/nn" +) + +func TestHostPLELookupMatchesResidentEmbedding(t *testing.T) { + const ( + prefix = "model.language_model.layers.1.ple" + shardCount = 2 + rowsPerShard = 3 + rowWidth = 16 + ) + weights := []uint32{ + 0x01234567, 0x89abcdef, + 0x11111111, 0x22222222, + 0x33333333, 0x44444444, + 0x55555555, 0x66666666, + 0x77777777, 0x88888888, + 0x99999999, 0xaaaaaaaa, + } + scales := []uint8{0x38, 0x40, 0x48, 0x38, 0x40, 0x48} + + t.Setenv("OLLAMA_MODELS", t.TempDir()) + layers := make([]manifest.Layer, 0, shardCount) + for shard := range shardCount { + name := fmt.Sprintf("%s.ple_embedding.ngram_embedding.shard_%d.weight", prefix, shard) + start := shard * rowsPerShard + end := start + rowsPerShard + weightData := safetensors.NewTensorDataFromBytes(name, "U32", []int32{rowsPerShard, rowWidth / 8}, uint32Bytes(weights[start*2:end*2])) + scaleData := safetensors.NewTensorDataFromBytes(name+".scale", "U8", []int32{rowsPerShard, rowWidth / 16}, scales[start:end]) + blob, err := io.ReadAll(safetensors.BuildPackedSafetensorsReaderWithMetadata([]*safetensors.TensorData{weightData, scaleData}, map[string]string{ + "group_size": "16", + "quant_type": "nvfp4", + })) + if err != nil { + t.Fatal(err) + } + layer, err := manifest.NewLayer(bytes.NewReader(blob), manifest.MediaTypeImageTensor) + if err != nil { + t.Fatal(err) + } + layer.Name = name + layers = append(layers, layer) + } + layerByName := make(map[string]manifest.Layer, len(layers)) + for _, layer := range layers { + layerByName[layer.Name] = layer + } + table, err := openHostPLETable(layerByName, prefix, shardCount, rowWidth) + if err != nil { + t.Fatal(err) + } + t.Cleanup(table.close) + + mlxtest.Run(t, func(t *mlxtest.T) { + ids := mlx.FromValues([]int64{5, 0, 3, 3, 2, 4}, 2, 3) + got := table.lookup(ids).AsType(mlx.DTypeFloat32) + want := (&nn.QuantizedEmbedding{ + Weight: mlx.FromValues(weights, shardCount*rowsPerShard, rowWidth/8), + Scales: mlx.FromValues(scales, shardCount*rowsPerShard, rowWidth/16), + GroupSize: 16, + Bits: 4, + Mode: "nvfp4", + }).Forward(ids).AsType(mlx.DTypeFloat32) + mlx.Eval(got, want) + if !slices.Equal(got.Floats(), want.Floats()) { + t.Fatalf("host lookup = %v, want %v", got.Floats(), want.Floats()) + } + }) +} + +func uint32Bytes(values []uint32) []byte { + data := make([]byte, len(values)*4) + for i, value := range values { + binary.LittleEndian.PutUint32(data[i*4:], value) + } + return data +} diff --git a/mlxrunner/model/qwen4_exp/ple_mapping_other.go b/mlxrunner/model/qwen4_exp/ple_mapping_other.go new file mode 100644 index 000000000..c2f90ceb3 --- /dev/null +++ b/mlxrunner/model/qwen4_exp/ple_mapping_other.go @@ -0,0 +1,42 @@ +//go:build !darwin && !linux + +package qwen4_exp + +import ( + "fmt" + "io" + "os" +) + +type hostMapping struct { + file *os.File + size int64 +} + +func openHostMapping(path string) (*hostMapping, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + info, err := f.Stat() + if err != nil { + f.Close() + return nil, err + } + return &hostMapping{file: f, size: info.Size()}, nil +} + +func (m *hostMapping) bytes(start, end int64) ([]byte, error) { + if start < 0 || end < start || end > m.size { + return nil, fmt.Errorf("mapping range [%d:%d] outside %d bytes", start, end, m.size) + } + data := make([]byte, end-start) + if _, err := m.file.ReadAt(data, start); err != nil && err != io.EOF { + return nil, err + } + return data, nil +} + +func (m *hostMapping) close() error { + return m.file.Close() +} diff --git a/mlxrunner/model/qwen4_exp/ple_mapping_unix.go b/mlxrunner/model/qwen4_exp/ple_mapping_unix.go new file mode 100644 index 000000000..d937e7d20 --- /dev/null +++ b/mlxrunner/model/qwen4_exp/ple_mapping_unix.go @@ -0,0 +1,55 @@ +//go:build darwin || linux + +package qwen4_exp + +import ( + "fmt" + "os" + + "golang.org/x/sys/unix" +) + +type hostMapping struct { + data []byte +} + +func openHostMapping(path string) (*hostMapping, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + info, err := f.Stat() + if err != nil { + return nil, err + } + if info.Size() <= 0 || info.Size() > int64(^uint(0)>>1) { + return nil, fmt.Errorf("invalid mapping size %d", info.Size()) + } + + data, err := unix.Mmap(int(f.Fd()), 0, int(info.Size()), unix.PROT_READ, unix.MAP_SHARED) + if err != nil { + return nil, err + } + // PLE row IDs are hashes spread across the table. Disabling sequential + // readahead keeps untouched pages reclaimable under memory pressure. + _ = unix.Madvise(data, unix.MADV_RANDOM) + return &hostMapping{data: data}, nil +} + +func (m *hostMapping) bytes(start, end int64) ([]byte, error) { + if start < 0 || end < start || end > int64(len(m.data)) { + return nil, fmt.Errorf("mapping range [%d:%d] outside %d bytes", start, end, len(m.data)) + } + return m.data[start:end], nil +} + +func (m *hostMapping) close() error { + if len(m.data) == 0 { + return nil + } + err := unix.Munmap(m.data) + m.data = nil + return err +} diff --git a/mlxrunner/model/qwen4_exp/qsa.go b/mlxrunner/model/qwen4_exp/qsa.go index 5f8e3cf13..cc3eb0bb0 100644 --- a/mlxrunner/model/qwen4_exp/qsa.go +++ b/mlxrunner/model/qwen4_exp/qsa.go @@ -110,9 +110,11 @@ func qsaCompressedKeys(rawKeys, positions *mlx.Array, indexer *attentionIndexer, 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) { +// qsaLogicalBlocks reproduces the reference's compressed-block top-k. +// Blocks outside a query's causal history are represented by -1; queryEnds +// lets attention append the incomplete causal block without expanding the +// selected blocks into per-token indices. +func qsaLogicalBlocks(scores *mlx.Array, b *batch.Batch, keyLength int32, cfg *Config) (selected, queryEnds *mlx.Array) { B, L := int32(scores.Dim(0)), int32(scores.Dim(1)) ratio := cfg.IndexerCompressRatio blocks := keyLength / ratio @@ -132,73 +134,22 @@ func qsaLogicalIndices(scores *mlx.Array, b *batch.Batch, keyLength int32, cfg * 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, + partitioned := mlx.Argpartition(mlx.Neg(masked), int(selectedCount)-1, -1) + selected = mlx.SliceStartStop(partitioned, []int32{0, 0, 0}, []int32{B, L, selectedCount}, ) } else { - indices = mlx.Tile(blockIDs, []int32{B, L, 1}) + selected = mlx.Tile(blockIDs, []int32{B, L, 1}) } - blockValid = indices.Less(visibleBlocks) + blockValid = selected.Less(visibleBlocks) + invalid := mlx.AddScalar(mlx.Zeros(mlx.DTypeInt32, int(B), int(L), int(selectedCount)), -1) + selected = mlx.Where(blockValid, selected, invalid).AsType(mlx.DTypeInt32) + return selected, mlx.Reshape(visibleTokens, B, L).AsType(mlx.DTypeInt32) +} - 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), +func qsaSparseAttention(q *mlx.Array, history *nn.KVHistory, blocks, queryEnds *mlx.Array, cfg *Config) *mlx.Array { + return mlx.IndexedBlockScaledDotProductAttention( + q, history.K(), history.V(), blocks, queryEnds, int(cfg.IndexerCompressRatio), cfg.Scale, ) - expanded = mlx.Reshape(expanded, B, L, selectedCount*ratio) - expandedValid = mlx.Reshape(expandedValid, B, L, selectedCount*ratio) - - // The current partial compression group is always attended causally. - tailWidth := ratio - 1 - tailOffsets := mlx.Reshape(mlx.Arange(0, float64(tailWidth), 1, mlx.DTypeInt32), 1, 1, tailWidth) - tailStart := mlx.MulScalar(mlx.FloorDivideScalar(visibleTokens, ratio), float32(ratio)) - tail := mlx.Add(tailStart, tailOffsets) - tailCount := mlx.Sub(visibleTokens, tailStart) - tailValid := tailOffsets.Less(tailCount).AsType(mlx.DTypeInt32) - - indices = mlx.Concatenate([]*mlx.Array{expanded, tail}, -1) - validInt := mlx.Concatenate([]*mlx.Array{expandedValid, tailValid}, -1) - valid = validInt.Greater(mlx.Zeros(mlx.DTypeInt32, int(B), int(L), validInt.Dim(2))) - indices = mlx.Where(valid, indices, mlx.Zeros(mlx.DTypeInt32, int(B), int(L), indices.Dim(2))) - return indices, valid -} - -func qsaSparseAttention(q *mlx.Array, history *nn.KVHistory, indices, valid *mlx.Array, cfg *Config) *mlx.Array { - outputType := q.DType() - B, queryHeads, L, D := int32(q.Dim(0)), int32(q.Dim(1)), int32(q.Dim(2)), int32(q.Dim(3)) - kvHeads := cfg.NumKeyValueHeads - repeats := queryHeads / kvHeads - - k := qsaGatherHistory(history.K(), indices) - v := qsaGatherHistory(history.V(), indices) - - qr := mlx.Reshape(q, B, kvHeads, repeats, L, 1, D) - kr := mlx.Transpose(mlx.ExpandDims(k, 2), 0, 1, 2, 3, 5, 4) - scores := mlx.Squeeze(mlx.Matmul(qr.AsType(mlx.DTypeFloat32), kr.AsType(mlx.DTypeFloat32)), 4) - scores = mlx.MulScalar(scores, cfg.Scale) - mask := mlx.ExpandDims(mlx.ExpandDims(valid, 1), 1) - fill := mlx.AddScalar(mlx.Zeros(scores.DType(), scores.Dims()...), -float32(math.MaxFloat32)) - scores = mlx.Where(mask, scores, fill) - probs := mlx.SoftmaxAxis(scores, -1, true) - - vr := mlx.ExpandDims(v.AsType(mlx.DTypeFloat32), 2) - out := mlx.Matmul(mlx.ExpandDims(probs, 4), vr) - out = mlx.Squeeze(out, 4) - return mlx.Reshape(out, B, queryHeads, L, D).AsType(outputType) -} - -// qsaGatherHistory selects each batch row's logical token indices from a -// [B, H, K, D] cache history and returns [B, H, L, S, D]. -func qsaGatherHistory(history, indices *mlx.Array) *mlx.Array { - B, H, K, D := int32(history.Dim(0)), int32(history.Dim(1)), int32(history.Dim(2)), int32(history.Dim(3)) - offsets := make([]int32, B) - for i := range offsets { - offsets[i] = int32(i) * K - } - logical := mlx.Add(indices, mlx.FromValues(offsets, int(B), 1, 1)) - flattened := mlx.Reshape(mlx.Transpose(history, 1, 0, 2, 3), H, B*K, D) - return mlx.Transpose(mlx.Take(flattened, logical, 1), 1, 0, 2, 3, 4) } diff --git a/mlxrunner/model/qwen4_exp/qsa_test.go b/mlxrunner/model/qwen4_exp/qsa_test.go index 10046a52a..6977d3a5b 100644 --- a/mlxrunner/model/qwen4_exp/qsa_test.go +++ b/mlxrunner/model/qwen4_exp/qsa_test.go @@ -17,22 +17,19 @@ func TestQSASelectsCompressedBlocksAndCausalTail(t *testing.T) { 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) - } + blocks, queryEnds := qsaLogicalBlocks(scores, b, 16, cfg) + if blocks.DType() != mlx.DTypeInt32 { + t.Fatalf("blocks dtype = %v, want int32", blocks.DType()) } - 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) + mlx.Eval(blocks, queryEnds) + + got := blocks.Ints() + slices.Sort(got) + if want := []int32{1, 3}; !slices.Equal(got, want) { + t.Fatalf("selected blocks = %v, want %v", got, want) + } + if got, want := queryEnds.Ints(), []int32{17}; !slices.Equal(got, want) { + t.Fatalf("query ends = %v, want %v", got, want) } }) } @@ -45,34 +42,29 @@ func TestQSASelectionMasksFutureBlocks(t *testing.T) { 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) + blocks, queryEnds := qsaLogicalBlocks(scores, b, 20, cfg) + mlx.Eval(blocks, queryEnds) - var selected []int32 - for i, value := range indices.Ints() { - if valid.Ints()[i] != 0 { - selected = append(selected, value) - } + got := blocks.Ints() + if want := []int32{0, -1}; !slices.Equal(got, want) { + t.Fatalf("selected blocks = %v, want %v", got, want) } - slices.Sort(selected) - if want := []int32{0, 1, 2, 3, 4}; !slices.Equal(selected, want) { - t.Fatalf("selected indices = %v, want %v", selected, want) + if got, want := queryEnds.Ints(), []int32{5}; !slices.Equal(got, want) { + t.Fatalf("query ends = %v, want %v", got, want) } }) } func TestQSASparseAttentionMatchesReference(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { - cfg := &Config{NumKeyValueHeads: 1, Scale: 1} + cfg := &Config{NumKeyValueHeads: 1, IndexerCompressRatio: 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) + blocks := mlx.FromValues([]int32{2, 0}, 1, 1, 2) + queryEnds := mlx.FromValues([]int32{3}, 1, 1) - out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg) + out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), blocks, queryEnds, cfg) out = out.AsType(mlx.DTypeFloat32) mlx.Eval(out) got := out.Floats() @@ -88,15 +80,15 @@ func TestQSASparseAttentionMatchesReference(t *testing.T) { func TestQSASparseAttentionIgnoresInvalidRows(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { - cfg := &Config{NumKeyValueHeads: 1, Scale: 1} + cfg := &Config{NumKeyValueHeads: 1, IndexerCompressRatio: 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) + blocks := mlx.FromValues([]int32{-1, 1}, 1, 1, 2) + queryEnds := mlx.FromValues([]int32{2}, 1, 1) - out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg) + out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), blocks, queryEnds, cfg) out = out.AsType(mlx.DTypeFloat32) mlx.Eval(out) if got, want := out.Floats(), []float32{7, 3}; !slices.Equal(got, want) { @@ -107,7 +99,7 @@ func TestQSASparseAttentionIgnoresInvalidRows(t *testing.T) { func TestQSASparseAttentionKeepsBatchRowsIndependent(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { - cfg := &Config{NumKeyValueHeads: 1, Scale: 1} + cfg := &Config{NumKeyValueHeads: 1, IndexerCompressRatio: 1, Scale: 1} q := mlx.FromValues([]float32{1, 0, 1, 0}, 2, 1, 1, 2) k := mlx.FromValues([]float32{ 1, 0, 0, 1, @@ -117,10 +109,10 @@ func TestQSASparseAttentionKeepsBatchRowsIndependent(t *testing.T) { 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) + blocks := mlx.FromValues([]int32{0, 1}, 2, 1, 1) + queryEnds := mlx.FromValues([]int32{2, 2}, 2, 1) - out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg) + out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), blocks, queryEnds, cfg) out = out.AsType(mlx.DTypeFloat32) mlx.Eval(out) if got, want := out.Floats(), []float32{10, 1, 40, 4}; !slices.Equal(got, want) { diff --git a/mlxrunner/model/qwen4_exp/qwen4_exp.go b/mlxrunner/model/qwen4_exp/qwen4_exp.go index 5b5935bc6..592f9cb04 100644 --- a/mlxrunner/model/qwen4_exp/qwen4_exp.go +++ b/mlxrunner/model/qwen4_exp/qwen4_exp.go @@ -40,6 +40,7 @@ type Model struct { quantMode string tensorQuant map[string]*model.TensorQuantInfo configData []byte + hostPLE map[string]*hostPLETable } // NewModel validates the publisher config before constructing runtime state. @@ -72,6 +73,10 @@ func NewModel(root *model.Root) (model.Model, error) { if root.GroupSize() > 0 { group = root.GroupSize() } + hostPLE, err := openHostPLETables(root.Manifest, &cfg) + if err != nil { + return nil, fmt.Errorf("open host-backed PLE: %w", err) + } return &Model{ Config: &cfg, Layers: make([]*Layer, cfg.NumHiddenLayers), @@ -81,6 +86,7 @@ func NewModel(root *model.Root) (model.Model, error) { quantMode: mode, tensorQuant: root.AllTensorQuant(), configData: configData, + hostPLE: hostPLE, }, nil } diff --git a/mlxrunner/model/qwen4_exp/weights.go b/mlxrunner/model/qwen4_exp/weights.go index 9efc91b3c..b2321ead2 100644 --- a/mlxrunner/model/qwen4_exp/weights.go +++ b/mlxrunner/model/qwen4_exp/weights.go @@ -120,6 +120,7 @@ type PLE struct { HeadVocabSizes *mlx.Array EmbeddingShards []nn.EmbeddingLayer ShardRows int + hostEmbedding *hostPLETable } type streamRMSNorm struct { @@ -339,6 +340,12 @@ func (m *Model) loadPLE(linears model.LinearFactory, tensors map[string]*mlx.Arr return nil, err } + if table := m.hostPLE[prefix]; table != nil { + p.hostEmbedding = table + p.ShardRows = table.rows + return p, nil + } + p.EmbeddingShards = make([]nn.EmbeddingLayer, m.SplitNGramParts) for i := range m.SplitNGramParts { path := fmt.Sprintf("%s.ple_embedding.ngram_embedding.shard_%d", prefix, i) From 257aa8ca63165b95d2cba6c2ad9a10a14d9e1830 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Tue, 22 Sep 2026 17:42:45 -0700 Subject: [PATCH 2/2] mlx: speed up Qwen 3.8 Flash Next Fuse small graph operations and pack hyper-connection projections to reduce prompt and generation overhead. --- mlxrunner/model/qwen4_exp/blocks.go | 12 +++- mlxrunner/model/qwen4_exp/engram.go | 20 ++++-- mlxrunner/model/qwen4_exp/hyper_connection.go | 64 ++++++++++++++++--- .../model/qwen4_exp/hyper_connection_test.go | 13 ++-- mlxrunner/model/qwen4_exp/weights.go | 39 +++++++++++ 5 files changed, 129 insertions(+), 19 deletions(-) diff --git a/mlxrunner/model/qwen4_exp/blocks.go b/mlxrunner/model/qwen4_exp/blocks.go index 1ad2511cc..c0a97e1f6 100644 --- a/mlxrunner/model/qwen4_exp/blocks.go +++ b/mlxrunner/model/qwen4_exp/blocks.go @@ -9,6 +9,15 @@ import ( "github.com/ollama/ollama/mlxrunner/nn" ) +var linearAttentionOutputGate = mlx.Compile2( + "Qwen4LinearAttentionOutputGate", + func(out, gate *mlx.Array) *mlx.Array { + dtype := out.DType() + return mlx.Mul(out.AsType(mlx.DTypeFloat32), mlx.Sigmoid(gate.AsType(mlx.DTypeFloat32))).AsType(dtype) + }, + mlx.Shapeless(), +) + 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 { @@ -54,8 +63,7 @@ func (a *linearAttention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, c 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 = linearAttentionOutputGate(out.AsType(outType), z) out = a.Out.Forward(mlx.Reshape(out, B, L, valueDim)) if recurrent != nil { recurrent.Put(b, convState, deltaState) diff --git a/mlxrunner/model/qwen4_exp/engram.go b/mlxrunner/model/qwen4_exp/engram.go index d441e61db..e5b7aa149 100644 --- a/mlxrunner/model/qwen4_exp/engram.go +++ b/mlxrunner/model/qwen4_exp/engram.go @@ -7,6 +7,19 @@ import ( "github.com/ollama/ollama/mlxrunner/batch" ) +var streamRMSNormFn = mlx.Compile( + "Qwen4StreamRMSNorm", + func(in ...*mlx.Array) []*mlx.Array { + x, weight, eps := in[0], in[1], in[2] + dtype := x.DType() + values := x.AsType(mlx.DTypeFloat32) + meanSquare := mlx.Mean(mlx.Mul(values, values), -1, true) + invRMS := mlx.Div(mlx.NewScalarArray(1), mlx.Add(meanSquare, eps).Sqrt()) + return []*mlx.Array{mlx.Mul(mlx.Mul(values, invRMS), weight.AsType(mlx.DTypeFloat32)).AsType(dtype)} + }, + mlx.Shapeless(), +) + func (n *streamRMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array { originalShape := x.Dims() if originalShape[len(originalShape)-1] == n.Weight.Size() { @@ -14,11 +27,8 @@ func (n *streamRMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array { 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) + epsArray := mlx.NewScalarArray(eps) + normalized := streamRMSNormFn(x, n.Weight, epsArray)[0] return normalized.Reshape(originalShape...) } diff --git a/mlxrunner/model/qwen4_exp/hyper_connection.go b/mlxrunner/model/qwen4_exp/hyper_connection.go index 333b3af9f..098b15b6c 100644 --- a/mlxrunner/model/qwen4_exp/hyper_connection.go +++ b/mlxrunner/model/qwen4_exp/hyper_connection.go @@ -2,9 +2,39 @@ package qwen4_exp import "github.com/ollama/ollama/mlx" +var ( + hyperConnectionSiLU = mlx.Compile2( + "Qwen4HyperConnectionSiLU", + func(x, count *mlx.Array) *mlx.Array { + return mlx.SiLU(mlx.Div(x, count)) + }, + mlx.Shapeless(), + ) + hyperConnectionMix = mlx.Compile3( + "Qwen4HyperConnectionMix", + func(streams, mix, count *mlx.Array) *mlx.Array { + return mlx.Div(mlx.Sum(mlx.Mul(streams, mlx.Sigmoid(mix)), 2, false), count) + }, + mlx.Shapeless(), + ) + hyperConnectionInject = mlx.Compile( + "Qwen4HyperConnectionInject", + func(in ...*mlx.Array) []*mlx.Array { + residual, branch, logits, count := in[0], in[1], in[2], in[3] + weight := mlx.Sigmoid(mlx.Div(logits, count)) + two := mlx.FromValue(float32(2)).AsType(weight.DType()) + weight = mlx.Mul(weight, two) + update := mlx.Mul(mlx.ExpandDims(weight, -1), mlx.ExpandDims(branch, -2)) + return []*mlx.Array{mlx.Add(residual, update)} + }, + mlx.Shapeless(), + ) +) + type hyperConnectionState struct { residual *mlx.Array normed *mlx.Array + inject *mlx.Array } func expandStreams(x *mlx.Array, cfg *Config) *mlx.Array { @@ -20,12 +50,28 @@ func (h *hyperConnection) Prepare(residual *mlx.Array, cfg *Config) (*mlx.Array, 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))) + var down, inject *mlx.Array + if h.PackedInput != nil { + packed := h.PackedInput.Forward(normed) + down = mlx.SliceStartStop(packed, + []int32{0, 0, 0}, + []int32{B, L, h.MixDownDim}, + ) + inject = mlx.SliceStartStop(packed, + []int32{0, 0, h.MixDownDim}, + []int32{B, L, int32(packed.Dim(2))}, + ) + } else { + down = h.InputMixDown.Forward(normed) + } + count := mlx.NewScalarArray(float32(cfg.HCCount)).AsType(down.DType()) + mixHidden := hyperConnectionSiLU(down, count) 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} + count = mlx.NewScalarArray(float32(cfg.HCCount)).AsType(streams.DType()) + branch := hyperConnectionMix(streams, mix, count) + return branch, hyperConnectionState{residual: residual, normed: normed, inject: inject} } // Inject broadcasts one branch result back into the residual streams. The @@ -34,12 +80,14 @@ func (h *hyperConnection) Prepare(residual *mlx.Array, cfg *Config) (*mlx.Array, 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) + logits := state.inject + if logits == nil { + logits = h.BlockInject.Forward(state.normed) + } + count := mlx.NewScalarArray(float32(cfg.HCCount)).AsType(logits.DType()) residual := mlx.Reshape(state.residual, B, L, cfg.HCCount, cfg.HiddenSize) - return mlx.Reshape(mlx.Add(residual, update), B, L, cfg.HCCount*cfg.HiddenSize) + out := hyperConnectionInject(residual, branch, logits, count)[0] + return mlx.Reshape(out, B, L, cfg.HCCount*cfg.HiddenSize) } // Reduce applies the trained input mixer as the final stream reduction. The diff --git a/mlxrunner/model/qwen4_exp/hyper_connection_test.go b/mlxrunner/model/qwen4_exp/hyper_connection_test.go index 4564da917..1a7893105 100644 --- a/mlxrunner/model/qwen4_exp/hyper_connection_test.go +++ b/mlxrunner/model/qwen4_exp/hyper_connection_test.go @@ -11,6 +11,9 @@ import ( func TestHyperConnectionMatchesReferenceFormula(t *testing.T) { mlxtest.Run(t, func(t *mlxtest.T) { + mlx.EnableCompile() + defer mlx.DisableCompile() + 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}} @@ -18,11 +21,13 @@ func TestHyperConnectionMatchesReferenceFormula(t *testing.T) { injectWeight := [][]float32{{0.2, -0.1, 0.3, 0.4}, {-0.3, 0.2, 0.1, 0.5}} input := []float32{1, 2, 3, 4} + downProj := nn.NewLinear(matrix(downWeight), nil) + injectProj := nn.NewLinear(matrix(injectWeight), nil) 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), + Norm: &streamRMSNorm{Weight: mlx.FromValues(normWeight, 2, 2)}, + InputMixUp: nn.NewLinear(matrix(upWeight), nil), + PackedInput: packHyperConnectionInput(downProj, injectProj), + MixDownDim: downProj.OutputDim(), } residual := mlx.FromValues(input, 1, 1, 4) branch, state := h.Prepare(residual, cfg) diff --git a/mlxrunner/model/qwen4_exp/weights.go b/mlxrunner/model/qwen4_exp/weights.go index b2321ead2..5a84842e8 100644 --- a/mlxrunner/model/qwen4_exp/weights.go +++ b/mlxrunner/model/qwen4_exp/weights.go @@ -14,6 +14,8 @@ type hyperConnection struct { InputMixDown nn.LinearLayer InputMixUp nn.LinearLayer BlockInject nn.LinearLayer + PackedInput nn.LinearLayer + MixDownDim int32 } type linearAttention struct { @@ -192,6 +194,7 @@ func loadHyperConnection(linears model.LinearFactory, tensors map[string]*mlx.Ar if h.InputMixDown, err = requiredLinear(linears, prefix+".input_mix_weight_down"); err != nil { return nil, err } + h.MixDownDim = h.InputMixDown.OutputDim() if h.InputMixUp, err = requiredLinear(linears, prefix+".input_mix_weight_up"); err != nil { return nil, err } @@ -199,10 +202,46 @@ func loadHyperConnection(linears model.LinearFactory, tensors map[string]*mlx.Ar if h.BlockInject, err = requiredLinear(linears, prefix+".block_inject_weight"); err != nil { return nil, err } + if packed := packHyperConnectionInput(h.InputMixDown, h.BlockInject); packed != nil { + h.PackedInput = packed + h.InputMixDown = nil + h.BlockInject = nil + } } return h, nil } +func packHyperConnectionInput(down, inject nn.LinearLayer) nn.LinearLayer { + switch down := down.(type) { + case *nn.Linear: + inject, ok := inject.(*nn.Linear) + if !ok || down.Bias != nil || inject.Bias != nil || down.Weight.Dim(1) != inject.Weight.Dim(1) { + return nil + } + return nn.NewLinear(mlx.Concatenate([]*mlx.Array{down.Weight, inject.Weight}, 0), nil) + case *nn.QuantizedLinear: + inject, ok := inject.(*nn.QuantizedLinear) + if !ok || down.Bias != nil || inject.Bias != nil || + down.Scales == nil || inject.Scales == nil || + down.QBiases != nil || inject.QBiases != nil || + down.GlobalScale != nil || inject.GlobalScale != nil || + down.GroupSize != inject.GroupSize || down.Bits != inject.Bits || down.Mode != inject.Mode || + down.Weight.NumDims() != 2 || inject.Weight.NumDims() != 2 || + down.Scales.NumDims() != 2 || inject.Scales.NumDims() != 2 || + down.Weight.Dim(1) != inject.Weight.Dim(1) || down.Scales.Dim(1) != inject.Scales.Dim(1) { + return nil + } + return &nn.QuantizedLinear{ + Weight: mlx.Concatenate([]*mlx.Array{down.Weight, inject.Weight}, 0), + Scales: mlx.Concatenate([]*mlx.Array{down.Scales, inject.Scales}, 0), + GroupSize: down.GroupSize, + Bits: down.Bits, + Mode: down.Mode, + } + } + return nil +} + func loadLinearAttention(linears model.LinearFactory, tensors map[string]*mlx.Array, prefix string) (*linearAttention, error) { a := &linearAttention{} var err error