File size: 7,743 Bytes
cab2ef1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | #!/usr/bin/env python3
"""Streaming MoE Layer Implementation for Apple Silicon MLX with Quantized Metal Matmul.
Replaces standard heavy MoE layer. Keeps routing gate and shared experts resident in RAM,
while streaming routed experts on-demand from NVMe storage with an LRU cache.
"""
from collections import OrderedDict
from typing import Dict, List, Optional, Tuple, Any
import mlx.core as mx
import mlx.nn as nn
from expert_registry import SafetensorsIndex, ModelSafetensorsIndex
class ExpertLRUCache:
"""Least-Recently-Used (LRU) In-Memory Cache for MoE Experts."""
def __init__(self, max_experts_in_ram: int = 16):
self.max_experts = max_experts_in_ram
self.cache: OrderedDict[int, Dict[str, mx.array]] = OrderedDict()
def get(self, expert_id: int) -> Optional[Dict[str, mx.array]]:
if expert_id in self.cache:
self.cache.move_to_end(expert_id)
return self.cache[expert_id]
return None
def put(self, expert_id: int, weights: Dict[str, mx.array]):
if expert_id in self.cache:
self.cache.move_to_end(expert_id)
else:
if len(self.cache) >= self.max_experts:
self.cache.popitem(last=False)
self.cache[expert_id] = weights
class StreamingMoELayer(nn.Module):
"""Memory-efficient streaming MoE layer for Apple Silicon MLX.
Dynamically routes tokens to experts, streaming weights directly from
NVMe mmap indices without loading entire multi-gigabyte models into RAM.
"""
def __init__(
self,
layer_idx: int,
gate_weight: mx.array,
shared_experts: Optional[Dict[str, mx.array]],
index_reader: Any,
top_k: int = 6,
gate_scales: Optional[mx.array] = None,
gate_biases: Optional[mx.array] = None,
routed_scaling_factor: float = 1.0,
topk_method: str = "greedy",
n_group: int = 1,
topk_group: int = 1,
max_cache_experts: int = 16,
moe_prefix: str = "mlp",
softmax_before_topk: bool = True,
):
super().__init__()
self.layer_idx = layer_idx
self.gate_weight = gate_weight
self.gate_scales = gate_scales
self.gate_biases = gate_biases
self.shared = shared_experts
self.index_reader = index_reader
self.top_k = top_k
self.routed_scaling_factor = routed_scaling_factor
self.topk_method = topk_method
self.n_group = n_group
self.topk_group = topk_group
self.moe_prefix = moe_prefix
self.softmax_before_topk = softmax_before_topk
self.cache = ExpertLRUCache(max_experts_in_ram=max_cache_experts)
def _get_expert_tensors(self, expert_id: int) -> Dict[str, mx.array]:
cached = self.cache.get(expert_id)
if cached is not None:
return cached
l_str = f"model.layers.{self.layer_idx}.{self.moe_prefix}.switch_mlp"
weights = {
"gate_w": self.index_reader.get_expert_slice(f"{l_str}.gate_proj.weight", expert_id),
"gate_s": self.index_reader.get_expert_slice(f"{l_str}.gate_proj.scales", expert_id),
"gate_b": self.index_reader.get_expert_slice(f"{l_str}.gate_proj.biases", expert_id),
"up_w": self.index_reader.get_expert_slice(f"{l_str}.up_proj.weight", expert_id),
"up_s": self.index_reader.get_expert_slice(f"{l_str}.up_proj.scales", expert_id),
"up_b": self.index_reader.get_expert_slice(f"{l_str}.up_proj.biases", expert_id),
"down_w": self.index_reader.get_expert_slice(f"{l_str}.down_proj.weight", expert_id),
"down_s": self.index_reader.get_expert_slice(f"{l_str}.down_proj.scales", expert_id),
"down_b": self.index_reader.get_expert_slice(f"{l_str}.down_proj.biases", expert_id),
}
self.cache.put(expert_id, weights)
return weights
@staticmethod
def _qmatmul(x: mx.array, w: mx.array, s: mx.array, b: Optional[mx.array] = None) -> mx.array:
"""Dynamic quantized matrix multiplication auto-detecting bits and group size."""
k = x.shape[-1]
bits = (w.shape[-1] * 32) // k
group_size = k // s.shape[-1]
return mx.quantized_matmul(
x, w, scales=s, biases=b, transpose=True, group_size=group_size, bits=bits
)
def forward_expert(self, x: mx.array, exp_id: int) -> mx.array:
"""Run forward pass for a single expert."""
exp_t = self._get_expert_tensors(exp_id)
h_gate = self._qmatmul(x, exp_t["gate_w"], exp_t["gate_s"], exp_t["gate_b"])
h_up = self._qmatmul(x, exp_t["up_w"], exp_t["up_s"], exp_t["up_b"])
h_expert = nn.silu(h_gate) * h_up
return self._qmatmul(h_expert, exp_t["down_w"], exp_t["down_s"], exp_t["down_b"])
def __call__(self, x: mx.array) -> mx.array:
orig_shape = x.shape
x_flat = x.reshape(-1, orig_shape[-1])
num_tokens, hidden_dim = x_flat.shape
# 1. Router Gate (resident in RAM)
if self.gate_scales is not None:
logits = self._qmatmul(x_flat, self.gate_weight, self.gate_scales, self.gate_biases)
else:
logits = x_flat @ self.gate_weight.T
k = self.top_k
if self.softmax_before_topk:
scores = mx.softmax(logits, axis=-1, precise=True)
if self.topk_method == "group_limited_greedy" and self.n_group > 1:
grouped_scores = scores.reshape(num_tokens, self.n_group, -1)
group_scores = grouped_scores.max(axis=-1, keepdims=True)
k_grp = self.n_group - self.topk_group
group_idx = mx.argpartition(group_scores, kth=k_grp - 1, axis=-2)[..., :k_grp, :]
scores = mx.put_along_axis(
grouped_scores, group_idx, mx.array(0.0, scores.dtype), axis=-2
).reshape(num_tokens, -1)
top_k_indices = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
top_k_scores = mx.take_along_axis(scores, top_k_indices, axis=-1)
if self.routed_scaling_factor != 1.0:
top_k_scores = top_k_scores * self.routed_scaling_factor
else:
# Mixtral style: Top-k on logits first, then softmax across the top-k chosen logits
top_k_indices = mx.stop_gradient(mx.argpartition(-logits, kth=k - 1, axis=-1)[..., :k])
chosen_logits = mx.take_along_axis(logits, top_k_indices, axis=-1)
top_k_scores = mx.softmax(chosen_logits, axis=-1, precise=True)
# 2. Shared Experts (resident in RAM)
if self.shared is not None:
sh_gate = self._qmatmul(x_flat, self.shared["gate_w"], self.shared["gate_s"], self.shared.get("gate_b"))
sh_up = self._qmatmul(x_flat, self.shared["up_w"], self.shared["up_s"], self.shared.get("up_b"))
h_shared = nn.silu(sh_gate) * sh_up
out_shared = self._qmatmul(h_shared, self.shared["down_w"], self.shared["down_s"], self.shared.get("down_b"))
else:
out_shared = mx.zeros_like(x_flat)
# 3. Dynamic Routed Experts Streaming (Vectorized on Metal GPU)
out_routed = mx.zeros_like(x_flat)
unique_experts = list(set(top_k_indices.flatten().tolist()))
for exp_id in unique_experts:
weight = mx.where(
top_k_indices == exp_id,
top_k_scores,
mx.zeros_like(top_k_scores)
).sum(axis=-1, keepdims=True)
exp_out = self.forward_expert(x_flat, exp_id)
out_routed = out_routed + exp_out * weight
final_out = (out_shared + out_routed).reshape(orig_shape)
return final_out
|