Merge 257aa8ca63 into c1c536297f

2026-10-20 18:41:17 +00:00 · 2026-10-12 09:51:08 +03:00 · 2026-10-12 09:51:08 +03:00 · 07c5ade47b
commit 07c5ade47b
parent c1c536297f 257aa8ca63
16 changed files with 1053 additions and 130 deletions
--- 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"))
+}
--- 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 {
--- a/mlx/indexed_attention.go
+++ 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<float>(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<float>(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<float>(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<float>(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<float>(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<InT>(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 <metal_simdgroup>\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]
+}
--- a/mlx/indexed_attention_test.go
+++ 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)
+				}
+			}
+		})
+	})
+}
--- 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)
@ -126,8 +134,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))
--- 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...)
 }

@ -62,6 +72,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)
--- 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
--- 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)
--- a/mlxrunner/model/qwen4_exp/ple_host.go
+++ 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)
+}
--- a/mlxrunner/model/qwen4_exp/ple_host_test.go
+++ 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
+}
--- a/mlxrunner/model/qwen4_exp/ple_mapping_other.go
+++ 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()
+}
--- a/mlxrunner/model/qwen4_exp/ple_mapping_unix.go
+++ 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
+}
--- 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)
 }
--- 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) {
--- 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
 }

--- 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 {
@ -120,6 +122,7 @@ type PLE struct {
 	HeadVocabSizes   *mlx.Array
 	EmbeddingShards  []nn.EmbeddingLayer
 	ShardRows        int
+	hostEmbedding    *hostPLETable
 }

 type streamRMSNorm struct {
@ -191,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
 	}
@ -198,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
@ -339,6 +379,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)