| """YuE2 AR/NAR Mixture-of-Transformers backbone in MLX. |
| |
| Mirrors m-a-p/YuE2-3B ``modeling_yue2.py``: a Qwen3-style decoder whose every |
| layer carries two attention/MLP paths (AR for text+codec tokens, NAR for the |
| acoustic flow-matching state). Weight names follow the checkpoint except |
| ``time_embedder.mlp.{0,2}`` which ``convert.py`` renames to ``fc1``/``fc2``. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import math |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
|
|
| class KVCache: |
| """Append-only per-layer cache, grown in 256-slot steps (mlx_lm pattern).""" |
|
|
| step = 256 |
|
|
| def __init__(self): |
| self.keys = self.values = None |
| self.offset = 0 |
|
|
| def update(self, k, v): |
| prev = self.offset |
| B, H, T, D = k.shape |
| if self.keys is None or prev + T > self.keys.shape[2]: |
| n = ((T + self.step - 1) // self.step) * self.step |
| new_k, new_v = mx.zeros((B, H, n, D), k.dtype), mx.zeros((B, H, n, D), v.dtype) |
| if self.keys is None: |
| self.keys, self.values = new_k, new_v |
| else: |
| if prev % self.step: |
| self.keys, self.values = self.keys[..., :prev, :], self.values[..., :prev, :] |
| self.keys = mx.concatenate([self.keys, new_k], axis=2) |
| self.values = mx.concatenate([self.values, new_v], axis=2) |
| self.offset += T |
| self.keys[..., prev:self.offset, :] = k |
| self.values[..., prev:self.offset, :] = v |
| return self.keys[..., :self.offset, :], self.values[..., :self.offset, :] |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.n_heads, self.n_kv, self.head_dim = cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] |
| h, d = cfg["hidden_size"], self.head_dim |
| self.q_proj = nn.Linear(h, self.n_heads * d, bias=False) |
| self.k_proj = nn.Linear(h, self.n_kv * d, bias=False) |
| self.v_proj = nn.Linear(h, self.n_kv * d, bias=False) |
| self.o_proj = nn.Linear(self.n_heads * d, h, bias=False) |
| self.q_norm = nn.RMSNorm(d, eps=cfg["rms_norm_eps"]) |
| self.k_norm = nn.RMSNorm(d, eps=cfg["rms_norm_eps"]) |
| self.rope_base = float(cfg["rope_theta"]) |
| self.scale = d ** -0.5 |
|
|
| def project_qkv(self, x, offset): |
| """Return q [B,H,T,D], k/v [B,KV,T,D] after QK-norm and RoPE at ``offset``.""" |
| B, T, _ = x.shape |
| q = self.q_norm(self.q_proj(x).reshape(B, T, self.n_heads, -1)).transpose(0, 2, 1, 3) |
| k = self.k_norm(self.k_proj(x).reshape(B, T, self.n_kv, -1)).transpose(0, 2, 1, 3) |
| v = self.v_proj(x).reshape(B, T, self.n_kv, -1).transpose(0, 2, 1, 3) |
| rope = lambda a: mx.fast.rope(a, self.head_dim, traditional=False, base=self.rope_base, scale=1.0, offset=offset) |
| return rope(q), rope(k), v |
|
|
| def attend(self, q, k, v, mask): |
| B, _, T, _ = q.shape |
| o = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask) |
| return self.o_proj(o.transpose(0, 2, 1, 3).reshape(B, T, -1)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| h, i = cfg["hidden_size"], cfg["intermediate_size"] |
| self.gate_proj = nn.Linear(h, i, bias=False) |
| self.up_proj = nn.Linear(h, i, bias=False) |
| self.down_proj = nn.Linear(i, h, bias=False) |
|
|
| def __call__(self, x): |
| return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class DecoderLayer(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| h, eps = cfg["hidden_size"], cfg["rms_norm_eps"] |
| self.input_layernorm = nn.RMSNorm(h, eps=eps) |
| self.self_attn = Attention(cfg) |
| self.post_attention_layernorm = nn.RMSNorm(h, eps=eps) |
| self.mlp = MLP(cfg) |
| self.nar_input_layernorm = nn.RMSNorm(h, eps=eps) |
| self.nar_self_attn = Attention(cfg) |
| self.nar_pre_mlp_layernorm = nn.RMSNorm(h, eps=eps) |
| self.nar_mlp = MLP(cfg) |
|
|
| def ar(self, x, cache: KVCache, mask): |
| q, k, v = self.self_attn.project_qkv(self.input_layernorm(x), cache.offset) |
| k, v = cache.update(k, v) |
| x = x + self.self_attn.attend(q, k, v, mask) |
| return x + self.mlp(self.post_attention_layernorm(x)) |
|
|
| def ar_prefill_kv(self, x): |
| """AR pass that also returns this layer's (k, v) for the NAR to attend over.""" |
| q, k, v = self.self_attn.project_qkv(self.input_layernorm(x), 0) |
| x = x + self.self_attn.attend(q, k, v, "causal") |
| return x + self.mlp(self.post_attention_layernorm(x)), (k, v) |
|
|
| def nar(self, x, ar_k, ar_v, offset): |
| q, k, v = self.nar_self_attn.project_qkv(self.nar_input_layernorm(x), offset) |
| k, v = mx.concatenate([ar_k, k], axis=2), mx.concatenate([ar_v, v], axis=2) |
| x = x + self.nar_self_attn.attend(q, k, v, None) |
| return x + self.nar_mlp(self.nar_pre_mlp_layernorm(x)) |
|
|
|
|
| class Backbone(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.embed_tokens = nn.Embedding(cfg["vocab_size"], cfg["hidden_size"]) |
| self.layers = [DecoderLayer(cfg) for _ in range(cfg["num_hidden_layers"])] |
| self.norm = nn.RMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"]) |
|
|
|
|
| class TimestepEmbedder(nn.Module): |
| def __init__(self, hidden, freq_size=256): |
| super().__init__() |
| self.freq_size = freq_size |
| self.fc1 = nn.Linear(freq_size, hidden) |
| self.fc2 = nn.Linear(hidden, hidden) |
|
|
| def __call__(self, t, dtype): |
| half = self.freq_size // 2 |
| freqs = mx.exp(-math.log(10000) * mx.arange(half, dtype=mx.float32) / half) |
| args = t.astype(mx.float32).reshape(-1, 1) * freqs[None] |
| emb = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1).astype(dtype) |
| return self.fc2(nn.silu(self.fc1(emb))) |
|
|
|
|
| class LatentPosEmbed(nn.Module): |
| """Holds the checkpoint's precomputed sinusoidal table ``pe`` [max_frames, hidden].""" |
|
|
| def __init__(self, max_frames, hidden): |
| super().__init__() |
| self.pe = mx.zeros((max_frames, hidden)) |
|
|
|
|
| class Yue2Model(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.config = cfg |
| h, d = cfg["hidden_size"], cfg["latent_dim"] |
| self.model = Backbone(cfg) |
| self.lm_head = nn.Linear(h, cfg["vocab_size"], bias=False) |
| self.llm2vae = nn.Linear(h, d) |
| self.vae2llm = nn.Linear(d, h) |
| self.time_embedder = TimestepEmbedder(h) |
| self.latent_pos_embed = LatentPosEmbed(cfg["max_latent_frames"], h) |
|
|
| |
| def ar_step(self, tokens, caches, all_positions=False): |
| """tokens [B,T] → logits [B,V] of the last position (or [B,T,V]).""" |
| x = self.model.embed_tokens(tokens) |
| mask = "causal" if tokens.shape[1] > 1 else None |
| for layer, cache in zip(self.model.layers, caches): |
| x = layer.ar(x, cache, mask) |
| x = self.model.norm(x if all_positions else x[:, -1:]) |
| logits = self.lm_head(x) |
| return logits if all_positions else logits[:, -1] |
|
|
| |
| def nar_prefill(self, ar_tokens): |
| """Run the AR path once over the chunk prefix; return per-layer (k, v).""" |
| x = self.model.embed_tokens(mx.array([ar_tokens])) |
| cache = [] |
| for layer in self.model.layers: |
| x, kv = layer.ar_prefill_kv(x) |
| cache.append(kv) |
| mx.eval(cache) |
| return cache |
|
|
| def nar_velocity(self, state, raw_t, ar_cache, ar_length): |
| """Flow-matching velocity for ``state`` [T,64] at logit-time ``raw_t``.""" |
| shift = self.config["timestep_shift"] |
| sig = mx.sigmoid(mx.array(raw_t, dtype=state.dtype)) |
| t = shift * sig / (1 + (shift - 1) * sig) |
| n = state.shape[0] + 2 |
| x = self.vae2llm(mx.pad(state, [(1, 1), (0, 0)])[None]) |
| x = x + self.time_embedder(t, state.dtype)[None] |
| pos = mx.minimum(mx.arange(n), self.config["max_latent_frames"] - 1) |
| x = x + self.latent_pos_embed.pe[pos][None] |
| for layer, (k, v) in zip(self.model.layers, ar_cache): |
| x = layer.nar(x, k, v, ar_length) |
| return self.llm2vae(self.model.norm(x))[0, 1:-1] |
|
|
|
|
| def load_model(path: Path) -> Yue2Model: |
| path = Path(path) |
| cfg = json.loads((path / "config.json").read_text()) |
| model = Yue2Model(cfg) |
| weights = mx.load(str(path / "model.safetensors")) |
| q = cfg.get("quantization") |
| if q: |
| nar_bits = q.get("nar_bits", q["bits"]) |
|
|
| def predicate(path, module): |
| if not isinstance(module, nn.Linear) or f"{path}.scales" not in weights: |
| return False |
| return {"group_size": q["group_size"], "bits": nar_bits if ".nar_" in path else q["bits"]} |
|
|
| nn.quantize(model, q["group_size"], q["bits"], class_predicate=predicate) |
| model.load_weights(list(weights.items())) |
| mx.eval(model.parameters()) |
| return model |
|
|
