| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import os |
| from dataclasses import dataclass, field |
| from functools import partial |
| from typing import Any, Dict, List, Optional, Tuple |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
| try: |
| from . import qwen3_5 as q35 |
| from .base import BaseModelArgs, create_attention_mask, create_ssm_mask |
| from .cache import ArraysCache, KVCache |
| except ImportError: |
| from mlx_lm.models import qwen3_5 as q35 |
| from mlx_lm.models.base import BaseModelArgs, create_attention_mask, create_ssm_mask |
| from mlx_lm.models.cache import ArraysCache, KVCache |
|
|
|
|
| @dataclass |
| class ModelArgs(BaseModelArgs): |
| model_type: str |
| text_config: dict |
| loop_settings: dict = field(default_factory=dict) |
| mhc: dict = field(default_factory=dict) |
|
|
|
|
| def execution_plan(num_layers: int, loop: dict, mhc: dict) -> Tuple[int, ...]: |
| if not loop: |
| return tuple(range(num_layers)) |
| start, end = int(loop["start"]), int(loop["end"]) |
| count = int(loop.get("eval_count", (loop.get("counts") or [1])[0])) |
| if "counts" in loop and list(loop["counts"]) != [count]: |
| raise ValueError(f"loop_settings.counts {loop['counts']} must be [{count}]") |
| if not (0 <= start < end <= num_layers) or count < 1: |
| raise ValueError(f"bad loop_settings {loop}") |
| if mhc: |
| ms, me = int(mhc["start"]), int(mhc["end"]) |
| inside = (start <= ms < me <= end) or (me <= start) or (ms >= end) |
| if not (0 <= ms < me <= num_layers) or not inside: |
| raise ValueError(f"mhc range [{ms}, {me}) must sit inside one loop phase") |
| return tuple(range(start)) + tuple(range(start, end)) * count + tuple(range(end, num_layers)) |
|
|
|
|
| def sinkhorn(logits: mx.array, iterations: int) -> mx.array: |
| """Log-space Sinkhorn: rows then columns, `iterations` times, then exp (fp32).""" |
| logits = logits.astype(mx.float32) |
| for _ in range(iterations): |
| logits = logits - mx.logsumexp(logits, axis=-1, keepdims=True) |
| logits = logits - mx.logsumexp(logits, axis=-2, keepdims=True) |
| return mx.exp(logits) |
|
|
|
|
| class MHCRouter(nn.Module): |
| """fp32 router. Kept as raw parameters (not nn.Linear) so convert never quantizes it.""" |
|
|
| def __init__(self, dim: int, streams: int, iterations: int, fixed_identity: bool, pre_scale: float): |
| super().__init__() |
| self.streams = streams |
| self.iterations = iterations |
| self.fixed_identity = fixed_identity |
| self.pre_scale = pre_scale |
| out = 2 * streams + (0 if fixed_identity else streams * streams) |
| self.weight = mx.zeros((out, streams * dim), dtype=mx.float32) |
| self.bias = mx.zeros((out,), dtype=mx.float32) |
|
|
| def __call__(self, residual: mx.array): |
| |
| fn = _route_fn(self.streams, self.iterations, self.fixed_identity, self.pre_scale) |
| branch, post, mix = fn(residual, self.weight, self.bias) |
| return branch, post, (None if self.fixed_identity else mix) |
|
|
| @staticmethod |
| def merge(residual: mx.array, branch: mx.array, post: mx.array, mix: Optional[mx.array]) -> mx.array: |
| if mix is None: |
| return residual + post.astype(branch.dtype)[..., None] * branch[..., None, :] |
| return _merge_fn()(residual, branch, post, mix) |
|
|
|
|
| def _route(residual, weight, bias, s, iterations, fixed_identity, pre_scale): |
| flat = residual.reshape(*residual.shape[:-2], -1).astype(mx.float32) |
| normalized = flat * mx.rsqrt(mx.mean(mx.square(flat), axis=-1, keepdims=True) + 1e-6) |
| logits = normalized @ weight.astype(mx.float32).T + bias.astype(mx.float32) |
| pre = mx.sigmoid(logits[..., :s]) |
| post = 2 * mx.sigmoid(logits[..., s : 2 * s]) |
| if fixed_identity: |
| mix = mx.zeros((1,), dtype=mx.float32) |
| else: |
| mix = sinkhorn(logits[..., 2 * s :].reshape(*logits.shape[:-1], s, s), iterations) |
| branch = (pre.astype(residual.dtype)[..., None, :] @ residual).squeeze(-2) |
| if pre_scale != 1.0: |
| branch = branch * pre_scale |
| return branch, post, mix |
|
|
|
|
| def _merge(residual, branch, post, mix): |
| return mix.astype(residual.dtype) @ residual + post.astype(branch.dtype)[..., None] * branch[..., None, :] |
|
|
|
|
| |
| |
| _COMPILE = os.environ.get("SN_NO_COMPILE", "0") != "1" |
| _ROUTE_CACHE: Dict[tuple, Any] = {} |
|
|
|
|
| def _route_fn(s, iterations, fixed_identity, pre_scale): |
| key = (s, iterations, fixed_identity, pre_scale) |
| if key not in _ROUTE_CACHE: |
| f = partial(_route, s=s, iterations=iterations, fixed_identity=fixed_identity, pre_scale=pre_scale) |
| _ROUTE_CACHE[key] = mx.compile(f) if _COMPILE else f |
| return _ROUTE_CACHE[key] |
|
|
|
|
| _MERGE = [] |
|
|
|
|
| def _merge_fn(): |
| if not _MERGE: |
| _MERGE.append(mx.compile(_merge) if _COMPILE else _merge) |
| return _MERGE[0] |
|
|
|
|
| class MHCDecoderLayer(q35.DecoderLayer): |
| def __init__(self, args: q35.TextModelArgs, layer_idx: int, mhc: dict): |
| super().__init__(args, layer_idx) |
| self.mhc_start, self.mhc_end = int(mhc["start"]), int(mhc["end"]) |
| self.layer_idx = layer_idx |
| self.streams = int(mhc["streams"]) |
| kw = dict( |
| dim=args.hidden_size, |
| streams=self.streams, |
| iterations=int(mhc["iterations"]), |
| fixed_identity=bool(mhc.get("fixed_identity", False)), |
| pre_scale=float(mhc.get("pre_scale", 1.0)), |
| ) |
| self.mhc_attn = MHCRouter(**kw) |
| self.mhc_ffn = MHCRouter(**kw) |
|
|
| def __call__(self, x: mx.array, mask=None, cache=None) -> mx.array: |
| |
| if self.layer_idx == self.mhc_start: |
| residual = mx.broadcast_to(x[..., None, :], (*x.shape[:-1], self.streams, x.shape[-1])) |
| else: |
| residual = x |
| branch, post, mix = self.mhc_attn(residual) |
| branch = self.input_layernorm(branch) |
| if self.is_linear: |
| branch = self.linear_attn(branch, mask, cache) |
| else: |
| branch = self.self_attn(branch, mask, cache) |
| residual = MHCRouter.merge(residual, branch, post, mix) |
| branch, post, mix = self.mhc_ffn(residual) |
| branch = self.mlp(self.post_attention_layernorm(branch)) |
| out = MHCRouter.merge(residual, branch, post, mix) |
| |
| return out.mean(axis=-2) if self.layer_idx == self.mhc_end - 1 else out |
|
|
|
|
| class SenseNovaLoopTextModel(nn.Module): |
| def __init__(self, args: q35.TextModelArgs, loop: dict, mhc: dict): |
| super().__init__() |
| self.args = args |
| self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) |
| ms, me = (int(mhc["start"]), int(mhc["end"])) if mhc else (0, 0) |
| self.layers = [ |
| MHCDecoderLayer(args, i, mhc) if ms <= i < me else q35.DecoderLayer(args, i) |
| for i in range(args.num_hidden_layers) |
| ] |
| self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) |
| self.plan = execution_plan(args.num_hidden_layers, loop, mhc) |
| self.fa_visit = next((e for e, p in enumerate(self.plan) if not self.layers[p].is_linear), None) |
| self.ssm_visit = next((e for e, p in enumerate(self.plan) if self.layers[p].is_linear), None) |
|
|
| def __call__(self, inputs: mx.array, cache=None, input_embeddings: Optional[mx.array] = None) -> mx.array: |
| h = input_embeddings if input_embeddings is not None else self.embed_tokens(inputs) |
| if cache is None: |
| cache = [None] * len(self.plan) |
| if len(cache) != len(self.plan): |
| raise ValueError(f"cache has {len(cache)} slots, the execution plan needs {len(self.plan)}") |
| fa_mask = create_attention_mask(h, cache[self.fa_visit]) if self.fa_visit is not None else None |
| ssm_mask = create_ssm_mask(h, cache[self.ssm_visit]) if self.ssm_visit is not None else None |
| for physical, c in zip(self.plan, cache): |
| layer = self.layers[physical] |
| h = layer(h, mask=ssm_mask if layer.is_linear else fa_mask, cache=c) |
| return self.norm(h) |
|
|
|
|
| class TextModel(q35.TextModel): |
| def __init__(self, args: q35.TextModelArgs, loop: dict, mhc: dict): |
| nn.Module.__init__(self) |
| self.args = args |
| self.model_type = args.model_type |
| self.model = SenseNovaLoopTextModel(args, loop, mhc) |
| if not args.tie_word_embeddings: |
| self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) |
| self.sharding_group = None |
|
|
| @property |
| def layers(self): |
| return self.model.layers |
|
|
| def make_cache(self): |
| return [ArraysCache(size=2) if self.model.layers[p].is_linear else KVCache() for p in self.model.plan] |
|
|
| @property |
| def cast_predicate(self): |
| def predicate(path: str): |
| return not (path.endswith("A_log") or ".mhc_attn." in path or ".mhc_ffn." in path) |
|
|
| return predicate |
|
|
|
|
| class Model(nn.Module): |
| def __init__(self, args: ModelArgs): |
| super().__init__() |
| self.args = args |
| self.model_type = args.model_type |
| self.language_model = TextModel(q35.TextModelArgs.from_dict(args.text_config), args.loop_settings, args.mhc) |
|
|
| def __call__(self, inputs: mx.array, cache=None, input_embeddings: Optional[mx.array] = None): |
| return self.language_model(inputs, cache=cache, input_embeddings=input_embeddings) |
|
|
| @property |
| def model(self): |
| return self.language_model.model |
|
|
| def sanitize(self, weights): |
| |
| |
| weights = q35.Model.sanitize(self, weights) |
| return {k: (v.astype(mx.float32) if (".mhc_attn." in k or ".mhc_ffn." in k) else v) for k, v in weights.items()} |
|
|
| @property |
| def layers(self): |
| return self.language_model.model.layers |
|
|
| def make_cache(self): |
| return self.language_model.make_cache() |
|
|
| @property |
| def quant_predicate(self): |
| return None |
|
|
| @property |
| def cast_predicate(self): |
| return self.language_model.cast_predicate |
|
|
