| """Packed serving: compressed tensors stay compressed; base loaders must never run.""" |
| import json |
|
|
| import numpy as np |
| import pytest |
| import torch |
|
|
| mx = pytest.importorskip("mlx.core") |
| nn = pytest.importorskip("mlx.nn") |
| pytest.importorskip("mlx_lm") |
|
|
| from mlx.utils import tree_flatten |
| from mlx_lm.models import qwen3_5 |
| from tokenizers import Tokenizer, models, pre_tokenizers |
| from transformers import PreTrainedTokenizerFast |
|
|
| from kev.checkpoint import Checkpoint, LoadOptions |
| from kev.model import PointerHead, SPECIAL |
| from kev.packed_model import PackedLinear |
| from scripts.build_packed_model import build, pack_native |
|
|
|
|
| def test_native_codes_and_compensation(): |
| rng = np.random.default_rng(8) |
| for bits in (2, 3, 4): |
| codes = rng.integers(0, 3 if bits == 2 else 1 << bits, size=(8, 64), dtype=np.uint8) |
| scale = mx.array(rng.uniform(.01, .1, (8, 1)).astype(np.float16)) |
| bias = -scale |
| q = PackedLinear(64, 8, bits, 64, rank=2) |
| q.weight, q.scales, q.biases = mx.array(pack_native(codes, bits)), scale, bias |
| q.comp_a = mx.array(rng.normal(0, .01, (2, 64))).astype(mx.bfloat16) |
| q.comp_b = mx.array(rng.normal(0, .01, (8, 2))).astype(mx.bfloat16) |
| x = mx.array(rng.normal(size=(3, 64))).astype(mx.bfloat16) |
| dense = mx.dequantize(q.weight, scale, bias, group_size=64, bits=bits) |
| decoded = (mx.array(codes).astype(mx.float32) * scale.astype(mx.float32) + bias.astype(mx.float32)).astype(mx.float16) |
| np.testing.assert_array_equal(np.asarray(dense), np.asarray(decoded)) |
| expected = (x.astype(mx.float32) @ dense.astype(mx.float32).T + |
| (x.astype(mx.float32) @ q.comp_a.astype(mx.float32).T) @ q.comp_b.astype(mx.float32).T).astype(mx.bfloat16) |
| np.testing.assert_allclose(np.asarray(q(x).astype(mx.float32)), np.asarray(expected.astype(mx.float32)), atol=.008, rtol=.008) |
| assert q.weight.dtype == mx.uint32 and q.weight.nbytes == codes.size * bits // 8 |
|
|
|
|
| def tiny_bundle(tmp_path): |
| source = tmp_path / "source" |
| source.mkdir() |
| config = {"model_type": "qwen3_5_text", "hidden_size": 64, "intermediate_size": 128, "num_hidden_layers": 2, |
| "num_attention_heads": 2, "num_key_value_heads": 1, "head_dim": 32, "vocab_size": 32, |
| "linear_num_value_heads": 2, "linear_num_key_heads": 2, "linear_key_head_dim": 128, |
| "linear_value_head_dim": 128, "linear_conv_kernel_dim": 4, "full_attention_interval": 2, |
| "tie_word_embeddings": True, "layer_types": ["linear_attention", "full_attention"]} |
| (source / "config.json").write_text(json.dumps(config)) |
| vocab = {w: i for i, w in enumerate(["<unk>", "<pad>", *SPECIAL, "billing", "shipping", "charged", "twice"])} |
| tokenizer = Tokenizer(models.WordLevel(vocab, unk_token="<unk>")) |
| tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() |
| PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="<unk>", pad_token="<pad>", additional_special_tokens=SPECIAL).save_pretrained(source) |
| head = PointerHead(64, 8) |
| torch.save({"base": "offline/does-not-exist", "head": head.state_dict(), "head_dim": 8, "temperature": 2.4}, source / "head.pt") |
| lm = qwen3_5.Model(qwen3_5.ModelArgs.from_dict(config)) |
| rng = np.random.default_rng(12) |
| tensors, kinds, groups = {}, {}, {} |
| modules = dict(lm.named_modules()) |
| for name, value in tree_flatten(lm.parameters()): |
| module_name = name.removesuffix(".weight") |
| if name.endswith(".weight") and isinstance(modules.get(module_name), nn.Linear): |
| n, k = value.shape |
| group = min(64, k) |
| codes = rng.integers(0, 3, size=n * k, dtype=np.uint8) |
| pad = np.pad(codes, (0, (-len(codes)) % 5), constant_values=1).reshape(-1, 5) |
| tensors[name + ".trits"] = mx.array((pad.astype(np.uint16) * [1, 3, 9, 27, 81]).sum(-1).astype(np.uint8)) |
| tensors[name + ".scales"] = mx.full((n, k // group), .05, dtype=mx.float16) |
| tensors[name + ".shape"] = mx.array([n, k], dtype=mx.int32) |
| kinds[name], groups[name] = "tern", group |
| elif name == "language_model.model.embed_tokens.weight": |
| w, s, b = mx.quantize(value.astype(mx.bfloat16), group_size=64, bits=4) |
| tensors[name + ".q"], tensors[name + ".scales"], tensors[name + ".biases"] = w, s, b |
| else: |
| tensors[name] = value.astype(mx.bfloat16) |
| packed = tmp_path / "tiny.safetensors" |
| mx.save_safetensors(str(packed), tensors, metadata={"format": "kev-ternary-v1", "rot": "none", "emb_bits": "4", "seed": "0", "run": str(source), "kinds": json.dumps(kinds), "groups": json.dumps(groups)}) |
| output = tmp_path / "bundle" |
| build(packed, source, source / "config.json", output) |
| return output |
|
|
|
|
| def test_offline_standalone_and_strict_loading(tmp_path, monkeypatch): |
| bundle = tiny_bundle(tmp_path) |
|
|
| def forbidden(*args, **kwargs): |
| raise AssertionError("original model loader was called") |
|
|
| import kev.mlx_model |
| monkeypatch.setattr(kev.mlx_model, "load_base", forbidden) |
| monkeypatch.setattr(kev.mlx_model, "load_full", forbidden) |
| monkeypatch.setattr(kev.mlx_model, "merge_lora", forbidden) |
| monkeypatch.setattr("huggingface_hub.snapshot_download", forbidden) |
| checkpoint = Checkpoint(str(bundle)) |
| assert checkpoint.packed and not checkpoint.full and checkpoint.hybrid_base() |
| tok, model = checkpoint.load("mps", LoadOptions(backend="auto")) |
| assert model.dtype == "bfloat16" and model.head.q.in_features == 64 and model.head.temperature == 2.4 |
| hidden = model._hidden([[tok.convert_tokens_to_ids("billing"), tok.convert_tokens_to_ids("shipping")]]) |
| assert hidden.shape == (1, 2, 64) and np.isfinite(np.asarray(hidden.astype(mx.float32))).all() |
| from fastapi.testclient import TestClient |
| from kev.serve import Server, app |
| server = Server(checkpoint, tok, model, "mps") |
| monkeypatch.setattr(app.state, "server", server, raising=False) |
| try: |
| with TestClient(app) as client: |
| response = client.post("/v1/systemone", json={"state": "charged twice", "questions": { |
| "team": {"type": "choice", "instructions": "billing shipping", "criteria": {"billing": "billing", "shipping": "shipping"}}}}) |
| assert response.status_code == 200 |
| assert "answers" in response.json() and "usage" in response.json() |
| assert client.get("/v1/models").json()["models"][0]["backend"] == "mlx" |
| finally: |
| server.close() |
| for name, module in model.lm.named_modules(): |
| if isinstance(module, (nn.QuantizedLinear, nn.QuantizedEmbedding)): |
| assert module.weight.dtype == mx.uint32 |
| assert len(checkpoint.weights_sha256()) == 64 |
| with pytest.raises(ValueError, match="MLX backend"): |
| checkpoint.load("cpu", LoadOptions(backend="torch")) |
| weights = mx.load(str(bundle / "model.safetensors")) |
| mx.eval(weights) |
| weights.pop("language_model.model.layers.0.input_layernorm.weight") |
| mx.save_safetensors(str(bundle / "model.safetensors"), weights) |
| with pytest.raises(ValueError, match="Missing"): |
| checkpoint.load("mps", LoadOptions(backend="mlx")) |
|
|
