#!/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