| """Experimental MLX runtime for folded dense qwen35 GGUF checkpoints.""" |
|
|
| import math |
| import sys |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import numpy as np |
| from codec import transcode |
| from mlx.utils import tree_flatten |
| from mlx_lm.models.qwen3_5 import TextModel, TextModelArgs |
|
|
| from mlx import nn |
|
|
|
|
| def fwht(x, block, signs, inverse=False): |
| shape, dtype = x.shape, x.dtype |
| if shape[-1] % block: |
| raise ValueError("Hadamard block does not divide activation width") |
| x = x.astype(mx.float32) |
| if not inverse: |
| x = x * signs |
| x = mx.hadamard_transform(x.reshape(-1, block), scale=1 / math.sqrt(block)).reshape( |
| shape |
| ) |
| if inverse: |
| x = x * signs |
| return x.astype(dtype) |
|
|
|
|
| class Packed(nn.Module): |
| def __init__(self, arrays, block=0, signs=None, embedding=False, dtype=mx.float16): |
| super().__init__() |
| self.weight, self.scales, self.biases = [mx.array(a) for a in arrays] |
| self.block, self.signs, self.embedding, self.dtype = ( |
| block, |
| signs, |
| embedding, |
| dtype, |
| ) |
|
|
| def __call__(self, x): |
| if self.embedding: |
| shape = x.shape |
| indices = x.reshape(-1) |
| out = ( |
| mx.dequantize( |
| self.weight[indices], |
| self.scales[indices], |
| self.biases[indices], |
| group_size=128, |
| bits=2, |
| ) |
| .reshape(*shape, -1) |
| .astype(self.dtype) |
| ) |
| return ( |
| fwht(out, self.block, self.signs, inverse=True) if self.block else out |
| ) |
| if self.block: |
| x = fwht(x, self.block, self.signs) |
| return mx.quantized_matmul( |
| x, |
| self.weight, |
| self.scales, |
| self.biases, |
| transpose=True, |
| group_size=128, |
| bits=2, |
| ) |
|
|
|
|
| def load(gguf_path, gguf_python, dtype=mx.float16): |
| sys.path.insert(0, str(gguf_python)) |
| from gguf import GGUFReader |
|
|
| reader = GGUFReader(str(gguf_path)) |
| fields = {k: v.contents() for k, v in reader.fields.items()} |
| if fields.get("general.architecture") != "qwen35": |
| raise ValueError("Only dense qwen35 is supported by this experimental loader") |
| g = lambda k, default=None: fields.get("qwen35." + k, default) |
| tensors = {t.name: t for t in reader.tensors} |
| if "output.weight" not in tensors: |
| raise ValueError( |
| "Explicit output head required; tied-head policy must be resolved separately" |
| ) |
| block = int(fields["prism.hadamard.block_size"]) |
| if block not in (512, 1024, 2048, 4096): |
| raise ValueError("Unvalidated Hadamard block size") |
| if fields.get("prism.hadamard.sign_mode") != "explicit": |
| raise ValueError("Explicit signs required") |
| widths = fields["prism.hadamard.sign_widths"] |
| values = fields["prism.hadamard.sign_values"] |
| signs, offset = {}, 0 |
| for width in widths: |
| a = np.asarray(values[offset : offset + width], dtype=np.float32) |
| if len(a) != width or not np.isin(a, [-1, 1]).all(): |
| raise ValueError("Invalid sign vector") |
| signs[width] = mx.array(a) |
| offset += width |
| if offset != len(values): |
| raise ValueError("Trailing sign values") |
| folded = set(fields["prism.hadamard.weight_names"]) |
| inverse = set(fields.get("prism.hadamard.inverse_weight_names", [])) |
| if inverse != {"token_embd.weight"}: |
| raise ValueError("Unexpected inverse-transform manifest") |
| if folded & inverse: |
| raise ValueError("Forward and inverse transform manifests overlap") |
| if not (folded | inverse) <= tensors.keys(): |
| raise ValueError("Transform manifest references missing tensors") |
| nv, nk = int(g("ssm.time_step_rank")), int(g("ssm.group_count")) |
| if nv <= 0 or nk <= 0 or nv % nk or int(g("ssm.inner_size")) % nv: |
| raise ValueError("Invalid GDN head dimensions") |
| hd = int(g("ssm.inner_size")) // nv |
| hk = int(g("ssm.state_size")) |
| cfg = { |
| "model_type": "qwen3_5_text", |
| "hidden_size": int(g("embedding_length")), |
| "intermediate_size": int(g("feed_forward_length")), |
| "num_hidden_layers": int(g("block_count")), |
| "num_attention_heads": int(g("attention.head_count")), |
| "num_key_value_heads": int(g("attention.head_count_kv")), |
| "head_dim": int(g("attention.key_length")), |
| "rms_norm_eps": float(g("attention.layer_norm_rms_epsilon")), |
| "vocab_size": int(tensors["token_embd.weight"].shape[1]), |
| "max_position_embeddings": int(g("context_length")), |
| "linear_num_value_heads": nv, |
| "linear_num_key_heads": nk, |
| "linear_value_head_dim": hd, |
| "linear_key_head_dim": hk, |
| "linear_conv_kernel_dim": int(g("ssm.conv_kernel")), |
| "full_attention_interval": int(g("full_attention_interval")), |
| "tie_word_embeddings": False, |
| "rope_parameters": { |
| "type": "default", |
| "rope_theta": float(g("rope.freq_base")), |
| "partial_rotary_factor": int(g("rope.dimension_count")) |
| / int(g("attention.key_length")), |
| }, |
| } |
| model = TextModel(TextModelArgs.from_dict(cfg)) |
| expected = {name: value.shape for name, value in tree_flatten(model.parameters())} |
| mapping = { |
| "attn_norm.weight": "input_layernorm.weight", |
| "post_attention_norm.weight": "post_attention_layernorm.weight", |
| "ffn_gate.weight": "mlp.gate_proj.weight", |
| "ffn_up.weight": "mlp.up_proj.weight", |
| "ffn_down.weight": "mlp.down_proj.weight", |
| "attn_q.weight": "self_attn.q_proj.weight", |
| "attn_k.weight": "self_attn.k_proj.weight", |
| "attn_v.weight": "self_attn.v_proj.weight", |
| "attn_output.weight": "self_attn.o_proj.weight", |
| "attn_q_norm.weight": "self_attn.q_norm.weight", |
| "attn_k_norm.weight": "self_attn.k_norm.weight", |
| "attn_qkv.weight": "linear_attn.in_proj_qkv.weight", |
| "attn_gate.weight": "linear_attn.in_proj_z.weight", |
| "ssm_alpha.weight": "linear_attn.in_proj_a.weight", |
| "ssm_beta.weight": "linear_attn.in_proj_b.weight", |
| "ssm_out.weight": "linear_attn.out_proj.weight", |
| "ssm_norm.weight": "linear_attn.norm.weight", |
| "ssm_a": "linear_attn.A_log", |
| "ssm_dt.bias": "linear_attn.dt_bias", |
| "ssm_conv1d.weight": "linear_attn.conv1d.weight", |
| } |
| global_map = { |
| "output.weight": "lm_head.weight", |
| "output_norm.weight": "model.norm.weight", |
| "token_embd.weight": "model.embed_tokens.weight", |
| } |
|
|
| def vperm(unit): |
| return ( |
| np.arange(nv * unit) |
| .reshape(nv // nk, nk, unit) |
| .transpose(1, 0, 2) |
| .reshape(-1) |
| ) |
|
|
| def reorder(a, stem): |
| if nv == nk: |
| return a |
| if stem == "attn_qkv.weight": |
| qk = 2 * nk * hk |
| return np.concatenate([a[:qk], a[qk:][vperm(hd)]], axis=0) |
| if stem == "attn_gate.weight": |
| return a[vperm(hd)] |
| if stem in ("ssm_alpha.weight", "ssm_beta.weight", "ssm_a", "ssm_dt.bias"): |
| return a[vperm(1)] |
| if stem == "ssm_conv1d.weight": |
| qk = 2 * nk * hk |
| return np.concatenate([a[:qk], a[qk:][vperm(hd)]], axis=0) |
| return a |
|
|
| assigned = [] |
| for name, t in tensors.items(): |
| stem = name |
| if name in global_map: |
| target = global_map[name] |
| elif name.startswith("blk."): |
| _, layer, stem = name.split(".", 2) |
| if stem not in mapping: |
| raise ValueError("Unmapped tensor " + name) |
| target = f"model.layers.{layer}." + mapping[stem] |
| else: |
| raise ValueError("Unmapped tensor " + name) |
| if target not in expected or target in assigned: |
| raise ValueError("Unexpected or duplicate target " + target) |
| shape = tuple(int(n) for n in t.shape[::-1]) |
| if stem == "ssm_conv1d.weight": |
| shape += (1,) |
| if shape != expected[target]: |
| raise ValueError( |
| f"Shape mismatch for {name}: {shape} != {expected[target]}" |
| ) |
| parent = model |
| parts = target.split(".") |
| for part in ( |
| parts[:-2] if t.tensor_type.name in ("PQ2_0", "PTQ1_0") else parts[:-1] |
| ): |
| parent = parent[int(part)] if part.isdigit() else getattr(parent, part) |
| if t.tensor_type.name in ("PQ2_0", "PTQ1_0"): |
| if len(t.shape) != 2: |
| raise ValueError("Quantized tensor is not a matrix") |
| shape = tuple(int(n) for n in t.shape[::-1]) |
| arrays = transcode(t.data.tobytes(), shape, t.tensor_type.name) |
| arrays = tuple(reorder(a, stem) for a in arrays) |
| if ( |
| stem == "ssm_out.weight" |
| and nv != nk |
| and not fields.get("prism.hadamard.gdn_v_grouped", False) |
| ): |
| raise ValueError("Unimplemented ungrouped folded GDN output") |
| active = name in folded or name in inverse |
| module = Packed( |
| arrays, |
| block if active else 0, |
| signs[shape[1]] if active else None, |
| name in inverse, |
| dtype, |
| ) |
| setattr(parent, parts[-2], module) |
| mx.eval(module.parameters()) |
| else: |
| if t.tensor_type.name not in ("F32", "F16"): |
| raise ValueError("Unsupported auxiliary type " + t.tensor_type.name) |
| a = np.asarray(t.data).copy() |
| a = reorder(a, stem) |
| if stem == "ssm_a": |
| if not (a < 0).all(): |
| raise ValueError("Invalid stored SSM A") |
| a = np.log(-a) |
| if stem == "ssm_conv1d.weight": |
| a = a[..., None] |
| if name in folded or name in inverse: |
| raise ValueError("Unimplemented transformed float matrix") |
| setattr(parent, parts[-1], mx.array(a)) |
| assigned.append(target) |
| missing = expected.keys() - set(assigned) |
| if missing: |
| raise ValueError(f"Missing model parameters: {sorted(missing)}") |
| model.eval() |
| mx.eval(model.parameters()) |
| return ( |
| model, |
| { |
| "config": cfg, |
| "block": block, |
| "source": str(Path(gguf_path).resolve()), |
| "assigned_tensors": len(assigned), |
| "folded": len(folded), |
| "inverse": len(inverse), |
| }, |
| reader, |
| ) |
|
|
