Release Project Pyrite Phase 2: MoE-Stream Engine with Mixtral-8x22B & DeepSeek support
cab2ef1 verified | #!/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 | |
| 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 | |